AISTATS 2026 Papers

Reinforcement Learning from Human Feedback (RLHF) has become central to aligning large language models with human values, typically by first learning a reward model from preference data which is then used to update the model with reinforcement learning. Recent alternatives such as Direct Preference Optimization (DPO) simplify this pipeline by directly optimizing on preferences. However, both approaches often assume uniform annotator preferences and rely on binary comparisons, overlooking two key limitations: the diversity of human evaluators and the limitations of pairwise feedback. In this work, we address both these issues. First, we connect preference learning in RLHF with the econometrics literature and show that binary comparisons are insufficient for identifying latent user preferences from finite user data and infinite users, while (even incomplete) rankings over three or more responses ensure identifiability. Second, we introduce methods to incorporate heterogeneous preferences into alignment algorithms. We develop an Expectation-Maximization adaptation of DPO that discovers latent annotator types and trains a mixture of LLMs accordingly. Then we propose an aggregation algorithm using a min-max regret fairness criterion to produce a single generative policy with equitable performance guarantees. Together, these contributions establish a theoretical and algorithmic framework for fairness and personalization for diverse users in generative model alignment.

We introduce LAMP (Local Attribution Mapping Probe), a method that shines light onto a black-box language model’s decision surface and studies how reliably a model maps its stated reasons to its reported predictions by approximating a decision surface. LAMP treats the model’s own self-reported explanations as a coordinate system and fits a locally linear surrogate that links those weights to the model’s output. By doing so, it reveals how much the stated factors steer the model’s decisions. We apply LAMP to three tasks: sentiment analysis, controversial-topic detection, and safety-prompt auditing. Across these tasks, LAMP reveals that many language models’ locally approximated linear decision landscapes overall agree with human judgments on explanation quality and, on a clinical case-file data set, align with expert assessments. Since LAMP operates without requiring access to model gradients, logits, or internal activations, it serves as a practical and lightweight framework for auditing proprietary language models, and enabling assessment of whether a model appears to behave consistently with the explanations it provides.

Knowledge Tracing (KT) aims to predict students’ academic performance by modeling their knowledge mastery over time, based on their historical learning interactions. However, current KT models often oversimplify student interactions by treating them as standard time series rather than as cognitive processes. Consequently, modeling student learning as a process of cognitive transformation rather than as a mere sequence of time-stamped events remains a fundamental challenge in KT research. To address this issue, we propose ACE-KT (cAscaded Cognitive modEling for Knowledge Tracing), a novel framework inspired by cognitive process theory, which shifts the focus from purely sequential modeling to cognitive representation learning. Specifically, we design a cascaded cognitive framework inspired by human cognitive processes in three sequential stages: convolution-based rhythm perception module, Transformer encoder-based contextual structuring module, and cognitive integration module implemented via a selective structured state space model. Extensive experiments on five real-world datasets demonstrate that ACE-KT consistently outperforms 20 SOTA KT baselines, demonstrating its effectiveness. The source code is publicly available at our GitHub repository (<https://github.com/AWord992/ACEKT.git).>

The vulnerabilities of deep neural networks against singularities have raised serious concerns regarding their deployment in the physical world. One of the most prominent and impactful physical-world adversarial perturbations is the attachment of patches to clean images, known as an adversarial patch attack. Similarly, natural noises such as Gaussian and Salt Pepper are highly prevalent in the real world. The current research need arises from the above vulnerabilities and the lack of efforts to tackle these two singularities independently and, especially, in combination. In this research, we have, for the first time, combined these two prominent singularities and proposed a novel dataset. Using this dataset, we have conducted a benchmark study of singularity data-point detection using features from several convolutional neural networks. For classification, rather than the popular neural network-based parameter tuning, we have used traditional yet effective machine learning classifiers. The extensive experiments across various in- and out-of-distribution (OOD) singularities reveal several interesting findings about the effectiveness of classifiers and show that it is hard to defend against adversaries when they are treated independently, and inefficient classifiers are selected.

Approximate inference is central to Bayesian learning, with variational inference (VI) providing a scalable framework for posterior approximation. While mean-field VI often fails in high dimensions, the more refined Bethe approximation, equivalent to the Thouless-Anderson-Palmer (TAP) free energy in statistical physics, has long been conjectured to capture Bayes-optimal behavior. We prove that the TAP formula holds for Bayesian linear regression with a uniform spherical prior at all noise levels (Δ>0), extending the result of Qiu and Sen (2022) in the high-noise regime. Our argument constructs a ridge regression functional that dominates the TAP free energy, yielding the first rigorous analysis of the global optimizer of the non-concave TAP functional for a planted inference model at an arbitrary noise level. This verifies that TAP, rather than mean-field, is the correct variational description in this setting.

As models, particularly large language models (LLMs), are deployed on increasingly challenging tasks, correctly evaluating their performance is growing in importance and difficulty. Expert human labelers are high-quality but scarce and resource-intensive to obtain, while crowd-sourced labels are more readily accessible at scale but lower in quality. We propose Maven (Model And Voter EvaluatioN), a hierarchical Bayesian model that combines these two label sources to produce model performance estimates on binary tasks that are less biased than using crowd-sourced labels alone and have lower variance than using expert labels alone. By modeling the ranking of model scores, Maven is robust to a range of prediction distributions and achieves constant inference time regardless of dataset size. We validate our approach on both simulated and real-world data, and deploy it to measure production models at Meta.

Spatio-temporal event data—such as crime incidents or shared-mobility usage—are generated by human decisions. Yet most existing models focus on statistical dependencies in time and space, overlooking the cognitive and social factors that shape behavior. We argue that uncovering underlying preferences is essential, as they provide a structured link between observed event data and decision processes. We introduce a preference-driven framework that models event distributions through a two-stage “consider–then–choose” process: sparse gating captures limited attention, and utility functions guide selection within the consideration set. To capture heterogeneity, we employ a mixture-of-experts design that reveals distinct preference patterns across groups and contexts. The framework incorporates sparse structural design, and we analyze its theoretical properties by establishing approximation and generalization guarantees. Empirical studies on crime and bike-sharing datasets demonstrate competitive predictive accuracy while providing interpretable insights into behavioral drivers. By shifting the focus from counts to preferences, our approach offers a behaviorally grounded and socially meaningful perspective for modeling event data.

Diffusion models can learn rich representations during data generation, showing potential for Self-Supervised Learning (SSL), but they face a trade-off between generative quality and discriminative performance. Their iterative sampling also incurs substantial computational and energy costs, hindering industrial and edge AI applications. To address these issues, we propose the Flow Matching-based Sensor Foundation Model (SenFlow), which jointly trains a representation encoder and a conditional flow matching generator. This decoupled design achieves both high-fidelity generation and effective recognition. By using flow matching to learn a simpler velocity field, SenFlow accelerates and stabilizes training, improving its efficiency for representation learning. Experiments on wearable sensor data show SenFlow reduces training time by 50.4

Minimum-Norm Interpolators (MNI) in overparameterized linear models have gained attention as a tractable framework for studying interpolation phenomena that resemble empirical observations in neural networks. Most prior work on these interpolators either exploits closed-form solutions when available or relies heavily on Gaussian comparison results, such as the convex Gaussian Min-Max Theorem (CGMT). In this paper, we introduce a new perspective on MNI under isotropic Gaussian covariates by leveraging tools from high-dimensional geometry. First, we obtain a “localized” bound on the MNI’s shrinkage of the original ground truth that occurs under isotropic Gaussian covariates when the norm is in an isotropic position. Then, we prove a sharp bound on the Mean Squared Error (MSE) of the ℓ_1-MNI, as obtained by Wang 22’ via a geometric proof, which avoids invoking the CGMT and instead relies on the work of Fleury 12’ on Gaussian polytopes.

Uniform bounds on sketched inner products underpin several important computational and statistical results in machine learning and randomized algorithms, including the Johnson-Lindenstrauss (J-L) lemma, the Restricted Isometry Property (RIP), randomized sketching, etc. However, many modern analyses involve sketched bilinear forms, for which existing uniform bounds either do not apply or are not sharp on general sets. In this work, we develop a general framework to analyze such sketched bilinear forms, and derive uniform bounds in terms of geometric complexities of the associated sets. Our approach relies on generic chaining and introduces new techniques for handling suprema over pairs of sets. We further extend our results to (i) sketch matrices with conditionally independent entries, e.g., as in CountSketch and SRHT (Subsampled Randomized Hadamard Transform), and (ii) bilinear forms involving a sum of T sketch matrices, showing that the deviation scales as √(T). This unified analysis recovers known results such as the J-L lemma as special cases, while extending RIP guarantees. Using our new bounds, we give tighter convergence bounds for sketched federated learning, and develop sketched bandits whose regret depends on the geometric complexity of the action and parameter sets rather than the ambient dimension.

Policy learning algorithms are regularly leveraged in domains such as personalized medicine and advertising to develop individualized treatment regimes. However, a critical deficit of existing algorithms is that they force a decision even when predictions are uncertain, a risky approach in high-stakes settings. The ability to abstain, that is, to defer to a safe default or an expert, is crucial but largely unexplored in this context. To remedy this, we introduce a framework for policy learning with abstention, in which policies that choose not to assign a treatment to some customers/patients receive a small, additive reward on top of the value of a random guess. We propose a two-stage learner that first identifies a set of near-optimal policies and then constructs an abstention class based on disagreements between the policies. We establish fast O(1/n)-type regret guarantees for the abstaining policy when propensities are known, and show how to extend these guarantees to the unknown-propensity case via a doubly robust (DR) objective. Furthermore, we demonstrate that our abstention framework is a versatile tool with direct applications to several other core problems in policy learning. We use our algorithm as a black box to obtain improved guarantees under margin conditions without the common realizability assumption. We also show that abstention provides a natural connection to both distributionally robust policy learning, where it acts as a hedge against small data shifts, and safe policy improvement, where the goal is to improve upon a baseline policy with high probability.

We study sampling from posterior distributions in Bayesian linear inverse problems where 𝐀, the parameters to observables operator, is computationally expensive. In many applications 𝐀 can be factored in a manner that facilitates the construction of a cost-effective approximation 𝐀. In this framework, we introduce Latent-IMH, a sampling method based on the Metropolis-Hastings independence (IMH) sampler. Latent-IMH first generates intermediate latent variables using the approximate 𝐀, and then refines them using the exact 𝐀. Its primary benefit is that it shifts the computational cost to an offline phase. We theoretically analyze the performance of Latent-IMH using KL divergence and mixing time bounds. Using numerical experiments on several model problems, we show that, under reasonable assumptions, it outperforms state-of-the-art methods such as the No-U-Turn sampler (NUTS) in computational efficiency. In some cases Latent-IMH can be orders of magnitude faster than existing schemes.

Quantum Machine Learning (QML) promises significant computational advantages, but preserving training data privacy remains challenging. Classical approaches like differentially private stochastic gradient descent (DP-SGD) add noise to gradients but fail to exploit the unique properties of quantum gradient estimation. In this work, we introduce the Differentially Private Parameter-Shift Rule (Q-ShiftDP), the first privacy mechanism tailored to QML. By leveraging the inherent boundedness and stochasticity of quantum gradients computed via the parameter-shift rule, Q-ShiftDP enables tighter sensitivity analysis and reduces noise requirements. We combine carefully calibrated Gaussian noise with intrinsic quantum noise to provide formal privacy and utility guarantees, and show that harnessing quantum noise further improves the privacy–utility trade-off. Experiments on benchmark datasets demonstrate that Q-ShiftDP consistently outperforms classical DP methods in QML.

Positional encoding (PE) is a core architectural component of Transformers, yet its impact on the Transformer’s generalization and robustness remains unclear. In this work, we provide the first generalization analysis for single-layer Transformer under in-context regression that explicitly accounts for a trainable PE module. Our result shows that PE systematically enlarges the generalization gap. Extending to the adversarial setting, we derive the adversarial Rademacher generalization bound. We find that the gap between models with and without PE is magnified under attack, demonstrating that PE amplifies the vulnerability of models. Our bounds are empirically validated by a simulation study. Together, this work establishes a new framework for understanding the clean and adversarial generalization in ICL with PE.

We study the decentralized multi-player stochastic bandit problem over a continuous, Lipschitz-structured action space where hard collisions yield zero reward. Our objective is to design a communication-free policy that maximizes collective reward, while separating coordination costs from learning costs. We propose a modular protocol that first solves the multi-agent coordination problem by identifying and seating players on distinct, high-value regions via a novel maxima-directed search and then decouples the problem into N independent single-player Lipschitz bandits. In the consensus regime, we obtain an end-to-end regret bound whose dominant learning term is Õ(T^(d+1)/(d+2)), matching the single-player Lipschitz rate; the upfront coordination cost is horizon-independent at fixed confidence and only polylogarithmic in Tin the expected-regret form. Under an additional public coverage/scheduling assumption for the epochic extension, we also obtain a gap-free Õ(T^(d+1)/(d+2))guarantee. We further derive a matching lower bound for the dominant learning term and extend the framework to general distance-threshold collision models.

We show that minimizing the expected Wasserstein loss between empirical distributions can lead to biased parameter estimates in the finite-sample regime. Remarkably, such bias arises even in well-specified settings where both empirical distributions are drawn from the same parametric family: unlike maximum likelihood estimation—understood here as maximizing the expected log-likelihood—optimizing one parameter while fixing another fails to recover the true fixed value. We derive closed-form expressions for the expected Wasserstein loss in one dimension and, focusing on location–scale models, provide an analytic characterization of the bias. This analysis reveals that finite-sample bias occurs whenever the expected loss varies along the diagonal subspace where parameter values coincide, and we propose a simple correction scheme that removes this effect. We extend our analysis to misspecified models and the Sinkhorn divergence, demonstrating that finite-sample bias persists in more practical settings. Experiments on synthetic and real data confirm that stochastic optimization of Wasserstein-based objectives converges to biased solutions, and validate the effectiveness of the proposed correction scheme.

We study a multi-agent multi-armed bandit problem (MA-MAB) in open systems, where multiple agents can enter and leave at any time and face multiple bandit problems to minimize the group-wise cumulative regret. To our knowledge, this is the first work to consider a dynamic set of agents that arrive and depart according to stochastic processes, systematically evolving over time. We also extend to a permissionless blockchain-based MA-MAB (PB-MA-MAB) problem, where agents may behave either honestly or maliciously depending on compliance with the mechanism, and malicious agents may disrupt honest ones. These formulations pose new challenges, as regret grows with the increasing number of agents. To this end, we design new UCB-based methodologies for both MA-MAB and PB-MA-MAB, introducing information-integration rules for existing agents and information-access mechanisms for new agents to fully leverage available information. We derive regret bounds for our algorithms and characterize the complexity of the formulation via regret lower bounds in both settings. We establish regret upper bounds of order maxO(M_0), O(log T), O(log^2 T(M_0))1_λ > 0 (a significant improvement over the naïve bound (M_0 + T)log T), where M_0 is the initial number of agents and C reflects the arrival/departure rate. We also prove lower bounds of O(log T) and O(M_0) for all consistent algorithms, and tighter bounds of O(log T + M_0) or O(log^2 T) for a subset including ours. These imply that our algorithm is nearly optimal in general and optimal in certain cases.

We study fixed-confidence Best Arm Identification (BAI) in semiparametric bandits, where rewards are linear in arm features plus an unknown additive baseline shift. Unlike linear-bandit BAI, this setting requires orthogonalized regression, and its instance-optimal sample complexity has remained open. For the transductive setting, we establish an attainable instance-dependent lower bound characterized by the corresponding linear-bandit complexity on shifted features. We then propose a computationally efficient phase-elimination algorithm based on a new 𝒳𝒴-design for orthogonalized regression. Our analysis yields a nearly optimal high-probability sample-complexity upper bound, up to log factors and an additive d^2 term, and experiments on synthetic instances and the Jester dataset show clear gains over prior baselines.

Techniques for concept extraction, such as sparse autoencoders and transcoders, aim to extract high-level symbolic concepts from low-level nonsymbolic representations. When these extracted concepts are used for downstream tasks such as model steering and unlearning, it is essential to understand their guarantees, or lack thereof. In this work, we present a unified theoretical framework for unsupervised concept extraction, in which we frame the task of concept extraction as identifying a generative model. We present a general meta-theorem for identifiability, which reduces the problem of establishing identifiability guarantees to the problem of characterizing the intersection of two sets. As we demonstrate on a range of widely-used approaches, this meta-theorem substantially simplifies the task of proving such guarantees, thus paving the way for the development of new, principled approaches for concept extraction.