ICML 2026 Papers

A variety of widely used optimization methods like SignSGD and Muon can be interpreted as instances of steepest descent under different norm-induced geometries. In this work, we study the implicit bias of mini-batch stochastic steepest descent in multi-class classification, characterizing how batch size, momentum, and variance reduction shape the limiting max-margin behavior and convergence rates under general entry-wise and Schatten-$p$ norms. We show that, without momentum, worst-case convergence and successful classification can only be guaranteed with full-batch gradient. In contrast, momentum enables small-batch convergence to an approximate max-margin solution through a batch-momentum trade-off, though it slows convergence. This approach provides fully explicit, dimension-free rates that improve upon prior results. Moreover, we prove that variance reduction can recover the exact full-batch implicit bias for any batch size, albeit at a slower convergence rate. Finally, we further investigate the batch-size-one steepest descent without momentum, and reveal its convergence to a fundamentally different bias via a concrete data example, which reveals a key limitation of purely stochastic updates. Overall, our unified analysis clarifies when stochastic optimization aligns with full‑batch behavior, and paves the way for perform deeper explorations of the training behavior of stochastic gradient steepest descent algorithms.

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.

Parity functions are fundamental Boolean operations with critical applications across machine learning, cryptography, and error correction. Yet, learning high-dimensional parity functions poses significant challenges: in a general setting, standard neural network architectures typically require exponential sample complexity, making gradient-based optimization intractable for large number of inputs $N$. We demonstrate that compact product-based neural architectures combined with stochastic data sparsity (Bernoulli inputs with $p_e \leq 1/N$) and appropriate hyperparameter choice enable efficient parity learning, with theoretical guarantees of convergence. Experiments validate our theory across dimensions up to $N = 100{,}000$, with empirical evidence showing optimal hyperparameter choices for $p_e$ and learning rate $\alpha$, as well as polynomial complexity scaling laws. This work establishes fundamental connections between architectural inductive bias and data sparsity, opening new possibilities for neural arithmetic, structured reasoning, binary neural networks, and machine learning applied to automated protocol discovery.

We study online maximization of non-monotone Diminishing-Return(DR)-submodular functions over down-closed convex sets, a regime where existing projection-free online methods suffer from suboptimal regret and limited feedback guarantees. Our main contribution is a new structural result showing that this class is $1/e$-Upper-Linearizable under carefully designed exponential reparametrization, scaling parameter, and surrogate potential, enabling a reduction to online linear optimization. As a result, we obtain optimal $O(T^{1/2})$ static regret with a single gradient query per round and unlock adaptive and dynamic regret guarantees, together with improved rates under semi-bandit, bandit, and zeroth-order feedback. Across all feedback models, our bounds strictly improve the state of the art.

In this paper, we study a class of non-convex optimization problems known as multi-affine quadratic equality constrained problems, which appear in various applications--from generating feasible force trajectories in robotic locomotion and manipulation to training neural networks. Although these problems are generally non-convex, they exhibit convexity or related properties when all variables except one are fixed. Under mild assumptions, we prove that the alternating direction method of multipliers (ADMM) converges when applied to this class of problems. Furthermore, when the "degree" of non-convexity in the constraints remains within certain bounds, we show that ADMM achieves a linear convergence rate. We validate our theoretical results through practical examples in robotic locomotion.

Softmax attention is a central component of transformer architectures, yet its nonlinear structure poses significant challenges for theoretical analysis. We develop a unified, measure-based framework for studying single-layer softmax attention under both finite and infinite prompts. For i.i.d. centered Gaussian inputs, we lean on the fact that the softmax operator converges in the infinite-prompt limit to a linear operator acting on the underlying input-token measure. Building on this insight, we establish non-asymptotic concentration bounds for the output and gradient of softmax attention, quantifying how rapidly the finite-prompt model approaches its infinite-prompt counterpart, and prove that this concentration remains stable along the entire training trajectory in general in-context learning settings with sub-Gaussian tokens. In the case of in-context linear regression, we use the tractable infinite-prompt dynamics to analyze training at finite prompt length. Our results allow optimization analyses developed for linear attention to transfer directly to softmax attention when prompts are sufficiently long, showing that large-prompt softmax attention inherits the analytical structure of its linear counterpart. This, in turn, provides a principled and broadly applicable toolkit for studying the training dynamics and statistical behavior of softmax attention layers in large prompt regimes.

Model merging has emerged as a lightweight paradigm for enhancing Large Language Models (LLMs), yet its underlying mechanisms remain poorly understood. In this work, we analyze late-stage pre-training trajectories and uncover a \textbf{Rank-1 Subspace} phenomenon: while raw optimization steps oscillate violently, consecutive \emph{merged} checkpoints collapse onto a stable, approximately one-dimensional linear manifold. We theoretically ground this observation in a \emph{river-valley} landscape analysis: averaging acts as a geometric low-pass filter that dampens high-curvature noise to reveal the optimal descent direction. Capitalizing on this insight, we propose \textbf{Extra-Merge}, a training-free strategy that extrapolates along this subspace to minimize loss without additional gradient updates. Extensive experiments across GPT-2 and LLaMA families (124M to 2B) demonstrate that Extra-Merge consistently outperforms standard merging baselines. Notably, it yields consistent zero-shot accuracy gains on Pythia-12B downstream tasks and generalizes effectively to the Muon optimizer (Jordan et al., 2024).

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.

We study the sample complexity of the *best-case* Empirical Risk Minimizer in the setting of stochastic convex optimization. We show that there exists an instance in which the sample size is linear in the dimension, learning is possible, but the Empirical Risk Minimizer is likely to be *unique* and to *overfit*. This resolves an open question by Feldman. We also extend this to approximate ERMs. Building on our construction we also show that (constrained) Gradient Descent potentially overfits when horizon and learning rate grow w.r.t sample size. Specifically we provide a novel generalization lower bound of $\Omega\left(\sqrt{\eta T/m^{1.5}}\right)$ for Gradient Descent, where $\eta$ is the learning rate, $T$ is the horizon and $m$ is the sample size. This narrows down, exponentially, the gap between the best known upper bound of $O(\eta T/m)$ and existing lower bounds from previous constructions.

Flatness of the loss landscape has been widely studied as an important perspective for understanding the behavior and generalization of deep learning algorithms. Motivated by this view, we propose Flatness-Aware Stochastic Gradient Langevin Dynamics (fSGLD), a first-order optimization method that biases learning its dynamics toward flat basins while retaining the computational and memory efficiency of SGD and SGLD. We provide a non-asymptotic theoretical analysis showing that fSGLD converges to a flatness-biased Gibbs distribution under a theoretically prescribed coupling between the noise scale $\sigma$ and the inverse temperature $\beta$, together with explicit excess risk guarantees. We empirically evaluate fSGLD across standard optimizer benchmarks, Bayesian image classification, uncertainty quantification, and out-of-distribution detection, demonstrating consistently strong performance and reliable uncertainty estimates. Additional experiments confirm the effectiveness of the theoretically prescribed $\beta$–$\sigma$ coupling compared to decoupled choices.

Sparse PCA is one of the most well-studied problems in high-dimensional statistics. In this problem, we are given samples from a distribution with covariance $\\mathbf{\Sigma}$, whose top eigenvector $\\mathbf{v} \in \\mathbb{R}^d$ is $s$-sparse. Existing sparse PCA algorithms can be broadly categorized into (1) combinatorial algorithms (e.g., diagonal or elementwise covariance thresholding) and (2) SDP-based algorithms. While combinatorial algorithms are much simpler, they are typically only analyzed under the spiked identity model (where $\\mathbf{\Sigma} \propto \\mathbf{I}_d + \gamma \\mathbf{v}\\mathbf{v}^\top$ for some $\gamma > 0$), whereas SDP-based algorithms require no additional assumptions on $\\mathbf{\Sigma}$. We demonstrate explicit counterexample covariances $\\mathbf{\Sigma}$ against the success of standard combinatorial algorithms for sparse PCA, when moving beyond the spiked identity model. In light of this discrepancy, we give the first combinatorial method for sparse PCA that provably succeeds for general $\\mathbf{\Sigma}$ using $\\mathsf{poly}(s, \log(d))$ samples and $d^2 \cdot \\mathsf{poly}(s, \log(d))$ time, by providing a global convergence guarantee on the truncated power method of Yuan and Zhang (JMLR, 2013). We provide a natural generalization of our method to recovering $k$ sparse principal components. Finally, we evaluate our method on synthetic and real-world sparse PCA datasets.

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.

Momentum Stochastic Gradient Descent (mSGD) relies on a fixed momentum coefficient shared across all parameters, failing to account for the heterogeneous structure of modern loss landscapes. In this work, we adopt a continuous-time formulation to introduce individual, adaptive momentum coefficients regulated by the kinetic energy of each model parameter. This mechanism automatically adjusts to evolving training dynamics to maintain stability without sacrificing convergence speed. We demonstrate that this adaptive friction is inextricably linked to cubic damping, a suppression mechanism from structural dynamics. We additionally introduce two optimization schemes by augmenting the continuous dynamics of mSGD and Adam with a cubic damping term. Empirically, our methods demonstrate robustness and match or outperform Adam on training ViT, BERT, and GPT2 tasks where mSGD typically struggles. We further provide theoretical results establishing the exponential convergence of the proposed schemes.

Cross-domain class-incremental learning (CD-CIL) requires models to continuously acquire new classes across shifting domains while retaining previously learned knowledge. Existing approaches often entangle what to update with how to update, resulting in unstable adaptation and severe forgetting under domain shifts. Inspired by the hippocampal learning mechanism that separates rapid adaptation from stable consolidation, we propose Parameter-Masked Decoupled Optimization (PMDO) that disentangles what knowledge is adapted from how learning proceeds in cross-domain class-incremental learning. We introduce a domain-aware knowledge decoupler that selectively adapts domain-relevant shared parameters, constraining incremental updates while preserving prior representations. To regulate how learning proceeds, we further design a stability-aware trajectory regulation that guides optimization along transferable and stable optimization trajectories, thereby reducing interference across domain transitions. PMDO enables effective cross-domain adaptation while mitigating catastrophic forgetting and maintaining long-term learnability. Extensive experiments across multiple benchmarks demonstrate the effectiveness of PMDO and its superiority over state-of-the-art methods.

We present learning-augmented algorithms for two general classes of online minimization problems: metrical task systems and laminar set cover. Both algorithms achieve improved theoretical guarantees using machine-learned predictions of an optimal solution to the dual linear program. Unlike optimal primal solutions, which can change drastically under tiny instance perturbations, these dual solutions are much more stable, which ensures the existence of good (and learnable) predictions for families of similar instances. While previous work has used dual predictions in offline settings and for online maximization problems, our algorithms are, to the best of our knowledge, the first demonstration that such dual predictions can be effective for online minimization. Our theoretical results are complemented by experiments on the $k$-server problem and the parking permit problem.

Ill-conditioned kernel matrices can make first-order methods for training neural networks converge slowly. We establish non-asymptotic convergence bounds for the Gauss–Newton method in both under- and overparameterized regimes, showing it avoids these conditioning bottlenecks. In the underparameterized setting, Gauss–Newton gradient flow in parameter space induces a Riemannian gradient flow on a low-dimensional submanifold of function space. Using tools from Riemannian optimization, we show that, under an appropriate output scaling, the loss satisfies geodesic Polyak–Lojasiewicz and Lipschitz-smoothness conditions, implying geometric convergence to the optimal in-class predictor at an explicit rate independent of Gram-matrix conditioning. In the overparameterized setting, we identify adaptive, curvature-aware regularization schedules and prove fast geometric convergence to a global optimum for both Gauss–Newton gradient flow and discrete-time Gauss–Newton iterates, with rates independent of the minimum eigenvalue of the neural tangent kernel. Overall, Gauss–Newton can be provably faster in ill-conditioned regimes where first-order methods slow down.

One of the most pressing challenges in artificial intelligence is to make models more transparent to their users. Recently, explainable artificial intelligence has come up with numerous methods to tackle this challenge. A promising avenue is to use concept-based explanations, that is, high-level concepts instead of plain feature importance scores. Among this class of methods, Concept Activation Vectors (CAVs, Kim et al., 2018) stand out as one of the main protagonists. One interesting aspect of CAVs is that their computation requires sampling random examples from the train set. Therefore, the actual vectors obtained may vary depending on the randomness of this sampling. In this paper, we propose a fine-grained theoretical analysis of CAV construction in order to quantify their variability. Our results, confirmed by experiments on several real-life datasets of four different modalities, point to an universal result: the variance of CAVs declines roughly as $1/N$, where $N$ is the number of random examples. Based on this, we give practical recommendations for a resource-efficient application of the method.

Recent developments in deep learning optimization have brought about radically new algorithms based on the Linear Minimization Oracle (LMO) framework, such as Muon (Jordan et al., 2024b) and Scion (Pethick et al., 2025b). After over a decade of Adam’s dominance, these LMO-based methods are emerging as viable replacements, offering several practical advantages such as improved memory efficiency, better hyperparameter transferability, and most importantly, superior empirical performance on large-scale tasks, including LLM training. However, a significant gap remains between their practical use and our current theoretical understanding: prior analyses (1) overlook the layer-wise LMO application of these optimizers in practice, and (2) rely on an unrealistic smoothness assumption, leading to impractically small stepsizes. To address both, we propose a new LMO-based framework called Gluon, capturing prior theoretically analyzed methods as special cases, and introduce a new refined generalized smoothness model that captures the layer-wise geometry of neural networks, matches the layer-wise practical implementation of Muon and Scion, and leads to state-of-the-art convergence guarantees. Our experiments with NanoGPT and CNN confirm that our assumption holds along the optimization trajectory, ultimately narrowing the gap between theory and practice.

Reinforcement learning (RL) often has a hierarchical structure, where an upper-level (UL) learner selects model parameters and a lower-level (LL) decision-making process responds, naturally leading to a bilevel optimization problem. Most existing bilevel RL methods assume a single-policy LL Markov decision process (MDP), and therefore fail to capture competitive structures arising in applications such as incentive design, where multiple policies interact. We study bilevel optimization problems in which the LL problem is a regularized min–max zero-sum Markov game and the UL objective is optimized through the saddle-point equilibrium induced by the LL game. In this work, we propose penalty-augmented Nikaido–Isoda descent–ascent (PANDA), a penalty-based first-order policy-gradient method based on the Nikaido–Isoda function. By exploiting the min–max game structure, PANDA avoids computing UL hypergradients and does not require second-order information. We prove that PANDA converges to stationary points without convexity assumptions on either the UL or LL objectives. Moreover, PANDA reaches an $\epsilon$-stationary point in $\tilde{\mathcal{O}}(\epsilon^{-1})$ iterations with sample complexity $\tilde{\mathcal{O}}(\epsilon^{-3})$, matching the best-known rates for bilevel RL with single-policy LL MDPs. Experiments demonstrate the superior performance of PANDA over closely related baselines.