Adaptive LayerNorm

Replaces LayerNorm’s learned $\gamma ,\beta$ with values predicted from a conditioning vector (timestep, class, text embedding), the FiLM-style mechanism that DiT and most diffusion transformers use to inject conditioning.

Equation

A linear head maps the condition to shift, scale, and gate, $(\beta ,\gamma ,\alpha )=Wc+b$:

$$\mathrm{adaLN}(x,c)=\gamma \odot \frac{x-\mu (x)}{\sigma (x)}+\beta ,x\leftarrow x+\alpha \odot \mathrm{Block}(\mathrm{adaLN}(x,c))$$

In adaLN-Zero, $W,b$ are initialized to zero so $\alpha =0$ and every block starts as the identity.

From scratch

import torch
import torch.nn as nn
from einops import rearrange, reduce
 
class AdaLNZero(nn.Module):
    def __init__(self, d, d_cond):
        super().__init__()
        self.proj = nn.Linear(d_cond, 3 * d)
        nn.init.zeros_(self.proj.weight)
        nn.init.zeros_(self.proj.bias)
 
    def forward(self, x, c, block, eps=1e-6):
        # x: (b, n, d); c: (b, d_cond)
        shift, scale, gate = rearrange(self.proj(c), "b (p d) -> p b 1 d", p=3)
        mu, var = x.mean(-1, keepdim=True), x.var(-1, keepdim=True, unbiased=False)
        h = (x - mu) * torch.rsqrt(var + eps) * (1 + scale) + shift
        return x + gate * block(h)

Notes

• Predict 1 + scale rather than scale so zero-init gives an identity modulation.
• Zero-init of gate makes deep stacks trainable at high LR: the main reason adaLN-Zero beats plain adaLN.
• Costs $3d\cdot d_{\text{cond}}$ params per block; DiT spends ~30% of params here.
• FiLM is the same idea without the normalization step.