Architecture Considerations for Consistency Models

While the core idea of consistency models lies in the training objective and the sampling process, the choice of neural network architecture remains a significant factor influencing performance, training stability, and inference speed. Often, architectures successfully employed for standard diffusion models serve as a strong starting point for consistency models, leveraging well-understood design patterns.

Leveraging Existing Diffusion Architectures

Consistency models, whether trained via distillation or standalone, typically adopt the same backbone architectures proven effective in diffusion modeling. This often means using:

1. U-Net Variants: Deep convolutional networks with skip connections, adept at capturing multi-scale spatial information in images. Advanced U-Nets incorporating self-attention and cross-attention mechanisms (as discussed in Chapter 2) are commonly used, especially for high-resolution image generation.
2. Transformer Architectures: Models like the Diffusion Transformer (DiT, covered in Chapter 3), which replace convolutional layers with transformer blocks, are also viable. These can excel at modeling long-range dependencies and scaling model capacity.

The rationale for reusing these architectures is straightforward: they are already optimized for the task of predicting outputs (like noise $\epsilon$ or data $x_0$) based on noisy inputs $x_t$ and time $t$. Consistency models simply repurpose this predictive capability towards learning the consistency mapping $f_{\theta}(x_t, t) \approx x_0$.

Time Embedding Integration

Just like standard diffusion models, consistency models require information about the current point in the "denoising" process, represented by time $t$ (or noise level $\sigma$). The methods for embedding $t$ are usually inherited from the base diffusion architecture:

• Sinusoidal Embeddings: The standard approach using fixed sinusoidal positional embeddings.
• Learnable Fourier Features: Variations where frequency and phase are learned.
• MLP Projection: Projecting the time embedding into a feature space suitable for integration into the network layers (e.g., via addition or adaptive normalization).

For continuous-time consistency models, ensuring the network can smoothly handle continuous $t$ values is important. The embedding mechanism must effectively represent this continuous variable.

Conditioning Mechanisms

Conditional generation (e.g., based on text prompts or class labels) in consistency models typically follows the same strategies used in the parent diffusion models:

• Cross-Attention: Integrating conditioning information (like text embeddings) into transformer or U-Net attention layers.
• Adaptive Normalization (AdaLN/AdaGN): Modulating normalization layer parameters based on conditioning embeddings.
• Concatenation: Directly concatenating conditioning embeddings to input feature maps (less common in state-of-the-art models).

The choice of conditioning mechanism depends heavily on the base architecture (U-Net vs. Transformer) and the type of conditioning signal. The consistency training objective itself doesn't necessitate fundamental changes to how conditioning is incorporated into the network's forward pass.

Network Capacity and Distillation

A significant consideration arises when using consistency distillation: Does the consistency model (student) need the same architectural complexity and parameter count as the pre-trained diffusion model (teacher)?

• Distillation: Often, a smaller student network can effectively learn the consistency mapping from a larger teacher. This is a primary advantage, as it leads to a faster and more memory-efficient model for inference. The distillation process transfers the essential mapping, potentially allowing for architectural simplification (e.g., fewer layers, smaller hidden dimensions).
• Standalone Training: When training from scratch, the network must have sufficient capacity to learn the complex mapping from noisy inputs $x_t$ across various $t$ directly to $x_0$. The required capacity will depend on the data complexity and desired output quality, potentially mirroring that of a standard diffusion model trained for the same task.

Diagram illustrating the relationship between a larger teacher diffusion model and a potentially smaller student consistency model during distillation. Both often share architectural patterns but differ in size and prediction target.

Output Layer Configuration

Standard diffusion models are often parameterized to predict the noise $\epsilon$ added at step $t$. The network output $\epsilon_{\theta}(x_t, t)$ is then used to estimate $x_0$. Consistency models, by definition, aim to directly map any point $x_t$ on a trajectory to its origin $x_0$. Therefore, the output layer(s) of a consistency model $f_{\theta}(x_t, t)$ are typically configured to directly produce an output with the same dimensions and value range as the input data $x_0$. This might involve adjustments to the final activation function (e.g., using tanh if data is normalized to [-1, 1]) compared to an $\epsilon$-predicting diffusion model.

Summary of Architectural Points

• Foundation: U-Nets and Transformers from standard diffusion models are excellent starting points.
• Time/Conditioning: Reuse effective embedding and conditioning techniques (sinusoidal/MLP for time, cross-attention/AdaLN for conditioning).
• Capacity: Distillation allows for potentially smaller, faster student models compared to the teacher. Standalone training requires sufficient capacity.
• Output: The network's final layer must be adapted to directly predict $x_0$ (or a value representing it).

While consistency models introduce a new approach for faster sampling, their architectural building blocks often remain familiar, allowing practitioners to use existing knowledge of diffusion model architectures while focusing on the details of the consistency objective and its implications for model size and output representation.