NeurIPS 2026 Papers

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

Whenever LLM training involves multiple outputs from the same input---agentic multi-turn trajectories, self-consistency reasoning, beam search distillation, or multi-rollout SFT/RL---the tokens form a tree-structured trajectory with shared prefixes. Existing pipelines linearize such data and treat each branch independently, causing substantial redundant computation that scales with the number of samples. We derive that averaging the loss over all branches is algebraically identical to a per-token weighted loss, reducing the problem to computing each token's log-probability exactly once. We propose DFS serialization of the tree, which visits every token exactly once, and adapt full-attention and SSM layers to ensure the resulting log-probabilities match independent per-branch computation exactly. For memory-constrained settings where the full tree exceeds GPU capacity, we propose Redundancy-Free Tree Partitioning, which achieves zero redundant computation with peak memory bounded by a single root-to-leaf path. Together, these contributions form Tree Training, achieving up to 6.2× end-to-end training speedup on dense and MoE models for both supervised fine-tuning and reinforcement learning.

Algorithms with predictions, or learning-augmented algorithms, has proved to be an extremely useful paradigm for combining machine learning with traditional algorithms. One of the textbook settings for this is searching a sorted array. Without a prediction classical binary search takes $O(\log n)$ queries, while with a prediction we can use "doubling binary search" to find the target key using $O(\log \eta)$ queries, where $\eta$ is the error of the prediction measured as the absolute value of the difference between the true location and the predicted location. Since an array is just a path graph, in this paper we ask whether similar bounds can be achieved for search on even slightly more general graphs: trees. We show first that the high-level answer is "no": there is no search algorithm that uses $O(\log \eta)$ queries, where now $\eta$ is the graph distance between the predicted location and the true location. However, as our main result, we show that such bounds can be achieved on trees which are "path-like" in that they have low pathwidth. In particular, we prove that there is a search algorithm which uses at most $O(k \log \eta)$ queries, where $k$ is the pathwidth of the tree. We also prove a lower bound showing that our algorithm has existentially optimal query complexity. Finally, we show experimentally, on real-life inputs, that our algorithm has query complexity which is notably better than the simple non-prediction based algorithm.

Large language models rely on positional representations for long context modeling, yet are typically trained on short sequences and evaluated at much longer lengths. Bridging this gap requires positional encodings remaining distinguishable at large distances while generalizing beyond the training length. In this paper, we identify a structural limitation of existing rotary-based positional encodings: all positional dimensions are controlled by a single global index, which couples training-time positional variation with long-range extrapolation. To address this, we propose Tree Rotary Positional Encoding (TPE), a structured multi-scale positional encoding designed for training from scratch that represents absolute positions through com positions of different components, avoiding single-index extrapolation. We further devise an Offset-based Positional Training (OPT) strategy to improve learning interactions across TPE components, enabling better generalization to longer con texts. We theoretically analyze that its multi-scale decomposition enables effective combinatorial position coverage. Experimental results show strong gains under extreme length extrapolation: with a 1.2B model trained from scratch on length 512, our TPE achieves a 64× extrapolation to 32K tokens, reducing perplexity from 4532.21 (RoPE) to 149.96, with comparable inference-time efficiency.

Large Language Model (LLM) agents increasingly excel at complex tasks, yet a persistent bottleneck emerges at the boundary of rule-based systems. Here, the core challenge is often not comprehending intent, but routing user intent through a prescribed, executable interface. This paper posits that this routing problem is best addressed not by scaling model reasoning alone, but by treating the interface itself—specifically, Domain-Specific Languages (DSLs) and their attached verifiers—as first-class boundary infrastructure. We argue that an effective DSL functions as a cognitive shortcut: by providing a high-abstraction, strongly constrained representation, it can compress the intent-to-action path, prune invalid actions before they are attempted, and mechanistically verify correctness via its harness. Our central contribution is to frame the deliberate selection, design, and evaluation of such language--harness pairings as an under-studied research variable that directly shapes an agent's effective action space, failure modes, and learning signals. We develop this position through a dual lens of lifecycle and domain analyses, and ground it with targeted case studies.

Large language models (LLMs) increasingly serve as reasoners and are being considered as automated evaluators, yet they remain susceptible to cognitive biases---often altering their reasoning when faced with spurious prompt-level cues such as consensus claims or authority appeals. Existing mitigations via prompting or supervised fine-tuning fail to generalize, as they modify surface behavior without changing the optimization objective that makes bias cues attractive. We propose \textbf{Epistemic Independence Training (EIT)}, a reinforcement learning framework built around a simple principle: models should learn that bias cues are \emph{unreliable} rather than learning to either follow or reject them. EIT trains on balanced conflict examples where each injected cue is equally likely to support the correct or incorrect answer, and uses a reward that penalizes bias-following errors without rewarding agreement with a cue that happens to be correct---making the cue non-predictive of reward. On controlled MMLU-Pro reasoning tasks with bias injection, EIT improves accuracy and robustness on both Qwen3-1.7B and Qwen3-4B when bias points to wrong answers, while preserving performance when bias aligns with truth. Trained only on bandwagon bias, EIT generalizes along two out-of-domain axes: held-out MMLU-Pro subjects and unseen bias types (authority, distraction, verbosity). EIT-trained Qwen3-4B further outperforms untrained Qwen3-8B and Qwen3-14B on bias resistance, indicating that targeted training is more effective than model scaling alone. Code and data are available at https://anonymous.4open.science/r/bias-mitigation-with-rl-BC47.

Many systems represent graphs as sequences by visiting nodes one at a time, for example, a molecule written as a SMILES string or a program written as a token stream from an AST traversal. The same graph can be traversed in many ways, producing different sequences for the same underlying object. Sequence Transformers treat these as distinct inputs and produce distinct embeddings, even though the graph remains unchanged. Existing graph-native Transformers address this by operating directly on the graph, but many practical pipelines receive serialized inputs and cannot be redesigned. We provide a theoretical characterization in three parts. Tree-positional encoding achieves structural invariance for acyclic graphs under canonical-root serialization but is insufficient for cyclic graphs due to spanning-tree non-uniqueness. Augmenting with cycle-membership encoding derived from the canonical cycle basis introduces a traversal-invariant anchor on cycle vertices. The resulting Tree+Cycle Positional Encoding is a drop-in replacement for sinusoidal PE within graph-aware serialization pipelines (those preserving token-to-atom alignment for structural tokens), requiring no architectural changes to the sequence model. It applies to any domain where labeled graphs are given as sequences with this alignment property. We validate the approach on two testbeds. On synthetic cycle counting and ring membership tasks, it outperforms the GIN baseline by $2\times$ on cycle counting and matches it on ring membership, demonstrating that pre-computed cycle structure encoded as PE matches message-passing GNNs at end-task performance. On molecular property prediction, it achieves $0.977 \pm 0.029$ mean cosine similarity across re-rootings ($n = 24{,}960$), outperforms a parameter-matched graph-native Transformer (Graphormer) on BACE and BBBP, and enables fragment-level interpretability with $+5$ points on alert alignment over a methodologically matched sinusoidal baseline.

Effective discrete representation learning with VQ-VAE serves as a fundamental component of modern autoregressive generative models. However, it is widely known to suffer from codebook collapse and inefficient code utilization. To address this issue, we propose a simple yet effective alternative to standard vector quantization by formulating neural quantization as a joint modeling and clustering process. Specifically, we learn the latent space using a lightweight flow-matching model alongside training, corresponding to a probability flow ODE that transports an often fixed well-behaved source distribution while preserving probability mass. Based on this, we define the codebook $\mathcal{C}_1$ by transporting even-quantile partitioned prior codes $\mathcal{C}_0$ from the source space through the learned flow, and obtaining the Voronoi centroids under the induced distribution, leading a structured discretization. We evaluated on image tasks, and two tasks require high-precision tokenization, demonstrating scalability and improved expressivity of the codebook with full utilization.

High-contrast multiscale PDE problems are common in real-world applications, yet current neural PDE solvers struggle to achieve sufficient accuracy on such tasks. The Generalized Multiscale Finite Element Method (GMsFEM) addresses this by compressing fine-scale heterogeneity into localized basis functions which are used to obtain more accurate solutions than neural PDE solvers. However, constructing these basis functions requires solving many local eigenvalue problems—the major computational bottleneck. We address this issue by proposing **Transolver-GMsFEM**, a new hybrid framework that utilizes the efficiency of neural PDE solver for predicting multiscale basis functions while preserving the solution quality of GMsFEM. Experiments on 2D/3D steady-state and time-dependent high-contrast multiscale PDEs on irregular grids show proposed method achieves over $100\times$ speedup for basis construction. Crucially, our experiments demonstrate that Transolver-GMsFEM significantly outperforms state-of-the-art neural PDE solvers in both accuracy and robustness, especially in out-of-distribution tasks where pure neural PDE solvers dramatically fail.

Public transit route planning traditionally depends on structured map infrastructure and complex routing engines, and no existing dataset supports training models to bypass this dependency. We present TransitLM, a large-scale dataset of over 13 million transit route planning records from four Chinese cities covering 120,845 stations and 13,666 lines, released as a continual pre-training corpus and benchmark data for three evaluation tasks with complementary metrics. Experiments show that an LLM trained on TransitLM produces structurally valid routes at high accuracy and implicitly grounds arbitrary GPS coordinates to appropriate stations without any explicit mapping. These results demonstrate that transit route planning can be learned entirely from data, enabling end-to-end, map-free route generation directly from origin-destination information. The dataset and benchmark are available at https://huggingface.co/datasets/GD-ML/TransitLM, with evaluation code at https://github.com/HotTricker/TransitLM.

Recent image editing systems have achieved impressive semantic understanding, visual fidelity, and instruction-following ability, while video editing remains substantially more difficult and costly. In this paper, we present a simple alternative to end-to-end video editing: instead of training a monolithic video editor, we transform a strong image editor into a video editor through anchor-based generation. Our key insight is that video editing can be decomposed into two subproblems: editing a sparse set of keyframes and propagating those edits across time. Based on this observation, we propose Anchor-based Video Editing (AVE), a two-stage framework in which a powerful image editor first performs composed editing on selected keyframes, and a motion-guided image-to-video diffusion model then generates the final video by treating the edited keyframes as fixed anchors. This design directly inherits the strengths of modern image editors while avoiding expensive end-to-end video editing training. Experiments on IVEBench and VIE-Bench show that AVE achieves strong performance in instruction following, temporal consistency, and content fidelity. Further ablations reveal that final video editing quality is strongly correlated with the quality of the image editor, suggesting that future progress in video editing may come from stronger image editing foundations and lightweight transfer to video.

Large Language Models (LLMs) have achieved remarkable progress on complex reasoning tasks, yet the mechanisms by which they acquire reasoning capabilities and the efficiency of their reasoning remain poorly understood. In this paper, we investigate the path-finding problem over directed graphs—a symbolic abstraction of multi-step reasoning. We first characterize the architectural requirements for this task, proving that one-layer transformers require $\Omega(N^2)$ size on dense graphs to implement the key DFS child-selection primitive. While a two-layer transformer implements full DFS with $O(N\log(N))$ size. Thus, a second layer is necessary and sufficient for near-linear DFS reasoning. We then extend this expressivity analysis to deeper models, proving the existence of an $L$-layer transformer that performs DFS with a look-ahead horizon of $2^{L-3}$, establishing that increasing depth yields an exponential gain in search efficiency. Finally, going beyond the expressivity, given appropriately curated training data, we show that two-layer transformers can provably learn to execute DFS via gradient flow and generalize to unseen graphs. Together, our results provide a mechanistic account of how transformers implement graph search, how depth improves reasoning efficiency, and how such reasoning algorithms can emerge through training.

Training selects for behavior, not circuitry: many weight configurations can implement the same function. Studying any single trained neural network thus risks describing accidents of one training run rather than the computation itself. This work shifts focus from what transformers happen to do to what they must do by extracting algorithmic cores, compact subspaces that are necessary and sufficient for a task and that recur across independently trained models. Here, Algorithmic Core Extraction (ACE) is introduced to isolate these subspaces, causally validate them, and recover the algorithms they implement across settings ranging from synthetic tasks to large-scale pretrained models. Markov-chain transformers embed three-dimensional cores in nearly orthogonal subspaces yet recover identical transition spectra. Modular-addition transformers form compact cyclic cores at grokking that later inflate under continued regularization, redundantly distributing the same computation across many functionally equivalent modes. This functional redundancy is found to accelerate the transition from memorization to generalization, yielding an inverse scaling law for grokking time. In six language models spanning two orders of magnitude in scale (GPT-2 Small/Medium/Large, LLaMA-3.1, Gemma-2, and Qwen2.5), subject-verb agreement is governed by a single, steerable axis that aligns across architectures. Flipping this axis inverts grammatical number throughout open-ended generation. Together these results suggest that beneath the apparent complexity of trained transformers lies a simpler, shared computational structure, and that targeting invariants rather than parameterizations may offer a more tractable path to mechanistic understanding and control.

In-context learning plays a central role in transformer-based large language models, yet its theoretical understanding remains limited. In this work, we study multiclass in-context classification under more realistic settings, including anisotropic class centers and label imbalance, by introducing a spectral data generation framework that constructs class-center matrices with a prescribed singular spectrum. We first show that the isotropy of class-center vectors, as quantified by the stable rank, improves ICL generalization performance. From a meta-learning perspective, our theorem shows that linear transformers can learn multiclass in-context classification with near-optimal per-label sample complexity, extending prior guarantees beyond the binary setting. In the test label imbalance regime, our analysis reveals that queries from majority classes are easier to classify, while those from minority classes are more error-prone; moreover, robustness to this bias improves with the stable rank. Finally, we empirically demonstrate that our theory is consistent with observations on transformers and pretrained large language models.

In image classification scenarios where both prediction and explanation efficiency are required, self-explaining models that perform both tasks in a single inference are effective. However, for users who already have prediction-only models, training a new self-explaining model from scratch imposes significant costs in terms of both labeling and computation. This study proposes a method to transfer the visual explanation capability of self-explaining Vision Transformer (ViT) models learned in a source domain to prediction-only VLM-based ViT models in a target domain via task arithmetic, without target-side explanation supervision. The proposed method endows explanation capability by adding an \emph{explainability vector} induced by explanation supervision in the source domain based on task arithmetic framework. Experiments on ten diverse target datasets show that a single explainability vector learned on ImageNet-1k augmented with patch-level explanation supervision transfers consistently, improving explanation quality while largely preserving classification accuracy. Beyond the transfer itself, we further investigate whether transfer success can be anticipated \emph{a priori} from model-internal features alone (without target-side explanation supervision), and find that a small set of such features carries non-trivial signal about the transfer outcome.

Time-parallel neural operators (FNO and successors) have transformed PDE surrogates on regular grids. Mesh-based operators (MeshGraphNets, Transolver) extend operator learning to irregular geometries but remain autoregressive in time and trained one geometry at a time, while classical parallel-in-time methods converge poorly on nonlinear advection-dominated PDEs. We propose space-time factorization: partition the model into local temporal operators (learned, graph-independent) and a spatial mixer that propagates information across nodes without graph-dependent learned parameters. This structural condition restores time-parallelism on irregular graphs and yields cross-graph zero-shot transfer once three information leaks—graph-dependent weights, graph-specific features, and graph-correlated training distributions—are controlled. The same protocol applies to two distinct spatial-mixer families tested here—fixed-spectral and learned-attention—suggesting the principle is not tied to a single mixer choice. We validate the principle in three domains. On simulating floodwave propagation in river networks, a 27K-parameter model trained on the Yangtze transfers zero-shot to the Mississippi at $R^2=0.998$, training $28\times$ faster than a recurrent message-passing baseline and scaling to $5\times$ larger graphs. On cross-geometry Navier–Stokes, both Chebyshev (SepFNO) and Physics-Attention (SepTransolver) variants reach within ~10% (relative $L_2$) of a locally-trained oracle using only 15 training geometries—to our knowledge the first demonstration of knowledge accumulation in this regime across fundamentally different topologies. On cross-city traffic forecasting, zero-shot transfer reaches within 12.9% of an oracle trained on the target city. A consistent pattern emerges across these settings: cross-geometry error decreases with both the number of training geometries and their relatedness to the test, with the rate of accumulation depending on the spatial mixer. The result opens the path for pooled-data, foundation-style models for graph dynamics on irregular topologies.

The cost of Kohn-Sham density functional theory (KS-DFT) calculations scales with the number of solver iterations, which depends on the quality of the initial guess. Machine learning methods that predict initial guesses from molecular geometry can reduce this cost, but matrix-prediction models fail when extrapolating to larger molecules, degrading rather than accelerating convergence [Liu et al., 2025]. We show that this failure is a supervision problem, not an extrapolation problem: models trained on ground-state targets fit those targets well out of distribution, yet produce initial guesses that slow convergence. Solver-Aligned Initialization Learning (SAIL) resolves this for both Hamiltonian and density matrix models by differentiating through the self-consistent field (SCF) solver end-to-end. We introduce the Effective Relative Iteration Count (ERIC), a correction to the commonly used RIC that accounts for hidden Fock-build overhead. On QM40, which contains molecules up to 4$\times$ larger than the training distribution, SAIL reduces ERIC by 37\% (PBE), 33\% (SCAN), and 28\% (B3LYP), more than doubling the previous state-of-the-art reduction on B3LYP. On QMugs molecules 10$\times$ larger than the training set, SAIL delivers a 1.35$\times$ wall-time speedup at the hybrid level of theory, extending ML SCF acceleration to large drug-like molecules.

We study transfer learning for a linear regression task using several least-squares pretrained models that can be overparameterized. We formulate the target learning task as optimization that minimizes squared errors on the target dataset with penalty on the distance of the learned model from the pretrained models. We analytically formulate the test error of the learned target model and provide the corresponding empirical evaluations. Our results elucidate when using more pretrained models can improve transfer learning. Specifically, if the pretrained models are overparameterized, using sufficiently many of them is important for beneficial transfer learning. However, the learning may be compromised by overparameterization bias of pretrained models, i.e., the minimum L2-norm solution's restriction to a small subspace spanned by the training examples in the high-dimensional parameter space. We propose a simple debiasing via multiplicative correction factor that can reduce the overparameterization bias and leverage more pretrained models to learn a target predictor.

Understanding how information propagates within vision-language models (VLMs) and large language models (LLMs) is central to interpretability and efficient deployment. Attention weights indicate where probability mass is allocated, but do not directly measure directional influence, redundancy, or functional importance. We introduce transfer entropy (TE) as an information-theoretic measure of directed information flow in Transformers, and develop tractable TE estimators for high-dimensional hidden states at the granularity of layers, tokens, and attention heads. We first analyze multimodal models, beginning with LLaVA-1.5-7B and then CLIP ViT-B/16. In LLaVA-1.5-7B, vision$\rightarrow$text TE increases with depth, suggesting that multimodal fusion is localized in the upper decoder layers. In CLIP, TE reveals late-layer redundancy in the vision tower and broader elevated TE in the text tower. We then analyze unimodal models, including RoBERTa, T5, Llama-3.2-3B, and Qwen-2.5-7B. Across these LLMs, TE reveals consistent depth-dependent structure: RoBERTa exhibits a mid-layer peak, decoder-only LLMs concentrate task-dependent computation in broad mid-depth bands, and T5 shows distinct encoder and decoder dynamics. Beyond layerwise analysis, TE also exposes token- and head-level communication roles, distinguishing sinks, broadcasters, and benign components. These signals support TE-guided layer, token, and attention-head pruning across VLMs and LLMs, consistently outperforming attention- or saliency-based heuristics under matched budgets. Overall, TE provides a unified directed-information lens for diagnosing, interpreting, and compressing Transformer-based models.