A diffusion network can predict clean data, added noise, or velocity. Salimans and Ho study these parameterizations for distillation into fewer sampling steps. Changing the prediction target for the same training sample changes endpoint stability and the weighting of prediction errors.
This module selects the supporting parameterization technique from Appendix D. Implement only make_v, v_to_x0, and v_to_eps; the other conversion helpers are supplied. Completing this module covers these rotation identities. Teacher–student targets and progressive training remain separate parts of the paper.
Return one Tensor of clean data recovered from a noise prediction, matching x_t shape [B,...]; B is batch size and trailing axes are sample coordinates. alpha/sigma are already broadcastable by the input convention; do not add coefficient axes. Preserve the documented matching floating dtype/device and gradients. Do not return multiple parameterizations as a tuple or mutate inputs.
Return one Tensor of noise recovered from a clean-data prediction, matching x_t shape [B,...]; B is batch size and trailing axes are sample coordinates. alpha/sigma are already broadcastable by the input convention; do not add coefficient axes. Preserve the documented matching floating dtype/device and gradients. Do not return multiple parameterizations as a tuple or mutate inputs.
Return one Tensor of angular velocity v, matching x0 shape [B,...]; B is batch size and trailing axes are sample coordinates. alpha/sigma are already broadcastable by the input convention; do not add coefficient axes. Preserve the documented matching floating dtype/device and gradients. Do not return multiple parameterizations as a tuple or mutate inputs.
Return one Tensor of clean data recovered from v, matching x_t shape [B,...]; B is batch size and trailing axes are sample coordinates. alpha/sigma are already broadcastable by the input convention; do not add coefficient axes. Preserve the documented matching floating dtype/device and gradients. Do not return multiple parameterizations as a tuple or mutate inputs.
Return one Tensor of noise recovered from v, matching x_t shape [B,...]; B is batch size and trailing axes are sample coordinates. alpha/sigma are already broadcastable by the input convention; do not add coefficient axes. Preserve the documented matching floating dtype/device and gradients. Do not return multiple parameterizations as a tuple or mutate inputs.
Return one Tensor in the dst parameterization, with pred/x_t shape [B,...], preserving the documented matching floating dtype/device and gradients. When src=dst, return pred directly (the same object); otherwise return the converted prediction. Do not return a dictionary of x0/eps/v or mutate inputs.
Let
Set
The prediction target therefore also implies a loss weighting. See Progressive Distillation, Appendix D, Equations 26 and 30.
Implement three core functions:
| Function | Inputs and output | Key idea |
|---|---|---|
make_v |
x0, eps, alpha, sigma → v |
Compute the rotation’s tangent |
v_to_x0 |
x_t, v, alpha, sigma → x0 |
Apply the inverse rotation directly |
v_to_eps |
x_t, v, alpha, sigma → epsilon |
Apply the inverse rotation directly |
Tensors have shape [B,...]; coefficients broadcast to that shape. Preserve shape and dtype without aggregating over the batch. Neither inverse requires division, preserving stability at the pure-signal and pure-noise endpoints.
eps_to_x0, x0_to_eps, and convert are provided. They handle the 1e-8 denominator floor, string validation, and identity conversions. convert calls your three core functions.
| Item | This exercise | Common scheduler interfaces |
|---|---|---|
| Input coefficients | Already square-rooted alpha and sigma | Often store alpha_cumprod and take square roots internally |
| Output | Return a Tensor directly | diffusers schedulers usually return a step-output object |
| Endpoint division | Clamp denominator to 1e-8 | Depends on the schedule and stability strategy |
| Data dimensions | Any broadcastable shape | Images usually use B×C×H×W |
Adjust the rotation angle phi to inspect the position and tangent vectors. The rotation direction determines the tangent's sign. This plot shows the coefficient plane; when high-dimensional independent noise and data are approximately orthogonal, it also suggests a sample-space interpretation.
Move through time to inspect amplification factors
Open module 7 in the browser, complete the three v-parameterization functions, and run the tests.
Tests use the supplied conversion entry point to check six round trips, then independently check finite differences, endpoints, and the error-weight identity. Distillation training is outside this exercise.
For teacher–student targets and progressive step reduction, continue reading Progressive Distillation. Flow Matching uses a different velocity, differentiated with respect to linear time; this module differentiates with respect to a rotation angle.