How a Diffusion Transformer Generates Shapes from Text

The model works in three stages: embedding, attention, and projection.

Diffusion Transformer (DiT) architecture

The model receives three types of inputs. Each is mapped to a 64-dimensional vector:

• 1 Timestep: Tells the model where it is on the path from noise to shape. Encoded with sinusoidal encoding.
• 4 Text tokens: Contain the prompt words. Encoded with learned word embeddings and positional encodings.
• 24 Coordinate points: Represent the (x, y) positions at the current timestep. Projected using a linear layer.

The model merges these vectors into a single 29-token sequence and passes them through 2 transformer layers. Each layer runs self-attention with 2 heads, followed by an MLP.

The output layer turns each of the 24 point tokens into a velocity, the direction from shape to noise. The points are moved slightly in the opposite direction. Repeating this process forms the shape.

To predict the direction, the model must know where the points are on their path from noise to shape. The timestep $t$ indicates this progress as it decreases from 1 to 0. The value is encoded with sine and cosine waves, called sinusoidal encoding. Most waves barely move over such a small range, so $t$ is scaled by 100 first. The result is a 64-dimensional vector:

Here, D = 64 is the dimension, k runs from 0 to 31 for each sine-cosine pair, and base b = 500 sets the frequency range.

Plotted in 3D, the components form a landscape of waves across time. The sliders show how the base and the dimension reshape it:

Sinusoidal encoding. Each ridge shows how an individual component changes over 100 timesteps.

Close timesteps should get close vectors. In vector space, similarity is measured with the dot product. Pairing sine with cosine makes this similarity depend only on the time difference. Consider two 2D encodings $\mathbf{a} = (\sin t_1, \, \cos t_1)$ and $\mathbf{b} = (\sin t_2, \, \cos t_2)$, in this case with $\omega = 1$. Their dot product simplifies via the angle difference identity:

Both vectors have length 1, so the dot product is their cosine similarity. While a single wave repeats every few steps, combining 32 frequencies makes the similarity fall as the timesteps move apart.

The encoding works like a clock with 32 hands. Each sine and cosine pair sets the position of one hand. Just as humans use hours, minutes, and seconds instead of counting only seconds, the network processes time more easily when split across multiple frequencies.

Shorter hands represent high frequencies and spin quickly, while longer hands represent low frequencies and turn slowly.

The model uses the prompt to know which shapes to draw and where to place them. Each prompt has a fixed length of 4 tokens from a 12-word vocabulary. A learned embedding table maps each token to a 64-dimensional vector. Unused slots are filled with PAD.

A positional encoding is added to each vector so the model can tell the difference between triangle on circle and circle on triangle.

Instead of learned positional embeddings (as in How a Transformer Plays Tic-Tac-Toe), the model follows the original Attention Is All You Need paper and uses sinusoidal encoding (with base b = 10).

The third input is the $(x, y)$ coordinates of 24 points at timestep $t$. A learned linear layer with weights and biases projects each coordinate pair into a 64-dimensional vector:

Note: Matrix multiplications visualized as a row–column diagram. Rows are separated to show each token is processed independently.

Attention is where tokens exchange information. The mechanics of queries (Q), keys (K), and values (V) are covered in How a Transformer Plays Tic-Tac-Toe.

The model has 2 layers with 2 heads each, a width of D = 64, and a context length of 29 tokens. Every point can attend to all other points, the prompt words, and the timestep. The only mask hides the empty PAD slots in the prompt.

In the last steps of generation, the second layer shows a distinct role for each head:

• Head 1 assigns points to shapes. Points attending to "triangle" form the triangle on the canvas, while points attending to "circle" form the circle.
• Head 2 tracks neighboring points. Each row shows three or four activations: the point itself on the diagonal, plus its left and right neighbors along the outline.

Click inside the green box to highlight where the corresponding points are located.

The timestep and prompt tokens have conditioned the points in attention and are no longer needed. The model keeps only the 24 point vectors.

A final linear layer with learned weights and biases projects each point vector to a 2D velocity: the direction from the shape toward noise.

Each point then moves the opposite way, from noise toward the shape, and the timestep $t$ decreases. The model runs again on the updated points and timestep. The loop repeats until $t = 0$.

Flow Matching connects clean shapes and random noise with straight lines. The forward process moves from shape to noise and is used in training. The reverse process moves from noise to shape and is used in generation. Both processes are described by the same formula:

Here, $x_0$ is the clean shape, $\epsilon$ is random noise, and $x_t$ represents the points at timestep $t$.

Forward Process

In the forward process, time $t$ runs from 0 to 1. As $t$ increases, the shape signal scales by $(1 - t)$ while the noise contribution scales by $t$.

Solid black dots represent the current points, while faint gray circles show their destination noise positions.

The points move in straight lines. The model learns to predict the direction of these lines, called the velocity:

Arrows represent the velocity vectors, pointing from each shape point to its corresponding noise position.

Reverse Process

The main formula can be rewritten in terms of the velocity $v$:

This formula can be rewritten in recursive form to calculate the position of the points at the previous step, $x_{t - \Delta t}$:

Generation starts at $t = 1$, where $x_1 = \epsilon$ is pure noise. The model predicts the velocity $v_\theta$, and the recursive formula calculates the next positions of the points, $x_{t - \Delta t}$. Repeating this step gradually forms the shape.

Move the vertical slider to inspect different trajectories.

Time Shifting

Generation can use any number of steps: more steps trace the path more accurately, while fewer steps run faster.

Early models spaced these steps evenly from 1 to 0. Later systems introduced time shifting to concentrate steps where the model needs the most precision.

A shift factor $s$ maps the evenly spaced times $t$ to new times $t'$:

At $s = 1$, nothing changes. Setting $s > 1$ puts more steps near noise, while $s < 1$ puts more steps near the shape.

Instead of memorizing fixed scenes, the model learns shapes and spatial relations as independent concepts.

For example, the prompt circle on triangle generates the scene correctly, even though the model never saw it during training.

The model can also blend shapes. For instance, the prompt square triangle produces a trapezoid.

Here are some examples not present in the training set:

Out-of-distribution prompts can be tested in the demo above.

This ability to combine separate concepts is called compositional generalization. It allows real image models to generate scenes never seen during training, like "a pelican riding a bicycle" or "an astronaut riding a horse".