# D3. Latents and DiT > Why does image generation come back to the Transformer? LLM by Hand · Theory · side trip: Diffusion · runs in your browser · last part on your computer · interactive page: https://llm.liko.page/learn/latents-and-dit/ **Boss challenge.** Write the full sampler yourself and generate a chosen shape in your browser. Then finish boss.py on your computer: your own noise, guidance and reverse step must make a tiny Transformer draw digits that a separate classifier can read. So far every point had 2 numbers. A small gray picture of a handwritten digit is 28 × 28 = 784 numbers. Nothing in levels D1 or D2 depended on how many numbers a point has. Adding noise, guessing it, stepping back: all the same formulas, now on 784 numbers at once. What changes is the noise-guessing network. It has to look at a whole picture and say, for every pixel, how much of it is noise. A stroke that starts in one corner continues somewhere else, so every part of the picture needs to know about the other parts. You have already met the tool built for exactly that: attention, from level 14. > **What you need:** levels D1–D2, and attention and Transformer blocks (levels 14–16). > Section 3 uses the autoencoders of side trip N6; it explains what it needs, but N6 shows it in full. ## 1. A picture as a sentence A Transformer reads a sequence of tokens. So we cut the picture into square **patches** and put them in a row, left to right, then top to bottom. Each patch becomes one token: its pixels in a row. **Patch size 2** means each patch is 2 × 2 pixels. Here is a 4 × 4 picture holding the numbers 0 to 15, cut into patches of size 2: ``` 0 1 | 2 3 token 0: [ 0, 1, 4, 5] 4 5 | 6 7 token 1: [ 2, 3, 6, 7] ------+------ → token 2: [ 8, 9, 12, 13] 8 9 | 10 11 token 3: [10, 11, 14, 15] 12 13 | 14 15 ``` Four patches, four pixels each, so the token matrix is (4, 4). Careful: “patch size 2” does not mean “cut into a 2 × 2 grid of pieces”. On this 4 × 4 picture both readings happen to give four patches of four pixels, so the example can’t distinguish them. On a 6 × 6 picture they differ. Try it. **Question.** A 6 × 6 picture is cut into patches of size 2 (each patch 2 × 2 pixels). What shape is the token matrix (tokens, numbers per token)? *Answer it on the page to check your work.* Real digits are bigger, so the same counting gets more interesting. **Question.** A 28 × 28 digit is cut into patches of size 7 (each patch 7 × 7 pixels). How many tokens is that? *Answer it on the page to check your work.* **Predict.** On a 28 × 28 picture, go from patch size 4 to patch size 2. What happens to the number of attention scores in each layer? A. Doubles B. 4 times as many C. 16 times as many *Answer it on the page to check your work.* Now a real digit. Point at or tap a patch to find its token, or a token to find its patch. Change the patch size and watch the token count. *[Interactive lab: Patch — open the page to use it]* Write the cutting yourself, in two steps. First, `reshape` can split each axis of the picture into (which block, which pixel inside the block). For the 4 × 4 picture and patch size 2: ```python x4 = img.reshape(2, 2, 2, 2) # axes: (block row, row in block, block column, column in block) x4[0, :, 1, :] # block row 0, block column 1 → [[2, 3], [6, 7]], the top-right patch ``` Write the sizes for any picture and patch size: **Code question.** Split each axis of an (H, W) picture into (which block, which pixel in the block), for patch size p. Fill in the blank (`____`): ```python def split(x, p): H, W = x.shape return x.reshape(____) img = np.arange(16).reshape(4, 4) x4 = split(img, 2) print(x4[0, :, 1, :]) # should be the top-right patch [[2, 3], [6, 7]] ``` *Answer it on the page to check your work.* Now one patch is `x4[i, :, j, :]`, but the two block axes, `i` and `j`, sit in positions 0 and 2. To list patches in order they must come first. `transpose` reorders axes: `a.transpose(1, 0)` swaps the two axes of a table, and `x4.transpose(0, 2, 1, 3)` gives the new order axis 0, axis 2, axis 1, axis 3, so the shape becomes (block row, block column, row in block, column in block). After that, a `reshape` to (number of patches, p · p) flattens each patch. **Code question.** Write patchify: an (H, W) picture → (number of patches, p·p), patches left to right, then top to bottom. Fill in the blank (`____`): ```python def patchify(x, p): H, W = x.shape return ____ img = np.arange(16).reshape(4, 4) print(patchify(img, 2)) ``` *Answer it on the page to check your work.* Once a picture is a sequence of tokens, any Transformer can read it. ## 2. The DiT: a Transformer that guesses noise The table under the lab is the whole model. Patches in, a noise guess for every patch out. The middle is the same pre-norm block you built in level 16: attention, feed-forward, LayerNorm, residuals. It has no causal mask, because it writes no sequence one token at a time: every patch may look at every other patch. Only two things are new: the step `t` and the label (which digit to draw) are turned into vectors and added to every token. **Question.** Patch size 4 on a 28 × 28 digit gives 49 tokens of 16 numbers. Wᵢₙ is (16, 96). What is the shape of tokens @ Wᵢₙ? *Answer it on the page to check your work.* **Question.** With 49 tokens, how many attention scores does one head compute in one layer? *Answer it on the page to check your work.* **If you are stuck: Why a Transformer? The small network in level D1 worked fine.** In level D1 each point was 2 numbers, and the network could see both at once. A picture is 784 numbers. A plain network that connects every pixel to every pixel needs a huge first layer and does not know that pixel 30 is right below pixel 2. Patches keep nearby pixels together. Attention then lets every patch ask every other patch “what is drawn over there?” That is exactly the kind of question a noise guess needs: a faint gray pixel is ink if the stroke continues next to it, and noise if it doesn't. **Deeper: How real models give t and the label to the network** Our tiny model simply adds one vector, made from t and the label, to every token. That works. Larger models do something stronger. From t and the label they compute a scale and a shift for each LayerNorm inside every block, so the step and the label can make whole features larger or smaller everywhere in the network. Many models also start those scales at zero, so each block begins as “do nothing” and training slowly makes it active. The rest is plain Transformer: the same blocks, stacked deeper, with more patches. Patches work for small digits. Large pictures need one more idea. ## 3. Diffusing in a smaller space Big pictures make this expensive. A 512 × 512 color picture has 3 numbers per pixel (red, green, blue), so 512 × 512 × 3 = 786,432 numbers. Cut into patches of size 8, it becomes 64 × 64 = 4,096 tokens, and attention compares every one of them with every other. The fix is to squeeze the picture first. You built the tool for this in level N6: an **autoencoder**. It is two networks trained together. The encoder squeezes a picture into a few numbers, the code. The decoder turns the code back into a picture. They are trained so the picture that comes back matches the one that went in. Then you run all of levels D1 and D2 on the codes instead of the pixels, and decode only once, at the very end. The lab below uses a small autoencoder that squeezes each digit into 16 numbers. **Question.** The autoencoder squeezes 784 pixels into 16 numbers. How many times fewer numbers does diffusion have to work on? *Answer it on the page to check your work.* One detail matters before you add noise to codes. Level D1’s formula keeps the variance at 1 only if the clean data has variance 1. The codes from this plain autoencoder are much bigger: their std is about 4, so some codes reach ±10. **Question.** The 16 code numbers have a std of about 4. Level D1’s formula needs clean data with variance 1, so std 1. You divide every code number by the same number first. Which number? *Answer it on the page to check your work.* So the recipe is: encode, divide by the codes’ std, run D1 and D2 on the scaled codes, multiply back, decode. Level N6 also showed why the kind of autoencoder matters. A plain autoencoder leaves holes between codes; a VAE, trained with a small KL term, keeps the codes smooth and close to size 1. Image generators use VAE-style encoders for that reason, so diffusion can land anywhere in the code space and still decode to a sensible picture. **Deeper: What real latents look like** Our code is a flat list of 16 numbers. The encoders inside large image models keep a small grid instead: for example, a 512 × 512 color picture becomes a 64 × 64 grid with 4 numbers per cell. That grid is cut into patches exactly like the pixels in section 1, and a DiT reads those patches as tokens. So the two ideas of this level work together: squeeze first, then patchify the squeezed grid. Diffusion directly on pixels also works, as the boss below shows; it is just more expensive for big pictures. The lab also walks from one digit to another: once by mixing pixels, once by mixing the 16-number codes. Guess first. **Predict.** An autoencoder turns each 28 × 28 digit into a 16-number code, and its decoder turns a code back into a picture. Take a 3 and a 7 and go halfway from one to the other. What does mixing the pixels give, compared with mixing the two codes and decoding the result? A. Both look like the same smooth blend of a 3 and a 7 B. Pixels: a faint 3 and 7 on top of each other. Codes: a single in-between digit C. Pixels: one in-between digit. Codes: a faint 3 and 7 on top of each other *Answer it on the page to check your work.* *[Interactive lab: Latent — open the page to use it]* **Try it** Slide the walk from the 3 to the 7 slowly. On the left, two faint digits on top of each other. On the right, one digit that changes shape. The 16 numbers describe what is drawn, not where each pixel is. That is why a diffusion model working on codes can make clean pictures with far less work. You now have every piece. The boss puts them together. ## 4. Boss, part 1: write the whole sampler Here is a dataset of 8 points on a ring. For a dataset this small you can write down the **perfect** noise-guessing function, `eps_exact`, with no training at all. Your job is the loop that turns random points into points on the ring (the code prints 8 samples; the tests draw 300): the reverse step from level D2, every step from T down to 1, with no fresh noise on the last one. **Code question.** Write the sampling loop. `eps_exact` is the perfect noise guess for this ring of 8 points. Every sample must land on one of the 8 points, and all 8 must appear. Fill in the blank (`____`): ```python T = 50 betas = np.linspace(1e-3, 0.2, T) # beta_1 .. beta_T (index 0 is t = 1) alphas = 1 - betas alpha_bar = np.cumprod(alphas) angles = np.arange(8) * np.pi / 4 targets = np.stack([np.cos(angles), np.sin(angles)], 1) # 8 points on a ring def eps_exact(x, t): """The best possible noise guess for this dataset. You don't need to change it.""" ab = alpha_bar[t - 1] d = ((x[:, None, :] - np.sqrt(ab) * targets[None]) ** 2).sum(-1) # (n, 8) # how likely each target is w = softmax(-d / (2 * (1 - ab)), axis=1) # the expected clean point x0 = w @ targets return (x - np.sqrt(ab) * x0) / np.sqrt(1 - ab) def sample(n, seed=0): rng = np.random.default_rng(seed) x = rng.normal(size=(n, 2)) # pure noise for t in range(T, 0, -1): # t = 50, 49, ..., 1 eps = eps_exact(x, t) ____ return x x = sample(8) print(np.round(x, 3)) ``` *Answer it on the page to check your work.* ## 5. Boss, part 2: draw digits on your computer This part runs on your computer. New to running Python locally? Start with the [setup page](/setup/). Download [boss.py](/files/latents-and-dit/boss.py) into a folder of its own, and run every command in that folder. Open `boss.py`. Three functions are marked TODO, and each is a step you already did in your browser: - `add_noise`: the one line from level D1, - `guide`: the guidance mix from level D2, - `step_back`: the reverse step from level D2, with no fresh noise on the last step. You just wrote it in NumPy in part 1; in PyTorch the fresh noise is `torch.randn_like(x)`. Everything else is written for you: the digit pictures (downloaded once, about 10 MB), the tiny DiT from section 2, the training loop, and the check. ```bash python boss.py ``` The script first tests your three functions on the small numbers from levels D1 and D2, so you see a mistake in seconds. Then it trains for `STEPS` = 8000 steps with your `add_noise` (about two minutes on our laptop, longer on a slow one), draws 8 pictures of every digit with your `guide` and `step_back`, saves them as `my_digits.png` (one row per digit, 0 at the top), and gives them to a separate digit classifier. You pass the boss when the classifier reads at least 70% of your 80 pictures as the digit you asked for. The script then prints a line that starts with `D3 PASS`. A correct solution scores about 85–90% on a laptop; the exact number depends on the run. If the small checks pass but the score stays under 70%, open `my_digits.png`, then raise `STEPS` near the top of `boss.py` and run again. Paste the `D3 PASS` line here. The code at the end of the line only shows that you pasted it unchanged; this part relies on your honesty. The level counts as cleared once part 1 (the sampler on this page) and this line both pass. **Code question.** The boss, part 2. Paste the PASS line that boss.py printed for your own `add_noise`, guide and `step_back`, between the quotes. Fill in the blank (`____`): ```python line = "____" # it looks like: D3 PASS 0.900 1a2b3c4d print(line) ``` *Answer it on the page to check your work.* There is no reference solution to download: the boss is the test. If you are stuck, go back to the small checks at the top of `boss.py`. Each of your three functions is one formula from level D1 or D2. To finish in minutes, this boss diffuses the 784 pixels directly. Doing the same loop on scaled codes from section 3 and decoding at the end is how large image models handle big pictures. The training hides the label 10% of the time; D2’s model hid 20%. Anything from about 10% to 20% works. The loss falls from about 0.14 to about 0.064 over the 8000 steps. That is far below the 0.32 lower limit from level D1, because that limit depends on the data. Neighboring pixels of a digit predict each other strongly, so a noisy digit shows much more about its noise than a noisy spiral point does. **Try it: Then try** Change `3.0` in the call to `guide` near the end of `main` to 0, then to 8, and open `my_digits.png` after each run. At 0 a row shows the wrong digits, and the score drops. At 8 every digit in a row looks the same, and thicker. That is the trade-off between accuracy and variety from level D2, now on pictures. ## You can now - Cut a picture into patch tokens with `reshape` and `transpose`. - Count the tokens and attention scores for a picture size and a patch size. - Write the whole sampler: reverse steps from T down to 1, with no fresh noise on the last one.