AISTATS 2026 Papers

We introduce a novel reinforcement learning algorithm (AGRO, for Any-Generation Reward Optimization) for finetuning Large Language Models. AGRO leverages the concept of response consistency, which states that the optimal policy satisfies a notion of consistency across any possible generation of the model. We derive algorithms that find optimal solutions via sample-based policy gradient and provide theoretical guarantees on their convergence. Our experiments demonstrate the effectiveness of AGRO in both on-policy and off-policy settings, showing improved performance on the MATH dataset over baseline methods.

We introduce off-policy distributional Q(λ), a new addition to the family of off-policy distributional evaluation algorithms. Off-policy distributional Q(λ) does not apply importance sampling for off-policy learning, which introduces intriguing interactions with signed measures. Such unique properties distributional Q(λ) from other existing alternatives such as distributional Retrace. We characterize the algorithmic properties of distributional Q(λ) and validate theoretical insights with tabular experiments. We show how distributional Q(λ)-C51, a combination of Q(λ) with the C51 agent, exhibits promising results on deep RL benchmarks.

Ensuring that classifiers are well-calibrated, i.e., their predictions align with observed frequencies, is a minimal and fundamental requirement for classifiers to be viewed as trustworthy. Existing methods for assessing multiclass calibration often focus on specific aspects associated with prediction (e.g., top-class confidence, class-wise calibration) or utilize computationally challenging variational formulations. In this work, we study scalable evaluation of multiclass calibration. To this end, we propose utility calibration, a general framework which measures the calibration error relative to a specific utility function that encapsulates the goals or decision criteria relevant to the end user. We demonstrate how this framework can unify and re-interpret several existing calibration metrics, particularly allowing for more robust versions of the top-class and class-wise calibration metrics, and going beyond such binarized approaches, towards assessing calibration for richer classes of downstream utilities.

Diffusion models achieve strong generative performance but often rely on large datasets that may include sensitive content. This challenge is compounded by the models’ tendency to memorize training data, raising privacy concerns. SFBD (Lu et al., 2025) addresses this by training on corrupted data and using limited clean samples to capture local structure and improve convergence. However, its iterative denoising and fine-tuning loop requires manual coordination, making it burdensome to implement. We reinterpret SFBD as an alternating projection algorithm and introduce a continuous variant, SFBD flow, that removes the need for alternating steps. We further show its connection to consistency constraint-based methods, and demonstrate that its practical instantiation, Online SFBD, consistently outperforms strong baselines across benchmarks.

Many graph learning applications involve analyzing geometric graphs (e.g., nearest neighbor graphs over embeddings) built over sensitive data, thus requiring formal privacy protections. In this paper, we study benchmark problems in privately analyzing geometric graphs obtained from high dimensional embeddings. We provide several new results for the differentially private approximation of minimum spanning trees and hierarchical clustering in Euclidean graphs. Our algorithms achieve a near optimal privacy-utility trade-off (up to constants), providing a (1+η)-multiplicative approximation with Õ(ρ/η^2) additive error per edge of the tree under ρ-dist privacy (a generalization of DP in geometric data where neighboring datasets different in a single point moved by at most ρ distance). Furthermore, we establish a separation between Euclidean and general graphs by proving a lower bound of Ω(ρ√(n)) additive error per edge of the tree for general graphs under a similar privacy notion, demonstrating that better utility is indeed achievable (allowing also multiplicative approximation) for geometric data. Our algorithm can also be directly applied to widely used clustering algorithm based on MST, incurring only a small loss in the approximation guarantee compared to its non-private counterpart.

Gaussian Processes (GPs) provide a flexible and statistically principled foundation for modelling spatiotemporal phenomena, but their 𝒪(N^3) scaling makes them intractable for large datasets. Approximate methods such as variational inference (VI), inducing-point (sparse) GPs, low-rank kernel approximations (e.g., Nystrom methods and random Fourier features), and approximations such as INLA improve scalability but typically trade off accuracy, calibration, or modelling flexibility. We introduce DeepRV, a neural-network surrogate that replaces GP prior sampling, while closely matching full GP accuracy at inference including hyperparameter estimates, and reducing computational complexity to 𝒪(N^2), increasing scalability and inference speed. DeepRV serves as a drop-in replacement for GP prior realisations in e.g. MCMC-based probabilistic programming pipelines, preserving full model flexibility. Across simulated benchmarks, non-separable spatiotemporal GPs, and a real-world application to education deprivation in London (n = 4,994 locations), DeepRV achieves the highest fidelity to exact GPs while substantially accelerating inference. Code is provided in the dl4bi Python package, with all experiments run on a single consumer-grade GPU to ensure accessibility for practitioners.

The core of generalization theory was developed for independent observations. Some PAC and PAC-Bayes bounds are available for data that exhibit a temporal dependence. However, there are constants in these bounds that depend on properties of the data-generating process: mixing coefficients, mixing time, spectral gap... Such constants are unknown in practice. In this paper, we prove a new PAC-Bayes bound for Markov chains. This bound depends on a quantity called the pseudo-spectral gap, γ_ps. The main novelty is that we can provide an empirical bound on γ_ps when the state space is finite. Thus, we obtain the first fully empirical PAC-Bayes bound for Markov chains. This extends beyond the finite case, although this requires additional assumptions. On simulated experiments, the empirical version of the bound is essentially as tight as the one that depends on γ_ps.

The classic secretary problem, a cornerstone of optimal stopping theory, assumes a random, immutable arrival order of candidates. While recent work has integrated machine-learned predictions to improve selection, the power to set the arrival order based on these predictions remains largely untapped. This paper introduces a novel framework for the secretary problem that leverages predictions for both valuation and strategic scheduling. We propose an algorithm that strategically controls the arrival time of the top-predicted candidate and dynamically adapts its hiring policy based on observed prediction accuracy. Our analysis shows that this approach achieves a worst-case competitive ratio of 0.229, surpassing the 0.215 bound of state-of-the-art algorithms that do not control ordering, bringing it closer to the upper bound of 1/e ≈ 0.368 while maintaining consistency guarantees. Furthermore, we demonstrate that our ordering framework can be adapted to improve fairness guarantees, doubling the success probability in a known fair algorithm from 1/16 to 1/8. Our results highlight that controlling the sequence is a powerful tool for building more robust and fair learning-augmented online algorithms.

The growing focus on distributed data and privacy has spurred the rise of Federated Learning (FL). Empirical studies show that, under equal resources, FL often underperforms centralized training, but the reasons behind this gap remain theoretically unclear. This lack of understanding leaves open whether FL is inherently inferior in generalization and how the gap might be closed. We address this by formulating FL as a server-based SGD optimization problem over distributed data and analyzing the generalization gap within the PAC-Bayesian framework. Our analysis derives non-vacuous bounds on this gap, showing that such a gap necessarily exists under equal resources and depends on training parameters. We further prove that the gap can be fully eliminated only by introducing new clients or adding new data to existing clients, with the latter being more efficient. In contrast, allowing FL to have advantages in other resources, such as larger models or more communication rounds, cannot close the gap. As a complementary analysis, we also confirm from a stability perspective that centralized FL holds a generalization advantage over decentralized FL, justifying our FL formulation choice. Extensive experiments across different model architectures and datasets validate our theory.

We propose a scalable framework for solving the Maximum Cut (MaxCut) problem in large graphs using projected gradient ascent on quadratic objectives. Our approach is differentiable and leverages GPUs for gradient-based optimization. It is not a machine learning method and does not require training data. Starting from a continuous relaxation of the classical quadratic binary formulation, we present a parallelized strategy that explores multiple initialization vectors in batch. We analyze the relaxed objective, showing it is convex and has fixed-points corresponding to local optima—particularly at boundary points—highlighting a key challenge in non-convex optimization. To improve exploration, we introduce a lifted quadratic formulation that over-parameterizes the solution space. We also provide a theoretical characterization of these lifted fixed-points. Finally, we propose DECO, a dimension-alternating algorithm that switches between the unlifted and lifted formulations, combined with importance-based degree initialization and a population-based evolutionary hyper-parameter search. Experiments on diverse graph families show that our methods attain comparable or superior performance relative to recent neural networks and GPU-accelerated sampling approaches.

Large language models (LLMs) have demonstrated remarkable performance across a wide range of tasks. However, the quadratic complexity of softmax attention remains a central bottleneck that limits their scalability. Alman and Song (NeurIPS 2023a; NeurIPS 2024a) proposed sub-quadratic time algorithms for attention inference and training, respectively, but they rely on the restrictive bounded-entry assumption. We show that this assumption rarely holds in practice, which significantly limits their applicability to modern LLMs. In this paper, we introduce support-basis decomposition, a new technique for accurate and efficient attention inference and training without the bounded-entry assumption. We empirically show that the entries of the query and key matrices exhibit sub-Gaussian behavior. Leveraging this widely observed property, we perform exact computation on sparse components and polynomial approximation on dense components. Without relying on restrictive assumptions, we theoretically show that our algorithm achieves sub-quadratic runtime while matching the approximation error of prior work, and we empirically validate its computational efficiency and downstream task performance. We further generalize our method to a multi-threshold setting that eliminates all distributional assumptions, providing the first theoretical justification for the empirical success of polynomial attention. Moreover, we show that softmax attention can be closely approximated by multiple polynomial attentions with significantly smaller ℓ_p error.

Score-based query attacks pose a serious threat to deep learning models by crafting adversarial examples (AEs) using only black-box access to model output scores, iteratively optimizing inputs based on observed loss values. While recent runtime defenses attempt to disrupt this process via output perturbation, most either require access to model parameters or fail when attackers adapt their tactics. In this paper, we first reveal that even the state-of-the-art plug-and-play defense can be bypassed by adaptive attacks, exposing a critical limitation of existing runtime defenses. We then propose Dashed Line Defense (DLD), a plug-and-play post-processing method specifically designed to withstand adaptive query strategies. By introducing ambiguity in how the observed loss reflects the true adversarial strength of candidate examples, DLD prevents attackers from reliably analyzing and adapting their queries, effectively disrupting the AE generation process. We provide theoretical guarantees of DLD’s defense capability and validate its effectiveness through experiments on ImageNet, demonstrating that DLD consistently outperforms prior defenses—even under worst-case adaptive attacks—while preserving the model’s predicted labels.

Thompson sampling is one of the most popular learning algorithms for online sequential decision-making problems and has rich real-world applications. However, traditional Thompson sampling algorithms are limited by the assumption that the rewards received are uncorrupted, which may not hold in real-world applications where adversarial reward poisoning exists. To make Thompson sampling more reliable, our goal is to make it robust against adversarial reward poisoning. Particularly, we consider a strong attack threat model where an adversary applies corruption after observing the agent’s actions. The main challenge is that one can no longer compute the actual posteriors for the true reward, as the agent can only observe the rewards after corruption. In this work, we solve this problem by computing pseudo-posteriors that are less likely to be manipulated by the attack. Particularly, we focus on two popular settings: stochastic bandits and contextual linear bandits with priors as Gaussian distributions. We are the first to propose robust algorithms based on Thompson sampling for the two bandit settings in both cases where the agent is aware or unaware of the attacker’s budget. We theoretically show that our algorithms guarantee near-optimal regret under any attack strategy.

We study a generalization of the Online Convex Optimization (OCO) framework with time-varying adversarial constraints. In this setting, at each round, the learner selects an action from a convex decision set 𝒳, after which both a convex cost function and a convex constraint function are revealed. The objective is to design a computationally efficient learning policy that simultaneously achieves low regret with respect to the cost functions and low cumulative constraint violation (CCV) over a horizon of length T. A major computational bottleneck in standard OCO algorithms is the projection operation onto the decision set 𝒳. However, for many structured decision sets, linear optimization can be performed efficiently. Motivated by this, we propose a projection-free online conditional gradient (OCG)-based algorithm that requires only a single call to a linear optimization oracle over 𝒳 per round. Our approach improves upon the state of the art for projection-free online learning with adversarial constraints, achieving Õ(T^3/4) bounds for both regret and CCV. Our algorithm is conceptually simple. It constructs a surrogate cost function as a nonnegative linear combination of the cost and constraint functions, and feeds these surrogate costs into a novel adaptive online conditional gradient subroutine introduced in this paper. We further extend our framework to the bandit setting, where we show that a new form of surrogate loss is necessary to properly handle bandit feedback—an issue overlooked in prior work. Finally, we develop an efficient Follow-the-Perturbed-Leader (FTPL)-based algorithm, particularly well-suited for online combinatorial optimization problems with discrete actions, which also achieves O(T^3/4) regret and CCV.

Calibration is central to reliable semantic uncertainty quantification, yet prior work has largely focused on discrimination, neglecting calibration. As calibration and discrimination capture distinct aspects of uncertainty, focusing on discrimination alone yields an incomplete picture. We address this gap by systematically evaluating both aspects across a broad set of confidence measures. We show that current approaches, particularly fixed-temperature heuristics, produce systematically miscalibrated and poorly discriminative semantic confidence distributions. We demonstrate that optimising a single scalar temperature, which, we argue, provides a suitable inductive bias, is a surprisingly simple yet effective solution. Our exhaustive evaluation confirms that temperature scaling consistently improves semantic calibration, discrimination, and downstream entropy, outperforming both heuristic baselines and more expressive token-level recalibration methods on question-answering tasks.

We propose a method for learning and sampling from probability distributions supported on the simplex. Our approach maps the open simplex to Euclidean space via smooth bijections, leveraging the Aitchison geometry to define the mappings, and supports modeling categorical data by a Dirichlet interpolation that dequantizes discrete observations into continuous ones. This enables density modeling in Euclidean space through the bijection while still allowing exact recovery of the original discrete distribution. Compared to previous methods that operate on the simplex using Riemannian geometry or custom noise processes, our approach works in Euclidean space while respecting the Aitchison geometry, and achieves competitive performance on both synthetic and real-world data sets.

Transformer-based models, despite their promise for long-term time series forecasting (LTSF), suffer from an inherent low-pass filtering effect that limits their effectiveness. This issue arises due to undifferentiated propagation of frequency components across layers, causing a progressive attenuation of high-frequency information crucial for capturing fine-grained temporal variations. To address this limitation, we propose Dualformer, a principled dual-domain framework that rethinks frequency modeling from a layer-wise perspective. Dualformer introduces three key components: (1) a dual-branch architecture that concurrently models complementary temporal patterns in both time and frequency domains; (2) a hierarchical frequency sampling module that allocates distinct frequency bands to different layers, preserving high-frequency details in lower layers while modeling low-frequency trends in deeper layers; and (3) a periodicity-aware weighting mechanism that dynamically balances contributions from the dual branches based on the harmonic energy ratio of inputs, supported theoretically by a derived lower bound. This design enables structured frequency modeling and adaptive integration of time-frequency features, effectively preserving high-frequency information and enhancing generalization. Extensive experiments conducted on eight widely used benchmarks demonstrate Dualformer’s robustness and superior performance, particularly on heterogeneous or weakly periodic data.

Explaining a single model can be misleading when many near-optimal models (a Rashomon set) yield different feature attributions. We frame this as a Rashomon set sampling problem and propose two practical axioms that any Rashomon sampler should satisfy: generalizability (meaning it must accept arbitrary reference models and loss functions) and Implementation Sparsity (meaning it should return a small, attribution-diverse subset of valid models). These two axioms are not satisfied by most known attribution methods, which we consider to be a fundamental weakness. Building on these axioms, we propose an ϵ-subgradient-based sampling framework and quantify effectiveness with Search Efficiency Ratio (SER) and Functional Explanation Range (FER). Experiments on a synthetic quadratic task and five real-world datasets show that our sampler achieves comparable or higher FER with up to ∼100× fewer models than exhaustive baselines such as TreeFARMS, while remaining agnostic to model class and loss. Even when the reference model is sub-optimal in practice, the resulting attributions align with ground truth and accepted domain knowledge.

Adapting pre-trained models to unseen feature modalities has become increasingly important due to the growing need for cross-disciplinary knowledge integration. A key challenge here is how to align the representation of new modalities with the most relevant parts of the pre-trained model’s representation space to enable accurate knowledge transfer. This requires combining feature alignment with target fine-tuning, but uncalibrated combinations can exacerbate misalignment between the source and target feature-label structures and reduce target generalization. Existing work however lacks a theoretical understanding of this critical interaction between feature alignment and target fitting. To bridge this gap, we develop a principled framework that establishes a provable generalization bound on the target error, which explains the interaction between feature alignment and target fitting through a novel concept of feature-label distortion. This bound offers actionable insights into how this interaction should be optimized for practical algorithm design. The resulting approach achieves significantly improved performance over state-of-the-art methods across a wide range of benchmark datasets.