Modern diffusion and flow models can generate high-quality samples, but inference requires repeatedly evaluating a large neural network while numerically solving an ordinary differential equation. Flow-map distillation addresses this cost by training a student to jump directly across a long segment of the teacher's trajectory, potentially generating an image in only one or a few network evaluations.
Most flow-map distillation methods, however, rely on a dataset to construct intermediate noisy states. FreeFlow: Flow Map Distillation Without Data asks a more fundamental question:
Can we distill the complete generative behavior of a pretrained flow model without using real data or precomputed teacher trajectories?
FreeFlow's answer is yes. Its key observation is that the prior distribution is the one location where the teacher's generating process is guaranteed to be correctly represented. Every teacher trajectory begins from the prior, regardless of whether the original training data are available or whether classifier-free guidance and post-training have changed the teacher's output distribution.
Starting from this observation, FreeFlow builds a data-free Predict-and-Correct framework:
- Predict with generating flows: learn a prior-anchored flow map whose local motion agrees with the teacher's velocity field.
- Correct with noising flows: align the distribution produced by the student with the teacher distribution, compensating for errors that accumulate along imperfect student trajectories.
The first component provides trajectory-level supervision; the second provides distribution-level correction. Together, they allow the student to learn from the teacher using only prior noise, teacher queries, and student-generated samples.
This article first explains the paper at a high level and then dives into its mathematical core: the transition from the squared regression objective in Equation 8 to the velocity-alignment gradient in Equation 9. That derivation reveals what the prediction objective is really doing and provides a useful lesson in chain rules, Jacobians, stop-gradient operators, and gradient-equivalent surrogate losses.
From PID to FreeFlow: a more complete formulation
This work is especially interesting to me because it revisits a question closely connected to my previous work, Physics Informed Distillation for Diffusion Models (PID). PID asked whether a student could learn the solution of a teacher's probability-flow ODE directly from the governing dynamics, rather than from stored teacher trajectories. It parameterized a trajectory function \(x_\theta(z,t)\), imposed the boundary condition \(x_\theta(z,T)=z\), and minimized a PINN-style ODE residual. Numerical differentiation between neighboring student predictions made that residual practical to optimize.
From my perspective, FreeFlow is a better-developed and more complete formulation of this core idea. It is not presented as a direct extension of PID, but it gives a strong modern answer to essentially the same underlying question:
Can a student learn the teacher's global generation map using only prior noise and access to the teacher's vector field, without first constructing a dataset of teacher trajectories?
The connection is clearest in FreeFlow's Predict branch. PID requires the student's trajectory function to satisfy the teacher ODE at the student's own predicted states. FreeFlow similarly parameterizes a flow map and requires its generating velocity to match the teacher's instantaneous velocity at student-predicted states. Both approaches therefore learn through a local differential constraint instead of directly regressing precomputed teacher endpoints.
FreeFlow then adds an important piece that PID did not explicitly model: distribution-level correction. A small local trajectory error changes the next state at which the teacher is queried, so even a student with a low local ODE residual can gradually drift away from the teacher's intended output distribution. FreeFlow names this compounding-error problem and introduces its Correct branch to align the noising velocity induced by student samples with the teacher field.
The conceptual progression can be summarized as:
- PID: learn the teacher's solution map by minimizing a physics-informed ODE residual.
- FreeFlow Predict: express the same trajectory-learning principle as prior-anchored generating-velocity alignment.
- FreeFlow Correct: repair the distribution mismatch that can remain after local trajectory errors accumulate.
FreeFlow also states the research question more sharply. PID framed the idea through the analogy to physics-informed neural networks; FreeFlow frames it directly as data-free flow-map distillation, identifies the prior as the guaranteed shared boundary, and cleanly separates trajectory fidelity from distribution fidelity. This is why, in my view, FreeFlow feels like a better version of the idea explored in PID: it preserves the insight that the generator can be learned as an ODE solution map, while adding clearer flow-map language, an explicit error-correction mechanism, and stronger evidence at larger image scales.
FreeFlow in one picture
The overall logic of the paper can be summarized as
\[ z\sim\pi \xrightarrow{\text{student flow map}} f_\theta(z,\delta) \xrightarrow{\text{teacher velocity query}} \text{trajectory supervision}, \]
supplemented by
\[ f_\theta(z,1) \xrightarrow{\text{re-noising}} x_r \xrightarrow{\text{teacher--student field comparison}} \text{distribution correction}. \]
No real sample is required in either branch.
The problem with conventional data-based distillation
Suppose the teacher defines a generating distribution \(\hat p_t\) at every time \(t\). Conventional distillation commonly constructs training states by noising samples from an available dataset, producing another family of distributions \(\tilde p_t\).
The usual training procedure implicitly assumes
\[ \tilde p_t\approx\hat p_t. \]
FreeFlow argues that this assumption can fail. The teacher may generalize beyond its original dataset, may have been modified by post-training, or may use classifier-free guidance that shifts its sampling distribution. The original data may also be private or completely unavailable. In these cases,
\[ \tilde p_t\neq\hat p_t, \]
so the student is trained on states that do not faithfully represent the teacher's actual sampling trajectories. FreeFlow calls this Teacher--Data Mismatch.
Why the prior solves the starting-point problem
Although the dataset-noising and teacher-generating distributions can differ at intermediate times, both processes meet at the prior:
\[ \tilde p_1=\hat p_1=\pi. \]
The prior is therefore the only distribution that is guaranteed to be on-distribution for the teacher and can be sampled without data. FreeFlow anchors every student trajectory at
\[ z\sim\pi, \qquad t=1. \]
Instead of learning a general flow map between arbitrary starting and ending times, the student learns where a fixed prior sample should move after a duration \(\delta\).
Predict: follow the teacher's generating field
FreeFlow parameterizes the student as
\[ f_\theta(z,\delta) = z+\delta F_\theta(z,\delta), \]
where \(F_\theta\) is the average velocity over the interval. Differentiating the predicted state with respect to duration gives its instantaneous generating velocity:
\[ v_G = \partial_\delta f_\theta = F_\theta+\delta\partial_\delta F_\theta. \]
The prediction objective trains this student velocity to agree with the teacher velocity evaluated at the student's current prediction:
\[ v_G\!\left(f_\theta(z,\delta),1-\delta\right) \approx u\!\left(f_\theta(z,\delta),1-\delta\right). \]
If this identity holds for all \(\delta\), and the boundary condition \(f_\theta(z,0)=z\) holds, then the student follows the same initial-value problem as the teacher and recovers its trajectory.
Why prediction alone is insufficient
The student queries the teacher at its own predicted state. If the student makes a small error, the next teacher velocity is evaluated at that erroneous state. Although the teacher velocity is locally correct there, it does not necessarily return the student to the original teacher trajectory. Errors can therefore accumulate as \(\delta\) grows.
Prediction answers:
Is the student moving according to the teacher's local vector field?
But it does not directly answer:
Does the student's final output distribution still match the teacher's distribution after all trajectory errors accumulate?
Correct: repair the generated distribution
FreeFlow therefore adds a correction objective. It first generates a student endpoint
\[ x_0^\theta=f_\theta(z,1), \]
then re-noises it to an intermediate state \(x_r\). An auxiliary model estimates the marginal noising velocity \(v_N\) induced by the student's generated distribution. FreeFlow compares this field with the teacher field:
\[ v_N(x_r,r) \approx u(x_r,r). \]
This branch is related to variational score distillation, but FreeFlow expresses the update through velocity fields. Conceptually,
\[ \boxed{ \text{prediction preserves trajectory dynamics}, \qquad \text{correction preserves the output distribution}. } \]
The paper finds that neither component is sufficient by itself: prediction suffers from accumulated trajectory error, while correction alone can lose the noise-to-sample correspondence and collapse. Their combination is the complete method.
Why Equation 9 is the mathematical core
The Predict branch is where FreeFlow converts its intuitive goal—“make the student follow the teacher”—into a trainable, data-free objective. Equation 8 initially looks like an ordinary regression loss on the average velocity \(F_\theta\). Equation 9 reveals that its actual update aligns the student's generating velocity \(v_G\) with the teacher velocity \(u\).
That transition is more than cosmetic. It explains:
- why the prediction loss is dynamically meaningful;
- why the teacher is queried at a student-predicted state;
- why the term \(\delta\partial_\delta F_\theta\) is necessary;
- how stop-gradient turns a constructed target into an external update direction;
- why the final expression is a gradient surrogate rather than a conventional regression loss.
Flow-map distillation papers often compress these vector-calculus steps into one line. To understand the derivation completely, we must answer several basic questions:
- Why is the gradient of a squared error not simply twice the network output?
- Why does a Jacobian appear when differentiating with respect to model parameters?
- What exactly does
stop-gradientdo? - How does the average-velocity residual become a generating-velocity residual?
- Is Equation 9 a new loss, or only a different representation of the gradient of Equation 8?
The rest of this article derives every step carefully, from scalar calculus to the full FreeFlow expression.
1. Formal setup for the Equation 9 derivation
Let
\[ z\sim\pi \]
be a sample from the prior distribution, and let \(\delta\in[0,1]\) denote the integration duration from the fixed prior time. FreeFlow parameterizes its student flow map as
\[ f_\theta(z,\delta) = z+\delta F_\theta(z,\delta), \]
where:
- \(f_\theta(z,\delta)\) is the predicted state after integration duration \(\delta\);
- \(F_\theta(z,\delta)\) is the student's average velocity over that interval;
- \(\theta\) denotes all trainable student parameters.
The boundary condition follows automatically:
\[ f_\theta(z,0)=z. \]
Let
\[ u(x,t) \]
be the instantaneous velocity supplied by the pretrained teacher. For an optimal student, FreeFlow derives the identity
\[ F_{\theta^*}(z,\delta) + \delta\,\partial_\delta F_{\theta^*}(z,\delta) = u\!\left(f_{\theta^*}(z,\delta),1-\delta\right). \]
This identity motivates the prediction target
\[ u_{\mathrm{target}} = u\!\left(f_\theta(z,\delta),1-\delta\right) - \delta\,\partial_\delta F_\theta(z,\delta). \]
FreeFlow Equation 8 then takes the form
\[ \mathcal L_{\mathrm{pred}} = \mathbb E_{z,\delta} \left[ \left\| F_\theta(z,\delta) - \operatorname{sg}\!\left(u_{\mathrm{target}}\right) \right\|_2^2 \right], \tag{1} \]
where \(\operatorname{sg}(\cdot)\) denotes stop-gradient.
Our goal is to understand how the gradient of this squared regression loss becomes the velocity-alignment gradient used in Equation 9.
2. First distinction: the forward value and the gradient
Stop-gradient does not change a tensor's numerical value during the forward pass:
\[ \operatorname{sg}(y)=y \qquad\text{in value}. \]
It changes only differentiation:
\[ \nabla_\theta\operatorname{sg}(y)=0. \]
For example, suppose
\[ y=y(\theta). \]
The two expressions
\[ F_\theta-y(\theta) \]
and
\[ F_\theta-\operatorname{sg}(y(\theta)) \]
have exactly the same forward value. Their gradients are different:
\[ \nabla_\theta(F_\theta-y) = \nabla_\theta F_\theta-\nabla_\theta y, \]
whereas
\[ \nabla_\theta\left(F_\theta-\operatorname{sg}(y)\right) = \nabla_\theta F_\theta. \]
In FreeFlow, the teacher query and the derivative correction are used to construct a target. Stop-gradient prevents the squared regression loss from differentiating through that target.
This means the optimization behaves like ordinary supervised regression:
\[ \text{student prediction} \longrightarrow \text{fixed target for this update}. \]
3. Warm-up: the scalar squared-error derivative
Before considering vectors, let the student output be a scalar \(F_\theta\in\mathbb R\), and let \(\bar y=\operatorname{sg}(y)\). Consider
\[ \mathcal L(\theta) = (F_\theta-\bar y)^2. \]
Expanding the square gives
\[ \mathcal L = F_\theta^2 -2F_\theta\bar y +\bar y^2. \]
Differentiate with respect to \(\theta\):
\[ \frac{d\mathcal L}{d\theta} = 2F_\theta\frac{dF_\theta}{d\theta} -2\bar y\frac{dF_\theta}{d\theta} +2\bar y\frac{d\bar y}{d\theta}. \]
Because \(\bar y\) is detached,
\[ \frac{d\bar y}{d\theta}=0. \]
Therefore,
\[ \frac{d\mathcal L}{d\theta} = 2\frac{dF_\theta}{d\theta} (F_\theta-\bar y). \tag{2} \]
This is simply the chain rule:
\[ \frac{d\mathcal L}{d\theta} = \frac{d\mathcal L}{dF_\theta} \frac{dF_\theta}{d\theta}, \]
with
\[ \frac{d\mathcal L}{dF_\theta} = 2(F_\theta-\bar y). \]
The distinction is important:
\[ 2(F_\theta-\bar y) \]
is the gradient with respect to the network output, whereas
\[ 2\frac{dF_\theta}{d\theta}(F_\theta-\bar y) \]
is the gradient with respect to the network parameters.
4. Why the answer is not just \(2F_\theta\)
Consider only the term
\[ F_\theta^2. \]
If we differentiate with respect to \(F_\theta\), then
\[ \frac{d}{dF_\theta}F_\theta^2 = 2F_\theta. \]
But FreeFlow optimizes the parameters \(\theta\), not the output as an independent variable. Because \(F_\theta\) depends on \(\theta\), the chain rule gives
\[ \frac{d}{d\theta}F_\theta^2 = 2F_\theta\frac{dF_\theta}{d\theta}. \]
For example, suppose
\[ F_\theta=\theta^2. \]
Then
\[ F_\theta^2=\theta^4, \]
and direct differentiation gives
\[ \frac{d}{d\theta}\theta^4=4\theta^3. \]
Using only \(2F_\theta\) would give \(2\theta^2\), which is incorrect. Applying the chain rule gives
\[ 2F_\theta\frac{dF_\theta}{d\theta} = 2\theta^2\cdot2\theta = 4\theta^3. \]
The general computational path is
\[ \theta \longrightarrow F_\theta \longrightarrow \mathcal L. \]
Backpropagation first computes the sensitivity of the loss to the output and then propagates it through the sensitivity of the output to the parameters.
5. Move from scalars to vectors
In a generative model, \(F_\theta\) is normally a high-dimensional vector or tensor. Flatten it conceptually into
\[ F_\theta = \begin{bmatrix} F_{\theta,1}\ \vdots\ F_{\theta,d} \end{bmatrix} \in\mathbb R^d. \]
Let the parameter vector contain \(p\) values:
\[ \theta\in\mathbb R^p. \]
The Jacobian of the output with respect to the parameters is
\[ J_\theta F_\theta = \frac{\partial F_\theta}{\partial\theta} = \begin{bmatrix} \frac{\partial F_1}{\partial\theta_1}&\cdots&\frac{\partial F_1}{\partial\theta_p}\ \vdots&\ddots&\vdots\ \frac{\partial F_d}{\partial\theta_1}&\cdots&\frac{\partial F_d}{\partial\theta_p} \end{bmatrix} \in\mathbb R^{d\times p}. \]
The loss is scalar, so its parameter gradient must have dimension \(p\):
\[ \nabla_\theta\mathcal L\in\mathbb R^p. \]
The output-space gradient has dimension \(d\):
\[ \nabla_{F_\theta}\mathcal L\in\mathbb R^d. \]
The chain rule therefore takes the form
\[ \nabla_\theta\mathcal L = (J_\theta F_\theta)^\top \nabla_{F_\theta}\mathcal L. \tag{3} \]
The dimensions confirm the transpose:
\[ (p\times d)(d\times1)=p\times1. \]
This operation is a vector-Jacobian product in the convention typically used by reverse-mode automatic differentiation.
6. Expand the vector squared error
Let
\[ \bar y=\operatorname{sg}(y). \]
The vector squared error is
\[ \mathcal L = \|F_\theta-\bar y\|_2^2. \]
Using \(\|a\|_2^2=a^\top a\),
\[ \mathcal L = (F_\theta-\bar y)^\top(F_\theta-\bar y). \]
Expand the product:
\[ \mathcal L = F_\theta^\top F_\theta -F_\theta^\top\bar y -\bar y^\top F_\theta +\bar y^\top\bar y. \]
The cross terms are scalars, so
\[ F_\theta^\top\bar y = \bar y^\top F_\theta. \]
Therefore,
\[ \mathcal L = F_\theta^\top F_\theta -2F_\theta^\top\bar y +\bar y^\top\bar y. \tag{4} \]
Now differentiate term by term.
6.1 First term
The output-space gradient is
\[ \nabla_{F_\theta}(F_\theta^\top F_\theta) = 2F_\theta. \]
Applying the parameter-space chain rule from Equation 3,
\[ \nabla_\theta(F_\theta^\top F_\theta) = 2(J_\theta F_\theta)^\top F_\theta. \tag{5} \]
6.2 Cross term
Because \(\bar y\) is constant for differentiation,
\[ \nabla_{F_\theta}\left(-2F_\theta^\top\bar y\right) = -2\bar y. \]
Hence,
\[ \nabla_\theta\left(-2F_\theta^\top\bar y\right) = -2(J_\theta F_\theta)^\top\bar y. \tag{6} \]
6.3 Target-only term
Since \(\bar y\) is detached,
\[ \nabla_\theta(\bar y^\top\bar y)=0. \tag{7} \]
6.4 Combine the terms
Combining Equations 5--7 gives
\[ \nabla_\theta\mathcal L = 2(J_\theta F_\theta)^\top F_\theta -2(J_\theta F_\theta)^\top\bar y. \]
Factor out the transposed Jacobian:
\[ \boxed{ \nabla_\theta\mathcal L = 2(J_\theta F_\theta)^\top (F_\theta-\bar y) }. \tag{8} \]
Substituting \(\bar y=\operatorname{sg}(y)\),
\[ \boxed{ \nabla_\theta \left\| F_\theta-\operatorname{sg}(y) \right\|_2^2 = 2(J_\theta F_\theta)^\top \left(F_\theta-\operatorname{sg}(y)\right) }. \tag{9} \]
Because stop-gradient leaves the forward value unchanged, the numerical residual is
\[ F_\theta-\operatorname{sg}(y) \overset{\mathrm{value}}{=} F_\theta-y. \]
7. Substitute the FreeFlow target
Return now to FreeFlow. To keep the notation compact, define
\[ F_\theta \equiv F_\theta(z,\delta), \]
and
\[ u \equiv u\!\left(f_\theta(z,\delta),1-\delta\right). \]
The target is
\[ y = u-\delta\partial_\delta F_\theta. \]
The residual appearing in the gradient of Equation 8 is therefore
\[ F_\theta-y = F_\theta-\left(u-\delta\partial_\delta F_\theta\right). \]
Rearrange it:
\[ F_\theta-y = F_\theta +\delta\partial_\delta F_\theta -u. \tag{10} \]
Thus, from Equation 9,
\[ \nabla_\theta\mathcal L_{\mathrm{pred}} = 2\mathbb E_{z,\delta} \left[ (J_\theta F_\theta)^\top \operatorname{sg}\!\left( F_\theta +\delta\partial_\delta F_\theta -u \right) \right]. \tag{11} \]
Why is it useful to write the residual this way? Because the first two terms have a direct dynamical interpretation.
8. From average velocity to generating velocity
Recall the flow-map parameterization
\[ f_\theta(z,\delta) = z+\delta F_\theta(z,\delta). \]
Differentiate with respect to \(\delta\). The prior sample \(z\) is fixed, so
\[ \partial_\delta z=0. \]
Using the product rule,
\[ \partial_\delta f_\theta(z,\delta) = F_\theta(z,\delta) + \delta\partial_\delta F_\theta(z,\delta). \tag{12} \]
FreeFlow defines this derivative as the student's generating velocity:
\[ v_G\!\left(f_\theta(z,\delta),1-\delta\right) \coloneqq \partial_\delta f_\theta(z,\delta). \]
Consequently,
\[ \boxed{ v_G = F_\theta+\delta\partial_\delta F_\theta }. \tag{13} \]
Equation 10 now becomes
\[ F_\theta-y = v_G-u. \tag{14} \]
Define the velocity-alignment residual
\[ \Delta_{v_G,u} \coloneqq v_G\!\left(f_\theta(z,\delta),1-\delta\right) - u\!\left(f_\theta(z,\delta),1-\delta\right). \tag{15} \]
Then Equation 11 becomes
\[ \boxed{ \nabla_\theta\mathcal L_{\mathrm{pred}} = 2\mathbb E_{z,\delta} \left[ (J_\theta F_\theta)^\top \operatorname{sg}\!\left( \Delta_{v_G,u} \right) \right] }. \tag{16} \]
This is the central velocity-alignment interpretation: FreeFlow updates the student according to the difference between how its predicted trajectory currently moves and how the teacher says that same state should move.
9. Construct the Equation 9 surrogate objective
Consider the scalar surrogate
\[ \widetilde{\mathcal L}_{\mathrm{pred}} = \mathbb E_{z,\delta} \left[ F_\theta(z,\delta)^\top \operatorname{sg}\!\left( \Delta_{v_G,u} \right) \right]. \tag{17} \]
Because the residual is detached,
\[ \nabla_\theta \operatorname{sg}(\Delta_{v_G,u})=0. \]
Differentiating Equation 17 gives
\[ \nabla_\theta\widetilde{\mathcal L}_{\mathrm{pred}} = \mathbb E_{z,\delta} \left[ (J_\theta F_\theta)^\top \operatorname{sg}(\Delta_{v_G,u}) \right]. \tag{18} \]
Comparing Equations 16 and 18,
\[ \nabla_\theta\mathcal L_{\mathrm{pred}} = 2\nabla_\theta\widetilde{\mathcal L}_{\mathrm{pred}}. \]
The factor \(2\) is an overall positive scale and can be absorbed into the learning rate or loss weight. Therefore, up to this constant factor, FreeFlow can express the Equation 8 update as
\[ \boxed{ \nabla_\theta \mathbb E_{z,\delta} \left[ F_\theta(z,\delta)^\top \operatorname{sg}\!\left( v_G\!\left(f_\theta(z,\delta),1-\delta\right) - u\!\left(f_\theta(z,\delta),1-\delta\right) \right) \right] }. \tag{19} \]
This is the form of FreeFlow Equation 9.
10. Equation 8 and the Equation 9 surrogate are not the same scalar loss
This is the most important conceptual caveat.
Equation 8 is a squared regression loss:
\[ \mathcal L_{\mathrm{pred}} = \|F_\theta-\operatorname{sg}(y)\|^2. \]
The surrogate behind Equation 9 is a dot product:
\[ \widetilde{\mathcal L}_{\mathrm{pred}} = F_\theta^\top \operatorname{sg}(v_G-u). \]
Their numerical values are generally different. They are used because their first-order parameter gradients point in the same direction, up to a constant factor:
\[ \nabla_\theta\mathcal L_{\mathrm{pred}} = 2\nabla_\theta\widetilde{\mathcal L}_{\mathrm{pred}}. \]
Thus, Equation 9 should be read as a gradient representation or gradient-producing surrogate, not as an algebraic claim that
\[ \|F_\theta-y\|^2 = F_\theta^\top(v_G-u). \]
That scalar equality is false in general.
The surrogate is useful because it exposes the optimization signal:
\[ \boxed{ \text{student parameter update} = \text{student output Jacobian}^\top \times \text{velocity mismatch}. } \]
11. Why does Equation 9 contain \(F_\theta\), not \(f_\theta\)?
Equation 8 directly regresses the average velocity \(F_\theta\). Its parameter Jacobian is therefore
\[ J_\theta F_\theta. \]
The flow-map output is
\[ f_\theta=z+\delta F_\theta, \]
so
\[ J_\theta f_\theta = \delta J_\theta F_\theta. \]
Using \(f_\theta\) in the surrogate would introduce an additional \(\delta\) factor unless it were compensated explicitly. The dot product with \(F_\theta\) is exactly the simple scalar whose gradient produces the Jacobian-vector product required by the original regression loss.
Importantly, Equation 9 does not train
\[ F_\theta\approx v_G-u. \]
The detached residual \(v_G-u\) acts as an externally computed output-space update direction. The role of the dot product is to propagate that direction through the network's parameter Jacobian.
12. A fully worked scalar FreeFlow example
Consider a simplified scalar model
\[ F_\theta(z,\delta)=\theta\delta. \]
Then
\[ f_\theta(z,\delta) = z+\delta F_\theta = z+\theta\delta^2. \]
The derivative of the average velocity is
\[ \partial_\delta F_\theta=\theta, \]
and the generating velocity is
\[ v_G = F_\theta+\delta\partial_\delta F_\theta = \theta\delta+\delta\theta = 2\theta\delta. \]
Suppose the teacher velocity at the predicted state is the scalar \(u\). The Equation 8 target is
\[ y = u-\delta\partial_\delta F_\theta = u-\theta\delta, \]
but it is detached when used as the target. The squared loss is
\[ \mathcal L = \left( heta\delta-operatorname{sg}(u-\theta\delta)\right)^2. \]
Its residual value is
\[ \theta\delta-(u-\theta\delta) = 2\theta\delta-u = v_G-u. \]
Because
\[ \frac{\partial F_\theta}{\partial\theta}=\delta, \]
the gradient is
\[ \frac{\partial\mathcal L}{\partial\theta} = 2\delta(v_G-u). \]
Now consider the surrogate
\[ \widetilde{\mathcal L} = F_\theta\operatorname{sg}(v_G-u) = \theta\delta\operatorname{sg}(v_G-u). \]
Its gradient is
\[ \frac{\partial\widetilde{\mathcal L}}{\partial\theta} = \delta(v_G-u). \]
The gradients differ only by the factor \(2\), exactly as the general derivation predicts.
13. The role of automatic differentiation
Two different derivative operations appear in this derivation and should not be confused.
13.1 Time derivative used to compute \(v_G\)
FreeFlow requires
\[ \partial_\delta F_\theta(z,\delta). \]
This is the derivative of the network output with respect to its scalar time or duration input. It can be computed with a Jacobian-vector product, or approximated using finite differences.
13.2 Parameter gradient used to optimize the student
Training requires
\[ (J_\theta F_\theta)^\top(v_G-u). \]
This is a vector-Jacobian product with respect to the model parameters, normally computed by ordinary reverse-mode backpropagation.
Therefore the computation involves two different Jacobians:
\[ \underbrace{\partial_\delta F_\theta}_{\text{time derivative for the target}} \qquad\text{and}\qquad \underbrace{J_\theta F_\theta}_{\text{parameter Jacobian for learning}}. \]
The first describes how the predicted average velocity changes over the trajectory. The second describes how the network output changes when its weights are updated.
14. PyTorch-style implementation
The following pseudocode mirrors the continuous-time derivation:
from torch.func import jvp
def student_average_velocity(params, z, delta, condition):
return student(params, z, delta, condition)
def predict_loss(params, z, delta, condition):
def forward_at_duration(current_delta):
return student_average_velocity(
params, z, current_delta, condition
)
# Primal output F and time derivative dF / ddelta.
F, dF_ddelta = jvp(
forward_at_duration,
(delta,),
(torch.ones_like(delta),),
)
predicted_state = z + delta * F
# The target branch is detached.
with torch.no_grad():
teacher_u = teacher(
predicted_state,
1.0 - delta,
condition,
)
target = teacher_u - delta * dF_ddelta
return ((F - target) ** 2).sum()
An explicitly gradient-equivalent surrogate can be written as
def surrogate_loss(params, z, delta, condition):
def forward_at_duration(current_delta):
return student_average_velocity(
params, z, current_delta, condition
)
F, dF_ddelta = jvp(
forward_at_duration,
(delta,),
(torch.ones_like(delta),),
)
predicted_state = z + delta * F
with torch.no_grad():
teacher_u = teacher(
predicted_state,
1.0 - delta,
condition,
)
v_G = F + delta * dF_ddelta
velocity_residual = v_G - teacher_u
return (F * velocity_residual).sum()
The second loss produces half the gradient of the unnormalized squared-error version. Multiplying it by \(2\) gives the same first-order gradient scale:
return 2.0 * (F * velocity_residual).sum()
If the original implementation uses a mean rather than a sum, its normalization factors must also be matched.
15. A subtle implementation point: what exactly should be detached?
Mathematically, the target branch in Equation 8 is detached:
\[ \operatorname{sg}\!\left( u(f_\theta,1-\delta) - \delta\partial_\delta F_\theta \right). \]
This means no gradient from this loss should pass through:
- the teacher output \(u\);
- the teacher input \(f_\theta\) inside the teacher query;
- the time derivative \(\partial_\delta F_\theta\) used in the target.
The student prediction outside the target remains differentiable:
\[ F_\theta \longrightarrow \theta. \]
In code, it is safest to construct the entire target or velocity residual under no_grad() or apply .detach() after constructing it.
A common error is to detach only the teacher parameters but allow the teacher query to backpropagate into the predicted state. That would add an unwanted term involving
\[ J_xu\;J_\theta f_\theta. \]
Another common error is to leave \(\partial_\delta F_\theta\) attached. That introduces higher-order derivatives of the student, changing both the objective and its computational cost.
16. First-order equivalence versus higher-order equivalence
It is common to write the squared-loss gradient as
\[ 2(J_\theta F_\theta)^\top \operatorname{sg}(F_\theta-y). \]
For ordinary first-order optimization, this gives the same numerical parameter gradient as
\[ 2(J_\theta F_\theta)^\top \left(F_\theta-\operatorname{sg}(y)\right). \]
However, these computational graphs are not identical for higher-order differentiation. Detaching the full residual explicitly declares that it is an external gradient signal. If one later differentiates the parameter gradient again—for meta-learning, Hessian computation, or gradient penalties—the difference matters.
Thus the safest statement is:
Equation 9 reproduces the intended first-order update of Equation 8, up to an overall scale, under the specified stop-gradient convention.
17. Geometric interpretation
The velocity residual
\[ \Delta_{v_G,u}=v_G-u \]
lives in output space. It says how the student's local motion differs from the teacher's vector field at the student's predicted state.
The parameter update requires translating this output-space error into parameter space:
\[ (J_\theta F_\theta)^\top\Delta_{v_G,u}. \]
Each parameter receives an update proportional to how changing that parameter would move the output along the velocity-error direction.
If a parameter barely changes \(F_\theta\), its corresponding Jacobian column is small, so it receives little update. If changing a parameter strongly affects the output in the direction of the residual, it receives a larger update.
This is the same basic mechanism underlying standard neural-network regression. FreeFlow's novelty lies not in the chain rule itself, but in identifying that its carefully constructed regression residual equals a mismatch between two meaningful velocity fields.
18. Common misunderstandings
Misunderstanding 1: “The derivative of \(F_\theta^\top F_\theta\) is just \(2F_\theta\)”
That is true only when differentiating with respect to \(F_\theta\):
\[ \nabla_{F_\theta}(F_\theta^\top F_\theta)=2F_\theta. \]
When differentiating with respect to \(\theta\), the chain rule adds the output Jacobian:
\[ \nabla_\theta(F_\theta^\top F_\theta) = 2(J_\theta F_\theta)^\top F_\theta. \]
Misunderstanding 2: “Stop-gradient makes the target zero”
No. Stop-gradient preserves the forward value and sets only its derivative to zero:
\[ \operatorname{sg}(y)=y, \qquad \nabla_\theta\operatorname{sg}(y)=0. \]
Misunderstanding 3: “Equation 9 regresses \(F_\theta\) toward \(v_G-u\)”
No. The dot product is a surrogate that produces the desired parameter gradient. It is not a squared regression target.
Misunderstanding 4: “Equation 8 and Equation 9 have equal numerical values”
They generally do not. Their intended first-order gradients are proportional.
Misunderstanding 5: “The teacher velocity is evaluated on the true teacher path”
In the FreeFlow prediction objective, the teacher is queried at the student-predicted state
\[ f_\theta(z,\delta). \]
This is central to FreeFlow's data-free construction, but it is also why student errors can accumulate.
Misunderstanding 6: “The time JVP and parameter backpropagation are the same operation”
They are not. The time derivative constructs \(v_G\); reverse-mode backpropagation maps the detached velocity residual into a parameter update.
19. The entire derivation in one chain
Starting from FreeFlow Equation 8,
\[ \mathcal L_{\mathrm{pred}} = \mathbb E \left\| F_\theta - \operatorname{sg}\!\left( u-\delta\partial_\delta F_\theta \right) \right\|^2, \]
differentiate the squared regression loss:
\[ \nabla_\theta\mathcal L_{\mathrm{pred}} = 2\mathbb E \left[ (J_\theta F_\theta)^\top \operatorname{sg}\!\left( F_\theta+delta\partial_\delta F_\theta-u \right) \right]. \]
Use the flow-map derivative
\[ v_G = \partial_\delta f_\theta = F_\theta+delta\partial_\delta F_\theta \]
to obtain
\[ \nabla_\theta\mathcal L_{\mathrm{pred}} = 2\mathbb E \left[ (J_\theta F_\theta)^\top \operatorname{sg}(v_G-u) \right]. \]
Finally, recognize that
\[ \nabla_\theta \left[ F_\theta^\top\operatorname{sg}(v_G-u) \right] = (J_\theta F_\theta)^\top\operatorname{sg}(v_G-u). \]
Therefore, up to the constant factor \(2\),
\[ \boxed{ \nabla_\theta\mathcal L_{\mathrm{pred}} \propto \nabla_\theta \mathbb E \left[ F_\theta^\top \operatorname{sg}(v_G-u) \right] }. \]
That is the path from Equation 8 to Equation 9.
20. Final takeaway
The derivation is built from three ordinary ideas:
- Chain rule: Differentiating with respect to network parameters introduces the network-output Jacobian.
- Stop-gradient regression: The constructed target contributes a numerical value but no derivative.
- Flow-map product rule: The combination \(F_\theta+\delta\partial_\delta F_\theta\) is exactly the student's generating velocity \(\partial_\delta f_\theta\).
Together, these transform a conventional squared regression update into a physically meaningful velocity-alignment update:
\[ \boxed{ \text{squared regression on average velocity} \quad\Longleftrightarrow_{\text{gradient}}\quad \text{alignment of student and teacher generating velocities}. } \]
The key word is gradient. Equation 9 does not claim that the two scalar losses are identical. It exposes the dynamical meaning of the update produced by Equation 8.
Reference
- Shangyuan Tong, Nanye Ma, Saining Xie, and Tommi Jaakkola. Flow Map Distillation Without Data, 2025.
- Joshua Tian Jin Tee, Kang Zhang, Hee Suk Yoon, Dhananjaya Nagaraja Gowda, Chanwoo Kim, and Chang D. Yoo. Physics Informed Distillation for Diffusion Models, TMLR 2024. Code.