NeurIPS 2026 Papers

Ensemble sampling is a practically attractive approach to randomized exploration, but existing theoretical guarantees for linear bandits require ensembles much larger than what its practical motivation would suggest. In particular, the sharpest existing analysis achieves the ensemble complexity of $\Theta(d\log T)$, leaving a $\log T$ gap from the intrinsic $\Omega(d)$ ensemble-size barrier. We aim to narrow this gap by proposing an ensemble sampling algorithm that refreshes the ensemble only when the regularized Gram matrix changes substantially. This mechanism localizes the perturbation analysis to epochs with controlled Gram-matrix drift and reduces the sufficient ensemble size to $\Theta(d\log d+d\log\log T)$, while preserving the state-of-the-art $\tilde{\mathcal O}(d^{3/2}\sqrt T)$ regret for ensemble sampling with arbitrary bounded arm sets. We further show that, when the arm set is finite of cardinality $K$, the proposed algorithm achieves the sharper regret bound $\tilde{\mathcal{O}}(d\sqrt{T\log K})$. To the best of our knowledge, this is the first ensemble-sampling guarantee that simultaneously recovers both canonical regret scalings known for randomized linear bandit algorithms: the $\tilde{\mathcal{O}}(d^{3/2}\sqrt{T})$ rate for arbitrary bounded arm sets and the $\tilde{\mathcal{O}}(d\sqrt{T\log K})$ rate for finite arm sets. The algorithm also admits an anytime implementation without resetting past data, and experiments show that it remains competitive with baselines while using substantially smaller ensembles.

Recent work has pointed out current models' limitations in sampling according to prespecified target distributions. In this work, we show that current models nevertheless possess a coarse form of control over their output distributions: by instructing a range of open-source models to maximize or minimize output certainty in a two-alternative forced-choice task with no correct answer, we find that several models can shift entropy in the instructed direction. Through contrastive activation analysis, we further identify a linear direction in activation space that models recruit under uncertainty-modulation instructions and that can be used to causally modulate entropy via activation steering in the absence of any such instructions, suggesting that deliberate entropy control has an identifiable internal representation.

We study the problem of efficient online reinforcement learning in the infinite horizon setting when there is an offline dataset to start with. We assume that the offline dataset is generated by an expert but with unknown level of competence, i.e., it is not perfect and not necessarily using the optimal policy. We show that if the learning agent models the behavioral policy (parameterized by a competence parameter) used by the expert, it can do substantially better in terms of minimizing cumulative regret, than if it doesn't do that. We establish an upper bound on regret of the exact informed PSRL (iPSRL) algorithm that scales as $\tilde{\mathcal{O}}(\sqrt{T})$. This requires a novel prior-dependent regret analysis of Bayesian online learning algorithms for the infinite horizon setting. We then propose the informed RLSVI (iRLSVI) algorithm to efficiently approximate the iPSRL algorithm. Empirical results show that small offline expert datasets are sufficient for iRLSVI to outperform online-only algorithms on hard exploration domains.

Speculative sampling is a widely used method for losslessly accelerating Large Language Model (LLM) inference. State-of-the-art speculative sampling methods (e.g. EAGLE series) use a ``dynamic draft tree'' with drafter temperature implicitly set to $T=0$. Greedy drafts, however, perform poorly as temperature increases. Moreover, we show that naively instantiating the dynamic draft-tree with drafter temperature $T>0$ leads to an incorrect distribution. Fixes exist, but are computationally prohibitive as they require $O(|\mathcal{V}|^k)$ computation, where $\mathcal{V}$ and $k$ denote vocab and number of child nodes. We address these limitations by introducing Gumbo, a new speculative sampling algorithm designed for high-temperature settings that uses only $O(|\mathcal{V}|)$ compute. Gumbo is drafter-invariant in that it does not require the draft distribution during verification. Additionally, Gumbo increases the average acceptance length and speedup by expanding and reranking the dynamic draft tree based on Gumbel scores rather than draft probabilities. At the heart of Gumbo is a multi-draft generalization of the communication-free coupling principles and Gumbel sampling developed by Daliri et al., which is of independent interest. Replacing EAGLE-3's speculative sampler with Gumbo results in up to a +26.7\% increase in speedup at a temperature of 1.0 across four models and five datasets, without altering the draft model or requiring any other code changes.

Quantization has become a standard tool for efficient LLM deployment, especially for local inference, where models are now routinely served at 2-3 bits per parameter. The state of the art is currently split into simple scalar quantization techniques, such as GPTQ or AWQ, which are widely deployed but plateau in accuracy at 3-4 bits per parameter (bpp), and "second-generation" vector- or trellis-quantized methods, such as QTIP, GPTVQ and AQLM, which push the accuracy frontier but are notoriously hard to implement and to scale. In this paper, we ask whether this gap is fundamental, or whether a carefully optimized $\textit{scalar}$ quantizer can recover most of it. We answer in the affirmative, by introducing GSQ (Gumbel-Softmax Quantization), a post-training scalar quantization method which jointly learns the per-coordinate grid assignments and the per-group scales using a Gumbel-Softmax relaxation of the discrete grid. GSQ matches the cardinality of the relaxation to the small number of levels available in the target bit-width regime (e.g., 3-8 levels for ternary and 3 bpp, respectively), making optimization tractable. Practically, on the standard Llama-3.1-8B/70B-Instruct models, GSQ closes most of the gap between scalar quantization and the QTIP frontier at 2 and 3 bits, while using a symmetric scalar grid with group-wise quantization, and thus remains compatible with existing scalar inference kernels. We further show that the same discrete-assignment optimization can be applied to practical GGUF K-Quant checkpoints: starting from publicly released GGUF models, GSQ improves accuracy while projecting the result back into the same deployment format. Finally, GSQ scales to trillion-scale Mixture-of-Experts models such as Kimi-K2.5, where vector-quantized methods are difficult to apply.

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.

Generative models such as diffusion and flow matching have become dominant paradigms for visuomotor policy learning, yet their reliance on iterative denoising incurs high inference latency incompatible with real-time robotic control. We present a Fast Legendre-polynomial Action policy via Sparse History-anchored flow (FLASH), which replaces discrete action-chunk generation with continuous Legendre polynomial trajectory representation. Specifically, by fitting expert demonstrations under sparse temporal sampling, FLASH enables a single inference to cover a significantly extended action horizon. To further accelerate generation, FLASH initiates the flow matching process from history polynomial coefficients rather than uninformative Gaussian noise, shortening the transport distance and enabling accurate single-step inference. Moreover, analytic polynomial differentiation directly provides desired velocity feed-forward signals to the torque controller without numerical approximation. Extensive experiments on five simulated and two real-world manipulation tasks demonstrate that FLASH achieves state-of-the-art success rates ($\ge 92\%$ across all tasks), a per-episode inference time of $31.40\,ms$ (up to $175\times$ faster than diffusion policies and $18\times$ faster than prior flow matching policies), up to $4\times$ faster training convergence than ACT, and $5\times$ to $7\times$ reduction in controller tracking error compared to discrete-action baselines.

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.

Learning the natural parameters $z \in \mathbb{R}^n$ of discrete distributions $\mu_z$ from independent samples constrained to a subset $S \subseteq \\{0,1\\}^n$ is a foundational challenge in high-dimensional statistics. Existing methods for efficiently estimating truncated Boolean product distributions, notably the work of [Fotakis et al' COLT'20, Algorithmica '22], require either strong local connectivity assumptions on $S$ -- a property denoted *fatness* -- or stringent anti-concentration assumptions and necessitate the total mass of the truncation set to be a constant with respect to $n$. Moreover, the results in [Fotakis et al' COLT'20, Algorithmica '22] suffer from sample complexities that scale as $\Omega(2^n)$ if the mass of $S$ is exponentially small in $n$. In this work, we circumvent these limitations by analyzing the geometry of $S$ under the measure $\mu_z$. We refine the existing parameter estimation guarantees under the fatness assumption, improving the prior sample complexity to $\mathcal{O}( \log n / \epsilon^2)$ for $\ell_\infty$-recovery, matching the untruncated minimax rate. We further generalize fatness using the notion of influence utilized in the analysis of Boolean functions and provide sufficient conditions for efficient inference. Notably, unlike previous work, our method does not require sampling at arbitrary parameterizations of the model and instead relies on gradient descent. Lastly, we establish a theoretical lower bound demonstrating the sample complexity exhibits an intrinsic exponential dependence on the width of the model and the minimum distance between elements in the set.

*Minority collapse*, where minority classes become indistinguishable, is a significant challenge in imbalanced learning. This challenge is addressed by methods such as Mixup with class-balanced sampling. Although minority collapse has been mathematically analyzed using the layer-peeled model alongside Neural Collapse, no prior work has analyzed minority collapse under Mixup, particularly from the perspective of mixed labels. We investigate this overlooked factor and raise the question: *Is mixed label balance important for alleviating minority collapse?* Our analysis reveals that (i) mixed labels should be balanced, and (ii) in this setting, interpreting mixed labels as singletons is beneficial. Motivated by this analysis, we propose a *Balanced Mixed Label Sampler* and a *Mixed-Singleton classifier*, which balance mixed labels and treat them as singleton labels. Through theoretical analysis and experimental results, we highlight the importance of balancing mixed labels in imbalanced learning.

We introduce a novel technique for scalable sampling of spin-system states with continuous symmetries using diffusion models. By applying our approach to the XY model, a fundamental continuous-spin model in condensed matter physics, we show that our technique addresses the shortfalls of the Markov chain Monte Carlo (MCMC) in generalization to varying system sizes. More specifically, we show that training a temperature-conditioned diffusion model on smaller-size XY model lattices enables the generation of accurate samples in larger lattice sizes. By tracking physically important observables of the model, such as spin correlations, our experiments demonstrate that diffusion sampling followed by few MCMC steps reduces the thermalization time by an order of magnitude relative to the standard MCMC with random initialization. Our study provides valuable insight as to how generative models can be used to study continuous-state condensed matter systems at scale.

Scientific measurements inherently contain variance. Standard regression, however, treats noisy data as absolute ground truth, leading models to memorize noise rather than physical laws. This overfitting hampers generalization, particularly when transferring from abundant theoretical proxies to scarce experiments. We introduce DETS, a framework replacing rigid point estimation with Physical Tolerance Modeling. Our Interval-Censored Evidential Engine (ICEE) maximizes probability mass within acceptable error margins, explicitly decoupling aleatoric noise from epistemic uncertainty. Using this filtered uncertainty signal, a Thermodynamic Sampling strategy dynamically selects theoretical data to align source and target domains. Experiments across thermodynamics, drug affinity, and bandgap prediction show DETS outperforms state-of-the-art methods. Crucially, it exhibits superior robustness in high-noise, data-scarce settings.

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.

Diffusion models achieve state-of-the-art image synthesis, with their generative trajectories fundamentally exhibit a spectral bias, resolving low-frequency global structures early and high-frequency fine details later. Conventional stochastic differential equation (SDE) solvers fail to account for this dynamic, naively injecting uniform white noise throughout the entire process and misusing the finite energy budget. In this work, we establish a mathematical framework that reconsiders SDE inference as a targeted, frequency-decoupled energy transfer. Leveraging this framework, we introduce Colored Noise Sampling (CNS), a novel, training-free stochastic solver. Rather than injecting uniform white noise, CNS utilizes a dynamic, timestep- and frequency-dependent schedule that more-efficiently allocates injected energy toward structurally unresolved frequency bands. By actively exploiting the model's inherent spectral bias, CNS systematically steers the generated distribution toward the true data manifold. Extensive experiments demonstrate that CNS significantly outperforms standard ODE and SDE baselines as a strictly plug-and-play, inference-time sampler substitution across diverse architectures (SiT, JiT, FLUX). Compared to standard sampling on ImageNet-256, CNS achieves substantial unguided FID reductions, improving from 8.26 to 6.27 on SiT-XL/2, 32.39 to 26.69 on JiT-B/16, and 11.88 to 8.31 on JiT-H/16, while yielding consistent relative FID improvements with Classifier-Free Guidance.

Autoregressive language models (ARMs) have been shown to memorize and occasionally reproduce training data verbatim, raising concerns about privacy and copyright liability. Diffusion language models (DLMs) have recently emerged as a competitive alternative, yet their memorization behavior remains largely unexplored due to fundamental differences in generation dynamics. To address this gap, we present a systematic theoretical and empirical characterization of memorization in DLMs. We propose a generalized probabilistic extraction framework that unifies prefix-conditioned decoding and diffusion-based generation under arbitrary masking patterns and stochastic sampling trajectories. Theorem 4.3 establishes a monotonic relationship between sampling resolution and memorization: increasing resolution strictly increases the probability of exact training data extraction, implying that autoregressive decoding corresponds to a limiting case of diffusion-based generation by setting the sampling resolution maximal. Extensive experiments across model scales and sampling strategies validate our theoretical predictions. Under aligned prefix-conditioned evaluations, we further demonstrate that DLMs exhibit substantially lower memorization-based leakage of personally identifiable information (PII) compared to ARMs.

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.

We introduce a new algorithm to sample from a probability density $\pi \propto e^{-V}$ with $V(x): \mathbb{R}^d \rightarrow \mathbb{R}$ a strongly multi-modal potential. Such distributions are challenging to sample from using standard gradient based methods, as exploration typically relies on inefficient diffusions. Inspired by connections between Schrodinger operators, Langevin diffusions and Morse theory, we develop a particle based method that exploits fast relaxation to transition pathways of the so-called Witten PDE on 1-forms (vector fields). Once pathways joining modes are sufficiently well explored, the method relaxes to inexpensive Langevin dynamics for sampling within approximately convex basins. Numerical experiments demonstrate superior performance over cheap gradient based sampling methods and competitive performance at lower cost compared to a state of the art method, parallel tempering.

Gaussian process (GP) bandits provide a powerful framework for performing blackbox optimization of unknown functions. The characteristics of the unknown function depend heavily on the assumed GP prior. Most work in the literature assume that this prior is known but in practice this seldom holds. Instead, practitioners often rely on maximum likelihood estimation to select the hyperparameters of the prior - which lacks theoretical guarantees. In this work, we study two algorithms for joint prior selection and regret minimization in GP bandits based on GP Thompson sampling (GP-TS): Prior-Elimination GP-TS (PE-GP-TS) that disqualifies priors with poor predictive performance, and HyperPrior GP-TS (HP-GP-TS) that utilizes a bi-level Thompson sampling scheme. We theoretically analyze the algorithms and establish a sublinear regret bound for HP-GP-TS. In addition, we demonstrate the effectiveness of these algorithms compared to the alternatives through extensive experiments with synthetic and real-world data.