NeurIPS 2026 Papers

Inverse Reinforcement Learning (IRL) is classically solved using Lagrangian optimization, by jointly optimizing a primal-dual problem with gradient descent to obtain the optimal reward function parameters and corresponding optimal policy. While this algorithmic view of IRL has been highly influential, gradient-based reward updates only paint a partial picture, largely ignoring the geometry of the problem in occupancy space. In this paper, we present an occupancy-space view of reward updates in IRL. We leverage the insight that both the optimal solution and its geometry in occupancy space are known, and show that the geometry along the e-geodesic (straight line in occupancy space) toward the optimal occupancy directly affects the gradient of the IRL dual problem in reward parameter space. This perspective naturally suggests preconditioning the dual gradient with the average Fisher information along the occupancy trajectory. We show that this preconditioned gradient is just a convex interpolation of the old reward and a new target reward, yielding efficient natural-gradient-style updates with little additional computational overhead. We then analyse these occupancy-space reward updates from a Mirror Descent (MD) perspective, and show that policy search over the resulting rewards induces approximate entropic MD on the policy. Finally, we conclude with convergence results of such geometry-aware IRL algorithms.

We propose the Tikhonov layer, a graph neural network layer that is interpretable by design: once trained, its learned parameters directly reveal which node features and which aspects of the graph topology were leveraged for prediction. In practice, the layer's propagation matrix takes the closed-form $R = (p(L)+Q)^{-1}Q$, where $L$ is the normalized graph Laplacian, $Q = \mathrm{diag}(q_1,\ldots,q_n)$ a learnable diagonal matrix of positive node-importance scores, and $p(\cdot)$ a learnable polynomial. For any input feature $x$, the layer output $Rx$ is the exact minimizer of a generalized graph Tikhonov problem that trades off node-level data fidelity against a topology-driven regularization penalty. The learned pair $(\{q_i\},p)$ constitutes a built-in explanation: large $q_i$ indicates that node $i$'s own features drive the prediction, while small $q_i$ signals reliance on the local graph topology; the shape of $p$ reveals whether homophily, heterophily, or a band-pass response is exploited. Expressivity is preserved by routing complexity through a dedicated, arbitrarily deep $Q$-network that produces the importance scores, while the Tikhonov layer itself remains transparent. We prove that distinct explanations necessarily yield distinct propagation operators, preventing degenerate explanations by design. Additionally, the Tikhonov layer provides, in a single layer, a global receptive field under mild support-connectivity conditions, mitigating both oversmoothing and oversquashing. Experiments on standard graph classification benchmarks confirm that the model matches (and sometimes outperforms) opaque baselines while producing interpretable and faithful explanations.

Vision-language models are designed to capture the compositional structure of the world through multi-modal alignment; however, it remains unclear whether their representations truly align with the way a human would naturally describe the same visual input. Prior works attempt to decompose visual embeddings into their concept-level contributions, but their instance-level optimization is inherently local and does not capture the model’s global semantic structure, which is essential for compositional generalization across different inputs, tasks and domains. We introduce GRACE, a method that learns this concept-based structure at the feature space level, producing semantic and general concept decompositions. To achieve this, we propose a novel training objective that preserves the geometric relations of different samples in the concept space, enabling a better understanding of how concepts compose and relate to each other. We validate GRACE on several image and video benchmarks, showing that its concept decomposition is aligned with human perception and achieves faithful decompositions across different tasks and domains. With minimal overhead, GRACE makes existing Vision-Language Models more interpretable, while preserving most of their original performance.

Post-training pruning can substantially reduce LLM inference costs, but it often degrades quality unless the remaining weights are adapted. Since global retraining is expensive at LLM scale, recent work has largely focused on increasingly sophisticated pruning criteria that aim to select better sparsity patterns without adaptation. We revisit this trade-off through local reconstruction: after pruning, we adapt one subset of the model parameters at a time on a calibration set, training it to match the corresponding intermediate activations of the dense model. We evaluate local reconstruction across model families and scales, up to 72B parameters, and establish three main findings. First, local reconstruction is an effective adaptation mechanism for LLMs: it matches post-pruning retraining while using over an order of magnitude less data and compute, even when using PEFT techniques. Second, reconstruction exhibits a broad "free-lunch" regime in granularity, i.e., the reconstruction parameter window: as long as the reconstructed region contains at least a nonlinear submodule, final quality is largely insensitive to the window size, allowing granularity to be chosen primarily based on memory constraints. In contrast, reconstructing individual matrices, despite being the natural approach often proposed in the literature, consistently underperforms, as small matrix-level errors accumulate into larger activation drift. Lastly, reconstruction reduces the relative importance of the pruning criterion: performance gaps between sophisticated criteria and simple baselines shrink with model scale, making simple methods competitive again. Overall, our results challenge the prevailing view that post-pruning adaptation is impractical for LLMs.

The rapid scale-up of AI infrastructure has made datacenter reliability a cornerstone of the cloud business economy. Hardware failures in hyper-scale AI clusters do not merely disrupt individual machines; they stall multi-million dollar training workloads, and so the downtime of individual machines can cause outsized economic losses. Minimizing machine downtime, therefore, is more important than ever. To address this, we present DrWatson, an AI-driven hardware diagnosis system that pioneers a foundation-model approach to automated datacenter maintenance. Deployed at production scale within a major AI hyperscaler, DrWatson yields substantial efficiency gains across diverse compute architectures. Real-world evaluations demonstrate that it reduces production downtime by 18.3% on GPU platforms and 16.1% on custom AI accelerators. Furthermore, it significantly accelerates hardware deployment, cutting quality assurance (QA) time by 28.3% for GPUs and 25.1% for custom accelerators. These results establish DrWatson as a field-tested solution for maximizing the reliability and cost efficiency of next generation AI fleets.

Modern training pipelines for language model agents begin with a supervised fine-tuning stage in which a small student imitates a costly teacher. Recent work mitigates the covariate shift of pure imitation learning by collecting teacher feedback at states the student itself reaches, with a prevailing trend toward elaborate filters on the teacher's responses. We frame this design choice as a budget-allocation problem and compare three constructions of supervised training data: short unfiltered teacher continuations at learner-induced states; full teacher trajectories filtered for success (the rejection-sampling step in recent on-policy expert-correction work); or those further restricted to tasks the student cannot solve on its own. Across three agentic benchmarks (HotpotQA, ALFWorld, and Terminal-Bench-Dev), short unfiltered teacher continuations beat pure behavioral cloning at matched supervision budgets, and on HotpotQA also match or exceed the filtered alternatives. On Terminal-Bench-Dev, this construction at one-tenth of the corpus budget and with no reinforcement-learning stage matches the OpenThoughts-Agent baseline that uses the full corpus together with reinforcement learning; the optimal continuation length is task-dependent. The same supervised-learning checkpoints yield faster early-stage gains under subsequent reinforcement learning, with mixed evidence on whether this advantage persists later in training.

Few-step diffusion sampling sharpens the tension between ODE and SDE dynamics. ODE samplers are stable but can lose diversity, while the standard SDE endpoint can become unreliable at low numbers of function evaluations (NFE). The challenge is therefore not simply whether stochasticity should be added, but which driver properties allow stochastic correction to contribute diffusion without losing finite-step control. We show that this tradeoff is governed by two driver-level quantities, the asymptotic covariance $Q$, which sets the effective diffusion level on the ODE-to-SDE spectrum, and the Poisson solution norm $C_\chi$, which controls the finite-step remainder. Gaussian drivers have $C_\chi = \infty$, so this finite-step remainder is not uniformly controlled by the same bound. Balancing the diffusion level set by $Q$ with the finite-step control set by $C_\chi$ formalizes a favorable regime for bounded structured drivers with intermediate diffusion strength. This regime directly gives a structured correction method, instantiated as Fast-Driver Sampler (FDS) with the shuffled-Lorenz driver. Across DiT-XL/2 and EDM2, FDS establishes a stronger low-NFE frontier among evaluated training-free samplers, improving ODE baselines while avoiding the degradation of the SDE endpoint.

Linear structural causal models (SCMs) are used to analyze relationships among random variables. In these models, directed edges encode direct causal effects, while bidirected edges capture latent confounding. The problem of generically identifying the hidden parameters from observed correlations remains open in causal inference. The half-trek criterion by Foygel, Draisma, and Drton [2012] and its generalization to the edge-wise criterion are important criteria for generic identification in linear SCMs, since the half-trek criterion can be efficiently decided by an algorithm with running time $\mathcal{O}(n^5)$. We here develop a new criterion that is equivalent to the edgewise criterion. Utilizing this criterion, we design a faster randomized algorithm for deciding the half-trek criterion with running time $\mathcal{O}(n^4)$. For sparse graphs, our algorithm even works in $\mathcal{O}(n^3)$. We show that the algorithm is useful in practice by providing an implementation that outperforms the implementation of HTC in the state-of-the art R package SEMID.

The MOEA/D is one of the most successful multi-objective evolutionary algorithms. It operates by decomposing the optimization problem into single-objective subproblems, which are then solved in a co-evolutionary manner. The performance of the MOEA/D depends critically on this decomposition. In this work, we propose a novel way to adjust this decomposition on the fly, for the case of discrete bi-objective optimization. Different from all previous works in this direction, we support our new algorithm with a proven performance guarantee. This is the first work that rigorously proves the runtime of a dynamic multi-objective evolutionary algorithm. For monotonic distortions of the simple LOTZ benchmark, it is known that classic decompositions leave large holes in the Pareto front, making it very hard for the MOEA/D to cover the entire Pareto front. In contrast, we prove that our self-adjusting MOEA/D solves all these LOTZ problems in an expected number of $O(n^4)$ function evaluations. This relatively low runtime as well as the details of our mathematical proof show that the dynamic decomposition strategy indeed evolves a decomposition adjusted to each of these problem instances.

Positional encoding in Vision Transformers presents a recurring trade-off: spatial relative-position tables (Swin) carry thousands of parameters per block and require ad-hoc interpolation to transfer across resolutions, while no-bias designs leave the attention logits geometry-agnostic. We propose a relative positional bias parameterized directly in the frequency domain. This parameterization yields three benefits over the standard spatial-table position encoding: (i) it expresses the same translation-equivariant bias with $3.25\times$ fewer parameters than Swin's spatial table at the same head count; (ii) it admits a closed-form, structurally exact transfer to a different patch grid, in contrast to bilinear interpolation of spatial tables which smooths the learned structure; and (iii) it shares the same parameterization style as a frequency-domain spectral mixer, enabling co-design of positional encoding and token mixing in a common substrate. We demonstrate the bias inside our **Dual-Domain Residual (DDR)** framework, in which a spectral mixer and a gated attention head act as parallel residuals within each block. Head-to-head, the frequency-domain bias recovers $+1.98$ pp over Swin RelPos and $+2.74$ pp over no positional bias; removing the entire attention branch (which carries the bias) drops accuracy by $5.35$ pp, while removing the spectral branch drops it by $2.22$ pp. On ImageNet-1K from scratch, DDR outperforms SpectFormer at every comparable parameter budget, by $+2.0$ pp at Tiny (DDR-Ti-SF, 10.7M, 78.9\%) and $+0.4$ pp at Base (DDR-B-Deep, 63.1M, 82.5\%); at sub-parity parameters, DDR-Ti+ (8.9M, 77.9\%) exceeds SpectFormer-Ti (9.0M, 76.9\%) by $+1.0$ pp. The learned per-layer gates concentrate spectral processing in early blocks and attention in late blocks, qualitatively matching the schedules that are hand-coded in the previous hybrid models.

Decentralized stochastic minimax optimization has recently attracted significant attention due to its applications in machine learning. However, existing state-of-the-art methods use learning rates of different scales for the primal and dual variables, making them difficult to tune in practice. To address this problem, this paper proposes a novel doubly smoothed decentralized stochastic minimax algorithm. Specifically, in terms of algorithm design, we update both the primal and dual variables using smoothed gradients and introduce novel approaches to handle the computation and communication of the auxiliary variables introduced by the smoothing technique. On the theoretical side, for nonconvex-PL problems, our convergence analysis reveals that the learning rates for the primal and dual variables are of the same scale. Moreover, the order of the condition number in our convergence rate is improved to $O(\kappa^{3/2})$. To the best of our knowledge, this is the first time it has been improved to such a favorable order. Finally, extensive experimental results validate the effectiveness of our algorithm.

3D scene graph prediction is commonly supervised with local object and predicate classification losses. Although effective for slot-wise label prediction, this formulation does not provide a graph-level semantic distance between a predicted scene graph and its ground truth: semantically mild label errors and structurally disruptive relational errors can be treated similarly, and subject-predicate-object facts are not compared as scene-level units. We propose Contextual Hellinger Triplet Geometry, a graph-level supervision objective that measures the semantic and structural discrepancy between 3D scene graphs. Our key idea is to reinterpret the conventional classifier outputs of a 3D scene graph predictor as an aligned probabilistic scene graph field, where object and predicate predictions jointly define contextual subject-predicate-object facts. This formulation yields a differentiable metric-based training objective with an efficient factorized implementation and can be applied to existing 3D scene graph prediction architectures without modifying their network design. Experiments on 3DSSG demonstrate that our supervision improves downstream graph-conditioned 3D scene generation, while ablation studies confirm the complementary roles of the proposed graph-level terms and user studies show better graph descriptiveness, relation plausibility, and structural preservation.

Graphic-design layouts are structured, multi-layered artifacts: a single composition specifies the position, type, z-order, and styling of dozens of components, and a single template typically begets a family of sibling variants that share its structural theme but differ in content, palette, or imagery -- properties that flat-canvas layout benchmarks discard. We introduce \textsc{GraphicDesignBench} (GDB), a diagnostic benchmark of $16$ tasks along two axes: \emph{spatial composition} ($8$ understanding + $4$ generation) and \emph{template variants} ($3$ understanding + $2$ generation), grounded in a new dataset of $1{,}148$ real-world templates with full component hierarchy, per-element styling, and sibling groupings. Across four frontier VLMs (GPT-5.4, Gemini-3.1-Pro / Flash-Lite, Claude-Opus-4.6) and two image generators (GPT-Image-1.5, Gemini-3.1-Flash-Image), GDB exposes concrete and dissociated gaps: component detection sits an order of magnitude below natural-image baselines, layer-order reasoning decouples from other spatial skills, partial completion collapses from single- to multi-element settings, and a simple font-Jaccard baseline matches or exceeds frontier VLMs on sibling matching and clustering. We release the dataset, tasks, metrics, and per-model outputs to track where current models are, and are not, ready to act as design collaborators.

We propose a deep learning framework for scalar-on-function regression with possibly multiple functional predictors observed on irregular grids. The method represents each functional covariate by finitely many basis coefficients computed from the observed trajectories and fits a deep neural network using these coefficients as features. We establish a general nonasymptotic excess risk bound for the neural network estimator that decomposes the error into statistical, truncation, discretization, and network approximation terms. We then specialize the theory to two important settings. (i) For continuous functional models, an equal-width neural network estimator achieves a nearly minimax optimal rate up to a log-logarithm factor. (ii) For functional multiple-index models, we design a bottleneck-equal-width neural network estimator that exploits the low-dimensional index-link structure and attains a nearly minimax optimal rate up to a logarithm factor. The results provide theoretical support for neural network-based scalar-on-function regression and highlight the importance of architecture design.

Submodular functions have applications in several domains, \eg, text, vision, speech, \etc. Recent works have proposed neural architectures that are submodular functions by design. However, they do not explicitly capture pairwise interactions between set elements. In this work, we begin with the characterization of pairwise submodular functions, which capture the pairwise interaction across the elements of the input set. We observe that pairwise interactions yield monotone supermodular functions, since the number of terms grows quadratically in terms of the input set size, which poses a significant challenge in converting it into a submodular function. To address it, we provide a novel result that a derivative rate-controlled concave function can transform a monotone supermodular function into a monotone submodular function. We also extend our results to monotone $\alpha$-submodular functions. Leveraging this characterization, we design multi-layer neural cross-interaction architectures for monotone submodular functions and analyze their expressivity. Finally, we perform several experiments which show that our model performs better than existing baselines.

Posterior collapse is a common failure mode in variational autoencoders, but existing theory has mostly focused on Gaussian latent variables. Dirichlet VAEs are comparatively less understood, despite their use in settings where latent variables represent proportions, abundances, memberships, or convex mixture weights. We derive a closed-form local stability threshold $\beta_c$ for collapse in symmetric Dirichlet $\beta$-VAEs with linear Gaussian decoders. This threshold shows that collapse in those models depends not only on the data spectrum, but also on the simplex dimension and prior concentration through a trigamma curvature factor induced by the Dirichlet KL. Although $\beta_c$ is derived from a local analysis, we find that it accurately tracks the empirical onset of collapse beyond the local regime. We show its practical relevance in two real-world linear-mixture settings: hyperspectral unmixing and convex archetypal analysis. Across experiments, $\beta/\beta_c$ consistently organizes collapse-related behavior, providing an interpretable coordinate and a practical way to narrow the search over the $\beta$ hyperparameter before training. Taken together, these empirical findings suggest that $\beta_c$ can serve, beyond its local derivation, as a practical reference point for anticipating collapse and guiding $\beta$ selection in Dirichlet $\beta$-VAEs.

Backpropagation remains the dominant method for training deep neural networks, but its reliance on non-local computations has motivated the search for biologically plausible alternatives. Two prominent frameworks, predictive coding (PC) and deep feedback control (DFC), offer distinct approaches to local learning, yet their relationship remains unclear. Here, we show that DFC arises as a local approximation to predictive coding dynamics near fixed points, with exact equivalence in shallow linear networks and controlled deviations arising with depth, nonlinearity, and distance from equilibrium. Building on this connection, we introduce Approx PC, a practical algorithm that replaces layer-wise error propagation in PC with a centralized feedback controller operating over a sequence of intermediate subgoals. Empirically, Approx PC recovers solutions consistent with predictive coding, while converging faster and achieving improved accuracy in our experiments. Our results unify energy-based and control-theoretic perspectives on learning and suggest a computationally efficient approach to training predictive coding networks through control-inspired dynamics.

This paper studies the problem of jointly learning a reward function, a cost function, and a policy from an expert's preferences. We formulate the problem as a constrained bi-level optimization problem, where the upper level infers the reward and cost functions from preferences, while the lower level optimizes a policy to best align with those preferences. To solve this problem, we propose a double-loop algorithm, Constrained Bi-level Optimization for Preference-Based Reinforcement Learning (CB-PbRL), which solves the lower-level optimization problem in the inner loop and the upper-level optimization problem in the outer loop. We establish a theoretical guarantee that CB-PbRL converges at a rate of $\mathcal{O}(1/\sqrt{K})$, and we demonstrate its effectiveness across multiple simulation environments.

\emph{Scaling} has been the central driver of progress in modern AI, yet data curation remains an exception: it still relies on manual heuristics rather than systematic, scalable methods. In practice, data recipes (e.g., filter thresholds, deduplication strength, domain mixing ratios) are determined through many small-scale ablation training runs. This work demonstrates that the performance of the best-discovered recipe improves \emph{predictably}--as a \emph{power law}--in the number of such ablation runs. Theoretically, the power-law form arises naturally from regret bounds in zero-order optimization theory. Empirically, we validate it across diverse training stages, including pretraining, supervised fine-tuning, preference optimization, and reinforcement learning. Building on this finding, we derive a principled approach to \emph{data quality--quantity tradeoff} under a fixed compute budget, yielding a Pareto frontier that dominates standard heuristic pipelines (e.g., DCLM). Together, this work provides a general framework for scaling data-curation compute, establishing it as a new axis alongside the well-studied pretraining and test-time scaling.