Integrating Timestep Information

As we outlined previously, the core of the reverse diffusion process is a neural network, typically a U-Net, tasked with predicting the noise $\epsilon$ that was added to an image $x_0$ to produce a noisy version $x_t$ at a specific timestep $t$. A significant aspect of this process is that the same network parameters $\theta$ must be used for all possible timesteps $t$ (from $0$ to $T$).

How can a single network adapt its behavior based on the current timestep? It needs to be informed about which timestep $t$ it is currently processing. Simply feeding the integer $t$ directly into the network is often suboptimal. Neural networks generally work better with normalized inputs, and a raw integer doesn't easily convey the "position" within the diffusion process in a way the network can effectively use. Furthermore, different timesteps require vastly different denoising behavior; the noise pattern and magnitude at $t=50$ are distinct from those at $t=800$.

Therefore, we need an effective way to represent the timestep $t$ and inject this information into the model. The most common and effective technique, inspired by the Transformer architecture (Vaswani et al., 2017), is to use sinusoidal positional embeddings.

Sinusoidal Timestep Embeddings

Instead of using the raw integer $t$, we transform it into a high-dimensional vector $e(t)$. This embedding vector provides a richer, more structured representation of the timestep that the network can better utilize. The standard approach uses sine and cosine functions of varying frequencies:

Let $d$ be the desired dimensionality of the embedding vector. For each dimension $i$ from $0$ to $d/2 - 1$, the embedding components are calculated as:

$$e(t)_{2i} = \sin\left( \frac{t}{10000^{2i/d}} \right)$$ $$e(t)_{2i+1} = \cos\left( \frac{t}{10000^{2i/d}} \right)$$

Here:

• $t$ is the current timestep.
• $i$ indexes the dimension within the embedding vector (up to $d/2$).
• $d$ is the total dimensionality of the embedding (e.g., 128, 256, or 512).

This formulation creates a vector where each pair of dimensions $(2i, 2i+1)$ corresponds to a sinusoid with a specific frequency. Lower dimensions (small $i$) vary slowly with $t$, while higher dimensions (large $i$) vary rapidly. This multi-frequency representation allows the network to easily distinguish between different timesteps and potentially generalize better to unseen timesteps if needed.

Example values of sinusoidal embeddings for different dimensions across timesteps 0 to 1000, assuming an embedding dimension $d$ greater than 51. Lower dimensions change slowly, while higher dimensions oscillate rapidly.

Injecting Embeddings into the U-Net

Once we have the timestep embedding vector $e(t)$ (which typically has a dimension like 256 or 512), we need to integrate it into the U-Net architecture. The U-Net processes the noisy image $x_t$ through a series of convolutional layers, downsampling, and upsampling blocks. The timestep information needs to influence these computations.

A common strategy involves:

1. Projection: The embedding vector $e(t)$ is often passed through a small Multi-Layer Perceptron (MLP) - typically one or two linear layers with an activation function like SiLU (Sigmoid Linear Unit) or ReLU. This MLP learns to transform the fixed sinusoidal embedding into a representation that is most useful for the network's internal workings.
2. Broadcasting and Addition: The output of this MLP (let's call it the processed embedding) is then typically broadcast to match the spatial dimensions (Height and Width) of the feature maps at various points within the U-Net. It's then added element-wise to these feature maps. This usually happens within the residual blocks that form the backbone of the U-Net.
3. Multi-Resolution Injection: This addition of the processed timestep embedding is often performed at multiple resolution levels within the U-Net, typically within each residual block in both the downsampling and upsampling paths. This ensures the timestep information influences processing across different spatial scales.

Diagram illustrating how a processed time embedding is typically added to feature maps within a U-Net block. The sinusoidal embedding is first transformed by an MLP before being injected.

By incorporating timestep information in this way, the single U-Net model $\epsilon_\theta(x_t, t)$ learns to adapt its noise prediction based on the level of noise present, indicated by $t$. This conditioning on time is fundamental to the operation of diffusion models, allowing the network to perform the appropriate denoising operation at each step of the reverse process.