Boltzmannator

Boltzmann generators and normalizing flows

This application illustrates, in one dimension, the principle of Boltzmann generators: an invertible transformation is trained to map a simple, easily sampled distribution onto the Boltzmann distribution of a physical system.

1 · The sampling problem

In equilibrium statistical mechanics a configuration x occurs with probability

p^*(x) = e^−U(x)/k_BT / Z ,

where U(x) is the potential energy, k_BT the thermal energy and Z the partition function, which is generally unknown. If U has several deep wells, p^* is concentrated in well-separated regions, and molecular dynamics or Monte Carlo simulations require long trajectories to move between them. A Boltzmann generator avoids this problem: it produces statistically independent samples by transforming samples of a simple distribution.

2 · Change of variables

Consider a latent variable z with a simple density p_z(z) (top left panel) and an invertible, differentiable map x = f_θ(z) (centre panel). The transformed variable has the density

p_x(x) = p_z(z) |df_θ/dz|⁻¹ , z = f_θ⁻¹(x) .

Where the map expands intervals (Jacobian |J| = |df_θ/dz| > 1), probability is spread over a larger range and the density decreases; where the map contracts, the density increases. The bottom left panel shows the factor |J|⁻¹. Since f_θ is invertible, p_x is normalized by construction. This construction is called a normalizing flow.

3 · Training the flow

The parameters θ are adjusted by gradient descent such that p_x approaches the target p^*. Two standard objectives are available.

Energy-based training minimises the reverse Kullback–Leibler divergence,

L = ⟨ U(f_θ(z))/k_BT − log |J(z)| ⟩_{z∼p_z} = KL(p_x ‖ p^*) + const .

Only the potential energy function is required; samples of the target are not needed. The first term drives the samples towards low energies; the second term, an entropy, prevents the distribution from collapsing onto the energy minimum.

Example-based training minimises the negative log-likelihood of data points x_i distributed according to p^* (forward Kullback–Leibler divergence),

NLL = ⟨ −log p_z( f_θ⁻¹(x)) + log |J| ⟩_data = KL(p^* ‖ p_x) + const .

The loss panel shows both contributions separately (blue and orange) together with their sum (dark line).

While the target is displayed, the figure reports two free energies. The variational free energy of the flow, βF_flow = ⟨ln p_z⟩ + ⟨U/k_BT − ln |J|⟩, is the free energy of the model distribution p_x in the target potential; βF_exact = −ln Z is the free energy of the target itself. The two are related by the Gibbs–Bogoliubov inequality,

F_flow = F_exact + k_BT KL(p_x ‖ p^*) ≥ F_exact ,

which holds because the Kullback–Leibler divergence is non-negative. The gap between the two values shown in the figure is therefore exactly k_BT KL(p_x ‖ p^*); it closes when the flow reproduces the target, and training can be read as the minimisation of the variational free energy towards the exact one.

4 · Reweighting

A trained flow approximates the target but does not reproduce it exactly. Exact averages can nevertheless be obtained by importance sampling, in which each sample x is weighted with

w(x) = p^*(x) / p_x(x) , ⟨A⟩ = Σ w_iA(x_i) / Σ w_i .

The effective sample size N_eff = (Σw)²/Σw² quantifies the overlap between p_x and p^* (activate Show importance weights after sampling).

5 · The transformation families

Polynomial maps are simple but not necessarily invertible; the application marks turning points and estimates p_x by a kernel density estimate in this case. The single-layer perceptron adds K sigmoid steps with positive weights to an increasing linear function and is therefore strictly monotonic and invertible. The rational-quadratic spline is a flexible monotonic spline on [−B, B] and is the building block of neural spline flows.

In one dimension an exact transformation exists, x^*(z) = F^*−1(F_z(z)), constructed from the two cumulative distribution functions (Show exact transformation). In higher dimensions no such formula exists and the transformation must be learned; this is the setting in which normalizing flows are applied in practice.

What you can do here

A simple latent density is transformed into a target Boltzmann distribution, either manually with the sliders or by training. Hovering over a control displays a short explanation.

The figure

Latent density p_z(z) (top left): selected and shaped on the Densities tab.
Transformation x = f_θ(z) (centre): orange dots mark points where it is not invertible.
Jacobian |J|⁻¹ (bottom left): the local stretch factor that reshapes the density.
Transformed density p_x(x): drawn sideways next to the map, sharing its x axis.
Target p^*(x): overlaid when Show target is active (automatically during training).
Right column: histograms of the latent and transformed samples, and the training loss.
While the target is shown, the top left corner reports the variational free energy βF_flow, the exact βF_exact = −ln Z and their difference, the Kullback–Leibler divergence KL(p_x ‖ p^*).

Densities tab

The latent distribution and its mean and width are chosen here. The target p^*(x) ∝ e^−U(x)/kT is defined by the coefficients of the potential U(x) = u₁x + u₂x² + u₃x³ + u₄x⁴ and the temperature kT. Show exact transformation displays the map that reproduces the target exactly. Show potential U(x) draws the potential next to the transformed density, and Show CDF construction illustrates how the exact map is obtained from the two cumulative distribution functions; the quantile follows the pointer over the map panel.

Map tab

The transformation family and its parameters are set here; the figure is updated continuously. Show mapping lines indicates how individual points z are mapped to x = f_θ(z). Show transport bands divides the latent density into N bands of equal probability mass and follows them through the map; each band keeps its mass, so its width shows the local stretching. Moving the pointer over the map panel (on touch screens: tapping it) displays z, x, the Jacobian and the densities at that point. The Animation control morphs the whole figure between the identity map (t = 0) and the current map (t = 1), or, with Path: Training, replays the parameter history of the last training run. The geodesic path follows the displacement interpolation f_t(z) = (1−t)z + t f_θ(z), which in one dimension is the optimal-transport path between p_z and p_x; coloured dots of equal probability mass ride the curve, and their crowding shows the density. Defaults restores the identity map; Randomize assigns random parameter values.

Training tab

Sample draws N latent points and passes them through the map; Data generates example points from the target distribution. The objective (energy-based requires only the potential, example-based requires the data), the optimiser, the number of epochs, the learning rate and the batch size are selected here. Every and Delay control the frequency of figure updates; a larger delay slows the visualisation. Stop interrupts a running training; Reset training clears the loss history and restores the identity map. Show training trails keeps faded copies of the map from earlier epochs, and the Loss slice selector displays the loss as a function of a single parameter around the current value.

Things to try

Double well. Set u₂ = −2 and u₄ ≈ 0.3 on the Densities tab, activate Show target and press Train. The initially unimodal density splits into two peaks.
Transport of probability mass. Activate Show transport bands on the Map tab and drag the map parameters. Each coloured band carries an equal share of the probability; the map changes its width but never its mass.
Slow motion. Set Delay to 200 ms and Every to 1 to display every epoch.
Non-invertible map. Select the polynomial map and increase θ₂ or θ₃ until turning points appear; the density is then estimated by a kernel density estimate (KDE).
Bimodal latent distribution. A bimodal latent density matched to a double-well target requires a much simpler transformation.
Optimiser comparison. Train the same problem with SGD and with Adam at different learning rates.
Learning from data. Press Data, select Example-based and train from samples instead of energies.
Mode collapse. Prepare a deep double well at low temperature and train energy-based: the flow concentrates its probability in a single well, since the reverse Kullback–Leibler divergence penalises probability in unlikely regions but tolerates missing modes. Then press Data, switch to Example-based, press Reset training and train again: the forward divergence places mass in both wells.
Quality control. After training, press Sample and activate Show importance weights; N_eff close to 100% indicates p_x ≈ p^*.

General

Auto-rescale recomputes the axis ranges after every change; deactivate it to keep the axes fixed for comparisons. The switch at the bottom toggles between light and dark mode; the dots next to the portrait select the header picture. The buttons above the figure save the current figure as a PNG image and copy a link that reproduces the current settings. Reset restores the initial state of the application.