adaLN-Zero and Classifier-Free Guidance in PyTorch

In this post:

• Introduction
• Step 1/5: Give the Model a Label
• Step 2/5: Sometimes Hide the Label
• Step 3/5: Classifier-Free Guidance at Sampling Time
• Step 4/5: adaLN-Zero: Start Every Block as "Do Nothing"
• Step 5/5: The Before/After on TinyDiT
• Recap

Introduction

So far TinyDiT makes some image. You can’t ask it for a cat or a ship. Today it learns to follow an instruction: a CIFAR-10 class label.

This is a classifier-free guidance PyTorch tutorial in three parts: how to feed a label into the model, how one training trick (sometimes hiding the label) gives us a free “guidance” knob at sampling time, and adaLN-Zero, the small initialisation change that DiT found works best for conditioning. Same as RMSNorm, QK-Norm, SwiGLU, and 2D RoPE in a Tiny DiT, every claim gets a measured before/after on TinyDiT.

Step 1/5: Give the Model a Label

In this step: We turn a class id into a learned vector and add it to the time vector, so adaLN carries it into every block.

A class label is just an integer from 0 to 9. We turn it into a learned vector with nn.Embedding (a lookup table, one row per class) and add it to the time vector c.

Every block already reads c through adaLN (post 9), so the label reaches every block without any new wiring:

\begin{aligned} c &= \underbrace{\operatorname{TimestepEmbedding}(t)}_{\text{when}} \\ &\quad + \underbrace{\operatorname{LabelEmbedding}(y)}_{\text{what}} \end{aligned}

The table has 11 rows, not 10. Row 10 means “no label”, and it’s the key to everything below.

class LabelEmbedding(nn.Module):
    """One learned vector per class, plus one extra "no class" vector (index num_classes) for CFG."""

    def __init__(self, num_classes: int, dim: int):
        super().__init__()
        self.null_class = num_classes
        self.table = nn.Embedding(num_classes + 1, dim)

    def forward(self, y: torch.Tensor) -> torch.Tensor:
        return self.table(y)

class LabelEmbedding(nn.Module):
    """One learned vector per class, plus one extra "no class" vector (index num_classes) for CFG."""

    def __init__(self, num_classes: int, dim: int):
        super().__init__()
        self.null_class = num_classes
        self.table = nn.Embedding(num_classes + 1, dim)

    def forward(self, y: torch.Tensor) -> torch.Tensor:
        return self.table(y)

LabelEmbedding (with a null class): what, why, how

What it is: a layer that turns a class id into a vector the model can use.

Why it works: an embedding table learns one vector per class. Adding it to the time vector reuses adaLN to steer every block. The extra “no label” row lets one model be both conditional and unconditional.

How it works: nn.Embedding(num_classes + 1, D); c = t_embed(t) + y_embed(y); id num_classes means “no label”.

Step 2/5: Sometimes Hide the Label

In this step: We hide the label for 10% of training images, which trains a conditional and an unconditional model in one.

During training, we replace the label with “no label” (id 10) for a random 10% of the images in every batch. That one line, inside our flow-matching loss, trains two models in one: a conditional model (“make a cat”) and an unconditional one (“make any image”), sharing all their weights.

import torch
from tinydit.expert import ExpertConfig, ExpertTinyDiT

torch.manual_seed(0)
model = ExpertTinyDiT(ExpertConfig(dim=128, depth=4, num_heads=4, class_cond=True))
print("label table:", tuple(model.y_embed.table.weight.shape), "-> 10 classes + 1 'no class' row, null id =", model.null_class)

y = torch.randint(0, 10, (20,))
drop = torch.rand(20) < 0.1                       # 10% of the batch forgets its label
y_train = torch.where(drop, torch.full_like(y, model.null_class), y)
print("labels      :", y.tolist())
print("after drop  :", y_train.tolist())

x, t = torch.randn(20, 3, 32, 32), torch.rand(20)
print("v shape:", tuple(model(x, t, y_train).shape))

import torch
from tinydit.expert import ExpertConfig, ExpertTinyDiT

torch.manual_seed(0)
model = ExpertTinyDiT(ExpertConfig(dim=128, depth=4, num_heads=4, class_cond=True))
print("label table:", tuple(model.y_embed.table.weight.shape), "-> 10 classes + 1 'no class' row, null id =", model.null_class)

y = torch.randint(0, 10, (20,))
drop = torch.rand(20) < 0.1                       # 10% of the batch forgets its label
y_train = torch.where(drop, torch.full_like(y, model.null_class), y)
print("labels      :", y.tolist())
print("after drop  :", y_train.tolist())

x, t = torch.randn(20, 3, 32, 32), torch.rand(20)
print("v shape:", tuple(model(x, t, y_train).shape))

label table: (11, 128) -> 10 classes + 1 'no class' row, null id = 10
labels      : [5, 3, 5, 0, 3, 5, 0, 7, 3, 6, 1, 2, 7, 1, 4, 1, 9, 8, 7, 6]
after drop  : [5, 3, 5, 0, 3, 5, 0, 7, 3, 6, 1, 2, 7, 1, 4, 10, 9, 8, 7, 6]
v shape: (20, 3, 32, 32)

One of these 20 labels (the 16th, a 1) became a 10. With conditional=True and drop_prob=0.1, FlowMatchingLoss from Flow Matching in PyTorch: Train a Model to Follow a Straight Line does exactly this before calling the model.

Step 3/5: Classifier-Free Guidance at Sampling Time

In this step: We ask the model twice per sampling step, with and without the label, and push further in the direction the label adds.

Now at every sampling step we can ask the model twice, once with the label and once without, and push further in the direction the label adds. This is classifier-free guidance (CFG; Ho and Salimans, Classifier-Free Diffusion Guidance, 2022):

\begin{aligned} v_{\text{cfg}} &= \underbrace{v(x_t, t, \varnothing)}_{\text{``any image''}} \\ &\quad + w \cdot \big(\underbrace{v(x_t, t, y)}_{\text{``a cat''}} - v(x_t, t, \varnothing)\big) \end{aligned}

• w = 0: ignore the label (unconditional).
• w = 1: plain conditional model, no guidance.
• w > 1: exaggerate whatever makes it a cat. Samples look more “on class” but less varied.

import torch

# One pixel's velocity from the same noisy image, asked two ways:
v_uncond = torch.tensor([0.20, -0.10])   # "some image"
v_cond = torch.tensor([0.50, -0.40])     # "a cat"
for w in [0.0, 1.0, 1.5, 3.0]:
    v = v_uncond + w * (v_cond - v_uncond)
    print(f"w = {w:3.1f}: v = [{v[0]:+.2f}, {v[1]:+.2f}]")

import torch

# One pixel's velocity from the same noisy image, asked two ways:
v_uncond = torch.tensor([0.20, -0.10])   # "some image"
v_cond = torch.tensor([0.50, -0.40])     # "a cat"
for w in [0.0, 1.0, 1.5, 3.0]:
    v = v_uncond + w * (v_cond - v_uncond)
    print(f"w = {w:3.1f}: v = [{v[0]:+.2f}, {v[1]:+.2f}]")

w = 0.0: v = [+0.20, -0.10]
w = 1.0: v = [+0.50, -0.40]
w = 1.5: v = [+0.65, -0.55]
w = 3.0: v = [+1.10, -1.00]

In our euler_sample it’s three lines (cfg_scale and null_class arguments):

@torch.no_grad()
def euler_sample(model: nn.Module, noise: torch.Tensor, steps: int = 20, y: torch.Tensor | None = None,
                 cfg_scale: float = 1.0, null_class: int | None = None,
                 return_path: bool = False):
    """Integrate from t=1 (noise) to t=0 (data) with `steps` equal Euler steps.

    y=None: unconditional. With y and cfg_scale != 1: classifier-free guidance (post 12),
    v = v_uncond + cfg_scale * (v_cond - v_uncond).
    """
    was_training = model.training
    model.eval()
    x = noise
    ts = torch.linspace(1.0, 0.0, steps + 1, device=noise.device)
    path = [x]
    for i in range(steps):
        t = ts[i].expand(x.shape[0])
        if y is None:
            v = model(x, t)
        elif cfg_scale == 1.0:
            v = model(x, t, y)
        else:
            v_c = model(x, t, y)
            v_u = model(x, t, torch.full_like(y, null_class))
            v = v_u + cfg_scale * (v_c - v_u)
        x = x + (ts[i + 1] - ts[i]) * v          # dt is negative: we walk back toward data
        if return_path:
            path.append(x)
    model.train(was_training)
    return (x, path) if return_path else x

@torch.no_grad()
def euler_sample(model: nn.Module, noise: torch.Tensor, steps: int = 20, y: torch.Tensor | None = None,
                 cfg_scale: float = 1.0, null_class: int | None = None,
                 return_path: bool = False):
    """Integrate from t=1 (noise) to t=0 (data) with `steps` equal Euler steps.

    y=None: unconditional. With y and cfg_scale != 1: classifier-free guidance (post 12),
    v = v_uncond + cfg_scale * (v_cond - v_uncond).
    """
    was_training = model.training
    model.eval()
    x = noise
    ts = torch.linspace(1.0, 0.0, steps + 1, device=noise.device)
    path = [x]
    for i in range(steps):
        t = ts[i].expand(x.shape[0])
        if y is None:
            v = model(x, t)
        elif cfg_scale == 1.0:
            v = model(x, t, y)
        else:
            v_c = model(x, t, y)
            v_u = model(x, t, torch.full_like(y, null_class))
            v = v_u + cfg_scale * (v_c - v_u)
        x = x + (ts[i + 1] - ts[i]) * v          # dt is negative: we walk back toward data
        if return_path:
            path.append(x)
    model.train(was_training)
    return (x, path) if return_path else x

Label dropout and classifier-free guidance: what, why, how

What it is: a knob at sampling time that trades variety for “looks like the class”.

Why it works: training with 10% hidden labels teaches both v(x_t, t, y) and v(x_t, t, \varnothing). Their difference points toward “more like class y“, and guidance follows it further (Ho and Salimans, 2022).

How it works: training: y = where(rand < 0.1, null, y). Sampling: v = v_u + w * (v_c - v_u), two forward passes per step.

Step 4/5: adaLN-Zero: Start Every Block as “Do Nothing”

In this step: We add zero-initialised gates to every block so a fresh network starts as the identity function.

Our post-9 block always adds its attention and MLP output to the residual stream. At initialisation, those outputs are random, so a fresh deep network starts out scrambling its input.

adaLN-Zero (Peebles and Xie, DiT, 2023) makes two changes:

1. The adaLN Linear produces two extra vectors, gates g_1 and g_2, one per branch: x = x + g_1 \cdot \operatorname{attn}(\dots), x = x + g_2 \cdot \operatorname{MLP}(\dots).
2. That Linear (and the final output layer) is initialised to zero. So every gate starts at 0, and every block starts as the identity function.

\begin{gathered} x \leftarrow x + \underbrace{g_1(c)}_{=\,0 \text{ at init}} \cdot \operatorname{Attn}\big(\operatorname{adaLN}(x, c)\big) \\ x \leftarrow x + \underbrace{g_2(c)}_{=\,0 \text{ at init}} \cdot \operatorname{MLP}\big(\operatorname{adaLN}(x, c)\big) \end{gathered}

Here is the block, with the gate lines kept sharp:

class ExpertDiTBlock(nn.Module):
    """DiT block with switchable norm / QK-norm / FFN / RoPE, and adaLN or adaLN-Zero.

    adaLN      (post 9):  x = x + attn(mod(norm(x)))                 modulation: shift, scale  x2
    adaLN-Zero (post 12): x = x + gate * attn(mod(norm(x)))          modulation: shift, scale, gate x2,
                          and the modulation Linear starts at zero, so the block starts as the identity.
    """

    def __init__(self, cfg: ExpertConfig):
        super().__init__()
        dim = cfg.dim
        self.zero = cfg.adaln_zero
        self.norm1 = make_norm(cfg.norm, dim)
        self.attn = ModernAttention(dim, cfg.num_heads, cfg.use_sdpa, cfg.qk_norm)
        self.norm2 = make_norm(cfg.norm, dim)
        self.mlp = make_mlp(cfg.mlp, dim, cfg.mlp_ratio)
        self.adaLN = nn.Sequential(nn.SiLU(), nn.Linear(dim, (6 if self.zero else 4) * dim))
        if self.zero:
            nn.init.zeros_(self.adaLN[-1].weight)
            nn.init.zeros_(self.adaLN[-1].bias)

    def forward(self, x: torch.Tensor, c: torch.Tensor, rope=None) -> torch.Tensor:
        if self.zero:
            sh1, sc1, g1, sh2, sc2, g2 = self.adaLN(c)[:, None].chunk(6, dim=-1)
        else:
            sh1, sc1, sh2, sc2 = self.adaLN(c)[:, None].chunk(4, dim=-1)
            g1 = g2 = 1.0
        x = x + g1 * self.attn(self.norm1(x) * (1 + sc1) + sh1, rope)
        x = x + g2 * self.mlp(self.norm2(x) * (1 + sc2) + sh2)
        return x

class ExpertDiTBlock(nn.Module):
    """DiT block with switchable norm / QK-norm / FFN / RoPE, and adaLN or adaLN-Zero.

    adaLN      (post 9):  x = x + attn(mod(norm(x)))                 modulation: shift, scale  x2
    adaLN-Zero (post 12): x = x + gate * attn(mod(norm(x)))          modulation: shift, scale, gate x2,
                          and the modulation Linear starts at zero, so the block starts as the identity.
    """

    def __init__(self, cfg: ExpertConfig):
        super().__init__()
        dim = cfg.dim
        self.zero = cfg.adaln_zero
        self.norm1 = make_norm(cfg.norm, dim)
        self.attn = ModernAttention(dim, cfg.num_heads, cfg.use_sdpa, cfg.qk_norm)
        self.norm2 = make_norm(cfg.norm, dim)
        self.mlp = make_mlp(cfg.mlp, dim, cfg.mlp_ratio)
        self.adaLN = nn.Sequential(nn.SiLU(), nn.Linear(dim, (6 if self.zero else 4) * dim))
        if self.zero:
            nn.init.zeros_(self.adaLN[-1].weight)
            nn.init.zeros_(self.adaLN[-1].bias)

    def forward(self, x: torch.Tensor, c: torch.Tensor, rope=None) -> torch.Tensor:
        if self.zero:
            sh1, sc1, g1, sh2, sc2, g2 = self.adaLN(c)[:, None].chunk(6, dim=-1)
        else:
            sh1, sc1, sh2, sc2 = self.adaLN(c)[:, None].chunk(4, dim=-1)
            g1 = g2 = 1.0
        x = x + g1 * self.attn(self.norm1(x) * (1 + sc1) + sh1, rope)
        x = x + g2 * self.mlp(self.norm2(x) * (1 + sc2) + sh2)
        return x

And a check of what each block does to its input at initialisation:

import torch
from tinydit.expert import ExpertConfig, ExpertDiTBlock

torch.manual_seed(0)
x, c = torch.randn(2, 64, 128), torch.randn(2, 128)
for zero in [False, True]:
    block = ExpertDiTBlock(ExpertConfig(dim=128, num_heads=4, adaln_zero=zero))
    out = block(x, c)
    name = "adaLN-Zero" if zero else "adaLN     "
    print(f"{name}: modulation outputs {block.adaLN[-1].out_features:4d} | "
          f"block(x) - x at init: {(out - x).abs().mean():.4f}")

import torch
from tinydit.expert import ExpertConfig, ExpertDiTBlock

torch.manual_seed(0)
x, c = torch.randn(2, 64, 128), torch.randn(2, 128)
for zero in [False, True]:
    block = ExpertDiTBlock(ExpertConfig(dim=128, num_heads=4, adaln_zero=zero))
    out = block(x, c)
    name = "adaLN-Zero" if zero else "adaLN     "
    print(f"{name}: modulation outputs {block.adaLN[-1].out_features:4d} | "
          f"block(x) - x at init: {(out - x).abs().mean():.4f}")

adaLN     : modulation outputs  512 | block(x) - x at init: 0.2212
adaLN-Zero: modulation outputs  768 | block(x) - x at init: 0.0000

At initialisation, the adaLN block changes its input by 0.22 on average, while the adaLN-Zero block changes it by exactly 0.0000.

adaLN-Zero: what, why, how

What it is: a safer start for a stack of conditioned blocks.

Why it works: with zero-initialised gates, every residual block starts as the identity, so early training isn’t fighting random noise from deep layers; each block opens its gate as it becomes useful (DiT, 2023).

How it works: the adaLN Linear outputs 6 vectors (shift, scale, gate for attention and for the MLP) and starts at zero, as does the final projection: x = x + gate * branch(modulate(norm(x))). In TinyDiT: --set adaln_zero=true.

Step 5/5: The Before/After on TinyDiT

In this step: We train adaLN and adaLN-Zero class-conditional models and score them at four guidance scales.

Two class-conditional runs, identical except for the block: the post-9 adaLN block, and adaLN-Zero. Same seed, same 3,000 steps, same 10% label dropout. After training, each model drew 1,000 images (100 per class) at four guidance scales.

Besides val MSE and probe-FD, a conditional model gets one more score: class accuracy, the share of samples that our small CIFAR-10 classifier recognises as the class we asked for. Random guessing gets 10%.

import json
from pathlib import Path

def load(name):
    return json.loads(Path(f"../runs/{name}/metrics.json").read_text())

runs = {"adaLN (post 9 block)": load("p12_cond_adaln"), "adaLN-Zero": load("p12_cond_zero")}
print(f"{'model':22s} {'params':>9s} {'val MSE':>8s} | {'guidance':>8s} {'probe-FD':>9s} {'class acc':>9s}")
for name, m in runs.items():
    r = m["raw"]                                  # the trained weights (not EMA; EMA is post 14)
    rows = [("none (w=0)", r["uncond"])] + [(f"w={w}", r[f"cfg{w}"]) for w in ("1.0", "1.5", "3.0")]
    for i, (w, s) in enumerate(rows):
        head = f"{name:22s} {m['params']:9,d} {r['val_mse']:8.4f}" if i == 0 else " " * 41
        print(f"{head} | {w:>10s} {s['probe_fd']:7.1f} {s['probe_acc']:9.1%}")
    print(f"{'':22s} loss curve: {[c[1] for c in m['curve'][::3]]}")

import json
from pathlib import Path

def load(name):
    return json.loads(Path(f"../runs/{name}/metrics.json").read_text())

runs = {"adaLN (post 9 block)": load("p12_cond_adaln"), "adaLN-Zero": load("p12_cond_zero")}
print(f"{'model':22s} {'params':>9s} {'val MSE':>8s} | {'guidance':>8s} {'probe-FD':>9s} {'class acc':>9s}")
for name, m in runs.items():
    r = m["raw"]                                  # the trained weights (not EMA; EMA is post 14)
    rows = [("none (w=0)", r["uncond"])] + [(f"w={w}", r[f"cfg{w}"]) for w in ("1.0", "1.5", "3.0")]
    for i, (w, s) in enumerate(rows):
        head = f"{name:22s} {m['params']:9,d} {r['val_mse']:8.4f}" if i == 0 else " " * 41
        print(f"{head} | {w:>10s} {s['probe_fd']:7.1f} {s['probe_acc']:9.1%}")
    print(f"{'':22s} loss curve: {[c[1] for c in m['curve'][::3]]}")

model                     params  val MSE | guidance  probe-FD class acc
adaLN (post 9 block)   1,151,792   0.2308 | none (w=0)    58.8      8.9%
                                          |      w=1.0    58.7     24.6%
                                          |      w=1.5    58.1     31.1%
                                          |      w=3.0    60.6     44.8%
                       loss curve: [0.51611, 0.27276, 0.24948, 0.23618]
adaLN-Zero             1,283,888   0.2157 | none (w=0)    60.7      8.9%
                                          |      w=1.0    53.2     25.3%
                                          |      w=1.5    50.3     32.7%
                                          |      w=3.0    52.3     45.2%
                       loss curve: [0.46049, 0.24348, 0.22631, 0.21828]

What the numbers say:

• Guidance works. Without the label (w = 0) samples hit the right class 8.9% of the time, which is chance. The plain conditional model (w = 1) gets about 25%, and guidance pushes that to 31 to 33% at w = 1.5 and 45% at w = 3.
• It’s a trade-off. Probe-FD is best at w = 1.5 and gets worse at w = 3: stronger guidance makes each image more “on class” but the set of images less varied, and probe-FD (like FID) punishes missing variety. That’s exactly the trade-off the CFG paper describes, and why image models expose a “guidance scale” setting.
• adaLN-Zero clearly helped. Val MSE dropped from 0.2308 to 0.2157 (-0.015, about 4× the seed gap we measured in RMSNorm, QK-Norm, SwiGLU, and 2D RoPE in a Tiny DiT), its loss was lower at every point of training, and its guided samples scored 50.3 instead of 58.1 at w = 1.5.
• With the label ignored (w = 0), the two models are about the same (60.7 vs 58.8). The gain is in the conditional samples, which is what the conditioning block is for.

Recap

• A class label enters TinyDiT through an nn.Embedding with one extra “no label” row, added to the time vector so adaLN carries it into every block.
• Hiding the label for 10% of training images trains a conditional and an unconditional model in one, and classifier-free guidance combines them at sampling time: v = v_u + w\,(v_c – v_u), at the cost of two forward passes per step.
• Measured: class accuracy goes from 8.9% (no label) to about 25% (label, w = 1), 33% (w = 1.5) and 45% (w = 3). Probe-FD is best at w = 1.5 and worse at w = 3, the usual quality-versus-variety trade-off.
• adaLN-Zero starts every block as the identity with zero-initialised gates. It improved val MSE from 0.2308 to 0.2157 and guided probe-FD from 58.1 to 50.3, with 11% more parameters, so part of the gain may come from size.
• From here on, TinyDiT uses adaLN-Zero and guidance scale 1.5.

Next up: What if the condition were a whole sentence instead of one label? In MMDiT Joint Attention in PyTorch: Two Streams, One Generator we build the MMDiT joint-attention block behind SD3 and FLUX and test it against adaLN.

← Previous: Post 11: RMSNorm, QK-Norm, SwiGLU, and 2D RoPE in a Tiny DiT
Next → Post 13: MMDiT Joint Attention in PyTorch: Two Streams, One Generator