NeurIPS 2026 Papers

Conference on Neural Information Processing Systems

Decision breakdown

6,437papers

Counted from the official paper list; may differ slightly from the number the conference announced.

  • Oral
    991.5%
  • Spotlight
    2403.7%
  • Poster
    6,09894.7%

Browse all papers

How Data Augmentation Shapes Neural Representations

Tianxiao He, Alex Williams, Sarah Harvey

Expand abstract

Data augmentation is widely recognized for improving generalization in deep networks, yet its impact on the geometry of learned representations remains poorly understood. In this work, we characterize how different data augmentation strategies reshape internal representations in neural networks. Using tools from shape analysis, we embed network hidden representations into a metric space where distance is invariant to scaling, translation, rotation and reflection. We show that increasing augmentation strength leads to well-behaved trajectories in this space, and that different augmentation types steer representations in distinct directions. Moreover, we investigate how neural representation shapes are distorted along data augmentation trajectories, and show that insights from neural geometry can predict which representations provide the most improvement when ensembling models. Our results reveal shared geometric patterns across architectures and seeds, and suggest that analyzing shape-space trajectories offers a principled tool for understanding and comparing data augmentation methods.

Poster
How Complete Should a Reference Be? A Benchmark Audit for Fluorescence Spot Detection

Jiyeong Kong, Pan-Gyun Jeong, Kyung-Tae Kang

Expand abstract

Fluorescence spot-detection benchmarks in biology and medicine often rely on reference annotations that are incomplete, ambiguous, or inconsistently defined. In such settings, detector performance is not determined solely by model predictions; it can also depend on which faint or ambiguous signals are counted as reference positives. We formulate this issue as a reference-audit problem: instead of treating the evaluation reference as fixed ground truth, we test whether benchmark outcomes remain stable under plausible changes in reference definition and annotation completeness. Using the public FISH_spots dataset, a benchmark for fluorescence in situ hybridization (FISH) spot detection, we construct Raw, Lenient, and Strict audit references and evaluate diverse classical and learning-based detectors under a unified point-matching protocol. We further apply annotation-retention stress tests that simulate random and structured missing-label mechanisms. Across these audits, changing only the evaluation reference can alter measured F1 scores and model rankings, including the selected top-1 detector, even when aggregate rank correlation remains high. These findings suggest that model selection in fluorescence spot-detection benchmarks should be tested across plausible reference definitions, rather than based on a single reference annotation set.

Poster
How Can SignSGD Outperform SGD? A Functional Scaling Law Perspective

Zilin Wang, Binghui Li, Jia-Nan Wang, Lean Wang, Jinbo Wang, Lei Wu

Expand abstract

We study signSGD in linear regression with diagonal features and power-law spectra, a toy but expressive setting where the effect of coordinate-wise sign normalization can be analyzed sharply. Specifically, we establish a functional scaling law (FSL) for signSGD under general learning rate schedules, which decomposes the loss dynamics into two components: learning of the target signal and noise accumulation. Compared with SGD, signSGD learns the signal faster but exhibits slower noise forgetting. This signal-noise tradeoff yields a regime-dependent comparison between signSGD and SGD: signSGD achieves better data-scaling efficiency in hard-task regimes, while they perform similarly in easy-task regimes. The key mechanism is that gradient noise is {\it curvature-aligned}: its directional variance is of the same order as the directional curvature. As a result, sign operation effectively acts as a \(\diag(\bH)^{-1/2}\) preconditioner in the noise-dominated regime, where $\bH$ is the Hessian matrix. Moreover, we demonstrate that the same square-root preconditioning mechanism can arise in RMSprop and Adam, where coordinate-wise normalizers essentially track the coordinate curvatures. Synthetic experiments validate the proposed mechanism and the predicted scaling behavior. Language model experiments further suggest that, beyond the controlled setting, the resulting FSL can serve as a surrogate model for fitting and predicting practical training loss curves.

Spotlight
Horizontal Diffusion Models: Score-based Generative Modeling on Frame-Connection Geometry

Sungwoo Park

Expand abstract

Geometric modeling has become an essential paradigm in AI for Science, where data often inhabit intrinsically curved or structured spaces. However, existing manifold generative models often rely on closed-form or tractable access to global geometric primitives, such as geodesic distances, parallel transport, or spectral decompositions, which are rarely available beyond canonical geometries. To overcome this limitation, we propose Horizontal Diffusion Models (HDMs), a score-based generative framework defined on the orthonormal frame bundle of a Riemannian manifold. Instead of requiring global geometric operations, HDM lifts Euclidean diffusion processes through the Levi-Civita horizontal distribution, using only local metric and connection information while preserving intrinsic manifold geometry. This construction allows standard Euclidean score networks to be lifted into geometry-consistent and gauge-equivariant horizontal vector fields, bypassing the need for manifold-specific neural architectures. We further derive a horizontal KL objective that reduces to a Euclidean score-matching loss and analyze the curvature-dependent cost of the lift through a generalization bound. Experiments on parametric surfaces and scientific datasets demonstrate high-fidelity generation across diverse manifolds with nontrivial geometry and topology.

Poster
Horizon Adaptive Offline Policy Learning via Value Stitching

Kexin Zheng, Xianyuan Zhan, Xintao Yan

Expand abstract

Learning accurate value functions plays a decisive role for reinforcement learning (RL) agents to solve long-horizon, complex tasks. Conventional temporal-difference (TD) learning objectives suffer from value-estimation bias that accumulates over the horizon, while extended-horizon modeling methods, such as $n$-step TD backups and Q-chunking, adopt a rigid, fixed-horizon value-modeling recipe that is often not flexible enough to capture complex value structures in long-horizon, multi-stage tasks. In this paper, we show that enabling value updates with dynamic horizon composition can yield a strong offline policy learning scheme. Our method, _Horizon Adaptive Offline Policy Learning via **VA**lue **ST**itching_ (**_VAST_**), replaces fixed-horizon backups with recursive, horizon-adaptive value composition. Its key ingredient is to couple value optimization with a future state- and horizon-length-conditioned **_auxiliary value function_** that is learned through direct data supervision, and a **_stitching policy_** that optimally selects the reward-maximizing horizon length and future sub-goal to achieve horizon-adaptive value stitching. This design enables direct estimation and compositional “stitching” of variable-length returns grounded in actionable sub-goal states, providing an accurate and greedily exploitable value-supervision signal for offline policy optimization. Across 50 tasks on OGBench, _VAST_ outperforms fixed-step, extended-horizon methods, and generative-value offline RL baselines, achieving strong performance particularly in high-complexity, long-horizon decision-making tasks.

Poster
HopChain: Multi-Hop Data Synthesis for Generalizable Vision-Language Reasoning

Shenzhi Wang, Shixuan Liu, Jing Zhou, Chang Gao, Xiong-Hui Chen, Binghai Wang, An Yang, Shiji Song, Bowen Yu, Gao Huang, Junyang Lin

Expand abstract

Vision-language models (VLMs) show strong multimodal capabilities but still struggle with fine-grained vision-language reasoning. We find that long chain-of-thought (CoT) reasoning exposes diverse failure modes, including perception, reasoning, knowledge, and hallucination errors, which can compound across intermediate steps. However, most existing data for reinforcement learning with verifiable rewards (RLVR) does not involve complex reasoning chains grounded in visual evidence throughout, leaving these failure modes understressed during training. We therefore propose HopChain, a scalable framework that constructs an out-of-distribution (OOD) proxy task: OOD only in query format, while strengthening foundational abilities shared with downstream benchmarks to support generalizable gains. Concretely, HopChain synthesizes multi-hop vision-language reasoning chains, each a logically dependent chain of instance-grounded hops where earlier hops establish the instances, sets, or conditions needed for later hops, ending in an unambiguous numerical answer for verifiable rewards. In experiments, we train Qwen3.5-35B-A3B and Qwen3.5-397B-A17B under two RLVR settings (with and without HopChain's multi-hop data) and compare them across 24 benchmarks spanning STEM and Puzzle, General VQA, Text Recognition and Document Understanding, and Video Understanding. Although the multi-hop data is not designed for any specific benchmark, it improves 20 of 24 benchmarks on both models, indicating broad and generalizable gains. Consistently, replacing full chained queries with half-multi-hop or single-hop variants reduces the average score across five representative benchmarks from 70.4 to 66.7 and 64.3, respectively. Improvements are particularly substantial in long-CoT and ultra-long-CoT regimes, peaking at more than 50 accuracy points in the ultra-long-CoT regime. These results show that OOD proxy tasks for long-chain visual reasoning are a scalable RLVR supervision source for generalizable VLM reasoning.

Poster
HoloGene: Learning to Lift Sliced Spatial Transcriptomics to Holistic 3D Gene Fields

Yumin Zheng, Qingtian Zhu, Ziyan Zhu, Kailu Song, Yinqiang Zheng, Jun Ding

Expand abstract

Spatial transcriptomics (ST) provides spatially resolved gene expression measurements, but most 3D ST datasets are acquired as independently processed 2D sections. This sliced acquisition creates a difficult reconstruction problem: measurements are dense within each section but sparse along the axial direction, and inter-section technical variation can be confounded with true biological change. Existing coordinate-based implicit neural representations (INRs) provide a continuous modeling framework, but standard isotropic coordinate encodings do not explicitly reflect this axial-to-lateral imbalance. We present HoloGene, an anisotropic conditional INR framework for estimating continuous 3D gene-expression fields from pre-registered ST sections. HoloGene first compresses high-dimensional expression profiles into a graph-autoencoder latent manifold. It then uses pseudo-label-derived biological signature conditioning to provide coarse intra-section biological context, while depth-dependent FiLM modulation models axial variation. To reduce slice-specific covariance shifts, we introduce a signature-stratified correlation-alignment regularizer over latent features. Using held-out measured sections as cross-sectional references, HoloGene improves reconstruction accuracy and biological preservation metrics compared with general-purpose INRs and adapted spatial-transcriptomics baselines. Ablation analyses show that decoupled depth conditioning improves axial generalization, while correlation alignment provides the strongest reduction in slice-specific representation leakage. These results support anisotropic conditional modeling as a practical strategy for estimating continuous 3D expression fields from discontinuous ST sections.

Poster
Holo4D: Holistic 4D Reconstruction as Geometric Control for Video Diffusion

Yushi Lan, Zeren Jiang, Koichi Namekata, Xingang Pan, Chuanxia Zheng, Andrea Vedaldi

Expand abstract

Reconstructing a dynamic 4D scene along a novel camera trajectory requires both metric geometry from the source video and generative completion of disoccluded target-view regions. Feed-forward 4D reconstruction models recover increasingly comprehensive input-view geometry, including depth, point maps, and dense 3D tracks, but remain tied to the observed camera stream. Camera-controlled video diffusion models (VDMs) provide strong generative priors for novel-view synthesis, yet their camera control is typically not precise enough for metric novel-view reconstruction. We argue that comprehensive input-view 4D reconstruction is an effective geometric interface between these two capabilities: properly exposed to a VDM, it turns reconstruction outputs into camera-accurate novel-view RGB videos that remain geometrically consistent under downstream 4D reconstruction. We introduce HOLO4D, a geometry-aware video-to-video framework that controls a pretrained VDM with a hybrid 4D geometric cache built from the source video. A small set of source views covering the camera motion forms the static cache, while dynamic regions are rendered from the matching source frame. The resulting hybrid RGB-D representation is combined with explicit valid masks, target-camera ray embeddings, and dense 3D tracking tokens. Conditioned on this 4D scaffold, the diffusion backbone synthesizes the target-view RGB video. Under a common downstream 4D reconstructor, our generated videos yield target-view depth more consistent with the held-out reconstruction than that of the baselines, and an additional feature-level readout shows that intermediate VAE features already encode target-view depth. Experiments and ablations further show that progressively richer 3D/4D reconstruction signals improve camera adherence.

Poster
Holistic Scaling Laws for Optimal Mixture-of-Experts Architecture Optimization

Weilin Wan, Jingtao Han, Debing Zhang, WEIZHONG ZHANG, Cheng Jin

Expand abstract

Scaling laws govern macroscopic resource allocation for LLMs, yet precise architectural configurations for Mixture-of-Experts (MoE) models remain guided by heuristics. Existing MoE scaling studies either incorporate MoE-specific variables into scaling formulas, causing fitted coefficients to grow rapidly without proportional experimental support, or fix all non-MoE factors, implicitly assuming global architecture does not influence local MoE scaling. We propose a reusable framework for holistic MoE architectural optimization. We first reveal that relying solely on FLOPs per token ($M$) biases MoE evaluation, as heterogeneous Attention/FFN computational densities enable parameter inflation without effective compute gains. We therefore establish a joint constraint triad of $M$, active parameters ($N_a$), and total parameters ($N$) for rigorous MoE characterization. To tame the resulting $\mathcal{O}(n^{16})$ search space, we employ mathematical decoupling: structural constraints and a rank-preserving property of the hidden dimension factorize the optimization into an $\mathcal{O}(n^3) + \mathcal{O}(n^2)$ two-phase search. Through extensive validation across 670+ MoE models spanning $10^{18}$ to $3 \times 10^{20}$ FLOPs, we derive globally applicable scaling laws mapping any compute budget to its optimal architecture. A key finding is that the near-optimal configuration band widens with scale, providing a quantitative basis for trade-offs between scaling law recommendations and engineering constraints. Our work delivers actionable blueprints for optimal MoE design under arbitrary compute budgets.

Poster
Hodge Laplacian Quasi-Harmonic Flows for Option Discovery

Andreas Theophilou, Özgür Şimşek

Expand abstract

Prior Laplacian-based option-discovery methods build options from 0-form Laplacians defined on nodes, where the resulting eigenvectors induce routes toward isolated boundaries or extrema, thereby skewing visitation toward specific states. In contrast, kernel eigenvectors of the 1-form Hodge Laplacian represent interpretable, topology-aware edge flows and naturally capture cyclic behaviour known as harmonic flows. These kernel eigenvectors are particularly attractive because they define zero-divergence edge flows, offering a natural solution for non-isolated visitation to sources and sinks in regions in which they are active. Despite this appeal, harmonic flows have not previously been explored as a basis for option discovery in reinforcement learning. This is what we do here. We first introduce harmonic-flow skills and show that, on state graphs with tree-like “dead-end” appendages, purely harmonic flows can become nearly inactive in those regions, which can limit coverage. To address this limitation, we propose quasi-harmonic flow options: skill-inducing edge flows that preserve harmonic cycles where harmonic activity is strong, while injecting drift into appendages using the lowest-magnitude non-harmonic eigenvector of the Hodge Laplacian, a low divergence solution, where harmonic activity is weak. Our approach is eigenvector-only, relying solely on spectral components of the Hodge Laplacian. We compare quasi-harmonic flow options against a harmonic-flow options baseline and state-of-the-art skill-discovery methods from the literature, demonstrating empirical gains and analytical improvements on hitting-time-based metrics.

Poster
HoTS: Homophily-Aware Temperature Scaling for Graph Neural Network Calibration

Inwoo Tae, Yoontae Hwang, Yongjae Lee

Expand abstract

For graph node classification, calibrated class probabilities are need when confidence scores, usually the maximum predicted class probability, are used to rank prediction, defer uncertain nodes to human review or control risk. Existing post-hoc calibrators either apply one global temperature or use graph-awared modules without a derived strctural form. We study node-level calibration for graph node-classification problem under structural heterogeneity. Our first results show that a logit-only temperature rule is insufficient when nodes with identical logits but different local homophily require different optimal temperatures. We then analyze a population-concentration contextual stochastic block model with Gaussian features and a one-layer linear GCN. Under equidistant class means, the Bayes posterior over class-template scores is a temperature-scaled softmax whose inverse temperature is governed by homophily-dependent signal strength. In the positive-signal homophilic regime, the resulting temperature decreases approximately with normalized local homophily. This law motivates Homophily-aware Temperature Scaling (HoTS), a simple post-hoc calibrator that assigns each node a positive scalar temperature from entropy-based logit concentration and estimated local homophily. HoTS has three parameters, preserving class, and learn the strength of the structural correction from calibraion data. Across 18 node-classification benchmarks, two GNN backbones, and eight calibration baselines, HoTS achieves the best mean Expected Calibration Error (ECE) of $4.66\%$, the best average rank, and the strongest confidence ranking in selective classification.

Poster
Hitting a Moving Target: Test-Time Adaptation for AI Text Detection under Continual Distribution Shift

Kevin Ren, Manish Raghavan, Nikhil Garg

Expand abstract

Deployed approaches for AI text detection often rely on training-time access to labeled datasets of both human-written and AI-generated text. This approach is vulnerable to three types of distribution shifts that occur continually post-deployment, and for which labeled data is often unavailable: adversarial humanization, new LLMs being released, and temporal drift in human writing. Simultaneously, existing approaches do not leverage a key signal of LLM usage: inference-time homogeneity. We propose a test-time adaptation (TTA) approach, using semi-supervised learning, that adapts to distribution shifts by leveraging homogeneity among unlabeled samples observed at inference time. Empirically, we find that state-of-the-art supervised detectors systematically fail when they encounter distribution shifts in AI-generated and human writing, both adversarial and natural, while test-time adaptation with semi-supervised learning is largely robust; e.g., the commercial model Pangram detects just 24.1\% of our adversarial AI-generated text, compared to 90.5\% for our test-time approach. We establish that test-time adaptation is a promising framework for AI text detection in the wild.

Poster
Hitting Time Isomorphism for Multi-Stage Planning with Foundation Policies

Magnus Victor Boock, Abdullah Akgül, Mustafa Mert Celikok, Melih Kandemir

Expand abstract

We present a new operator-theoretic representation learning framework for offline reinforcement learning that recovers the directed temporal geometry of a controlled Markov process from hitting time observations. While prior art often produces symmetric distances or fails to satisfy the triangle inequality, our framework learns a Hilbert-space displacement geometry where expected hitting times are realized as linear functionals of latent displacements. We prove that this representation exists under latent linear closure and is uniquely identifiable up to a bounded linear isomorphism. For finite-dimensional implementations, we show that global hitting-time error is bounded by one-step transition error amplified by the environment's transient spectral radius. Furthermore, we provide finite-sample guarantees accounting for approximation, statistical complexity, and trajectory-label mismatch. Derived from this theory, we curate Isomorphic Embedding Learning (IEL) as a new goal-agnostic foundation policy learning algorithm that anchors a HILP-style consistency objective with explicit hitting-time regression to ensure that the learned geometry reflects actual decision-time progress. This asymmetric and compositional structure enables robust graph-based multi-stage planning for long-horizon navigation. Our experiments demonstrate that IEL improves the state of the art of learning foundation policy policies from offline maze locomotion data.

Poster
Hit Expansion via Localized Exploration of Synthesizable Chemical Space

Walter Virany, Yidong Jin, Andrew Lian, Dmytro Shevchuk, Amir Mehdi Soufi Enayati, Mike Tyers, Robert Batey, Piotr Gaiński, Michał Koziarski

Expand abstract

Generative models for drug design which directly produce synthetic pathways have gained significant popularity due to their ability to constrain the search space to synthetically accessible molecules. However, existing methods have focused primarily on de novo molecular design, and rarely start the generation process from known binders. In this paper, we present HELiX: a template-based GFlowNet for localized exploration of chemical space. HELiX learns to partially decompose a given synthetic trajectory to an intermediate state, and then perform forward synthesis in a manner that preserves synthetic accessibility, leading to diverse, high-reward analog generation. Moreover, we prioritize sample efficiency by incorporating Bayesian optimization into the training procedure. We diagnose problems inherent to training GFlowNets with Bayesian optimization and introduce a greedy acquisition strategy which effectively balances between exploration and exploitation without the need for reward shaping. Finally, we show that local exploration is inherently robust to noisy oracle evaluations, a common problem in drug development when using unreliable proxies for binding affinity.

Poster
History-Aware Conformal Prediction Sets for Censored Time-to-Event Outcomes

Yuyao Wang, Alexander W Levis, Shu Yang, Larry Han

Expand abstract

Existing conformal prediction methods for time-to-event outcomes leverage only baseline covariates, producing prediction intervals that are insufficiently informative to facilitate decision making. We propose History-Aware Prediction Sets (HAPS), a conformal framework that constructs prediction sets for individual event times using covariate histories observed up to a decision time, targeting coverage among individuals who have survived to this time. HAPS handles right censoring adjusted for time-varying confounders via inverse probability of censoring weighting. When the censoring weights are consistently estimated, it achieves PAAC (probably asymptotically approximately correct) coverage among survivors. We further propose two doubly robust extensions of HAPS to weaken reliance on consistent estimation of the censoring distribution. In simulations, HAPS and its extensions reduce median prediction interval length by up to 75\% relative to baseline comparators while maintaining close to nominal coverage. On two public benchmark data sets, HAPS reduces the median interval length by up to 60\% for predictions at year 5, compared to the baseline comparators.

Poster
HireWatch: Evaluating LLM Compliance with U.S. Employment Law Under Contextual Pressure

Huy Nghiem, Maria Kostylew, Noam Kolt

Expand abstract

The deployment of large language models (LLMs) in hiring pipelines is growing, yet existing evaluations generally overlook an important aspect of this development: the compliance of LLMs with applicable legal rules. To fill this gap, we introduce HireWatch, a benchmark that assesses whether LLM hiring assistants comply with U.S. employment discrimination law (e.g., the Civil Rights Act of 1964). We construct a 200-item bank consisting of interview questions that are either lawful or unlawful under U.S. federal statute, regulation, and enforcement guidance. We evaluate the propensity of 18 proprietary and open-source LLMs to select lawful vs. unlawful questions from this bank across ~100,000 trials under a novel seven-tier contextual escalation ladder, crossing occupational details with organizational pressure. Our findings reveal that LLMs' compliance with employment discrimination law is brittle. While LLMs ordinarily comply with legal rules under different job contexts, they readily violate the law when subject to institutional pressure (e.g., an unlawful manager instruction or company policy). Meanwhile, although the inclusion of explicit legal reminders in prompts can improve legal compliance, it does not altogether eliminate violations of employment discrimination law.

Poster
Hint Tuning: Less Data Makes Better Reasoners

Siqi Fan, Minghao Li, Xiaoqian Ma, Xiusheng Huang, Zhuo Chen, Bowen Qin, ZhangLiujie, Shuo Shang, Weihang Chen

Expand abstract

Large reasoning models achieve high accuracy through extended chain-of-thought but generate 5--8 more tokens than necessary, applying verbose reasoning uniformly regardless of problem difficulty. We propose Hint Tuning, a data-efficient approach that teaches models to calibrate reasoning depth. Our key insight: the corresponding instruct model serves as an ideal difficulty probe. By testing what the instruct model can solve with varying guidance, we automatically construct training data across three states: No-Hint (direct answer), Sparse-Hint (minimal prefix), and Full-Hint (complete reasoning). This converts the abstract challenge of difficulty labeling into a measurable consistency check between the instruct and reasoning models. With only 1K self-annotated samples, Hint Tuning achieves 24--66% token reduction (31.5% average) across mainstream reasoning models (Qwen3-Thinking, DeepSeek-R1-Distill) at multiple scales (4B--32B) while maintaining competitive accuracy on five benchmarks. Unlike methods requiring massive distillation datasets or expensive RL, we achieve superior efficiency through simple alignment with the instruct model's capabilities.

Poster
HilbertGen-3D:Hilbert--Multifractal Conditioning for Topology-Aware 3D Generation

Ziyi Zhang, Shixuan Li, Yorie Nakahira, Paul Bogdan, Guannan Qu

Expand abstract

High-fidelity 3D generative modeling underpins game asset creation, medical imaging, and VR/AR environmental design, where geometric structure directly influences perception and decision-making. Diffusion models represent a powerful paradigm for point cloud generation, yet existing approaches rely heavily on local voxel or patch features and often miss global structural organization. Recent topology-aware methods show that global conditioning can improve generation quality, yet persistent-homology pipelines add substantial preprocessing complexity. We introduce HilbertGen-3D, a topology-aware conditional diffusion framework based on Hilbert–Multifractal Conditioning (HMC). HMC computes multi-scale voxel measures, serializes them via a 3D Hilbert curve to expose structured spatial dependencies, and estimates the generalized fractal dimensions from scale-dependent partition functions. We encode these signals as global embeddings and multi-scale tokens, and inject them into a diffusion denoiser through an efficient bottleneck fusion module. On ShapeNet dataset, HilbertGen-3D reduces the 1-Nearest Neighbor Accuracy by up to 8.74 points and improves coverage by up to +21.45 over prior topology-aware methods. This shows that combining local density cues with global Hilbert–multifractal structure provides a simple, efficient, and effective recipe for high-fidelity and diverse 3D generation.

Poster
Higher-Order Action Supervision Makes A Strong Policy Class

Peng Cheng, Yunxian Hou, Zhi Zhou, Qian Zhang, Chang Huang, Xianyuan Zhan

Expand abstract

Modern data-driven decision-making methods, such as imitation learning (IL) and reinforcement learning (RL), have achieved great success in solving many complex tasks. However, these methods often suffer from serious control instability and robustness issues when applied in real-world applications such as robotics and autonomous driving, posing notable challenges for their practical deployment. We argue that this instability issue stems largely from their limitations on solely supervising and optimizing zeroth-order actions (i.e., the action labels), failing to account for higher-order action dynamics and temporal consistency. In this paper, we show that simultaneously supervising both zeroth- and first-order actions can dramatically enhance policies' performance and control robustness. To achieve this, we introduce a novel and elegant loss scheme supported by formal theoretical guarantees that can equip any off-the-shelf policy model (e.g., deterministic, stochastic, or flow policies) with the capability for higher-order action supervision, without requiring any structural modifications. Moreover, our proposed method can serve as a lightweight plug-and-play module that seamlessly integrates with a broad spectrum of existing offline RL frameworks. Extensive evaluations on OGBench and D4RL demonstrate that our approach yields substantial performance and robustness improvements across a wide range of continuous control environments. Notably, our method can also enhance policies' out-of-distribution (OOD) generalization capability in the challenging low-data regime, making it an ideal tool in tackling many real-world control problems.

Poster
High-dimensional Limit of SGD for Diagonal Linear Networks

Begoña García Malaxechebarría, Courtney Paquette, Maryam Fazel, Dmitriy Drusvyatskiy

Expand abstract

Understanding the behavior of stochastic gradient methods is a central problem in modern machine learning. Recent work has highlighted diagonal linear networks as a simplified yet expressive setting for analyzing the optimization and generalization properties of neural models. In this work, we show that in the high-dimensional regime, stochastic gradient descent on diagonal linear networks is well-approximated by continuous dynamics governed by a stochastic differential equation (SDE), which explicitly decouples the drift from the gradient noise. We further derive a deterministic partial differential equation whose solution propagates the relevant state of the iterates and characterizes the time evolution of a broad class of observable statistics, including the risk, curvature, and other metrics for optimality. To the best of our knowledge, this is the first deterministic high-dimensional description of SGD for a genuinely nonlinearly parametrized model. Finally, we show that, under a suitable parametrization, the stochastic dynamics are globally well posed and converge exponentially fast to zero risk with high probability, yielding a fully explicit non-asymptotic description of their long-time behavior. Numerical simulations corroborate our theoretical findings.

Poster