ICML 2026 Papers

A key challenge in probabilistic regression is ensuring that predictive distributions accurately reflect true empirical uncertainty. Minimizing overall prediction error often encourages models to prioritize informativeness over calibration, producing narrow but overconfident predictions. However, in safety-critical settings, trustworthy uncertainty estimates are often more valuable than narrow intervals. Realizing the problem, several recent works have focused on post-hoc corrections; however, existing methods either rely on weak notions of calibration (such as PIT uniformity) or impose restrictive parametric assumptions on the nature of the error. To address these limitations, we propose a novel nonparametric re-calibration algorithm based on conditional kernel mean embeddings, capable of correcting calibration error without restrictive modeling assumptions. For efficient inference with real-valued targets, we introduce a novel characteristic kernel over distributions that can be evaluated in $\mathcal{O}(n \log n)$ time for empirical distributions of size $n$. We demonstrate that our method consistently outperforms prior re-calibration approaches across a diverse set of regression benchmarks and model classes.

Empirical successes show that regularization improves the stability and efficiency of reinforcement learning (RL), with applications in robotics and post-training of large language models. Yet, theoretical analyses of regularized Markov decision processes (MDPs) have mostly been confined to the standard RL setting. In this work, we investigate regularized MDPs through the lens of robust RL. We introduce a doubly regularized MDP framework that combines policy and dynamics regularizations, enabling robust policy learning while naturally accommodating continuous action spaces. Within this framework, we develop an optimism-based online algorithm and provide the first finite-sample regret guarantees in both tabular and linear settings. Our results show that algorithms for doubly regularized MDPs are as sample-efficient as well-studied robust MDP algorithms, while additionally benefiting from the flexibility of soft policies. We further design practical algorithmic variants for both settings and demonstrate empirically that our approach efficiently and effectively handles function approximation and exploration in large state-action spaces, achieving robust performances.

Logistic regression is widely used in applications; however, when the dimension scales with the sample size, theory reveals that the asymptotic behavior of common M-estimators depends on nonzero bias and variance factors, which are functions of the signal strength. To leverage the theory to design valid statistical methodologies, it is essential to obtain accurate estimates of the signal strength. In this work, we utilize a data-splitting strategy to efficiently estimate the signal strength. To alleviate issues caused by separable data, we analyze the exact asymptotics of an M-estimator with a data-driven, non-decomposable regularizer that adapts to the true covariance structure. We justify the validity of our method through both theoretical analysis and numerical experiments.

Neural Controlled Differential Equations (Neural CDEs) provide a powerful continuous-time framework for sequence modeling, yet the roughness of the driving control path often restricts their efficiency. Standard splines introduce high-frequency variations that force adaptive solvers to take excessively small steps, driving up the Number of Function Evaluations (NFE). We propose a novel approach to Neural CDE path construction that replaces exact interpolation with Kernel and Gaussian Process (GP) smoothing, enabling explicit control over trajectory regularity. To recover details lost during smoothing, we propose an attention-based Multi-View CDE (MV-CDE) and its convolutional extension (MVC-CDE), which employ learnable queries to inform path reconstruction. This framework allows the model to distribute representational capacity across multiple trajectories, each capturing distinct temporal patterns. Empirical results demonstrate that our method, MVC-CDE with GP, achieves state-of-the-art accuracy while significantly reducing NFEs and total inference time compared to spline-based baselines.

Bayesian Last Layers (BLLs) provide a convenient and computationally efficient way to estimate uncertainty in neural networks. However, they underestimate epistemic uncertainty because they apply a Bayesian treatment only to the final layer, ignoring uncertainty induced by earlier layers. We propose a method that improves BLLs by leveraging a projection of Neural Tangent Kernel (NTK) features onto the space spanned by the last-layer features. This enables posterior inference that accounts for variability of the full network while retaining the low computational cost of inference of a standard BLL. We show that our method yields posterior variances that are provably greater or equal to those of a standard BLL, correcting its tendency to underestimate epistemic uncertainty. To further reduce computational cost, we introduce a uniform subsampling scheme for estimating the projection matrix and for posterior inference. We derive approximation bounds for both types of subsampling. Empirical evaluations on UCI regression, contextual bandits, image classification, and out-of-distribution detection tasks in image and tabular datasets, demonstrate improved calibration and uncertainty estimates compared to standard BLLs and competitive baselines, while reducing computational cost.

Neural Processes (NPs), and specifically Transformer Neural Processes (TNPs), have demonstrated remarkable performance across tasks ranging from spatiotemporal forecasting to tabular data modelling. However, many of these applications are inherently sequential, involving continuous data streams such as real-time sensor readings or database updates. In such settings, models should support cheap, incremental updates rather than recomputing internal representations from scratch for every new observation—a capability existing TNP variants lack. Drawing inspiration from Large Language Models, we introduce the Incremental TNP ($\texttt{incTNP}$). By leveraging causal masking, Key-Value (KV) caching, and a data-efficient autoregressive training strategy, $\texttt{incTNP}$ matches the predictive performance of standard TNPs while reducing the computational cost of updates from quadratic to linear time complexity. We empirically evaluate our model on a range of synthetic and real-world tasks, including tabular regression and temperature prediction. Our results show that, surprisingly, $\texttt{incTNP}$ delivers performance comparable to—or better than—non-causal TNPs while unlocking orders-of-magnitude speedups for sequential inference. Finally, we assess the consistency of the model's updates---by adapting a metric of "implicit Bayesianness", we show that under a one-at-a-time streaming protocol, $\texttt{incTNP}$ retains a prediction rule as implicitly Bayesian as standard non-causal TNPs, demonstrating that $\texttt{incTNP}$ achieves the computational benefits of causal masking without sacrificing the consistency required for streaming inference.

We present the first comprehensive Lean 4 formalization of statistical learning theory (SLT) grounded in empirical process theory. Our end-to-end formal infrastructure implement the missing contents in latest Lean library, including a complete development of Gaussian Lipschitz concentration, Dudley’s entropy integral theorem for sub-Gaussian processes, and an application to least-squares (sparse) regression with a sharp rate. The project was carried out using a human-AI collaborative workflow, in which humans design proof strategies and AI agents execute tactical proof construction, leading to the human-verified Lean 4 toolbox for SLT. Beyond implementation, the formalization process exposes and resolves implicit assumptions and missing details in standard SLT textbooks, enforcing a granular, line-by-line understanding of the theory. This work establishes a reusable formal foundation and opens the door for future developments in machine learning theory. The code is provided in https://github.com/YuanheZ/lean-stat-learning-theory.

Many geometric statistics and manifold learning pipelines routinely produce observations---such as tangent vectors or local frames---whose natural home is a varying family of fibers attached to different points of a base manifold, rather than a single shared vector space. Forming empirical averages requires transporting these observations to a common reference fiber, introducing curvature- and holonomy-driven effects absent from classical concentration theory. We develop a non-asymptotic concentration theory for such transported empirical means, deriving finite-sample, dimension-free Hoeffding- and Bernstein-type bounds via sharp Hilbert-space inequalities. When shortest paths to the reference point are non-unique, transport becomes path-dependent and introduces a deterministic holonomy bias; we isolate and quantify this bias through bundle curvature and loop geometry, with sharp closed-form formulas for the tangent bundle of a round sphere. The resulting bias--variance decomposition separates the stochastic fluctuation decaying at the classical $n^{-1/2}$ rate in sample size $n$, from a curvature-driven error floor that no amount of additional data can eliminate; minimax lower bounds confirm both terms are unavoidable. We further establish a robust median-of-means estimator achieving optimal rates under heavy tails, and a central limit theorem in the reference fiber. Controlled experiments on the sphere validate all theoretical predictions.

Evaluating generative AI models is increasingly resource-intensive due to slow inference, expensive raters, and a rapidly growing landscape of models and benchmarks. We propose ProEval, a proactive evaluation framework that leverages transfer learning to efficiently estimate performance and identify failure cases. ProEval employs pre-trained Gaussian Processes (GPs) as surrogates for the performance score function, mapping model inputs to metrics such as the severity of errors or safety violations. By framing performance estimation as Bayesian quadrature (BQ) and failure discovery as superlevel set sampling, we develop uncertainty-aware decision strategies that actively select or synthesize highly informative inputs for testing. Theoretically, we prove that our pre-trained GP-based BQ estimator is unbiased and bounded. Empirically, extensive experiments on reasoning, safety alignment, and classification benchmarks demonstrate that ProEval is significantly more efficient than competitive baselines. It requires 8–65x fewer samples to achieve estimates within $\pm1\%$ of the ground truth, while simultaneously revealing more diverse failure cases under a stricter evaluation budget. Our open-sourced code and data can be found at https://github.com/google-deepmind/proeval.

Stability is a property of fundamental importance in real-world systems. Although it has been widely studied and well understood in *control theory* (CT) for deterministic systems, it is largely overlooked in stochastic systems such as *Markov decision processes* (MDPs). In this paper, we aim to translate the steady-state control problem, well established in CT, where the goal is to synthesize a controller with prescribed asymptotic stability properties, into the MDP framework. To this end, we propose the novel *fast-mixing steady-state* (FMSS) problem. Given an ergodic MDP and a target steady-state distribution, the objective is to synthesize a Markovian policy that induces this distribution with the fastest possible convergence rate. Addressing this problem requires controlling the spectral properties of the induced *Markov chain* (MC) transition matrix, which generally leads to non-convex programs. Thus, we derive a tractable surrogate objective that leads to a convex program, whose properties we study in terms of approximation quality, feasibility, and computational complexity. We then move to the learning setting and propose an "offline" sample-based algorithm for FMSS (FMSS-SV), designed for tabular MDPs, in which the environment’s transition model is estimated from data. We quantify the impact of transition model estimation errors on both the objective value and the learned policy, and provide a finite-sample complexity analysis.

We study the task of learning the structure of a $d$-sparse Gaussian graphical model on $n$ variables from a single trajectory of Glauber dynamics. Beyond algorithmic considerations, many applications present temporally correlated observations rather than i.i.d. samples. Moreover, in the classical i.i.d. setting, polynomial-time structure learning from a sublinear in $n$ number of samples is suspected to be computationally hard without additional assumptions on the precision matrix. Motivated in part by this, we design the first polynomial-time algorithm that recovers the conditional-independence graph from a single Glauber trajectory, with a trajectory-length guarantee that does not depend on the mixing time. Technically, our algorithm has three components. First, we estimate the conditional variances and rescale the trajectory to reduce to the unit-diagonal case, without changing the underlying graph. Second, we design a local edge test that extracts adjacency information from short update windows by isolating pairwise influence. Third, we aggregate these local statistics using a robust median-based estimator, and prove accuracy despite contamination and temporal dependence arising from a single trajectory.

Standard Bayesian Optimization (BO) assumes uniform smoothness across the search space—an assumption violated in multi-regime problems such as molecular conformation search through distinct energy basins or drug discovery across heterogeneous molecular scaffolds. A single GP either oversmooths sharp transitions or hallucinates noise in smooth regions, yielding miscalibrated uncertainty. We propose RAMBO, a Dirichlet Process Mixture of Gaussian Processes that automatically discovers latent regimes during optimization, each modeled by an independent GP with locally-optimized hyperparameters. We derive collapsed Gibbs sampling that analytically marginalizes latent functions for efficient inference, and introduce adaptive concentration parameter scheduling for coarse-to-fine regime discovery. Our acquisition functions decompose uncertainty into intra-regime and inter-regime components. Experiments on synthetic benchmarks and real-world applications—including molecular conformer optimization, virtual screening for drug discovery, and fusion reactor design—demonstrate consistent improvements over state-of-the-art baselines on multi-regime objectives. Code is available at https://github.com/AnthonyZhangYan/RAMBO.

Graph Masked Auto-Encoder (GMAE) has emerged as a prevalent self-supervised paradigm, showing superior performance in graph learning. However, existing methods mainly emphasize reconstruction outcomes and give limited specification to how neighborhood context is used for reconstruction. Our experimental investigation shows that, even when reconstruction outcomes are similar, the ways of using neighborhood context differ substantially, resulting in performance shifts. To address this issue, we propose Consistency- and Coverage-aware Graph Masked Auto-Encoder (C2-GMAE), which encourages more consistent use of neighborhood context and promotes broader training coverage in the graph. Specifically, C2-GMAE leverages positional encoding as an observable structural reference, introduces density-partitioned masking to improve coverage across regions, and amplifies heterophilic edges to reduce the attenuation of discriminative relational information during reconstruction. Extensive experiments on multiple benchmarks demonstrate that C2-GMAE improves downstream performance against GMAE baselines.

Bayesian additive spanning tree (BAST) is an useful tool for interpretable, non-parametric regression on complex constrained domains. It improves upon the performance of Bayesian additive regression trees (BART) by replacing axis-aligned splits through binary tree components by cuts on a spanning tree components, enabling the formation of contiguous splits that respect the underlying complex structure. While BAST is effective for learning on constrained spaces, it still relies on hard partitions, albeit on spanning trees, which limits its ability to represent smoothly varying functions on constrained domains. We propose Probabilistically-routed Bayesian additive spanning trees (PR-BAST), a principled relaxation that replaces hard cuts on spanning tree components with probabilistic routing along spanning tree components. PR-BAST represents the regression surface as an additive ensemble of such spanning tree-aligned smooth components. Conditional on a fixed spanning tree, each component in PR-BAST induces a Gaussian random field with a sparse, tree-structured precision matrix, enabling scalable posterior computation via sparse linear algebra. We theoretically establish that PR-BAST yields strictly faster posterior contraction rates compared to BAST under graph-smooth truth. Experiments on synthetic and real datasets demonstrate that PR-BAST consistently improves accuracy over BAST and other baselines, while retaining the interpretability of tree-based models.

Sampling from unnormalized densities is analogous to the generative modeling problem, but the target distribution is defined by a known energy function instead of data samples. Because evaluating the energy function is often costly, a primary challenge is to learn an efficient sampler. We introduce Flow Sampling, a framework built on diffusion models and flow matching for the data-free setting. Our training objective is conditioned on a noise sample and regresses onto a denoising diffusion drift constructed from the energy function. In contrast, diffusion models' objective is conditioned on a data sample and regresses onto a noising diffusion drift. We utilize the interpolant process to minimize the number of energy function evaluations during training, resulting in an efficient and scalable method for sampling unnormalized densities. Furthermore, our formulation naturally extends to Riemannian manifolds, enabling diffusion-based sampling in geometries beyond Euclidean space. We derive a closed-form formula for the conditional drift on constant curvature manifolds, including hyperspheres and hyperbolic spaces. We evaluate Flow Sampling on synthetic energy benchmarks, small peptides, large-scale amortized molecular conformer generation, and distributions supported on the sphere, demonstrating strong empirical performance.

Nonlinear conservation laws are at the heart of many of the most important dynamical systems in science and engineering. In practical applications, such systems are often subject to various sources of uncertainty, e.g. due to sparse or noisy measurements. Inferring physical quantities and fields of interest then becomes an ill-posed problem which both classical numerical methods and modern deep learning-based methods struggle to treat appropriately. Recent work has framed classical numerical methods as Bayesian inference under Gaussian process priors, resulting in a physics-aware treatment of uncertainties. Following this line of work, we develop a novel numerically conservative method for uncertainty-aware simulations of nonlinear conservation laws. We use recent sparse approximation techniques to scale up to large-scale forward and inverse problems. For forward simulation, we inherit the accuracy of classical solvers while providing structured uncertainty quantification. On inverse problems, we recover posteriors over nonparametric source fields in seconds --- outperforming neural baselines that take minutes to produce a less accurate point estimate.

Gaussian processes (GPs) provide a principled Bayesian framework for uncertainty estimation, but their computational complexity severely limits scalability to large datasets. We propose SIKA-GP, which accelerates GP inference using sparse inducing kernel approximations based on a dyadic ordered template basis, incurring only ${O}(\log M)$ complexity dependence on the number of inducing points. Our approach constructs compact and expressive kernel representations from sparsely activated bases, enabling efficient tensorized GPU computation and seamless integration with modern large-scale models. SIKA-GP can be naturally embedded into Bayesian neural networks (BNNs) with sparse activations, yielding significant speedups in both training and inference without sacrificing predictive performance. The method naturally extends to deep feature learning, addressing the scalability challenges introduced by deep architectures and high-dimensional feature representations. Empirical results on vision and transformer-based language benchmarks demonstrate that our approach consistently delivers fast and accurate GP models, providing a principled path toward scalable kernel learning.

Importance-Weighted Neural Processes (IWNPs) provide a principled framework for probabilistic meta-learning by using multi-particle latent representations to approximate the marginal log-likelihood of task data tightly. However, this work reveals that the standard optimization of IWNPs suffers from the Matthew effect in the latent space, where high-likelihood particles dominate gradient signals. The neglect of lower-likelihood regions leads to poor tail-risk generation and unstable fast adaptation. While robust objectives such as $\text{CVaR}_\alpha$ can mitigate these risks, they often entail a trade-off that degrades average-case performance. This work proposes \underline{O}rder-\underline{S}tatistics Aligned \underline{N}eural \underline{P}rocesses (OS-NPs) to achieve latent space robust optimization without sacrificing average result. Specifically, we stratify multiple inference particles into disjoint difficulty bins based on order statistics and derive the regularized worst-case optimization framework for OS-NPs. Our method aligns the reduction of stratified order-statistic losses in IWNPs and provides a computationally efficient pipeline to implement. Extensive experiments demonstrate that the OS-NP constitutes stable, reliable probabilistic meta-learning that significantly enhances tail-risk robustness while maintaining or even improving average performance.

We study a cooperative form of multi-agent reinforcement learning with state space dynamics and agent interaction controlled by an underlying graph. Each agent has a local state and action, the evolution of the local state depends only on the states and actions in the $1$-hop neighborhood defined by the graph. Structured dynamics of this type arise in various applications, including network resource allocation, co-operative games, epidemic control, and wireless scheduling. The global state-action space scales exponentially in the number of agents, so that computing global optimal policies is intractable in the worst-case. We study conditions under which it is possible to approximate the optimal policies by a local policy for each agent that depends only on states associated with nodes within its $m$-hop neighborhood. By controlling the propagation of influences via a Dobrushin-type stability matrix, we establish that globally optimal policies can approximated by local policies with sub-optimality gap decaying exponentially in $m$.