NeurIPS 2026 Papers

Masked diffusion language models generate text by iteratively unmasking positions, with practical samplers often selecting reveal positions from the model's own logits. The reveal order is therefore an endogenous latent variable of the rollout process with an enormous discrete latent space. Policy optimization objectives relying on pre-defined schedules or sampled trajectories introduce either bias or high variance. In this paper, we propose \texttt{FEROM}, \emph{Frontier Endogenous Reveal-Order Marginal Policy Optimization}, which targets the rollout-induced marginal response policy of masked diffusion LMs. \texttt{FEROM} derives a Rao--Blackwellized policy-gradient identity over latent reveal orders and expresses the resulting estimator as posterior edge occupancy on a reveal-state DAG. To make marginalization practical, we introduce Frontier Reveal Marginalization, a budgeted estimator that combines scorelaw local reveal estimation with frontier expansion of high-mass partial states. Integrated into a GRPO-style objective, \texttt{FEROM} replaces single-path log-scores with a locally marginalized edge-based surrogate. Experiments on math and coding tasks show comparable or improved results over existing methods under matched compute budgets. Offline proxy study further shows potential in gains with increased budgets.

Safety-critical multi-agent systems require agents to learn coordinated behaviors while avoiding unsafe actions throughout learning. Existing safe reinforcement learning theory has established zero-violation guarantees for single-agent Markov decision processes with instantaneous hard constraints, while much of safe multi-agent reinforcement learning focuses on cumulative constraints, policy optimization, or shielding. We study a stricter setting: cooperative Markov games with unknown dynamics and coupled instantaneous hard constraints, where a joint action must be safe at every time step of every episode. This setting introduces challenges absent from the single-agent case: actions that are locally safe for individual agents may be unsafe jointly, unsafe joint actions can affect future feasible regions, and naive reductions to single-agent safe RL suffer exponential dependence on the number of agents. We propose a graph-structured safe learning algorithm that constructs conservative multi-agent safety certificates and explores optimistically only within certified safe joint subgraphs. Under structured coupling assumptions, the algorithm guarantees zero constraint violation with high probability and achieves sublinear regret against the optimal safe joint policy, with complexity depending on local interaction structure rather than the full joint action space.

We study instruction following in multi-task reinforcement learning, where an agent must zero-shot execute novel tasks not seen during training. In this setting, linear temporal logic (LTL) has been adopted as a powerful framework for specifying structured, temporally extended tasks. While existing approaches successfully train generalist policies, they often struggle to effectively capture the rich logical and temporal structure inherent in LTL specifications. In this work, we address these concerns with a novel approach to learn *structured* task representations that facilitate training and generalisation. Our method conditions the policy on sequences of Boolean formulae constructed from a finite automaton of the task. We propose a hierarchical neural architecture to encode the logical structure of these formulae, and introduce an attention mechanism that enables the policy to reason about future subgoals. Experiments in a variety of complex domains demonstrate the strong generalisation capabilities and superior performance of our approach.

Reinforcement learning with verifiable rewards (RLVR) for Large Reasoning Models hinges on baseline estimation for variance reduction, but existing approaches pay a heavy price: PPO requires a policy-model scale critic, while GRPO needs multiple rollouts per prompt to keep its empirical group mean stable. We introduce $\textbf{POISE}$ ($\underline{P}$olicy $\underline{O}$ptimization with $\underline{I}$nternal $\underline{S}$tate Value $\underline{E}$stimation), which obtains a baseline at negligible cost by using the policy model's internal signals already computed during the policy forward pass. A lightweight probe predicts the expected verifiable reward from the hidden states of the prompt and generated trajectory, as well as token-entropy statistics, and is trained online alongside the policy. To preserve gradient unbiasedness despite using trajectory-conditioned features, we introduce a cross-rollout construction that predicts each rollout's value from an independent rollout's internal states. Because POISE estimates prompt value using only a single rollout, it enables higher prompt diversity for a fixed compute budget during training. This reduces gradient variance for more stable learning and also eliminates the compute overhead of sampling costs for detecting zero-advantage prompts. On Qwen3-4B and DeepSeek-R1-Distill-Qwen-1.5B across math reasoning benchmarks, POISE matches DAPO while requiring less compute. Moreover, its value estimator shows similar performance to a separate LLM-scale value model and generalizes to various verifiable tasks. By leveraging the model's own internal representations, POISE enables more stable and efficient policy optimization.

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.

Aerial vision-language navigation requires agents to follow natural-language instructions through closed-loop perception and action in 3D environments. We argue that aerial VLN can be formulated as a prediction-driven world-action problem: the agent should anticipate latent world evolution and act according to the predicted consequences. To this end, we propose WorldVLN, the first autoregressive world action model for aerial VLN. Unlike full-sequence video-generation world models that generate an entire visual clip, WorldVLN adapts a latent autoregressive video backbone to predict short-horizon world-state transitions and directly decodes them into executable waypoint actions. After each action segment is executed, newly received observations are encoded back into the autoregressive context, enabling closed-loop world-action prediction. We further introduce a two-stage training framework that first grounds the video prior in instruction-conditioned navigation dynamics and then develops Action-aware GRPO, the first reinforcement learning method tailored to autoregressive WAMs, to optimize waypoint decisions through their downstream rollout consequences. On public outdoor and indoor benchmarks, WorldVLN consistently outperforms existing Vision-Language-Action baselines with $12\%+$ success-rate gains and larger advantages on challenging cases. It further transfers zero-shot to real drone deployment, suggesting that the proposed WorldVLN offers a promising route for spatial action tasks. Demos and code are available at \url{https://worldvln.com} (anonymous).

Video world models aim to generate temporally coherent visual observations conditioned on actions, but often suffer from 3D geometric inconsistencies such as spatial drift, object deformation, and perspective violations, which limit their reliability in downstream tasks. We propose WorldPrism, a post-training framework that improves 3D consistency via reinforcement learning without modifying the base model architecture. The central challenge lies in designing a reliable reward signal. To address this, we introduce Bidirectional Geometry-Flow Verification, a bidirectional cross-space mechanism that evaluates consistency through two complementary paths: (i) lifting 2D optical flow into 3D to measure position-space agreement, and (ii) projecting reconstructed 3D geometry to match the point correspondence. By leveraging these two independent and complementary signals, this design mitigates the blind spots of unilateral 2D evaluation. We further incorporate hierarchical temporal evaluation to capture both local and long-range geometric consistency. Experiments on diverse datasets demonstrate that WorldPrism achieves state-of-the-art 3D consistency while preserving visual quality and generative diversity.

Training transformer-based architectures with data augmentation has become an increasingly popular approach for equivariant machine learning. Despite its empirical success, the interplay between the transformer architecture, invariance to different symmetries, and augmentation budgets remains underexplored. In this paper, we investigate the ability of a vanilla transformer to learn a wide range of symmetries through limited data augmentation. We identify that it performs strongly for common angle-preserving symmetries, while it struggles with non-angle-preserving ones. Within the angle-preserving family, we further find that base subgroups such as translation, rotation, and scale are learned more easily and accurately than compositional groups that combine them. Motivated by this observation, we perform a structural analysis of the trained models and identify interpretable mechanisms that induce invariance to each base angle-preserving symmetry. These findings support what we coin winning the *Symmetry Lottery*: the transformer architecture is aligned with learning mechanisms invariant to symmetries that happen to be prevalent in scientific domains.

Large language models achieve high reasoning performance via explicit chain-of-thought and reinforcement learning, but require long output sequences and extended inference time. Latent reasoning reduces this cost by shifting computation into a latent space; however, continuous latent methods are hard to train, suffering from unstable and uninterpretable reasoning trajectories. We argue these issues stem from a misalignment between continuous-space reasoning and discrete symbolic supervision, as continuous states lack explicit anchors for step-by-step alignment. To resolve this, we propose **Discrete Latent Reasoning (DLR)**, the first method that converts continuous latent states into explicit discrete tokens. Inspired by render-based compression, we render textual chains of thought into images, extract visual features, and construct a discrete latent vocabulary via clustering-based fine-tuning. Expanding the vocabulary and output head enables standard autoregressive modeling over both natural language and latent tokens, supporting pretraining alignment, SFT, and RL. Experiments on five reasoning benchmarks and two model series (Qwen3-VL and LLaMA-3) confirm that **DLR** outperforms prior latent reasoning baselines with up to **20$\times$ compression**. Furthermore, the learned latent trajectories retain an interpretable semantic structure. Overall, discrete latent tokens provide a controllable and interpretable basis for efficient latent reasoning.

When is it better to rely on information from others rather than experiment to discover the answer yourself? Social learning refers to the ability of humans to learn by observing, imitating, or interacting with others. In this work, we use multi-agent reinforcement learning (MARL) to investigate the conditions under which social learning achieves higher returns and improved learning efficiency relative to individual learning. We begin with a theoretical analysis and discover that, under coordinated group exploration and low skill transmission costs, social learning achieves lower sample complexity than individual learning. Motivated by this insight, we propose Selective Social Learning (SSL), a novel social learning algorithm that reduces the cost of skill transmission, and we design experiments across three MARL environments to test these predictions. Empirical results show that SSL outperforms baselines under conditions consistent with our theory. Surprisingly, we also find that in non-stationary environments with sparse rewards, which mirror the real world, social learning can emerge in standard MARL without an explicitly designed social learning mechanism. This suggests that the apparent absence of social learning in standard RL stems from limitations in training environment design rather than in RL itself. We therefore call for future MARL training environments to inherently incorporate sparsity and non-stationarity, enabling agents to naturally develop social-learning behaviors that transfer to complex real-world settings.

Reinforcement learning from trajectory-level outcomes has become a standard approach for training role-partitioned LLM agent pipelines, in which multiple specialized roles such as an orchestrator, a retriever, and a synthesizer collaborate to solve complex reasoning tasks. Upon trajectory failure, existing training methods do not explicitly determine which role should receive updates, applying outcome signals to all roles indiscriminately or diffusely across the pipeline. However, this paper identifies that failures can often be traced to a primary responsible role, which we term the *bottleneck role*. This suggests that updating all roles may contaminate non-bottleneck roles with irrelevant updates. Therefore, an intuitive solution is to exclusively update the bottleneck role for each failed trajectory. Unfortunately, accurately *identifying this role is highly challenging*, as we show that roles exhibiting visible symptoms of failure are often merely downstream victims of errors originating from other roles, rather than being the true bottleneck. To tackle this challenge, we propose **R**ole **I**nference for **S**elective **E**volution (*RISE*), the first method that formulates per-trajectory update-target selection as role bottleneck inference. Specifically, RISE leverages Beta-binomial lower-confidence scores to estimate role responsibility, effectively discounting uncertain evidence and suppressing updates when attribution confidence is insufficient. Furthermore, RISE decouples the visible symptom locus from the support-constrained responsibility target prior to applying selective role-level training. Empirically, across HotpotQA, 2WikiMultihopQA, and MuSiQue with Qwen2.5-7B and Llama3-8B backbones, RISE consistently outperforms the strongest learning-based baseline, with up to **11.1%** relative F1 and **14.1%** relative EM improvements over GiGPO on MuSiQue with Llama3-8B. Code is available at: [https://anonymous.4open.science/status/RISE-C0F0](https://anonymous.4open.science/status/RISE-C0F0)

Detailed image captioning is usually framed as an amount problem: a model should mention more visual facts while avoiding hallucination. We argue that this framing hides a basic structure of long-form captioning: a caption is an ordered sequence, and different failures appear at different positions. We introduce a position-aware rubric representation that decomposes dense human references into atomic visual claims, verifies generated captions against those claims, and induces a quartile-wise composition of supported, hallucinated, and generic content. This lens reveals two qualitatively different failures in long captions: models tend to produce unsupported specificity in the middle, while the final portion often escapes into vague, non-evidential language rather than simply accumulating more hallucination. Dense human references do not exhibit the same position-dependent pattern under the identical pipeline, suggesting that the effect is a model behavior rather than an artifact of long descriptive text. Motivated by this diagnosis, we propose *PAGER* (**P**osition-**A**ware **G**rounded **E**vidence **R**einforcement), a GRPO training objective with separate signals for supported fine-grained middle evidence, evidence-bearing tail continuation, claim coverage, and content-aware length control. The result is a closed loop: the same ordered claim representation identifies where long captions fail and supplies the position-specific reward needed to repair those failures. Experiments on detailed-caption benchmarks show that PAGER improves position-wise grounding and pairwise caption preference while preserving competitive global caption quality.

Offline RL can learn effective policies from fixed datasets, but deployment objectives often shift after training, while the trained actor may be frozen due to data, cost, or governance constraints. We study deployment-time adaptation of frozen offline actors using Product-of-Experts (PoE) composition with a goal-conditioned prior. Our main practical finding is graceful degradation rather than universal improvement: with degraded or random priors, precision-weighted PoE remains anchored to the frozen actor, whereas additive and prior-only adaptation often collapse, and a KL-budget selector typically reaches a near-oracle operating point. We also clarify that, for diagonal-Gaussian actors and priors, PoE with coefficient $\alpha$ induces the same deterministic policy as KL-regularized adaptation with $\beta=\alpha/(1-\alpha)$, with posterior covariance differences cancelling under deterministic action selection. Across D4RL MuJoCo and AntMaze diagnostics, results reveal an actor-competence ceiling: composition helps or preserves performance when the frozen actor is sufficiently competent, but cannot recover when the actor itself fails. Swept CQL/IQL-guided and SF/GPI baselines are competitive in some settings, but do not eliminate this regime boundary. Thus, PoE/KL-Reg is best viewed as an actor-anchored safety layer for deployment-time objective shift, not as a general return-maximization method, with effectiveness governed by actor competence rather than knob choice.

Large-scale AI training runs span days within fixed HPC allocations, making checkpointing essential for bounded recovery after failure. Asynchronous checkpointing hides checkpoint latency from the training loop, but does not guarantee that the most recent state has reached a durable target: data may still be in flight to the shared parallel filesystem (PFS) when a failure occurs, forcing rollback to an older checkpoint. We formalize this durability gap through a distributional model that combines hardware crossover analysis with statistical models of PFS completion time under scale and contention. Rank-local completion times follow a lognormal body distribution, while wall-clock completion is governed by an extreme-value process determined by the slowest writer. The resulting analysis yields a contention-dependent threshold,~$C^\star_p(n)$, beyond which the PFS is no longer the fastest path to durable persistence. Guided by this analysis, we design CacheX, a transparent staging layer that composes with PyTorch Distributed Checkpoint~(DCP). CacheX stages bulk data to local DRAM and replicates it to a neighboring node's NVMe, providing allocation-local durability while preserving eventual persistence to the PFS. On Frontier, across IOR runs up to 256~nodes and end-to-end FSDP+DCP training up to 64~nodes, CacheX reduces write-time variability by up to~$43\times$, achieves a replica-durable effective throughput of~${\sim}\,2.5$ to $2.8$\,GiB/s per node decoupled from PFS contention, and reduces worst-case rollback exposure. Our results show that staging is not a universal accelerator; rather, it is a conditional policy for selecting the first durable checkpoint boundary under measured scale, load, and tail behavior.

For $K=2$ regimes, scalar thresholding of the Bayesian posterior is Bayes-optimal. For $K \geq 3$, the posterior evolves on a $(K-1)$-dimensional simplex $\Delta^{K-1}$, yet virtually all deployed detection systems reduce it to a scalar score without theoretical justification. We resolve this question completely. We introduce **Extended Decision Sufficiency (EDS)**—three linear constraints on emission densities and transition dynamics (Rank-One Emissions, Markov Factorisation, Normal Factorisation), verifiable in $O(K^2)$ operations from model parameters alone—and prove the following exact characterisation: **Main result.** EDS holds if and only if the Bayes-optimal stopping value function is constant on every level set of a scalar statistic $\phi : \Delta^{K-1} \to \mathbb{R}$. Sufficiency follows from algebraic closure of the Bellman operator $\mathcal{B}$ under EDS. Necessity is proved by a coupling-based fixed-point contradiction: for each failure mode we exhibit an explicit belief pair with equal $\phi$-values but a strictly positive total-variation gap on predictive densities, which propagates to a value-function gap via the Lipschitz bound on $\mathcal{B}$ and the $\rho$-contraction of $\mathcal{T}$, without assuming any structure on $V$ beyond continuity. Three consequences are sharp and quantitative: * **(i)** When EDS fails, *every* scalar rule—regardless of architecture, capacity, or training data—incurs a strictly positive, computable, algorithm-independent lead-time loss. * **(ii)** When EDS holds, optimal expected lead time is given in closed form by a hitting-time functional of a scalar Markov chain. * **(iii)** Under approximate EDS with deviation $\eta$, performance degrades at rate $O(\eta/(1-\rho))$, making $\eta$ a practical model-diagnostic with direct performance implications. Across 72 synthetic configurations and the CICIDS2017 benchmark, predicted and observed lead times agree within 0.4 steps (MAE). A controlled experiment with matched parameter counts confirms that the lead-time advantage is structural, not a regularisation artefact, thereby validating the theory as the sole causal mechanism.

Modern off-policy reinforcement learning algorithms often rely on simple uniform replay sampling and it remains unclear when and why non-uniform replay improves over this strong baseline. Across diverse RL settings, we show that the effectiveness of non-uniform replay is governed by three factors: replay volume, the number of replayed transitions per environment step; expected recency, how recent sampled transitions are; and the entropy of the replay sampling distribution. Our main contribution is clarifying when non-uniform replay is beneficial and providing practical guidance for replay design in modern off-policy RL. Namely, we find that non-uniform replay is most beneficial when replay volume is low, and that high-entropy sampling is important even at comparable expected recency. Motivated by these findings, we adopt a simple Truncated Geometric replay that biases sampling toward recent experience while preserving high entropy and incurring negligible computational overhead. Across large-scale parallel simulation, single-task, and multi-task settings, including three modern algorithms evaluated on five RL benchmark suites, this replay sampling strategy improves sample efficiency in low-volume regimes while remaining competitive when replay volume is high.

Imitation learning (IL)—training an agent to replicate expert behavior from demonstrations—underpins applications from robotics to language model training. Standard approaches such as Behavior Cloning (BC) are known to suffer from compounding errors and performance plateaus, particularly when the learner cannot perfectly represent the expert’s policy (e.g., as is typical in distillation). Two interventions are widely understood empirically to improve performance: querying the expert interactively along the learner’s own trajectories, and using value function estimation en route to generating a policy. We investigate the nature of these improvements and their potentially surprising interplay. Our main finding is that expert interaction relaxes the representational demands on the learner: one only needs a model capable of realizing the expert’s value function, bypassing the (often stricter) requirement of realizing the expert’s policy itself. Concretely, we introduce OVI, an interactive IL algorithm that is statistically and computationally efficient whenever the learner can represent the expert’s value function. We complement this with a negative result showing that interaction is necessary: without significantly stronger representational assumptions than expert-value realizability alone, a broad class of value-based IL algorithms cannot succeed in the offline setting. These findings bear out empirically: OVI outperforms offline policy-based (BC), interactive policy-based (DAgger), and offline value-based IL methods, with the largest gains when the learner network is substantially less expressive than the expert’s.

Causal sensitivity analysis bounds the policy value of off-policy evaluation under unobserved confounding. In offline reinforcement learning, however, sensitivity bounds are typically applied not to raw states but to learned representations. We show that this two-stage pipeline can silently break the coverage guarantee: state aggregation under hidden confounding amplifies the effective sensitivity parameter in the latent space, so bounds computed at the nominal confounding level may exclude the true policy value. We characterise the amplification mechanism, give a sufficient condition for preservation, and prove a continuity result showing the safety condition is not an isolated point: amplification grows continuously with the conditional-independence violation. Neither unsupervised nor standard task-aware representation objectives target the right quantity. We propose Sensitivity-Preserving Representation Learning (SPRL), which augments latent dynamics prediction with a kernel-weighted within-cell propensity homogeneity penalty, an observable surrogate for the unobservable safety condition, and show across synthetic and semi-synthetic benchmarks that SPRL is the only method tested that simultaneously preserves coverage in regimes where unsupervised representations break it and prevents amplification in regimes where task-aware baselines fail.

Many problems are inherently compositional: solving them requires a sequence of intermediate decisions that jointly determine the final answer. A central question is when such a compositional structure can be learned from outcome-level feedback alone, where supervision of the intermediate steps is unavailable. We study this question for autoregressive models trained using reinforcement learning with verifiable rewards (RLVR). We identify the \emph{task-advantage ratio}, a joint property of the task and the model, which measures whether intermediate decisions present an advantage in reaching a correct final solution. We show that this ratio governs learnability in our setting: when the advantage is present, RLVR efficiently learns the target composition, while when it is absent, training can converge to suboptimal compositions. We further show that the required advantage arises naturally in several structured problems, but may depend critically on the quality of the initial model. Our results help to clarify when compositional problems can be learned from final rewards alone.