AISTATS 2026 Papers

Functional bilevel optimization (FBO) provides a powerful framework for hierarchical learning in function spaces, yet current methods are limited to static offline settings and perform suboptimally in online, non-stationary scenarios. We propose SmoothFBO, the first algorithm for non-stationary FBO with both theoretical guarantees and practical scalability. SmoothFBO introduces a time-smoothed stochastic hypergradient estimator that reduces variance through a window parameter, enabling stable outer-loop updates with sublinear regret. Importantly, the classical parametric bilevel case is a special reduction of our framework, making SmoothFBO a natural extension to online, non-stationary settings. Empirically, SmoothFBO consistently outperforms existing FBO methods in non-stationary hyperparameter optimization and model-based reinforcement learning, demonstrating its practical effectiveness. Together, these results establish SmoothFBO as a general, theoretically grounded, and practically viable foundation for bilevel optimization in online, non-stationary scenarios.

Learning in games is a central topic in both learning theory and game theory, and learning dynamics based on online learning have made significant theoretical and practical advances in recent years. In particular, in two-player zero-sum games, it has been shown that using optimistic follow-the-regularized-leader (OFTRL) as the online learning algorithm allows us to upper bound the individual regret for each player by a constant independent of the time horizon. However, in realistic game scenarios, players are not always able to observe the outcomes of their interactions immediately. Motivated by this, very recently, the problem of learning from delayed feedback in games has been proposed, and it has been shown that by using a variant of OFTRL, one can achieve a social regret upper bound of Õ(D^2 + 1) for a fixed delay time D. This study investigates the optimal dependence on the delay parameter D in the setting of learning from delayed feedback in games. In particular, we show that a simple algorithm that runs D+1 independent copies of the standard OFTRL designed for the non-delayed setting achieves social and individual regret upper bounds of Õ(D + 1), thereby improving the existing bounds by a factor of D. Moreover, we provide a matching lower bound: for any learning dynamic, there exists a payoff matrix such that the regret of every player is at least Ω(D + 1).

The parity problem—deciding whether the number of ones in a binary vector is odd or even—remains challenging for standard neural networks due to linear inseparability and the need for global interactions. We propose TESLA, an activation defined as a learnable combination of sine and cosine terms, enabling explicit control over polynomial degree and selective amplification of high-order components. Theoretically, we show that constraining TESLA’s coefficients yields Lipschitz/Rademacher complexity bounds and shapes the training dynamics to emphasize higher-frequency structure. Empirically, on parity with input length n=32, TESLA attains strong generalization with 100K training samples (≈ 0.002

Accurate dose-response forecasting under sparse sampling is central to precision pharmacotherapy. We present the Amortized In-Context Mixed-Effect Transformer (AICMET) model, a transformer-based, latent-variable framework that unifies mechanistic compartmental priors with amortized, in-context Bayesian inference. AICMET is pre-trained on hundreds of thousands of synthetic pharmacokinetic trajectories with Ornstein-Uhlenbeck priors over the parameters of compartment models, endowing the model with strong inductive biases and enabling zero-shot adaptation to new compounds. At inference time, AICMET is conditioned on the collective context of previously profiled trial participants, generating calibrated posterior predictions for newly enrolled patients after a few early drug concentration measurements. This capability collapses traditional model development cycles from weeks to seconds, while preserving some degree of expert modelling. Experiments across public datasets show that AICMET attains state-of-the-art predictive accuracy, and faithfully quantifies inter-patient variability - outperforming both nonlinear mixed-effects baselines and recent neural ODE variants. Our code repository, pretrained model and tutorials are available online.

Generating samples from exponentially tilting a given distribution of random vectors when samples from the given distribution are available finds applications in fields such as finance and climate science and in the broad area of rare event simulation. In this article, we discuss the asymptotic efficiency of an estimator obtained by exponentially tilting the empirical distribution. We provide a sharp characterization of how much one can accurately tilt distributions given a certain number of samples. Our findings reveal a surprising dichotomy: While twisting unbounded distributions is a fundamentally hard task, for bounded distributions, one can accurately tilt by a large amount using much fewer samples.

In this paper, we study the finite-time behavior of the TD(0) temporal-difference method with linear function approximation (LFA). We consider on-policy independent and identically distributed (i.i.d.) samples, a constant learning step, and the Polyak-Juditsky averaging method. We establish a new convergence rate, for the Mean-Square Error (MSE) on the approximated function, that is (i) fast in the sense that it admits an optimal dependency in the number of iterations k (i.e., of order 1/k), (ii) is robust to ill-conditioning: it only depends on an initial error and model-independent constants and (iii) is sharp up to a multiplicative constant lower than 11. In particular, it does not depend on the smallest eigenvalue of the uncentered covariance matrix of the linear parametrization, unlike all pre-existing O(1/k) rates in the TD(0) literature. We also introduce PCTD(0), a variant of TD(0), which benefits from better convergence properties under an additional assumption of strong mixing on the Markov Chain.

We develop a unified statistical framework for softmax-gated Gaussian mixture of experts (SGMoE) that addresses three long-standing obstacles in parameter estimation and model selection: (i) non-identifiability of gating parameters up to common translations, (ii) intrinsic gate-expert interactions that induce coupled differential relations in the likelihood, and (iii) the tight numerator-denominator coupling in the softmax-induced conditional density. Our approach introduces Voronoi-type loss functions aligned with the gate-partition geometry and establishes finite-sample convergence rates for the maximum likelihood estimator (MLE). In over-specified models, we reveal a link between the MLE’s convergence rate and the solvability of an associated system of polynomial equations characterizing near-nonidentifiable directions. For model selection, we adapt dendrograms of mixing measures to SGMoE, yielding a consistent, sweep-free selector of the number of experts that attains pointwise-optimal parameter rates under overfitting while avoiding multi-size training. Simulations on synthetic data corroborate the theory, accurately recovering the expert count and achieving the predicted rates for parameter estimation while closely approximating the regression function. Under model misspecification (e.g., ϵ-contamination), the dendrogram selection criterion is robust, recovering the true number of mixture components, while the Akaike information criterion, the Bayesian information criterion, and the integrated completed likelihood tend to overselect as sample size grows. On a maize proteomics dataset of drought-responsive traits, our dendrogram-guided SGMoE selects two experts, exposes a clear mixing-measure hierarchy, stabilizes the likelihood early, and yields interpretable genotype-phenotype maps, outperforming standard criteria without multi-size training.

We study resource allocation problems in which a central planner allocates resources among strategic agents with private cost functions in order to minimize a social cost, defined as an aggregate of the agents’ costs. This setting poses two main challenges: (i) the agents’ cost functions may be unknown to them or difficult to specify explicitly, and (ii) agents may misreport their costs strategically. To address these challenges, we propose an algorithm that combines preference-based learning with Vickrey–Clarke–Groves (VCG) payments to incentivize truthful reporting. Our algorithm selects informative preference queries via D-optimal design, estimates cost parameters through maximum likelihood, and computes VCG allocations and payments based on these estimates. In a one-shot setting, we prove that the mechanism is approximately truthful, individually rational, and efficient up to an error of 𝒪̃(K^-1/2) for K preference queries per agent. In an online setting, these guarantees hold asymptotically with sublinear regret at a rate of 𝒪̃(T^2/3) after T rounds. Finally, we validate our approach through a numerical case study on demand response in local electricity markets.

Modern cryptographic methods for implementing privacy-preserving LLMs such as HE require the LLMs to have a polynomial form. Forming such a representation is challenging because transformers include non-polynomial components, such as Softmax and layer normalization. Previous approaches have either directly approximated pre-trained models with large-degree polynomials, which are less efficient over HE, or replaced non-polynomial components with easier-to-approximate primitives before training, e.g., Softmax with pointwise attention. The latter approach might introduce scalability challenges. We present a new HE-friendly variant of self-attention that offers a stable form for training and is easy to approximate with polynomials for secure inference. Our work introduces the first polynomial LLMs over a billion parameters, exceeding the size of previous models by more than tenfold. The resulting models demonstrate reasoning and in-context learning (ICL) capabilities comparable to standard transformers of the same size, representing a breakthrough in the field. Finally, we provide a detailed latency breakdown for each computation over encrypted data, paving the way for further optimization, and explore the differences in inductive bias between models relying on our HE-friendly variant and standard transformers. Our code is attached as a supplement.

Conditional independence tests (CIT) are widely used for causal discovery and feature selection. Even with false discovery rate (FDR) control procedures, they often fail to provide frequentist guarantees in practice. We highlight two common failure modes: (i) in small samples, asymptotic guarantees for many CITs can be inaccurate and even correctly specified models fail to estimate the noise levels and control the error, and (ii) when sample sizes are large but models are misspecified, unaccounted dependencies skew the test’s behavior and fail to return uniform p-values under the null. We propose Empirically Calibrated Conditional Independence Tests (ECCIT), a method that measures and corrects for miscalibration. For a chosen base CIT (e.g., GCM, HRT), ECCIT optimizes an adversary that selects features and response functions to maximize a miscalibration metric. ECCIT then fits a monotone calibration map that adjusts the base-test p-values in proportion to the observed miscalibration. Across empirical benchmarks on synthetic and real data, ECCIT achieves valid FDR with higher power than existing calibration strategies while remaining test agnostic. Code is available at <https://github.com/tansey-lab/ECCIT.>

Variational resampling (VR) is a method for deterministically resampling the N particles in sequential Monte Carlo (SMC) algorithms (also known as particle filters), by minimising the Kullback–Leibler divergence from the empirical measure of the N weighted original particles to the empirical measure of M unweighted resampled particles. The combination of VR with a weight transformation (called smoothing weights) has shown to often yield a smaller mean-square error (MSE) than standard resampling schemes in the literature. However, its bias has never been investigated. In this paper, we first show that VR incurs a weighting bias and a truncation bias. We then propose a mechanism to alleviate the weighting bias through an uneven weighting of the resampled particles. We also show that the truncation bias implies that the particle approximation of the target distribution is restricted to a region in which the unnormalised weights are larger than some threshold with high probability. We prove that this probability approaches 1 if M = O(N) as N →∞. Finally, we empirically illustrate that the smaller MSE of VR observed in the literature may be attributable to an underestimation of uncertainty caused by the use of the smoothing weights.

Machine learning models are often trained on proprietary or private datasets that cannot be openly shared. However, the trained model weights are frequently distributed under the assumption that sharing model parameters does not compromise the confidentiality or privacy of the training data. In this work, we challenge this assumption by presenting Training-Like Data Reconstruction (TLDR), as a general-purpose and architecture-agnostic framework for reconstructing training data from a fully trained classifier. Our approach leverages network inversion techniques to recover data that closely resembles the original training samples by exploiting key properties of the classifier with respect to the training data, without requiring access to training dynamics, gradients, pre-trained models, auxiliary datasets, or unobvious priors. Operating in this extreme setting, we demonstrate successful reconstruction of samples with high similarity to the original training data from diverse classifier architectures highlighting critical privacy concerns associated with sharing model parameters. While prior work in this extreme setting has been limited to binary MLP classifiers trained on small datasets, our framework extends to multi-class classification tasks for models based on diverse architectures trained on significantly larger and more complex datasets. Furthermore, we provide quantitative evaluation using the Structural Similarity Index Measure (SSIM) to compare the reconstructed samples with the training samples.

The Muon optimizer has rapidly emerged as a powerful, geometry-aware alternative to AdamW, demonstrating strong performance in large-scale training of neural networks. However, a critical theory-practice disconnect exists: Muon’s efficiency relies on fast, approximate orthogonalization, while most theoretical analyses study idealized exact-SVD updates. This work moves beyond the ideal by providing a general analysis of the inexact orthogonalized update at Muon’s core. We develop our analysis within the general framework of Linear Minimization Oracle (LMO)-based optimization, introducing a realistic additive error model to capture the inexactness of practical approximation schemes. Our analysis yields explicit bounds that quantify performance degradation as a function of the LMO inexactness/error, δ. We reveal a fundamental coupling between this inexactness and the optimal step size and momentum: lower oracle precision requires a smaller step size but larger momentum parameter. These findings elevate the approximation procedure, such as the number of Newton-Schulz steps, from an implementation detail to a critical parameter that must be co-tuned with the learning schedule. NanoGPT experiments directly confirm the predicted coupling, with optimal learning rates clearly shifting as approximation precision changes.

Vision transformers (ViT) rely on attention mechanism to weigh input features, and therefore attention scores have naturally been considered as explanations for its decision-making process. However, attention scores are almost always non-zero, resulting in noisy attention maps and limiting interpretability. Can we quantify uncertainty measures of attention scores and obtain regularized attention scores? To this end, we consider attention scores of ViT in a statistical framework where, e.g., noise would lead to insignificant yet non-zero scores. Leveraging statistical learning techniques, we introduce the bootstrapping for attention scores which generates a baseline distribution of attention scores by resampling input features. Such a bootstrap distribution is then used to estimate significances and posterior probabilities of attention scores. In natural and medical images, the proposed Attention Regularization approach demonstrates a straightforward removal of spurious attention arising from noise, drastically improving shrinkage and sparsity. Quantitative evaluations are conducted using both simulation and real-world datasets. Our study highlights bootstrapping as a practical regularization tool when using attention scores as explanations for ViT.

Classical mixture models (MMs) are widely used tractable proposals for approximate inference settings such as variational inference (VI) and importance sampling (IS). Recently, mixture models with negative coefficients, called subtractive mixture models (SMMs), have been proposed as a potentially more expressive alternative. However, how to effectively use SMMs for VI and IS is still an open question as they do not provide latent variable semantics and therefore cannot use sampling schemes for classical MMs. In this work, we study how to circumvent this issue by designing several expectation estimators for IS and learning schemes for VI with SMMs, and we empirically evaluate them for distribution approximation. Finally, we discuss the additional challenges in estimation stability and learning efficiency that they carry and propose ways to overcome them. Code is available at <https://github.com/april-tools/delta-vi.>

The central goal of active learning is to gather data that maximises downstream predictive performance, but popular approaches have limited flexibility in customising this data acquisition to different downstream problems and losses. We propose a rigorous loss-driven approach to Bayesian active learning that allows data acquisition to directly target the loss associated with a given decision problem. In particular, we show how any loss can be used to derive a unique objective for optimal data acquisition. Critically, we then show that any loss taking the form of a weighted Bregman divergence permits analytic computation of a central component of its corresponding objective, making the approach applicable in practice. In regression and classification experiments with a range of different losses, we find our approach reduces test losses relative to existing techniques.

We study counterfactual prediction under assignment bias and propose a mathematically grounded, information-theoretic approach that removes treatment–covariate dependence without adversarial training. Starting from a bound that links the counterfactual–factual risk gap to mutual information, we learn a stochastic representation Z that is predictive of outcomes while minimizing I(Z;T). We derive a tractable variational objective that upper-bounds the information term and couples it with a supervised decoder, yielding a stable, provably motivated training criterion. The framework extends naturally to dynamic settings by applying the information penalty to sequential representations at each decision time. We evaluate the method on controlled numerical simulations and a real-world clinical dataset, comparing against recent state-of-the-art balancing, reweighting, and adversarial baselines. Across metrics of likelihood, counterfactual error, and policy evaluation, our approach performs favorably while avoiding the training instabilities and tuning burden of adversarial schemes.

This work studies and develops projection-free algorithms for online learning with linear optimization oracles (a.k.a. Frank–Wolfe) for handling the constraint set, and for convex loss functions. More precisely, this work (i) shows how to exploit semidefinite programming to jointly design and analyze online Frank–Wolfe-type algorithms numerically in a variety of settings, (ii) leverages those design techniques to propose an improved (optimized) variant of an online Frank–Wolfe algorithm along with its conceptually simple potential-based proof, and (iii) extends this proof to its anytime version, which benefits from a similar O(T^3/4) regret rate without requiring knowledge of the time horizon T in advance. We are not aware of other direct regret guarantees for an anytime version of online Frank–Wolfe without using the classical doubling trick. Based on the semidefinite technique, we conclude with strong numerical evidence suggesting that no pure online Frank–Wolfe algorithm within our model class can have a regret guarantee better than O(T^3/4) without additional assumptions, that the current algorithms do not have optimal constants, and that multiple linear optimization rounds do not generally help to obtain better regret bounds.

Vector quantization is a prevalent technique for reducing the memory footprint in a wide range of computational problems, such as the training and deployment of deep learning models and vector search systems. We introduce PolarQuant, a novel online vector quantization method that leverages random preconditioning and polar transformation. Our approach efficiently transforms Euclidean vectors into polar coordinates using a recursive algorithm, and then quantizes the resulting angles. A key insight is that, following random preconditioning, the angles in the polar representation exhibit a tightly bounded, highly concentrated, and analytically computable distribution, independent of the input data. This nice distribution eliminates the need for explicit normalization and learned data-dependent quantization codebooks, steps that introduce significant memory and runtime overhead in traditional product/scalar quantization methods. By circumventing this data-dependent step, PolarQuant achieves substantial memory and runtime savings, making it highly suitable for online scenarios such as KV cache compression. The long-context evaluation demonstrates that PolarQuant compresses the KV cache by over 4.2X while achieving the best quality scores compared to the state-of-the-art methods.