NeurIPS 2026 Papers

Bilevel optimization is a core tool in practical machine learning workflows, including hyperparameter tuning, data reweighting, robust learning, reinforcement learning, and imbalanced classification. Hypergradient computation requires the inverse of the inner Hessian, typically estimated via a truncated Neumann series. While existing analyses focus on bias and variance, we show that the product structure of the Neumann series introduces a third source of error: skewness. Unlike ordinary gradients, whose asymmetry is suppressed by mini-batch averaging, the Neumann estimator amplifies skewness across factors, with the third-moment error growing with truncation depth, inner condition number, and per-factor noise asymmetry. This arises in many practical settings such as corrupted labels, outliers, heterogeneous or class-imbalanced datasets. To address this, we introduce the Antithetic Neumann Estimator (ANE), which pairs each Hessian sample with a mirrored counterpart around a reference Hessian and averages the resulting inverse-Hessian estimates. With an exact reference, the antithetic product retains only even-order noise, eliminating all skewness; with an approximate reference, the residual skewness scales linearly with the reference error. ANE requires only a modification of the Hessian-vector step and integrates with existing stochastic or variance-reduced bilevel methods. We instantiate ANE-SGD and ANE-VR, the latter adding recursive variance reduction. Across robust regression, corrupted-label learning, offline reinforcement learning, and imbalanced classification, ANE-based methods achieve strictly better bias–variance–skewness tradeoffs than stochastic and variance-reduced baselines, reducing loss by up to 24% and improving predictive metrics by up to 27%. This establishes skewness correction as a practical and complementary mechanism for improving hypergradient quality.

Pass@$k$ is a widely used performance metric for verifiable large language model tasks, including mathematical reasoning, code generation, and short-answer reasoning. It defines success if any of $k$ independently sampled solutions passes a verifier. This multi-sample inference metric has motivated inference-aware fine-tuning methods that directly optimize pass@$k$. However, prior work reports a recurring trade-off: pass@$k$ improves while pass@1 degrades under such methods. This trade-off is practically important because pass@1 often remains a hard operational constraint due to latency and cost budgets, imperfect verifier coverage, and the need for a reliable single-shot fallback. We study the origin of this trade-off and provide a theoretical characterization of when pass@$k$ policy optimization can reduce pass@1 through gradient conflict induced by prompt interference. We show that pass@$k$ policy gradients can conflict with pass@1 gradients because pass@$k$ optimization implicitly reweights prompts toward low-success prompts; when these prompts are what we term \emph{negatively interfering}, their upweighting can rotate the pass@k update direction away from the pass@1 direction. We illustrate our theoretical findings with large language model experiments on verifiable mathematical reasoning and coding tasks.

Double Oracle (DO) and PSRO approximate Nash equilibria in two-player zero-sum games by iteratively expanding a strategy population with best responses. Online Double Oracle (ODO) interleaves multiplicative weights updates (MWU) with discovery and obtains an $O(\sqrt{k \ln k / T})$ rate. We identify the \emph{reinitialisation} step as an underexplored design choice: at each expansion, MWU-based methods reset newly discovered strategies to a uniform prior. We introduce Virtual Double Oracle (VDO), which exploits a bilinear identity---a strategy's full hindsight loss history collapses to its utility against the time-averaged opponent---to initialise each new strategy with the MWU state it would have had from the start. Conditional on the discovered population, VDO recovers the fixed-population rate $O(\sqrt{\ln k / T})$ asymptotically, removing ODO's leading-order $\sqrt{k}$ window-decomposition penalty. Discounted VDO adds a single parameter $\beta$ to control finite-horizon prior strength. Experiments across matrix games, poker, and Goofspiel validate both contributions: pure VDO confirms the asymptotic rate, and discounted VDO gives the strongest finite-horizon performance.

We introduce a variational framework for diffusion models with anisotropic noise schedules with a matrix-valued path $M_t(\theta)$ that allocates noise across subspaces. Central to our framework is a trajectory-level objective that jointly trains the score network and learns $M_t(\theta)$, which encompasses general parameterization classes of matrix-valued noise schedules. We further derive an estimator for $\partial_\theta \nabla \log p_t$ that enables efficient optimization of the $M_t(\theta)$ schedule. For inference, we develop an efficiently-implementable reverse-ODE solver that is an anisotropic generalization of the second-order Heun discretization algorithm. Across CIFAR-10, AFHQv2, FFHQ, and ImageNet-64, our method consistently improves upon the baseline EDM model in all NFE regimes.

Many scientific design problems require identifying rare high-fitness objects in vast discrete spaces using costly black-box feedback. Variational active-generation methods, including variational search distributions (VSD) and conditioning by adaptive sampling (CbAS), address this by learning a distribution over high-fitness designs via inference on a level-set posterior. However, these methods rely on generators with explicit likelihoods, most commonly autoregressive models, creating a mismatch with modern non-autoregressive approaches, such as discrete diffusion and flow-based models, which refine designs in parallel and better capture multi-modal structure but have implicit marginal likelihoods. We introduce \textit{Active Flow Matching} (AFM), a variational active-generation framework that removes this likelihood bottleneck by reformulating level-set KL objectives over the conditional endpoint distributions of discrete flow models. AFM yields forward-, reverse-, and symmetric-KL objectives without requiring density evaluation. We show that forward-KL AFM is target-consistent, recovering the level-set posterior without access to marginal likelihoods. Empirically, across protein and small-molecule design tasks, AFM outperforms strong baselines, including autoregressive VSD and CbAS, inference-time guidance, KL-regularised fine-tuning, and Top-$K$ retraining.

Modern language models often need to optimize a primary accuracy objective while also accommodating secondary behavioral preferences, such as verbosity, agreeableness, or the level of technical expertise in its response. In practice, a base model may exhibit a desired behavior very rarely or not at all. Thus, endowing the model with a target behavior creates a sparse behavioral reward bottleneck. To address such multi-objective problems, we introduce Vector-Steered Policy Optimization (VSPO) which employs a steering vector associated with the target behavior to control the behavior intensity of the generated rollouts. VSPO is obtained by modifying GRPO to sample rollouts with varying steering intensities. This process can be interpreted as an on-policy latent self-distillation procedure where the model internalizes its steering vector. By varying steering intensities, VSPO upsamples rare behaviors and enriches rollout diversity, which alleviates the sparse reward issue and provably accelerates the policy optimization. Through comprehensive theory and experiments, we establish that VSPO has favorable properties compared to vanilla reward shaping and other alternative approaches. Specifically, under a bandit abstraction, VSPO provably achieves better iteration complexity than reward-shaped GRPO when the steering-induced distributions are sufficiently aligned with the target behavior. We evaluate VSPO across multiple reasoning benchmarks, including MATH and MMLU-Pro, for four target behaviors: explanation expertise, confidence expression, robustness to misleading context, and response verbosity. Our results show that VSPO consistently improves the control along target behavior while maintaining or improving task accuracy compared with reward shaping, teacher-trace distillation, and guidance-based baselines.

Off-policy updates are inevitable in reinforcement learning (RL) for large language models (LLMs) due to rollout staleness from asynchronous training and mismatches between training and inference engines. Naive importance sampling gives an unbiased correction but suffers from high variance, which is amplified by unbounded ratios and autoregressive generation. Prior remedies either rely on scenario-specific engineering, or trade bias for variance via token-level clipping or sequence-level normalization, yet these approaches remain largely heuristic. We propose **V**ariational s**E**quence-level **S**oft **P**olicy **O**ptimization (**VESPO**). By explicitly incorporating variance reduction into a variational formulation, we derive a principled closed-form reshaping kernel that operates directly on sequence-level importance weights, avoids token-level approximation and length normalization, and admits an explicit variance bound for the deployed kernel. Experiments on math reasoning and code generation show that VESPO maintains stable training under severe off-policy conditions (staleness up to $64\times$) and delivers consistent gains across both dense and Mixture-of-Experts (MoE) models, outperforming recent reshaping baselines under matched setup.

We propose an adaptive proximal gradient method for minimizing the sum of two functions, where one is a simple convex function, and the other belongs to one of the three classes: nonconvex smooth, convex nonsmooth, or convex smooth. The key feature of the method is an adaptive step size that accumulates historical gradient mapping norms in the denominator. Without any modification or knowledge of problem parameters, the method converges across all three problem classes under mild bounded-iterates and bounded-variance assumptions, with rates matching those of the proximal gradient method up to logarithmic factors, in both deterministic and stochastic settings. For the convex setting, we further propose an accelerated variant. It retains a similar near-optimal convergence rate for the nonsmooth case and achieves an improved rate of order $\widetilde{O}\big(1/t^2 + \sigma/\sqrt{t}\big)$ for the smooth case, which is optimal up to logarithmic factors. Notably, we develop new techniques for controlling the effect of stochastic noise, which are applicable across all three problem classes in the stochastic setting and enable simplified analysis.

Pruning and quantization are two dominant techniques that effectively address the computational and storage burdens of large language model (LLM) inference on edge devices. Recently, one-shot pruning has gained particular attention for its ability to identify weight supports via optimization without retraining. However, integrating such methods seamlessly with quantization for further model lightweighting is non-trivial. The discreteness of quantization shatters their required variable continuity, inevitably collapsing this integration into a decoupled pruning-quantization pipeline with inherently suboptimal outcomes. To tackle this problem, we propose Quantization-aware One-shot Pruning (QOP), which directly optimizes the quantized weights under sparsity and discreteness constraints, thereby explicitly capturing the impact of quantization on pruning within its objective. Specifically, QOP generalizes the alternating direction method of multipliers to sparse-constrained discrete optimization, enabling the identification of the high-quality support and the update of quantized weights. Theoretically, via the well-conditioned approximation obtained by a slight perturbation, we establish convergence to a stationary feasible point and provide the convergence rate. Extensive experiments across various LLMs, sparsity levels, and quantization settings demonstrate that QOP consistently outperforms existing baselines in terms of both average accuracy and perplexity under weight quantization.

Despite the promise of Neural Combinatorial Optimization (NCO), recent literature has increasingly coupled *decision paradigms*, *training objectives*, and *problem classes* with specific backbone architectures, limiting generic model design and fair cross-paradigm evaluation. We propose UECO, a simple yet effective **U**nified **E**ncoder for **C**ombinatorial **O**ptimization with a structure-aware attention mixture scheme, where a lightweight module defined upon graph-based locality is interleaved with conventional (global) node-focused attention blocks to dynamically evolve edge representations and inject local topological information into attention scores. Instead of fusing static edge features in a single step, UECO thereby integrates local and global messages in a progressive and problem-agnostic manner. UECO is orthogonal to task-specific decoders, and works seamlessly with *Local Construction* (LC), *Global Prediction* (GP), and *Adaptive Expansion* (AE) paradigms, across node- or edge-oriented problems, symmetric or asymmetric tasks, and sparse or dense graphs, for supervised, reinforcement, or generative learning. Extensive experiments show that, via plug-in substitution of encoders, UECO consistently improves solution and generalization quality over diverse backbones and task scales on TSP, ATSP, CVRP, ACVRP, MIS, MaxClique, and MaxCut.

$(1+1)$-evolution strategy (ES) is a classic variant of evolution strategy for black-box optimization via direct and random search while keeping only a single solution per generation in $\mathbb{R}^n$. Despite convergence with different step-size adaptations, it often remains impractical due to dependence on the unfavorable high dimension of the search space $\mathbb{R}^n$. In fact, many modern optimization problems are naturally constrained to low-dimensional manifolds embedded in very high-dimensional ambient spaces. To see the potential of intrinsic dimensionality reduction for improving practicality, we consider the truncated Riemannian $(1+1)$-ES on a compact Riemannian submanifold $\mathcal{M}\subset\mathbb{R}^n$ and establish its non-asymptotic stationarity guarantee at the standard nonconvex zeroth-order rate of $\widetilde{\mathcal{O}}(\varepsilon^{-4})$ function evaluations. The explicit Gaussian second-moment term in the bound scales with the intrinsic tangent dimension $d=\dim\mathcal{M}$ rather than the ambient dimension $n$, while the remaining truncation dependence is captured by an explicit small-ball drift constant $\alpha_d(R/\sigma_0)$. To the best of our knowledge, this gives the first non-asymptotic stationarity analysis for a comparison-based ES-type method on compact nonconvex Riemannian submanifolds under localized pullback smoothness. Numerical illustrations on sparse PCA over spheres and smooth-surrogate stationarity diagnostics illustrate behavior consistent with the theory and its limitations, while additional experimental studies on truncation radius diagnostics, black-box tuning, and classical evolutionary computation benchmarks support the practical relevance of geometry-aware comparison-based search under limited function-evaluation budgets.

It is widely recognized that in many applications, strong predictive performance is not enough to warrant use of a learned model. Additionally, it is required to ascertain whether the model conforms to certain properties such as fairness, robustness or safety. Unfortunately, most learning approaches are not guaranteed to produce a model that adheres to these properties. Model repair addresses this through a post-hoc debugging loop. Approaches such as verification are used to check if the desired property holds for a model; if not they return a counterexample which is used to modify or repair the ML model. We propose a novel repair method for tree ensembles. Our approach is based on a MILP encoding where the objective is to modify the leaf values in the tree ensemble such that it fixes the counterexamples while at the same time preserves predictive performance. The use of constrained optimization over the model parameters is compelling as it is guaranteed to fix a given set of counterexamples. We implement and evaluate our method in two use-case scenarios, classification requiring fairness and action policies in MDPs requiring safety. Results show that our MILP repair of tree ensembles can be feasible computationally, and is often successful in improving the desired property while not degrading performance.

We introduce Quokka, the first comprehensive scaling law for diffusion language models (DLMs), encompassing both compute- and data-constrained regimes alongside key modeling and optimization designs. Under compute constraints, we find that optimal parameter and dataset sizes scale equally with compute ($N_{\mathrm{opt}} \propto C^{0.5}$, $D_{\mathrm{opt}} \propto C^{0.5}$); however, DLMs are 2--5$\times$ more data-hungry than autoregressive (AR) models, requiring larger corpora and proportionally smaller models for a given FLOP budget. Under data constraints, validation loss follows a U-shaped curve across epochs, with the onset of overfitting scaling as $e_{\mathrm{opt}} \propto U_D^{0.39}/N^{0.55}$. This indicates that optimally leveraging a fixed unique data budget ($U_D$) requires allocating both modestly larger models and more training epochs. Beyond data and parameter allocation, we provide actionable guidance on DLM design choices: absorbing-mask transitions outperform uniform diffusion, linear noise schedules are the most stable and performant, and principled diffusion ELBO objectives ultimately surpass MaskGIT losses despite slower initial convergence. Furthermore, we demonstrate that easy-to-hard noise curricula accelerate early learning, established AR scaling laws for batch size and learning rate transfer seamlessly to DLMs, and weight decay, while unhelpful for single-epoch runs, is essential for multi-epoch regularization.

In this paper, we provide improved bounds on the complexity of solving $\ell_p$-Lipchitz convex optimization problem in parallel. We lower bound the depth of highly-parallel algorithms, i.e., the number of rounds required by a parallel algorithm making $O(\mathrm{poly}(d))$ queries per round to compute an $\epsilon$-approximate minimizer of a $\ell_p$-Lipschitz, $d$-dimensional convex function over the unit $\ell_p$ ball. We provide a framework for proving lower bounds for parallel $\ell_p$-Lipschitz convex optimization that generalizes prior lower bound proofs and for all $p \in (1,\infty)$ (other than 2) improves the depth for which it can be shown that no super-polylogarithmic improvement over sequential algorithms is possible. We also prove a variety of new lower bounds for $p>2$. Notably, for $p=\infty$, we establish the near optimal parallel complexity of the class of ball-acceleration methods up to depth $\tilde{O}(\sqrt{d})$, which is the first near-optimality result in the parallel convex optimization setting for an algorithm other than mirror descent. Additionally, by providing new algorithms we characterize the optimal depth needed to solve to accuracy $\epsilon = d^{-1/4}$ for all $p \geq 2$ (up to polylogarithmic factors). We also provide extensions of our lower bounds to parallel quantum $\ell_p$-Lipschitz convex optimization.

Scaling laws and hyperparameter tuning remain central challenges in training large-scale models, particularly under differential privacy (DP), where optimization is further complicated by noise injection and per-sample gradient clipping. While Maximal Update Parametrization ($\mu$P) allows for zero-shot hyperparameter transfer in standard training, we show that it fails under DP constraints due to the distinct spectral properties of high-dimensional noise versus low-rank gradients. In this work, we introduce DP-$\mu$P, a theoretically grounded framework that extends $\mu$P to the privacy-preserving setting. By analyzing the interaction between gradient clipping, noise injection, and feature learning, we derive a Signal-to-Noise Ratio (SNR) condition that governs layer-wise learning rate scaling. Combining this with our proposed spectral scaling formulation, we provide specific scaling rules for DP-SGD, DP-Adam, and DP-Muon. Empirically, we demonstrate that DP-$\mu$P enables effective zero-shot transfer of learning rates across model widths on diverse architectures (MLP, ViT, GPT-2), matching or exceeding the performance of individually tuned baselines while significantly reducing computational costs.

Proximal Policy Optimization (PPO) and it's variants are the most widely used policy-gradient methods in reinforcement learning, yet the role of its actor update mechanism is still poorly understood theoretically. In particular, standard PPO combines clipped surrogate gradients, multiple epochs of minibatch updates, and reuse of rollout data, but existing convergence analyses do not fully capture this update structure. In this work, we theoretically study PPO policy updates with symmetric clipping from a policy-gradient perspective and interpret its actor updates as a cyclic biased-gradient method with sample reuse and random reshuffling. Our first contribution is a clean formalization of PPO clipped actor updates through surrogate gradients that approximate the true policy gradient. Using the performance difference lemma, we prove a linear bias bound which quantifies how the surrogate gradient drifts as the policy moves away from the sampling policy. Our second contribution is a convergence analysis of cyclic surrogate-gradient ascent, showing that additional PPO-style biased updates can improve progress under conservative learning rates without requiring additional samples. Finally, we analyze the stochastic minibatch version with reshuffling and obtain convergence-to-stationarity guarantees under standard smoothness assumptions and a bounded critic-bias condition. Overall, our results provide a theoretical interpretation of PPO’s multi-epoch actor updates: extra clipped surrogate steps introduce bias, but can still improve optimization efficiency by compensating for small, stable step sizes through sample reuse.

Inverse optimization (IO) seeks to infer the parameters of a decision-maker's objective from observed context-action data. We study *noiseless IO*, where demonstrations are generated by a ground-truth objective. We provide a high-probability $\mathcal{O}(\frac{d}{T})$ generalization bound for the induced action set, where $d$ is the number of unknown parameters and $T$ is the size of the training dataset. We strengthen these guarantees under additional conditions that ensure uniqueness of the chosen action, bringing our IO guarantees in line with best-arm identification results in the bandit literature. We further show that the $\mathcal{O}(\frac{d}{T})$ rate is tight over all consistent estimators considered here, and extend the result to both instantaneous and cumulative regret. Notably, the resulting regret lower bound matches the corresponding upper bounds in the adversarial setting, indicating that the stochastic IO setting is effectively adversarial for the class of estimators studied here. Finally, we propose a parameter-free algorithm with lower per-iteration complexity than generic solvers. Experiments validate the predicted rates and illustrate the tightness of our bounds.

Initialization determines the optimization ``fate'' of a model, going far beyond just setting the scale of its weights. It establishes a fundamental initial geometry over the data, permanently dictating which examples are considered close, which directions are easy to change, and which representations gradient descent will naturally refine first. Because this starting point dictates the model's future trajectory, we ask: can this initial geometry be explicitly chosen by referencing a completely different architecture? To achieve this, we introduce representational similarity as an optimization prior. This is an initialization-only procedure that aligns a target network to the representational geometry of a randomly initialized guide network before any downstream training occurs. Crucially, the guide network transfers no learned knowledge: it is frozen, never sees labels. We demonstrate that this cross-architecture transfer of fate works in two key domains: instilling the optimization benefits of depth into shallow targets by aligning them to deeper, randomly initialized guides, and transferring representational granularity by using fine-resolution guides to shape the initialization geometry of networks built with coarser representations. Ultimately, this establishes that a model's architectural destiny can be decoupled from its physical structure and explicitly programmed at initialization.

In Low-Rank Adaptation (LoRA), the scaling factor $\alpha$ is often treated as a mere complement to the learning rate, yet its role in optimization remains poorly understood. In this paper, we reveal that the scaling factor $\alpha$ and the learning rate function differently, with $\alpha$ emerging as the dominant driver of effective optimization, delivering gains that cannot be replicated by learning rate scaling alone. Through the synergy of extensive empirical analysis and a theoretical Signal-Drift framework, we uncover three findings into LoRA’s scaling mechanism: First, LoRA’s spectral suppression smooths the optimization landscape, rendering standard hyperparameters overly conservative and creating an optimization gap. Second, when leveraging this smoothness to accelerate convergence, $\alpha$ outperforms the learning rate by amplifying the task signal without increasing the drift ratio. Third, the optimal scaling factor follows a sublinear relationship with the rank, well characterized by a square-root law with an unexpectedly large coefficient, revealing the insufficient scaling of existing rank-tied heuristics. Based on these insights, we propose LoRA-$\alpha$, a minimalist framework that restores $\alpha$ to its principled regime, making LoRA compatible with standard small learning rates. Extensive evaluations across diverse tasks demonstrate that LoRA-$\alpha$ consistently improves performance while streamlining hyperparameter search, unleashing the learning potential of LoRA.