NeurIPS 2026 Papers

Animals generate complex behaviors by flexibly recombining a finite set of motor primitives, but existing behavior segmentation methods oversimplify this process by imposing discrete syllables under restrictive generative assumptions. Here, we introduce Motif-based Continuous Dynamics (MCD), a framework that models behavior as a continuous, compositional process driven by reusable motor motifs. MCD leverages reinforcement learning (RL) to (1) discover interpretable motif representations via transition-based representation learning, and (2) model behavior as a time-varying mixture of these motifs through motif-based policies. This formulation avoids restrictive dynamics assumptions while capturing continuity, compositionality, and long-term dependencies in behavior. Across simulated gridworld tasks, maze navigation, and animal behavior datasets, MCD identifies reusable motifs and interprets the trajectories with them. These results provide a generative account of behavior as combinations of fundamental motifs, offering a flexible and interpretable framework for studying natural behavior.

Deep neural networks are highly vulnerable to bit corruption in stored weights, whether from naturally occurring memory faults or deliberate fault-injection attacks. Existing defenses remain fundamentally limited: fault-tolerance methods degrade under cumulative corruption and remain vulnerable to strong white-box attacks, while ECC-based recovery methods are often tied to specific precisions or architectures and expose concentrated vulnerable surfaces. We propose a simple training-based recovery framework that endows model weights with recoverable structure through block-wise zero-sum constraints enforced by differentiable projection. The resulting method incurs no parameter-space overhead and generalizes naturally across architectures and numerical precisions. We further develop adaptive adversarial bit-flip attacks tailored to recovery-based defenses, covering both our method and prior ECC-based approaches. Across diverse architectures, precisions, datasets, and threat models, our method consistently delivers strong robustness and recoverability while preserving clean performance.

Neural network verification is key to certifying robustness in safety-critical systems. Sound verifiers commit to abstract domains that propagate soundly through the network's operations, paying for that commitment with looseness at scale and applicability largely confined to piecewise-linear architectures. Probabilistic verifiers relax soundness, but they too commit to a fixed family of shapes for the network's reachable outputs. When the network's true output geometry differs from the chosen shape, the verifier over-approximates the reach-set and fails to certify networks that are in fact safe. We propose instead to learn the reach-set geometry from the network's behavior. A flow-matching model trained on input-output samples and calibrated by conformal prediction yields a probabilistic reach set whose shape adapts to the true shape of the reach-set rather than to an a priori choice. The resulting set has no closed-form description, so we reframe the specification check as a rare-event estimation problem on the learned distribution, restricted to the calibration's coverage region, and solve it with bounded adaptive multilevel splitting. The pipeline produces a probabilistic safety certificate combining the conformal coverage with a rare-event upper bound. Our pipeline performs comparably to sound verifiers on a subset of VNN-COMP 2025 benchmarks where sound verification is tractable and matches or exceeds fixed-shape probabilistic baselines on most benchmarks in the same suite. On a synthetic family of growing-depth networks, it scales sub-exponentially where sound verifiers time out or abstain.

The Normalized Transformer, or nGPT (Loshchilov et al., 2025) achieves impressive training speedups and does not require weight decay or learning rate warmup. However, despite having hyperparameters that explicitly scale with model size, we observe that nGPT does not exhibit learning rate transfer across model dimension and token horizon. To rectify this, we combine numerical experiments with a principled use of alignment exponents (Everett et al., 2024) to revisit and modify the $\mu$P approach to yperparameter transfer (Yang and Hu, 2021). The result is a novel nGPT parameterization we call $\nu$GPT. Through extensive empirical validation, we find $\nu$GPT exhibits learning rate transfer across width, depth, and token horizon.

We study learning rate (LR) scaling for hybrid architectures combining Transformer and State-Space Model (SSM) blocks, a class adopted by several recent production language models. In particular, we focus on the gap between the theoretical scaling rules derived for SSMs under zero-order-hold (ZOH) discretization at infinite width with proportional state size, and the field-standard practical implementations using simplified-ZOH Mamba at fixed state size. Surprisingly, in this practical regime hybrid architectures achieve a near-zero LR transfer gap across widths 256-2048 and depths 4-32 at sub-billion-parameter scale using only the original $\mu$P prescription, even though SSM operations fall outside its Tensor Programs representability conditions and every parameterization we test fails the standard coordinate-check diagnostic of $\mu$P correctness. We attribute this to a two-condition decomposition of LR transfer in hybrid architectures: a global update-to-weight invariance, enforced by $\mu$P's initialization and LR scaling; and a local per-component balance, provided by AdamW's $g/\sqrt{v}$ per-parameter normalization. Our observations show that the optimal LR is invariant to width up to $8\times$, as well as to depth, sequence length, batch size, and Transformer-to-SSM ratio, and transfers to Nemotron-H, a production hybrid outside our custom architecture set. We hope these findings fill the gap between theoretical scaling rules and practical hybrid implementations, and stimulate further research toward bridging it.

Low-Rank Adaptation (LoRA) is the prevailing approach for efficient large language model (LLM) fine-tuning. Building on this paradigm, recent studies have proposed alternative initialization strategies and architectural modifications, reporting substantial improvements over vanilla LoRA. However, these gains are often demonstrated under fixed or narrowly tuned hyperparameter settings, despite the known sensitivity of neural networks to training configurations. In this work, we systematically re-evaluate nine representative LoRA variants alongside vanilla LoRA through extensive hyperparameter searches. Across tasks spanning mathematical reasoning, commonsense reasoning, code generation, and instruction following at diverse model scales, we find that different LoRA methods favor distinct learning rate ranges. Crucially, once learning rates are properly tuned, all methods achieve similar peak performance (within 1-2\%), with only subtle rank-dependent behaviors. These results suggest that vanilla LoRA remains a competitive baseline and that improvements reported under a single training configuration may not reflect consistent methodological advantages. Finally, a second-order analysis attributes the differing optimal learning rate ranges to variations in the largest Hessian eigenvalue, aligning with classical learning theories.

Stochastic Gradient Descent (SGD) is a workhorse algorithm for continuous optimization. It is a folklore belief among deep learning practitioners that learning rate decay in SGD substantially improves performance. However, it has remained unclear whether learning rate decay offers an asymptotic advantage for \emph{global} nonconvex optimization. We answer this question in the affirmative by demonstrating a natural class of nonconvex functions for which we prove that SGD with learning rate decay requires \emph{exponentially fewer} gradient queries for global optimization than SGD with any fixed learning rate. To the best of our knowledge, this is the first exponential advantage that has been demonstrated for learning rate decay. The class of functions we consider includes many popular benchmark nonconvex functions. Our results provide a robust explanatory theory for the empirically observed benefits of learning rate decay. Our technical results are built upon a new discretization analysis of SGD with decaying step size, that allows us to establish a clean correspondence with continuous-time annealing of Langevin diffusions. We then show that this annealed diffusion performs the non-logconcave sampling task associated to the optimization problem in polynomial time, via a novel application of weak Poincar\'{e} inequalities.

Graph Neural Networks (GNNs) operate under the implicit assumption that the underlying data manifold is Riemannian, where the metric tensor is strictly positive-definite. While effective for homophilic graphs, this inductive bias creates a fundamental geometric mismatch for heterophilic graphs, where edges often signify dissimilarity or structural repulsion. In this work, we propose a paradigm shift from learning on fixed manifolds to learning the manifold itself. We postulate that heterophilic graphs are naturally embedded in pseudo Riemannian manifolds endowed with an indefinite metric, allowing for negative squared distances to model repulsive interactions. To formalize this, we introduce the Graph Signature Index, a spectral invariant that diagnoses the geometric nature of a graph. This index enables our proposed Pseudo-Riemannian Attention Network (PRAT) to dynamically learn an indefinite metric tensor, effectively capturing both attractive and repulsive interactions. By unifying these opposing forces within a single framework, PRAT demonstrates how a subtle shift in the metric signature yields substantial performance gains on heterophilic benchmarks. Our analysis reveals that PRAT’s learned geometry emergently aligns with the spectral signature of the graph, validating our geometric hypothesis.

In many safety-critical applications, control of uncertain dynamical systems relies on observers that estimate states and external disturbances. Neural network observers can improve estimation accuracy, but certifying their Lyapunov stability via Linear Matrix Inequality (LMI) constraints leads to large-scale semidefinite programs (SDPs) that are difficult to solve for large networks. To overcome this scalability bottleneck, we propose a novel two-stage training framework for provably stable neural network observers. Our approach decouples the optimization into a point-guided Lyapunov pre-training phase, which rapidly achieves high estimation accuracy and local stability over sampled states, followed by an LMI fine-tuning phase that efficiently satisfies a strict global Lyapunov stability certificate. We provide formal theoretical guarantees for local stability radii and probabilistic coverage over a prescribed compact error-state domain under specified regularity and sampling assumptions. Experiments on nonlinear control benchmarks and X-29 aircraft ablations show that our LMI-certified neural network observers train significantly faster than direct LMI-based methods and generalize robustly across diverse systems, achieving improved tracking accuracy over a range of observer baselines. The code is available at https://anonymous.4open.science/r/LearningNeuralNetworkObserver-2E1E.

In many modern applications of reinforcement learning (RL), the natural reward for a task of interest is inherently sparse: a reward of 0 is given everywhere except when the task is completed, when a reward of 1 is given. Training a policy to maximize such a sparse reward requires solving a challenging credit assignment problem, leading to slow or ineffective RL improvement. We propose a simple approach to transform a sparse outcome reward into a dense process reward. Our approach relies on training a discriminator to distinguish between previous successful and unsuccessful trajectories, and using this discriminator to incentivize the RL-learned policy to match the state-action visitations of successful trajectories, while avoiding those of unsuccessful trajectories. By incentivizing the policy to match the visitations over all states, not just those that correspond to task success, this reward provides dense feedback on whether progress is being made towards task completion, and, we show, provably achieves this without changing the optimal policy. Focusing on finetuning of robotic control policies, we demonstrate that our approach leads to significantly faster RL finetuning performance on both simulated and real-world manipulation tasks, as compared to simply maximizing the sparse outcome reward.

Robust optimization is a widely used framework for decision-making under uncertainty, particularly in high-stakes applications where reliability is critical. A key challenge in this paradigm lies in constructing uncertainty sets that balance robustness and performance: overly conservative sets lead to pessimistic decisions, while insufficient coverage risks failure in practice. Recent approaches based on conformal prediction provide finite-sample, distribution-free guarantees for uncertainty sets, but remain largely task-agnostic and disconnected from downstream decision objectives. In this paper, we propose a decision-aware conformal prediction framework that directly learns the geometry of uncertainty sets to improve robust decision-making. Our approach introduces a polyhedral nonconformity score that induces feature-dependent uncertainty sets, and a three-step procedure that integrates conformal calibration, robust-decision-aware learning, and re-calibration to correct for post-selection bias. We establish finite-sample coverage guarantees for the final, data-dependent uncertainty set, while achieving improved decision performance by reducing unnecessary conservativeness. This work bridges the gap between statistical validity and decision optimality, providing a principled framework for data-driven robust optimization.

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.