V-Prediction Conversions
Problem Statement
Implement the v-parameterization: the velocity target used by progressive distillation, Imagen Video and Stable Diffusion 2.x, plus its inverse back to and .
Background
Write and , so that traces a circular arc between the clean sample and pure noise. The velocity is the derivative of that arc:
and are orthogonal unit-scaled directions, which makes the inverse a rotation by the same angle:
Note there is no division by anywhere. That is the whole point: -prediction degenerates at (the target becomes trivially itself, and recovering divides by something tiny), while -prediction degenerates at . The target is well-conditioned at both ends, which is what makes few-step distillation and zero-terminal-SNR schedules work.
Your Task
Implement two functions:
def compute_v(x0, eps, alpha_bar):
def v_to_x0_eps(x_t, v, alpha_bar):
compute_v returns the velocity target; v_to_x0_eps returns the tuple (x0, eps).
Input Format
- x0, eps, x_t, v: NumPy arrays of matching shape.
- alpha_bar (float): the scalar in (0, 1].
Output Format
compute_v returns one array; v_to_x0_eps returns a tuple of two arrays.
Sample
x0 = np.array([1.0, -1.0])
eps = np.array([0.5, 0.25])
ab = 0.36
v = compute_v(x0, eps, ab)
print(np.round(v, 4).tolist())
x_t = np.sqrt(ab) * x0 + np.sqrt(1 - ab) * eps
r0, re = v_to_x0_eps(x_t, v, ab)
print(np.round(r0, 4).tolist())
print(np.round(re, 4).tolist())
With sqrt(0.36) = 0.6 and sqrt(0.64) = 0.8, v = 0.6eps - 0.8x0, and the inverse rotation returns exactly the original x0 and eps.
Example
x0 = np.array([1.0, -1.0]) eps = np.array([0.5, 0.25]) ab = 0.36 v = compute_v(x0, eps, ab) print(np.round(v, 4).tolist()) x_t = np.sqrt(ab) * x0 + np.sqrt(1 - ab) * eps r0, re = v_to_x0_eps(x_t, v, ab) print(np.round(r0, 4).tolist()) print(np.round(re, 4).tolist())
[-0.5, 0.95] [1.0, -1.0] [0.5, 0.25]
sqrt(0.36) = 0.6 and sqrt(0.64) = 0.8. So v = 0.6*[0.5, 0.25] - 0.8*[1.0, -1.0] = [-0.5, 0.95]. The noisy sample is x_t = 0.6*x0 + 0.8*eps = [1.0, -0.4]. Rotating back, x0_hat = 0.6*x_t - 0.8*v and eps_hat = 0.8*x_t + 0.6*v return the originals exactly, since the two maps are transposed rotations.
Constraints
0 < alpha_bar <= 1.- Both functions use the same two coefficients and -- no divisions.
- Mind the signs:
compute_vsubtracts the term, and the recovery subtracts the term while the recovery adds it. - The two maps must round-trip: with
x_t = sqrt(ab)*x0 + sqrt(1-ab)*eps, callingv_to_x0_eps(x_t, compute_v(x0, eps, ab), ab)returns(x0, eps). - Do not round inside either function.