Derivation of the Flow-SDE step from the flow-matching ODE

Setting. Fast-WAM's forward process is $x_\sigma = (1-\sigma ) x_0 + \sigma \varepsilon$ with $\varepsilon \sim \mathcal{N}(0, I)$ and $\sigma \in [0, 1]$, and the network is trained to predict $v = \varepsilon - x_0$ (training_target in the scheduler). Differentiating the forward process in σ gives $dx/d\sigma = \varepsilon - x_0 = v$, so the deterministic sampler $x_{k+1} = x_k + v \Delta _k$ is the Euler discretisation of the ODE $dx = v d\sigma$, integrated from σ = 1 (noise) down to σ = 0 (data). This is the same convention as π₀ and πRL, whose $\tau$ plays the role of σ.

Step 1 — the score in terms of the velocity. Write $m = \mathbb{E}[x_0 \mid x_\sigma ]$ and $e = \mathbb{E}[\varepsilon \mid x_\sigma ]$ for the posterior means. The learned velocity is the posterior mean of $\varepsilon - x_0$, so $v = e - m$, and taking the posterior mean of the forward process gives $x_\sigma = (1-\sigma ) m + \sigma e$. Solving these two linear equations for $e$:

$$\begin{aligned} m &= (x_\sigma - \sigma e) / (1-\sigma) \\ v &= e - (x_\sigma - \sigma e)/(1-\sigma) = (e - x_\sigma)/(1-\sigma) \quad \Rightarrow \quad e = x_\sigma + (1-\sigma) v . \end{aligned}$$

Tweedie's formula for a Gaussian perturbation of scale σ states $\mathbb{E}[\varepsilon \mid x_\sigma ] = -\sigma \nabla \log p_\sigma (x_\sigma )$, hence

$$\nabla_x \log p_\sigma(x) = -\,e / \sigma = -\big( x + (1-\sigma) v \big) / \sigma = -\,x/\sigma - \big((1-\sigma)/\sigma\big) v .$$

This matches the expression πRL prints, $\nabla \log q_\tau(A^\tau) = -A^\tau/\tau - ((1-\tau)/\tau)\, v^\tau$.

Step 2 — the equal-marginal SDE. Song et al. (2021) show that if a forward-in-σ diffusion $dx = f d\sigma + g dw$ has marginals $p_\sigma$, then the reverse-direction SDE $dx = [f - g^2 \nabla \log p_\sigma ] d\sigma + g d\bar{w}$ and the probability-flow ODE $dx = [\,f - (g^2/2)\, \nabla \log p_\sigma\,]\, d\sigma$ share those same marginals. Given the ODE drift $v = f - (g^2/2)\, \nabla \log p$, the reverse SDE drift is therefore $v - (g^2/2)\, \nabla \log p$, for any non-negative $g(\sigma)$:

$$dx = \big[\, v - (g(\sigma)^2/2)\, \nabla \log p_\sigma(x) \,\big]\, d\sigma + g(\sigma)\, d\bar{w} .$$

Substituting the score from Step 1:

$$dx = \big[\, v + \big(g(\sigma)^2 / (2\sigma)\big) \big( x + (1-\sigma) v \big) \,\big]\, d\sigma + g(\sigma)\, d\bar{w} .$$

This is precisely the SDE both πRL ($dA^\tau = [\,v^\tau + \sigma_\tau^2/(2\tau)(A^\tau + (1-\tau)v^\tau)\,]\, d\tau + \sigma_\tau\, dw_\tau$) and Flow-GRPO state, with their $\sigma_\tau$ equal to $g(\sigma)$ here.

Step 3 — the noise schedule. πRL and Flow-GRPO choose $g(\sigma) = a \sqrt{\sigma/(1-\sigma)}$. With that choice the drift coefficient simplifies: $g^2/(2\sigma) = a^2/(2(1-\sigma))$. The choice is a convention rather than a necessity; any $g$ preserves the marginals. It diverges at σ = 1, which is why the first denoising step must be deterministic.

Step 4 — discretisation. Euler–Maruyama with Fast-WAM's step $\Delta _k = \sigma _{k+1} - \sigma _k < 0$ gives

$$\begin{aligned} \mu_k &= x_k + \big[\, v_k + \big(g_k^2 / (2\sigma_k)\big) \big( x_k + (1-\sigma_k) v_k \big) \,\big] \Delta_k \\ \mathrm{std}_k &= g_k \sqrt{|\Delta_k|} \\ x_{k+1} &\sim \mathcal{N}\big( \mu_k, \mathrm{std}_k^2 I \big) , \end{aligned}$$

which is the update in Chapter 5. Setting $a = 0$ recovers the deterministic step. The Gaussian transition is what makes the log-probability tractable; the parameters enter only through v_k, so $\partial \log \mathcal{N} / \partial \varphi = \big((x_{k+1} - \mu_k)/\mathrm{std}_k^2\big) \cdot \partial \mu_k / \partial \varphi$, and $\partial \mu_k / \partial v_k = [\,1 + (g_k^2/(2\sigma_k))(1-\sigma_k)\,]\, \Delta_k$.

Closed-form KL used in Chapter 9. For diagonal Gaussians, per dimension,

$$\mathrm{KL}\big( \mathcal{N}(\mu_1, s_1^2) \,\|\, \mathcal{N}(\mu_2, s_2^2) \big) = \tfrac{1}{2} \big[ \log(s_2^2/s_1^2) + (s_1^2 + (\mu_1-\mu_2)^2)/s_2^2 - 1 \big] ,$$

and substituting $\mu _1 = z^\theta , s_1^2 = \Sigma _\theta , \mu _2 = z^{\mathrm{old}}, s_2^2 = \Sigma _{\mathrm{old}}$ gives the expression in §9.4.