AISTATS 2026 论文列表

In this work, we investigate the generalization properties of random feature methods. Our analysis extends prior results for Tikhonov regularization to a broad class of spectral regularization techniques and further generalizes the setting to operator-valued kernels. This unified framework enables, for the first time, a rigorous theoretical analysis of neural operators and neural networks through the lens of the Neural Tangent Kernel (NTK). In particular, it allows us to establish optimal learning rates and provides a good understanding of how many neurons are required to achieve a given accuracy. Furthermore, we establish minimax rates in the well-specified case and also in the misspecified case, where the target is not contained in the reproducing kernel Hilbert space. These results sharpen and complete earlier findings for specific kernel algorithms.

Meta-learning is a popular approach for learning new tasks with limited data by leveraging the commonalities among different tasks. However, meta-learned models can perform poorly when context data is too limited, or when data is drawn from an out-of-distribution (OoD) task. Especially in safety-critical settings, this necessitates an uncertainty-aware approach to meta-learning. In addition, the often multimodal nature of task distributions can pose unique challenges to meta-learning methods. To this end, we present LUMA, a meta-learning method for regression that (1) makes probabilistic predictions on in-distribution tasks efficiently, (2) is capable of detecting OoD context data, and (3) handles heterogeneous, multimodal task distributions effectively. The strength of our framework lies in its solid theoretical basis, enabling analytically tractable Bayesian inference on a linearized model for principled uncertainty estimation and robust generalization. We achieve this by adopting a probabilistic perspective and learning a parametric, tunable task distribution via Bayesian inference on a linearized neural network, leveraging Gaussian process theory. Moreover, we make our approach computationally tractable by leveraging a low-rank prior covariance learning scheme based on the Fisher Information Matrix. Our numerical analysis demonstrates that LUMA quickly adapts to new tasks and remains accurate even in low-data regimes, it effectively detects OoD tasks, and that both of these properties continue to hold for multimodal task distributions.

We study sparse linear bandits in both adversarial and stochastic settings. While existing literature has extensively explored sparse linear bandits in the stochastic regime, the adversarial setting, particularly for general l_p-ball action sets (p>1), remains poorly understood. Our work addresses this gap by showing that the curse of dimensionality in adversarial linear bandits can be broken under a natural fixed sparsity support assumption. Specifically, we design algorithms for the l_∞- and l_2-balls that integrate sparsity support identification with the OSMD algorithm, achieving regret bounds O(s√(T)log T ) and O(√(sT)log T ), respectively. These results nearly match the optimal results when the sparsity support is known, and significantly improve upon the O(d√(T)) regret of algorithms ignoring sparsity. Furthermore, in the stochastic setting, we show how the geometry of the l_p-ball action set influences both exploration and regret. Our work highlights fundamental contrasts between adversarial and stochastic regimes, and establishes the first regret guarantees for sparse adversarial linear bandits beyond the l_1-ball action set.

We introduce GKK+, a new design for multi-arm randomized controlled trials. Standard Bernoulli randomization is robust but often yields poor covariate balance, while existing restricted-randomness designs mainly address two-arm settings. GKK+ extends the Karmarkar–Karp (KK) differencing method to multiple arms. When covariates are smooth and well-behaved, GKK+ achieves an exponentially better covariate balance than the standard Bernoulli design while preserving sufficient randomness. GKK+ improves efficiency in estimating treatment effects and supports standard asymptotic inference. Simulations on synthetic and real datasets demonstrate improved balance and lower estimator variance compared to existing methods.

Causal discovery algorithms often perform poorly with limited samples. While integrating expert knowledge (including from LLMs) as constraints promises to improve performance, guarantees for existing methods require perfect predictions or uncertainty estimates, making them unreliable for practical use. We propose the Guess2Graph (G2G) framework, which uses expert guesses to guide the sequence of statistical tests rather than replacing them. This maintains statistical consistency while enabling performance improvements. We develop two instantiations of G2G: PC-Guess, which augments the PC algorithm, and gPC-Guess, a learning-augmented variant designed to better leverage high-quality expert input. Theoretically, both preserve correctness regardless of expert error, with gPC-Guess provably outperforming its non-augmented counterpart in finite samples when experts are "better than random". Empirically, both show monotonic improvement with expert accuracy, with gPC-Guess achieving significantly stronger gains.

Offline reinforcement learning is important in many settings with available observational data but the inability to deploy new policies online due to safety, cost, and other concerns. Many recent advances in causal inference and machine learning target estimation of “causal contrast" functions such as CATE, which is sufficient for optimizing decisions and can adapt to potentially smoother structure. We develop a dynamic generalization of the R-learner (Nie et al 2021, Lewis and Syrgkanis 2021) for estimating and optimizing the difference of Q^π-functions, Q^π(s,a)-Q^π(s,a_0), for potential discrete-valued actions a,a_0, which can be used to optimize multiple-valued actions without loss of generality. We leverage orthogonal estimation to improve convergence rates, even if Q and behavior policy (so-called nuisance functions) converge at slower rates and prove consistency of policy optimization under a margin condition. The method can leverage black-box estimators of the Q-function and behavior policy to target estimation of a more structured Q-function contrast, and comprises of simple squared-loss minimization.

Recovering the random graph model from an observed collection of networks is known to present significant challenges in the setting, where the networks do not share a common node set and have different sizes. More specifically, the goal is the estimation of the graphon function that parametrizes the nonparametric exchangeable random graph model. Existing methods typically suffer from either limited accuracy or high computational complexity. We introduce a new histogram-based estimator with low algorithmic complexity that achieves high accuracy by jointly aligning the nodes of all graphs, in contrast to most conventional methods that order nodes graph by graph. Consistency results of the proposed graphon estimator are established. A numerical study shows that the proposed estimator outperforms existing methods in terms of accuracy, especially when the dataset comprises only small and variable-size networks. Moreover, the computing time of the new method is considerably shorter than that of other consistent methodologies. Additionally, when applied to a graph neural network classification task, the proposed estimator enables more effective data augmentation, yielding improved performance across diverse real-world datasets.

We study algorithms for construction of composable coresets for the task of Determinant Maximization under partition constraint. Given a point set V ⊂ℝ^d that is partitioned into s groups V_1,⋯, V_s, and integers k_1,…,k_s, where k=∑_i k_i, the goal is to pick k_i points from group V_i such that the overall determinant of the picked k points is maximized. Determinant Maximization and its constrained variants have gained a lot of interest for modeling diversity, and have found applications in the context of data summarization. When the cardinality k of the selected set is greater than the dimension d, we show a peeling algorithm that gives us a composable coreset of size kd with a provably optimal approximation factor of d^O(d). When k≤ d, we show a simple coreset construction with optimal size and approximation factor. As a further application of our technique, we get a composable coreset for determinant maximization under the more general laminar matroid constraints, and a composable coreset for unconstrained determinant maximization in a previously unresolved regime. Our results generalize to all strongly Rayleigh distributions and to several other experimental design problems. As an application, we improve the runtime of the practical local-search based algorithm of [Anari-Vuong–COLT’22] for determinantal maximization under partition constraint from O(n^2^sk^2^s) to O(n k^2^s), making it only linear on the number of points n.

In off-policy evaluation (OPE) for partially observable Markov decision processes (POMDPs), an agent must infer hidden states from past observations, which exacerbates both the curse of horizon and the curse of memory in existing OPE methods. This paper introduces a novel covering analysis framework that exploits the intrinsic metric structure of the belief space (distributions over latent states) to relax traditional coverage assumptions. By focusing on the policies with stability property, we derive error bounds that mitigate exponential blow-ups in horizon and memory length. Our unified analysis technique applies to a broad class of OPE algorithms, yielding concrete error bounds and coverage requirements expressed in terms of belief space metrics rather than raw history coverage. We illustrate the improved sample efficiency of this framework via case studies: the double sampling Bellman error minimization algorithm, and the memory-based future-dependent value functions (FDVF). In both cases, our coverage definition based on the belief-space metric yields tighter bounds.

Conformal prediction quantifies the uncertainty of machine learning models by augmenting point predictions with valid prediction sets. For complex scenarios involving multiple trials, models, or data sources, conformal prediction sets can be aggregated to create a prediction set that captures the overall uncertainty, often improving precision. However, aggregating multiple prediction sets with individual 1-α coverage inevitably weakens the overall guarantee, typically resulting in 1-2α worst-case coverage. In this work, we propose a framework for the weighted aggregation of prediction sets, where weights are assigned to each prediction set based on their contribution. Our framework offers flexible control over how the sets are aggregated, achieving tighter coverage bounds that interpolate between the 1-2α guarantee of the combined models and the 1-α guarantee of an individual model depending on the distribution of weights. Importantly, our framework generalizes to data-dependent weights, as we derive a procedure for weighted aggregation that maintains finite-sample validity even when the weights depend on the data. This extension makes our framework broadly applicable to settings where weights are learned, such as mixture-of-experts (MoE), and we demonstrate through experiments in the MoE setting that our methods achieve adaptive coverage.

We theoretically investigate the phenomena of generalization and memorization in diffusion models. Empirical studies suggest that these phenomena are influenced by model complexity and the size of the training dataset. In our experiments, we further observe that the number of noise samples per data sample (m) used during Denoising Score Matching (DSM) plays a significant and non-trivial role. We capture these behaviors and shed insights into their mechanisms by deriving asymptotically precise expressions for test and train errors of DSM under a simple theoretical setting. The score function is parameterized by random features neural networks, with the target distribution being d-dimensional Gaussian. We operate in a regime where the dimension d, number of data samples n, and number of features p tend to infinity while keeping the ratios ψ_n=n/d and ψ_p=p/d fixed. By characterizing the test and train errors, we identify regimes of generalization and memorization as a function of ψ_n,ψ_p, and m. Our theoretical findings are consistent with the empirical observations.

Influence Functions are a standard tool for attributing predictions to training data in a principled manner and are widely used in applications such as data valuation and fairness. In this work, we present realistic incentives to manipulate influence-based attributions and investigate whether these attributions can be systematically altered by an adversary. We show that small systemic perturbations to models can indeed alter influence-based attributions as desired. We work on logistic regression models trained on ResNet feature embeddings and standard tabular fairness datasets and provide efficient attacks with backward-friendly implementations. Our work raises questions on the reliability of influence-based attributions in adversarial circumstances.

The bilevel variational inequality (BVI) problem is a broad framework covering optimal equilibrium selection and equilibrium problems with equilibrium constraints (EPECs). We propose Regularized Operator Extrapolation (R-OpEx), a single-loop first-order algorithm for smooth and nonsmooth BVIs with stochastic monotone operators. R-OpEx combines Tikhonov regularization with operator extrapolation, requires only one operator evaluation per iteration, and tracks a single sequence of iterates. We show that R-OpEx obtains an ϵ-solution in 𝒪(ϵ^-4) iterations for nonsmooth stochastic BVIs. If the inner operator is smooth and stochastic, we show an improved complexity of 𝒪(ϵ^-2) for the outer level operator while maintaining 𝒪(ϵ^-4) complexity for the inner level operator. For a smooth deterministic inner level operator, the overall complexity reduces to 𝒪(ϵ^-2). Finally, we improve the complexities substantially when the outer level is strongly monotone. To our knowledge, this is the first work to establish such guarantees for nonsmooth stochastic BVIs. We validate our results through numerical studies.

In this paper, we establish the non-asymptotic validity of the multiplier bootstrap procedure for constructing the confidence sets using the Stochastic Gradient Descent (SGD) algorithm. Under appropriate regularity conditions, our approach avoids the need to approximate the limiting covariance of Polyak-Ruppert SGD iterates, which allows us to derive approximation rates in convex distance of order up to 1/√(n). Notably, this rate can be faster than the one that can be proven in the Polyak-Juditsky central limit theorem. To our knowledge, this provides the first fully non-asymptotic bound on the accuracy of bootstrap approximations in SGD algorithms. Our analysis builds on the Gaussian approximation results for nonlinear statistics of independent random variables.

Generative models guided by text prompts are widely evaluated for fidelity and prompt alignment, yet their ability to produce diverse outputs remains underexplored. Existing diversity metrics such as Vendi and RKE, which are based on the von Neumann and Rényi entropies of kernel matrices, were developed for unconditional models and cannot distinguish prompt-induced from model-induced variability. We address this gap by introducing Conditional-Vendi and Conditional-RKE, diversity measures derived from the conditional entropy of positive semidefinite matrices. These scores isolate model-induced diversity in prompt-guided generation, with Conditional-RKE enjoying an O(1/√(n)) convergence rate. For Conditional-Vendi, we introduce a truncated-spectrum approximation that yields scalable and consistent estimates. Experiments on text-to-image, image-captioning, and language generation tasks demonstrate that the conditional scores recover ground-truth diversity orderings and can also guide diffusion models toward more diverse generations. The codebase is available at <https://github.com/mjalali/conditional-vendi.>

Clustering tasks such as k-means and k-median are central in unsupervised learning, and streaming algorithms for these tasks are widely used to handle large or evolving datasets. When applied in sensitive domains, however, such algorithms must also provide rigorous privacy guarantees. In this work, we provide the first differentially private (DP) algorithms for k-means and k-median clustering of d-dimensional Euclidean data points over a stream of length at most T, using space that is sublinear in T, in the continual release setting where the algorithm is required to output a clustering at every timestep. We achieve (1) an O(1)-multiplicative approximation with O(k^1.5 poly(d, log T)) space and poly(k,d,log T) additive error, or (2) a (1+γ)-multiplicative approximation with O_γ(poly(k, 2^O_γ(d), log T)) space for any γ>0, with additive error poly(k, 2^O_γ(d), log T). Our main technical contribution is a DP clustering framework for data streams that only requires an offline DP coreset or clustering algorithm as a blackbox.

We study transformers’ in-context learning of variable-length Markov chains (VOMCs), focusing on the finite-sample accuracy as the number of in-context examples increases. Compared to fixed-order Markov chains (FOMCs), learning VOMCs is substantially more challenging due to the additional structural learning component. The problem is naturally suited to a Bayesian formulation, where the context-tree weighting (CTW) algorithm, originally developed in the information theory community for universal data compression, provides an optimal solution. Empirically, we find that single-layer transformers fail to learn VOMCs in context, whereas transformers with two or more layers can succeed, with additional layers yielding modest but noticeable improvements. In contrast to prior results on FOMCs, attention-only networks appear insufficient for VOMCs. To explain these findings, we provide explicit transformer constructions: one with D+2 layers that can exactly implement CTW for VOMCs of maximum order D, and a simplified two-layer construction that uses partial information for approximate blending, shedding light on why two-layer transformers can perform well.

Post-hoc recalibration methods are widely used to ensure that classifiers provide faithful probability estimates. We argue that parametric recalibration functions based on logistic regression can be motivated from a simple theoretical setting for both binary and multi-class classification. This insight motivates the use of more expressive calibration methods beyond standard temperature scaling. For multi-class calibration however, a key challenge lies in the increasing number of parameters introduced by more complex models, often coupled with limited calibration data, which can lead to overfitting. Through extensive experiments, we demonstrate that the resulting bias-variance tradeoff can be effectively managed by structured regularization, robust preprocessing and efficient optimization. The resulting methods lead to substantial gains over existing logistic-based calibration techniques. We provide efficient and easy-to-use open-source implementations of our methods, making them an attractive alternative to common temperature, vector, and matrix scaling implementations.

Regret in stochastic multi-armed bandits traditionally measures the difference between the highest reward and either the arithmetic mean of accumulated rewards or the final reward. These conventional metrics often fail to address fairness among agents receiving rewards, particularly in settings where rewards are distributed across a population, such as patients in clinical trials. To address this, a recent body of work has introduced Nash regret, which evaluates performance via the geometric mean of accumulated rewards, aligning with the Nash social welfare function known for satisfying fairness axioms. To minimize Nash regret, existing approaches require specialized algorithm designs and strong assumptions, such as multiplicative concentration inequalities and bounded, non-negative rewards, making them unsuitable for even Gaussian reward distributions. We demonstrate that an initial uniform exploration phase followed by a standard Upper Confidence Bound (UCB) algorithm achieves near-optimal Nash regret, while relying only on additive Hoeffding bounds, and naturally extending to sub-Gaussian rewards. Furthermore, we generalize the algorithm to a broad class of fairness metrics called the p-mean regret, proving (nearly) optimal regret bounds uniformly across all p values. This is in contrast to prior work, which made extremely restrictive assumptions on the bandit instances and even then achieved suboptimal regret bounds. Numerical simulations validate our method’s practical efficacy, broadening the accessibility of fairness in bandit algorithms.