FALCON: Few-step Accurate Likelihoods for Continuous Flows

Train a few-step flow map to stay invertible, and its likelihoods become cheap and accurate enough to reweight molecular samples to the Boltzmann distribution.

A Boltzmann generator samples molecular shapes from a neural network and corrects them with importance weights, which need the network's exact probability for every sample. FALCON gets that probability from 4 to 16 network steps instead of the hundreds an ODE solver takes.

The paper combines three existing pieces: importance sampling, which needs exact densities; flow maps, which generate in a few large jumps; and a cycle loss, which keeps those jumps invertible so the densities can be computed. This page works each one in a one-dimensional toy you can drive, derives the training target (correcting a sign in the paper's Eq. 10), and reads the peptide results, including the rows FALCON loses.

Sampling molecules from the Boltzmann distribution

A molecule at a fixed temperature does not sit in one shape. Its atoms jiggle and its backbone rotates between a handful of stable shapes, and statistical physics says how often each shape occurs. If $x \in \mathbb{R}^d$ lists every atom's coordinates and $\mathcal{E}(x)$ is the potential energy in units of $k_BT$, the probability density of finding the molecule at $x$ is the Boltzmann distribution

Free energies and heat capacities are averages under this distribution, so a cheap way to draw independent samples from it would speed up drug design and materials work. The energy is cheap to evaluate at any $x$, since a force field is a sum of bond, angle and electrostatic terms. The normalizer $\mathcal{Z}$ is an integral over all of $\mathbb{R}^d$ and is never computed.

The standard tool is molecular dynamics: integrate Newton's equations with femtosecond time steps and record where the molecule goes. Stable shapes are separated by energy barriers, and a trajectory can spend microseconds in one basin (the low-energy region around one stable shape) before it crosses to the next. The samples come out correlated and the rare basins are visited too seldom. The paper's datasets show this: they come from single 1 µs simulations, and the training split is the first $10^5$ frames after warm-up, so it over- or under-samples some basins. On tri-alanine the training set misses one backbone conformation entirely.

The four test molecules are short alanine peptides: alanine dipeptide (22 atoms, so $d = 66$), tri-alanine, alanine tetrapeptide and hexa-alanine. Their slow coordinates are the backbone dihedral angles $\phi$ and $\psi$, two per residue, and a Ramachandran plot of those angles shows the basins.

How a Boltzmann generator corrects its samples

A Boltzmann generator (Noé et al., 2019) trains a generative model on the biased simulation data, draws $K$ independent samples from it, and fixes the bias with importance sampling. Each sample gets a weight equal to the target density over the model density, and any average under the target is estimated as a weighted average:

In (2), $p_1^\theta$ is the density of the trained model (the subscript 1 marks the data end of the flow, introduced below) and $o$ is any observable, such as an indicator of one basin. Dividing by the sum of the weights is the self-normalized part: the unknown $\mathcal{Z}$ multiplies every weight equally and cancels. The estimate converges to the true average as $K$ grows, whatever the model got wrong, provided the model puts some probability wherever the target does. More samples buy more accuracy without retraining, which the paper calls inference-time scaling.

Take three samples whose log-weights $-\mathcal{E}(x) - \log p_1^\theta(x)$ come out as 2.0, 0.0 and −1.0. The weights are $e^{2} = 7.39$, 1 and 0.37, which normalize to 0.844, 0.114 and 0.042. The first sample sits in a basin the model under-samples, so it counts for most of the estimate. The effective sample size (Kish, 1957), normalized by the sample count as in the paper's Appendix E,

is 0.459 here: the three weighted samples carry about as much information as 1.4 unweighted ones. ESS is 1 when the model equals the target and falls toward $1/N$ when one sample takes all the weight.

Now suppose the model reports a log-density for the first sample that is 0.5 too high. Its log-weight drops to 1.5 and the normalized weights become 0.766, 0.171 and 0.063, so the estimate of that basin's population moves by 8 percentage points because of one wrong number. An error of one nat (one unit of natural logarithm) in $\log p_1^\theta$ is a factor of $e \approx 2.72$ in that sample's weight.

Figure 1 runs the same estimator on 2,000 samples from a double well. The target puts 25% of its mass in the right well and the proposal puts 50% there. With the error slider at zero, the reweighted bars match the amber curve. Then add error two ways: random noise per sample, and an error that grows with $x$.

Random error inflates the variance of the weights, so ESS falls and more samples are needed for the same precision. Systematic error biases the estimate, and no number of samples removes it. ESS does not detect it either: at 2 nats of error that grows with $x$, Figure 1 reports an ESS of 0.353 while the reweighted share of the right well is 0.011 instead of 0.250, because ESS measures how evenly the weights are spread, not whether they are correct. The paper reports that the usual shortcuts for continuous flows produce both kinds: the stochastic trace estimator adds noise, and fixed-step integrators “systematically overestimate likelihood” (Appendix D.2.1).

Why likelihoods from continuous flows are expensive

A generator usable here must give samples and their exact densities. Normalizing flows do this through the change of variables formula. If a sample is made by pushing base noise $x_0 \sim p_0$ through an invertible, differentiable map $x_1 = X(x_0)$, the density at the output is the input density divided by how much the map stretches volume there:

The Jacobian $\partial X / \partial x_0$ is a $d \times d$ matrix. Computing it with automatic differentiation costs about $d$ network passes (one per input or output coordinate, or one batched pass that does the same work), and its determinant costs $O(d^3)$, which the paper reports is small next to the passes at these sizes. Classic discrete flows, such as the TarFlow model inside the SBG baseline, restrict the architecture so the determinant is cheap. That makes them fast to sample but, per the paper, slow and unstable to train by maximum likelihood.

Recent Boltzmann generators use continuous normalizing flows (CNFs) instead. A CNF defines the map as the solution of an ODE $dx_s/ds = v_\theta(x_s, s)$ from $s = 0$ (noise) to $s = 1$ (data), as in neural ODEs. The velocity network $v_\theta$ can be any architecture, and it can be trained without simulating the ODE by flow matching: draw noise $x_0$ and a data point $x_1$, form $x_s = s x_1 + (1 - s)x_0$, and regress

The density of a CNF sample comes from the instantaneous change of variables (Chen et al., 2018): along the trajectory, $\log p_s$ changes at the rate $-\operatorname{tr}(\partial v_\theta/\partial x)$, so the log-density is integrated alongside the sample. Written in integrated form, the paper's Eq. (3) is

The exact trace needs the diagonal of a $d \times d$ Jacobian, about $d$ passes at every solver step; Hutchinson's estimator replaces it with one random projection, but the paper finds the resulting noise destabilizes the importance weights. The integral also needs many steps. Boltzmann generators use the adaptive Dormand–Prince 4(5) solver (Dopri5), which shrinks its step until an error estimate falls below a tolerance, here $10^{-5}$, because fixed steps bias the log-density. In Table 5 the adaptive solver takes 200 to 265 function evaluations per sample. The two counts multiply: for alanine dipeptide that is 257 evaluations, each with a 66-column trace, on the order of $257 \times 66 \approx 17{,}000$ network passes for one sample's likelihood. The paper also notes that the solver takes 2 to 3 times more steps when it integrates the log-density than when it only generates.

Flow maps: one network call per jump

A flow map learns the jumps instead of the velocity. For two times $s$ and $t$, the exact flow map $X_v(x_s, s, t)$ sends a point at time $s$ to where the ODE trajectory through it is at time $t$ (Eq. 6). Dividing the displacement by the elapsed time gives the average velocity over the interval:

The instantaneous velocity $v$ is the slope of the trajectory at one instant; $u$ is the slope of the chord from $s$ to $t$. A network $u_\theta(x_s, s, t)$ that takes both times as inputs gives the discrete flow map

and sampling on any grid $0 = t_0 < t_1 < \dots < t_N = 1$ takes $N$ network calls. When $t = s$, the average velocity equals the instantaneous one, so the same network also serves as an ordinary flow-matching velocity.

Why not take $N$ large Euler steps with a flow-matching velocity? Because an Euler step moves along the tangent at the start of the interval, and trajectories curve. Even when every training pair is joined by a straight line, the velocity at a point averages over all the pairs passing through it, and the resulting trajectories bend. Figure 2 uses a one-dimensional toy where the velocity field is known exactly: noise $\mathcal{N}(0, 1)$ at time 0, a mixture with 65% of its mass near −1.5 and 35% near +1.5 at time 1.

At $N = 1$, Euler collapses everything onto one point. At time 0 the noise carries no information about which data point it will become, so the velocity is the mean data point minus the current position, $v(x, 0) = \bar{x}_1 - x$, and one step of size 1 lands on $\bar{x}_1 = -0.45$ from every start. The Jacobian of that step is zero, so its density is infinite and no importance weight can be computed. The flow map lands on the correct trajectory at every $N$ in this toy, because its jumps are exact. A trained network's jumps are only approximately exact, which is why FALCON still uses 4 to 16 steps.

The likelihood of a flow-map sample is the discrete change of variables applied once per jump: $N$ Jacobians of size $d \times d$, about $N \cdot d$ network passes in the paper's accounting (Appendix D.2.1), with no trace estimator and no solver tolerance. For alanine dipeptide at $N = 4$ that is $4 \times 66 = 264$ passes, against the roughly 17,000 above. This shortcut has one condition, which the next two sections deal with: the formula is valid only if every jump is invertible.

How FALCON trains the flow map

Regressing $u_\theta$ on its definition (4) would need the integral, meaning an ODE solve per training example. MeanFlow (Geng et al., 2025) avoids it with an identity that relates $u$ to $v$ and a derivative of $u$, and FALCON uses the same loss with its own time convention. Multiply (4) by $t - s$:

and move the start point $s$ along its trajectory, holding $t$ fixed. On the left, the product rule gives $-u + (t - s)\,du/ds$. On the right, moving the lower limit of an integral forward removes the integrand at that limit, giving $-v(x_s, s)$. Setting the two equal and rearranging:

The derivative is total because $x_s$ itself moves with velocity $v$ as $s$ changes. It equals a Jacobian-vector product (the derivative of the output along one chosen direction in input space) with tangent $(v, 1, 0)$ on the inputs $(x_s, s, t)$, which forward-mode autodiff computes in one pass alongside $u_\theta$ itself. Putting the network's own derivative into the right-hand side, and stopping its gradient (treating the target as a constant during backpropagation, written $\mathrm{sg}$), gives the loss (the paper's Eq. 10, with the sign corrected):

with $v = x_1 - x_0$, the velocity of the straight line through the training pair, standing in for the marginal velocity exactly as in flow matching. The printed Eq. (10) and Algorithm 1 have a minus sign in front of $(t - s)$. That is MeanFlow's sign, and MeanFlow puts the network's position input at the upper limit of the interval, so it differentiates with respect to that limit. FALCON puts the position at the lower limit, which flips the sign. Figure 3 lets you check the plus sign on one trajectory of the toy, computing $du/ds$ by finite differences along the path.

At the default $s = 0.4$, $t = 0.9$, the trajectory is still near the middle and almost flat ($v = 0.036$) but about to climb into the upper mode, so the chord slope is $u = 1.575$. The derivative of $u$ along the path is 3.08, and $0.036 + 0.5 \times 3.08 = 1.575$ matches the chord. With the printed minus sign the target would be −1.504, pointing the network toward the wrong mode.

Two other choices complete the loss. First, the flow-matching term is folded in by sampling some fraction of pairs with $s = t$, where the target reduces to $v$; the paper does not state the fraction or the distribution of $(s, t)$. Second, FALCON-A replaces (10) with an additive loss (Appendix B): the jump from $s$ to $t$ should equal a jump to the midpoint followed by a jump from the midpoint, the SplitMeanFlow objective of Guo et al. (2025), which Appendix B equates with the progressive self-distillation loss of Boffi et al. (2025).

Why the flow map must be invertible

The exact flow map $X_v$ is always invertible when $v$ is Lipschitz (it changes at a bounded rate in $x$): two ODE trajectories cannot cross, so running the ODE backward undoes it (the Picard–Lindelöf theorem, in the paper's Appendix A). Proposition 1 says a perfectly trained $u_\theta$ inherits this, since at the optimum $X_u = X_v$. A network partway through training has no such guarantee. Nothing in (10) stops two nearby inputs from being sent to the same output, or the map from folding back on itself.

If the map goes up, back down, then up again over some interval of inputs, every output in the corresponding range is reached from three inputs, and its true density is the sum of three terms $p_0(x_j)/|X'(x_j)|$. The formula used at sampling time evaluates only the term for the input the sample actually came from. Where the map turns around, $X' = 0$ and that term goes to infinity. Figure 4 starts from the exact one-step map of the toy and adds a bump in the left mode.

If the map folds, no function of $y$ can return the right $x$ for all three preimages, so any backward map has a round-trip error, at least the variance of $x$ among the preimages of each $y$ (the bottom readout of Figure 4). If a backward map undoes the forward map exactly, the forward map is one-to-one on the data it is tested on. FALCON already has a backward map available at no extra cost: the same network called with the times swapped, $X_u(\cdot, t, s)$. The cycle loss (8) and the full objective (9) are

Proposition 2 says a smooth minimizer of (8) is invertible and satisfies the change of variables, whether or not it matches $X_v$. FALCON needs the map to be invertible; it does not need the map to be the true flow, because importance sampling corrects any proposal whose density it knows. The guarantee holds at zero loss, and only on the region the training points cover. The proof (Appendix A) is two lines: if $X_u(\cdot, t, s)$ undoes $X_u(\cdot, s, t)$ exactly, the latter has an inverse by definition.

The weight $\lambda_r$ trades sample quality against invertibility. The paper sweeps $\lambda_r \in \{1, 10, 10^2, \dots, 10^5\}$ on alanine dipeptide (Fig. 6): weak regularization leaves the map poorly invertible and lowers ESS, strong regularization costs sample quality, and $\lambda_r = 10$ is used everywhere. Even then the round trip through the network is not exact: the forward-backward reconstruction error is on the order of $10^{-2}$ (Appendix F.2). To show that an inverse exists, the authors freeze a trained forward model, generate $2 \times 10^5$ noise-sample pairs from it, train a second network to map samples back to noise, and reach reconstruction errors within $10^{-4}$ after 1,000 epochs, which they compare to discrete flows that are invertible by construction.

Section 3 also parameterizes the network as $u_\theta = \operatorname{sign}(t - s)\,h_\theta$, on the stated grounds that $u$ tends to $+v$ as $t \to s$ from above and to $-v$ from below. Under definition (4), both one-sided limits are $+v$, as Figure 3 shows when $t$ approaches $s$ from either side; it is the displacement $(t - s)u$ that changes sign. With the sign factor, the learned $h_\theta$ must jump at $s = t$. The Provenance panel records this together with the sign in (10).

The gradient reaches the network through $u$ in the regression term and through both the forward and the backward call in the cycle term; only the target is detached. One training step for alanine dipeptide, with the corrected target:

# One FALCON training step (Algorithm 1, with the sign of u_tgt corrected).
# Alanine dipeptide has 22 atoms, so every x below has shape (B, 22, 3).
x1 = random_rotation(md_batch())         # data, centre of mass removed
x0 = randn_like(x1)
x0 = x0 - x0.mean(dim=1, keepdim=True)   # noise, also centred
s, t = sample_times()                    # a fraction of pairs have s == t
xs = s * x1 + (1 - s) * x0               # point on the straight path
v = x1 - x0                              # conditional velocity

# u and its total derivative along the path, one forward-mode pass
u, du_ds = jvp(u_net, (xs, s, t), (v, 1.0, 0.0))
u_tgt = v + (t - s) * du_ds              # paper prints "-"; see Provenance

xt_hat = xs + (t - s) * u                # forward jump s -> t
xs_hat = xt_hat + (s - t) * u_net(xt_hat, t, s)   # backward jump t -> s

loss = mse(u, stopgrad(u_tgt)) + lam_r * mse(xs, xs_hat)   # lam_r = 10
loss.backward()   # reaches u_net through u and through both jumps

Training uses AdamW with learning rate $5 \times 10^{-4}$, weight decay $10^{-4}$, a cosine schedule with a 5% warm-up, and an exponential moving average of the weights with decay 0.999 (Appendix D.1).

Sampling with likelihoods in N steps

Algorithm 2 runs the jumps and carries the log-density along. The Jacobian of one jump (7) is $I + (t - s)\,\partial u_\theta/\partial x$, so each step subtracts one log-determinant:

The grid matters when $N$ is small. The paper compares five schedules for 8-step sampling of alanine dipeptide (Fig. 7) and finds the EDM schedule of Karras et al. (2022) substantially better than the other four, with $\rho = 7$, $\sigma_{\min} = 10^{-3}$, $\sigma_{\max} = 1$:

Appendix D.2.2 writes every schedule with $t = 1$ at the noise end and $t = 0$ at the data end ($t_0 = 1$, $t_N = 0$), the reverse of the training convention, so the grid in FALCON's direction is $1 - t_i$. For $N = 4$ that is 0, 0.698, 0.929, 0.989, 1: the first jump covers 70% of the interval and the last covers 1.1%. The authors attribute the gain to the flow field varying more near the data. The EDM toggle in Figure 2 uses the same grid.

In code, for $K$ samples of alanine dipeptide at $N = 4$:

# Samples plus exact log-densities (Algorithm 2), alanine dipeptide, d = 66.
taus = [0.0, 0.698, 0.929, 0.989, 1.0]     # N = 4, EDM grid, rho = 7, flipped
x = sample_base(K, 66)                     # K samples of 22 atoms x 3 coords
logp = log_base_density(x)
for s, t in zip(taus[:-1], taus[1:]):
    step = lambda y: y + (t - s) * u_net(y, s, t)
    J = batched_jacobian(step, x)          # (K, 66, 66): 66 passes or one vmap
    logp = logp - slogdet(J).logabsdet     # change of variables, this step
    x = step(x)
log_w = -energy(x) - logp                  # energy in units of kT
w = softmax(log_w)                         # self-normalized weights, Eq. (2)

Each step costs one network call for the jump and a $66 \times 66$ Jacobian. The paper uses $N$ = 4 for alanine dipeptide, 8 for tri-alanine and the tetrapeptide, and 16 for hexa-alanine, and it evaluates every method on $10^4$ samples.

A transformer instead of an equivariant network

Molecular energies do not change when the whole molecule is rotated or translated, and earlier CNF Boltzmann generators build that symmetry into equivariant graph networks. Those networks are slow, and the cost of hundreds of solver steps per likelihood kept them small: the paper counts at most 2.3 million parameters for the CNF baselines. With likelihoods at 4 to 16 steps, FALCON uses a standard diffusion transformer (DiT) with an extra embedding for the second time input: hidden size 192, 6 blocks, 6 heads, a 64-dimensional conditioning vector, 3.2 million parameters, the same for all four peptides (Table 6). Symmetry is approximate: training samples are randomly rotated, and the mean position is subtracted to remove translations, following SBG.

The authors also train a continuous-time version of the same transformer, “DiT CNF”, as a baseline (Fig. 2 and Table 4). Comparing FALCON with it separates the effect of few-step likelihoods from the effect of the architecture.

What the experiments show

Three metrics are reported, each on $10^4$ samples after reweighting. ESS measures how efficiently the weights use the samples. $\mathcal{E}$-$\mathcal{W}_2$ is the 2-Wasserstein distance (roughly, the root-mean-square distance values must be moved to turn one histogram into the other) between the energy histograms of generated and reference samples; energies are sensitive to small errors in bond lengths, so this metric checks local structure. $\mathbb{T}$-$\mathcal{W}_2$ is the 2-Wasserstein distance on the torus of backbone dihedral angles, which checks whether the basins are covered in the right proportions. Lower is better for both distances. Figure 5 has Tables 2, 3 and 5 of the paper.

On alanine dipeptide, FALCON is not the strongest method. ECNF++, an equivariant CNF with a Dopri5 solver, has the best ESS (0.275 against FALCON's 0.067) and nearly the best dihedral distance (0.189 against 0.402; FALCON-A's 0.180 is lower). FALCON has the best energy distance, 0.225. The paper's case rests on the larger molecules, where ECNF++ degrades: on hexa-alanine its ESS is 0.006 and its energy distance 10.668, against FALCON's 0.060 and 0.892.

Against SBG, a discrete flow built on TarFlow (a transformer-based autoregressive flow that is invertible by construction), FALCON is better on every metric for every molecule, with or without SBG's sequential Monte Carlo (SMC) refinement, which moves the samples with annealed Langevin steps. On the tetrapeptide, for example, FALCON's energy distance is 0.686 against 1.068 for SBG with importance sampling and 1.007 with SMC. Figure 4 of the paper gives SBG up to $5 \times 10^6$ samples, and its energy distance still stays above that of a 4-step FALCON evaluated on $10^4$.

FALCON and FALCON-A split the metrics. FALCON-A has the higher ESS on every molecule (0.104 against 0.077 on tri-alanine) and the best dihedral distance on tri-alanine and hexa-alanine, while the default FALCON has the best energy distance in Table 3 on all three larger peptides and is best or second best on every Table 3 metric there, the reason the authors pick it as the default.

The violet rows show what the few steps cost. Integrated with Dopri5 at 200 to 265 function evaluations, the same family of model has 1.6 to 3.9 times the ESS (0.264 against 0.067 on alanine dipeptide, 0.128 against 0.060 on hexa-alanine) and better dihedral distances on the three smallest systems. The step count is set after training, so a user can choose a point between the two. Measured in function evaluations, the Table 5 gap is 64× for alanine dipeptide (257 against 4) and 13× for hexa-alanine (207 against 16). The abstract's “two orders of magnitude” comes from the paper's Figure 2, which compares wall-clock time at matched dihedral accuracy on alanine dipeptide.

Training plus inference time (Table 4, units not stated) is lowest for FALCON on all four molecules: 7.65 against 12.52 for ECNF++ and 16.83 for SBG on alanine dipeptide, and 25.76 against 137.4 and 57.50 on hexa-alanine. SBG samples faster than FALCON, but its maximum-likelihood training takes longer than FALCON's regression losses.

Limitations

The authors list three. First, the likelihoods are accurate in practice but not guaranteed: invertibility is encouraged by a loss, checked after the fact, and holds exactly only at zero loss. A step that folds somewhere the checks did not look would give wrong weights without any error being raised. Second, one-step generation does not work yet. Third, the work covers only molecular conformation sampling, on peptides of up to six residues; the paper names Bayesian inference and robotics as future applications, and models with structured Jacobians as a route to faster sampling.

Two more follow from the paper's numbers; they are this page's observations, not the authors'. A dense Jacobian costs about $d$ passes per step and per sample, which is 66 to 189 here (hexa-alanine has 63 atoms) but would be 3,000 for a 1,000-atom protein, before the $O(d^3)$ determinant. And the reported ESS values (0.055 to 0.077 for FALCON on the larger peptides) mean the weighted estimate is worth about 6 to 8% of the samples drawn, so most of the generated samples, each with its own Jacobian, contribute little to it.