NeurIPS 2026 Papers

Deliberating takes time. In real-time settings, that time is not free. Standard reinforcement learning (RL) sidesteps this as the environment waits indefinitely for the agent's decision. Instead, we study real-time RL environments where the environment progresses while waiting for the agent's action. Building on prior real-time formalizations, we introduce variable-delay real-time RL, where the agent chooses how long to deliberate at each decision point since the environment progresses. For the planning agents we use, the right delay is state-dependent, and naively planning how long to plan can paralyze the agent. We instead approach this setting by training a lightweight gating policy on top of a planner to select state-dependent planning budgets. Across real-time Pac-Man, Tetris, Snake, Speed Hex, and Speed Go, our gating policy outperforms fixed-budget and heuristic baselines, and transfers to a real-time setup where the environment and agent run on two different GPUs.

Whether navigating a building, operating a robot, or playing a game, an agent that acts effectively in an environment must first learn an internal model of how that environment works. Partially observable Markov decision processes (POMDPs) provide a flexible modeling class for such internal world models, but learning them from observation-action trajectories alone is challenging and typically requires extensive environment interaction. We ask whether language-model priors can reduce costly interaction by leveraging prior knowledge, and introduce Pinductor (POMDP-inductor): an LLM proposes candidate POMDP models from a few observation-action trajectories and iteratively refines them to optimize a belief-based likelihood score. Despite using strictly less information, Pinductor matches the performance and sample efficiency of LLM-based POMDP learning methods that assume privileged access to the hidden state, while significantly surpassing the sample efficiency of tabular POMDP baselines. Further results show that performance scales with LLM capability and degrades gracefully as semantic information about the environment is withheld. Together, these results position language-model priors as a practical tool for sample-efficient world-model learning under partial observability, and a step toward generalist agents in real-world environments.

Infinite-dimensional orthonormal basis expansions play a central role in representing and computing with function spaces due to their favorable linear algebraic properties. However, common bases such as Fourier or wavelets are fixed and do not adapt to the structure of a given problem or dataset. In this paper, we aim to represent these bases with neural networks and optimize them. Our key idea is that any target infinite-dimensional orthonormal basis can be viewed either as a point on the Lie manifold of the orthogonal group, or equivalently, as the endpoint of a continuous path on that manifold that connects a reference basis, e.g. Fourier, to that target. Paths on the Lie manifold satisfy ordinary differential equations (ODEs) governed by skew-adjoint integral operators. Using neural networks to define finite-rank generators of such ODEs allows us to parameterize and optimize orthonormal bases in function space. While relying on finite-rank generators to model infinite operators might seem restrictive, we prove a universality result: even with a rank-2 generator, the integrated solutions of the ODE are dense in the orthogonal group under the appropriate operator topology. In other words, for any target orthonormal basis, there exists a path originating from a reference basis and driven by finite-rank generators that gets arbitrarily close to that target basis. We demonstrate the flexibility of our framework by transforming the Fourier basis into the principal components of a functional dataset, eigenfunctions of linear operators, or dynamic modes of energy-preserving physical simulations.

Recent work has identified incremental learning in shallow networks trained on single-index and multi-index models. However, existing analyses often rely on simplifying settings, such as small initialization, correlation loss, or layer-wise training. These choices reduce neuron interactions and leave some feature learning dynamics under standard initialization unexplored. We study gradient flow dynamics for polynomial-width two-layer networks learning orthogonal multi-index targets under standard initialization using polynomially many samples. We first prove that incremental learning still occurs: the loss decreases sequentially according to the Hermite expansion of the target, with lower-order components learned before higher-order components recover the individual target directions. In this standard initialization regime, training also shows a competitive reallocation of parameter mass: after the total mass fits the target mean and stabilizes, mass shifts into the target subspace and then concentrates on aligned neurons. Technically, we introduce a symmetry-based finite-width approximation via symmetrized networks, rather than comparing directly with an infinite-width limit. This yields better control of approximation errors and may be of independent interest.

Learning flexible motor primitives is a hallmark of skilled motor control. Recent neuroscience theory proposes that motor primitives may be implemented as low-rank perturbations of a shared recurrent network, but leaves open how such a system is learned. We translate this principle into a novel architecture for learning motor skills end-to-end: a shared recurrent core modulated by a bank of residual adapters, each selected by a discrete latent code. Trained on closed-loop biomechanical control, the adapters develop emergent low-rank perturbations of the recurrent dynamics despite no architectural rank constraint, placing task representations in disparate subspaces of the shared core network. A simple high-level policy over the learned options, optimized while the whole network is frozen, sequences the low-rank adapters to produce novel out-of-distribution movements. We demonstrate the ability to generalize to novel motor sequences within the closed-loop control setting, improving on the generalization error of a task-input-conditioned multitask baseline by up to order of magnitude.

Optimal Transport (OT) provides a principled framework for comparing probability distributions. Classical solvers remain the gold standard for this task, providing exact solutions. At the same time, OT is an ideal testbed for studying whether learning-based amortized models can internalize the structure of OT type problems. Here, ground-truth solutions are available, constraints are explicit, and exact discrete solutions have well-characterized sparsity. We describe ToR-OT, a transformer-based model that learns to predict discrete OT plans through autoregressive token regression. We reframe OT into a structured sequence generation task. ToR-OT decodes sparse transport plans as sequences rather than solving an optimization task at test time. We focus on three properties: (a) Zero-shot generalization: pretraining on a large synthetic corpus of OT instances allows a single model to generalize to unseen distribution pairs and heterogeneous cost metrics without retraining. (b) Sparse primal decoding: by exploiting the inherent sparsity in exact discrete OT solutions, the decoder predicts only nonzero entries (avoids dense coupling outputs). (c) Amortized inference: on low-dimensional discrete problems, ToR-OT gives accurate approximate plans and compares favorably with per-instance neural OT baselines. Overall, ToR-OT is not yet a replacement for classical solvers, but best viewed as a step toward understanding how autoregressive transformers can learn structure of constrained optimization problems like OT.

Does a neuron’s shape predict whom it connects to? Peters’ rule, a canonical geo- metric principle in connectomics, predicts connectivity from spatial arbor overlap, but true wiring also depends on molecular and morphological specificity. We ask how far neuron skeletons, with identity and graph features withheld, predict directed connections in the FlyWire Drosophila connectome. Because most non- connected pairs are far apart, headline metrics are dominated by easy negatives that geometry can reject. We therefore evaluate controlled contrasts: plausible non-connections whose arbors are nearby, overlapping, or cell-type matched, but that still do not synapse. Across a controlled model ladder, from arbor overlap and hand-crafted morphology baselines to Connectoformer (a bidirectional Trans- former with pooled and pairwise cross-attention scoring), performance improves when the model gains the corresponding signal: global shape, local branch com- patibility, and tree context. Morphology and arbor co-occupancy are complemen- tary: morphology discriminates similar-looking unconnected pairs but can score pairs whose arbors never touch; arbor co-occupancy does the inverse. By suppress- ing essentially-zero-overlap pairs without boosting high-overlap pairs, a fixed one- sided arbor prior performs well in both regimes. The encoder also transfers across Drosophila individuals: zero-shot evaluation on the male CNS connectome beats arbor overlap on hard-negative discrimination. Stratified evaluation reveals a clean division of labor: geometry rules out pairs that cannot physically connect, and morphology picks the actual partners from what remains.

Offline reinforcement learning aims to learn a policy solely from fixed datasets, which often contain multimodal action distributions. Flow policies can naturally represent such multimodal behaviors, but learning an efficient one-step flow policy remains challenging: standard value guidance often leads to mode collapse or exploits overestimation bias in out-of-distribution regions. To address this, we introduce One-step Flow policy via Optimal Transport (OptiFlow), a framework for one-step flow policy learning as a structured sample-allocation problem. OptiFlow jointly trains a value-aware reference flow policy and an efficient one-step policy, coupling their action samples through state-wise entropic optimal transport. For each state, critic-estimated values define the priority of distillation target actions, while the action-distance cost ensures geometrically compatible pairings. By avoiding direct critic maximization, our transport-guided approach enables in-distribution exploitation by anchoring the one-step policy to high-value, dataset-supported modes without the risk of out-of-distribution divergence. Experimental results demonstrate that OptiFlow effectively captures optimal multimodal behaviors and achieves strong performance across diverse offline RL benchmarks.

Solving inverse problems on video requires priors over both frame appearance and temporal dynamics. Many recent approaches combine pretrained image diffusion priors with motion guidance, sidestepping the costs of training large video models. Temporal consistency is often enforced via photometric losses or warping-based constraints, which implicitly assume that motion explains frame-to-frame or latent-to-latent changes up to small deviations, an assumption that fails under occlusions, lighting changes, and non-rigid deformation. In this paper, we propose to learn the coupling between motion and frame appearance as a generative prior. We learn the distribution of motion-induced appearance residuals with a conditional diffusion model and use it to regularize reconstruction in video inverse problems. Our motion-appearance prior acts as a plug-in regularizer for diffusion-based video reconstruction, preserving the benefits of image diffusion models while remaining compatible with existing appearance regularization techniques. Experiments on dynamic video inverse problems show that, depending on the task, reconstruction with a learned motion-appearance coupling prior matches or outperforms feature-consistency and noise-warping baselines.

Learning parity functions, more general modular addition, is a challenging machine learning task due to its input sensitivity. A recent study substantially scaled modular addition learning in both the number of summands and the modulus. Its key idea is to increase zeros in training sequences, reducing the effective number of summands and thus controlling training difficulty; however, this induces covariate shift between training and test input distributions. This study theoretically and empirically analyzes this side effect and proposes a covariate-shift-free method for modular addition. Specifically, we introduce an auxiliary modulus $Kq$ during training, which reduces wrap-around frequency and problem difficulty while preserving the same input distribution across training and testing. Experiments show strong scalability and sample efficiency: even for large input length $N$, large modulus $q$, and small datasets---where the sparse method fails to learn---our method achieves equal or better match accuracy and relaxed $\tau$-accuracy. For example, at $N=64$ and $q=974269$, our method trained on 100K samples achieves $97.0\%$ $\tau$-accuracy at $\tau=0.05$, while the sparse method achieves only $9.5\%$ with the same data size and $93.9\%$ even when extended to 1M samples.

GraphRAG enhances Large Language Model reasoning by organizing retrieved knowledge into structured evidence graphs, enabling inference over connected evidence rather than isolated text fragments. Yet, existing GraphRAG methods often either miss query-critical links or introduce noisy or conflicting evidence that distracts reasoning. We propose CONSIST, a Conflict-aware Nash-optimized Minimal Sufficient graph learning algorithm that regulates evidence graphs to resolve conflicts during reasoning. Specifically, we cast graph learning as an exploration–editing process, where a candidate graph is expanded via structural and semantic connections and then refined by pruning unreliable edges. We further introduce Gumbel-based edge editing, framing it as a Nash-guided bargaining process, where current and future players negotiate over a customized utility to determine reliable edits that improve both immediate answer support and long-term graph quality. Experiments on multi-hop question answering and conflict-aware benchmarks demonstrate that CONSIST consistently outperforms recent state-of-the-art baselines, achieving average performance gains up to 7.0% EM and 6.3% F1. Moreover, CONSIST also produces substantially more compact evidence graphs, reducing retained nodes by 42.9% and edges by 54.7% on average.

Budget-feasible mechanisms (BFMs) are a fundamental tool for budget-constrained procurement, but designing strong BFMs remains challenging, particularly for general valuation classes. Classical approaches rely on problem-specific analysis and worst-case approximations, often yielding conservative mechanisms in practice. This motivates data-driven methods for discovering mechanisms with stronger empirical performance. However, strict budget constraints expose a key limitation of existing neural automated mechanism design (AMD): they either fail to provide exact economic guarantees or suffer from severe scalability bottlenecks. We propose BFMNet, a scalable menu-based neural framework for learning truthful budget-feasible procurement mechanisms. BFMNet first learns *self-bid-independent* menus from market context, ensuring dominant-strategy incentive compatibility (DSIC) and individual rationality (IR) by construction, and then applies an ex-post, instance-level menu transformation to enforce budget feasibility. We prove that the resulting mechanism satisfies exact DSIC, IR, and budget feasibility. By avoiding intractable global discretization grids and restrictive Lipschitz constraints, BFMNet remains both expressive and scalable. Extensive experiments show that BFMNet improves utility over classical baselines by up to $34.1$%, remains competitive with state-of-the-art neural baselines that provide only approximate economic guarantees, and scales robustly to larger auction instances.

We study how to choose a probability distribution, or lottery, over $m$ alternatives for $n$ agents when each agent's acceptable lotteries form an unknown linear halfspace of the simplex. The learner observes only binary accept/reject membership queries. Prior work asks whether there exists a lottery accepted by every agent; we study the relaxation that minimizes the number of agents who reject the lottery. We first show that this problem is APX-hard even when all halfspaces are explicitly known and scores are binary. We then give exact algorithms under additional structure: fixed $m$, and a margin condition under which an optimal lottery can be replaced by one supported on $\mathcal{O}(\log n/\gamma^2)$ alternatives. With limited queries, uniform sampling of agents gives additive population guarantees after recovering the sampled agents' normalized halfspaces. For arbitrary $m$, we give a polynomial-time LP based on a one-sided margin penalty and prove a population guarantee depending on the number of true violations and the number of agents near their thresholds. Finally, we prove a tight deterministic lower bound: without a stochastic link between queried and unqueried agents, querying only $K$ distinct agents cannot guarantee additive gap below $n-K$, and this gap is achievable.

While multi-agent reinforcement learning (MARL) has shown strong empirical performance, existing methods struggle with large number of agents, due to the combinatorial growth of joint interactions. Mean-field (MF) approximations address this by replacing pairwise interactions with interactions against population distributions, yielding tractable large-population policies. However, existing MF formulations focus on fully cooperative or purely competitive settings and fail to capture the mixed cooperative–competitive structure of team-based games. We extend mean-field learning to _zero-sum team games_, where agents cooperate within teams and compete at the team level. We show that such games admit $\epsilon$-optimal decentralized policies that depend only on local states and population distributions. Building on this structure, we propose MF-MAPPO, a scalable algorithm with a shared actor and a minimally informed critic per team. MF-MAPPO is trained directly in finite-population simulators rather than using mean-field oracles, thereby enabling deployment to realistic scenarios with thousands of agents. We further extend MF-MAPPO to partially observable settings via a simple gradient-regularized training scheme. Experiments on large-scale benchmarks in our simulation platform $\texttt{MFEnv}$, including population games with analytical solutions and high-dimensional battlefield scenarios, demonstrate that MF-MAPPO outperforms existing MARL baselines and yields rich, heterogeneous behaviors.

Modern recording technologies enable simultaneous measurement of high-dimensional neural activity and rich behavioral variables, creating both an opportunity and a modeling challenge: identifying the shared latent structure that mediates the brain-behavior relationship and the underlying dynamics through which it evolves across behavioral states. Existing approaches focus on different parts of this challenge: shared representation learning methods identify neural-behavioral subspaces but typically do not model dynamical structure, while dynamical models capture temporal variation and discrete state transitions in neural activity, but do not explicitly disentangle shared neural-behavioral structure from modality-specific variability. We introduce Switching Shared Latent Dynamics (SSLD), a unified framework that learns a latent representation shared between neural activity and behavior and endows this shared space with nonlinear switching recurrent dynamics to reflect the hypothesis that behaviorally relevant neural representations exhibit structured dynamical changes across behavioral states. Complementary private latent variables capture modality-specific variability, providing a window into neural activity that does not directly interact with behavior. We evaluate SSLD on a simulated dataset and four diverse experimental datasets spanning species, brain regions, and recording modalities: monkey motor and premotor cortex during reaching, somatosensory cortex during a bump task, widefield calcium imaging across mouse dorsal cortex during (a) self-initiated decision-making and (b) spontaneous movements. Across all four datasets, SSLD accurately reconstructs neural and behavioral signals, recovers discrete states in the dynamics that align with experimentally defined behavioral epochs, and isolates behaviorally-relevant neural information in the shared latent while preserving private neural variability. Ablations confirm that shared representation learning, behavioral supervision, and switching dynamics each contribute to performance. SSLD offers an interpretable approach to modeling shared neural-behavioral dynamics that complements ongoing efforts toward foundation models for neuroscience.

Domain generalization aims to learn models that generalize to unseen target domains using labeled data from multiple source domains. Many existing approaches focus on learning domain-invariant representations, but enforcing invariance alone may fail to preserve sufficient task-relevant information under domain shift. We propose BLENDER, a reconstruction-aware framework for domain generalization that employs a hierarchical latent decomposition to separate domain-invariant and domain-specific factors. This structure prioritizes informativeness in the domain-invariant representation and captures residual domain-specific variation through conditional modeling. We further provide a theoretical analysis establishing a bound on the reconstruction risk for unseen target domains, revealing how domain invariance and latent disentanglement contribute to generalization beyond the observed domains. Experiments on standard domain generalization datasets demonstrate strong performance under substantial domain shifts, and qualitative analyses show that BLENDER promotes a structured allocation of information between domain-invariant and domain-specific latent representations.

Human videos contain rich manipulation priors, but using them for robot learning remains difficult because raw observations entangle scene understanding, human motion, and embodiment-specific action. We introduce MoT-HRA, a hierarchical vision-language-action framework that learns human-intention priors from large-scale human demonstrations. We first curate HA-2.2M, a 2.2M-episode action-language dataset reconstructed from heterogeneous human videos through hand-centric filtering, spatial reconstruction, temporal segmentation, and language alignment. On top of this dataset, MoT-HRA factorizes manipulation into three coupled experts: a vision-language expert predicts an embodiment-agnostic 3D trajectory, an intention expert models MANO-style hand motion as a latent human-motion prior, and a fine expert maps the intention-aware representation to robot action chunks. A shared-attention trunk and read-only key-value transfer allow downstream control to use human priors while limiting interference with upstream representations. Experiments on hand motion generation, simulated manipulation, and real-world robot tasks show that MoT-HRA improves motion plausibility and robust control under distribution shift.

Despite the effectiveness of Cube-and-Conquer (C&C) for solving challenging Boolean Satisfiability (SAT) problems, no prior work has shown that transformer-based models can learn effective cubing heuristics. We introduce a neuro-symbolic post-training framework for this task. We design an MCTS-based data curation pipeline that uses symbolic heuristics to explore splitting decisions over SAT competition formulas, producing preference data grounded in solver statistics and augmented with reasoning traces from a teacher model. Our two-stage post-training, supervised fine-tuning (SFT) followed by direct preference optimization (DPO), enables a 4B-parameter model to achieve a pass@5 score of 53 on 100 SAT competition benchmarks, surpassing frontier LLMs such as Claude-Sonnet-4 (50) and matching the best symbolic heuristic (53). Ablations show that SFT alone improves pass@5 from 46 to 51, with DPO adding 2 additional benchmarks; an entropy/agreement ablation on realized first-cube decisions further shows that SFT, not DPO, accounts for the root-level decision diversity that produces complementary per-run coverage over deterministic symbolic methods. This demonstrates that transformers can be trained to make effective cubing decisions in a domain traditionally dominated by symbolic methods.

Vision Transformers (ViTs) are bottlenecked by the length of their input token sequence. Standard patchification fixes this length with a single patch size applied across the image, regardless of regional content. Prior work on adaptive patch sizing reduces token count by assigning coarser patches to visually simple regions, but relies on noisy pixel-level heuristics such as local entropy to make patchification decisions. We argue that this compute allocation should be grounded in the model’s own semantic representations, not in raw pixel statistics. We introduce Split Policy Learned Image Tokenization which learns to allocate token resolution based on what the model finds meaningful, producing shorter token sequences that are better matched to task-relevant content. SPLIT matches the accuracy of full-resolution ViT classification baselines while using 30% fewer tokens. The same gains hold dense prediction: SPLIT achieves competitive accuracy on COCO and ADE20K, demonstrating that semantically grounded patchification generalises well beyond classification.