Stanford CME296 Diffusion & Large Vision Models | Lecture 5 - Architectures — Annotated Lecture Notes & Key Takeaways

Course Information

Institution: Stanford

Original Course: Stanford CME296 Diffusion & Large Vision Models ｜ Spring 2026 ｜ Lecture 5 - Architectures

Instructor Bio: This lecture is co-delivered by **Afshine Amidi** and **Shervine Amidi**, adjunct faculty at Stanford University’s Institute for Computational and Mathematical Engineering (ICME). Afshine Amidi earned his engineering degree from École Centrale Paris and a Master of Science in Operations Research from the Massachusetts Institute of Technology, where he worked under Prof. Dimitris Bertsimas and received the Jean Gaillard Scholarship and MIT Dean’s Fellowship. He has held professional roles at Amazon and McKinsey & Company, leading AI and business strategy initiatives, and co-authors widely adopted technical learning guides on machine learning, algorithms, and transformer models. Shervine Amidi holds B.S. and M.S. degrees in engineering from École Centrale Paris, as well as an M.S. in Computational and Mathematical Engineering from Stanford University. He has conducted computer vision research at the Stanford Vision Lab and the École Centrale Paris Visual Computing Center. Currently a Senior Machine Learning Engineer at Netflix, he previously worked at Google DeepMind on the Gemini team, Google Assistant, and Uber’s data science division. He has served as a teaching assistant for multiple core Stanford CS courses and has been an adjunct lecturer at the university since 2021.

Course Description: This lecture dives into the neural network architectures that underpin modern diffusion and large vision models. It reviews foundational building blocks including convolutional layers and U-Net structures, then transitions to self-attention mechanisms and the rise of Diffusion Transformers (DiT) as the dominant architectural paradigm. Additional topics include multimodal DiT designs that integrate text and visual embeddings, and practical architectural optimizations for scaling generative vision models efficiently.

Lecture 5 is where the course stops treating the image generator as a black box and starts asking what should actually sit inside it. The first four sessions built the training objective from three angles — denoising, score matching, and flow matching — then moved the problem into a compressed latent space produced by a VAE, patched the blur with adversarial and perceptual losses, and added classifier-free guidance so a prompt could steer the output. None of that said anything about the network doing the work. This session is about architectures, and the lecturer says up front that the aim is not a catalog of every variant but an explanation of why the same handful of components keeps reappearing. That framing is worth carrying through the whole notes set: almost everything here is motivation, not enumeration.
 

Before any architecture appears, he pins down the interface. Three things go in: the noisy latent $x_t$, a noise-level indicator $t$, and a condition $c$. One thing comes out — for consistency he commits to predicting velocity, on the grounds that most current models do, while noting that noise or score parameterizations are roughly equivalent. Text-only conditioning gives you the familiar text-to-image family; text plus an input image gives you editing models. Then comes the wish list, and it quietly organizes the rest of the lecture: understand global structure, preserve local detail, accept external signals, and scale.
 

U-Net: How One Family of Architectures Encodes the Way Humans Scan Images

 

The first candidate is motivated anthropomorphically. Humans look at a picture by fixating on one region and scanning across it (10:02), so the lecturer proposes building in a bias that the model should behave the same way. He names the term explicitly — inductive bias — and polls the room; only a few hands go up. The operation that instantiates this bias is convolution, and he spends a little time on vocabulary: the output is a feature map, the filters are three-dimensional $F \times F \times C$ rather than flat squares, there are $K$ of them, and the step size is the stride. The learned parameters are the filter weights (plus optional biases); pooling, by contrast, learns nothing at all.
 

The interesting move is where he turns that bias against itself. Checked against the wish list, convolution clearly handles local detail and tractable computation. Global structure is the problem, and he formalizes it through the receptive field — how much of the input a single output activation can see (14:42). With a $5 \times 5$ filter stacked twice you have seen a handful of pixels; real images are $500 \times 500$ or larger, so you would need an absurd number of layers. His fix is the one the field settled on: shrink the spatial grid so receptive fields grow early, then grow it back. Pooling or strided convolution on the way down, transpose convolution on the way up, and skip connections that crop and copy feature maps from the down path and concatenate them into the up path. The result is the U-Net, named for the shape of the diagram rather than for any principle.
 

Two student questions at this point are worth recording, because they target exactly the confusion the U-Net invites. First: isn't this just an autoencoder? No — an autoencoder's job is reconstruction, whereas here the network takes a noisy latent and predicts the quantity that will remove noise from it. Second: must the output match the input's dimensions? Yes, and the justification is purely arithmetic — the flow-matching update adds velocity times a step size to the current latent, so the shapes have to line up.
 

Conditioning is then handled in two halves: representation and injection. For timesteps he offers the wristwatch analogy — hours, minutes and seconds run at different rates, which motivates encoding $t$ as a vector of sinusoids at different frequencies (28:29). For text he lists options without picking one: a learned vector per predefined class, per-token representations, embeddings lifted from a pretrained LLM, or the CLS token of a ViT. Injection has three routes — add the vector straight onto the feature map, use it to modulate the features, or run cross-attention — and he states plainly that there is no consensus (33:09). The historical note closes the section: U-Net appeared in 2015 for medical segmentation, entered generation through DDPM and latent diffusion, and was scaled up aggressively in Stable Diffusion XL in 2023.
 

Why Convolutions Hit a Wall, and What Attention Buys These Architectures

 

The pivot to transformers is staged as a puzzle. Given that the U-Net covers the wish list, why change anything? His answer is a scene: a teddy bear standing next to a mirror (36:14). Rendering it correctly requires nearly identical low-level texture in two distant regions of the same image, and convolutions mix information progressively as you move away from a pixel, so distant patches never meet cleanly. I found this the most memorable moment of the lecture, and also the least supported — it's asserted with a diagram, no failure case shown, and I'd have liked to see an actual U-Net output next to a transformer output on that prompt.
 

Attention is introduced through text first: queries, keys, values, a normalized dot product, and the associated value — with the explicit point that this mechanism exists to remove inductive bias, the opposite impulse from the convolution section. For images we only borrow the encoder half of the original transformer, since nothing is being translated. ViT's contribution is the preprocessing trick: slice the image into patches and treat patches as tokens. The generation analogue, the diffusion transformer, arrived in 2022 and does the same thing to a noisy latent. One detail worth noting from the patchify step: an $I \times I \times C$ latent cut into $P \times P \times C$ patches yields a sequence of length $I^2/P^2$, so smaller patches mean longer sequences and more compute. Granularity and cost are the same knob.
 

Inside the Diffusion Transformer: Patching, Conditioning, and Scaling

 

Three ways to inject, and the one that won

 

The DiT paper compared three conditioning routes: adaptive layer norm, which modulates tokens using the external signal; an extra cross-attention sub-block where image tokens query the timestep and label embeddings; and simply concatenating the condition vectors into the input sequence (45:19). Adaptive layer norm won. He previews FID here as the lower-is-better metric used to judge this, with the honest caveat that evaluation is lecture seven's topic — which means, as a listener, you're being asked to accept "adaLN performs best" on a yardstick that hasn't been defined yet. That's not a flaw in the argument so much as a scheduling artifact, but it's worth flagging when you review.
 

The brown fluffy teddy bear, and why per-dimension modulation makes sense

 

Rather than writing the adaLN formula and moving on, he builds intuition with a deliberately fake example (48:22). Suppose a patch embedding has interpretable dimensions — dimension 10 encodes brownness, 24 encodes roundness, 59 encodes fluffiness, 130 encodes whiteness. Early in sampling, at high noise, you want brownness and roundness to surface and you want neither fine fluffiness (too early) nor whiteness (wrong bear). Late in sampling, at low noise, the global shape is settled, so you want fluffiness up and brownness toned down. The desire is therefore a per-dimension intensity dial driven jointly by the prompt and by the noise level.
 

The mechanism follows naturally (52:58). Sum the timestep and condition embeddings, push them through a small MLP, and produce three vectors: $\alpha$ as a gate, $\gamma$ as a scale, $\beta$ as a shift. Normalize the token, scale by $1 + \gamma$, shift by $\beta$, apply the sub-layer, then gate the whole contribution by $\alpha$ before adding it back to the residual stream. Setting $\gamma = 0$ makes the scaling a no-op; setting $\alpha = 0$ blocks the entire update. The "zero" in adaLN-zero means exactly that: initialize the gates to zero so that at the first training step the condition has no influence and the block behaves unconditionally (56:02). He notes that implementations typically zero all three, though zeroing the gate alone is sufficient.
 

A full forward pass, block by block

 

The walkthrough is the densest stretch of the session and probably the most useful to rewatch. Sample $z_0$ from a Gaussian in the VAE's latent space. Slice it into patches, linearly project each to dimension $D$, and add a position embedding — added here, he warns, only because it's the simplest option, not the current standard. Separately, add the timestep embedding to the label embedding, run the sum through an MLP, and get six vectors rather than three: two gates, two scales, two shifts, one triple for the attention sub-layer and one for the feed-forward sub-layer (1:07:02).
 

Inside a block: layer-norm the tokens, apply $\gamma_1$ and $\beta_1$, run self-attention, gate the result with $\alpha_1$ and add the residual. Repeat the same pattern around the feed-forward network with the second triple. The output matrix is what he calls contextualized — every patch embedding has now been computed in the presence of every other. A final normalization, linear layer and reshape produce the predicted velocity in the input's shape. Then the Euler step from flow matching advances the latent, the loop repeats until $z_1$, and the VAE decoder converts the final latent into pixels, "because we humans understand pixels."
 

Scaling is two knobs, not one

 

Parameter count alone is misleading, he argues, because you can hold parameters fixed and still change patch size, and patch size feeds straight into sequence length. FLOPs capture this, and the punchline is that neither model size nor the other factor scaled alone delivers much; scaled together, results improve sharply (59:49). One transcription caveat: the second knob is consistently rendered as "batch size," but given that the surrounding argument is entirely about patch granularity and FLOPs, I'm fairly confident he means patch size. Worth confirming against the slide before you rely on it.
 

MM-DiT and the Per-Token Conditioning Problem

 

This is the section where the lecturer attacks his own architecture, and it's the sharpest limitation in the lecture. AdaLN modulates every token identically. Change the prompt to "a brown fluffy teddy bear surrounded by white walls" and you now want brownness emphasized where the bear is and whiteness emphasized where the wall is — one global dial cannot do both. His fix keeps self-attention intact and changes only how text enters.
 

Two families emerge. Cross-attention treats text as fixed: image patches are queries, text supplies keys and values, and each patch asks which words are relevant to it. Joint attention puts both modalities into the same self-attention operation. His analogy is a painter and a poet — cross-attention is the poet handing over a finished instruction sheet, joint attention is the two of them in the same room negotiating, so the poet can decide on smaller lettering when the painter runs out of space. Most current work, he says, uses joint attention, and the name coined for the resulting architecture is MM-DiT — multimodal diffusion transformer — introduced by Stable Diffusion 3 in 2024.
 

A second axis cuts across this: whether each modality gets its own feed-forward path or shares one. Single stream treats everything identically; double stream gives image and text separate projections; hybrid mixes both layer types. He points at Qwen-Image as a double-stream example (the transcript spells it "Gwen image"), Z-Image as single stream, and Flux.1 Kontext as a hybrid, all 2025 releases, with 2026 work continuing the trend. He also explicitly says these papers are not​ on the final — useful to know if you're triaging.
 

Position Embeddings: RoPE, 2D Tricks, and an Admittedly Open Problem

 

Roughly an hour in, a second speaker takes over (the transcript names him Shervin, referring back to the first speaker as Afshin) and stays with one narrow question: how do you tell a transformer where a patch is?
 

He starts in 1D. Write an embedding as content plus position, take the attention dot product, and a position-position term falls out. The original transformer hardcoded that term with sines and cosines whose denominator is a large base raised to an exponent that sweeps across the embedding dimension (59:49 onward). Low indices oscillate fast, high indices barely move. Because $\sin a \sin b + \cos a \cos b = \cos(a-b)$, the dot product between two positions collapses into a sum of cosines of their difference, which is exactly the relative notion you wanted; the plotted curve drifts downward as distance grows, though not monotonically, and he explains the shape — an initial drop as every cosine leaves zero, then growing wobble as slower terms start to oscillate.
 

Then he turns critical, and this is the second genuine limitation of the session. The motivation for positional information lives in the attention computation, yet absolute embeddings are added to the input, "at the place where that information... doesn't matter most." Worse, those added vectors produce cross terms with the content embeddings that nobody asked for. Not fatal, he concedes, just misaligned with the stated intent. Everyone did it anyway: the original transformer hardcoded it, ViT learned it, and DiT hardcoded a 2D version at the input level.
 

RoPE, from 2021, relocates the information: rotate queries by their position and keys by theirs before the dot product, so the dependence on $n - m$ appears naturally inside attention. Extending to 2D is where things get unsettled. Axial RoPE dedicates parts of the vector to the x-axis and other parts to y; mixed RoPE interleaves rotations for both axes within the same vector. A Fourier-transform analysis in the paper he cites found that axial encoding produces artifacts traceable to the missing x-y interaction, while mixed encoding spends the same frequency budget more effectively. Beyond that, resolution changes the meaning of coordinates, so one approach recenters the grid with zero at the image center (Seedream 2.0, rendered "SeaDream 2"). And once text tokens share a joint attention block, they need positions too — placing them along the diagonal keeps the model from reading text as a spatial continuation of the image, a choice attributed to the Qwen-Image paper.
 

He closes by calling positional encoding for images an open problem: many papers, many variations, no single standard, and real trade-offs between them. That's a fair summary of the whole second half — the lecture is more honest about what isn't settled than most survey talks manage to be. It ends, pleasingly, with the teddy bear from the editing example dancing, and a good weekend.
 

If you're reviewing, the two things I'd rewatch are the six-coefficient block walkthrough and the axial-versus-mixed comparison; both move fast and both are the kind of thing that looks obvious on a slide and stops being obvious five minutes later.
 

Content Disclaimer:

This article is for general reference only and does not constitute professional R&D guidance, production process advice or quality certification. All material performance data has specific test premises; readers should verify parameters against actual equipment and working conditions.