NeurIPS 2026 Papers

We introduce a new tool, Express, for converting a non-causal attention approximation into a causal approximation with matching approximation guarantees. When combined with the state-of-the-art Thinformer approximation, Express improves upon the best known causal attention guarantees, delivering $\log^{3/2}(n)/s$ approximation error with only $O(s)$ memory and $O(s^2 \log^2(n))$ compression overhead for a sequence of length $n$. We pair these developments with an efficient I/O-aware Triton implementation, demonstrate substantial speedups over FlashAttention 2, and use Express to accelerate three distinct components of the language modeling pipeline: long-context prefill, KV cache compression, and long-form decoding.

Reinforcement learning (RL) agents face a continuous trade-off between exploiting known strategies and exploring novel states to obtain higher rewards. While exploration is often undirected, it can also be guided by intrinsic signals such as novelty or uncertainty. However, combining these signals with task rewards is challenging, as they can sometimes drive the agent away from optimal trajectories and reduce final performance. We find that this degradation arises from a structural occupancy conflict: as the policy improves, exploration signals naturally begin to anti-correlate with task advantages. In this work, we introduce Kantorovich Dual Shaping (KDS), a framework that resolves this conflict through constrained policy-space optimization. KDS formulates exploration as an Optimal Transport problem and applies the Sinkhorn-Knopp algorithm on the empirical batch manifold to reweight intrinsic rewards in a geometry-aware, point-by-point manner. When combined with standard RL objectives, it acts as a spatial filter that suppresses exploration in conflicting regions while enhancing it elsewhere. We evaluate this general approach across multiple off-policy and on-policy algorithms, achieving stable exploration that consistently avoids task degradation in complex continuous control environments.

Explainability, the capability of an artificial intelligence system (AIS) to explain its outcomes in a manner that is comprehensible to human beings at an acceptable level, has been deemed essential for critical sectors, such as healthcare. Is it really the case? In this perspective, we consider two extreme cases, Oracle (without explainability) versus AI Colleague (with explainability) for a thorough analysis. We discuss how the level of automation and explainability of AIS can affect the determination of liability among the medical practitioner/facility and manufacturer of AIS. We argue that explainability plays a crucial role in setting a responsibility framework in healthcare, from a legal standpoint, to shape the behavior of all involved parties and mitigate the risk of potential defensive medicine practices.

Pre-trained diffusion and flow policies have emerged as powerful backbones for robotics, yet adapting them to novel tasks via reinforcement learning (RL) remains computationally unstable, sample inefficient and lacks theoretical guarantees. Prior approaches that fine-tune the latent noise space bypass costly backpropagation through time but struggle with "noise aliasing," where determining the optimal noise input to the pre-trained generative model requires repeatedly refitting separate Q-value networks. This paper resolves these challenges by introducing a novel theoretical framework demonstrating that imposing a linear structure on the environment induces a Latent Noise Linear Markov Decision Process (LNL-MDP). By treating the pre-trained model as a virtual dataset, our LNL-MDP fine tune framework falls under the FineTuneRL theory with near-optimal guarantee. To translate this idealized theory into a practical methodology, we derive the Bridge Equation which eliminates the noise aliasing problem. Leveraging this insight, we propose Latent Representation Bridging for Diffusion (LARBRID), an algorithm that efficiently learns the representation space from the LNL-MDP and optimal noises to steer the pre-trained model to act optimally in given tasks. We demonstrate that LARBRID is exceptionally sample efficient, yielding up to $2\text{-}3\times$ performance increases on challenging control problems across a diverse suite of robotic locomotion and manipulation tasks.

We propose Drifting Field Policy (DFP), a non-ODE one-step generative policy built on the drifting model paradigm. Under online RL finetuning, we frame the policy update as a reverse-KL Wasserstein-2 gradient flow toward a soft target policy, realized by DFP as a steepest-descent step on probability space. By construction, this gradient step is decomposed into an ascent toward higher action-value regions and a score matching with the anchor policy as a trust region. We further derive a simple, tractable surrogate of the otherwise intractable update loss, akin to behavior cloning on top-K critic-selected actions. We find empirically that this mechanism uniquely benefits the drifting backbone owing to its non-ODE parameterization. With one-step inference, DFP achieves state-of-the-art performance on several manipulation tasks across Robomimic and OGBench, outperforming ODE-based policies.

Offline reinforcement learning requires improving a policy from fixed data while avoiding out-of-distribution actions with unreliable value estimates. Diffusion and flow policies handle this trade-off by modeling the behavior distribution to regularize the RL objective, but they require iterative denoising, solver integrations, and in more efficient variants, distillation or other approximations at inference. We propose DriftQL, which combines a drift-based behavioral regularizer with critic-driven policy improvement. The value signal biases the policy toward high-value regions of the data support, while attraction and repulsion together keep generated actions near the data and prevent collapse onto a single mode. DriftQL is implemented as a single network with a unified training objective and generates actions in a single forward pass. On D4RL and OGBench, DriftQL consistently outperforms diffusion and flow methods, advancing the state of the art. Under degraded data quality, where the baselines visibly struggle, DriftQL remains close to its clean-data performance, positioning it as a promising alternative to diffusion and flow-based methods while maintaining the simplicity and efficiency of deterministic approaches.

Reinforcement-learning-based post-training for large language models vastly improves capabilities on verifiable tasks however it suffers from training instability and diversity collapse. As a solution, many new advantage functions have emerged. Each method simultaneously changes which problems receive gradient, the balance between positive and negative updates, and the overall gradient scale making their comparison difficult. We propose a unifying framework that decomposes each function by their induced positive and negative gradient mass. Analyzing policy performance and the geometry of model weights reveals: symmetric methods are better but slow and asymmetric methods are fast. Combining the advantages of both, we build the FADE: Focal Adaptive Dynamic Entropy advantage which is fast, diverse and accurate at various model scales (7B, 32B).

Policy gradient computes a backward pass for every sample, even though the backward pass is expensive and most samples carry little learning value. The Delightful Policy Gradient (DG) provides a forward-pass signal of learning value: \emph{delight}, the product of advantage and surprisal (negative log-probability). We introduce the \emph{Kondo gate}, which compares delight against a compute price and pays for a backward pass only when the sample is worth it, thereby tracing a quality--cost Pareto frontier. In bandits, zero-price gating preserves useful gradient signal while removing perpendicular noise, and delight is a more reliable screening signal than additive combinations of value and surprise. On MNIST and transformer token reversal, the Kondo gate skips most backward passes while retaining nearly all of DG's learning quality, with gains that grow as problems get harder and backward passes become more expensive. Because the gate tolerates approximate delight, a cheap forward pass can screen samples before expensive backpropagation, suggesting a speculative-decoding-for-training paradigm.

Policy learning from observational data is fundamental to applications ranging from personalized medicine to welfare-program targeting, but standard approaches optimize empirical performance and may fail when the deployment distribution differs from the training distribution. We study \emph{distributionally-robust policy learning}: given observational data $\\{(X_i,W_i,Y_i)\\}_{i=1}^n$ with binary treatment $W$, we seek a policy $\pi \in \Pi$ that maximizes the worst-case policy value over an $f$-divergence ball of radius $\rho$ around the data-generating distribution. We propose an estimator built on cross-fitted doubly-robust scores for the conditional treatment effect, paired with a DRO solver that handles \texttt{KL}, $\chi^2$, and \texttt{CVaR} ambiguity sets. We establish two convergence guarantees. First, the estimated policy attains the standard $\widetilde{O}(n^{-1/2})$ rate on its DRO regret, and inherits the doubly-robust property from its non-DRO counterpart -- the rate continues to hold whenever either the outcome model or the propensity model is consistently estimated, even at slow nonparametric rates. We then propose a sample-splitting bias-correction algorithm derived from the Lagrangian dual, and show that under a Tsybakov-style margin condition on the worst-case treatment-effect contrast, the estimator attains a fast rate of $\widetilde{O}(n^{-(1+\alpha)/(2+\alpha)})$ , interpolating between $\widetilde{O}(1/\sqrt{n})$ and $\widetilde{O}(1/n)$ as the margin parameter $\alpha$ ranges over $[0,\infty)$. To our knowledge, this is the first work to incorporate a Tsybakov margin condition into the analysis of distributionally-robust policy learning, and the first to obtain fast rates in this setting. Empirically, on semi-synthetic experiments calibrated to a welfare-to-work program, distributional robustness yields meaningful improvements over empirical welfare maximization when training and test distributions differ structurally, while remaining competitive when no shift is present.

EHR foundation models create new opportunities for scalable clinical prediction across health systems, yet their deployment remains constrained by a basic tokenization bottleneck. The same clinical concept can appear as different ontology codes, local identifiers, synonyms, or free-text descriptions, while standard EHR tokenizers treat these variants as unrelated symbols. We introduce MEDIATOR, a discrete ontology-agnostic medical concept tokenizer that turns EHR harmonization into an open-form tokenization problem. Instead of relying on manual mappings, a fixed ontology vocabulary, or continuous event embeddings, MEDIATOR maps any clinical concept expressed as text into a shared vector-quantized codebook. MEDIATOR is trained with synonym supervision from paired concept descriptions, and further guided by enriched concept descriptions, empirical code co-occurrence, and curated relation graphs so that equivalent expressions collapse while non-equivalent concepts remain organized in a medically meaningful discrete space. The resulting vocabulary is compact, inspectable, reusable by standard EHR Transformers, and open to previously unseen code descriptions or text strings. Used as the token interface of an EHR Transformer, MEDIATOR improves mean AUPRC on 20 MIMIC-III tasks by 4.0% and improves zero-shot transfer to 14 eICU tasks by 8.8%, while supporting transferable attribution across hospitals.

Distributed reinforcement learning trains on data from stale, buggy, or mismatched actors, producing actions with high surprisal (negative log-probability) under the learner's policy. The core difficulty is not surprising data per se, but \emph{negative learning from surprising data}. High-surprisal failures can dominate finite-batch updates through large perpendicular components, while high-surprisal successes reveal opportunities the current policy would otherwise miss. The \textit{Delightful Policy Gradient} (DG) separates these cases by gating each update with delight, the product of advantage and surprisal, suppressing rare failures and preserving rare successes without behavior probabilities. In a tabular analysis, DG suppresses the perpendicular second moment of high-surprisal failures by a policy-overlap factor that vanishes as the learner improves. The advantage sign is essential for surprisal-based filtering: any learner-probability-only gate that suppresses rare failures also suppresses rare successes. On MNIST with simulated staleness, DG without off-policy correction outperforms importance-weighted PG with exact behavior probabilities. On a transformer sequence task with staleness, actor bugs, reward corruption, and rare discovery, DG often achieves nearly order-of-magnitude lower error. When all four frictions act simultaneously, its sample-efficiency advantage is order-of-magnitude and grows with task complexity.

Imitation learning has achieved remarkable progress in robot manipulation, but its success heavily depends on large-scale demonstration data that are costly to collect. Recent work has therefore explored demonstration augmentation, but they are fundamentally limited in deformable manipulation, where task-relevant object variation is governed by high-dimensional deformations and physics-induced internal constraints rather than low-dimensional pose changes. We present **DeformGen**, a dynamics-based augmentation framework that achieves *topological diversity* for deformable objects. Instead of perturbing object pose, DeformGen expands the valid state distribution by applying localized physical disturbances, forward-simulating the dynamics, and stabilizing the result to obtain topology-coherent deformable states with physical plausibility. Given these synthesized states, DeformGen further transfers source manipulation trajectories via deformation-field warping, which lifts per-particle displacements into a continuous spatial function to adapt the end-effector trajectory consistently with the deformed geometry. In this way, our method augments both the state distribution and its associated manipulation behavior. Experiments on high-fidelity deformable manipulation benchmarks show that DeformGen consistently improves policy learning over training on the original demonstrations alone and over rigid-style augmentation baselines.

In this position paper, we argue that successful decentralized AI needs trustworthy, adaptive, and portable governance at scale. The ML community must adopt an architecture that cleanly separates \emph{policy processing} (the evaluation of evidence against governance requirements) from \emph{capability enforcement} (the gating of access to protected digital assets). We argue that the prevailing approach, in which each organization manages bespoke policy logic, custom-built for its infrastructure, produces fragmentation that undermines interoperability, transparency, and auditability of federated systems. Drawing on a survey of governance mechanisms in federated learning and data collaboration frameworks and on a concrete reference architecture built around community-driven policy objects, we argue that this decoupling is technically feasible and architecturally necessary and supports a broad class of decentralized AI governance settings. We further suggest that the AI community should invest in open, standardized policy abstractions rather than proliferate siloed governance solutions. We address counterarguments concerning the costs of standardization and the feasibility of universal policy languages.

We present DeFlow, a decoupled offline RL framework for extracting high-value actions from a learned multi-step flow prior. Directly optimizing iterative generative policies typically requires backpropagation through ODE solvers, while shortcut policies trade off expressivity, inference cost, and policy-improvement objectives. DeFlow keeps the flow model as a behavior prior and trains a lightweight action-conditioned residual module for value improvement under an adaptive trust-region penalty. This design bypasses solver differentiation and encourages proximity to the learned prior without treating the trust region as exact data-support preservation. Empirically, DeFlow is competitive with recent generative offline RL baselines, with the clearest gains on selected multimodal manipulation and offline-to-online settings.

Accurate forecasting of Earth's dynamic systems is vital in geophysics, where neural operators offer a powerful data-driven paradigm. Standard operators often treat the Earth as a flat Euclidean plane, introducing geometric artifacts. In contrast, Spherical Neural Operators (SNOs) resolve this mismatch by operating within the spherical spectral space, naturally adapting to spherical geometries. However, for computational efficiency, existing SNOs typically rely on the rotation equivariance assumption. This assumption simplifies the transitions of physical states into isolated spectral modal evolutions, thereby over-idealizing the Earth as a homogeneous system. To solve this problem, we no longer rely on the rotation equivariance assumption that isolates spectral modal evolution, but directly derive an explicit modal interaction kernel and propose Diagonal Spherical Neural Operators named DIAGNO. DIAGNO explicitly parameterizes cross-modal interactions through spectral modal decomposition, thereby preserving the inherent heterogeneity of Earth dynamics. Furthermore, leveraging the unique diagonal structure of the spherical spectrum, we design inter- and intra-diagonal interaction mechanisms to natively capture essential zonal and meridional dynamics, respecting intrinsic geophysical regularities. Extensive experiments across simulated fluid dynamics and real-world Earth system reanalysis (including atmosphere and ocean) datasets demonstrate DIAGNO’s superior capability, consistently achieving state-of-the-art performance in multi-step forecasting. Beyond this, we validate DIAGNO’s practical analytical value in professional geophysical domains by assessing its fidelity in capturing climatological anomalies and maintaining physical consistency.

Subgraph mining aims to discover meaningful subgraphs from a large graph given a seed node, with applications in community detection, functional module identification, anti-money laundering, and knowledge discovery. Existing methods typically adopt static approaches that cannot dynamically and intelligently adapt to diverse graph structures. Moreover, training models for subgraph mining faces inherent challenges including exposure bias and the difficulty of learning when to stop expansion. To address these limitations, we propose DASM (Dynamic Autoregressive Subgraph Mining), a novel two-stage framework that models subgraph mining as a sequential decision process. In the first stage, we employ step-wise supervised learning with multi-label objectives and asymmetric F1 loss to teach the model how to select correct next nodes. In the second stage, we apply Group Relative Policy Optimization (GRPO) to align the model with end-to-end mining quality, enabling it to learn optimal stopping decisions through trajectory-level feedback. Our dynamic autoregressive approach allows the model to make state-aware decisions at each step, automatically determining subgraph boundaries without predefined size constraints. Experiments on multiple datasets demonstrate that DASM significantly outperforms existing baselines in terms of F1 score and IoU.

While recent work in Reinforcement Learning with Verifiable Rewards (RLVR) has shown that a small subset of critical tokens disproportionately drives reasoning gains, an analogous token-level understanding of On-Policy Distillation (OPD) remains largely unexplored. In this work, we investigate high-loss tokens, a token type that—as the most direct signal of student-teacher mismatch under OPD's per-token KL objective—should progressively diminish as training converges according to existing studies; however, our empirical analysis shows otherwise. Even after OPD training reaches apparent saturation, a substantial subset of tokens continues to exhibit persistently high loss; these tokens, which we term Rock Tokens, can account for up to 50\% of the tokens in generated outputs. Our investigation reveals two startling paradoxes. First, despite their high occurrence frequency providing a disproportionately large share of total gradient norms, Rock Tokens themselves remain stagnant throughout training, resisting teacher-driven corrections. Second, through causal intervention, we find that these tokens provide negligible functional contribution to the model's actual reasoning performance. These findings suggest that a vast amount of optimization bandwidth is spent on structural and discourse residuals that the student model cannot or need not internalize. By deconstructing these dynamics, we demonstrate that strategically bypassing these "stumbling blocks" can significantly streamline the alignment process, challenging the necessity of uniform token weighting and offering a more efficient paradigm for large-scale model distillation.

Learning from multiple imperfect oracles is a practical way to improve reinforcement learning (RL) under limited interaction. However, most contemporary robust multi-oracle methods are on-policy and equire frequent fresh rollouts, which limits sample efficiency in domains where data collection is expensive. We propose Oracle-Aggregated Policy Improvement (OPI), an off-policy framework for policy improvement with multiple suboptimal oracles. OPI augments the oracle set with the learner, constructs an oracle-guided behavior policy through candidate-action pooling and source-aware critic scoring, and trains the learner using a generalized advantage-weighted behavior cloning objective combined with deterministic policy gradients. The method is motivated by a tractable proxy for the max-advantage objective and a policy-improvement relative to a strong oracle baseline. We evaluate OPI on diverse tasks across MetaWorld, the DeepMind Control Suite, and customized Box2D environments with heterogeneous oracles. Across sparse-reward manipulation, dense-reward locomotion, and customized control settings, OPI is strongest in sparse-reward and complementary-oracle regimes and remains competitive with representative baselines in dense-reward environments. Additional analyses show that oracle composition is most beneficial when oracle skills are complementary and that estimating oracle-specific action values is particularly helpful in sparse-reward regimes

RLVR and OPD have become standard paradigms for post-training. We provide a unified analysis of these two paradigms in consolidating multiple expert capabilities into a single model, identifying capability loss in different ways: mixed RLVR suffers from inter-capability divergence cost, while the pipeline of first training experts and then performing OPD, though avoiding divergence, fails to fully absorb teacher capabilities due to large behavioral pattern gaps between teacher and student. We propose **Co-Evolving Policy Distillation (CoPD)**, which encourages parallel training of experts and introduces OPD during each expert's ongoing RLVR training rather than after complete expert training, with experts serving as mutual teachers (making OPD bidirectional) to co-evolve. This enables more consistent behavioral patterns among experts while maintaining sufficient complementary knowledge throughout. Experiments validate that CoPD achieves all-in-one integration of text, image, and video reasoning capabilities, outperforming strong baselines such as mixed RLVR and MOPD, and even surpassing domain-specific experts.