Normalization Techniques (GroupNorm, AdaLN)

Proper normalization is significant for stable and efficient training of deep neural networks, and diffusion models are no exception. While Batch Normalization (BatchNorm) is ubiquitous in many vision tasks, its reliance on batch statistics can be problematic for diffusion models due to:

1. Small Batch Sizes: Training large diffusion models often necessitates small batch sizes due to memory constraints, making batch statistics noisy and less reliable.
2. Varying Statistics: The statistics of activations can change significantly across different timesteps ($t$) in the diffusion process, potentially violating BatchNorm's assumptions.

Therefore, alternative normalization techniques that are independent of the current batch are preferred within the U-Net architectures used for diffusion. Two prominent methods are Group Normalization and Adaptive Layer Normalization.

Group Normalization (GroupNorm)

Group Normalization (GroupNorm) offers a middle ground between Layer Normalization (which normalizes across all channels) and Instance Normalization (which normalizes each channel independently). It operates by dividing the channels into a predefined number of groups and computing the mean and variance for normalization within each group.

Let $x$ be the input feature map with shape $(N, C, H, W)$, where $N$ is the batch size, $C$ is the number of channels, and $H, W$ are spatial dimensions. GroupNorm divides the $C$ channels into $G$ groups, each containing $C/G$ channels (assuming $C$ is divisible by $G$). Normalization is then applied independently to each group across the $(C/G, H, W)$ dimensions.

The mean $\mu_g$ and standard deviation $\sigma_g$ for a group $g$ are calculated as:

$$\mu_g = \frac{1}{(C/G)HW} \sum_{k \in S_g} \sum_{h=1}^H \sum_{w=1}^W x_{n,k,h,w}$$ $$\sigma_g = \sqrt{\frac{1}{(C/G)HW} \sum_{k \in S_g} \sum_{h=1}^H \sum_{w=1}^W (x_{n,k,h,w} - \mu_g)^2 + \epsilon}$$

where $S_g$ is the set of channel indices belonging to group $g$, and $\epsilon$ is a small constant for numerical stability. The normalized output $\hat{x}$ is then:

$$\hat{x}_{n,c,h,w} = \frac{x_{n,c,h,w} - \mu_g}{\sigma_g}$$

where $c$ is in group $g$. Finally, learnable scale ($\gamma$) and shift ($\beta$) parameters are applied per channel:

$$y_{n,c,h,w} = \gamma_c \hat{x}_{n,c,h,w} + \beta_c$$

The main advantage of GroupNorm is its independence from the batch size $N$. Its statistics are computed only over channel groups and spatial dimensions for each sample independently, making it effective to small or fluctuating batch sizes common in diffusion model training. It is typically inserted within residual blocks, often before the non-linear activation function. The number of groups $G$ is a hyperparameter, often set to a value like 32.

Data flow for Group Normalization, highlighting the grouping of channels before computing statistics.

Adaptive Layer Normalization (AdaLN)

While GroupNorm provides stable normalization, diffusion models often benefit from conditional normalization, where the normalization process itself is modulated by conditioning information, most importantly the diffusion timestep $t$. Adaptive Layer Normalization (AdaLN), often used in the form of AdaLN-Zero, achieves this.

AdaLN builds upon Layer Normalization (LayerNorm), which normalizes features across all channels for a given sample. The core idea is to make the learnable scale ($\gamma$) and shift ($\beta$) parameters functions of the conditioning input, such as a timestep embedding $\mathbf{e}_t$.

Standard LayerNorm applies:

$$y = \gamma \frac{x - \mu}{\sigma} + \beta$$

where $\mu$ and $\sigma$ are computed across all channels $C$ (and spatial dimensions if applied to feature maps).

In AdaLN, the scale and shift are predicted from the conditioning embedding:

1. Compute standard Layer Normalization (or sometimes Group Normalization) on the input features $x$.
2. Process the conditioning embedding $\mathbf{e}_t$ (often after passing it through a non-linearity like SiLU) using a linear layer (or a small MLP) to predict modulation parameters, typically separate scale and shift components. Let's call them $\gamma_t$ and $\beta_t$.
3. Apply these predicted parameters to the normalized features $\hat{x}$: $$y = \gamma_t \hat{x} + \beta_t$$

A common variant, particularly in diffusion models, is AdaLN-Zero. This initializes the modulation parameters such that they initially have no effect, promoting stability early in training. The output block applying the modulation is often structured as:

$$y = \hat{x} * (1 + \gamma_t) + \beta_t$$

Here, the linear layer predicting $\gamma_t$ and $\beta_t$ is initialized to output zeros. Thus, initially, $y = \hat{x}$, and the network learns to modulate the features gradually as training progresses.

AdaLN (or AdaLN-Zero) is particularly effective in diffusion U-Nets because it allows the network to tailor its feature processing based on the noise level (timestep $t$). It's frequently used within residual blocks, often applied after the main convolution and normalization (like GroupNorm), specifically controlling the output of the residual branch before it's added back to the identity connection.

Data flow for Adaptive Layer Normalization (AdaLN-Zero variant). The conditioning embedding is used to predict scale and shift parameters that modulate the normalized features.

Choosing and Combining Normalization

In advanced U-Net architectures for diffusion:

• GroupNorm is often used as the primary normalization layer within convolutional blocks due to its batch independence and proven effectiveness in vision tasks. It stabilizes the activations before non-linearities or attention layers.
• AdaLN (or AdaLN-Zero) is frequently employed specifically for injecting conditional information (like timestep $t$ or class embeddings). It typically modulates the output of residual blocks or specific feature maps where conditional control is desired. It can sometimes replace the standard learnable scale/shift of GroupNorm or LayerNorm, or be applied in addition.

Using both GroupNorm for general stability and AdaLN for conditional modulation provides a powerful combination for building high-performance diffusion models. The exact placement and configuration (number of groups for GroupNorm, architecture of the MLP predicting AdaLN parameters) are design choices that can impact model performance and training dynamics.