Let’s take a look at how we can introduce stochasticity into our flow matching sampling process and compare the results to the commonly used Euler discrete sampling. We’ll mainly cover 1 ODE sampler and 3 SDE samplers with the ODE sampler as our baseline for comparison. All results/generations use the Flux-Krea-Dev model with a standard 28 step schedule including the timestep shift.
The standard flow matching ODE is given as $dZ_{t} = v(Z_{t}, t).dt$ , which can be solved by taking successive Euler steps $Z_{t_{i+1}} = Z_{t_{i}} + (t_{i+1} - t_{i}).v(Z_{t_{i}},t_{i})$ in the time interval $[0,1]$ over N timesteps.
<aside> <img src="/icons/exclamation-mark_yellow.svg" alt="/icons/exclamation-mark_yellow.svg" width="40px" />
Diffusion and Flow use opposite time settings. i.e In diffusion, at $t=0$ we assume the sample is a clean image whereas at $t=1$ it is a random noise sample drawn from a Gaussian distribution $\mathcal{N}(0, I)$ whereas $t=0$ is where the sample is gaussian noise in the case of Flow matching. In practice we reformulate the samplers to match traditional diffusion assumption. So while the theory and diagrams below use the flow matching convention, the code is adapted for the diffusion convention. We’ll see how that is done below.
</aside>
Using a diffusion/flow model to solve the ODE numerically introduces errors both from the model approximation as well the discrete timestep updates. These errors can cause the final generation to deviate from the expected distribution. Adding a Langevin dynamics term acts as a feedback mechanism to correct the drift. This is given as
$$ dZ_{t} = \underbrace{v(Z_{t},t).dt}{Flow} + \underbrace{ \sigma^{2}{t}\nabla log{\rho_{t}(Z_{t})}dt + \sqrt{2\sigma^{2}{t}}dW{t}}_{Langevin} $$
where $\rho_{t}$ is the density function of $Z_{t}$ following the ODE model. $\sigma_t$ is a non-negative sequence and $W_{t}$ is standard Brownian motion. Though we add the Langevian term, the marginal of both the initial ODE and the modified SDE above are the same i.e the distribution of $Z_{t}$ is the same even though the processes are different. This is how stochastic samplers can fit into the flow matching framework for generative model sampling. For a more detailed explanation see: https://rectifiedflow.github.io/blog/2024/diffusion/
The stochastic samplers use different methods to approximate the Langevin term in the equation above to add a regularizing effect to reduce the solver error.
This is a slight modification of the usual Euler step where instead of predicting $z_{t+1}$ from $z_{t}$, we instead predict the clean image at $t=1$ based on the current value $z_t$. Then we randomly sample a gaussian noise $z^{`}{0}$ and interpolate between them to predict the value $z{t_{i+1}}$.
Stochastic Euler Step
The Stochastic curved Euler takes a similar approach. It too estimates the clean image from the current value $z_{t_{i}}$ and instead of picking a random noise to interpolate from, it estimates the original noise that resulted in the current value after $t_i$ timesteps and uses a noise replacement factor, $\alpha_i$ to arrive at a refreshed initial noise $z^{ref}{t_0}$. Then it interpolates to $t=t{i+1}$ between $[z^{ref}{t}, z^{`}{t_1}]$ to calculate $z_{t_{i+1}}$
Stochastic Curved Euler Sampler Step
