FUND: Density Flow for Sampling Unnormalised Distributions

Vikas Kanaujia; Vipul Arora

cs.LGSep 26, 2026

FUND: Density Flow for Sampling Unnormalised Distributions

Authors: Vikas Kanaujia, Vipul Arora

Organizations: Department of Electrical Engineering IIT Kanpur · Department of Electrical Engineering KU Leuven

Abstract

Efficient sampling from Boltzmann distributions is central to modelling complex physical systems. Markov Chain Monte Carlo (MCMC) methods suffer from critical slowing down, high autocorrelation, and poor mode-mixing, limiting their scalability. Recent advances, like Boltzmann Generators, offer a promising alternative but remain constrained by costly MCMC-based training, inefficient sampling, and poor ergodicity. We introduce an algorithm for learning Boltzmann distributions that does not require any true samples for training. Our approach draws inspiration from flow matching but departs fundamentally from sample-trajectory matching to distribution-trajectory matching. The algorithm iteratively reshapes the target distribution, using model generated samples to guide learning and ensure comprehensive mode coverage. We validate our method on standard benchmarks, including a 2D Gaussian mixture, Many-Well distributions, and high-dimensional scalar $φ^4$ theory. The proposed approach not only improves sampling performance and accuracy over traditional MCMC and flow-based baselines but also establishes a new method for sample-free learning of complex physical distributions.

Figures & tables

Figure 1: (A) Progressive learning of successive intermediate distributions within a unified flow model using iterative sample reuse. Here $t^{\prime}>t$ , and samples generated from $q_{t}$ are reused to train $q_{t^{\prime}}$ . Latent samples from $q_{z}(z)$ are used to compute $\mathcal{L_{RKL}}(\theta)$ , while replay-buffer samples together with $\log q_{t}(x)$ and $\log p_{t^{\prime}}(x)$ are used to compute $\mathcal{L}(\theta)$ (B) Illustrates sample generation procedure using the trained flow model.

Figure 2: Visualisation of intermediate annealed target distributions $p_{t}(\mathbf{x})\propto p(\mathbf{x})^{t}$ for a 1-D Gaussian mixture model. As $t$ increases, the density gradually transitions from a smoothed, low-contrast form to the fully resolved target distribution. The panel shows 1-D plot for: (A) $t_{1}=0.01$ , (B) $t_{2}=0.1$ , (C) $t_{3}=0.2$ , and (D) $t_{N}=1$ .

\mathcal{L}_{\text{final}}(\theta)=\lambda_{2}\mathcal{L}_{\mathrm{RKL}}(\theta)

Algorithm 1 Training of FUND

Figure 3: Sample plots for the MOG-40 distribution generated using different methods alongside the corresponding ground-truth samples. These plots illustrate the mode-coverage behavior discussed in the quantitative results and highlight the ability of different methods to preserve the multimodal structure of the target distribution during sampling. (A) Ground-truth samples. (B) FKL. (C) RKL. (D) IDEM. (E) FAB. (F) Proposed FUND(SCORE) method.

	MOG-40			MW-8			MW-16
Method	NLL $\downarrow$	RNLL $\downarrow$	ESS $\uparrow$	NLL $\downarrow$	RNLL $\downarrow$	ESS $\uparrow$	NLL $\downarrow$	RNLL $\downarrow$	ESS $\uparrow$
FKL	7.16 $\pm$ 0.02	7.87 $\pm$ 0.06	0.63 $\pm$ 0.01	7.60 $\pm$ 0.05	-31.60 $\pm$ 0.77	0.34 $\pm$ 0.02	17.14 $\pm$ 0.07	-48.57 $\pm$ 5.77	0.001 $\pm$ 0.004
RKL	3527.06 $\pm$ 106.44	7.07 $\pm$ 0.25	0.42 $\pm$ 0.26	76.72 $\pm$ 73.84	-35.40 $\pm$ 0.01	0.97 $\pm$ 0.03	604.83 $\pm$ 85.65	-70.80 $\pm$ 0.06	0.83 $\pm$ 0.06
FAB	7.14 $\pm$ 0.03	8.30 $\pm$ 0.29	0.63 $\pm$ 0.06	7.00 $\pm$ 0.02	-33.87 $\pm$ 0.10	0.81 $\pm$ 0.01	14.28 $\pm$ 0.01	-66.53 $\pm$ 0.13	0.31 $\pm$ 0.02
IDEM	7.09 $\pm$ 0.28	7.01 $\pm$ 0.30	0.42 $\pm$ 0.03	7.64 $\pm$ 0.09	-34.86 $\pm$ 0.02	0.55 $\pm$ 0.03	15.32 $\pm$ 0.43	-70.08 $\pm$ 0.36	0.25 $\pm$ 0.03
FUND(FKL)	7.11 $\pm$ 0.03	7.53 $\pm$ 0.10	0.68 $\pm$ 0.03	6.97 $\pm$ 0.01	-34.25 $\pm$ 0.04	0.87 $\pm$ 0.02	14.04 $\pm$ 0.05	-67.95 $\pm$ 0.25	0.71 $\pm$ 0.04

Table 1: Evaluation results for the MOG-40, MW-8, and MW-16 distributions, using the metrics described in Section 3 . Best results are highlighted in bold and second best are underlined.

Figure 4: Sample plots of intermediate distributions $q_{t}$ learned for the MOG-40 distribution, illustrating the gradual transformation induced by the flow model over $t$ . Panels show sample plots for distribution : (A) $q_{t_{1}}$ at $t_{1}=0.01$ (B) $q_{t_{2}}$ at $t_{2}=0.05$ (C) $q_{t_{N}}$ at $t_{N}=1.0$ and (D) the target distribution $p$ .

	$8\times 8$			$10\times 10$			$12\times 12$
Method	NLL $\downarrow$	RNLL $\downarrow$	ESS $\uparrow$	NLL $\downarrow$	RNLL $\downarrow$	ESS $\uparrow$	NLL $\downarrow$	RNLL $\downarrow$	ESS $\uparrow$
FKL	13.40 $\pm$ 0.07	-9.48 $\pm$ 0.59	0.06 $\pm$ 0.01	22.01 $\pm$ 0.08	-17.60 $\pm$ 0.19	0.01 $\pm$ 0.00	29.89 $\pm$ 0.07	-32.86 $\pm$ 0.58	0.02 $\pm$ 0.004
RKL	23.18 $\pm$ 0.19	-18.75 $\pm$ 0.44	0.08 $\pm$ 0.05	29.01 $\pm$ 1.01	-29.32 $\pm$ 0.78	0.70 $\pm$ 0.04	38.55 $\pm$ 8.29	-41.32 $\pm$ 0.23	0.29 $\pm$ 0.02
FAB	13.19 $\pm$ 0.21	-5.14 $\pm$ 0.59	0.10 $\pm$ 0.05	21.47 $\pm$ 0.22	-1.05 $\pm$ 0.86	0.0001 $\pm$ 0.00	31.44 $\pm$ 0.60	-31.06 $\pm$ 8.35	0.09 $\pm$ 0.07
IDEM	23.76 $\pm$ 1.05	-23.17 $\pm$ 2.70	0.004 $\pm$ 0.004	38.13 $\pm$ 2.13	-25.64 $\pm$ 7.12	0.004 $\pm$ 0.004	65.84 $\pm$ 4.44	-15.82 $\pm$ 3.63	0.003 $\pm$ 0.005
FUND(FKL)	12.48 $\pm$ 0.21	-13.94 $\pm$ 1.52	0.14 $\pm$ 0.04	19.59 $\pm$ 0.03	-17.54 $\pm$ 0.83	0.14 $\pm$ 0.01	28.74 $\pm$ 0.17	-35.49 $\pm$ 1.29	0.33 $\pm$ 0.14

Table 2: Scalar

\phi^{4}

distribution results, with metrics described in Section 3 . Best results are highlighted in bold and second best are underlined.

Figure 5: Test negative log-likelihood (NLL) trajectories over $t$ for the scalar $\phi^{4}$ and MW-16 benchmarks. The curves compare the performance of FUND(SCORE) and FUND(FKL) across intermediate target densities.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Modelling strategy	NLL( $\downarrow$ )
Incremental Modelling	7.69
Single Block Modelling	7.07

Appendix

Table 3: Comparison between incremental modelling and single block modelling on MOG-40.

Distribution	# Affine Coupling Blocks	F or D
MOG-40	16	320
MW-8	12	512
MW-16	12	512
Scalar $\phi^{4}$ ( $8\times 8$ )	10	256
Scalar $\phi^{4}$ ( $10\times 10$ )	12	256
Scalar $\phi^{4}$ ( $12\times 12$ )	12	256

Appendix

Table 4: Summary of the NF architecture configurations used across different distributions.

N	NLL( $\downarrow$ )	RNLL( $\downarrow$ )	ESS( $\uparrow$ )
4	15.04	3.25	0.02
6	13.13	-9.75	0.02
9	12.34	-15.68	0.15

Appendix

Table 5: Results for the scalar

\phi^{4}

model distribution on an

8\times 8

lattice, illustrating the effect of varying the number of intermediate distributions

N

on model performance.

Target Distribution	$N$	$t$
MOG-40	7	{0.01,0.05,0.1,0.25,0.50,0.75,1.0}
MW-8	6	{0.1,0.2,0.4,0.6,0.8,1.0}
MW-16	6	{0.1,0.2,0.4,0.6,0.8,1.00}
Scalar $\phi^{4}$ ( $8\times 8$ )	9	{0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0}
Scalar $\phi^{4}$ ( $10\times 10$ )	9	{0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0}
Scalar $\phi^{4}$ ( $12\times 12$ )	9	{0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0}

Appendix

Table 6: Intermediate distribution parameters

N

and

t

-schedules used in training.

Setting	Annealing Schedule $(\lambda_{1},\lambda_{2})$	NLL $\downarrow$	RNLL $\downarrow$	ESS $\uparrow$
Setting 1	$(1,1)\rightarrow(0.5,1)\rightarrow(0.25,1)\rightarrow(0.1,1)\rightarrow(0.01,1)$	7.07	-33.60	0.79
Setting 2	$(1,0.1)\rightarrow(1,0.25)\rightarrow(1,0.5)\rightarrow(1,1)$	7.05	-33.46	0.77
Setting 3	$(1,0.1)\rightarrow(1,0.5)\rightarrow(0.5,1)\rightarrow(0.1,1)$	7.03	-33.62	0.82
Setting 4	$(1,0.1)\rightarrow(1,0.25)\rightarrow(1,0.5)\rightarrow(1,1)\rightarrow(0.5,1)\rightarrow(0.25,1)\rightarrow(0.1,1)\rightarrow(0.01,1)$	6.96	-34.30	0.88

Appendix

Table 7: Sensitivity analysis of the annealing schedule for

(\lambda_{1},\lambda_{2})

on the MW-8 distribution using FUND(FKL).

$t_{1}$	NLL ( $\downarrow$ )	Number of Modes Recovered
0.010	7.07	40
0.025	20.03	33
0.050	53.20	20
0.075	990.99	15
0.100	1327.82	12

Appendix

Table 8: Sensitivity analysis of the initial annealing parameter

t_{1}

on the MOG-40 benchmark.

Figure 6: Effect of the initial annealing parameter $t_{1}$ on MOG-40. The figure shows sample plots of (A) the target distribution and (B)–(F) generated samples for $t_{1}=0.01$ , $0.025$ , $0.05$ , $0.075$ , and $0.1$ , respectively. Larger $t_{1}$ values lead to mode collapse.

	$\phi^{4}$ $(12\times 12)$		MW-16
Method	Training Time	Inference Time	Training Time	Inference Time
MCMC (HMC)	N/A	510 sec	–	–
RKL	15.16 hrs	2.2 sec	0.24 hrs	0.56 sec
FAB	51.20 hrs	2.2 sec	2.5 hrs	0.56 sec
IDEM	5.25 hrs	44.0 sec	4.3 hrs	34.0 sec
FUND(FKL)	35.69 hrs	2.2 sec	2.3 hrs	0.56 sec

Appendix

Table 9: Training and inference runtime comparison across methods for

\phi^{4}

(12\times 12)

and MW-16.

Explore similar work

Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows
Sep 25, 2026Louis Grenioux, RuiKang OuYang, Luhuan WuMolecular Generation Normalizing Flows
Center for Computational Mathematics, Flatiron Institute, New York, NY, USA · Department of Engineering, University of Cambridge, Cambridge, UK · Department of Applied Mathematics and Statistics, Johns Hopkins University, Baltimore, MD, USA
Jeffreys Flow: Robust Boltzmann Generators for Rare Event Sampling via Parallel Tempering Distillation
Apr 7, 2026Guang Lin, Christian Moya, Di Qi +1
Department of Mathematics, Purdue University, West Lafayette, IN 47907, USA
Leveraging generative models to assist Monte Carlo sampling
Aug 7, 2026Marylou GabriéNormalizing Flows Generative Modeling
Laboratoire de Physique Statistique, École normale supérieure, PSL Research University, 24 rue Lhomond, 75005 Paris, France

Sep 25, 2026J/K move · Enter open · S save

Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows

Louis Grenioux, RuiKang OuYang, Luhuan Wu

Center for Computational Mathematics, Flatiron Institute, New York, NY, USA · Department of Engineering, University of Cambridge, Cambridge, UK · Department of Applied Mathematics and Statistics, Johns Hopkins University, Baltimore, MD, USA

Molecular Generation Normalizing Flows Flow Map Learning Importance Sampling

Open paper AskPDF arXiv

Generating independent, equilibrium samples of molecular systems at scale remains a central obstacle in computational statistical mechanics. Boltzmann Generators address this by pairing a generative model with importance sampling to obtain consistent samples from the target distribution. We introduce Normalizing Flow Flow Maps (NF

^2

M), which combines the strengths of recent stochastic flow maps with the tractability of classic normalizing flows to build a Boltzmann Generator. Unlike most methods, which correct the generative model only at the end, NF

^2

M reweighs each denoising transition as generation proceeds, avoiding wasted compute on trajectories that are ultimately discarded. At each denoising step, a conditional normalizing flow proposes clean configurations given the current noisy state (a simpler task than sampling directly from the target) and its exact likelihood enables correcting each proposal toward the true denoising transition of the target Boltzmann distribution. This is in contrast to most existing methods, whose likelihoods are approximate or expensive to evaluate, undermining the statistical reliability of the correction. We establish consistency of the corrected transitions and bound how approximation errors propagate through the sampling chain. We evaluate NF

^2

M on peptide systems, demonstrating improved sampling efficiency and sample quality.

Abstract

Explore similar work

Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows

Jeffreys Flow: Robust Boltzmann Generators for Rare Event Sampling via Parallel Tempering Distillation

Leveraging generative models to assist Monte Carlo sampling