Consistency Model Training: Standalone Approach

Consistency Training (CT) offers a method to train consistency models from scratch, without relying on a pre-trained diffusion model. This contrasts with other techniques, such as consistency distillation, which often accelerate existing diffusion models. This standalone approach directly optimizes a neural network $f_ heta(x, t)$ to satisfy the consistency property: points along the same Probability Flow (PF) ODE trajectory should map to the same origin $x_0$.

The core challenge lies in enforcing this property without having explicit access to the ODE trajectories defined by a pre-trained model. Standalone consistency training cleverly addresses this by using numerical ODE solvers and the model's own evolving predictions during training.

The Consistency Training Objective

The objective is to learn a function $f_\theta(x, t)$ such that for any valid time $t \in [\epsilon, T]$ and any point $x_t$ on an ODE trajectory originating from $x_0$, we have $f_\theta(x_t, t) \approx x_0$. Here, $\epsilon$ is a small minimum time step close to 0, and $T$ is the maximum diffusion time.

To achieve this, CT minimizes a loss function that encourages consistency between outputs of the model evaluated at adjacent time steps along estimated ODE trajectories. Consider two consecutive time steps $t_n$ and $t_{n+1}$ from a discretized schedule $t_1, t_2, ..., t_N$, where $t_1 \approx \epsilon$ and $t_N = T$. Let $x_{t_n}$ and $x_{t_{n+1}}$ be points on the same (estimated) ODE trajectory. The consistency loss enforces that the model's predictions for $x_0$ from these two points are similar:

$$L_{CT}(\theta, \theta^-) = \mathbb{E}_{x_0 \sim p_{data}, n \sim \mathcal{U}(1, N-1), z \sim \mathcal{N}(0, I)} \left[ \lambda(t_{n+1}) d(f_\theta(x_{t_{n+1}}, t_{n+1}), f_{\theta^-}(x_{t_n}, t_n)) \right]$$

Let's break down this objective:

1. Sampling: We sample a real data point $x_0$, a random time step index $n$, and a noise vector $z$.
2. Trajectory Points: We need to obtain $x_{t_n}$ and $x_{t_{n+1}}$. These points are estimated by taking one step of a numerical ODE solver (like first-order Euler or second-order Heun) starting from a point perturbed from $x_0$. For example, using the Euler method for the PF ODE $\frac{dx}{dt} = \frac{1}{2} \beta(t) (\nabla_x \log p_t(x) + x)$:
   • $x_{t_n} \approx x_0 + \sqrt{\alpha(t_n)^2 / (1 - \alpha(t_n)^2)} z$ (Approximation based on diffusion process properties)
   • $x_{t_{n+1}}$ is obtained by taking one step from $x_{t_n}$ using the ODE solver. This requires estimating the score $\nabla_x \log p_t(x)$ at $(x_{t_n}, t_n)$. Crucially, CT often uses the score implied by the current model's estimate of $x_0$ or related techniques, effectively bootstrapping the process.
3. Consistency Function $f_\theta$: This is the neural network being trained. It takes a noisy input $x_t$ and time $t$ and predicts the corresponding $x_0$.
4. Target Network $f_{\theta^-}$: A slowly updated, exponential moving average (EMA) of the main network's parameters ($\theta$) is used for the "target" prediction $f_{\theta^-}(x_{t_n}, t_n)$. This stabilizes training, similar to techniques used in reinforcement learning or BYOL. The update rule is typically $\theta^- \leftarrow \mu \theta^- + (1 - \mu) \theta$, where $\mu$ is a momentum coefficient close to 1 (e.g., 0.999).
5. Distance Metric $d(\cdot, \cdot)$: This measures the difference between the two predictions. Common choices include the L1 loss, L2 loss (Mean Squared Error), or perceptual losses like LPIPS.
6. Weighting Function $\lambda(t_{n+1})$: This function weights the loss based on the time step. It often prioritizes matching at earlier time steps (closer to the data) or uses weighting schemes derived from diffusion model theory.

Estimating Trajectories and Scores

The most significant difference from distillation is how $x_{t_n}$ and $x_{t_{n+1}}$ (points on the same trajectory) are obtained. Since there's no teacher model providing the score $\nabla_x \log p_t(x)$, CT must estimate it.

One common approach involves using the relationship between the score function and the conditional expectation $\mathbb{E}[x_0 | x_t]$. If the model $f_\theta(x_t, t)$ estimates $x_0$, it can be used to approximate the score needed for the ODE solver step that generates $x_{t_{n+1}}$ from $x_{t_n}$. This creates a self-supervised loop where the model refines its understanding of the trajectories and the consistency mapping simultaneously.

Training Algorithm Overview

The standalone consistency training process typically follows these steps in each iteration:

1. Sample Data: Draw a mini-batch of data points $\{x_0^{(i)}\}$ from the true data distribution $p_{data}$.
2. Sample Timesteps: For each $x_0^{(i)}$, sample a time index $n^{(i)} \sim \mathcal{U}(1, N-1)$. Let $t_n = \text{schedule}[n]$ and $t_{n+1} = \text{schedule}[n+1]$.
3. Generate Trajectory Pair:
   • Generate noise $z^{(i)} \sim \mathcal{N}(0, I)$.
   • Estimate $x_{t_n}^{(i)}$ (e.g., using $x_0^{(i)}$ and $z^{(i)}$ based on the diffusion process definition).
   • Estimate the score at $(x_{t_n}^{(i)}, t_n)$, potentially using $f_\theta$ or $f_{\theta^-}$.
   • Use a numerical ODE solver (e.g., one step of Heun's method) with the estimated score to compute $x_{t_{n+1}}^{(i)}$ from $x_{t_n}^{(i)}$.
4. Compute Model Outputs:
   • Online network prediction: $y_{n+1}^{(i)} = f_\theta(x_{t_{n+1}}^{(i)}, t_{n+1})$
   • Target network prediction: $y_n^{(i)} = f_{\theta^-}(x_{t_n}^{(i)}, t_n)$ (Stop gradient through target network).
5. Calculate Loss: Compute the consistency loss $L_{CT}$ using the distance metric $d$ and weighting $\lambda(t_{n+1})$: $$\text{Loss} = \frac{1}{B} \sum_{i=1}^B \lambda(t_{n+1}^{(i)}) d(y_{n+1}^{(i)}, y_n^{(i)})$$ where $B$ is the batch size.
6. Gradient Update: Compute the gradient $\nabla_\theta L_{CT}$ and update the online network parameters $\theta$ using an optimizer (e.g., Adam).
7. Update Target Network: Update the target network parameters $\theta^-$ using EMA: $\theta^- \leftarrow \mu \theta^- + (1 - \mu) \theta$.

Diagram illustrating the standalone consistency training loop for a single data point. The process involves sampling data and time, generating a pair of points along an estimated ODE trajectory using score estimates, computing outputs from the online and target networks, and updating the networks based on the consistency loss.

Architectural Approaches

The network architecture $f_\theta(x, t)$ used in standalone CT is often similar to those employed in standard diffusion models or consistency distillation, such as U-Net variants or Transformers (like DiT). The main requirement is the ability to process a noisy input $x_t$ and a time embedding $t$ to produce an estimate of $x_0$. Adaptations might involve modifications to how time is embedded or how conditioning information (if any) is integrated.

Comparison to Distillation

• Pros of Standalone Training:
  • Independence: Does not require a pre-trained diffusion model, saving the computational cost and time associated with training one first.
  • Potential for Higher Fidelity: Might avoid limitations inherent in the distillation process where the student model tries to mimic a potentially imperfect teacher.
• Cons of Standalone Training:
  • Stability: Can be harder to stabilize than distillation, as it relies on bootstrapping the score estimation. Requires careful tuning of hyperparameters like the learning rate, EMA decay rate ($\mu$), and loss weighting ($\lambda(t)$).
  • Convergence Speed: May require more training iterations to converge compared to distillation, which starts with guidance from a strong teacher model.

Standalone consistency training represents a significant step towards efficient generative modeling, enabling the creation of fast, few-step or even single-step generative models directly from data. While it presents unique training challenges compared to distillation, its independence from pre-trained models makes it an appealing and powerful technique in the generative modeling toolkit.