Neural Flow Samplers with Shortcut Models
Wuhao Chen, Zijing Ou, Yingzhen Li
TL;DR
This paper introduces a novel neural flow sampler with an improved estimator and a shortcut consistency model, significantly enhancing sampling efficiency and accuracy for complex density estimation tasks.
Contribution
It presents a new estimator for intractable terms in neural flow samplers and a shortcut model to reduce sampling steps, advancing the state-of-the-art in neural density sampling.
Findings
Outperforms existing flow-based neural samplers on synthetic datasets.
Achieves better accuracy and efficiency in complex n-body system simulations.
Demonstrates robustness across diverse sampling scenarios.
Abstract
Sampling from unnormalized densities presents a fundamental challenge with wide-ranging applications, from posterior inference to molecular dynamics simulations. Continuous flow-based neural samplers offer a promising approach, learning a velocity field that satisfies key principles of marginal density evolution (e.g., the continuity equation) to generate samples. However, this learning procedure requires accurate estimation of intractable terms linked to the computationally challenging partition function, for which existing estimators often suffer from high variance or low accuracy. To overcome this, we introduce an improved estimator for these challenging quantities, employing a velocity-driven Sequential Monte Carlo method enhanced with control variates. Furthermore, we introduce a shortcut consistency model to boost the runtime efficiency of the flow-based neural sampler by…
Peer Reviews
Decision·ICLR 2026 Conference Withdrawn Submission
The method is principled in nature, and the few-step consistency approach is practically useful for improving sampler inference efficiency, especially when attempting to scale up to larger systems.
The performance of the method on energy-W2 appears quite variable across all tested potentials, particularly in the few-step regime (8 steps). To strengthen the evaluation, the method could be benchmarked on a larger and more complex system like LJ55, which is commonly used for comparison in prior works (e.g., iDEM and other methods use this as a baseline). It would also be valuable to compare the approach against newer state-of-the-art methods such as adjoint sampling, which has recently been s
1. Adapts the shortcut-modeling idea to a flow/SMC-based sampling setup, showing it can be used to reduce the number of transport steps. 2. Achieves consistently strong results on the reported toy and small N-body benchmarks (DW-4, LJ-13, GMM-40, MW-32). 3. Empirically shows that the proposed estimator yields lower-variance estimates of $\partial_t \log Z_t(x)$ during training, which is important for stabilizing PINN-style objectives.
1. The paper is hard to follow and poorly structured. It reads like a mix of several ideas and lacks a cohesive structure. The title focuses on shortcut models, but the first thing introduced is variance reduction for $\partial_t \log Z_t$, rather than the shortcut method, which seems to be the main topic according to the title. Reordering the sections would help where the method should be introduced first and then talk about improving stability by having a lower variance estimator of $\partial
1. The paper proposes novel improvements for a class recently proposed samplers that rely on ideas coming from generative modelling. 2. An ablation study of the proposed improvements is presented and encouraging although each test is only performed in one experiment on one system.
3. There is no discussion on the computational budget of the method compared to related methods in general and more specifically in each numerical experiment proposed neither at the stage of training nor at the stage of sampling once training is done. For instance, at least in the main text, there is no discussion of the number of integration steps for other methods. It is therefore difficult to interpret the performance comparison between methods and the benefits of the proposed improvements.
1. The paper is clearly written 2. The proposed control variate is shown to be effective, with a lower variance of the estimation. 3. It is a good contribution to explore the few-step generative models in the region of neural sampler, as most methods are more focus on the quality of samples and evaluations of energy functions, rather than the efficiency of the sampling process. 4. Training neural samplers that generate correct weights of modes are important but usually underestimated in many
1. The first concern is the estimation of $\partial_t\log Z_t$. It is well-known that estimating the log partition function, aka free energy, is a fundamental and challenging problem in statistical sciences. An accurate estimation of the free energy is usually hard and not scalable to high dimensional spaces. Intuitively speaking, if $\partial_t\log Z_t$ is accurate and $\partial_T\log Z_T$ is tractable, then $\log Z_0$ can be estimated through integration with small accumulated bias, which is c
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsExplainable Artificial Intelligence (XAI) · Model Reduction and Neural Networks · Reinforcement Learning in Robotics