NeurIPS 2026 Papers

Proximal Policy Optimization (PPO) is value-driven only after experience has been collected: Generalized Advantage Estimation reinforces sampled actions according to their bootstrapped advantage, but the exploration process that produced those samples is still governed mainly by policy-space noise and entropy regularization. As a result, PPO does not explicitly ask, before sampling, which uncertain transitions would most improve its advantage estimates. We argue that PPO exploration should be formulated as *value-priced uncertainty*: uncertainty should be explored when it is both large and able to change the advantage. We propose **VALU** (**V**alue-**A**ware **L**atent **U**ncertainty), a family of PPO-compatible exploration bonuses of the form b(s,a) = P(s,a) U(s,a). The price $ P(s,a) = \lVert \nabla_{z'} \tilde V_\psi(\hat z') \rVert $ measures how strongly uncertainty in the predicted next latent can perturb the advantage, while \(U(s,a)\) is an uncertainty quantity instantiated by non-parametric effective counts or learned prediction errors such as RND. Our theory derives the value-pricing term from a unified advantage-uncertainty view: an uncertainty coordinate induces an optimistic advantage correction through a dual norm. This yields count-based UCB and MCR members, corresponding respectively to optimism over latent-transition uncertainty and one-sample shrinkage of the optimistic advantage envelope. Across six standard MuJoCo continuous-control tasks, VALU substantially improves PPO's sample efficiency and final return, and achieves the best PPO-compatible performance on most tasks while remaining competitive on the others, compared with strong exploration baselines including RND, ICM, RE3, VCSE, and OPPO. Ablations over uncertainty estimators, schedules, latent encoders, and bonus components show that the gains come from pricing uncertainty by local value sensitivity rather than from any single implementation choice.

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.

$(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.

In this work, we introduce a novel training-free encoding framework, dubbed **XShapeEnc**, that encodes an arbitrary spatially grounded 2D geometric shape into a compact representation exhibiting five favorable properties, including invertibility, adaptivity, generality and controllability. Specifically, a 2D spatially grounded geometric shape is decomposed into its normalized geometry within the unit disk and its pose vector, where the pose is further transformed into a harmonic pose field that also lies within the unit disk. A set of orthogonal Zernike basis is constructed to encode shape geometry and pose either independently or jointly, with controllable relative emphasis on shape geometry or shape pose. We demonstrate the theoretical validity, efficiency, discriminability, and wide applicability of **XShapeEnc** via extensive analysis and experiments across a wide range of shape-aware tasks and our self-curated XShapeCorpus dataset. We envision **XShapeEnc** as a foundational tool for research beyond one-dimensional sequential data and data-driven, learning-based encoding paradigms, paving the way for a unified spatial encoding framework for frontier 2D spatial intelligence.

Transformers perform inference by iteratively transforming token representations across layers. This layerwise computation has been studied empirically, and recent mean-field theories of Transformer dynamics explain how attention can drive token distributions toward clustering. However, existing mean-field analyses largely treat model parameters as prescribed, leaving open how training reshapes this clustering picture. We study this question in a noisy mean-field Transformer in which only a parameter-linear FFN is trained under $L^2$ regularization. We find and analyze a training-induced phase in the dynamics: after initially following attention-driven clustering, the token distribution can leave the clustered regime near the final layers. Our mathematical analysis is based on an entropy-regularized interaction energy that captures the clustering bias of attention. More broadly, our results point toward a training-aware mean-field theory of Transformer dynamics, in which training and inference dynamics are treated together.

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.

Machine learning models, particularly large foundation models, are increasingly evaluated along multiple axes that are relevant for their usage. In addition to accuracy measures appropriate for the specific task of interest, measures capturing safety, calibration and fairness properties are often of interest. The classical approach to these problems assumes that the test set has been held out from the training process, an assumption that is increasingly harder to justify for large foundation models. In this work, we formalize a notion of evaluability without this assumption, and ask if it is feasible to evaluate black-box models in settings where our test samples may have been used during training. We study this problem under different settings, capturing different levels of access to the model. We first study the natural setting, where the evaluator only looks at the output of the model on the test set. We fully characterize the class of distributions where accurate evaluation is possible, and show tight upper and lower bounds on the sample complexity of evaluation. Our results show that black-box evaluation in this set up is feasible if and only if the distribution is close to being small support. We then consider evaluation algorithms that can query the model on additional inputs and demonstrate connections to self-correctors studied in program testing. Using these, we show that a small amount of query access can increase the power of the evaluator for natural distributions and concept classes.

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.

Transformers have demonstrated remarkable in-context learning (ICL) capabilities, enabling them to perform new tasks without additional fine-tuning. However, their performance often deteriorates when encountering out-of-distribution (OOD) inputs that deviate from the training distribution, and the underlying theory remains poorly understood. To fill this gap, we characterize the OOD error under the input distribution shift through the interplay between the dynamics of the so-called $\alpha$-type and $\beta$-type attention weights, which represent the transformer’s confidence in identifying the correct and incorrect features, respectively. Our results indicate that the OOD error for each feature depends on all pairwise interactions between the training features and OOD features, and under certain cases the transformer performs no better than random guessing. To improve the OOD generalization performance, we next investigate the impact of model finetuning with the OOD data, and particularly, characterize the model forgetting performance on the source domain. Interestingly, the performance on the source domain may not always degrade after finetuning, which highly depends on the nature of the feature shift: finetuning on OOD domain keeps enhancing the confidence of identifying correct features from the original distribution, while the interference from other incorrect features may either increase or decrease. Extensive experiments on both synthetic and real data are conducted to corroborate the theoretical insights.

We develop a framework for analyzing parameter symmetries in deep ReLU networks and obtain a complete characterization of the generic parameter fibers for three-layer bottleneck architectures. Our approach provides explicit semi-algebraic descriptions of these fibers and yields a polynomial time algorithm for deciding functional equivalence of two parameters. The symmetries include discrete and continuous transformations arising from layer composition, and depend on whether deeper layers hide or preserve geometric structure from preceding layers. Finally, we show that some of these symmetries induce local conservation laws along gradient flow, while others do not.

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.

Predictive models on structured tabular data often face imbalanced targets: training emphasizes common values while rare outcomes stay underrepresented. Distributionally robust optimization hedges against shift, but standard ambiguity sets use a single radius over the entire distribution, so adversarial mass concentrates on dense regions and rare target ranges remain weakly protected; worst-case statements then track bulk behavior rather than tails. We introduce target-conditional distributionally robust optimization (TC-DRO), which defines ambiguity per target region with radii that grow where data are scarce. We establish coverage, per-region excess risk that tightens with local sample size when standard DRO yields vacuous tail guarantees, and a dual as weighted empirical risk with theory-derived density-adaptive weights. Experiments on seven tabular benchmarks show that Wasserstein DRO collapses to empirical risk minimization under uniform weights, while global $\chi^2$-DRO can remain close to ERM in practice; TC-DRO improves tail MSE and SERA relative to both standard training and recent imbalanced-regression methods.

We study signal propagation at initialization in transformers through the averaged partial Jacobian norm (APJN), a measure of gradient amplification across layers. We extend APJN analysis to transformers with bidirectional attention and permutation-symmetric input token configurations by deriving recurrence relations for activation statistics and APJNs across layers. Our theory predicts how attention modifies the asymptotic behavior of the APJN at large depth and matches APJNs measured in deep vision transformers. The criticality picture known from residual networks carries over to transformers: the pre-LayerNorm architecture exhibits power-law APJN growth, whereas transformers with LayerNorm replaced by elementwise tanh-like nonlinearities have stretched-exponential APJN growth, indicating that the latter are subcritical. Applied to Dynamic Tanh (DyT) and Dynamic erf (Derf) transformers, the theory explains why these architectures can be more sensitive to initialization and optimization choices and require careful tuning for stable training.

Large language models (LLMs) deployed as agents remain vulnerable to prompt-based instruction injection, where user inputs induce the model to violate higher-priority system directives. Existing methods often treat violations as a coherent class or rely on surface text patterns, which can limit robustness as attack strategies evolve. In this paper, we investigate LLM hidden-state dynamics under attack and uncover a phenomenon we term Asymmetric Logic: the model's response generations concentrate in a relatively low-variability region of latent representation space when it's compliant to the system prompt, while compromised behaviors exhibit diverse deviations. Based on this observation, we propose Asymmetric Entropic Optimization (AEO), a novel framework that monitors the internal shifts in latent representations to detect and proactively rectify anomalous behaviors. Our method effectively identifies logic deviations by quantifying the entropic divergence between the compliant manifold and the chaotic violation space. Across a range of state-of-the-art LLM backbones, AEO demonstrates strong robustness to diverse instruction injection attacks. For instance, on Qwen2.5-14B-Instruct, AEO attains a detection AUC of 97.28% with FPR95 of only 10.38%. The same chaos score also gates a lightweight post-hoc correction module that lifts overall instruction-following accuracy from 63.94% to 76.00% on IHEval, without any model fine-tuning. The code is available at https://anonymous.4open.science/r/AEO-code-56D2/.

We study stochastic zeroth-order (ZO) optimization of smooth nonconvex objectives under heavy-tailed sample-gradient noise, a regime motivated by empirical evidence that gradient noise in modern machine learning can violate the bounded-variance assumption underlying classical ZO theory. While the first-order literature has established optimal rates under bounded $p$-th moment noise for $p \in (1,2]$, analogous high-probability guarantees for nonconvex ZO optimization remain largely unexplored. The ZO setting is not a direct corollary of first-order theory: first-order methods update using $\nabla F(x;\xi)$, the same object on which the noise assumption is imposed, whereas derivative-free methods only observe noisy function values and construct directional finite-difference estimates. Thus, weak-$L_p$ control of $\nabla F(x;\xi)-\nabla f(x)$ must first be transferred to scalar two-point directional estimates, a concentration step absent from standard first-order analyses. We propose the Robust Scalar-Clipped Zeroth-Order method (\textbf{RSC-ZO}), a two-point method that clips each scalar directional derivative before aggregation. Under sample-wise smoothness and a weak-$L_p$ tail bound on the sample-gradient noise, RSC-ZO finds an $\varepsilon$-stationary point with high probability using $ \widetilde{O}\\left( \frac{d^{\frac{p}{2(p-1)}}} {\varepsilon^{\frac{3p-2}{p-1}}} \right) $ noisy function evaluations, matching the optimal first-order $\varepsilon$-dependence. At $p=2$, our bound becomes $\widetilde{O}(d\varepsilon^{-4})$, matching the classical dimension-accuracy dependence for stochastic ZO methods under bounded-variance noise, which is typically established only in expectation. In contrast, our guarantee holds with high probability and under the strictly weaker weak-$L_2$ tail condition, which can allow infinite gradient-noise variance. We further extend the analysis to a momentum variant and quantify the resulting batch-size/stepsize tradeoff.

Second-order methods are provably faster than first-order methods, and their efficient implementations for large-scale optimization problems have attracted significant attention. Yet, optimization problems in ML often have nonsmooth derivatives, which makes the existing convergence rate theory of second-order methods inapplicable. In this paper, we propose a new semismooth Newton method (SSN) that enjoys both global convergence rates and asymptotic superlinear convergence without requiring second-order differentiability. Crucially, our method does not require (generalized) Hessians to be evaluated at each iteration but only periodically, and it reuses stale Hessians otherwise (i.e., it performs lazy Hessian updates), saving compute cost and often leading to significant speedups in time, whilst still maintaining strong global and local convergence rate guarantees. We develop our theory in an infinite-dimensional setting and illustrate it with numerical experiments on matrix factorization and neural networks with Lipschitz constraints.

While adaptive gradient methods are the workhorse of modern machine learning, sign-based optimization algorithms such as Lion and Muon have recently demonstrated superior empirical performance over AdamW in training large language models (LLM). However, a theoretical understanding of why sign-based updates outperform variance-adapted methods remains elusive. In this paper, we aim to bridge the gap between theory and practice through the lens of heavy-tailed gradient noise, a phenomenon frequently observed in language modeling tasks. Theoretically, we introduce a novel generalized heavy-tailed noise condition that captures the behavior of LLMs more accurately than standard finite variance assumptions. Under this noise model, we establish sharp convergence rates of SignSGD and Lion for generalized smooth function classes, matching or surpassing previous best-known bounds. Furthermore, we extend our analysis to Muon and Muonlight, providing what is, to our knowledge, the first rigorous analysis of matrix optimization under heavy-tailed stochasticity. These results offer a strong theoretical justification for the empirical superiority of sign-based optimizers, showcasing that they are naturally suited to handle the noisy gradients associated with heavy tails. Empirically, LLM pretraining experiments validate our theoretical insights and confirm that our proposed noise models are well-aligned with practice.

Top-two algorithms are simple and effective for fixed-confidence best-arm identification, but their sharp non-asymptotic behavior is still not well understood. We study this problem for Bernoulli bandits through $\beta$-EB-TCI, the empirical-best top-two rule of Jourdan et al. (2022), whose challenger is chosen using a Bernoulli transportation cost with a logarithmic count penalty. We prove that, after the empirical leader has become the true best arm and its sampling fraction stays close to $\beta$, the stopping time is $T_{\beta}^{\star}(\mu)\log(1/\delta)$ up to lower-order concentration terms. We also show that, in this regime, every challenger is sampled linearly often. Thus, for the original algorithm without forced exploration, the main remaining difficulty is to control when the empirical leader becomes permanently correct. These results imply a non-asymptotic high-probability bound for all Bernoulli instances with a unique best arm. If the algorithm satisfies a finite-mean sufficient-exploration condition, the bound further yields the sharp expected sample complexity. In particular, this gives the sharp expectation result for the unguarded Bernoulli rule when all arm means are pairwise distinct, using the sufficient-exploration result of Jourdan et al. (2022). Finally, if we add a mild forced-exploration rule that contributes only $O(\sqrt{Kt})$ pulls up to time $t$, we obtain a self-contained expected sample-complexity theorem for any number of arms under the unique-best-arm assumption. We also identify a limitation of proof strategies that try to handle equal suboptimal means through a single index-comparison argument.

We study optimistic bilevel optimization when the lower-level problem has a non-isolated manifold of minimizers. In this setting, the hyper-objective may be non-differentiable because the upper-level criterion must choose among multiple lower-level solutions. Under a local Polyak--{\L}ojasiewicz (P{\L}) condition, we show that differentiability does not require the lower-level solution set to be a singleton: uniqueness of the optimistic selection is sufficient. This yields an explicit pseudoinverse-based hyper-gradient formula extending the classical singleton-minimizer result. We further characterize the regularity of the hyper-objective: non-degeneracy of the selected minimizer along the solution manifold yields local smoothness, while failure of uniqueness can create many non-differentiable points and failure of non-degeneracy can destroy all positive H\"older regularity of the hyper-gradient. Motivated by this theory, we propose HG-MS, a select-then-differentiate method combining explicit optimistic selection with efficient pseudoinverse-based hyper-gradient computation. Despite the nonconvex nature of optimistic selection over the lower-level solution manifold, we show that HG-MS converges to a stationary point of the optimistic objective with complexity governed by the intrinsic dimension of the solution manifold rather than its ambient dimension. Empirically, we test a practical variant of HG-MS for matched-budget LLM source reweighting. This variant preserves the select-then-differentiate principle and obtains the best GSM8K/MATH scores across the tested backbones, along with competitive or best MT-Bench instruction-following results.