NeurIPS 2026 Papers

On-Policy Distillation (OPD) has emerged as a dominant post-training paradigm for large language models, especially for reasoning domains. However, OPD remains unstable in practice due to the high gradient variance of its single-sample Monte Carlo estimator, and recipes for stable training are still immature. We propose vOPD (On-Policy Distillation with a control variate baseline), which casts OPD as policy-gradient RL and stabilizes it by introducing a control variate baseline—canonically a value function—from the RL literature. We show that the OPD value function admits a closed form as the per-token negative reverse KL divergence between the student and the teacher, available directly from the already-computed forward pass with no additional critic or inference. Existing stabilization methods either compute the full token-level reverse KL over the entire vocabulary, adding significant overhead, or restrict it to a top-k support, biasing the objective. vOPD instead preserves the lightweight single-sample estimator, subtracting the value function as a detached baseline to keep the gradient unbiased while reducing variance. Furthermore, we show that a top-k approximation of the baseline further lowers cost without compromising performance. Across mathematical and scientific reasoning benchmarks, vOPD consistently outperforms vanilla OPD and matches the most expensive full-vocabulary baseline, offering an efficient stabilization of On-Policy Distillation through principled RL variance reduction.

Concept erasure is a prominent approach to achieving fairness in machine learning through the removal of sensitive attributes from learned representations. Prior concept erasure methods typically define debiasing indirectly through the failure of a chosen adversary, probe, or fairness penalty, making the target of erasure dependent on a particular decoder family or optimization setup. Rather than depending on a particular architecture, task type, or training paradigm, we introduce KINDER, which aims to directly remove sensitive-attribute information from a model's intermediate representations and is compatible with any deep model that contains an intermediate feature space. KINDER operates in a random Fourier feature space, where it estimates a sensitive subspace from protected group prototypes and projects representations onto its orthogonal complement, thereby forming an explicit representation cleaning mechanism. We conduct extensive experiments across diverse settings and show the effectiveness of KINDER on (i) unimodal and multimodal datasets, (ii) supervised and self-supervised settings, (iii) classification, regression, and image segmentation tasks, and (iv) diverse data modalities, including visual, textual, and tabular data.

Large language models (LLMs) are increasingly used for knowledge base question answering (KBQA), where answering requires selecting entities from a question-specific knowledge-graph subgraph. Yet LLMs are known to hallucinate across tasks, and KBQA is no exception: even when we provide a graph as the knowledge source, the model may rely on parametric knowledge instead of graph evidence or perform invalid reasoning over the given relations. Such hallucinated answer nodes can limit the practical deployment of KBQA systems, especially in high-stakes domains such as healthcare. We formulate hallucination detection in KBQA as an answer-node classification problem and propose a lightweight graph-based framework that treats the answering LLM as a black box. KG-Guard represents each KBQA instance as an augmented graph. It initializes node features with semantic representations of KG entities, marks topic entities and LLM-proposed answer nodes with learned vectors, and connect a virtual question node to the topic entities. A graph encoder then produces verification-oriented node representations, and a small MLP classifies each proposed answer node using its graph representation together with the question embedding. Experiments on WebQSP, ComplexWebQuestions, and PUGG show that our detector achieves the highest F1 on all three benchmarks ($82.0$, $87.4$, and $84.3$), outperforming LLM-judge and sampling-based baselines, while having $\sim305\times$ fewer parameters than the reference approaches. Beyond detection, the node-level feedback is actionable: when flagged answers are fed back to the KBQA system for iterative refinement, downstream KBQA F1 improves by $13.0$-$14.5$ points and Exact Match by $16.9$-$17.6$ points.

Liability adjudication tasks (e.g., legal charge determination), require analyzing event descriptions against domain-specific skills to determine whether a party should bear liability. While Large Language Models (LLMs) have shown promise in this area by leveraging executable skills for interpretable reasoning, existing approaches face three critical limitations: **skill libraries become stale as real-world scenarios dynamically evolve**, **skill descriptions often suffer from a semantic execution gap where LLMs deviate from intended expert logic**, and **growing skill collections exceed context budgets causing critical skill omission**. To address these challenges, we propose **K**nowledge-**D**ata Dual-Driven **A**utomated **S**kill **O**ptimization (KDASO), which continuously evolves a skill library through neighborhood-conditioned induction from recent adjudication data, execution-aligned refinement based on LLM trajectory analysis, and context-budgeted routing for scalable skill selection. Experiments across ride-hailing adjudication, legal liability determination, and logical reasoning judgment demonstrate that KDASO significantly outperforms current general-purpose and domain-specific baselines. KDASO further achieves an 8\% improvement when transferring skills to concept-similar scenarios and maintains consistent performance gains through continuous optimization, underscoring the high transferability and stability of generated skills.

Large language models (LLMs) are increasingly deployed in K–12 education, yet existing benchmarks such as C-Eval, CMMLU, GaokaoBench, and EduEval measure only whether a model can answer an exam question, i.e., factual recall. Effective educational AI further requires curriculum cognition: the structured understanding of how knowledge is organized, including prerequisite chains, concept taxonomies, experiment–concept links, and pedagogical sequencing. Curriculum cognition is neither probed by current benchmarks nor explicitly taught by current instruction-tuning data. To close this gap, we introduce K12-KGraph, a curriculum-aligned knowledge graph extracted from the official People's Education Press textbooks, covering mathematics, physics, chemistry, and biology across primary, middle, and high school, with seven node types (Concept, Skill, Experiment, Exercise, Section, Chapter, Book) and nine relation types spanning taxonomy, prerequisite, association, verification, assessment, location, and order. Building on this single graph, we derive two complementary resources: K12-Bench, a 23,640-question multi-select benchmark across five graph-derived task families (Ground, Prereq, Neighbor, Evidence, and Locate) that jointly probe curriculum cognition; and K12-Train, a KG-guided supervised fine-tuning corpus of ~2,300 QA pairs synthesized from node attributes and edge semantics. Experiments expose a clear gap and a clear remedy: on K12-Bench, even the strongest proprietary model (Gemini-3-Flash) reaches only 57% exact match and the best open-source model (Gemma-4-31B-IT) only 46%, with Prereq and Neighbor being the hardest; yet under a strictly matched 2,300-sample SFT budget on Qwen3-4B-Base and Llama3.1-8B-Base, K12-Train consistently outperforms equally sized subsets of eight mainstream instruction-tuning corpora (OpenHermes, Infinity, UltraChat, WizardLM, DataFlow, LMSYS, SmolTalk, Tulu-3) on both GaokaoBench and EduEval, showing that structural curriculum grounding is remarkably sample-efficient for educational SFT. We release the graph, benchmark, training data, and the full construction pipeline.

World models enable prediction, planning, and decision-making by learning internal simulators of environment dynamics. Although substantial progress has been made in world models from pixel-based observations, many physical systems are instead observed through vector-based measurements that are noisy and only partially informative of the underlying state. We introduce K-PWM, a Koopman-structured probabilistic world model for prediction and control under partial observation. K-PWM combines a parameterized Koopman-structured state-space model with a nonlinear observation decoder, and performs probabilistic latent-state inference directly through the learned model parameters rather than through a separately trained inference network. We train K-PWM using a generalized expectation conditional maximization (GECM) procedure with principled initialization. Across Gymnasium MuJoCo tasks, K-PWM improves sample efficiency, yields reliable long-horizon predictions, and supports downstream control through model predictive path integral (MPPI) control. These results suggest that structured probabilistic latent dynamics provide a useful route toward data-efficient world modeling and control in partially observable vector-valued settings.

LLM agents, which often comprise parallel inference tasks, are commonly adopted to solve real-world problems. When serving such task-parallel LLM agents in shared GPU servers, the scheduler is expected to attain fast agent completion with guaranteed worst-case performance. For that objective, our insight is to selectively pampering agents based on their completion order under idealized fair-sharing. We design Justitia, a fair and also efficient scheduler for task-parallel LLM agents. Noticing that memory is prevalently a bottleneck in LLM serving, Justitia quantifies the true agent cost in a memory-centric manner. It also adopts a light-weight yet accurate method to predict agent costs. Finally, Justitia adopts a virtual-time based fair queuing algorithm to reduce the overall performance with guaranteed worst-case delay. We have implemented Justitia atop vLLM, and experimental results involving diverse agents show that it can substantially enhance the scheduling efficiency with fairness preserved.

Network interference complicates A/B testing on online platforms, such as social networks and marketplaces, where causal methods based on single experiments often suffer from significant bias due to complex interference patterns. This paper demonstrates the statistical benefits of merging data from multiple experiments with varying treatment proportions. Sequential experimentation with increasing traffic, or ramp-up, is widely used in tech companies for risk management and cost control. Beyond operational benefits, we show that regression-based estimators trained on merged data achieve substantial bias reduction, even under simple randomization schemes and regression models. We focus on the global average treatment effect (GATE), a key estimand in the tech industry, and consider a general interference pattern that extends beyond the 1-hop setting. We present an exact bias–variance analysis of the linear regression estimator and show that, in practical settings, the bias term typically dominates. In addition, we characterize the substantial bias reduction achieved by our proposed approach, which merges experimental data collected at different ramp-up stages to improve the training of the regression function. We also offer an intuitive explanation for this reduction and highlight the synergy between cluster-level randomization and our approach. Furthermore, we consider a refined estimator based on graph neural networks (GNN). Extensive simulations across novel and challenging scenarios confirm that our methodology significantly improves the accuracy of regression-based estimators.

Reliable progress in offline policy learning depends on a small set of methodological practices, including careful reporting, well-tuned baselines, and evaluation across diverse conditions. Prior work has noted that results can be sensitive to reporting choices, hyperparameter tuning, and dataset properties independently, but these sources of variability have not been systematically investigated at the scale needed to understand how they shape conclusions. To address this gap, we present a large-scale empirical study of offline reinforcement and imitation learning, training over 145,000 policies across 114 datasets. At this scale, no algorithm dominates: aggregate performance across top methods is often close, but the leaders differ substantially across environments. We find that proper hyperparameter tuning frequently reshuffles perceived algorithm rankings; and simple baselines, including behavior cloning, are often stronger than is commonly assumed after extensive tuning. We also study hyperparameter transfer and sensitivity across environments, identifying a simple strategy that generalizes well. From these analyses we distill practical recommendations, and release JumpStart: a resource suite of trained models, per-model hyperparameter and reward data, strong baselines across all environments, and a website to make retrieval and analysis trivial. Together, these resources aim to make offline policy learning research more reliable and to open new directions for work beyond the scope of this study.

Two methodologies dominate current practices of benchmarking: rubric-based scoring evaluates items against predefined criteria, whereas comparative judgment elicits pairwise preferences between outputs. Although both methodologies are widely used, the choice between them is rarely justified. We release JudgmentBench, a benchmark of 30 real-world legal tasks, paired with 1,539 rubric scores and 1,530 pairwise preference judgments collected from practicing attorneys—including at major U.S. law firms—with substantial experience. The annotations constitute the first publicly available dataset in a high-expertise domain in which both supervision signals are elicited from the same experts on the same items. Using LLM-generated outputs at three constructed quality levels, we provide an initial empirical comparison: comparative judgments recover the intended quality ordering substantially better than rubrics (mean Spearman's rank correlation of 0.908 vs. 0.150, $\widehat{D}=0.758$ [0.494, 1.021]) while requiring less than half the annotation time. The patterns hold for human annotators and LLM autograders. Beyond this initial comparison, the paired structure of the dataset supports a broader research agenda on how expert judgment should be elicited, aggregated, and used as supervision in domains without verifiable ground truth.

As large language models (LLMs) are increasingly deployed to users around the world, they are integrated into everyday tasks across diverse cultural contexts, from drafting personal communications to brainstorming creative ideas. These tasks are inherently cultural: they require contextual appropriateness, symbolic resonance, and tacit cultural expectations that native speakers draw on instinctively, meaning that a response can be factually plausible yet unmistakably wrong to a local reader. Existing cultural benchmarks have treated culture as a flat set of facts via fact verification or norm entailment methods, and have adopted LLM-as-a-Judge without examining whether they can capture such thick cultural errors. To address this gap, we present JuICE (Benchmark for LLM-Judge in Identifying Cultural Errors), a multilingual dataset of 7,470 span-level annotations of cultural and linguistic errors, collected from native speakers in long-form LLM responses. It covers 1,050 query-response pairs from four countries (United States, South Korea, Indonesia, and Bangladesh), in both English and their countries' main languages. Using JuICE, we find that even the strongest LLM-judge achieves only an F1 of 0.52 against human-annotated errors, and that this gap is systematic. LLM-judges reliably detect surface-level linguistic and factual mistakes but consistently miss thick cultural errors that local residents readily identify. Our findings suggest that robust cultural evaluation must move beyond surface-level detection toward frameworks that account for the depth and situatedness of cultural meaning.

Fairness in machine learning typically relies on the assumption of clean training data. However, in high-stakes applications like healthcare and criminal justice, privacy mechanisms and proxy variables frequently introduce simultaneous noise into both sensitive attributes and target labels. While existing methods address label or attribute noise in isolation, they fail under this "double noise" regime, often inadvertently amplifying bias. In this paper, we propose $\textbf{Jointly Robust Fairness (JRF)}$, a unified framework that guarantees fairness under simultaneous data corruption. JRF integrates a Forward Loss correction to handle label noise with a novel Matrix-Inverse Robust Loss (MIRL) that algebraically recovers the true demographic parity gap from noisy attribute observations. We provide rigorous theoretical guarantees for our estimator, including a finite-sample concentration bound demonstrating that the sample complexity scales quadratically with the inverse of the noise transition determinant. Furthermore, our framework naturally generalizes to both symmetric and highly asymmetric noise distributions. Extensive experiments on the Adult, Bank, and COMPAS datasets demonstrate that JRF reduces fairness violations by over 91\% in high-noise regimes ($\rho \ge 0.3$) compared to standard robust baselines, maintaining strict fairness without significant degradation in predictive utility.

Post-training of Large Language Models (LMs) often prioritizes accuracy and helpfulness at the expense of diversity. This creates a tension: while post-training improves response quality, it also sharpens output distributions and reduces the range of ideas, limiting the usefulness of LMs in creative and exploratory tasks such as brainstorming, storytelling, or problem solving. We address this challenge with Diversity-Aware Reinforcement Learning (DARLING), a framework that jointly optimizes for response quality and semantic diversity. At its core, DARLING introduces a learned partition function to measure diversity beyond surface-level lexical variations. This diversity signal is then combined with a quality reward during online reinforcement learning, encouraging models to generate outputs that are both high-quality and distinct. Experiments across multiple model families and sizes show that DARLING generalizes to two regimes: non-verifiable tasks (instruction following and creative writing) and verifiable tasks (competition math). On five benchmarks in the first setting, DARLING consistently outperforms quality-only RL baselines, producing outputs that are simultaneously of higher quality and novelty. In the second setting, DARLING achieves higher pass@1 (solution quality) and pass@k (solution variety). Most strikingly, explicitly optimizing for diversity catalyzes exploration in online RL, which manifests itself as higher-quality responses.

Target trial emulation (TTE) provides a framework for answering causal questions using observational data when randomized controlled trials (RCTs) are infeasible. However, standard methods for treatment effect estimation have been developed in isolation, failing to jointly address the compounding challenges inherent to analyses of observational data such as electronic health records (EHRs). In particular, these challenges include time-varying confounding and missing-not-at-random (MNAR) missingness reaching 50\%-80\% for critical biomarkers. To address this gap, we propose a two-stage pipeline that jointly handles MNAR missingness and causal structure across time. CausalFlow-T, a Directed Acyclic Graph (DAG)-constrained normalizing flow with Long Short-Term Memory (LSTM)-encoded patient history, performs exact invertible counterfactual inference, eliminating the approximation errors and confounding biases where existing variational and adversarial methods fail silently. Ablations on four synthetic and one semi-synthetic dataset with known counterfactuals confirm that its two core design choices address strictly non-overlapping failure modes (DAG constraints for confounding separation, exact inference for structural propagation) with neither compensating for the absence of the other. To handle the incomplete data CausalFlow-T receives as input, we propose an LLM-driven evolutionary imputer and evaluate it with three LLM backends, including two open-source models. Across 30\%-80\% MNAR missingness, the imputer achieves the best pooled rank across biomarker and causal metrics, leading on point-wise accuracy and temporal extrapolation while maintaining average treatment effect (ATE) recovery where statistical methods progressively degrade. Applied to a cohort of adults with type 2 diabetes in Swiss primary care initiating a GLP-1 receptor agonist or SGLT-2 inhibitor, the pipeline recovers a per-protocol weight-loss difference of $-0.98$ kg [$95\%$ CI $-1.01$, $-0.96$] favoring GLP-1 receptor agonists, consistent with RCT evidence and estimated directly from realistically incomplete real-world EHR data.

Selective knowledge distillation has become increasingly important for compressing large language models, where dense logit-based knowledge transfer over long sequences and large vocabularies incurs substantial computational overhead. Distillation supervision is naturally organized over a sequence–vocabulary matrix, with its highest-information components often concentrated on a small subset of specific sequence–vocabulary pairs. Existing methods typically select sequence positions or vocabulary candidates independently, relying on coarse-grained one-dimensional heuristics that weaken the precision and effectiveness of knowledge transfer. To address this, we propose $\textbf{J}$oint $\textbf{S}$equence–$\textbf{V}$ocabulary $\textbf{D}$istillation ($\textbf{JSVD}$), a plug-and-play selective distillation framework that identifies and focuses supervision on informative sequence–vocabulary pairs. For each training example, JSVD constructs bidirectional teacher--student candidate supports, scores sequence--vocabulary pairs using prediction discrepancies, and progressively focuses distillation on unresolved informative pairs during training. Extensive experiments across diverse distillation objectives, model pairs, and benchmarks show that JSVD consistently improves downstream performance while reducing distillation overhead and accelerating convergence.

Tool-augmented language models are bounded by the APIs humans bothered to write; existing tool-creation systems patch this by prompting a frozen LLM at inference time, leaving the model that writes a tool decoupled from the one that uses it, with no signal that the schemas it produces are schemas it can invoke. We propose **SMITH** (Schema-grounded Multi-task Iterative Tool Honing), a reinforcement learning framework that jointly trains tool creation and tool use inside a single policy. Each rollout is either a **build** task (write a tool from a few examples) or a **use** task (invoke a pooled tool on a held-out question). Three separate reward axes catch schema, code, and outcome failures independently, so each failure mode contributes its own gradient. A 4B Qwen3 trained with SMITH on 13 procedural reasoning tasks with exact verifiers reaches $79.8$ macro-average accuracy on held-out tasks, the best across all evaluated methods and ahead of an untrained 30B-A3B tool-writer. It also reaches $40.4$ on TabMWP-Hard and $42.6$ on out-of-domain GQA ($+7.6$ over the best same-backbone inference-time baseline), without any visual or tabular training data. When invoked by a frozen 350M student, tools written by our 4B match those produced by a writer an order of magnitude larger. The same recipe also lifts Qwen3-8B and Granite-3.3-8B without modification.

Building agents that can perform new skills by composing existing skills is a long-standing goal of AI agent research. Towards this end, we investigate how to efficiently acquire a sequence of skills, formalized as hierarchical neural options. However, existing model-free hierarchical reinforcement algorithms need a lot of data. We propose a novel method, which we call "AgentOWL" (Option and World model Learning), that jointly learns --- in a sample efficient way --- an abstract world model (abstracting across both states and time) and a set of neural options. We show, on a subset of Object-Centric Atari games, that our method can learn more skills using less data than baseline methods and possesses learning and generalization capabilities that the baselines do not have.

Reconstructing the connectome of the brain is a central goal of modern neuroscience, but the required electron microscopy (EM) imaging is extremely time-consuming. EM image super-resolution (SR) accelerates this acquisition for downstream neuron segmentation, yet existing methods do not adequately couple SR with segmentation for EM: SR lacks segmentation guidance, and segmentation cannot adapt to super-resolved features. In this work, we propose a joint SR and segmentation framework specifically designed for EM images. In the super-resolution part, we introduce segmentation as an auxiliary task via a shared encoder with task-specific decoders to provide a semantic prior, improving reconstruction fidelity. In the segmentation part, we introduce a lightweight partial diffusion model that initializes denoising from a partially noised SR embedding, producing a structural prior that closely approximates real HR features and enhances segmentation accuracy on super-resolved EM images. Extensive experiments across multiple datasets and degradation settings demonstrate the superior performance of our framework.

This paper studies test-time aggregation, an approach that generates multiple reasoning traces and aggregates them into a final answer. Most existing methods rely on evaluation signals collected from candidate traces in isolation or answer frequencies, while ignoring comparative interactions among candidates. We propose *Joint Consistency* (JC), formulated as a constrained Ising-type energy minimization problem, where independent evaluation signals act as external fields and pairwise comparisons act as interactions. JC provides a unified framework for test-time aggregation that subsumes existing voting and weighted aggregation methods as special cases. Our construction of the interaction matrix leverages LLM-as-a-judge comparisons, and admits a theoretical interpretation under answer-level homogeneity assumptions. Moreover, we develop an efficient approximation strategy that makes interaction modeling practical for large-scale test-time aggregation. Experiments on math and code reasoning benchmarks show that JC consistently outperforms existing baselines across tasks, judge models, trace budgets, and trace-generation settings.