Consistency Model: Single-step generation via consistency function.

Implements the Consistency Model from "Consistency Models" (Song et al., ICML 2023). Learns a function f(x_t, t) that maps any point on a probability flow ODE trajectory directly to its origin, enabling single-step generation without iterative denoising.

Key Innovation: Self-Consistency Property

A consistency function f satisfies: for all t, t' on the same trajectory, f(xt, t) = f(x{t'}, t'). In particular, f(x_t, t) = x_0 for all t.

Diffusion (many steps):
  x_T -> x_{T-1} -> ... -> x_1 -> x_0    (N denoising steps)

Consistency Model (one step):
  x_T -> x_0    (single forward pass!)

Or few-step refinement:
  x_T -> x_0' -> add_noise(x_0', t') -> x_0''    (2 steps, better quality)

Training Approaches

1. Consistency Distillation (CD): Distill from a pre-trained diffusion model
2. Consistency Training (CT): Train from scratch without a teacher

Both enforce: f(x_{t+1}, t+1) = f(x_t, t) for adjacent timesteps.

Architecture

Input (x_t, sigma_t)
      |
      v
+-----------------------+
| Skip Connection       |
| c_skip(t) * x_t +    |
| c_out(t) * F(x_t, t) |
+-----------------------+
      |
      v
Output: predicted x_0

The skip connection ensures the boundary condition f(x, sigma_min) = x.

Usage

model = ConsistencyModel.build(
  input_dim: 64,
  hidden_size: 256,
  num_layers: 4
)

# Single-step generation
x_0 = ConsistencyModel.single_step_sample(model, params, noise)

# Multi-step refinement
x_0 = ConsistencyModel.multi_step_sample(model, params, noise, steps: 3)

Reference

• Paper: "Consistency Models"
• arXiv: https://arxiv.org/abs/2303.01469