NeurIPS 2026 Papers

Particle-based variational inference (ParVI) methods approximate an intractable target distribution by evolving an ensemble of interacting samples. Existing approaches rely predominantly on kernel-based repulsion (e.g., SVGD), which suffers from variance collapse in high dimensions and mode collapse on multimodal targets—pathologies caused by the absence of global transport structure. We introduce entropic transport descent (ETD), a ParVI family that frames each particle update as an entropy-regularized optimal transport problem. Derived from the JKO proximal scheme by lifting to the space of couplings and relaxing via the KL chain rule, each ETD iteration reduces to a Sinkhorn computation. The resulting transport plan provides global coordination, guiding each particle to nearby high-density proposals and naturally preserving multimodal structure. ETD can operate entirely score-free, requiring only pointwise evaluations of the unnormalized target density. Experiments on variance-collapse diagnostics, Bayesian logistic regression, neural networks, and molecular Boltzmann distributions show that ETD matches or outperforms SVGD, AGF-SVGD, and SGLD, with the largest gains in high-dimensional and multimodal settings.

Pretrained diffusion models increasingly serve as frozen teachers feeding downstream pipelines such as text-to-3D, single-step distillation, and data attribution. The teacher gradients these pipelines consume are Monte Carlo expectations over noise levels and Gaussian noise; their estimator variance dominates compute cost because each draw requires expensive upstream work, such as rendering, simulation, or encoding. We introduce CARV, a compute-aware variance-accounting framework that motivates a hierarchical Monte Carlo estimator: amortize the expensive upstream computation over cheap diffusion-noise resamples, sharpened by timestep importance sampling and a stratified inverse-CDF construction. Across diffusion-guided workloads, we obtain 2–3× effective compute multipliers, most from amortized reuse and approximately 25% additional gain from importance sampling plus stratification, without changing the objective. We also map regimes where these gains translate into improved downstream metrics versus regimes where they do not, such as DMD.

We study feature subsampling in greedy model search, a randomization mechanism motivated by random forests and other ensemble methods. While recent theory suggests that this randomization acts solely as a variance reduction mechanism analogous to ridge regularization, these results largely rely on base learners optimized via ordinary least squares (OLS). We investigate the effects of feature subsampling on greedy forward selection, a tractable abstraction of the adaptive split search used by decision trees. Assuming an orthogonal design, we prove that ensembling with feature subsampling can reduce both bias and variance, contrasting with the pure variance reduction of convex base learners. Specifically, we show that both the training error and degrees of freedom need not be monotone in the subsampling rate, breaking the analogy with standard shrinkage methods like the lasso or ridge regression. Furthermore, we characterize the exact asymptotic behavior of the estimator, showing that it adaptively reweights OLS coefficients based on their rank, with weights that are well approximated by a logistic function. These are mechanism level results for a tractable analogue of greedy selection that depends on the response, not performance guarantees for full random forests.

We study fixed-confidence best-action identification (BAI) in stochastic minimax trees. This problem is increasingly relevant in modern AI planning, where deep minimax search and Monte Carlo Tree Search (MCTS) with language model long rollouts face a fundamental tradeoff: heuristic evaluations are cheap but biased, while accurate rollouts are reliable but prohibitively expensive. We propose 2FFS, a two-fidelity tree-search algorithm that brings multi-fidelity flat bandit ideas into trees. The algorithm combines minimax-style fast expansion with MCTS-style stochastic sampling, adaptively deciding when to exploit cheap biased evaluations and when to invoke expensive accurate evaluations for local certification. We prove fixed-confidence correctness, establish finite stopping for exact identification, and give a polynomial-depth cost upper bound for general-depth trees. Across numerical stochastic-tree experiments, 2FFS uses substantially fewer samples and computational operations comparing to existing BAI-MCTS baseline.

Step-size selection remains a central challenge in large-scale neural network optimization: conservative steps slow convergence, while aggressive steps can destabilize training. We propose Zero-and-First-Order optimization (ZFO), a lightweight framework that decouples direction selection from step-size selection. ZFO uses a trusted first-order optimizer to determine the direction and performs zeroth-order evaluations only along this one-dimensional subspace to choose how far to move. Using the current forward-backward pass and two additional fixed-batch function evaluations, ZFO constructs a local model of the objective along the proposed direction and selects a curvature-aware step within a bounded search interval. This yields an adaptive step-selection mechanism with lower cost than full line search. We provide theoretical guarantees showing that shared-sample evaluations produce reliable finite-difference curvature estimates, that the induced local model selects a near-optimal step along the search interval, and that ZFO converges to a neighborhood of a stationary point. Across multiple language models and datasets, ZFO improves accuracy and stability over standard first-order fine-tuning algorithms.

Many prediction tasks evaluate a tree ensemble on row groups that share feature values: discrete-time survival models expand each patient into $G$ time steps, click-through-rate models score each item in a search session, and scenario analyses vary a few inputs while keeping others fixed. Standard tree-ensemble inference treats each row independently, ignoring this input-side correlation. We instantiate partial evaluation for grouped tree-ensemble inference: constant features are static, varying features dynamic. The algorithm walks each tree once per group, partitions a row bitmask at varying splits, and skips empty subtrees via an unsplit shortcut. We prove a structural work decomposition: per-tree work splits into the constant-projected subtree size $|T_c|$, $G$ leaf writes, and a predicate-mask provisioning cost $Q$. For the trace evaluator, per-row work approaches $(d_v+1)/(d+1)$ as $G \to \infty$, giving asymptotic speedup $1/(1-\bar\rho_G)$. Empirically, TreeWalker delivers ${\sim}3\times$ algorithmic speedup over a row-independent traversal at the reference model size ($T{=}500$, $L{=}8$; $G{=}16$ for SUPPORT/FLCHAIN and empirical groups for Expedia), rising to $6.8$--$7.8\times$ at $G{=}128$ on the survival datasets. For f64 models, outputs match treelite GTIL up to tree-ordering roundoff; for f32 models, TreeWalker's f64 leaf accumulation is close to a Kahan-compensated reference than native f32 on $99.98\%$ of rows, with the remaining $0.02\%$ tied.

To lower the expertise barrier in machine learning, the AutoML community has focused on the CASH problem, which jointly automates algorithm selection and hyperparameter tuning. While traditional methods like Bayesian Optimization (BO) struggle with cold-start issues, Large Language Models (LLMs) can mitigate these through semantic priors. However, existing LLM-based optimizers generalize poorly to high-dimensional, structured CASH spaces. In this paper, we propose LB-MCTS, a trajectory-structured optimization framework that uses a Monte Carlo Tree Search tree as a shared state for algorithm selection, hyperparameter refinement, and BO–LLM proposer synergy. Within this shared state, BO provides algorithm-specific surrogate modeling for quantitative search, while the LLM exploits path-aware selective memory to generate semantic proposals and reflections. As the surrogate model improves, a reliability-aware proposer policy adaptively shifts from LLM-driven to BO-driven proposals within a unified search trajectory. Experiments on 104 AMLB datasets demonstrate that LB-MCTS consistently outperforms BO-based, LLM-based, and hybrid baselines.

Tree-sliced distances have emerged as scalable alternatives to Sliced Wasserstein distances by projecting probability measures onto tree metric spaces and exploiting closed-form transport on trees. Existing tree-sliced constructions, however, are largely built around first-order Wasserstein geometry. While this choice yields efficient computation, it fixes the discrepancy to an $L^1$-type aggregation of subtree mass imbalances and provides limited control over the geometry used in downstream optimization. We propose Tree-Sliced Orlicz IPM (TS-Orlicz), a tree-sliced framework that replaces the tree-level $W_1$ discrepancy with an Orlicz integral probability metric. Building on the tractable formulation of Orlicz IPMs on trees, TS-Orlicz aggregates Orlicz-induced discrepancies over random tree systems while preserving the scalability of tree-sliced computation. By varying the underlying Orlicz function, the framework recovers $L^p$-type behavior and also supports non-polynomial, tail-sensitive geometries that emphasize large subtree discrepancies. We show that TS-Orlicz preserves key properties of tree-sliced constructions and admits efficient computation via closed-form special cases and one-dimensional scalar optimization. We further extend the framework to probability measures on hyperspheres. Experiments across Euclidean and spherical settings, including gradient flows, self-supervised learning, and diffusion-based generative modeling, show that Orlicz-induced tree-sliced discrepancies achieve competitive or improved performance over sliced and tree-sliced baselines while maintaining low computational cost.

We study $N$-armed stochastic dueling bandits under the Condorcet-winner assumption, where three widely adopted objectives are considered: best-arm identification (BAI), weak regret, and strong regret. We propose *Tree-Guided Identify-Then-Exploit (TG-ITE)*, the first unified framework to tackle all these objectives to our knowledge. Without requiring stronger assumptions, we propose a shared tree-guided identification approach to find a high-confidence incumbent within $O(N)$ comparisons. We further propose varied exploitation strategies to utilize this warm-start stage to optimize the specific objectives at hand. This methodology enables our approach to (1) achieve $O(N)$ sample complexity in BAI without commonly adopted stronger assumptions; (2) build the first winner-stays-style algorithm to achieve $O(N)$ weak regret; (3) enjoy the same $O(N \log T)$ guarantee as specialized strong-regret approaches; (4) realize the joint optimization of BAI and weak regret with $O(N)$ guarantees for both, eliminating the sub-optimal gap of $O(\log N)$ in the existing approach. Our results provide evidence that the trade-off between BAI and regret minimization is relatively benign in dueling bandits.

It is widely recognized that in many applications, strong predictive performance is not enough to warrant use of a learned model. Additionally, it is required to ascertain whether the model conforms to certain properties such as fairness, robustness or safety. Unfortunately, most learning approaches are not guaranteed to produce a model that adheres to these properties. Model repair addresses this through a post-hoc debugging loop. Approaches such as verification are used to check if the desired property holds for a model; if not they return a counterexample which is used to modify or repair the ML model. We propose a novel repair method for tree ensembles. Our approach is based on a MILP encoding where the objective is to modify the leaf values in the tree ensemble such that it fixes the counterexamples while at the same time preserves predictive performance. The use of constrained optimization over the model parameters is compelling as it is guaranteed to fix a given set of counterexamples. We implement and evaluate our method in two use-case scenarios, classification requiring fairness and action policies in MDPs requiring safety. Results show that our MILP repair of tree ensembles can be feasible computationally, and is often successful in improving the desired property while not degrading performance.

Whenever LLM training involves multiple outputs from the same input---agentic multi-turn trajectories, self-consistency reasoning, beam search distillation, or multi-rollout SFT/RL---the tokens form a tree-structured trajectory with shared prefixes. Existing pipelines linearize such data and treat each branch independently, causing substantial redundant computation that scales with the number of samples. We derive that averaging the loss over all branches is algebraically identical to a per-token weighted loss, reducing the problem to computing each token's log-probability exactly once. We propose DFS serialization of the tree, which visits every token exactly once, and adapt full-attention and SSM layers to ensure the resulting log-probabilities match independent per-branch computation exactly. For memory-constrained settings where the full tree exceeds GPU capacity, we propose Redundancy-Free Tree Partitioning, which achieves zero redundant computation with peak memory bounded by a single root-to-leaf path. Together, these contributions form Tree Training, achieving up to 6.2× end-to-end training speedup on dense and MoE models for both supervised fine-tuning and reinforcement learning.

Algorithms with predictions, or learning-augmented algorithms, has proved to be an extremely useful paradigm for combining machine learning with traditional algorithms. One of the textbook settings for this is searching a sorted array. Without a prediction classical binary search takes $O(\log n)$ queries, while with a prediction we can use "doubling binary search" to find the target key using $O(\log \eta)$ queries, where $\eta$ is the error of the prediction measured as the absolute value of the difference between the true location and the predicted location. Since an array is just a path graph, in this paper we ask whether similar bounds can be achieved for search on even slightly more general graphs: trees. We show first that the high-level answer is "no": there is no search algorithm that uses $O(\log \eta)$ queries, where now $\eta$ is the graph distance between the predicted location and the true location. However, as our main result, we show that such bounds can be achieved on trees which are "path-like" in that they have low pathwidth. In particular, we prove that there is a search algorithm which uses at most $O(k \log \eta)$ queries, where $k$ is the pathwidth of the tree. We also prove a lower bound showing that our algorithm has existentially optimal query complexity. Finally, we show experimentally, on real-life inputs, that our algorithm has query complexity which is notably better than the simple non-prediction based algorithm.

Large language models rely on positional representations for long context modeling, yet are typically trained on short sequences and evaluated at much longer lengths. Bridging this gap requires positional encodings remaining distinguishable at large distances while generalizing beyond the training length. In this paper, we identify a structural limitation of existing rotary-based positional encodings: all positional dimensions are controlled by a single global index, which couples training-time positional variation with long-range extrapolation. To address this, we propose Tree Rotary Positional Encoding (TPE), a structured multi-scale positional encoding designed for training from scratch that represents absolute positions through com positions of different components, avoiding single-index extrapolation. We further devise an Offset-based Positional Training (OPT) strategy to improve learning interactions across TPE components, enabling better generalization to longer con texts. We theoretically analyze that its multi-scale decomposition enables effective combinatorial position coverage. Experimental results show strong gains under extreme length extrapolation: with a 1.2B model trained from scratch on length 512, our TPE achieves a 64× extrapolation to 32K tokens, reducing perplexity from 4532.21 (RoPE) to 149.96, with comparable inference-time efficiency.

Large Language Models (LLMs) have achieved remarkable progress on complex reasoning tasks, yet the mechanisms by which they acquire reasoning capabilities and the efficiency of their reasoning remain poorly understood. In this paper, we investigate the path-finding problem over directed graphs—a symbolic abstraction of multi-step reasoning. We first characterize the architectural requirements for this task, proving that one-layer transformers require $\Omega(N^2)$ size on dense graphs to implement the key DFS child-selection primitive. While a two-layer transformer implements full DFS with $O(N\log(N))$ size. Thus, a second layer is necessary and sufficient for near-linear DFS reasoning. We then extend this expressivity analysis to deeper models, proving the existence of an $L$-layer transformer that performs DFS with a look-ahead horizon of $2^{L-3}$, establishing that increasing depth yields an exponential gain in search efficiency. Finally, going beyond the expressivity, given appropriately curated training data, we show that two-layer transformers can provably learn to execute DFS via gradient flow and generalize to unseen graphs. Together, our results provide a mechanistic account of how transformers implement graph search, how depth improves reasoning efficiency, and how such reasoning algorithms can emerge through training.

We analyze the query complexity of finding a local minimum in $t$ rounds on general graphs. More precisely, given a graph $G = (V,E)$ and oracle access to an unknown function $f : V \to \mathbb{R}$, the goal is to find a local minimum---a vertex $v$ such that $f(v) \leq f(u)$ for all $(u,v) \in E$---using at most $t$ rounds of interaction with the oracle. The query complexity is well understood on grids, but much less is known beyond. This abstract problem captures many optimization tasks, such as finding a local minimum of a loss function during neural network training. For each graph with $n \geq 2$ vertices and constant integer $t \geq 2$, we prove a deterministic upper bound of $O(n^{1/t} (s\Delta)^{1-1/t})$, where $s$ is the separation number and $\Delta$ is the maximum degree of the graph. We complement this result with a randomized lower bound of $\Omega(n^{1/t})$ that holds for any connected graph. We also find that parallel steepest descent with a warm start provides improved bounds for graphs with high separation number and bounded degree. To obtain our results, we utilized an advanced version of Gemini at various stages of our research. We discuss our experience in the Methodology section.

Language models are increasingly deployed alongside test-time search. We argue that we should therefore shift the role of RL post-training from converging on a single best response to producing a diverse pool of competent candidates. Standard policy gradient methods such as GRPO optimize a fixed scalar reward and drive the policy toward near-duplicate responses, erasing the diversity search needs. We propose Vector Policy Optimization (VPO), which exploits the fact that practical rewards are vector-valued, e.g., per-test-case correctness, per-criterion ratings, or per-hop credit. Instead of collapsing these into one scalarization, VPO samples randomized scalarizations from a distribution over the reward simplex, incentivizing candidates to specialize to different trade-offs along the Pareto front. The model emits multiple candidates per prompt, with each conditioning on the previous ones, allowing for directed diversification. VPO is a drop-in replacement for the GRPO advantage estimator. Across four domains, VPO consistently improves test-time best@$k$ over scalar baselines, with gains widening as the test-time budget grows.

Log-likelihood is a standard metric for evaluating generative models. Unfortunately, in contrast to autoregressive models (ARMs), discrete diffusion models generally do not admit exact computation of this quantity. Existing evaluations, therefore, rely on the evidence lower bound (ELBO), leaving unclear how much higher the true value may be. We address this by introducing the **Tangent Upper Bound on Evidence** (**TUBE**), a variational upper bound on log-likelihood that admits an unbiased Monte Carlo estimator. Our TUBE extends across latent-variable models, including masked diffusion models (MDMs), any-order ARMs (AO-ARMs), and block variants of both. Applied to block MDMs and block AO-ARMs, TUBE reveals our key empirical finding that these models lie strictly below the exact ARM baseline, showing that ARMs still dominate in likelihood.

We study stochastic random search for nonconvex optimization with only noisy function evaluations. We propose Mi2P, a two-point variant of stochastic three-point methods, and analyze it under generalized $(L_0,L_1)$-smoothness. Two regimes are considered. When the mean objective $f$ is $(L_0,L_1)$-smooth and component variance is bounded, Mi2P achieves $\widetilde{O}(d^3\sigma_0^2/\varepsilon^{6})$, matching the rate of \citet{boucherouite2024mistp} under a strictly weaker assumption (they require $f$ to be $L$-smooth, i.e., $L_1=0$). When an $L^1$-average $(L_0,L_1)$-smoothness holds across $\xi$, Mi2P attains the optimal $\widetilde{O}(d\sigma^2/\varepsilon^{4})$ rate; the per-sample $L$-smoothness used by gradient-estimation methods such as RSGF \citep{ghadimi2013stochastic} (i.e., each $f_\xi$ is $L$-smooth, $L_1 = 0$ pointwise in $\xi$) is a special case of this condition. For finite sums with $G$-Lipschitz components, a translation-symmetric variance-reduction scheme yields $\widetilde{O}(\min\{d^{4/3}n^{2/3}/\varepsilon^{8/3},\, dn/\varepsilon^2\})$ without storing snapshots. We also extend the framework to $\delta$-inexact comparison feedback (e.g., RLHF, A/B testing), giving an $O(\sqrt{d\delta})$ accuracy floor. Our analysis is presented for directional distributions with bounded support; the same rates hold for any rotation-invariant distribution with sub-Gaussian tails on the norm of $s$ (including Gaussian) under a step-size restriction that does not affect the asymptotic complexity.

Recovering executable CAD programs from 3D meshes is fundamentally challenging due to the long-horizon, compositional nature of CAD construction and the need for precise estimation of both discrete operations and continuous parameters. Existing learning-based methods primarily treat mesh-to-CAD reconstruction as one-shot sequence prediction and are largely limited to simple sketch-extrude pipelines, restricting operation diversity and preventing the use of intermediate geometric feedback, with no mechanism for refining generated programs. To address this limitation, we introduce StepCAD, a generative optimization approach that performs step-by-step, geometry-conditioned program synthesis followed by explicit refinement in program space. Given an input mesh, StepCAD first predicts construction sequences using a state-conditioned CAD policy conditioned on both target and intermediate geometry, then refines them through IoU-guided tree search over local program edits. To address the limited operation diversity in prior work, we further introduce ARCADE-1.5M, a large-scale dataset of 1.5M executable CAD programs with diverse operations, long-horizon sequences up to 150+ operations, and 12.5M intermediate state-action transitions. Experiments on multiple CAD reconstruction benchmarks show that StepCAD achieves state-of-the-art reconstruction accuracy, validity, and robustness, with up to 87.2\% relative IoU improvement over the best prior method and increasingly larger gains as shape complexity increases.