Why Instantaneous Velocity Is Not Enough for 8-Step Generation

Fine-tuning an image generator for a specific task has become much easier. With the right dataset and a LoRA, we can teach a model a new style, domain, or editing behavior at a relatively modest cost.

But making a model draw what we want and making it draw quickly are two different problems.

If the original model calls its network 30 or 50 times per image, attaching a LoRA usually leaves that cost almost unchanged. The style changes; the sampler does not. We call each network evaluation an NFE, or Number of Function Evaluations. A 50-step sampler uses roughly 50 NFE.

I was not looking for a one-step generator at any cost. I wanted to retain useful quality and control while turning a task-specific model into a fast generator that works at around 8 NFE.

My first question was simple:

Why not take the teacher we call 50 times and call it only eight times?

Following that question led me to the flow map in Align Your Flow (AYF) and its Eulerian Map Distillation objective.

What it means to ask the teacher 50 times

Strictly speaking, the teacher is not a “50-step model.” It is a continuous velocity field. Given the current state, it tells us which direction to move:

\frac{d x_t}{dt}=v_\phi(x_t,t).

$x_t$ is the noisy latent at time $t$ .
$v_\phi(x_t,t)$ is the current velocity predicted by the pretrained teacher.
$\phi$ denotes the frozen teacher parameters.

The number 50 belongs to the numerical solver that follows this field. In the simplest Euler form, one update is

x_{t+\Delta t} \approx x_t+\Delta t\,v_\phi(x_t,t).

A 50-step sampler never asks the teacher to describe a long future in one answer. It asks for the instantaneous velocity, moves a short distance, then asks again from its new position. If the trajectory curves, the next call can correct the direction.

It is like navigation that gives you a new instruction every hundred meters. Each instruction only has to be locally correct because another one is coming soon.

What changes when we ask only eight times?

Reducing 50 steps to eight makes every interval roughly six times longer. The velocity at the starting point does not tell us how the trajectory will curve across that whole interval.

The teacher’s instantaneous velocity is not wrong. The problem is that the gap between questions has changed while the answer still describes only the present moment.

Over a short interval, one tangent approximates a curve well. Over a long interval, following the starting tangent sends us farther from the true path. Simply evaluating the original teacher eight times therefore increases integration error and can disturb both image quality and control.

That raised the next question:

If the model must cross a long interval before it can ask again, what should it predict instead of instantaneous velocity?

A flow map sends the current state to a future state

A flow map directly transports a state at one time to its state at another time. If we write the exact flow map as $\Phi$ , then

\Phi_{t\to s}(x_t)=x_s.

Given the current state $x_t$ , the departure time $t$ , and the destination time $s$ , it returns the point $x_s$ on the teacher trajectory.

The AYF student parameterizes that map as

f_\theta(x_t,t,s) =x_t+(s-t)F_\theta(x_t,t,s).

There are two related objects here that are easy to conflate:

$f_\theta$ is the flow map that sends $x_t$ to $x_s$ .
$F_\theta$ is the interval-average velocity used to express that movement.

If the exact destination $x_s$ were known, the average velocity would be

F(x_t,t,s) =\frac{x_s-x_t}{s-t} =\frac{1}{s-t}\int_t^s v(x_u,u)\,du.

The teacher’s $v_\phi$ is the tangent at one point. The student’s $F_\theta$ summarizes the direction changes across an interval as one displacement. An eight-step student divides the trajectory into eight intervals and predicts the average movement needed before its next query.

Teacher trajectory

noise, t=1  ──→  x_t3  ──→  x_t2  ──→  x_t1  ──→  target x_s
                     ╰──────── flow map: t3 → s ────────╯
                                 ╰── flow map: t2 → s ──╯
                                              ╰─ t1 → s ─╯

We can now state the difference more precisely:

A 50-step sampler applies the teacher’s current velocity only over a short interval, then asks again.
The student predicts the movement over $(t,s)$ , because it has to travel farther before the next call.

But where does the correct average velocity come from?

Once I understood the flow map, another problem appeared:

Do we need to integrate the teacher all the way to $s$ for every interval $(t,s)$ , produce $x_s$ , and regress to that endpoint?

That is possible, but it requires many teacher rollouts. AYF-EMD uses a different condition.

Hold the destination time $s$ fixed and move the starting point $(x_t,t)$ along one teacher trajectory. No matter where we start on that trajectory, the flow map should arrive at the same $x_s$ :

f_\theta(x_t,t,s)=x_s.

As the departure time $t$ moves along the teacher trajectory, the predicted destination must not move:

\frac{d}{dt}f_\theta(x_t,t,s)=0.

This equation states the central question in AYF-EMD:

If I move the starting point slightly, does the student still predict the same destination?

How does the prediction change when the starting point moves?

Start again from the definition of the flow map:

f_\theta(x_t,t,s) =x_t+(s-t)F_\theta(x_t,t,s).

Along the teacher trajectory, $dx_t/dt=v_\phi$ . Differentiating with respect to $t$ gives

\frac{d f_\theta}{dt} =v_\phi-F_\theta +(s-t)\left(J_xF_\theta\,v_\phi+\partial_tF_\theta\right).

The two unfamiliar terms have concrete meanings:

$J_xF_\theta\,v_\phi$ measures how the student’s average-velocity prediction changes when $x$ moves in the teacher-velocity direction.
$\partial_tF_\theta$ measures how the prediction changes when only the time input $t$ moves.

Together, they give the rate of change of $F_\theta$ as $(x_t,t)$ travels along the teacher trajectory:

\tau =J_xF_\theta\,v_\phi+\partial_tF_\theta.

Do we really need the whole Jacobian?

Thinking about the implementation immediately produces another question:

$J_xF_\theta$ is enormous. Do we really need to construct the entire Jacobian matrix?

We do not. We only need the change in one direction: the teacher-trajectory direction $(v_\phi,1)$ . That is exactly what a Jacobian–Vector Product, or JVP, computes.

Define

h(\varepsilon) =F_\theta(x_t+\varepsilon v_\phi,t+\varepsilon,s).

Differentiating with respect to $\varepsilon$ gives

\left.\frac{dh}{d\varepsilon}\right|_{\varepsilon=0} =J_xF_\theta\,v_\phi+\partial_tF_\theta =\tau.

In PyTorch, this directional derivative can be evaluated with torch.func.jvp.

Checking the same condition with the quotient rule

The previous derivative is dense, so I wanted to check it from the definition of average velocity. Hold the destination $x_s$ and destination time $s$ fixed:

F(t)=\frac{x_s-x_t}{s-t}.

Apply the quotient rule,

\left(\frac{N}{D}\right)' =\frac{N'D-ND'}{D^2}.

Let $N=x_s-x_t$ and $D=s-t$ . Then $N'=-v_\phi$ and $D'=-1$ , so

\frac{dF}{dt} =\frac{(-v_\phi)(s-t)+(x_s-x_t)}{(s-t)^2}.

Substitute $x_s-x_t=(s-t)F$ :

\frac{dF}{dt} =\frac{v_\phi-F}{t-s}.

The left-hand side is the JVP $\tau$ from above. A correct flow map must therefore satisfy

F_\theta-v_\phi +(t-s)\left(J_xF_\theta\,v_\phi+\partial_tF_\theta\right)=0.

The two derivations expose the same fact from different directions:

Differentiating the flow map says that the destination must remain fixed.
Differentiating average velocity tells us how the interval velocity must change for that to be true.

A one-dimensional check

Suppose the teacher trajectory is $x_t=t^2$ . Its instantaneous velocity is $v(t)=2t$ . The exact average velocity from $t$ to $s$ is

F(t,s)=\frac{s^2-t^2}{s-t}=s+t.

Differentiating with respect to $t$ gives $dF/dt=1$ . The other side agrees:

\frac{v-F}{t-s} =\frac{2t-(s+t)}{t-s}=1.

Finally, apply the flow map:

x_t+(s-t)F =t^2+(s-t)(s+t) =s^2.

Every departure time $t$ predicts the same destination, $x_s=s^2$ .

The residual that measures misalignment

In AYF implementations, the quantity $g$ measures how far the student is from this condition:

g =F_\theta-v_\phi +(t-s)\left(J_xF_\theta\,v_\phi+\partial_tF_\theta\right).

When $g=0$ , the student is locally aligned with the teacher trajectory. When it is not zero, moving the starting point changes the student’s destination.

That leads to another implementation question:

$g$ is a correction direction in output space. How do we turn it into a gradient for the student parameters?

AYF-EMD normalizes $g$ and passes it to the student through a surrogate loss. A simplified implementation is

\tau =\operatorname{stop}\!\left( J_xF_\theta\,v_\phi+\partial_tF_\theta \right),

g =\operatorname{stop}(F_\theta)-v_\phi +r(t-s)\tau,

\hat g =\frac{\operatorname{stop}(g)}{\lVert g\rVert+c},

L =\left\| F_\theta-\operatorname{stop}(F_\theta)+\hat g \right\|^2.

$\operatorname{stop}$ means that a value participates in the forward calculation but receives no gradient through that path. The teacher velocity, JVP tangent, $g$ , $\hat g$ , and the second copy of $F_\theta$ behave like fixed targets. Gradients flow only through the first $F_\theta$ in the loss.

Numerically, $F_\theta-\operatorname{stop}(F_\theta)=0$ . Its derivative tells a different story:

\nabla_\theta L =2\bigl(J_\theta F_\theta\bigr)^\top\,\widehat{g}.

First, $g$ tells us the direction in which the student output is wrong. The surrogate then transmits that correction arrow to the student parameters.

In the implementation, torch.func.jvp treats the Student parameters and conditioning as fixed values and differentiates only with respect to $(x,t)$ . We detach the resulting tangent $\tau$ , then use a separate differentiable forward pass to obtain gradients with respect to $\theta$ . This avoids backpropagating through the JVP itself.

Isn’t the JVP meaningless early in training?

If the student were completely random, that concern would be justified. The input derivative of a random function is not an obvious source of useful supervision.

In practice, two mechanisms work together.

First, the student starts from the same pretrained flow model as the teacher. New flow-map conditioning modules or LoRAs are initialized so they do not severely disrupt the initial output. At the start of training, $F_\theta$ is at least close to the existing velocity model.

Second, AYF uses a tangent warmup coefficient $r$ . Early in training, when $r=0$ ,

g\approx F_\theta-v_\phi,

which resembles ordinary velocity matching. As the student adapts, $r$ gradually rises toward roughly 0.99 and activates the JVP term. Normalizing $g$ by $\|g\|+c$ also limits excessively large updates.

Why I chose AYF

There are already powerful approaches to few-step generation. LADD—Latent Adversarial Diffusion Distillation—uses adversarial training in latent space to compress high-resolution models into very few steps. Its results can be excellent, but generator–discriminator balance, discriminator features, learning rates, and regularization all become part of the problem.

What attracted me to AYF was that it does not force us to build one specialized one-step generator. It learns one model that connects arbitrary pairs of times.

Because $(t,s)$ is part of the input, the same student can run at 1, 2, 4, or 8 steps. For an 8-NFE target, it applies eight smaller flow maps. For a faster preview, it can use two or four larger intervals.

The non-adversarial AYF-EMD stage also has a useful diagnostic quality. It directly enforces that moving the starting point along a teacher trajectory should not change the destination. When training fails, that gives me a concrete condition to inspect.

It is not free:

The JVP makes every training step substantially more expensive than a conventional LoRA fine-tune.
Differentiating the time embedding can be unstable and may need dedicated stabilization.
The student inherits the teacher’s trajectories—and its errors.
Its one-step quality may trail a method optimized solely for one-step generation.
The JVP pass and the differentiable forward that produces parameter gradients must compute the same function.

The last constraint led to my next question.

Why can’t `torch.func.jvp` use Flash Attention?

One AYF-EMD training step evaluates the student along at least two paths:

A torch.func.jvp pass computes $F$ and its directional input derivative $\tau$ .
A differentiable forward pass builds the graph for parameter gradients and backward.

Some stock Flash SDPA and widely used Flash Attention paths do not provide the forward-mode AD rule required by torch.func.jvp. One practical workaround is to route the JVP pass through Math SDPA while using a Flash kernel for the ordinary forward and backward.

That can be faster, but it creates a new question:

Is it valid to compute the tangent with function A (Math) and update the parameters with the Jacobian of function B (Flash)?

The two kernels may implement the same attention equation, but BF16 reduction order, online softmax, and kernel fusion mean they are not bitwise the same numerical function. The hybrid update is approximately

\bigl(J_\theta F_{\text{Flash}}\bigr)^\top\,\widehat{g}_{\text{Math}}.

It mixes the tangent-derived correction from one implementation with the parameter Jacobian of another.

Whether that difference is negligible cannot be settled by inspection. In intervals where $F-v$ and $(t-s)\tau$ nearly cancel, the true residual $g$ is small. A minor backend difference could then rotate the normalized direction $\hat g$ . In other intervals, the same error may be irrelevant. We have to measure it.

I plan to compare at least:

primal relative and maximum error between Math and Flash backends;
agreement between the torch.func.jvp tangent and finite differences;
cosine similarity and relative L2 error of the final parameter gradients;
teacher alignment and generation quality in short, same-seed training runs; and
full AYF update time and peak memory.

I am also building a Flash-JVP attention path that supplies the missing forward-mode rule directly. The primal attention can use the Flash formulation, while the directional derivative with respect to Q, K, and V is computed by an explicit jvp rule and a Triton kernel.

This is more than a speed optimization. The experiment asks whether a faster kernel preserves the training function defined by the paper. Once the measurements are complete, I will share how the custom Flash-JVP was constructed, whether the Math/Flash hybrid was safe, and how much it changed the full training cost in the next Curiosity.

References

Sabour, Fidler, and Kreis. “Align Your Flow: Scaling Continuous-Time Flow Map Distillation.” 2025.
NVIDIA Toronto AI Lab. Align Your Flow project page.
Sauer et al. “Fast High-Resolution Image Synthesis with Latent Adversarial Diffusion Distillation.” 2024.
PyTorch. torch.func.jvp does not support flash_attn.
PyTorch. Extending torch.func with autograd.Function.