What numerical solvers (like Euler–Maruyama) are used in continuous-time diffusion models?

Numerical solvers like the Euler–Maruyama method are commonly used to simulate continuous-time diffusion models, which describe stochastic processes governed by differential equations with random noise. These solvers approximate the trajectories of stochastic differential equations (SDEs) by discretizing time into small steps. The Euler–Maruyama method, a straightforward extension of the deterministic Euler method, is popular due to its simplicity. It updates the system state by combining a deterministic drift term (based on the current state and time) and a stochastic diffusion term (scaled by a Wiener process increment). For example, given an SDE ( dX_t = \mu(X_t, t)dt + \sigma(X_t, t)dW_t ), the Euler–Maruyama update is ( X_{t+1} = X_t + \mu(X_t, t)\Delta t + \sigma(X_t, t)\sqrt{\Delta t} \cdot \mathcal{N}(0,1) ), where ( \Delta t ) is the step size. This method is computationally efficient but has a weak convergence order of 1, meaning errors grow linearly with step size. Despite its limitations, it’s widely used in diffusion models for tasks like training and sampling due to its ease of implementation.

Other solvers, such as the Milstein method and stochastic Runge-Kutta variants, improve accuracy by accounting for higher-order terms in the SDE expansion. The Milstein method adds a correction term involving the derivative of the diffusion coefficient ( \sigma ), which reduces the error caused by the stochastic term. For instance, its update rule includes ( \frac{1}{2}\sigma(X_t, t)\partial_x\sigma(X_t, t)(\Delta W_t^2 - \Delta t) ). This raises the convergence order to 1.0 in the strong sense (pathwise accuracy), making it more precise for systems with non-negligible noise dependencies. Stochastic Runge-Kutta methods, like those of orders 1.5 or 2.0, further generalize these ideas by using intermediate steps to approximate higher-order terms without requiring explicit derivatives. These methods are useful in diffusion models where precise trajectory simulation is critical, such as in physics-informed generative models or when fine-grained control over the noise process is needed.

When choosing a solver, developers balance computational cost, accuracy, and stability. Euler–Maruyama is often sufficient for basic diffusion model training, where exact paths matter less than overall distributional accuracy. For example, in diffusion-based image generation, the cumulative effect of many small steps can tolerate individual step errors. However, for applications like rare event simulation or systems with stiff dynamics (rapidly changing terms), higher-order methods like Milstein or adaptive step-size solvers (e.g., the Kloeden-Platen scheme) may be necessary. Adaptive methods adjust step sizes dynamically to control error, which is useful when the SDE’s behavior varies significantly across time. Libraries like torchsde or Diffrax implement these solvers, allowing developers to experiment with trade-offs. Ultimately, the choice depends on the problem’s sensitivity to approximation errors and the available computational resources.