ICML 2026 Papers

The stochastic Polyak step size (SPS) has proven to be a promising choice for stochastic gradient descent (SGD), delivering competitive performance relative to state-of-the-art methods on smooth convex and non-convex optimization problems, including deep neural network training. However, extensions of this approach to non-smooth settings remain in their early stages, often relying on interpolation assumptions or requiring knowledge of the optimal solution. In this work, we propose a novel SPS variant, Safeguarded SPS (SPS$_{safe}$), for the stochastic subgradient method, and provide rigorous convergence guarantees for non-smooth convex optimization with no need for strong assumptions. We further incorporate momentum into the update rule, yielding equally tight theoretical results. Comprehensive experiments on convex benchmarks and deep neural networks corroborate our theory: the proposed step size achieves competitive performance to existing adaptive baselines and exhibits stable behavior across a wide range of problem settings. Finally, in the context of deep neural network training, the gradient norms under our step size do not collapse to (near) zero, indicating robustness to vanishing gradients.

In-context learning (ICL) is a practical way to adapt large models, yet under strict context limits, it remains unclear how to spend scarce tokens without being misled by noisy, redundant, or conflicting demonstrations. We address this gap by targeting token-budgeted context construction: how to select and compress demonstrations so the prompt carries maximal task-relevant signal with minimal predictive distortion. We propose RDCO, a deterministic, training-free optimizer that scores demonstrations by marginal task information per token, penalizes redundancy and prefix-conditioned conflicts, and finally compacts the selected context under a bounded predictive-divergence constraint to control drift. Across a 10-dataset ICL suite spanning classification and structured generation, RDCO achieves the best average performance (63.26 Acc. on classification and 60.26 EM on generation), improving over the strongest classification and generation baselines by 2.20 and 2.26 points, respectively, and improving the 10-task overall average by 4.94 points under the same budget. Our results suggest that viewing prompts as finite-capacity messages yields a principled and effective path to more reliable and token-efficient ICL.

Differentiable optimization studies how to embed a mathematical program as a differentiable layer in machine learning pipelines. However, existing approaches typically rely on implicit differentiation, involving expensive Hessian computation while differentiating through optimality conditions. To address this challenge, we formulate the differentiable optimization problem as a bilevel optimization instance. We construct a new active-set Lagrangian as a proxy to compute an $\epsilon$-approximate hypergradient using only near-constant $O(\log (1/\epsilon))$ first-order information. We also show that applying this efficient hypergradient oracle to constrained bilevel optimization improves the overall gradient complexity to $\tilde{O}(\delta^{-1}\epsilon^{-3})$ to reach a $(\delta, \epsilon)$-Goldstein stationary point. We implement our method `FFOLayer`, as a drop-in Python library compatible with existing differentiable optimization solvers. Our algorithm shows significantly faster computation with similar convergence compared to other existing solvers. The source code is available at [https://github.com/GT-KOALA/FFOLayer](https://github.com/GT-KOALA/FFOLayer).

The Traveling Salesperson Problem (TSP) has long served as a benchmark for evaluating the strength of optimization techniques in the classical theory of algorithms. In recent efforts to apply ML to algorithmic problems, TSP has also become a natural testbed for the development of ML-based techniques. A common approach is to train a neural network to output a heatmap estimating the likelihood of each edge to be part of the optimal tour; however, converting such a heatmap into an actual tour remains a non-trivial and often computationally intensive step. In this work, we propose algorithms for transforming heatmaps into tours with theoretical guarantees linking the achieved approximation ratio to the quality of the provided heatmap. In the spirit of *algorithms with predictions*, our results can be described as $(1+2\eta/OPT)$-approximation algorithms, where $\eta$ denotes the L1 distance between the prediction (heatmap) and an optimal solution (tour). Since the previous works lack such explicit guarantees, we compare our approach against them experimentally.

This paper proposes asymmetric perturbation, where only one player's payoff function is perturbed, for solving bilinear saddle-point optimization problems, commonly arising in minimax problems, game theory, and constrained optimization. Symmetric perturbation is known to require decreasing its strength to ensure convergence to a solution, i.e., an equilibrium in the original game, resulting in a slower rate. First, with asymmetric perturbation, we show that, for a sufficiently small perturbation strength, the equilibrium strategy of the asymmetrically perturbed game coincides with an equilibrium strategy of the original unperturbed game. Second, building on this coincidence, we construct a learning algorithm with a linear last-iterate convergence rate. Third, motivated by the fact that the coincidence relies on the perturbation strength being sufficiently small, we also provide a parameter-free variant, retaining the linear rate. Finally, we empirically demonstrate fast convergence toward equilibria in both normal-form and extensive-form games.

Credal predictors are epistemic-uncertainty-aware models that produce a convex set of probabilistic predictions. They provide a principled framework for quantifying predictive epistemic uncertainty (EU) and have been shown to improve model robustness across a range of settings. However, most state-of-the-art (SOTA) methods primarily define EU as disagreement induced by random training initializations, which mainly reflects sensitivity to optimization randomness rather than uncertainty from more substantive sources. In response, we formulate EU as disagreement between models trained under different degrees of relaxation of the i.i.d. assumption between the training and test distributions. Building on this idea, we propose *CreDRO*, which learns an ensemble of plausible models via distributionally robust optimization. As a result, CreDRO captures EU arising not only from training randomness but also from informative disagreement due to potential train–test distribution shifts. Empirically, CreDRO consistently outperforms SOTA credal approaches on downstream tasks, including out-of-distribution detection on extensive benchmarks and selective classification in medical settings.

Modern optimizers combine gradients from the current mini-batch with historical optimization state, such as momentum or adaptive moments. While highly effective, aggregating across the batch and incorporating this history can produce parameter updates that increase the loss of individual samples. We term this effect *harm* and formalize the parameter update as an optimization problem that explicitly minimizes the conflicting impact of both batch averaging and past optimization state on current data. Because the exact formulation is intractable, we introduce a highly efficient proxy. We first reduce the problem's dimensionality to the batch size, and then drastically cut memory and speed bottlenecks by successfully restricting the optimization to the last linear layer. This hinges on the unexpected finding that *this layer alone reliably captures the second-order statistics of the per-sample gradients*. The resulting surrogate problem integrates readily into standard optimizers like SGD and AdamW, and can be solved using a small number of GPU-friendly iterations. Crucially, the method exhibits *favorable scaling properties*, as the relative computational overhead shrinks as the model size or input grows. Experiments on image classification benchmarks confirm *reduced per-sample interference* and *improved generalization*.

We study optimization over non-convex constraint sets that are homeomorphic to a ball, encompassing important problem classes such as star-shaped sets that frequently arise in machine learning and engineering applications. We propose **Hom-PGD$^+$**, a learning-based and projection-efficient first-order method that efficiently solves such problems without requiring expensive projection or optimization oracles. Our approach leverages an invertible neural network (INN) to learn the homeomorphism between the non-convex constraint set and a unit ball, transforming the original problem into an equivalent ball-constrained optimization where projections admit efficient solutions. We establish that Hom-PGD$^+$ achieves an $\mathcal{O}(\epsilon^{-2})$ convergence rate to an ($\epsilon + \mathcal{O}(\sqrt{\epsilon_{\text{inn}}})$)-approximate stationary solution, where $\epsilon_{\text{inn}}$ denotes the homeomorphism learning error. This rate significantly improves upon existing methods for optimization over non-convex sets, while maintaining a per-iteration complexity of only $\mathcal{O}(W)$ for $W$ INN parameters. Extensive experiments, including QCQP, chance-constrained power-system optimization, and non-uniform adversarial attacks, demonstrate that Hom-PGD$^+$ achieves competitive solution quality while delivering speedups of up to one order of magnitude.

We analyze algorithms for solving stochastic variational inequalities (VI) without the bounded variance or bounded domain assumptions, where our main focus is min-max optimization with possibly unbounded constraint sets. We focus on two classes of problems: monotone VIs; and structured nonmonotone VIs that admit a solution to the *weak Minty VI*. The latter assumption allows us to solve structured nonconvex-nonconcave min-max problems. For both classes of VIs, to make the expected residual norm less than $\varepsilon$, we show an oracle complexity of $\widetilde{O}(\varepsilon^{-4})$, which is the best-known for constrained VIs. In our setting, this complexity had been obtained with the bounded variance assumption in the literature, which is not even satisfied for bilinear min-max problems with an unbounded domain. We obtain this complexity for stochastic oracles whose variance can grow as fast as the squared norm of the optimization variable.

Many learning tasks in science and engineering are characterized by sparse datasets, which limits the effectiveness of purely data-driven approaches. At the same time, these problems are often accompanied by rich domain knowledge derived from physical laws, operational requirements, and expert heuristics. Such knowledge is frequently expressed as rules involving logical propositions and linear inequalities. Existing neuro-symbolic methods typically enforce these rules approximately through soft penalties, assume input-independent rules when designing specialized architectures, or rely on non-differentiable post-processing at inference time to achieve hard constraint satisfaction. While recent advances in differentiable optimization layers enable end-to-end feasibility enforcement within neural networks, extending these approaches to logical or mixed-integer rules remains challenging due to inherent nonconvexity. In this work, we propose a unified end-to-end framework for enforcing hard, input-dependent mixed integer linear constraints within neural networks. Our approach represents rules as disjunctive constraints and applies hierarchical convex relaxations to obtain convex hull formulations. These relaxations yield tractable linear constraints that can be embedded as differentiable optimization layers while enabling exact rule satisfaction. We demonstrate the effectiveness of the proposed framework on real-world datasets, achieving perfect rule satisfaction and strong predictive performance.

We introduce an order-invariant reinforcement learning framework for black-box combinatorial optimization. Classical estimation-of-distribution algorithms (EDAs) often rely on learning explicit variable dependency graphs, which can be costly and may fail to capture complex interactions efficiently. In contrast, we parameterize a multivariate autoregressive generative model trained without a fixed variable ordering. By sampling random generation orders during training, a form of information-preserving dropout, the model is encouraged to be invariant to variable order, promoting search-space diversity, and shaping the model to focus on the most relevant variable dependencies, improving sample efficiency. We adapt Group Relative Policy Optimization (GRPO) to this setting, providing stable policy-gradient updates from scale-invariant advantages. Across a wide range of benchmark problem instances of varying sizes, our method frequently achieves the best performance and consistently avoids catastrophic failures.

In this work, we revisit Transformer optimization through the lens of second-order geometry and establish a direct connection between architectural design, activation scale, the Hessian matrix, and the maximum tolerable learning rate. We introduce a simple normalization strategy, termed SimpleNorm, which stabilizes intermediate activation scales by construction. Then, by analyzing the Hessian of the loss with respect to network activations, we theoretically show that SimpleNorm significantly reduces the spectral norm of the Hessian, thereby permitting larger stable learning rates. We validate our theoretical findings through extensive experiments on large GPT models at parameter scales 1B, 1.4B, 7B and 8B. Empirically, SimpleGPT, our SimpleNorm-based network, tolerates learning rates 3$\times$-10$\times$ larger than standard convention, consistently demonstrates strong optimization stability, and achieves substantially better performance than well-established baselines. Specifically, when training 7B-scale models for 60K steps, SimpleGPT reduces the training loss from 2.290 to 2.208 compared with Llama2 with QKNorm. Our code is available at \url{https://github.com/Ocram7/SimpleGPT}.

Spectrally sparse signal reconstruction arises in a wide range of applications and can be formulated as a low-rank Hankel matrix completion problem. We develop a Jacobi-preconditioned gradient descent method that preserves the low per-iteration complexity of first-order algorithms while achieving linear convergence at a rate independent of the condition number. By introducing a generator that maps factor-based iterates to matrix space, we establish equivalence with manifold-based methods, enabling direct convergence analysis while avoiding the need to define distances under complex-symmetric factorization ambiguity. Extensive experiments demonstrate that the proposed algorithm outperforms state-of-the-art methods in both iteration count and computational time across a broad range of problem settings.

We introduce a constrained optimization framework for training transformers that behave like optimization descent algorithms. Specifically, we enforce layerwise descent constraints on the objective function and replace standard empirical risk minimization (ERM) with a primal-dual training scheme. This approach yields models whose intermediate representations decrease the loss monotonically in expectation across layers. We apply our method to both unrolled transformer architectures and conventional pretrained transformers on tasks of video denoising and text classification. Across these settings, we observe that constrained transformers achieve stronger robustness to perturbations and maintain higher out-of-distribution generalization, while preserving competitive in-distribution performance.

Sufficient dimension reduction (SDR) makes high-dimensional regression tractable by projecting the covariates onto a low-dimensional subspace that preserves the conditional mean of the response. Existing gradient-based estimators either operate in the ambient space and suffer from the curse of dimensionality, or localize in the reduced space at a per-outer-iteration cost at least quadratic in the sample size. We show that minimizers of the population Minimum Average Variance Estimation (MAVE) risk approximate the same Grassmannian target as the Outer Product of Gradients (OPG), and recast the empirical criterion as a smooth maximization on the Stiefel manifold with closed-form Riemannian gradient. The resulting algorithm, SMAVE, combines sparse projected-space nearest-neighbor localization with Riemannian stochastic gradient ascent. A simplified version comes with almost-sure convergence and a non-asymptotic rate matching the standard non-convex stochastic first-order scaling. Empirically, SMAVE matches or improves on RMAVE's synthetic subspace recovery at moderate-to-high ambient dimension, and on four real datasets it uniformly improves over OPG and is competitive with or outperforms RMAVE at orders of magnitude lower runtime.

Learning rate warm-up -- increasing the learning rate at the beginning of training -- has become a ubiquitous heuristic in modern deep learning, yet its theoretical foundations remain poorly understood. In this work, we provide a principled explanation for why warm-up improves training. We rely on a generalization of the $(L_0, L_1)$-smoothness condition, which bounds local curvature as a linear function of the loss suboptimality and exhibits desirable closure properties. We show -- both theoretically and empirically -- that this condition is satisfied by common neural architectures and accurately captures the curvature of the optimization landscape early in training. Adapting the learning rate in response to this curvature condition naturally induces a warm-up–like schedule, and we show that this choice yields provably faster convergence guarantees than using a fixed learning rate. Experiments on language and vision models show that the resulting one-parameter warm-up schedule can match tuned linear warm-up and improve over no warm-up.

Speculative decoding accelerates autoregressive large language model (LLM) inference by using a lightweight draft model to propose candidate tokens that are then verified in parallel by the target model. The speedup is significantly determined by the acceptance rate, yet standard training minimizes Kullback-Leibler (KL) divergence as a proxy objective. While KL divergence and acceptance rate share the same global optimum, small draft models, having limited capacity, typically converge to suboptimal solutions where minimizing KL does not guarantee maximizing acceptance rate. To address this issue, we propose **LK losses**, special training objectives that directly target acceptance rate. Comprehensive experiments across four draft architectures and six target models, ranging from 8B to 685B parameters, demonstrate consistent improvements in acceptance metrics across all configurations compared to the standard KL-based training. We evaluate our approach on general, coding and math domains and report gains of up to 8-10\% in average acceptance length. LK losses are easy to implement, introduce no computational overhead and can be directly integrated into any existing speculator training framework, making them a compelling alternative to the existing draft training objectives.

The DeepWalk algorithm has been widely used for learning node embeddings in graphs. Combined with the idea of _negative sampling_, the DeepWalk algorithm has been shown to be implementable at scale, easily handling graphs with millions of nodes. However, theoretical guarantees on the resulting embeddings are much less understood. Recent results have studied the minimizers of the objective and have shown interesting guarantees for certain graph classes. However, the optimization _trajectory_, i.e., what happens when we start at a random initialization and run gradient descent, remains poorly understood. This is especially true for the implementation of DeepWalk using Skip-gram with negative sampling (SGNS), since the variance of the stochastic updates turns out to be very large. In this work, we make progress on this question. We show that for "small norm" initialization, under a spectral gap assumption on the graph, the DeepWalk embeddings align with the column space of a fixed low-rank matrix. For graphs generated from Stochastic Block Models with certain separation conditions, our results imply that the DeepWalk embeddings recover cluster structure. To the best of our knowledge, our results give the first analysis of the optimization trajectory of DeepWalk with negative sampling on non-trivial graph classes.

Safety alignment of large language models is fragile: even small fine-tuning perturbations elastically revert behaviors toward those of the pre-training, with degradation inversely proportional to the size of the alignment set. We ask how to achieve safety alignment with minimal augmentation. To this end, we model augmentation as a set of group actions on sequences and formalize robustness gains as a normalized, monotone submodular function over transformations. We then leverage submodular optimization to select minimal augmentations that provably improve robustness. Experiments confirm that our approach efficiently restores safety alignment while minimizing the overhead of augmentation.