NeurIPS 2026 Papers

Tree-based Monte-Carlo Tree Search (MCTS) duplicates the same state when it is reached through different trajectories, which can waste simulations in stochastic MDPs. We introduce Graph-Based Stochastic-Power-UCT (GS-Power-UCT), which shares states reached at the same planning depth while keeping separate values for states reached at different depths. This design applies to general stochastic MDPs, including problems with cycles. We prove that for a fixed planning horizon, the root estimate converges to the finite-horizon value at rate $O(n^{-1/2})$, matching tree-based Stochastic-Power-UCT while reusing samples across shared states. We also study two full-state variants: GS-Power-UCT-F, which stores one node per physical state to increase sample sharing but may mix values from different remaining horizons, and GS-Power-UCT-F$^+$, which uses an adaptive horizon to control this bias. The latter converges to $V^{\star}(s_0)$, the optimal infinite-horizon discounted value at the root state $s_0$, when the remaining cross-depth gap vanishes. Experiments on stochastic planning benchmarks show improved sample efficiency over tree-based and graph-based baselines.

The recent breakthrough Score-Repellent Monte Carlo (SRMC) method improves MCMC sampling by tilting the target density $\pi$ in $\mathbb{R}^d$ with the factor $\exp(-\alpha\,\theta_n^\top s(x))$, where the score function $s(x) = \nabla_x \log \pi(x)$, and $\theta_n$ is the score average from past samples. However, a single scalar repellence strength $\alpha\ge0$ cannot handle both the steep and the flat directions in anisotropic energy landscapes at once. We introduce Geometry-Aware SRMC (GA-SRMC), which generalizes the Euclidean alignment to $\theta^\top M s(x)$ for any symmetric positive-definite matrix $M$. A stochastic-approximation central limit theorem under a trace constraint on $M$ identifies the trace-normalized inverse Fisher $M_{\rm opt}\propto S^{-1}$ as the unique minimizer of the worst-direction covariance bound, elevating the Fisher information $S=\text{Cov}_\pi(s,s)$ from a proposal-side preconditioner (e.g., FisherMALA, natural gradient) to the optimal repulsion-side matrix, with a single online estimate $\widehat S_n$ serving both roles at no additional cost. On an $8$-target simulation spanning the condition number $\kappa\in[1.3,100]$ and dimension $d\in[2,785]$, GA-SRMC dominates SRMC by $4.4\times$ on anisotropic Gaussian ($\kappa=100$) and $6.8\times$ on MNIST under the MALA baseline, by $18\text{ - }32\\%$ on the well-conditioned Bayesian logistic-regression posteriors, and by up to $+48\\%$ under FisherMALA; the advantage scales monotonically with the condition number $\kappa$.

Large Language Model-based Automated Algorithm Design (LLM-AAD) has shown promising results across diverse domains by coupling LLMs with iterative search. Some recent efforts have begun to fine-tune LLMs for algorithm design through reinforcement learning. However, these attempts still rely on outcome-level supervision, leaving intermediate ideas and plans without direct credit, even though they determine the conceptual direction of algorithm design. In this paper, we propose Algorithm Tree Policy Optimization (ATPO), which utilizes tree-structured rollouts to decouple the generation of conceptual ideas from their specific code implementations. By applying group-relative advantage estimation, ATPO assigns credit to both strategic ideas and algorithmic implementations, thereby optimizing the policy across multiple levels of abstraction. Integration with FunSearch, EoH, and OpenEvolve demonstrates improvements over their original counterparts. Notably, the learned policy transfers effectively to unseen problem instances and alternative search methods, indicating that it can improve the generalization of algorithm design policies.

Introducing new capabilities to frontier models has long been the goal for posttraining, which relies predominantly on supervised finetuning (SFT) and reinforcement learning (RL) to achieve this. Prevailing wisdom dictates that on-policy RL enables strong generalization on new tasks without losing existing capabilities, while SFT is prone to weak generalization and forgetting. At the same time, SFT enables learning from inherently off-policy expert data, whereas RL must rely on a model's ability to generate positive signal on a new task. In our work, we seek to bridge the strength of on-policy learning with the information signal of off-policy expert traces. To this end, we introduce a Markov chain Monte Carlo (MCMC) sampling algorithm that transforms off-policy traces into more on-policy ones given a reference model for finetuning. Across tasks like scientific skill acquisition, mathematical reasoning, and fact learning, we demonstrate that simple SFT on these boosted off-policy traces can generalize better and forget less than on-policy counterparts. More generally, our approach outlines a principled methodology for model-specific data refinement, suggesting broader utility as a plug-and-play component throughout the post-training pipeline.

Discovery of novel materials for physical memory is essential for superior memory systems enabling AI hardware scaling. Among them, ferroelectrics hold promise for highly scalable DRAM. Yet their discovery remains driven largely by human intuition and exhaustive database screening. We present Ferrogen, a generative pipeline for the targeted discovery of novel ferroelectric materials for memory by fine-tuning the diffusion-based crystal structure generator Mattergen on ferroelectric-relevant properties. To enable property-conditioned generation, we construct a ML-labeled dataset using a suite of fast machine learning estimators, including a novel two-stage ensemble polarization predictor, dramatically reducing reliance on high-throughput Density Field Theory (DFT) during training and screening. The base Mattergen model is fine-tuned via lightweight adapter layers conditioned on polarization, switching energy, and metal-probability embeddings, with classifier-free guidance enabling targeted generation of candidates with high polarization, low but finite switching barriers, and insulating character. Generated candidates are screened by the same ML estimators and validated through rigorous DFT calculations, including structural relaxation, band gap verification, and Berry phase polarization computation. Against the ML screens, Ferrogen achieves a roughly $80 \times$ improvement in search efficiency over exhaustive database search. The screened candidates are then finally validated using DFT and made publicly available as a database of theoretically confirmed novel ferroelectrics. The pipeline is readily extensible to other functional material classes, establishing generative models as a powerful paradigm for accelerating electronic materials discovery for high performance memory.

In quantum systems, physical quantities of interest can often be expressed as expectation values of a probability distribution, which allows their estimation by Markov Chain Monte Carlo (MCMC) methods. This approach, an instance of Quantum Monte Carlo (QMC), plays a central role in obtaining predictions of system properties relevant to high energy or condensed matter physics. The use of Hamiltonian dynamics to design a Markov kernel, known as Hybrid or Hamiltonian Monte Carlo (HMC), is widely used in quantum applications. In this context, it is standard to use an HMC kernel that is ``adjusted'' with the Metropolis-Hastings (MH) criterion, so that estimates of expectations converge to their true value in the limit of infinite sample size. However, recent work in computational statistics suggests that an \emph{unadjusted} HMC kernel can be significantly more efficient, while maintaining an asymptotic bias which is small relative to the error arising from the variance of the finite sample size. We adapt this approach to the QMC setting, focusing on a challenging and high-dimensional \emph{spin-fermion} model of a quantum phase transition. We find that unadjusted methods outperform adjusted HMC in terms of effective sample size per gradient call by around an order of magnitude, while retaining accurate results.

AI research agents accelerate ML research by automating hypothesis generation, experimentation, and empirical refinement. Existing agent strategies range from greedy hill-climbing to tree search and evolutionary optimization, yet which strategy choices drive performance remains unclear. Answering this question requires a benchmark that separates agent strategy (e.g., search topology) from execution infrastructure (e.g., code editor), so that performance differences are attributable to strategy rather than infrastructure, and that provides process-level metrics beyond final scores to analyze exploration behaviors. Existing benchmarks offer limited support. We propose FML-Bench, a benchmark of 18 fundamental ML research tasks across 10 domains that separates agent strategy from execution infrastructure and defines 12 process-level behavioral metrics. Evaluating six representative agents, we find that: (1) strategy complexity alone does not guarantee strong performance: a simple greedy hill-climber nearly matches the best-performing tree-search agent, both well above the remaining agents; (2) our analysis suggests this pattern relates to improvement opportunity structure: greedy search tends to be more effective when opportunities are dense, while tree-search and evolutionary strategies tend to be more effective when opportunities are sparse; an adaptive agent built on this insight switches to broader exploration upon detecting improvement stagnation and outperforms the other six agents, lending initial support to this observation; and (3) process-level analysis reveals that early convergence and directionally focused exploration are significantly associated with final performance, while solution diversity and compute cost are not. Our benchmark is available at: https://anonymous.4open.science/r/Anonymous-78B6.

Large language models (LLMs), especially reasoning models, generate extended chain-of-thought (CoT) reasoning that often contains explicit deliberation over future outcomes. Yet whether this deliberation constitutes genuine planning, how it is structured, and what aspects of it drive performance remain poorly understood. In this work, we introduce a new method to characterize LLM planning by extracting and quantifying search trees from reasoning traces during four-in-a-row gameplay. By fitting computational cognitive models on the extracted search trees, we characterize how plans are structured and how they influence move decisions. We find that LLMs' search is shallower than humans', and that performance is predicted by search breadth rather than depth. Most strikingly, although LLMs expand deep nodes in their traces, their move choices are best explained by a myopic model that ignores those nodes entirely. A causal intervention study in which we selectively prune CoT paragraphs further suggests that move selection is driven predominantly by shallow rather than deep search. These patterns contrast with human planning, where performance is driven primarily by deep search. Together, our findings reveal a key difference between LLM and human planning: while human expertise is driven by deeper search, LLMs do not act on deep lookahead. This dissociation offers targeted guidance for aligning LLM and human planning. More broadly, our framework provides a generalizable approach for interpreting the structure of LLM planning across strategic domains.

The asymptotic behaviour of Monte Carlo Exploring Starts (MCES) is a long-standing open question in reinforcement learning, even in the tabular setting. We investigated the convergence properties of tabular MCES by constructing examples in which the algorithm converges to suboptimal solutions. This paper presents new counterexamples for both initial-visit and first-visit MCES and gives a convergence-restoring modification for the initial-visit case. We show that stable suboptimal solutions may exist for initial-visit MCES with sample-average updates even when greedy actions are updated more often than non-greedy actions on average. However, by scaling learning rates inversely to update frequencies on a state-by-state basis, convergence to optimality is guaranteed. Unlike previous uniformisation methods, this modification is applicable to large-scale problems that require approximating the estimated value function. We then extend the example to show that sample-average first-visit MCES may also converge to suboptimal solutions. This largely settles a fundamental open problem and shows that exploring starts alone do not guarantee convergence to optimality. More broadly, these results highlight that convergence depends critically on the relative size and frequency of updates applied to different actions, making the choice of learning rates and the balance between exploration and exploitation central to the analysis of MCES and the implementation of scalable Monte Carlo control methods.

By far the most common way to estimate an expected loss in machine learning is to draw samples, compute the loss on each one, and take the empirical average. However, sampling is not necessarily optimal. Given an MLP at initialization, we show how to estimate its expected output over Gaussian inputs without running samples through the network at all. Instead, we produce approximate representations of the distributions of activations at each layer, leveraging tools such as cumulants and Hermite expansions. We show both theoretically and empirically that for sufficiently wide networks, our estimator achieves a target mean squared error using substantially fewer FLOPs than Monte Carlo sampling. We find moreover that our methods perform particularly well at estimating the probabilities of rare events, and additionally demonstrate how they can be used for model training. Together, these findings suggest a path to producing models with a greatly reduced probability of catastrophic tail risks.

Inference-time search (e.g., Tree of Thoughts, ToT) with large language models (LLMs) often concentrates on a small set of structurally or semantically similar trajectories, leaving alternatives underexplored—a failure mode we call \textit{reasoning basin collapse}. We introduce \textsc{BASIN}, a training-free selection method that penalizes repeated visits to the same reasoning basin, defined symbolically for arithmetic tasks or semantically via hypothesis clustering, thereby reallocating search across diverse reasoning strategies. Under matched inference budgets, \textsc{BASIN} improves over ToT by up to $+22$pp on Game of 24. To explain why this quality-agnostic penalty is effective, we introduce the \emph{redundancy gap} $\Delta$, which measures how differently search concentrates for correct versus incorrect predictions: standard ToT often operates near $\Delta \approx 0$, failing to distinguish correct from incorrect searches by their visit patterns, while \textsc{BASIN} consistently induces $\Delta > 0$, concentrating search on correct basins while dispersing incorrect ones. More broadly, \textsc{BASIN} suggests structure-aware selection as a simple and general approach to improving inference-time reasoning.

Combining the predictions of multiple models is one of the most effective strategies in machine learning. Thus, it is only natural to investigate the ensembling of language models (LMs). However, because LMs define distributions over strings, a proper ensemble requires aggregating predictions at the string level, which requires a computation that is generally intractable. In this work, we cast LM ensembling as an inference problem and propose the deployment of sequential Monte Carlo (SMC), an appropriate approximate inference scheme. Concretely, we introduce a unified framework for composing multiple LMs into ensemble distributions parameterized by a broad family of aggregation functions. To sample from these distributions, we introduce an SMC algorithm that operates in a shared character space, enabling ensembles of models with mismatching vocabularies and consistent sampling in the limit. We evaluate a range of ensembles across prompt and model combinations for various text generation tasks, finding that consensus-seeking strategies such as the product consistently outperform probability averaging, and that better posterior approximations can yield better ensemble performance.

Guiding pretrained flow-based generative models for conditional generation or to produce samples with desired target properties enables solving diverse tasks without retraining on paired data. We present ESS-Flow, a gradient-free method that leverages the Gaussian prior of the source distribution in flow-based models to perform Bayesian inference directly in the source space using Elliptical Slice Sampling. ESS-Flow only requires forward passes through the generative model and observation process, making it applicable even when gradients are unavailable, such as with non-differentiable simulation-based observations or quantization in the generation process. We demonstrate its effectiveness on scientific applications including designing materials with desired target properties and predicting protein structures from sparse inter-residue distance measurements.

Algorithmic recourse (AR) aims to provide a recourse action that alters an undesired prediction made by a machine learning model. While standard AR methods assume the model does not change over time, this assumption is often violated in practice due to distribution shifts or data updates, rendering the suggested actions invalid. To address this issue, recent studies have proposed several methods that provide actions robust to model changes. Most of them construct an uncertainty set of models by perturbing continuous model parameters, such as the weight vectors of neural networks, and then seek actions that remain valid for all models in the set. However, such an approach cannot be directly applied to tree ensembles because their model parameters include discrete tree structures, which cannot be perturbed in the same way as neural networks. In this paper, we introduce a new robust AR framework for tree ensembles through the lens of distributionally robust optimization (DRO). Our key idea is to express the model change in a tree ensemble as a perturbation to the distribution of predictions made by decision trees in the ensemble, rather than to the model parameters themselves. This formulation allows us to naturally cast the robust AR problem for tree ensembles as a DRO problem. We then show that, by choosing a specific discrepancy measure between distributions, our robust validity constraint can be reformulated into a tractable one that can be evaluated as efficiently as the standard validity constraint. Building on this reformulation, we propose an efficient optimization algorithm by extending the existing search-based method for tree ensembles. Experimental results demonstrated that our method achieved a better robustness-cost trade-off than existing robust methods applicable to tree ensembles, while maintaining comparable runtime to the standard method.

Constructing mechanistic models of neural circuits is a fundamental goal of neuroscience, yet verifying such models is limited by the lack of ground truth. To rigorously test model discovery, we establish an in silico testbed using neuromechanical simulations of a larval zebrafish as a transparent ground truth. We find that LLM-based tree search autonomously discovers predictive models that significantly outperform established forecasting baselines. However, in-distribution accuracy is a poor proxy for faithful system identification, as models exploit statistical shortcuts. Structural priors prove essential for enabling robust out-of-distribution generalization and recovery of interpretable mechanistic models. Our insights provide guidance for modeling real-world neural recordings and offer a broader template for AI-driven scientific discovery.

Diffusion-based foundation models for crystal generation have become competent priors over the manifold of plausible materials, but turning that prior into generation of promising candidates requires targeting stability while balancing several conflicting properties at once. Fine-tuning is expensive and produces a new model for every new objective, while differentiable guidance is not applicable to many real-world rewards. We instead adapt these models at \emph{inference time} using Diffusion Tree Search (DTS), a Monte Carlo tree search procedure over the denoising trajectory. We make two technical contributions on top of DTS that are particularly relevant for AI-for-science applications. First, we show how to use existing property-conditional foundation models (e.g., classifier-free guidance variants of MatterGen) as guided proposals inside DTS, with importance-corrected backups and selection that keep the unconditional base model as the target distribution. Second, we extend DTS to multi-objective sampling by maintaining per-objective soft value estimates and selecting expansions through a sampled scalarization, so a single search produces samples spanning a Pareto front while sharing computation across preferences. We instantiate both adaptations on de novo crystal generation: single-objective stability search with different underlying models, conditional generation of crystals belonging to rare space groups, and a multi-objective tasks using a pretrained conditional model. Across these, DTS-based inference-time adaptation improves stability rates and Pareto-front quality without modifying the base models.

We develop diffusion-based samplers for target distributions known up to a normalising constant. To this end, we rely on the well-known diffusion path that smoothly interpolates between a simple base distribution and the target, popularised by diffusion models. We tackle the score estimation problem by developing an efficient sequential Monte Carlo sampler that evolves auxiliary variables from conditional distributions along the path, providing principled score and density estimates for time-varying distributions. To control the variance of score estimates, we further propose practical control variate schedules that incur minimal overhead. We adapt this general framework to paths induced by the Ornstein–Uhlenbeck (OU) time-reversal process, stochastic interpolants, and diffusion annealed Langevin dynamics, outlining their trade-offs. Finally, we provide theoretical guarantees and empirically demonstrate the effectiveness of our method on several synthetic and real-world datasets.

Particle filters are a standard tool for nonlinear state estimation, but their resampling step is discrete, preventing gradient-based learning in variational sequential Monte Carlo. We introduce Differentiable Systematic Resampling (DSR), a temperature-controlled relaxation of systematic resampling, that preserves the low-variance structure of systematic resampling while enabling full gradient flow. DSR converges to exact systematic resampling as the temperature vanishes, and we prove exponential convergence rates for the induced bias. Compared to optimal-transport-based differentiable resampling, DSR avoids iterative solvers and has substantially lower computational overhead. Experiments on stochastic dynamical systems and real-world handwriting data show that DSR achieves comparable or superior filtering and dynamics learning performance.

The process of discovery requires active exploration---the act of collecting new and informative data. However, efficient autonomous exploration remains a major unsolved problem. The dominant paradigm addresses this challenge by using Reinforcement Learning (RL) to train agents with intrinsic motivation, maximizing a composite objective of extrinsic and intrinsic rewards. We suggest that this approach incurs unnecessary overhead: while policy optimization is necessary for precise task execution, employing such machinery solely to expand state coverage may be inefficient. In this paper, we propose a new approach that explicitly decouples exploration from policy optimization and bypasses RL entirely during the exploration phase. Our method uses a tree-search strategy inspired by the Go-With-The-Winner algorithm, paired with a measure of uncertainty to systematically drive exploration. By removing the overhead of policy optimization, our approach explores an order of magnitude more efficiently than standard intrinsic motivation baselines on hard exploration benchmarks. Further, we demonstrate that the trajectories discovered during exploration can be distilled into deployable policies using existing supervised backward learning algorithms, achieving state-of-the-art performance by a wide margin on Montezuma's Revenge, Pitfall!, and Venture without relying on domain-specific knowledge. Finally, we demonstrate the generality of our framework in high-dimensional continuous action spaces by solving the MuJoCo Adroit dexterous manipulation and AntMaze tasks in a \emph{sparse-reward} setting, directly from image observations and without expert demonstrations or offline datasets. To the best of our knowledge, this has not been achieved before.