Time conditioning: telling a network the time

The same point needs a different answer at a different time

Generative models such as flow matching blend a noise draw x₀ and a data example x₁: x_t = (1 − t) x₀ + t x₁. At t = 0 the input is pure noise; at t = 1 it is data. The network sees x_t and must predict how to move it toward the data. The right move depends on t, even for the same x_t.

Your turn. What is the velocity at position 0 at t = 0.9?

First, turn the time into many numbers

One number is a weak input for a network. Its first layer can only scale it and add to it, so small differences in t are hard to act on. The usual fix uses Fourier features: the sines and cosines of t at several frequencies. Nearby times then get similar lists of numbers, and the fast waves let the network tell close times apart.

A small MLP then maps those features to a vector e(t), the time embedding. It has as many numbers as each token has features. Other conditions can ride along. For example, a model that takes a class label can add the label's embedding to e(t) and use the sum wherever it would use e(t).

adaLN lets the time set each layer's scale and shift

Layer normalization takes each token's features, subtracts their mean and divides by their spread. It then multiplies by a learned scale and adds a learned shift. In plain layer norm, the scale and shift are fixed after training: the same for every input and every time.

Adaptive layer normalization computes the scale and shift from the time embedding instead. A linear layer maps e(t) to two vectors, γ(t) and β(t), one number per feature:

adaLN(h, t) = LN(h) ⊙ (1 + γ(t)) + β(t)

Here h is a token's feature vector, LN is layer normalization without its own scale and shift, and ⊙ multiplies feature by feature. The "1 +" means that γ = 0 leaves the scale at 1. Each block of the network has its own linear layer, so each block can respond to the time in its own way.

Your turn. LN(h) = (−1, 0, 1), γ = (1, 0, 0) and β = (0, 0, 2). What does adaLN output?

adaLN-Zero starts every block switched off

In a diffusion transformer, the time layer also makes a gate α for each branch of a block. The branch's output is multiplied by α before it is added back to the input, through a residual connection. The layer that makes γ, β and α starts with all its weights set to zero. So at the start α = 0, every branch adds nothing, and every block passes its input through unchanged. Training then opens each gate as far as it helps.

This variant is called adaLN-Zero. The paper that introduced diffusion transformers compared it with three other ways to feed in the time and a class label: adding them as extra tokens, reading them with cross-attention, and plain adaLN. adaLN-Zero gave the best images in that comparison.

The code writes the layer as a Python class named AdaLN, built on PyTorch's nn.Module. Its __init__ runs once, when AdaLN(D) makes the layer. It first calls super().__init__(), which lets nn.Module set itself up. Then it stores the layer's parts in self, the layer itself. Its forward runs each time the layer is called, as in layer(h, t). The Python tutorial explains classes.

adaln.py (PyTorch)

import torch
import torch.nn as nn

# fix the layer's random starting weights, so each run prints the same
torch.manual_seed(0)
D = 4                                      # features per token

def time_features(t):
    # Sines and cosines of t at four frequencies: 8 numbers per time.
    # torch.arange(4) is [0, 1, 2, 3], so freqs holds π × 1, 2, 4 and 8.
    freqs = torch.pi * 2.0 ** torch.arange(4)
    # torch.cat joins the 4 sines and the 4 cosines into one row of 8 numbers.
    return torch.cat([torch.sin(t * freqs), torch.cos(t * freqs)], dim=-1)

class AdaLN(nn.Module):
    def __init__(self, d, d_time=8):
        super().__init__()
        self.norm = nn.LayerNorm(d, elementwise_affine=False)    # no fixed scale or shift
        self.to_scale_shift = nn.Linear(d_time, 2 * d)           # the time chooses them

    def forward(self, h, t):
        # .chunk(2, dim=-1) cuts the 8 outputs into two halves of 4: scale and shift.
        scale, shift = self.to_scale_shift(time_features(t)).chunk(2, dim=-1)
        return self.norm(h) * (1 + scale) + shift

layer = AdaLN(D)
h = torch.tensor([2.0, 4.0, 6.0, 8.0])     # one token's features
# Inside torch.no_grad(), PyTorch records no gradients, so .numpy() can turn
# each result into a NumPy array.
with torch.no_grad():
    print("normalized    ", layer.norm(h).numpy().round(2))
    for t in [0.1, 0.9]:
        print(f"adaLN at t={t}", layer(h, torch.tensor([t])).numpy().round(2))

normalized     [-1.34 -0.45  0.45  1.34]
adaLN at t=0.1 [-1.21 -0.52  0.95  0.34]
adaLN at t=0.9 [-0.97  0.09 -0.08  1.25]

The layer here is untrained, so its outputs are arbitrary. The point is that one token leaves it differently at two times. For brevity, the code skips the small MLP and feeds the sines and cosines straight to the linear layer.