NeurIPS 2026 Papers

This paper studies the task of estimating heterogeneous treatment effects in causal panel data models with covariate effects. We propose a novel **Co**variate-Adjusted **DE**ep C**A**usal **L**earning (**CoDEAL**) framework, that cohesively deal with the underlying heterogeneity and nonlinearity of both panel units and covariate effects. CoDEAL integrates nonlinear covariate effect components (parameterized by a feed-forward neural network) with nonlinear factor structures (modeled by a multi-output autoencoder) to form a heterogeneous causal panel model. The nonlinear covariate component flexibly captures complex covariate influences on outcomes, and the nonlinear factor decomposition enables CoDEAL to effectively capture both cross-sectional and temporal dependencies inherent in the data panel. This latent structural information is subsequently integrated into a customized matrix completion algorithm, thereby facilitating more accurate counterfactual imputation. Moreover, the use of a multi-output autoencoder explicitly accounts for heterogeneity across units and enhances the model interpretability of the latent factors. We establish theoretical guarantees on the convergence of the estimated counterfactuals and demonstrate the compelling performance of the proposed method using extensive simulation studies and real-world data applications.

Physics-Informed Neural Networks (PINNs) solve forward PDEs by minimizing residual losses from the governing equations with initial and boundary conditions, but they often struggle with discontinuities such as shocks. In contrast, finite volume methods (FVM) handle discontinuities by enforcing integral conservation, which admits weak solutions. Motivated by this, we propose a \emph{Coupled Integral PINN (CI-PINN)} that augments a standard PINN with an auxiliary network for integral potentials and coupled integral constraints. This improves robustness near shocks while avoiding meshing and the numerical flux integration/reconstruction used in classical schemes. We validate CI-PINN on forward benchmarks including Burgers, Buckley--Leverett, the Euler system, and the Shallow-Water equations.

Guidance methods are essential tools for steering the generation process in diffusion models, improving the generalizability of stochastic robotic controllers and enabling high-quality controlled image synthesis. Despite strong parallels between flow-matching and diffusion, integrating guidance into flow-matching models remains quite challenging. Recent work has proposed general guidance frameworks for flow-matching, however, they are direct adaptations of diffusion guidance methods and can prove unstable. In this work, we conduct a theoretical analysis of guidance in flow-matching and show that a naive translation of diffusion-based guidance to flow-matching breaks the coupling distribution used during training, resulting in weaker control over sampling accuracy. Building on this insight, we propose **Coupled Guidance**, a flow-matching-specific guidance mechanism derived from a geometric perspective that leverages the coupling structure inherent to flow-matching training. To validate the effectiveness of our approach, we conduct extensive experimentation on a custom robotic *collision avoidance* task using Maniskill3, four robotic *dexterous manipulation* tasks from the Adroit benchmark, and three image generation tasks, namely *inpainting*, *super-resolution*, and *deblurring*, using the ImageNet-128 and CelebA-HQ-256 datasets. Our findings show that our framework improves sample quality and guidance stability across all tasks, without imposing any additional overhead.

Crowd counting is a vital computer vision task with wide-ranging applications, from public surveillance to urban planning. A persistent challenge is achieving scale invariance despite noisy labels, as humans often make counting errors with similar objects. Most existing methods struggle to maintain consistency and robustness when testing images differ in scale from the training data. We propose ScaLe-aware Invariance and Correcting Errors (SLICE) for robust crowd counting, a novel scale-aware loss function that improves scale invariance, enabling consistent and robust predictions under varying scales and annotation noise. The loss function combines two components: a scale-invariance term that applies multiscale supervision to enforce consistency across density maps of the same scene, and an error correction term that addresses inevitable annotation noise by modeling label uncertainty with probabilistic methods such as a mixture of Gaussians. SLICE is model-agnostic and integrates easily into existing crowd counting architectures without extra computational overhead during inference. Extensive experiments on five benchmark datasets show that SLICE greatly improves scale invariance and cross-dataset generalization, achieving the state-of-the-art results in challenging scenarios. Using CrowdDiff, Steerer, and CrowdFormer as baselines, SLICE lowers mean absolute error (MAE) by 5–7\% by addressing annotation noise and scale variation. Applying the scale-invariance term across multiple scales further cuts MAE by up to 15\%.

Decoding cognitive functions and brain diseases from functional magnetic resonance imaging (fMRI) is a foundational problem in neuroscience advanced by recent deep learning across task and resting-state (rs) settings. Existing region-level decoders capture spatial structure through region-of-interest (ROI)-aware graph or Transformer architectures, but apply a uniform temporal operator across all regions. Consequently, they do not explicitly model how cortical regions integrate information over heterogeneous timescales. We propose CoRTeX, a region-resolved recurrent architecture for fMRI in which each cortical region carries its own learnable parameters and evolves on its own learned timescale. At its core, per-region differentiable temporal integration windows allow each ROI to learn how far back in time to integrate information through a soft mask over a private buffer. Cross-region information is restricted to a single identity-aware attention channel, additionally biased by a co-activation-driven memory whose learning and forgetting rates are themselves region-specific. Experiments on stimulus classification using task-fMRI (NSD) and on two disease diagnosis tasks using rs-fMRI (ABIDE, COBRE) demonstrate that CoRTeX achieves competitive performance with strong baselines while enabling region-level analyses that are not well-defined for models with exchangeable units: the learned per-region timescales exhibit interpretable spatial organization, and class-specific regional attribution aligns with category-selective cortex. Code is available at: https://anonymous.4open.science/r/CoRTeX-F.

Network-level traffic signal control (TSC) involves coordination under partial observability and strict real-time constraints. In real-world urban networks with dozens or hundreds of intersections, direct coordination incurs an intractably large joint action space, making centralized cooperation difficult to deploy in practice. We propose CorridorLight, which tackles this challenge by casting cooperation as task negotiation. A high-level Corridor Agent represents the network as a direction-level graph and periodically generates a small set of junction-disjoint corridor tasks via a seed-and-grow procedure. These tasks provide lane-level preference guidance, while Intersection Agents maintain fast decentralized control and selectively support the active corridor task. By learning when to cooperate versus act selfishly, CorridorLight enables an adaptive trade-off between local delay minimization and corridor-level relief, leading to improved global performance. To this end, we derive a closed-form, Lipschitz-continuous gating function by maximizing a counterfactual net gain between cooperative and selfish behavior, yielding smooth, interpretable cooperation without brittle heuristics. Experiments on synthetic grids and real-world networks show consistent system-level improvements and outperform cooperative TSC baselines across diverse topologies and demand patterns.

Existing methods for Virtual Try-On (VTON) often struggle to accurately transfer garment appearance, especially in unpaired settings where accurate person-to-garment correspondence is required. These methods do not explicitly supervise person-to-garment alignment, leaving correspondence to be learned implicitly within the generation model. In this paper, we first analyze full self-attention in DiT-based architectures and show that person-to-garment query-key matching aligned with local geometric correspondence is closely related to try-on quality. Building on this insight, we introduce CORrespondence ALignment (CORAL), a DiT-based framework that explicitly aligns query-key matching with robust external correspondences. CORAL integrates two complementary components: a correspondence distillation loss that aligns reliable matches with person-to-garment attention, and an entropy minimization loss that sharpens the attention distribution. For evaluation, we further propose a VLM-based evaluation protocol to better reflect human preference. CORAL consistently improves over the baseline, enhancing both global shape transfer and local detail preservation. Extensive ablations validate our design choices. Code and weights will be publicly available.

We study contextual bandit problems with correlated arms and access to surrogate reward signals produced by a machine learning model, motivated by applications such as large language model (LLM) routing. Unlike classical contextual bandits that rely solely on bandit feedback and assume conditional independence across arms, our setting allows context-dependent inter-arm correlations and auxiliary reward information that may be noisy or misspecified. We propose algorithms that leverage such surrogate rewards through two complementary designs. A coupled reward-mixing approach pools true and surrogate rewards to accelerate learning when surrogate signals are reliable, while a decoupled prediction-mixing approach maintains separate estimators for bandit feedback and surrogate rewards and adaptively combines their predictions. This decoupling yields robustness to surrogate misspecification, recovering regret guarantees comparable to reward-only bandit methods in the worst case, while achieving improved regret when surrogate predictions are sufficiently informative. We provide theoretical regret analyses for both approaches and evaluate them on LLM routing benchmarks under varying accuracy versus cost trade-offs. The results demonstrate improved sample efficiency and consistently better accuracy–cost trade-offs compared to standard contextual bandit baselines and strong static routing methods.

While machine learning (ML) architectures have evolved rapidly to account for complex data, loss functions like cross-entropy remain mostly structure-agnostic in many real-world applications. However, the "class-symmetric" nature of these standard losses fundamentally limits the ability of ML models to exploit structural relationships between classes, particularly when facing structured noise. We propose Conveyance, a new classification approach and associated loss function tailored to structured class spaces. It allows users to encode graph-like relations between classes without having to define complex joint distributions or manually tune utility matrices. Technically, our loss function operates by maximizing two separate margins over distinct class partitions, while preserving formal properties such as monotonicity and partial convexity. We demonstrate the versatility and effectiveness of our method by applying it to hierarchical classification, ordinal regression, and multiple instance learning. Across these tasks, Conveyance either matches or exceeds the performance of specialized baselines, thereby offering a unified solution for structured class spaces.

Control variables are widely used in statistical modelling to account for omitted variable bias. However, they have largely been underexplored in deep learning. This is surprising, given that deep learning models encode image-inferable covariates, such as demographic variables, into their predictions when these covariates are correlated with the outcome---a form of omitted variable bias referred to as 'shortcut learning'. While many existing confound-control or fairness methods try to restrict the correlation of such covariates with model predictions, we show that this fails to correct for omitted variable bias. We therefore propose a control variable approach for deep learning models, based on generalised additive modelling of the effects of model inputs and covariates. As flexible additive models can suffer from concurvity, we introduce an estimation procedure that refits the final layer of a pre-trained network to include covariate effects, using cross-fitting with ridge penalisation. We show how these effects can be orthogonalised with respect to covariates to exclude their mediated effects and that model predictions can be marginalised over the covariate distribution to control for their effect. This yields unbiased, interpretable predictions and offers flexibility to model the desired effects depending on the scientific or fairness objective. We verify our approach using simulated images, where it recovers the true covariate effects, even in small samples. Existing methods either require more data or fail to recover the true effects. We apply our method to real neuroimaging data with experimentally induced confounding, where it recovers prediction performance to near the level of a model trained on unconfounded data.

Autoregressive neural simulators now match classical solvers on short-horizon prediction of physical systems, yet their accuracy degrades rapidly when rolled out over long horizons. In this work we identify *transient amplification* of perturbations around rollout trajectories as a structural mechanism driving rollout error. Using a linearization analysis we show that when the Jacobians along an autoregressive trajectory are non-normal and non-commuting, the model amplifies errors transiently, resulting in model rollout drift even when the overall system is asymptotically stable. Building on the analysis, we propose \emph{commutativity regularization}: a combination of two penalties designed to reduce the normality defect of individual Jacobians and the commutator norm of Jacobians across steps. The penalties are estimated with Jacobian-vector products and have no inference-time cost. We show a propagator bound that quantifies rollout error under approximate commutativity and normality. We evaluate UNet and FNO variants with commutativity regularization on 1D and 2D spatio-temporal data in synthetic and real settings, showing successful long-horizon rollouts over thousands of steps. We show that the method improves FourCastNet climate forecasts on ERA5 without using any new data. The gain is most pronounced out-of-distribution: trained on trajectories of a few hundred steps, regularized models remain in-distribution for thousands of rollout steps on initial conditions where baselines diverge.

Visual reinforcement learning (RL) is promising for real-world problems such as autonomous driving, robotic locomotion, and manipulation, but remains highly vulnerable to visual distribution shifts between training and deployment. Existing methods improve robustness through Q-consistency regularization, masking, and auxiliary objectives, yet most of them follow a coupled representation learning paradigm in which control-oriented and robustness-oriented representation learning are jointly optimized within the same training framework. We argue that this coupled structure can introduce harmful interference between control learning and robustness learning, degrading original-environment control performance and limiting generalization under visual shifts. To address this issue, we propose Separate Then Align Representations (STAR), which reformulates visual RL generalization as a problem of decoupled representation learning. Specifically, we first learn a control-relevant reference representation through online RL under weak augmentation, and then perform an offline post-representation alignment stage that maps strongly augmented inputs back to this reference representation. Experiments on RL-ViGen benchmarks spanning DMControl and Robosuite show that STAR achieves strong generalization under visual distribution shifts while better preserving performance in the original environment than prior coupled approaches. Our source codes are available in the supplementary material.

Generative agents have proven to be powerful assistants in a wide variety of contexts. Given this success, users are now deploying agents with minimal restrictions in open ended, multi-agent environments. Current methods for measuring the dynamics of open-ended multi-agent systems are limited to qualitative inspection. In this paper, we extend the process-theoretic notion of adaptive control charts to multi-agent systems to enable automated monitoring. Using simulation, we demonstrate that adaptive control charts are necessary for monitoring multi-agent systems that can learn from their environment. We further demonstrate, both empirically and theoretically, that adaptive control charts are susceptible to adversarial agents that defect sufficiently slowly. These results illustrate a fundamental tradeoff in multi-agent system control: either agents in a system cannot learn \textit{or} the system is susceptible to adversaries.

Model merging integrates task-specific fine-tuned models into a single multi-task model, but often suffers from parameter interference caused by conflicting task-vector updates. Existing methods typically mitigate conflicts by pruning task vectors based on weight magnitude or random heuristics, treating Transformers as unstructured ``bags of parameters'' and overlooking their inherent modularity. In this paper, we propose \textbf{C}ontribution-\textbf{A}ware \textbf{S}tructured \textbf{S}parsity (CASS), a unified framework that reduces parameter interference by identifying and preserving task-specific components. At the core of CASS is a contribution-aware structured mask that identifies task-relevant attention heads and FFN neurons. We instantiate this mask in two settings: CASS-Merging, the primary post-hoc setting where masks serve as a plug-and-play denoising filter for existing merging operators, and CASS-Tuning, an extension for scenarios with fine-tuning access where masks constrain gradients to reduce structural overlap between task vectors. Our analysis shows that task-relevant components are sparse and partially disjoint, supporting structured component-level filtering as an effective way to reduce merging interference. Extensive experiments across vision (ViT, 20 tasks) and language (RoBERTa, 8 tasks; Qwen2.5, 4 tasks) benchmarks demonstrate that CASS improves a range of representative merging baselines.

In many critical domains, features are not freely available at inference time: each measurement may come with a cost of time, money, and risk. Longitudinal prediction further complicates this setting because both features and labels evolve over time, and missing measurements at earlier timepoints may become permanently unavailable. We propose CAPE (Contrastive Acquisition-Plan Embeddings), a partitioning acquisition policy for longitudinal active feature acquisition. CAPE first defines NOCT, an oracle-based objective that scores a set of future feature-time acquisitions by its expected predictive loss together with its acquisition cost. Based on NOCT, CAPE learns a contrastive acquisition-utility manifold of masked partial observations, where proximity reflects the expected utility of future acquisitions. Offline, training states are embedded and partitioned at each rollout timepoint; candidate future plans are scored within each partition using plug-in NOCT, and the best partition-level acquisition plans are cached. During inference, CAPE retrieves the cached plan via nearest-centroid lookup, enabling low-latency adaptive acquisitions. Experiments on synthetic and real-world healthcare datasets demonstrate that CAPE outperforms nonparametric, greedy, and RL-based alternatives, achieving higher accuracy at lower acquisition costs.

While diffusion has drawn considerable recent attention from the language modeling community, continuous diffusion has appeared less scalable than discrete approaches. To challenge this belief we revisit Plaid, a likelihood-based continuous diffusion language model (DLM), and construct RePlaid by aligning the architecture of Plaid with modern discrete DLMs. In this unified setting, we establish the first scaling law for continuous DLMs that rivals discrete DLMs: RePlaid exhibits a compute gap of only $22\times$ compared to autoregressive models, matches Duo while using fewer parameters, and outperforms MDLM in the over-trained regime. We benchmark RePlaid against recent continuous DLMs: on OpenWebText, RePlaid achieves a new state-of-the-art PPL bound of $22.6$ among continuous DLMs and superior generation quality. These results suggest that continuous diffusion, when trained via likelihood, is a highly competitive and scalable alternative to discrete DLMs. Moreover, we offer theoretical insights to understand the advantage of likelihood-based training. We show that optimizing the noise schedule to minimize the ELBO's variance naturally yields linear cross-entropy (information loss) over time. This evenly distributes denoising difficulty without any case-specific time reparameterization. In addition, we find that optimizing embeddings via likelihood creates structured geometries and drives the most significant likelihood gain.

Modern object detectors achieve strong performance on standard benchmarks, yet their robustness to contextual variation remains insufficiently investigated. Prior evaluations largely rely on aggregate metrics such as AP on uncontrolled distribution shifts, which conflate multiple factors and can obscure how performance degrades under context change. We introduce \textsc{ContextShift}, a controlled benchmark for evaluating context dependence by systematically manipulating object--context relationships while preserving object appearance. Built on COCO 2017, it isolates context as an independent variable both implicitly, through geometric transformations and explicitly via synthetic and natural background substitutions - with a compatibility axis based on normalized pointwise mutual information (NPMI) for continuous evaluation. Across diverse detector architectures, we observe a consistent degradation pattern: false negatives increase up to 227\% and prediction volume decreases down to -44\%, while false positives remain relatively stable or decline. This observed suppression behavior is not revealed by aggregate metrics such as AP, which can mask substantial recall loss and changes in prediction dynamics. Further analysis suggests that degradation is not primarily explained by reduced confidence, but is instead consistent with a reduced formation of valid detection candidates. Moreover, we find that performance along the statistical compatibility axis is non-monotonic, peaking at intermediate NPMI and degrading toward both extremes, indicating that statistical co-occurrence does not correlate linearly with effective visual context. Finally, we show that a training methodology based on context-aware augmentations rather than the dataset alone, which yields overall better models: every augmented variant outperforms the dataset-only baseline on both original (unmanipulated) test images as well as manipulated images, indicating that data-oriented augmentation can partially recover performance lost to prediction-suppression failures by exposing the model to object--context decoupling during training.

In this work, we develop proximal preconditioned gradient methods with a focus on spectral gradient methods providing a proximal extension to the Muon and Scion optimizers. We introduce a family of stochastic algorithms that can handle a wide variety of convex and nonconvex constraints and study its convergence under heavy-tailed noise, through a novel analysis tailored to the geometry of the proposed methods. We further propose a variance-reduced version, which achieves faster convergence under standard noise assumptions. Finally, we show that the polynomial iterations used in Muon are more accurately captured by a nonlinear preconditioner than by the ideal matrix sign, leading to a convergence analysis that more faithfully reflects practical implementations.

Biological recurrent circuits can reuse a common substrate by changing its operating regime rather than replacing the circuit itself. This principle is difficult to isolate in fully adaptive neural models, where recurrence, controllers, readouts, and internal representations can all absorb task structure. As a result, successful multi-context behavior does not reveal which route made the computation possible. We introduce Constrained Modulatory Reservoirs (CMR), a fixed-substrate framework for studying such reconfiguration with the main bypasses removed. The primary dynamics and low-rank perturbation basis are fixed, and the readout is denied direct access to the cue, forcing input-dependent adaptation to act through a limited set of recurrent connectivity changes. In transient-cue channel-selection tasks, this restriction makes a capacity transition visible. Low-loss behavior appears only when the available perturbation directions are sufficient, the selected change is retained across the cue-to-signal gap, and the primary dynamics express it during the response window. Controls show that this effect is not explained by low-dimensional cue representation alone. Allowing the substrate to learn hides the bottleneck, while activity- and gain-level alternatives fail to reproduce the same regime. CMR therefore reframes flexible recurrent computation as a question of where adaptive influence is allowed to enter a fixed dynamical system, making capacity, memory, and expression experimentally separable.