AISTATS 2026 论文列表

Although actor-critic methods have been successful in practice, their theoretical analyses have several limitations. Specifically, existing theoretical work either sidesteps the exploration problem by making strong assumptions or analyzes impractical methods with complicated algorithmic modifications. Moreover, the actor-critic methods analyzed for linear MDPs often employ natural policy gradient and construct "implicit" policies without explicit parameterization. Such policies are computationally expensive to sample from, making the environment interactions inefficient. To that end, we focus on the finite-horizon linear MDPs and propose an optimistic actor-critic framework that uses parametric log-linear policies. In particular, we introduce a tractable logit-matching regression objective for the actor. For the critic, we use approximate Thompson sampling via Langevin Monte Carlo to obtain optimistic value estimates. We prove that the resulting algorithm achieves 𝒪(ϵ^-4) and 𝒪(ϵ^-2) sample complexity in the on-policy and off-policy setting, respectively. Our results match prior theoretical work in achieving the state-of-the-art sample complexity, while our algorithm is more aligned with practice.

With LLMs increasingly deployed in corporate data management, it is crucial to ensure that these models do not leak sensitive information. In the context of corporate data management, the concept of sensitivity awareness has been introduced, enabling LLMs to adhere to predefined access rights rules. However, it remains unclear how sensitivity awareness relates to established notions of privacy, such as differential privacy (DP), thereby making it difficult to deploy meaningfully in real-world applications. In this work, we formalize the notion of sensitivity awareness and theoretically establish its connection to DP. Additionally, we develop a supervised fine-tuning recipe to make existing, four-bit quantized LLMs more sensitivity-aware. With a performance boost of up to 21.7

Factual reliability in domain-specific Large Language Models (LLMs) is paramount in high-stakes applications where incorrect outputs carry significant risks. Current detection methodologies often rely on expensive retrieval validation or labor-intensive manual annotation, creating substantial barriers to scalable deployment. To bridge the gap, we propose SiGHT, a self-supervised graph framework designed for efficient hallucination detection in specialized contexts. SiGHT introduces an automated training pipeline that leverages prompt strategies to synthesize plausible hallucinated content from structured knowledge, effectively eliminating the need for human labeling. By mapping texts to high-resolution word-level relational graphs, the framework employs a Graph Attention Network (GAT) to model fine-grained semantic dependencies and identify structural inconsistencies. Empirical evaluations on the MSMARCO-QnA and RAGTruth-QA benchmarks demonstrate that SiGHT achieves a 46.94

We study stochastic nonconvex optimization under heavy-tailed noise. In this setting, the stochastic gradients only have bounded p–th central moment (p–BCM) for some p ∈ (1,2]. Building on the foundational work of Arjevani et al. (2022) in stochastic optimization, we establish tight sample complexity lower bounds for all first-order methods under relaxed mean-squared smoothness (q-WAS) and δ-similarity ((q,δ)-S) assumptions, allowing any exponent q∈[1,2] instead of the standard q= 2. These results substantially broaden the scope of existing lower bounds. To complement them, we show that Normalized Stochastic Gradient Descent with Momentum Variance Reduction (NSGD-MVR), a known algorithm, matches these bounds in expectation. Beyond expectation guarantees, we introduce a new algorithm, Double-Clipped NSGD-MVR, which allows the derivation of high-probability convergence rates under weaker assumptions than in previous works. Finally, for second-order methods with stochastic Hessians satisfying bounded q-th central moment assumptions for some exponent q ∈[1,2] (allowing q≠ p), we establish sharper lower bounds than previous works while improving over Sadiev et al. (2025) (where only p=q is considered) and yielding stronger convergence exponents. Together, these results provide a nearly complete complexity characterization of stochastic nonconvex optimization in heavy-tailed regimes.

Deep equilibrium models (DEQs) achieve infinitely deep network representations without stacking layers by exploring fixed points of layer transformations in neural networks. Such models constitute an innovative approach that achieves performance comparable to state-of-the-art methods in many large-scale numerical experiments, despite requiring significantly less memory. However, DEQs face the challenge of requiring vastly more computational time for training and inference than conventional methods, as they repeatedly perform fixed-point iterations with no convergence guarantee upon each input. Therefore, this study explored an approach to improve fixed-point convergence and consequently reduce computational time by restructuring the model architecture to guarantee fixed-point convergence. Our proposed approach for image classification, Lipschitz multiscale DEQ, has theoretically guaranteed fixed-point convergence for both forward and backward passes by hyperparameter adjustment, achieving up to a 4.75× speedup in numerical experiments on CIFAR-10 at the cost of a minor drop in accuracy.

Machine learning has become increasingly popular in informing data-driven policy-making. Policies influence behavior in individuals or populations, and ideally, through observational signals, policy-makers learn which policies are effective. However, in many settings, individual actions cannot be perfectly observed. This issue, known in economics as moral hazard, poses a significant challenge. In this work, we study the foundational multitasking principal–agent contract design problem and demonstrate how instrumental regression and the generalized method of moments (GMM) estimator can be used to estimate or learn a good contract. As a bonus result, we also give a uniformity characterization of the shape of the optimal contract.

We model neural training as a classical multi-time map from controllable interventions—batch choices, augmentations, and optimizer micro-steps—to model predictions on a fixed probe set. On this basis, we introduce a simple, model-agnostic witness of training memory based on back-flow of distinguishability. In a controlled two-step protocol, we compare predictive distributions after one intervention versus two; a positive increase Δ_BF = D_2 - D_1 > 0, with D∈TV, JS, H, certifies observable-level non-Markovianity. Across controlled SGD experiments, we observe consistent positive back-flow with tight bootstrap confidence intervals, stronger effects under higher momentum, larger batch overlap, and more micro-steps, and marked collapse under a causal break that resets optimizer state. The witness is inexpensive, requires no architectural changes, and is robust across TV/JS/Hellinger. We position this as a measurement contribution: a practical diagnostic, and empirical evidence, that real training often deviates from the Markov idealization in ways that matter for optimizer behavior, data order, and schedule design.

Self-attention layers have become fundamental building blocks of modern deep neural networks, yet their theoretical understanding remains limited, particularly from the perspective of random matrix theory. In this work, we provide a rigorous analysis of the singular value spectrum of the attention matrix and establish the first Gaussian equivalence result for attention. In a natural regime where the inverse temperature remains of constant order, we show that the singular value distribution of the attention matrix is asymptotically characterized by a tractable linear model. We further demonstrate that the distribution of squared singular values deviates from the Marchenko–Pastur law, which has been believed in previous work. Our proof relies on two key ingredients: precise control of fluctuations in the normalization term and a refined linearization that leverages favorable Taylor expansions of the exponential. This analysis also identifies a threshold for linearization and elucidates why attention, despite not being an entrywise operation, admits a rigorous Gaussian equivalence in this regime.

Large Language Models (LLMs) are increasingly deployed in safety-critical systems that rely heavily on reading comprehension—extracting and reasoning over exten- sive in-context information. However, existing evaluations of LLMs on reading comprehension are typically over limited test sets containing only a tiny fraction of the vast number of possible prompts. Empirical evaluations on these test sets have questionable reliability and generalizability. We propose a fundamentally different approach: rather than evaluating LLMs with fixed datasets, we introduce the first framework for certifying LLMs based on large probability distributions over realistic reading comprehension prompts. To create these distributions, we use knowledge graphs (KGs) as structured representations of real-world knowledge and define the distributions’ sample spaces with prompts based on directed acyclic subgraphs of the KGs. We also incorporate realistic noise designed to mimic real-world complexity, such as distractor texts and synonyms. Our prompt distributions have i.i.d. samplers represented as probabilistic programs. Our framework generates novel, formal probabilistic quantitative certificates that provide high-confidence, tight bounds on the probability that an LLM correctly answers any prompt drawn from these distributions. We enable formal certification for SOTA LLMs by using an input-output example-driven approach. We apply our framework to certify SOTA LLMs in precision medicine and general question-answering domains. Our results uncover previously unknown vulnerabilities caused by natural prompt noise and establish the first formal performance hierarchies among these models.

The Lasso has been widely used in a high-dimensional setting, but its estimation accuracy may become inadequate when the covariates are highly correlated or when the number of covariates is extremely large. To overcome this problem, we propose a novel preconditioner that adaptively induces a low-rank structure in the design matrix. The proposed preconditioner achieves a higher probability of sign correctness under some conditions. We establish theoretical guarantees showing that our method dominates the standard Lasso, and we further demonstrate its superiority over the correlation shifting. To validate its practical effectiveness, we conducted numerical experiments on synthetic and semi-real datasets, and the proposed method presented better performance than existing methods.

Unsupervised learning of high-dimensional data is challenging due to irrelevant or noisy features obscuring underlying structures. It’s common that only a few features, called the influential features, meaningfully define the clusters. Recovering these influential features is helpful in data interpretation and clustering. We propose i-IF-Learn, an iterative unsupervised framework that jointly performs feature selection and clustering. Our core innovation is an adaptive feature selection statistic that effectively combines pseudo-label supervision with unsupervised signals, dynamically adjusting based on intermediate label reliability to mitigate error propagation common in iterative frameworks. Leveraging low-dimensional embeddings (PCA or Laplacian eigenmaps) followed by k-means, i-IF-Learn simultaneously outputs influential feature subset and clustering labels. Numerical experiments on gene microarray and single-cell RNA-seq datasets show that i-IF-Learn significantly surpasses classical and deep clustering baselines. Furthermore, using our selected influential features as preprocessing substantially enhances downstream deep models such as DeepCluster, UMAP, and VAE, highlighting the importance and effectiveness of targeted feature selection.

Transformers play an important role in molecular representation learning, enabling unsupervised learning from large scale unlabeled molecule datasets. However, existing Transformer based methods suffer from heavy training computation and slow inference. To accelerate the computation and relieve the burdensome pre-training, we propose a Mamba-based framework that leverages selective state space models to learn molecular representations more efficiently. Unlike conventional methods, our model, GeoMamba-SE(3), offers streamlined computation with linear-time complexity. However, naively applying Mamba to molecules struggles with SE(3) symmetry, representations can drift under rotations/translations—leading to chemically inconsistent features. To address this, we introduce a geometry and statistics aware design: (i) complete local frames at atoms by converting geometric vectors into scalar channels suitable for SSMs; (ii) multi-stream Mamba blocks are modulated by SE(3)-invariant scalars to preserve geometric stability; and (iii) we impose statistical symmetry constraints via orbit-kernel losses and invariant risk minimization, treating SE(3) actions and conformers as environments. This yields practical SE(3) stability without heavy high-order tensor representations. Experiments show that our method achieves new state-of-the-art performance benchmarks on the MoleculeNet datasets, while using only one-sixth of the training computation and 57

Flow Matching (FM) underpins many state-of-the-art generative models, yet recent results indicate that Transition Matching (TM) can achieve higher quality with fewer sampling steps. This work answers the question of when and why TM outperforms FM. First, when the target is a unimodal Gaussian distribution, we prove that TM attains strictly lower KL divergence than FM for finite number of steps. The improvement arises from stochastic difference latent updates in TM, which preserve target covariance that deterministic FM underestimates. We then characterize convergence rates, showing that TM achieves faster convergence than FM under a fixed compute budget. Second, we extend the analysis to Gaussian mixtures and identify local–unimodality regimes in which the sampling dynamics approximate the unimodal case, where TM can outperform FM. The approximation error decreases as the minimal distance between component means increases, highlighting that TM is favored when the modes are well separated. However, when the target variance approaches zero, each TM update converges to the FM update, and the performance advantage of TM diminishes. In summary, we show that TM outperforms FM when the target distribution has well-separated modes and non-negligible variances. We validate our theoretical results with controlled experiments on Gaussian distributions, and extend the comparison to real-world applications in image and video generation.

Real-world tables are messy: column headers are inconsistent, cells contain errors or missing values, and crucial information is scattered across multiple tables and documents. These issues cause even state-of-the-art language models to fail at seemingly simple questions. We present a robust framework for table understanding that explicitly handles these challenges through three coordinated mechanisms: structure-aware encoders that learn invariance to common corruptions, trainable slots that compress evidence to a fixed-size representation, and grounding modules that align each slot to supporting text passages. Unlike prior work that treats tables as flat text or relies on clean schemas, our approach maintains strong performance even under schema corruption and structural perturbations. Across eight benchmarks spanning question answering, fact verification, and text generation, we achieve the best performance among methods without external tools on five tasks and remain competitive with systems using much larger models or SQL executors. Under schema corruption and row/column permutations, our method degrades by less than 2 points, while baselines drop by up to 6-22 points, confirming that explicit denoising and grounding are essential for robust table understanding.

Identifying the drift and diffusion of an SDE from its population dynamics is a notoriously challenging task. Researchers in machine learning and single-cell biology have only been able to prove a partial identifiability result: for potential-driven SDEs, the gradient-flow drift can be identified from temporal marginals if the Brownian diffusivity is already known. Existing methods therefore assume that the diffusivity is known a priori, despite it being unknown in practice. We dispel the need for this assumption by providing a complete characterization of identifiability: the gradient-flow drift and Brownian diffusivity are jointly identifiable from temporal marginals if and only if the process is observed outside of equilibrium. Given this fundamental result, we propose nn-APPEX, the first Schrödinger Bridge–based inference method that can simultaneously learn the drift and diffusion of a gradient-flow SDE solely from observed marginals. Extensive experiments show that nn-APPEX’s ability to adjust its diffusion estimate enables accurate inference, while previous Schrödinger Bridge methods obtain biased drift estimates due to their assumed, and likely incorrect, diffusion.

Stochastic linear bandits are a fundamental model for sequential decision making. Although well studied in the single-agent setting, many real-world scenarios involve multiple agents solving heterogeneous bandit problems, each with a different unknown parameter. This paper investigates the optimal regret achievable in collaborative personalized linear bandits. We derive an information-theoretic lower bound showing how the number of agents, the number of rounds, and the degree of heterogeneity jointly affect regret. We propose a two-stage collaborative algorithm that achieves the optimal regret. We model heterogeneity via a hierarchical Bayesian framework and introduces a novel information-theoretic technique for bounding regret. Our results offer a complete characterization of when and how collaboration helps with a optimal regret bound Õ(d√(mn)), Õ(dm^1-γ√(n)), Õ(dm√(n)) for the number of rounds n in the range of o ( d/m σ^2), Θ( d/m^2γσ^2) and ω( d/σ^2, ) respectively, where σ measures the level of heterogeneity, m is the number of agents, and γ∈[0, 1/2] is an absolute constant. In contrast, without collaboration achieves a regret bound O(dm√(n)) at best.

In this paper, we study last-iterate convergence of learning algorithms in bilinear saddle-point problems, a preferable notion of convergence that captures the day-to-day behavior of learning dynamics. We focus on the challenging setting where players select actions from compact convex sets and receive only bandit feedback. Our main contribution is the design of an uncoupled learning algorithm that guarantees last-iterate convergence to the Nash equilibrium with high probability. We establish a convergence rate of Õ(T^-1/4) up to polynomial factors in problem parameters. Crucially, our proposed algorithm is computationally efficient, requiring only an efficient linear optimization oracle over the players’ compact action sets. The algorithm is obtained by combining techniques from experimental design and the classic Follow-The-Regularized-Leader (FTRL) framework, with a carefully chosen regularizer function tailored to the geometry of the action set of each learner.

Switchback experiments are widely used in dynamic systems—such as ridesharing platforms and online marketplaces—to evaluate interventions under interference. However, the standard pipeline of estimating the average treatment effect (ATE) with a difference-in-means (DM) estimator can exhibit systematic bias in dynamic settings with evolving system state, due to intertemporal dependence (“carryover effects"). In this paper, we study this bias in a continuous-time Markov chain model of switchback experiments with stochastically monotone dynamics and state-monotone rewards; these are reasonable representations of mean-reverting and auto-regressive systems. We show the DM estimator systematically underestimates the true ATE, because it targets an average of transient treatment effects rather than the ATE itself. Using the Ornstein-Uhlenbeck process as a tractable example, we derive closed-form expressions for bias and variance; this analysis shows that standard approaches overestimate the true variance. Taken together, these effects mean that standard switchback experiment analysis yields overly conservative inference. We validate our theory using a ride-sharing simulation with real-world calibration.

We introduce CoPRIME (Contrastive Probabilistic Routing for IMbalanced tokens with ELBO-regularized mixture of experts), a probabilistic routing framework for multimodal representation learning that generalizes multimodal representation learning beyond vision-text by tackling the fundamental challenge of extreme token imbalance across modalities. This imbalanced-ness is particularly pronounced between spectrogram-tokenized audio and text. CoPRIME augments contrastive pretraining with an ELBO-regularized routing objective that jointly promotes 1) expert specialization, requiring experts to explain the tokens they receive, and 2) diverse utilization via KL regularization to a uniform prior. To stabilize routing, we further replace standard CoV-based regularizers with entropy-based importance and load losses, yielding smoother gradients and flexible, modality-aware routing without rigid uniformity constraints. On MOSEI and IEMOCAP datasets, CoPRIME achieves state-of-the-art zero- and few-shot emotion and sentiment results, outperforming dense Transformers and prior multimodal MoE variants while retaining the efficiency of sparse conditional computation. Ablations isolate the role of each loss and show that ELBO is the primary driver of stable specialization under modality imbalance, with entropy-based regularizers further improving convergence and utilization.