UAI 2026 Papers

Mutual information is one of the most general ways to measure relationships between random variables, but estimating this quantity for complex systems is challenging. Denoising diffusion models have recently set a new bar for density estimation, so it is natural to consider whether these methods could also be used to improve mutual information estimation. Using the recently introduced information-theoretic formulation of denoising diffusion models, we show that diffusion models can be used in a straightforward way to estimate mutual information. In particular, the mutual information corresponds to half the gap in the minimum mean square error between conditional and unconditional diffusion, integrated over all signal-to-noise ratios in the noising process. Our approach not only passes self-consistency tests but also outperforms traditional and score-based diffusion estimators. Furthermore, our method leverages adaptive importance sampling to achieve scalable estimation, while maintaining strong performance even when the mutual information is high. Code to reproduce our experiments is available at <https://github.com/fengsxy/Diffusion_MI>, and the unified mutual information estimation library is available at <https://github.com/fengsxy/Diffusion-MI>.

The value function of a POMDP exhibits the piecewise-linear-convex (PWLC) property and can be represented as a finite set of hyperplanes, known as α-vectors. Most state-of-the-art POMDP solvers (offline planners) follow the point-based value iteration scheme, which performs Bellman backups on α-vectors at reachable belief points until convergence. However, since each α-vector is |S|-dimensional, these methods quickly become intractable for large-scale problems due to the prohibitive computational cost of Bellman backups. In this work, we demonstrate that the PWLC representation of POMDP value functions can be generalized by replacing the alpha-vector sets with a finite set of neural networks. This insight enables a novel POMDP planning algorithm, called *Neural Value Iteration*, which combines the generalization capability of neural networks with the classical value iteration framework. Requiring only a black-box simulator, our approach scales to extremely large POMDPs that are intractable for existing offline solvers.

Estimating outcome distributions under time-varying treatments is an essential task for personalized decision-making, particularly in domains such as healthcare. Most prior work in this area focuses on point predictions, which fail to capture the inherent variability in outcomes. Recent efforts in causal inference have begun integrating generative models to address this limitation by estimating interventional distributions. However, existing approaches—including causal normalizing flows—are generally restricted to static settings and are not well suited to sequential, time-dependent data. In this work, we propose a novel framework that extends causal normalizing flows to time-series, enabling simulation-based interventional density estimation over time. Our method learns representations of treatment and covariate history that capture temporal dependencies. Conditioned on these representations and guided by a causal graph, our flow-based model generates interventional samples, allowing for the simulation of outcome trajectories under alternative treatment strategies. We evaluate our approach on both linear and non-linear synthetic time-series as well as on a simulated tumor growth dataset, demonstrating that it achieves performance competitive with state-of-the-art baselines, while accommodating a broader spectrum of causal queries.

Large language models based on the transformer architecture can solve highly complex tasks, yet their fundamental limitations on simple algorithmic problems remain poorly understood. In this work, we focus on basic counting tasks and investigate how the difficulty of these tasks scales with the transformer embedding dimension, the context length, and the vocabulary size. We reveal a sharp theoretical phase transition governed by the relationship between the embedding dimension and the vocabulary size. When the dimension is at least as large as the vocabulary, transformers can perfectly maintain token counts. However, when the vocabulary exceeds the embedding dimension, the interference between non-orthogonal token representations forces the network weights to scale polynomially. This renders the exact counting algorithm numerically unstable and practically unlearnable. We empirically validate this bottleneck by training transformers from scratch, demonstrating a strict performance drop at the theoretical threshold and catastrophic out of distribution failure when scaling the vocabulary or context length. Furthermore, we show that state-of-the-art pretrained models suffer from similar failure cases. Our work reveals a critical blind spot absent from the current literature regarding the connection among these three parameters, proving that vocabulary size fundamentally dictates the difficulty of counting tasks.

Convergence-rate analysis for classifiers is often conducted under either Tsybakov margin or Massart margin. The former is a relatively weak condition that typically yields polynomial rates, while the latter is substantially stronger but can guarantee exponential rates. In this paper, we introduce a new condition, called *Boltzmann margin*, that bridges the gap between these two regimes. It is weaker than Massart margin, generally stronger than Tsybakov margin, and can imply many of their properties under suitable conditions. We apply Boltzmann margin to the analysis of kNN classifiers and establish the first *near-exponential* convergence rates for kNN classification. We also present extensions of the main results and provide numerical evidence supporting the main theoretical implications.

Due to the customized requirements, few-shot diffusion models have attracted much attention. However, only a few works analyze few-shot models, and none involve the fast few-shot optimization process, which is important for quickly responding to users. In this work, we evaluate the role of each operation in the optimization process and prove the convergence guarantee for few-shot diffusion models. A standard operation for the few-shot model is only fine-tuning some key parameters to avoid overfitting the limited target dataset. We first show that this operation is insufficient from empirical and theoretical perspectives. Empirically, we conduct real-world few-shot fine-tuning experiments with underfitting and overfitting bad pre-trained models and show that the results are heavily influenced by these bad models. Theoretically, we also prove that the few-shot phase can not learn the ground-truth parameters and suffers from a small gradient when using a bad pre-trained model. Based on these results, we highlight the importance of a great pre-trained model by showing it can warm up few-shot models and lead to a strongly convex landscape for few-shot diffusion models. As a result, the few-shot model fast converges to the ground-truth parameters. In contrast, we show that with a bad initialization, the pretraining phase requires large optimization steps to converge.

Despite the growing demand for eliciting uncertainty from large language models (LLMs), empirical evidence suggests that LLM behavior is not always adequately captured by the elicitation techniques developed under the classical probabilistic uncertainty framework. This mismatch leads to systematic failure modes, particularly in settings that involve ambiguous question-answering, in-context learning, and self-reflection. To address this, we propose novel prompt-based uncertainty elicitation techniques grounded in *imprecise probabilities*, a principled framework for representing and eliciting higher-order uncertainty. Here, first-order uncertainty captures uncertainty over possible responses to a prompt, while second-order uncertainty (uncertainty about uncertainty) quantifies indeterminacy in the underlying probability model itself. We introduce general-purpose prompting and post-processing procedures to directly elicit and quantify both orders of uncertainty, and demonstrate their effectiveness across diverse settings. Our approach enables more faithful uncertainty reporting from LLMs, improving credibility and supporting downstream decision-making.

Vectorized end-to-end (E2E) autonomous driving models predominantly rely on reactive imitation learning. By failing to model underlying environmental dynamics, these agents lack counterfactual reasoning capabilities and struggle with uncertainty in out-of-distribution (OOD) scenarios. To bridge the gap between reactive imitation and deliberative decision-making, we introduce Reasoning-enhanced VAD (ReVAD), empowering agents with a “System 2” cognitive framework via Latent Space Search. Unlike computationally expensive pixel-level world models, ReVAD integrates a probabilistic Token Dynamics Model with Latent Monte Carlo Tree Search to perform efficient lookahead planning entirely within a sparse semantic token space. By simulating future states to evaluate risk and utilizing imitation policies solely as search priors, ReVAD effectively mitigates the distribution shift inherent in pure imitation learning, demonstrating significantly improved robustness and safety in high-uncertainty environments.

The transition from multi-rounds federated learning to one-shot federated learning (OFL) markedly alleviates communication burden and represents a major step toward realistic deployment. Most existing OFL approaches require clients to perform local training, imposing a substantial computational burden on client side, while others rely on pre-trained models. In this paper, we propose a novel one-shot federated learning framework based on random feature extractor (FedRFE). Unlike existing approaches, it does not require any local model training or pre-trained model, featuring superior resource efficiency. Through comprehensive experiments, we show FedRFE achieves competitive performance while being robust in challenging scenarios including data heterogeneity and client scalability. Our code is available at<https://github.com/luyuanxyang/FedRFE>

In optimization-based decision-making, when the objective coefficient vector is unknown, a common approach is to predict it from covariates using a machine learning model and solve the downstream optimization problem. However, improving predictive accuracy does not necessarily lead to better decisions. This has motivated Decision-Focused Learning (DFL), which trains predictive models to minimize decision loss measured by regret. Despite recent progress, existing DFL methods do not sufficiently address two sources of uncertainty: (i) observation errors in the measured coefficients, and (ii) distribution shift in the coefficient vector at deployment due to environmental changes. These uncertainties can degrade solution quality. In this paper, we propose two robust regret losses to address these uncertainties. For uncertainty (i), we introduce an uncertainty set around the observed coefficient vector that captures measurement errors and define the loss as the worst-case regret over this set. For uncertainty (ii), we construct a Wasserstein ambiguity set around the conditional empirical distribution and define the loss using the worst-case distribution within the set. We enable efficient training via Danskin-based subgradients. Experiments demonstrate that our method reduces regret and yields more stable solutions than existing robust DFL approaches.

Stochastic Gradient Descent (SGD) is a cornerstone of modern optimization. While its performance under heavy-tailed noise is often addressed through specialized modifications such as gradient clipping or normalization, we investigate a more fundamental question: how does vanilla SGD, particularly with momentum, perform in the presence of heavy-tailed noise? In this paper, we refine existing convergence results for vanilla SGD and, more importantly, provide the first comprehensive convergence analysis of vanilla SGD with momentum for strongly convex, convex, and nonconvex objectives, without employing any gradient control mechanisms. Our results demonstrate that the obtained convergence rates are inferior to the optimal rates achieved by clipped or normalized variants of SGD, thereby revealing inherent limitations of vanilla methods under heavy-tailed noise. The theoretical findings are supported by experiments on synthetic functions.

Variational Inference (VI) is a fundamental inference technique in Bayesian machine learning for approximating complex posterior distributions. Traditional VI often relies on the mean-field factorization, which can inadequately capture true posterior complexity. Recent advancements have leveraged neural networks to model implicit distributions, offering increased flexibility. However, the practical constraints of neural network architectures still produces inaccuracies. In this paper, we propose a method called Implicit Variational Rejection Sampling (IVRS), which integrates implicit distributions with rejection sampling to improve the posterior approximation. Our method uses neural networks to construct implicit proposal distributions, and rejection sampling with a discriminator network that estimates the density ratio between the implicit proposal and the true posterior for refining the approximation. Towards this end, we introduce the Implicit Resampling Evidence Lower Bound (IR-ELBO) as a metric to characterize the resampled distribution’s quality and derive a tighter variational lower bound. Experimental results demonstrate that our method outperforms traditional variational inference techniques.

We study model-free methods for distributionally robust infinite-horizon average-reward Markov decision processes (MDPs). We present non-asymptotic convergence analyses of Q-learning and actor–critic algorithms for robust average-reward MDPs under contamination, total-variation distance, and Wasserstein uncertainty sets. A key ingredient of our analysis is showing that the optimal robust Bellman operator is a strict contraction with respect to a carefully designed semi-norm. This property enables a stochastic approximation update that learns the optimal robust Q-function using 𝒪̃(ϵ^-2) samples. We also establish robust TD convergence bounds whose constants are uniform over all stationary policies, yielding an efficient data-driven routine for robust critic estimation. Building on this, we introduce an actor–critic algorithm that learns an ϵ-optimal robust policy within 𝒪̃(ϵ^-2) samples. We provide numerical simulations to illustrate the qualitative behavior of the proposed algorithms. Our results contribute to the theoretical foundations of robust planning under model misspecification, and to model-free approaches for building robust long-run policies directly from simulation data.

Events in spatiotemporal domains arise in numerous real-world applications, where uncovering event relationships and enabling accurate prediction are central challenges. Classical Poisson and Hawkes processes rely on restrictive parametric assumptions that limit their ability to capture complex interaction patterns, while recent neural point process models increase representational capacity but integrate event information in a black-box manner, hindering interpretable relationship discovery. To address these limitations, we propose a Kronecker-Structured Nonparametric Spatiotemporal Point Process (KSTPP) that enables transparent event-wise relationship discovery while retaining high modeling flexibility. We model the background intensity with a spatial Gaussian process (GP) and the influence kernel as a spatiotemporal GP, allowing rich interaction patterns including excitation, inhibition, neutrality, and time-varying effects. To enable scalable training and prediction, we adopt separable product kernels and represent the GPs on structured grids, inducing Kronecker-structured covariance matrices. Exploiting Kronecker algebra substantially reduces computational cost and allows the model to scale to large event collections. In addition, we develop a tensor-product Gauss-Legendre quadrature scheme to efficiently evaluate intractable likelihood integrals. Extensive experiments demonstrate the effectiveness of our framework.

Despite being widely adopted as a canonical framework for learning robust models, adversarial training suffers from robust overfitting. Existing empirical and theoretical explorations fail to provide a satisfactory mechanistic interpretation of the phenomenon. By modeling adversarial training with momentum SGD as a discrete-time dynamical system, we propose a PAC-Bayesian analytical framework that proves time-resolved robust generalization bounds. Specifically, our framework tracks the closed-form evolution of the posterior mean and covariance under both stationary and non-stationary transient regimes, connecting the model’s robust generalization performance to learning rate, local loss geometry, and mini-batch stochastic gradients. By estimating the key quantities associated with the bound, we illustrate the underlying mechanism of robust overfitting. Our framework also shows how adversarial weight perturbation reduces robust generalization gaps by suppressing dominant loss-curvature modes, while suggesting that excessive penalization can be sub-optimal for optimization.

In many scientific applications, hypotheses are generated and tested continuously in a stream. We develop a framework for improving online multiple testing procedures with false discovery rate (FDR) control under arbitrary dependence. Our approach is two-fold: we construct methods via the online e-closure principle, as well as a novel formulation in online compound e-values constructed via donations. This yields strict power improvements over state-of-the-art e-value and p-value procedures while retaining FDR control. We further derive algorithms that compute the rejection decision at time t in O(log t) time, and we demonstrate improved empirical performance on synthetic and real data.

When large language models produce natural-language rationales, those explanations are frequently unfaithful to the model’s actual reasoning—and no existing framework provides a principled way to repair them when new evidence arrives. We introduce ARGUS, a framework that structures LLM self-explanations as Dung-style abstract argumentation frameworks and verifies them under grounded and preferred semantics. When new or uncertain evidence renders an explanation inconsistent, ARGUS computes a minimum-cost set of edit operations that restores the desired acceptability status of the target argument. The repair operator satisfies adapted AGM revision postulates and is bidirectionally characterized by them (Representation Theorem): the decision problem is in P under grounded semantics, NP-complete under preferred and stable semantics, and Σ_2^P-complete under skeptical stable semantics. A k-neighborhood approximation and an answer set programming (ASP) encoding ensure scalability to practical framework sizes. We validate the framework on HotpotQA and FEVER, where ARGUS achieves relative improvements of 10.3

We explore the problem of federated combinatorial causal bandits (FedCCB), where multiple agents collaboratively select variables for intervention and gather feedback. The primary objective of FedCCB is to identify optimal interventions for each agent while minimizing the total cumulative regret associated with the target nodes. A key challenge in FedCCB stems from the inherent heterogeneity of local causal models, which often exhibit diverse causal influences. To address this challenge, we propose a novel Federated Subset-Clustered Bandit (FedSCuB) method. This method groups agents to tackle heterogeneity based on the similarity of their causal relationships with specified subsets of variables. FedSCuB incorporates an intervention-based exploration strategy to collect partial observations necessary for clustering, as well as an alternating minimization method to facilitate collaboration among agents within the same cluster. Theoretical analysis demonstrates that the proposed FedSCuB method achieves sub-linear regret and offers a better regret bound compared to baseline methods. Empirical evaluations on synthetic tasks further confirm the effectiveness and superiority of our method.

One of the main challenges in reinforcement learning (RL) is that the agent has to make decisions that would influence the future performance without having complete knowledge of the environment. Dynamically adjusting the level of epistemic risk during the learning process can help to achieve reliable policies in safety-critical settings with better efficiency. In this work, we propose a new framework, Distributional RL with Online Epistemic Risk Adaptation (DRL-ORA). This framework quantifies both epistemic and implicit aleatory uncertainties in a unified manner and dynamically adjusts the epistemic risk levels by solving a total variation minimization problem online. The framework generalizes the existing variants of risk adaptation approaches with better explainability and flexibility. The selection of risk levels is performed efficiently via a Follow-The-Leader-type algorithm. We show that DRL-ORA outperforms existing methods that rely on fixed risk levels or manually designed risk level adaptation in multiple classes of tasks.