Note Wisdom
Annotated notes from Stanford CME296 Lecture 6 on training text-to-image diffusion models: flow-matching loss, timestep and resolution scheduling, REPA, pre-training data and curriculum, post-training alignment, DreamBooth/LoRA personalization, and distillation for faster production inference.
Institution: Stanford
Original Course: Stanford CME296 Diffusion & Large Vision Models | Spring 2026 | Lecture 6 - Model Training
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 walks through the full training lifecycle of generative vision models, from initial pre-training to downstream adaptation. It covers pre-training strategies and curriculum learning, continued pre-training and supervised fine-tuning pipelines, preference optimization methods such as Flow-GRPO and Diffusion-DPO, parameter-efficient adaptation techniques including DreamBooth and LoRA fine-tuning, and knowledge distillation for improving deployment efficiency.
Lecture six is where this course stops being a theory class and starts being a shop class. The first half of the semester built the math — DDPM, score-based SDEs, flow matching, U-Nets and diffusion transformers — and this session is the payoff: what does model training actually look like when the goal is a text-to-image system someone will pay to call? It runs a touch over 100 minutes and, unusually for this series, it changes hands about three-quarters of the way in (1:16:34), when the first lecturer passes the floor to Sherwin for the distillation block. The register shifts noticeably there: the opening stretch is Socratic and keeps throwing questions at the room, the back half is more of a critical tour of papers.
There's also an explicit narrowing of scope right at the top, worth knowing before you start. A text-to-image stack has several trained pieces — the denoising backbone, the VAE that buys you a cheaper latent space, and the text encoders that turn a caption into conditioning. Today is only about the backbone. The embedding models get pushed to the following week, when multimodal LLMs come up. If you arrived hoping for detail on how the caption encoder is trained, you'll be waiting.
The organizing frame is a life cycle with four stops. Pre-training teaches the model to produce images at all — any images — and it's the expensive one, heavy on compute and heavier still on data work. Post-training takes a model that can already make pictures and teaches it to make good ones, with the promise that "good" will get defined later; the lecture mostly keeps that promise. An optional third stage, called tuning here, is personalization — bending the model toward one specific subject. The fourth stop isn't really training so much as a shipping constraint: distillation, so the system can generate in fewer steps and stop costing a fortune per query.
The recap that opens things is short and worth having even if you sat through lecture five. The U-Net dominated the early 2020s on the strength of a convolutional inductive bias that loosely mimics how people scan a picture — local detail on the way down, global structure in the bottleneck, skip connections ferrying fine detail back up the other side. Its weak spot is the absence of any direct wiring between distant regions, which is fatal when an image only makes sense if two far-apart patches agree; the example given is a person looking at themselves in a mirror. Diffusion transformers, from late 2022 onward, fix that by letting every patch attend to every other. Then come the variants — adaptive layer norm for injecting timestep and conditioning, and the multimodal DiT, which promotes text to a stand-alone input modality rather than smuggling it in as a uniform modulation across every patch.
I appreciated the candour about the previous session having aimed at intuition for why these components exist, not at completeness. It sets expectations for this one too.
Before the four stages get going there's a long preliminary block — the better part of forty minutes — on how the objective is set up in practice. It's the densest part of the lecture.
Start with the loss. The first three lectures each gave a different view of the same problem: DDPM parameterizes a network to predict the noise that was added and minimizes an L2 distance to it; the score-based view treats the score as a compass pointing back toward clean data and regresses against the true score of the perturbed distribution; flow matching reframes the whole business as transport and regresses a velocity field you integrate with something like Euler's method. The lecturer's position is that the field has settled, and from here on the default for the course is the flow-matching objective. Useful to say out loud, since the slides up to that point keep all three in play.
Then comes the stretch I'd flag as the most interesting twenty minutes of the session. The flow-matching loss takes an expectation over timesteps, and in textbook form the timestep is drawn uniformly. The lecturer stops and asks whether that's sensible, and walks the room through three cases. With almost no signal — maximum noise — the best any model can do is guess the average direction of the entire dataset. Easy, and no amount of extra training turns it into more than an average. Near the clean end the image is nearly finished and only small tweaks remain, which is also easy. In the middle, though, you know roughly what the object is and roughly where its parts go, but every consequential decision is still open: for a teddy bear that means fur texture, where the eyes sit, how the limbs are arranged.
So uniform sampling spends as much of the budget on the two easy regimes as on the hard one. The practical fix is to draw timesteps from a logit-normal distribution instead (22:46). The construction is a small trick: take the logit of t, which maps the unit interval onto the whole real line, draw that from a normal, then push it back through a sigmoid so the result is guaranteed to land between zero and one. You can't just sample t from a normal, because a normal will happily hand you values outside the interval. The payoff is a knob — mean and standard deviation — that lets you aim the sampling at whichever region of the trajectory matters for your use case.
The second adjustment is more surprising. Take one image at low resolution and the same image at high resolution, add identical amounts of noise, and the low-res version looks noisier. The explanation offered is spatial correlation: neighbouring pixels in a real image tend to share values, and the added noise is zero-mean, so a high-resolution image gives you many noisy samples of roughly the same underlying value that average out, while a low-resolution image has fewer neighbours to lean on and a corrupted pixel genuinely destroys information.
What follows is the one piece of live algebra in the lecture, and it was the hardest thing for me to follow in real time. The setup is a toy image, single channel, every pixel equal to the same unknown constant. Perturb it at some noise level and average across all N pixels; that average has mean (1−t)C and variance t²/N. Solve for the constant and the standard deviation of the estimate works out to t/(1−t) times one over the square root of N, so uncertainty scales as the inverse square root of resolution. Equate that uncertainty between a low-resolution N and a high-resolution M and you can solve for the timestep at M that makes the two look equally noisy. The algebra gets waved off as skippable, with the insistence that the mindset is what matters — fair, but the derivation is doing real work in that section, and a worked numeric example would have served me better than the symbolic result.
Two details that are easy to miss. The paper being followed uses the opposite timestep convention from the rest of the course — zero clean, one noisy — adopted temporarily so the formula comes out right. And in response to a question, M and N here are latent-space resolutions, not pixel counts.
REPA, short for representation alignment, is the first practical accelerator (44:54). Your diffusion transformer is already computing internal features, and there exist pre-trained encoders that are very good at representing images semantically. So project the hidden output of some DiT layer, compare it against the encoder's representation of the same patch in its clean form, and penalize the distance. The analogy used is handing a book to someone trying to learn a topic.
The claimed effect is a large speed-up in training — the number read off the quality-versus-iteration curve was on the order of eighteen times, which I'd treat as a property of that graph and that setup rather than a general constant. Two ablations get discussed. Aligning at earlier layers beats aligning at later ones, with the offered interpretation that late layers are busy with local detail while the pre-trained encoder deals in semantics, so the match should happen where the model is still thinking semantically. And larger models seem to gain more from the technique than smaller ones.
The pre-training section is short, deliberately. One slide taken from the Qwen technical report stands in for the entire data pipeline: filtering, resolution handling, cleaning and so on. The message is about proportion rather than technique — this work is time-consuming and it matters, even though it earns one slide.
The concept that does get airtime is curriculum learning, meaning the training material is ordered by difficulty instead of shuffled. Easy means low resolution, a fixed square aspect ratio, and a short simple caption — the running example for the whole lecture is a teddy bear. Hard means high resolution, arbitrary aspect ratios, and long tangled prompts. A natural question from the room, how you train on mixed resolutions at all, gets a clean answer: with a transformer backbone, resolution becomes sequence length. The latent is cut into patches, higher resolution means more patches and therefore more tokens, much like a longer context window in a language model.
Most of the lecture lives here, and the split is drawn along a distinction I hadn't seen made so sharply. Continued training adds knowledge: you had a model trained on nature scenes and cars, you want teddy bears, so you keep training on teddy bears. Supervised fine-tuning targets behaviour instead — better lighting, more pleasing composition, closer adherence to what the prompt actually asked for. Same machinery, different intent.
Then a third category, framed as the missing half of everything so far: until now the model has only been shown what to do, never what not to do. Preference methods work from comparisons — this image beats that one — collected from human raters or from some automated scorer, which is promised for the next lecture.
The rest is a tour of techniques imported from language modelling, and the lecturer says outright that this is the pattern to watch. Reward feedback learning trains a reward model that scores a prompt-image pair using a pairwise loss of the Bradley–Terry type — the math is skipped and pointed at the sister course, CME 295 — and then, since that reward model is differentiable, pushes its gradient back into the generator to raise expected reward. Flow GRPO adapts the GRPO algorithm: for each prompt you deliberately generate a diverse set of images, the authors framing that diversity through an SDE, score each one with a black-box reward model, and treat the reward relative to the group as an advantage driving a policy update. Diffusion DPO is the direct analogue of DPO, and the takeaway from its three lines of equations is unusually concrete — make the model better at predicting velocities that lead to winning images and worse at predicting the ones that lead to losers.
All three carry the same hazard, where the proxy you optimized drifts away from the thing you wanted, and the standard guardrail shows up as a KL term between the updated policy and the old one, keeping any single step from wandering too far.
My favourite moment in the session is a small piece of theatre near the end of this block (1:05:20). The lecturer puts up a lovely picture of a teddy bear reading and asks the room what prompt produced it. People guess; it's something of a trick. The actual prompt is a truncated monster — a warm intimate indoor reading scene featuring a plush teddy bear as the central object, going on about the objects and the lighting, roughly four times longer than what fits on the slide.
The point is a mismatch. Models get aligned on richly specified captions, while real users type five words. Prompt enhancement is the layer in between that rewrites the short request into the long, in-distribution one the model was trained against. The closing analogy for the whole pipeline makes it stick: pre-training is learning what food is, continued training is reading recipes, supervised fine-tuning is plating and presentation, preference tuning is a food inspector tasting several versions of the same dish and telling you which was better, and prompt enhancement is the waiter who translates a customer's order into the specification the kitchen actually works from.
Personalization is the optional third stage, and DreamBooth is the worked example. You have a handful of photos of one particular bear in a few poses, and you want that bear and no other. The recipe binds the subject to a rare token, written as [V], so the model associates the object with a word carrying no prior baggage.
The lecturer asks the room what goes wrong and gets the answer they wanted: the model over-specializes. Drive it hard on a few images and it forgets what it could previously draw, and diversity collapses. The fix is a prior preservation loss — keep training on prompt-image pairs the model already handles well, mixed in with the new subject, so the objective says learn this bear without losing the rest.
Then the practical problem: these models are big. The largest open text-to-image models are said to sit in the tens of billions of parameters, and you don't want to update all of that for one bear. So LoRA comes in from the LLM world — freeze the base weights and train a pair of low-rank matrices whose product gets added on, capturing most of the effect at a fraction of the cost. DreamBooth plus LoRA is what you'll actually see in the wild. The trade-off is stated plainly: excellent if you need thousands of images of one subject with high fidelity to the input photos, but you pay a training cost per subject and can't reuse what you trained for anything else.
After the handover the register changes. The motivating argument is economic before it's technical: a high-volume provider has no incentive to spend a thousand denoising steps per image, businesses often want good enough rather than best, real-time applications exist, and nobody wants to waste compute. Keep most of the quality, cut the number of function evaluations.
The framing opens in familiar territory. Classical distillation has a large teacher and a smaller student matching its outputs, typically through a KL divergence between distributions; BERT to DistilBERT is offered as the quantitative example, roughly halving parameters while retaining about 97% of performance, with the lecturer hedging the numbers from memory, which I appreciated. But for image models, shrinking the student isn't the knob that pays — it degrades quality badly — and the low-hanging knob is the step count instead.
Why not simply regress the teacher's output in one step? Because it's too hard; the borrowed metaphor is asking a student painter to reproduce a finished canvas in a single stroke. Progressive distillation (1:22:19) is the answer — halve the step count, train a student to do in one step what the teacher did in two, then repeat, each generation becoming the next one's teacher. The image used is a chain of identical brothers, log-two of the original step count of them, each painting one stroke for every two of the one before, visualized as secants drawn across the trajectory at each halving. Its practical charm is that you can pick a step budget and read the quality off the curve, and it reportedly beats DDIM at matched step counts.
Then the critique, and this is where the back half gets interesting. Progressive distillation takes the sampling path as given; it never tries to make the path easier, only to walk it faster. But lecture three introduced reflow, which straightens those paths. InstaFlow (1:25:16) combines the two: sample noise, integrate the ODE, collect input-output pairs, fit a new model on them to get straighter trajectories and fewer Euler steps, then distill on top to lock in a single step. A question from the room — why not just do more reflow rounds? — gets a neat answer: each round injects discretization error into the pairs you're fitting, so more isn't better, and the first round captures most of the straightening anyway.
The loss choice deserves a note because it recurs. Warm up with plain MSE between student and teacher output, then switch to LPIPS, which pushes both images through a frozen pre-trained network and compares feature maps layer by layer. The reasoning echoes an earlier lecture: per-pixel error tells you about values, not about whether the image looks right. There's also a pointed question about why straight paths aren't simply enforced during pre-training — the random noise-to-target pairings you train on aren't optimal, and the resulting curvature is a side effect of the MSE objective rather than a design goal, which is precisely what reflow exists to clean up.
Consistency models (1:30:32) attack from the other end. Rather than shortening the walk, demand that every point along a deterministic trajectory map straight to its endpoint, so you can jump from noise to clean image immediately. Two ways in: consistency training, which noises a clean image at two adjacent levels and forces the outputs to agree, and consistency distillation, where a teacher does the denoising and the student learns to be self-consistent along the teacher's path. Someone in the room immediately spots the failure mode — a student could satisfy consistency by always outputting zero. The two guards given are that the per-step prediction is a constrained parameterization rather than a free function, and that gradients aren't pushed through both students; one is held as a slow-moving average while the other does the predicting. Doing this in latent space is what the literature calls LCM.
The last stretch is the most compressed and the hardest for me to hold onto. Regressing to a teacher output produces regression to the mean and soft-looking images, the thing GANs are known to fix, so the lecture pivots to a distribution-level objective: give the same noise to teacher and student, re-noise the student's output, and ask both a teacher and a student-like model to denoise it, using the difference in the advice they give as the signal. Written as a KL between the student's distribution and the target, its gradient has a tidy reading as a difference of scores — or of velocities, in flow-matching terms — meaning do more of what the teacher suggests and less of what the student suggests, with a regression term alongside for stability. The Q&A confirms what you'd suspect: without the teacher-student comparison you collapse onto the teacher's mode, without the regression you may never reach the target distribution, and balancing the two is hyperparameter work.
Then a rapid-fire list of complaints about that setup — it compares against a teacher rather than ground truth, the student-advice branch needs an expensive diffusion loop and duplicates the very model being trained, and the teacher-minus-student signal is a bit soft. ADIDA (1:39:17) reframes the problem as classification, which by making it harder sharpens the student's outputs. LADD gets mentioned as the fix for the constant shuttling between pixel and latent space, moving the whole thing into latent space only. The session closes on trade-offs: consistency models are easy to tune with lightweight adapters, GAN-style training brings dynamics that are hard to manage, the MSE-versus-LPIPS tension never really goes away, and the trend is to do all of it in latent space.
A few threads felt unfinished, and not only because the session was running over.
The resolution-shifting derivation is the one I'd most want another pass at. The spatial-correlation intuition landed; the algebra that converts it into a timestep schedule was quick, self-described as skippable, and then immediately used. A before-and-after pair of images with the shifted schedule applied would have taught me more than the symbolic result.
The eighteen-times figure for REPA arrives with little scaffolding about what's held constant, and the explanation for why earlier layers work better is offered explicitly as one possible reading rather than a finding. I believe the direction of the result; I'd want the ablations before quoting the magnitude.
The distillation finale is genuinely tough, and the lecturer says as much — apologies for the big formulas, no good illustration available. Fair enough, but it means the most recent material in the session is also the least recoverable from these notes alone.
And there's a structural point worth naming. Nearly every technique in the post-training block carries a name imported from language modelling — GRPO, DPO, LoRA, continued training, supervised fine-tuning. The lecture treats this as an observation about where the field is, and it's a useful one, but it also means the session quietly inherits the LLM community's unresolved argument about how much of preference tuning is genuinely aligning a model versus overfitting a reward proxy. That doesn't get examined here.
Still, as a map of what happens between a randomly initialized denoiser and a product people actually use, the shape of model training has rarely been laid out this clearly for me. The four-stage frame is what I'd write on the inside of the notebook cover.
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.
All contents below are exclusive to the paid Word file, NOT available on this web page

