AISTATS 2026 Papers

Amortized meta-learning methods, such as neural processes, promise near-instantaneous inference on new labeled datasets encountered during downstream tasks. Recent adaptations of the transformer architecture have propelled these approaches to impressive performance in tasks like function estimation, parameter inference, and decision-making. Curiously, their success still primarily stems from the expressiveness of transformers, lacking a bias for modeling the functional structures between features and labels shared across datasets. We argue and show that this leads to training sample inefficiency and sub-optimal performance, and address this by introducing a novel set encoding technique called Predictive Deep Sets. Our approach exploits a strong bias towards functional structures by meta-learning an RKHS that captures domain-critical functional patterns, and by representing datasets as optimal fit functions within this space. Besides providing theoretical justification for this approach, we empirically demonstrate orders of magnitude increases in training data sample efficiency compared to strong baselines across various settings.

Federated Unlearning (FU) enables clients to remove the influence of specific data from a collaboratively trained shared global model, addressing regulatory requirements such as GDPR and CCPA. However, this unlearning process introduces a new privacy risk: A malicious server may exploit unlearning updates to reconstructthe data requested for removal, a form of Data Reconstruction Attack (DRA). While DRAs for machine unlearning have been studied extensively in centralized Machine Learning-as-a-Service (MLaaS) settings, their applicability to FU remains unclear due to the decentralized, client-driven nature of FU. This work presents DRAUN, the first attack framework to reconstruct unlearned data in FU systems. DRAUN targets optimization-based unlearning methods, which are widely adopted for their efficiency. We theoretically demonstrate why existing DRAs targeting machine unlearning in MLaaS fail in FU and show how DRAUN overcomes these limitations. We validate our approach through extensive experiments on five datasets and five model architectures, evaluating its performance against five popular unlearning methods, effectively demonstrating that state-of-the-art FU methods remain vulnerable to DRAs.

In this work, a tensor completion problem is studied, which aims to perfectly recover the tensor from partial observations. The existing theoretical guarantee requires the involved transform to be orthogonal, which hinders its applications. In this paper, jumping out of the constraints of isotropy and invertibility for the first time, the theoretical guarantee of exact tensor completion with arbitrary linear transforms is established by directly operating the tensors in the transform domain. With the enriched choices of transforms, we theoretically disclose why slim transforms outperform their square counterparts, providing support for existing works on experimental excellence of slim transforms. Our model and analysis greatly enhance the flexibility of tensor completion, supported by extensive experimental results.

Feature-based explanation methods aim to quantify how features influence the model’s behavior, either locally or globally, but different methods often disagree, producing conflicting explanations. This disagreement arises primarily from two sources: how feature interactions are handled and how feature dependencies are incorporated. We propose GRANITE, a generalized regional explanation framework that partitions the feature space into regions where interaction and distribution influences are minimized. This approach aligns different explanation methods, yielding more consistent and interpretable explanations. GRANITE unifies existing regional approaches, extends them to feature groups, and introduces a recursive partitioning algorithm to estimate such regions. We demonstrate its effectiveness on real-world datasets, providing a practical tool for consistent and interpretable feature explanations.

Learning the structure of a Bayesian network from decentralized data poses two major challenges: (i) ensuring rigorous privacy guarantees for participants, and (ii) avoiding communication costs that scale poorly with dimensionality. In this work, we introduce Fed-Sparse-BNSL, a novel federated method for learning linear Gaussian Bayesian network structures that addresses both challenges. By combining differential privacy with greedy updates that target only a few relevant edges per participant, Fed-Sparse-BNSL efficiently uses the privacy budget while keeping communication costs low. Our careful algorithmic design preserves model identifiability and enables accurate structure estimation. Experiments on synthetic and real datasets demonstrate that Fed-Sparse-BNSL achieves utility close to non-private baselines while offering substantially stronger privacy and communication efficiency.

In two-player zero-sum games, the learning dynamic based on optimistic Hedge achieves one of the best-known regret upper bounds among strongly-uncoupled learning dynamics. With an appropriately chosen learning rate, the social and individual regrets can be bounded by O(log(mn)) in terms of the numbers of actions m and n of the two players. This study investigates the optimality of the dependence on m and n in the regret of optimistic Hedge. To this end, we begin by refining existing regret analysis and show that, in the strongly-uncoupled setting where the opponent’s number of actions is known, both the social and individual regret bounds can be improved to O(√(log m log n)). In this analysis, we express the regret upper bound as an optimization problem with respect to the learning rates and the coefficients of certain negative terms, enabling refined analysis of the leading constants. We then show that the existing social regret bound as well as these new social and individual regret upper bounds cannot be further improved for optimistic Hedge by providing algorithm-dependent individual regret lower bounds. Importantly, these social regret upper and lower bounds match exactly including the constant factor in the leading term. Finally, building on these results, we improve the last-iterate convergence rate and the dynamic regret of a learning dynamic based on optimistic Hedge, and complement these bounds with algorithm-dependent dynamic regret lower bounds that match the improved bounds.

Deep Feedback Models (DFMs) are a new class of stateful neural networks that combine bottom up input with high level representations over time. This feedback mechanism introduces dynamics into otherwise static architectures, enabling DFMs to iteratively refine their internal state and mimic aspects of biological decision making. We model this process as a differential equation solved through a recurrent neural network, stabilized via exponential decay to ensure convergence. To evaluate their effectiveness, we measure DFMs under two key conditions: robustness to noise and generalization with limited data. In both object recognition and segmentation tasks, DFMs consistently outperform their feedforward counterparts, particularly in low data or high noise regimes. In addition, DFMs translate to medical imaging settings, while being robust against various types of noise corruption. These findings highlight the importance of feedback in achieving stable, robust, and generalizable learning. Code is available at github.com/DCalhas/deepfeedbackmodels.

In this paper, we provide a novel perspective on the underlying structure of real-world data with ground-truth clustering via characterization of an abundantly observed yet often overlooked density–geometry correlation. We leverage this correlation to design CoreSPECT (Core Space Projection based Enhancement of Clustering Techniques), a general framework that improves the performance of generic clustering algorithms. Our framework boosts the performance of clustering algorithms by applying them to strategically selected regions, then extending the partial partition to a complete partition for the dataset using a novel neighborhood graph based multi-layer propagation procedure. We provide initial theoretical support of the functionality of our framework under the assumption of our model, and then provide large-scale real-world experiments on 20 datasets that include standard image datasets as well as genomics datasets. We observe two notable improvements. First, CoreSPECT improves the NMI of K-Means by 20

Transfer entropy is a fundamental measure for quantifying directed information flow in time series, with applications spanning neuroscience, finance, and complex systems analysis. However, existing estimation methods suffer from the curse of dimensionality, require restrictive distributional assumptions, or need exponentially large datasets for reliable convergence. We address these limitations in the literature by proposing TENDE (Transfer Entropy Neural Diffusion Estimation), a novel approach that leverages score-based diffusion models to estimate transfer entropy through conditional mutual information. By learning score functions of the relevant conditional distributions, TENDE provides flexible, scalable estimation while making minimal assumptions about the underlying data-generating process. We demonstrate superior accuracy and robustness compared to existing neural estimators and other state-of-the-art approaches across synthetic benchmarks and real data.

State-of-the-art post-training pipelines for reasoning LLMs rely on bootstrapped reasoning loops: they sample many traces, score them, and reinforce the highest-scoring ones, typically by correctness. This can improve accuracy while still collapsing the distribution inside the correct set onto a narrow family of redundant strategies, reducing creative problem-solving. To diagnose this failure mode, we introduce Distributional Creative Reasoning (DCR), a variational framework that casts training as gradient flow on the simplex of reasoning traces. The framework yields three core results. First, a diversity-decay analysis shows that STaR-style rejection fine-tuning and exact mean-field GRPO amplify whichever correct trace is already larger, while DPO regresses pairwise correct-trace ratios toward the reference ratios. Second, it explains why entropy and KL can slow or tether collapse but do not reward semantically distinct correct strategies for being distinct, and how a creativity kernel supplies the missing relational term. Third, under mild conditions, the resulting dynamics converge to a unique, stable, and diverse equilibrium, yielding practical guidance for kernel and hyperparameter design. DCR thus offers a principled route to training reasoning LLMs that remain both correct and creative.

SimCLR is a popular contrastive learning method for vision tasks, renowned for its ability to pre-train neural networks to learn efficient representations. Despite its empirical effectiveness, the theoretical understanding of SimCLR is still very limited, even in the simplest learning scenarios. In this paper, we introduce a theoretical case study of SimCLR. Specifically, we consider training a two-layer convolutional neural network (CNN) to learn a toy image data model that has been considered in a series of recent works. For this particular learning task, we precisely characterize the label complexity under which SimCLR pre-training followed by supervised fine-tuning achieves approximately zero training loss and almost optimal test loss. Notably, the label complexity for SimCLR pre-training is far less demanding compared to direct supervised training, especially when the signal-to-noise ratio in the data is low. Our analysis sheds light on the benefits of SimCLR in learning with fewer labels.

Supervised disentanglement, that is, learning interpretable nonlinear latent representations of a target data view informed by an auxiliary data view, is a central challenge in interpretable machine learning. We formulate this problem as a partially linear invertible canonical correlation analysis (PLiCCA). Specifically, given two data views, (i) complex data lying near a potentially high-dimensional manifold, and (ii) auxiliary high-dimensional multivariate data, PLiCCA learns latent variables for the complex view that are maximally correlated with sparse linear combinations of the auxiliary variables. In contrast to regression-based approaches to supervised disentanglement, the proposed method yields a latent embedding whose coordinates are explicitly ordered by their interpretability with respect to the auxiliary variables. We formalize the population PLiCCA problem and establish existence results. We then show a close theoretical connection between PLiCCA and conditional latent variable models, in particular conditional variational autoencoders and conditional normalizing flows, which enables practical estimation. We demonstrate our approach on brain imaging data, where PLiCCA is used to learn embeddings informed by auxiliary demographic, psychometric, and behavioral variables.

Deep learning models are primarily trained through batchwise optimization, where each update can potentially be a tug-of-war among samples, shaping the overall trajectory of learning. Existing interpretability tools, most notably influence functions, have provided valuable insights into how individual training samples affect model predictions, primarily at test time. However, these methods were not intended to capture these inter-sample interactions that arise during training. Here, we ask a complementary question: How does optimizing the loss on one training sample affect the loss on the rest during learning? We introduce Influence Graphs (IGs), directed inter-sample graphs where each edge weight w_ij quantifies how optimizing on sample X_i influences the loss of sample X_j. We estimate these influences via simulated batch interventions and slope coefficients of loss changes, enabling scalable construction of IGs during training. We further define the Mean-of-Mean In-Degree Influence (MMDI) and prove it bounds generalization under practical assumptions. Empirically, MMDI correlates strongly with test accuracy in noisy-label settings, making it a useful diagnostic of model quality even before test metrics are available. Finally, we show that IGs reveal distinct, evolving training phases, offering a new lens on the dynamics of learning.

Disentanglement learning is central to understanding and reusing learned representations in variational autoencoders (VAEs). Although equivariance has been explored in this context, effectively exploiting it for disentanglement remains challenging. In this paper, we propose a novel method, called Multiple Invertible and Partial-Equivariant Transformation (MIPE-Transformation), which integrates two main parts: (1) Invertible and Partial-Equivariant Transformation (IPE-Transformation), guaranteeing an invertible latent-to–transformed-latent mapping while preserving partial input-to-latent equivariance in the transformed latent space; and (2) Exponential-Family Conversion (EF-Conversion) to extend the standard Gaussian prior to an approximate exponential family via a learnable conversion. In experiments on the 3D Cars, 3D Shapes, and dSprites datasets, MIPE-Transformation improves the disentanglement performance of state-of-the-art VAEs.

We present Flow-Reasoner, a Deliberate-When-Needed neuro-symbolic model that integrates continuous latent cognition with selective symbolic reasoning. The mental module is a latent state vector evolving smoothly under a first-order ordinary differential equation (ODE), capturing continuous thought that drifts and decays between interventions. The action module is a temporal point process whose intensities are modulated by symbolic rules. Crucially, reasoning is not constant: it is triggered only at irregular instants—when an observed action arrives or when a latent state crosses a threshold—at which point a bounded differentiable forward-chaining procedure updates beliefs and adjusts event likelihoods. Between these triggers, cognition evolves autonomously under the ODE without symbolic intervention. This design yields a model that (i) unifies continuous-time dynamics with selective logical reasoning, (ii) predicts both the type and timing of future actions, and (iii) produces concise rule traces that explain predictions. Empirical studies on synthetic benchmarks and real-world behavioral datasets demonstrate that Flow-Reasoner consistently outperforms strong temporal point process baselines, while providing interpretable, cognitively inspired explanations of decision dynamics. The code is publicly available at <https://github.com/shennnnwj/flow-reasoner.>

Convex Markov Games (cMGs) were recently introduced as a broad class of multi-agent learning problems that generalize Markov games to settings where strategic agents optimize general utilities beyond additive rewards. While cMGs expand the modeling frontier, their theoretical foundations, particularly the structure of Nash equilibria (NE) and guarantees for learning algorithms, are not yet well understood. In this work, we address these gaps for an extension of cMGs, which we term General Utility Markov Games (GUMGs), capturing new applications requiring coupling between agents’ occupancy measures. We prove that in GUMGs, Nash equilibria coincide with the fixed points of projected pseudo-gradient dynamics (i.e. first-order stationary points), enabled by a novel agent-wise gradient domination property. This insight also yields a simple proof of NE existence using Brouwer’s fixed-point theorem. We further show the existence of Markov perfect equilibria. Building on this characterization, we establish a policy gradient theorem for GUMGs and design a model-free policy gradient algorithm. For potential GUMGs, we establish iteration complexity guarantees for computing approximate-NE under exact gradients and provide sample complexity bounds in both the generative model and on-policy settings. Our results extend beyond prior work restricted to zero-sum cMGs, providing the first theoretical analysis of common-interest cMGs.

Biological machine learning is often bottlenecked by a lack of scaled data. One promising route to relieving data bottlenecks is through high-throughput screens, which can experimentally test the activity of 10^6-10^12 protein sequences in parallel. In this article, we introduce algorithms to optimize high throughput screens for data creation and model training. We focus on the large-scale regime, where dataset sizes are limited by the cost of measurement and sequencing. We show that when active sequences are rare, we maximize information gain if we only collect positive examples of active sequences, i.e. x with y>0. We can correct for the missing negative examples using a generative model of the library, producing a consistent and efficient estimate of the true p(y| x). We demonstrate this approach in simulation and on a large-scale screen of antibodies. Overall, co-design of experiments and inference lets us accelerate learning dramatically.

Gradient clipping is a fundamental tool in Deep Learning, improving the high-probability convergence of stochastic first-order methods like SGD, AdaGrad, and Adam under heavy-tailed noise, which is common in training large language models. It is also a crucial component of Differential Privacy (DP) mechanisms. However, existing high-probability convergence analyses typically require the clipping threshold to increase with the number of optimization steps, which is incompatible with standard DP mechanisms like the Gaussian mechanism. In this work, we close this gap by providing the first high-probability convergence analysis for DP-Clipped-SGD with a fixed clipping level, applicable to both convex and non-convex smooth optimization under heavy-tailed noise, characterized by a bounded central α-th moment assumption, α∈ (1,2]. Our results show that, with a fixed clipping level, the method converges to a neighborhood of the optimal solution with a faster rate than the existing ones. The neighborhood can be balanced against the noise introduced by DP, providing a refined trade-off between convergence speed and privacy guarantees.

Countless science and engineering applications in multi-objective optimization (MOO) necessitate that decision-makers (DMs) select a Pareto-optimal (PO) solution which aligns with their preferences. Evaluating individual solutions is often expensive, and the high-dimensional trade-off space makes exhaustive exploration of the full Pareto frontier (PF) infeasible. We introduce a novel, principled two-step process for obtaining a compact set of PO points that aligns with user preferences, which are specified a priori as general monotonic utility functions (MFs). Our process (1) densely samples the user’s region of interest on the PF, then (2) sparsifies the results into a small, diverse set for the DM. We instantiate this framework with soft-hard functions (SHFs), an intuitive class of MFs that operationalizes the common expert heuristic of imposing soft and hard bounds. We provide extensive empirical validation of our framework instantiated with SHFs on diverse domains, including brachytherapy, engineering design, and large language models. For brachytherapy, our approach returns a compact set of points with over 3