ICML 2026 论文列表

We introduce Context Tuning, a simple and effective method to significantly enhance few-shot adaptation of language models (LLMs) without fine-tuning model parameters. While prompt-based adaptation techniques have demonstrated the effectiveness of lightweight adaptation methods for LLMs, they typically initialize a trainable prompt or prefix with irrelevant tokens for the task at hand. In contrast, Context Tuning initializes the trainable prompt or prefix with task-specific demonstration examples, leveraging the model’s inherent In-Context Learning (ICL) ability to extract relevant information for improved few-shot learning performance. Extensive evaluations on benchmarks such as CrossFit, UnifiedQA, MMLU, BIG-Bench Hard, and ARC demonstrate that Context Tuning outperforms traditional prompt-based adaptation methods and achieves competitive accuracy with Test-Time Training at significantly higher training efficiency.

In this paper, we provide a comprehensive theoretical analysis of Stochastic Gradient Descent (SGD) and its momentum variants (Polyak Heavy-Ball and Nesterov) for tracking time-varying optima under strong convexity and smoothness. Our finite-time bounds reveal a sharp decomposition of tracking error into transient, noise-induced, and drift-induced components. This decomposition exposes a fundamental trade-off: while momentum is often used as a gradient-smoothing heuristic, under distribution shift it incurs an explicit drift-amplification penalty that diverges as the momentum parameter $\beta$ approaches 1, yielding systematic tracking lag. We complement these upper bounds with minimax lower bounds under gradient-variation constraints, proving this momentum-induced tracking penalty is not an analytical artifact but an information-theoretic barrier: in drift-dominated regimes, momentum is unavoidably worse because stale-gradient averaging forces systematic lag. Our results provide theoretical grounding for the empirical instability of momentum in dynamic settings and precisely delineate regime boundaries where vanilla SGD provably outperforms its accelerated counterparts.

The LogSumExp function, dual to the Kullback-Leibler (KL) divergence, plays a central role in many important optimization problems, including entropy-regularized optimal transport (OT) and distributionally robust optimization (DRO). In practice, when the number of exponential terms inside the logarithm is large or infinite, optimization becomes challenging since computing the gradient requires differentiating every term. We propose a novel convexity- and smoothness-preserving approximation to LogSumExp that can be efficiently optimized using stochastic gradient methods. This approximation is rooted in a sound modification of the KL divergence in the dual, resulting in a new $f$-divergence called the *Safe KL divergence*. Our experiments and theoretical analysis of the LogSumExp-based stochastic optimization, arising in DRO and continuous OT, demonstrate the advantages of our approach over existing baselines.

Automatic differentiation (AD) frameworks such as JAX and PyTorch have enabled gradient-based optimization for a wide range of scientific fields. Yet, many ''hard'' primitives in these libraries such as thresholding, Boolean logic, discrete indexing, and sorting operations yield zero or undefined gradients that are not useful for optimization. While numerous ''soft'' relaxations have been proposed that provide informative gradients, the respective implementations are fragmented across projects, making them difficult to combine and compare. This work introduces **SoftJAX** and **SoftTorch**, open-source, feature-complete libraries for *soft differentiable programming*. These libraries provide a variety of soft functions as drop-in replacements for their hard JAX and PyTorch counterparts. This includes (i) elementwise operators such as *clip* or *abs*, (ii) utility methods for manipulating Booleans and indices via fuzzy logic, (iii) axiswise operators such as *sort* or *rank* -- based on optimal transport or permutahedron projections, and (iv) offer full support for straight-through gradient estimation. Overall, SoftJAX and SoftTorch make the toolbox of soft relaxations easily accessible to differentiable programming, as demonstrated through benchmarking and a practical case study. Code is available at github.com/a-paulus/softjax and github.com/a-paulus/softtorch.

We propose a mathematically principled PDE gradient flow framework for distributionally robust optimization (DRO). Exploiting the recent advances in the intersection of Monte Carlo sampling and statistical optimal transport, we show that our theoretical framework can be implemented as practical algorithms for sampling from worst-case distributions and, consequently, DRO. While numerous previous works have relied on dual reformulation techniques, we contribute a sound and complete gradient flow view based on SDEs or PDEs that can be used to construct new algorithms for general, potentially non-convex, losses. Without loss of generality, we solve a class of Wasserstein and entropy-regularized DRO problems using the recently-discovered Wasserstein Fisher-Rao and Stein variational gradient flows. Notably, we also show some simple reductions of our framework recover exactly previously proposed popular DRO methods, and provide new insights into their theoretical limits and optimization dynamics of DRO. Numerical studies based on stochastic gradient descent on machine learning tasks provide empirical backing for our theoretical findings.

Federated stochastic multi-level compositional minimax optimization supports a growing number of machine learning applications. However, the interplay of multi-level compositional structure and minimax formulation in federated learning setting poses significant optimization challenges, resulting in slow convergence rates for existing algorithms. In this paper, we propose a novel federated learning algorithm, LS$^2$MC-GDA, that leverages smoothing techniques and variance-reduced stochastic compositional gradients. To support our theoretical analysis, we introduce a stage-wise extension of LS$^2$MC-GDA, which serves to bridge the gap between different stationarity measures. We establish that our algorithm achieves a sample complexity of $O(\kappa^{3/2}/N\epsilon^3)$ and a communication complexity of $O(\kappa/\epsilon^2)$, substantially improving existing theoretical results in terms of the condition number $\kappa$ and the solution accuracy $\epsilon$ and achieving a linear speedup with respect to the number of workers $N$. Finally, experimental results validate the effectiveness of our algorithm.

Deep Reinforcement Learning (DRL) has emerged as a promising approach for solving Combinatorial Optimization (CO) problems, such as the 3D Bin Packing Problem (3D-BPP), Traveling Salesman Problem (TSP), or Vehicle Routing Problem (VRP), but these neural solvers often exhibit brittleness when facing distribution shifts. To address this issue, we uncover the Satisficing Generalization Edge, which we validate both theoretically and experimentally: identifying a set of promising actions is inherently more generalizable than selecting the single optimal action. To exploit this property, we propose Adaptive Selection After Proposal (ASAP), a generic framework that decomposes the decision-making process into two distinct phases: a proposal policy that acts as a robust filter, and a selection policy as an adaptable decision maker. This architecture enables a highly effective online adaptation strategy where the selection policy can be rapidly fine-tuned on a new distribution. Concretely, we introduce a two-phase training framework enhanced by Model-Agnostic Meta-Learning (MAML) to prime the model for fast adaptation. Extensive experiments on 3D-BPP, TSP, and CVRP demonstrate that ASAP improves the generalization capability of state-of-the-art baselines and achieves superior online adaptation on out-of-distribution instances.

Recently, many empirical work has shown that, in machine learning, the noise distribution of stochastic gradients often exhibits heavy tails when stochastic optimization methods are employed. Most existing theoretical analyses of heavy-tailed stochastic methods rely on various convexity and smoothness assumptions and our knowledge of how heavy-tailed stochastic methods behave in the setting of weakly convex optimization is still limited. In the weakly convex setting, this paper derives new upper bounds on the convergence of the stochastic gradient method (SGD) under heavy-tailed noises. In particular, for vanilla SGD, we establish an in-expectation convergence guarantee on the bounded constrained domain under the assumption of bounded $p$-th central moment ($p$-BCM) of the gradient noise, and a high-probability guarantee on the unbounded domain when the noise follows a heavy-tailed sub-Weibull distribution. By equipping SGD with the gradient clipping (Clip-SGD), we demonstrate that it achieves high-probability convergence in the unbounded domain under the $p$-BCM gradient noise. All of our high-probability convergence bounds depend on the failure probability only through polynomial-logarithmic factors. Finally, we present experiments to validate our theoretical findings.

Decentralized stochastic bi-level optimization has been actively studied in recent years. However, existing studies assume that the lower-level loss function is strongly convex, which limits their applicability to many machine learning models. To address this limitation, in this paper, we propose a novel decentralized stochastic first-order optimization algorithm, which does not require second-order Hessian or Jacobian matrices, for the setting where the lower-level loss function is nonconvex but satisfies the Polyak–Łojasiewicz (PL) condition. Additionally, unlike existing single-agent methods that introduce a regularization term to the lower-level loss function to artificially enforce strong convexity, our algorithm does not require such modification. Moreover, our algorithm employs a constant single-timescale learning rate for updating variables, which is different from the time-dependent and two-timescale learning rate schedules used in prior work. To establish the convergence rate, we develop a new convergence analysis framework for the pure PL condition, rather than relying on the artificial strong convexity introduced through regularization in existing single-agent methods. To the best of our knowledge, this is the first algorithm for nonconvex decentralized bi-level optimization that offers theoretical convergence guarantees under mild conditions. Finally, our extensive experimental results on hyperparameter optimization and model pruning applications validate the efficacy of the proposed algorithm.

Distributionally Robust Optimization (DRO) is widely used to improve model robustness, with existing methods addressing either geometric perturbations (e.g., input shifts) or statistical contamination (e.g., heavy-tailed noise and outliers) effectively. However, these uncertainty sources often co-exist, and coupling them through a single divergence or optimal transport constraint conflates geometric displacement with loss-based outlierness, frequently discarding informative high-leverage samples. We introduce nested DRO, a bilevel formulation that combines an outer optimistic $\phi$-divergence cleaning step with an inner pessimistic optimal-transport robustification step, thereby decoupling geometric smoothing from statistical cleaning. We prove that this structure naturally induces a geometry-invariant, loss-based reweighting mechanism that separates outlier suppression from transport-induced regularization. We derive a tractable strong dual for the resulting non-convex problem and show its equivalence to variance-regularized risk minimization, leading to a clear statistical interpretation of the induced weights. Empirical results on synthetic and real datasets demonstrate that nested DRO consistently outperforms geometry-coupled DRO baselines, particularly under heavy-tailed contamination where preserving high-leverage structure is crucial.

Bayesian optimization (BO) selects evaluation points for expensive black-box objectives using Gaussian process (GP) predictive distributions. Kernel choice and hyperparameter selection can lead to miscalibrated predictive distributions and an inappropriate exploration--exploitation trade-off. For minimization, sampling criteria such as expected improvement (EI) depend on the predictive distribution below the current best value, so lower-tail miscalibration directly affects the sampling decision. This article studies goal-oriented calibration of GP predictive distributions below a low threshold $t$ in the noiseless setting, for standard GP models with hyperparameters selected by maximum likelihood. A framework for predictive reliability below $t$ is introduced, based on two notions of spatial calibration: occurrence calibration over the design space and thresholded $\mu$-calibration on sublevel sets of the form $\lbrace x\in\mathbb{X}, f(x)\le t \rbrace$. Building on this framework, we propose tcGP, a post-hoc method that calibrates GP predictive distributions below $t$, and we show that the resulting EI-based global optimization algorithm remains dense in the design space. Experiments on standard benchmarks show improved lower-tail calibration and BO performance relative to standard GP models and globally calibrated GP models.

The training of neural networks often entails objective functions that are not globally $L$-smooth. For these functions, it is both theoretically and practically difficult to reply to the question: what is the largest possible step size that ensures the convergence of gradient descent (GD)? We address this longstanding open question in deep learning by providing a unifying definition of "large'' step sizes that requires only local Lipschitz (or even Hölder) continuity of the gradient. We design first-order adaptive methods that provably yield large step sizes and show that they operate at the edge of stability (EoS) right from the start of the training. In particular, the loss decreases nonmonotonically and the product between the step size and sharpness, i.e., the largest eigenvalue of the hessian, stays above the EoS threshold of 2 throughout training. Using our method, we are also able to minimize the sharpness all the way down to its global minimum. Contrary to expectation, we find that encountering globally-flat regions too early in the training may both slow down convergence and jeopardize the generalization ability of the network. Exploiting a self-stabilization argument, we allow GD to enter slightly sharper valleys and turn unsuccessful training runs into very successful ones.

Modern neural networks often achieve high accuracy on their training distribution but degrade sharply under distribution shifts. We address this problem through *Robust Satisficing* (RS), an optimization objective that seeks parameters which attain a target level of in-distribution performance while minimizing *fragility*, defined as the rate at which performance deteriorates as the data distribution departs from training. We develop a gradient-based algorithm, *Iterative Robust Satisficing* (IRS), that directly optimizes this criterion. Across a range of synthetic and real-world distribution shifts, including long-tailed image classification, group shifts induced by spurious correlations, and natural shifts in tabular regression, IRS consistently improves performance on minority and worst-case groups without sacrificing overall accuracy. Notably, IRS achieves these robustness gains with a per-step computational cost similar to standard stochastic gradient descent and requires only a single forward and backward pass per update. Together, these results suggest that minimizing fragility provides a practical and effective alternative to existing robust training methods for learning models that remain reliable under distribution shift.

Sparse tangent portfolio optimization aims to learn an interpretable, low-cardinality portfolio in the tangency direction of the mean-variance frontier. However, the associated cardinality-constrained formulation is NP-hard, and standard predict-then-optimize pipelines often misalign forecasting accuracy with downstream portfolio quality. We propose an end-to-end decision-focused learning framework that reformulates Sharpe ratio maximization as a Disciplined Parametrized Programming (DPP)-compliant convex programming layer and replaces discrete selection with a smooth top-$k$ operator enforcing an exact cardinality $k$. This enables gradient flow through prediction, asset selection, and re-optimization, allowing the predictive model to directly optimize portfolio performance. Across four major equity markets, our method achieves competitive and often superior out-of-sample Sharpe ratios compared with historical and prediction-focused baselines, with particularly strong gains in larger asset universes. Our code is publicly available.

This paper considers decentralized optimization of convex functions with mixed affine equality constraints involving both local and global variables. Constraints on global variables may vary across different nodes in the network, while local variables are subject to coupled and node-specific constraints. Such problem formulations arise in machine learning applications, including federated learning and multi-task learning, as well as in resource allocation and distributed control. We analyze this problem under smooth and non-smooth assumptions, considering both strongly convex and general convex objective functions. Our main contribution is an optimal algorithm for the smooth, strongly convex regime, whose convergence rate matches established lower complexity bounds. We further provide optimal and near-optimal methods for the remaining cases.

In distributed machine learning, efficiently training across multiple agents with heterogeneous data distributions remains a central challenge. We address the problem of stochastic, strongly convex distributed optimization by applying accelerated gradient ascent to the dual variables and multi-step stochastic gradient descent (SGD) to the primal variables in the Lagrangian formulation. This approach naturally enables local computation, as the inner SGD loops require no inter-agent communication. We prove that the method converges for any number of local updates, attaining the optimal communication complexity when local computation is sufficient. Our analysis builds on an inexact accelerated gradient framework, where the partial gradient of the Lagrangian with respect to the dual variables is treated as an inexact gradient of the dual function. A notable byproduct of this framework is an algorithm that achieves optimal reproducibility guarantees under biased gradient estimates.

Optimal-transport distributionally robust optimization (OT-DRO) robustifies data-driven decision-making under uncertainty by capturing the sampling-induced statistical error via optimal transport ambiguity sets. The standard OT-DRO pipeline consists of a two-step procedure, where the ambiguity set is first designed and subsequently embedded into the downstream OT-DRO problem. However, this separation between uncertainty quantification and optimization may lead to excessive conservatism. We introduce an end-to-end pipeline to automatically learn decision-focused ambiguity sets for OT-DRO problems, where the loss function informs the shape of the ambiguity set, leading to less conservative decisions whose distributional robustness is enforced via data-driven bootstrapping. We formulate the learning problem as a bilevel optimization program and solve it via a hypergradient-based method. By leveraging the recently introduced nonsmooth conservative implicit function theorem, we establish convergence to a critical point of the bilevel problem. We present experiments validating our method on standard portfolio optimization and linear regression tasks.

We propose **GradPower**, a lightweight gradient-transformation technique for accelerating language model pre-training. Given a gradient vector $\boldsymbol{g}=(g_ {i})_ {i}$, GradPower first applies the elementwise `sign-power` transformation: $ \varphi_ p(\boldsymbol{g}) = \left({\rm sign}(g_ i)|g_ i|^p\right)_ {i} $ for a fixed $p>0$, and then feeds the transformed gradient into a base optimizer. Notably, GradPower requires only a **single-line code change** and no modifications to the base optimizer’s internal logic, including the hyperparameters. When applied to AdamW (termed **AdamWPower**), GradPower consistently achieves lower terminal loss across diverse architectures (LLaMA, Qwen2MoE), parameter scales (66M to 2B), datasets (C4, OpenWebText), and learning-rate schedules (cosine, warmup-stable-decay). The most pronounced gains are observed when training modern mixture-of-experts models with warmup-stable-decay schedules. GradPower also integrates seamlessly with other state-of-the-art optimizers, such as Muon, yielding further improvements. Finally, we provide theoretical analyses that reveal the underlying mechanism of GradPower and highlight the influence of gradient noise.

Antimicrobial peptide discovery is challenged by the astronomical size of peptide space and the relative scarcity of active peptides. While generative models provide latent maps of this space, they typically ignore decoder-induced geometry and rely on flat Euclidean metrics, making exploration distorted and inefficient. Existing manifold-based approaches assume fixed intrinsic dimensionality, which fails for real peptide data. We introduce **PepCompass**, a geometry-aware framework based on a **Union of $\kappa$-Stable Riemannian Manifolds** that captures local decoder geometry while maintaining computational stability. PepCompass performs global interpolation via **Potential-minimizing Geodesic Search (PoGS)** to bias discovery toward promising seeds and enables local exploration through **Second-Order Riemannian Brownian Efficient Sampling** and **Mutation Enumeration in Tangent Space**, which together form **Local Enumeration Bayesian Optimization (LE-BO)**. PepCompass achieves a 100% *in-vitro* validation rate: PoGS identifies four novel seeds and LE-BO optimizes them into 25 highly active, broad-spectrum peptides, demonstrating that geometry-informed exploration is a powerful paradigm for antimicrobial peptide design.