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 $x_0$ and $\varepsilon$.

Background

Write $\sqrt{\bar{\alpha}_t} = \cos\phi$ and $\sqrt{1-\bar{\alpha}_t} = \sin\phi$, so that $x_t = \cos\phi\,x_0 + \sin\phi\,\varepsilon$ traces a circular arc between the clean sample and pure noise. The velocity is the derivative of that arc:

$v = \sqrt{\bar{\alpha}_t}\,\varepsilon - \sqrt{1-\bar{\alpha}_t}\,x_0$

$v$ and $x_t$ are orthogonal unit-scaled directions, which makes the inverse a rotation by the same angle:

$\hat{x}_0 = \sqrt{\bar{\alpha}_t}\,x_t - \sqrt{1-\bar{\alpha}_t}\,v, \qquad \hat{\varepsilon} = \sqrt{1-\bar{\alpha}_t}\,x_t + \sqrt{\bar{\alpha}_t}\,v$

Note there is no division by $\sqrt{\bar{\alpha}_t}$ anywhere. That is the whole point: $\varepsilon$-prediction degenerates at $\bar{\alpha}_t \to 0$ (the target becomes trivially $x_t$ itself, and recovering $x_0$ divides by something tiny), while $x_0$-prediction degenerates at $\bar{\alpha}_t \to 1$. The $v$ 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 $\bar{\alpha}_t$ 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.