Conditioning by modulation:
adaLN-Zero and SPADE

adaLN — adaptive layer norm

Standard LayerNorm normalizes activations, then applies fixed learned affine parameters. adaLN makes those parameters dynamic: the scale and shift are produced on the fly from a conditioning vector c — typically a timestep embedding plus a class embedding, passed through a small MLP.

The (1 + γ) form is deliberate — it centers the multiplicative scale at \(1\) (identity), so training starts neutral instead of at zero scale. In DiT this happens twice per block (before attention, before the FFN), giving four parameters per block: \(\gamma_1,\beta_1\) and \(\gamma_2,\beta_2\).

# c: condition  [B, cond_dim]   x: tokens [B, T, D]
shift1, scale1, shift2, scale2 = MLP(SiLU(c)).chunk(4, -1)

x_mod = LayerNorm(x) * (1 + scale1) + shift1
x = x + Attention(x_mod)          # plain residual

x_mod = LayerNorm(x) * (1 + scale2) + shift2
x = x + FFN(x_mod)

adaLN-Zero — adding the gate

adaLN-Zero (Peebles & Xie, DiT 2023) adds one element: a per-block gating scalar \(\alpha\) that multiplies each sub-block's output before the residual add. The linear layer producing \(\alpha\) is initialized to zero.

The zero-initialization

self.adaLN_modulation = nn.Sequential(
    nn.SiLU(),
    nn.Linear(cond_dim, 6 * hidden_size),
)

# ★ the critical line
nn.init.zeros_(self.adaLN_modulation[-1].weight)
nn.init.zeros_(self.adaLN_modulation[-1].bias)

An SE-like side note

The \(\alpha\) pathway is structurally close to a Squeeze-and-Excitation block — a side branch that rescales the main features channel-wise. Omitting bias terms:

This buys extra expressivity even before zero-init enters the picture.

Why zero-init wins

The ICLR-2025 analysis (OpenReview E4roJSM9RM) decoupled three candidate explanations and ranked them:

As training runs, the \(\alpha\) weights drift from a spike at zero toward a Gaussian-like distribution centered at zero. Zero-init drops you near the center of that target already. By an entropy argument, adaLN-Zero's distribution simply spreads outward (entropy increases — the natural direction), while other inits must first contract toward zero, fighting the gradient.

# replace the zeros_ weight line with:
nn.init.normal_(self.adaLN_modulation[-1].weight, std=0.001)

SPADE — spatial modulation

SPADE (Park et al., CVPR 2019 — a.k.a. GauGAN) tackles a sibling problem: conditioning a generator on a segmentation mask without normalization erasing its spatial layout.

The problem

Run a mask through an encoder, then normalize in the decoder, and the semantic signal washes out — regions of the same label collapse to one normalized value, so spatial structure disappears. The fix must be spatially aware.

The mechanism

Instead of a fixed affine layer after normalization, SPADE predicts \(\gamma\) and \(\beta\) per spatial location from the mask, via a small CNN:

class SPADE(nn.Module):
    def __init__(self, ch, cond_ch):
        super().__init__()
        self.norm   = nn.InstanceNorm2d(ch, affine=False)
        self.shared = nn.Sequential(nn.Conv2d(cond_ch, 128, 3, padding=1), nn.ReLU())
        self.gamma  = nn.Conv2d(128, ch, 3, padding=1)
        self.beta   = nn.Conv2d(128, ch, 3, padding=1)

    def forward(self, x, mask):
        m = F.interpolate(mask, size=x.shape[2:], mode='nearest')
        f = self.shared(m)
        return self.gamma(f) * self.norm(x) + self.beta(f)

The connection

adaLN-Zero and SPADE are the same move at different spatial granularity: normalize to strip dataset statistics, then re-inject conditioning through a learned affine transform.

Synthesis — spatial adaLN-Zero

To condition a DiT on a binary mask while staying inside the adaLN family: replace the modulation MLP with a CNN that emits spatially-varying versions of all six parameters (including the \(\alpha\) gates), then modulate each patch token with the parameters at its own grid cell. This is SPADE's spatial idea carried into the transformer, with the gate retained.

1 · keep zero-init on the final conv

self.out_conv = nn.Conv2d(256, 6 * hidden_size, 1)
nn.init.zeros_(self.out_conv.weight)   # preserves the adaLN-Zero property
nn.init.zeros_(self.out_conv.bias)

2 · fuse timestep additively in feature space

spatial = self.mask_cnn(mask)                     # [B,256,H/P,W/P]
time    = self.time_proj(t_emb)[..., None, None] # [B,256,1,1]
fused   = spatial + time                          # broadcast -> global denoise signal

3 · spatial alignment is implicit

If mask and image share patch size \(P\), then CNN cell \((i,j)\) modulates image token \((i,j)\) exactly. No extra positional encoding needed on the modulation maps.

Choosing an approach