NeurIPS 2026 Papers

Latent Gaussian models (LGMs) are a popular class of Bayesian hierarchical models that include Gaussian processes, as well as certain spatial models and mixed-effect models. Efficient Bayesian inference of LGMs often requires marginalizing out the latent variables. For LGMs with a non-Gaussian likelihood, exact marginalization is not possible and a popular approach is to do approximate marginalization with an integrated Laplace approximation (ILA). Using ILA produces an approximate posterior which, in some settings, can differ significantly from the correct posterior, which impacts downstream applications. We propose an importance sampling scheme to correct the error introduced by ILA. By increasing the number of samples in importance sampling, the posterior with ILA converges to the correct posterior. This idea is realized with various techniques, including pseudo-marginalization, quasi-Monte Carlo and randomized quasi-Monte Carlo. We implement our methods in an automatic differentiation framework to support gradient-based algorithms when doing inference on the hyperparameters. For the latter, we specifically consider the use of Hamiltonian Monte Carlo. We demonstrate the benefits of reduced error in various applied models.

Efficient sampling from complex probability distributions is an important task across Bayesian inference and certain classes of deep generative models. Markov chain Monte Carlo (MCMC) remains one of the most widely used tools for this. However, even state-of-the-art samplers, such as the NUTS implementations in Stan and Turing.jl, can converge slowly on ill-conditioned, heavy-tailed, or near-degenerate targets. We identify that sampling error can be decomposed into two terms: mismatch in the potential energy (PE) marginal, and average conditional mismatch across PE level sets. We propose *Contour Monte Carlo (CMC)*, which corrects each of these in turn: first by drawing energies $u_i$ from the PE marginal, and then by sampling states on the corresponding level sets. For a broad family of radial targets, including the Gaussian and Student's $t$ distributions, both stages admit closed-form draws, so CMC is exact and requires no Markov chain. Beyond this, for general targets, we estimate the PE marginal using umbrella sampling with MBAR, then run constrained MCMC on the sampled level sets, initialised from energy-matched umbrella samples (which tend to lie in regions of high conditional mass). We show that this procedure essentially reduces the global sampling problem to approximating the one-dimensional PE marginal, after which posterior draws can be generated readily on demand. Furthermore, the level-set chains are embarrassingly parallel, and therefore well-suited to modern accelerator hardware. On a suite of challenging benchmarks -- composed of `posteriordb` models, synthetic targets, and real-world inference problems -- CMC converges substantially faster than HMC and NUTS baselines at a fixed computational budget.

Gaussian processes (GPs) offer a principled probabilistic model over functions, but exact inference is restricted to the linear-Gaussian regime. We establish an explicit equivalence between GPs and a class of linear diffusion models, recasting predictive sampling as an ODE with closed-form Gaussian dynamics and a likelihood-dependent guidance term that admits a simple Monte Carlo approximation. In the linear-Gaussian setting, we recover standard GP conditioning exactly; beyond conjugacy, the same machinery handles any conditioning statement admitting point-wise likelihood evaluation — including non-linear physics, and, for the first time, natural language via large language models. Whitening isolates the irreducible non-Gaussian dynamics, minimising Wasserstein-2 transport cost and eliminating numerical stiffness. The result is a general-purpose GP inference scheme requiring no bespoke derivations. Together, these results provide a general mechanism for incorporating the full richness of real-world knowledge as conditioning information, opening a new frontier for the probabilistic modelling of real-world problems.

Large-scale approximate nearest neighbor search commonly relies on partitions for indexing: database vectors are partitioned into clusters, and for each query a probing function selects the clusters to be scanned. The query probing function and the database partition are rarely treated as separate entities: most techniques assign queries with the same assignment function as the database vectors, which is suboptimal especially when database and query distributions differ. This paper introduces CwA (Cluster with Auctions), which addresses this limitation by jointly learning a balanced database partition and a neural probing function. CwA optimizes search performance directly for the query distribution. It minimizes its objective by alternating two steps: (i) gradient descent on the neural network of the probing function, and (ii) a large-scale combinatorial optimization of the cluster assignment for the database vectors. We solve the latter with a parallelizable auction algorithm that balances the partition by design. To further scale CwA, we extend the method to a Cartesian product of clusters that increases the partition’s granularity. When database and query distributions differ, CwA achieves up to 4.7$\times$ throughput over the state of the art at equal recall. In the in-distribution (ID) setting, even a simple linear probing function trained with CwA outperforms competing deep neural methods.

Safe Policy Improvement (SPI) aims to improve a baseline policy offline using fixed data while avoiding performance degradation with high probability. Existing SPI methods work in discrete domains but do not extend to continuous control. We propose Calibrated Safe Policy Improvement (Cal-SPI), a Monte Carlo Tree Search-based continuous SPI method that replaces count-based support with calibrated model trust. Cal-SPI learns a deep ensemble dynamics model and calibrates ensemble disagreement on held-out data through a Wilks-style tolerance construction, turning raw uncertainty into a statistical certificate of one-step model error. This certificate defines a hard gate within model-based tree search. We provide a theoretical analysis showing that calibrated model trust and a baseline-relative switching criterion yield a conditional PAC-style improvement bound for gate-certified root decisions. Experiments across three MuJoCo continuous domains show that Cal-SPI avoids performance degradation in low-data regimes and improves over the baseline once calibrated model trust becomes reliable.

We formalize and enable the task of open tree decomposition, which segments an image into hierarchical trees of visual components with unconstrained granularity and flexibility. Specifically, we provide the foundation benchmark for this new paradigm with the following three key contributions. First, we overcome the prohibitively high cognitive and physical bottlenecks of manual annotation by developing a fully automated generation pipeline that synergizes the semantic reasoning of Large Vision-Language Models (LVLMs) with the precise geometric grounding of SAM 3. Second, leveraging this pipeline, we construct COCOTree, a massive-scale benchmark featuring over 21K images and 1.8M structural nodes. By embracing an open-vocabulary space of over 3.5K unique labels, it successfully captures the long-tail distribution of complex physical assemblies. Notably, rigorous human evaluation confirms our generated annotations demonstrate strong alignment with human structural judgment. Third, we establish a standardized evaluation protocol by proposing the Open Tree Quality (OTQ) metric, which jointly assesses mask precision, label accuracy, and structural consistency. We release our dataset and benchmark code at https://github.com/melonkick3090/COCOTree.

Reinforcement learning with verifiable rewards (RLVR) has become the dominant paradigm for post-training large language models (LLMs) on complex reasoning and multi-step agentic tasks. However, the standard scalar reward signal is sparse and uninformative: it conveys whether a trajectory succeeded but not which reasoning steps failed or how to correct them, severely limiting sample efficiency on long-horizon tasks where successful trajectories are rare. We propose \textbf{CARD} (\textbf{C}ritique-\textbf{A}ugmented \textbf{R}einforcement \textbf{D}istillation), a training framework that amplifies the learning signal from a \texttt{verify\_answer} tool within the GRPO objective via two complementary signals. First, a token-level self-distillation reweighting measures how much each token's generation was shaped by the expert critique: tokens that merely echo the critique text are masked, while critique-driven revision tokens are amplified. Second, a rejection-sampling distillation loss maximizes the likelihood of critique-free trajectories for responses where the critique provably improved the answer, directly accelerating the internalization of critique-guided behavior. Together, these two signals distill critique-augmented reasoning into the model's base policy without requiring additional rollouts or separate reward models. Experiments on four long-form deep search benchmarks and three multi-hop question-answering benchmarks demonstrate consistent improvements over the GRPO baseline. On the deep search benchmarks, CARD improves average performance by 40\% over the Qwen3-8B base model and by 6\% over GRPO; with Qwen3-14B, the gains are 46\% and 7\%, respectively. These results establish expert critique as an effective and scalable source of rich text supervision for reinforcement learning of long-horizon research agents. Furthermore, we show that \texttt{verify\_answer} can serve as a test-time self-revision scaffold: at inference, the tool returns only task rubric criteria---without expert feedback---prompting the model to reflect on its draft and self-revise, yielding consistent performance gains over the no-VA baseline.

Generative models are often required to serve many downstream conditional tasks. In the training-free setting, a pretrained unconditional model is reused as a prior and combined with a task-specific likelihood to sample from the corresponding posterior, without retraining the generative model. While such conditional samplers have been developed for diffusion and flow generative models, training-free conditional sampling for Schrödinger bridge (SB) generative models remains unexplored, despite the growing appeal of SBs as flexible stochastic transport models between general source and target distributions. We propose BridgeTwist, a training-free conditional sampler for pretrained SB generative models, formulated as a tempered twisted sequential Monte Carlo scheme. At each step, the sampler uses a closed-form plug-in twist surrogate, which combines the velocity and score fields exposed by SB pretraining into an explicit endpoint predictor. Since this surrogate is less reliable near the source, we use a time-varying temperature to temper the twist, downweighting its influence at early steps. Telescoping importance weights then cancel all intermediate twist factors, yielding an asymptotically exact conditional sampler without any auxiliary twist network or pilot rollouts. Empirically, BridgeTwist consistently outperforms training-free baselines on class-conditional sampling and inpainting on MNIST and CIFAR-10, and on text-to-image generation on CelebA-HQ.

This work proposes the Adaptive Federated Hamiltonian Monte Carlo (AFHMC) for Bayesian federated learning (FL) tasks. The prior work on HMC applications for FL achieves $O(\sqrt{d}/\epsilon)$ communication cost under log-convex distributions, already demonstrating its advantage over the Langevin dynamic-based counterpart. Based on a refined analysis, this paper further reveals the benefit of HMC for the FL regime by incorporating a control mechanism for local node shifts. We prove that AFHMC at most requires $O((d/\epsilon^2)^{1/3})$ communication cost up to a logarithmic term, which is significantly better than the existing results. Moreover, the proposed method adapts to client heterogeneity. Under low-heterogeneity settings (i.e., the target sampling precision $\epsilon^2$ is higher than the heterogeneity level), the communication rate of AFHMC can be as good as $O(\log(d/\epsilon^2))$, resembling the communication costs for FL optimization tools. To the best of our knowledge, this provides the first logarithmic communication complexity result for federated sampling. Our code is available at https://anonymous.4open.science/r/Adaptive-Federated-HMC.

LLM-based search agents are widely used for information-seeking tasks, but their reliance on external tool returns introduces a critical security risk: web content retrieved during execution is untrusted, exposing agents to prompt injection and goal hijacking. Prior work on search-agent safety primarily focuses on static web-content injection, but modern agents issue follow-up queries and cross-check competing sources, so a single injected page is often diluted or rejected. We show that the channel delivering search and page observations is a fragile security boundary: beyond exposing the agent to a single poisoned page, a mediated search interface can repeatedly steer how the agent gathers evidence and forms its final answer. Under a constrained tool-intermediary threat model, appending only one controlled result per query can substantially increase attack success when the evidence is coordinated across the agent's trajectory. We study this setting with a strategy-driven long-horizon attack system and introduce Authority-Chain Hijack (ACH), an expert-refined strategy that turns isolated search-result and page-content manipulations into a coherent evidence chain across seemingly corroborating sources. ACH achieves the highest Overall ASR among all baselines, reaching 56.0\%/85.0\% ASR/MaxN\,ASR on the SafeSearch benchmark. We further introduce Trace-Guided Strategy Evolution (TGSE), which automatically improves reusable attacker strategies from execution traces, replacing manual redesign with trace-driven refinement and further raising these to 71.4\%/95.0\%.

Conformal prediction (CP) provides a distribution-free framework for uncertainty quantification with finite-sample marginal coverage guarantees. Yet, missing values introduce distributional heterogeneity, where marginal coverage can hide systematic undercoverage for certain missingness patterns. At the same time, existing methods for conditional guarantees often fail when missing values are present, as they rely on geometric continuity of the covariate space, which is broken by missingness. Although imputation restores continuity, it obscures the missingness patterns and introduces task-dependent overhead. We develop a new CP framework that is \emph{adaptive to missingness patterns}, without imputing covariates. We rely on data augmentation by feature embedding that accounts for missingness. Our framework represents each input with missing values by a tree embedding together with its raw missingness mask, and then learns a regularized quantile score function. Varying the choices for the penalty function and the augmentation enables us to characterize three practically relevant validity levels: exact coordinate-wise validity, exact validity on observed masks, and approximate missingness-conditional guarantees for general mechanisms (MCAR, MAR, MNAR). Experiments on synthetic and real datasets support these guarantees and show that the proposed method reduces miscoverage risk while maintaining informative intervals, especially in non-MCAR settings. These results highlight the practical value of uncertainty quantification that adapts to missingness.

Monte-Carlo Tree Search (MCTS) typically stores only scalar mean estimates, which makes it poorly suited for tail-risk objectives such as conditional value-at-risk (CVaR). We introduce **Categorical Thompson Sampling with Optimism (CATSO)** and **Particle Thompson Sampling with Optimism (PATSO)**, two distributional MCTS algorithms that maintain finite-support empirical backup laws at Q-edges. CATSO represents Q-edge laws with categorical atoms and Dirichlet Thompson sampling, while PATSO represents them with adaptive particles. Both algorithms select actions using the CVaR of a Thompson-sampled Q-edge law plus a polynomial optimism bonus, and propagate scalar continuation values through visit-weighted averages of child Q-edge CVaR scores. For the finite-depth nested CVaR planning target, we prove an $O(n^{-1/2})$ bound on the root value error after $n$ rollouts, which upper-bounds the risk-sensitive simple regret. CATSO adds a fixed-grid discretization bias, while capped PATSO adds a tunable Wasserstein compression bias. Experiments on stochastic planning benchmarks show that distributional Q-edges improve lower-tail return and reduce catastrophic outcomes when mean-optimal and risk-sensitive behavior differ.

We introduce the Thresholding Monte Carlo Tree Search problem, in which, given a tree $\mathcal{T}$ and a threshold $\theta$, a player must answer whether the root node value of $\mathcal{T}$ is at least $\theta$ or not. In the given tree, 'MAX' or 'MIN' is labeled on each internal node, and the value of a 'MAX'-labeled ('MIN'-labeled) internal node is the maximum (minimum) of its child values. The value of a leaf node is the mean reward of an unknown distribution, from which the player can sample rewards. For this problem, we develop a $\delta$-correct sequential sampling algorithm based on the Track-and-Stop strategy that has asymptotically optimal sample complexity. We show that a ratio-based modification of the D-Tracking arm-pulling strategy leads to a substantial improvement in empirical sample complexity, as well as reducing the per-round computational cost from linear to logarithmic in the number of arms.

Neural architecture search (NAS) methods have grown increasingly efficient, yet they remain bounded by manually engineered search spaces that require substantial domain expertise and must be rebuilt for every new task. Large language models (LLMs) can generate architectures in an open-ended space, but how to optimally divide labor between LLM-driven design and NAS-driven search remains unexplored. We propose a mechanism that bridges these two paradigms: an LLM produces a high-quality seed architecture, then decomposes it into a *slotted architecture*---a scaffold with named, interchangeable module slots that automatically defines a bounded, task-specific search space for conventional NAS to explore, without manual engineering. We instantiate this mechanism in **AgentNAS**, a modular three-phase pipeline in which each component's contribution can be measured independently. On 17 tasks spanning classification, dense regression, segmentation, and multi-label tagging across diverse modalities (NAS-Bench-360 and NAS-Unseen), AgentNAS establishes a new state of the art on 11 tasks, outperforming published baselines including task-specific expert designs. Ablation studies show that the two search mechanisms are broadly complementary: the LLM-generated seed already surpasses published baselines on the majority of tasks, and NAS delivers additional gains in most cases through combinatorial recombination across slots---a mode of search that independent LLM samples cannot replicate. These patterns hold across three LLMs of different capability levels, confirming that the division of labor is robust.

Speculative decoding accelerates generative inference of large language models (LLMs) by using a small draft model to propose multiple candidate tokens, which are then verified in parallel with the target model in a single decoding iteration. While the state-of-the-art method of tree-based speculative decoding helps improve generation throughput over non-speculative inference, deploying it in LLM serving systems often yields suboptimal performance due to dynamically changing serving conditions. In particular, our analysis shows that the optimal tree configuration---the one that maximizes performance---varies with two key serving conditions: request rate and per-request characteristics. Based on the analysis, we present AdaTree, a plug-in component for LLM serving systems that dynamically adapts tree configurations to varying serving conditions. AdaTree predicts model execution time and acceptance length across different tree configurations and selects the one that maximizes speculative decoding efficiency. To capture the non-linear relationship between tree configuration and model execution time, AdaTree employs decision trees as its core modeling primitive. Our evaluation shows that AdaTree consistently outperforms both chain-based and static tree-based speculative decoding across diverse serving conditions.

Bayesian neural networks (BNNs) offer principled uncertainty quantification, yet sampling from their posteriors via Hamiltonian Monte Carlo (HMC) remains challenging. While recent work theoretically identified the efficiency degradation of ReLU-based networks, smooth unbounded activations (_e.g.,_ GELU, Swish, Mish), which are prevalent in modern architectures, were assumed to be efficient due to their smoothness. In this work, we challenge this existing justification by deriving explicit local error formulas for the leapfrog integrator. Our analyses reveal that, in addition to differentiability, third-order curvatures of the potential energy also have a significant influence on HMC efficiency of BNNs inference. Empirically, we found that even smooth unbounded activations may degrade sampling performance as severely as piecewise linear activations. Thus, optimally tuning the step size for these networks may require a comparable empirical step size scaling as their piecewise counterparts. This result sheds light on the implications for architectural design in Bayesian deep learning.

The public belief state (PBS), which is the posterior over game histories conditioned on public information, is a fundamental abstraction for designing game-theoretically sound search algorithms for imperfect-information games. Existing sound PBS search techniques for adversarial games fall into two categories: gadget-game-based algorithms (including Libratus, DeepStack, Pluribus, Supremus, and Student of Games) and ReBeL-based algorithms. Both these categories possess disadvantages, including a reliance on discontinuous functions, an inability to re-solve subgames, and an implicit representation of policies. In this work, we propose Gravel, a new approach for PBS search---different from these previous two categories---that achieves game-theoretic soundness via regularization. Unlike the aforementioned approaches, Gravel relies on smooth functions, can initiate search at any point in the game, and outputs explicit policies directly. We empirically demonstrate the advantages of Gravel for no-limit Texas hold'em, where we show that it outperforms ReBeL both in head-to-head competitions against Slumbot and in endgame solving. We include our codebase in the submission, making Gravel the first open-source high-performance no-limit Texas hold'em AI.

Backpropagation remains the dominant method for training deep neural networks, but its reliance on non-local computations has motivated the search for biologically plausible alternatives. Two prominent frameworks, predictive coding (PC) and deep feedback control (DFC), offer distinct approaches to local learning, yet their relationship remains unclear. Here, we show that DFC arises as a local approximation to predictive coding dynamics near fixed points, with exact equivalence in shallow linear networks and controlled deviations arising with depth, nonlinearity, and distance from equilibrium. Building on this connection, we introduce Approx PC, a practical algorithm that replaces layer-wise error propagation in PC with a centralized feedback controller operating over a sequence of intermediate subgoals. Empirically, Approx PC recovers solutions consistent with predictive coding, while converging faster and achieving improved accuracy in our experiments. Our results unify energy-based and control-theoretic perspectives on learning and suggest a computationally efficient approach to training predictive coding networks through control-inspired dynamics.