[{"content":"Paper: Sparks of Weight-Space Programming: LLMs can natively read \u0026amp; write neural network weights (Takehi, Teoh, Zhang, Khattab, MIT CSAIL, Nov 2026)\nLooping signal in the residual stream012‖d‖2.630.16Per-layer increment along u00.20.4051015192022layerMLPs of layers 15 and 19:the two largest writersMLPattention The looping signal grows layer by layer; Astra edited the two largest MLP writers, layers 15 and 19. The paper uses Astra to edit the weights of Qwen3.5-0.8B directly, with no training or gradient descent, to make it better at math, which raises MATH-500 from 7.9% to 48.1%. Actually, it is not magic or uninterpretable. Astra uses a method from directional ablation, introduced by Arditi et al. (2024) to remove refusal from chat models. This note is mainly about directional ablation, with Astra\u0026rsquo;s edit as a worked example.\nPart 1 explains directional ablation as introduced by Arditi et al. (2024). Part 2 follows how Astra applied the same idea to Qwen3.5-0.8B. Part 3 covers how Astra went beyond plain ablation to reach the final edit.\n0. Setup: how an MLP writes into the residual stream Qwen3.5-0.8B has hidden size 1024 and FFN intermediate size 3584. For the MLP in layer $\\ell$:\n$a \\in \\mathbb{R}^{3584}$: the MLP\u0026rsquo;s intermediate activation at one token position (the elementwise product of gate and up in SwiGLU). $a_j$ is the activation of neuron $j$. $W \\in \\mathbb{R}^{1024 \\times 3584}$: the down_proj weight. Its $j$-th column $w_j \\in \\mathbb{R}^{1024}$ is the vector neuron $j$ writes into the residual stream. $m = W a = \\sum_j a_j w_j \\in \\mathbb{R}^{1024}$: the MLP output added to the residual stream. Think of the MLP as 3584 writers: writer $j$ adds its fixed vector $w_j$ to the residual stream, scaled by its activation $a_j$.\nPart 1. Directional ablation (Arditi et al., 2024) 1.1 Idea Find a unit vector $u$ in residual space that represents some behavior, then modify the weights so that they can no longer write anything along $u$, keeping only the component of each weight column orthogonal to $u$.\n1.2 The operation Let $u \\in \\mathbb{R}^{1024}$ with $\\|u\\| = 1$. The component of the MLP output along $u$ is\n$$u^\\top m = u^\\top W a = \\sum_j \\underbrace{(u^\\top w_j)}_{c_j}\\, a_j$$where $c_j = u^\\top w_j$ is a scalar: how much neuron $j$ writes along $u$ per unit of activation.\nEach column decomposes into a part along $u$ and a part orthogonal to it:\n$$w_j = \\underbrace{(u^\\top w_j)\\,u}_{\\text{along } u} + \\underbrace{\\big(w_j - (u^\\top w_j)\\,u\\big)}_{\\text{orthogonal to } u}$$Directional ablation keeps only the orthogonal part of every column:\n$$W' = (I - u u^\\top)\\,W = W - u\\,(u^\\top W)$$Here $uu^\\top$ is a $1024 \\times 1024$ projection matrix and $u^\\top W$ is the length-3584 row vector of all $c_j$. The projection acts on the output side (the 1024-dim residual space), not on the 3584-dim activation side. After the edit, $u^\\top W' = u^\\top W - (u^\\top u)(u^\\top W) = 0$, so for any activation $a$ the layer writes nothing along $u$, while the orthogonal part $(I - uu^\\top)W$ is unchanged. This is a rank-1 edit. There is no strength parameter: the component along $u$ is removed exactly once.\n1.3 Finding $u$: difference of means Arditi et al. take $u$ to be the difference between the mean activations of harmful and harmless prompts, normalized.\nWhy a difference. The underlying assumption is the linear representation hypothesis: a concept corresponds to a fixed direction, and the stronger the concept, the larger the hidden state\u0026rsquo;s component along it:\n$$h_i = b + z_i\\, v + \\epsilon_i$$ $h_i \\in \\mathbb{R}^d$: hidden state of prompt $i$ at some layer and position. $b$: the part shared by all prompts (chat template, overall activation mean). $v$: the concept direction we want. $z_i$: the concept\u0026rsquo;s strength on prompt $i$ (large for harmful, small for harmless). $\\epsilon_i$: content-specific variation (topic, wording) that varies unsystematically across prompts. Taking group means and subtracting:\n$$\\mu_{\\text{harm}} - \\mu_{\\text{safe}} = (\\bar z_{\\text{harm}} - \\bar z_{\\text{safe}})\\, v + (\\bar\\epsilon_{\\text{harm}} - \\bar\\epsilon_{\\text{safe}})$$Subtracting removes the shared part $b$, which is usually much larger than the $v$ component. Averaging makes the $\\epsilon$ terms cancel; the more samples, the closer $\\bar\\epsilon_{\\text{harm}} - \\bar\\epsilon_{\\text{safe}}$ is to zero. What remains is approximately a multiple of $v$, and normalizing gives $u$.\n1.4 Limitation: confounds A difference of means captures every systematic difference between the two groups, not just the concept of interest. If harmful prompts are, say, systematically longer, a length direction $g$ also differs between groups:\n$$h_i = b + z_i\\, v + \\ell_i\\, g + \\epsilon_i \\quad\\Rightarrow\\quad \\mu_{\\text{harm}} - \\mu_{\\text{safe}} = (\\bar z_{\\text{harm}} - \\bar z_{\\text{safe}})\\,v + (\\bar \\ell_{\\text{harm}} - \\bar \\ell_{\\text{safe}})\\,g$$The second term does not cancel, so $u$ is a mix of $v$ and $g$. Two consequences follow: ablating $u$ also damages the unrelated direction $g$ (collateral damage), and the part of $v$ orthogonal to $u$ is left in place (incomplete removal).\n1.5 What \u0026ldquo;never writing along $u$\u0026rdquo; means $u$ is an axis; the difference vector is only the tool used to pick it. Define a hidden state\u0026rsquo;s coordinate on this axis as $p(h) = u^\\top h$. By construction, all of the gap between the two group means lies on this axis: $u^\\top(\\mu_{\\text{harm}} - \\mu_{\\text{safe}}) = \\|\\mu_{\\text{harm}} - \\mu_{\\text{safe}}\\|$, and the gap is zero along every direction orthogonal to $u$.\nArditi et al. apply the projection to every matrix that writes into the residual stream: the embedding, every attention output projection, and every MLP down_proj. The residual at any position is the sum of these writers:\n$$h = e + \\sum_\\ell \\big(m_\\ell + o_\\ell\\big)$$Since each writer satisfies $u^\\top W' a = 0$ for any input, the sum also gives $u^\\top h = 0$ for every prompt and every token. So it is not just that the two groups\u0026rsquo; difference becomes zero; each group\u0026rsquo;s coordinate becomes zero. \u0026ldquo;Equal\u0026rdquo; would be a weaker condition: if the groups sat at 3 and 1 on the axis, moving both to 1 also makes the difference zero, but leaves content along $u$. Ablation instead empties the axis entirely. (Moving one group to the other\u0026rsquo;s position would be activation steering, which is a different operation.)\nThree things follow:\nThe gap between the group means along $u$ is gone. At this linear level, the average \u0026ldquo;harmful\u0026rdquo; and \u0026ldquo;harmless\u0026rdquo; representations can no longer be told apart. Any downstream component that reads the concept is computing something like $r^\\top h$. Whatever part of that readout depends on the $u$ component now reads 0, so the signal does not propagate. Everything orthogonal to $u$ (topic, syntax, and so on) is untouched, so the rest of the model\u0026rsquo;s abilities are largely preserved. 1.6 Result on refusal The model stops refusing harmful requests. Downstream components decide whether to refuse by reading the residual stream\u0026rsquo;s component along $u$; with $u^\\top h \\equiv 0$, a harmful prompt looks like a harmless one on this axis and the model answers normally. This held across more than a dozen open chat models up to the 70B scale, with only small changes on standard capability benchmarks, because only one direction was removed. For this reason the paper calls $u$ the refusal direction: it behaves like a \u0026ldquo;the model intends to refuse\u0026rdquo; signal rather than the model\u0026rsquo;s understanding that content is harmful.\nPart 2. Astra on Qwen3.5-0.8B: the same recipe Astra cannot run the model directly. Each forward pass is a request to the harness that specifies what to return. It had 50 forward passes in this run, and every training-set problem it used came from the GSM8K and MATH training sets; it never saw MATH-500.\n2.1 Diagnosis The first forward pass generated answers on 96 practice problems (32 GSM8K, 64 MATH). Only 11 were correct and 85 hit the 32,768-token limit, with the truncated tails about 95% repeated 8-grams. Astra concluded that the main problem was looping, not inability to do the math.\n2.2 Building four contrast pairs Astra picked four problems from those 96 and, for each, wrote two continuations that share the same prompt and the same correct opening reasoning, then diverge. Both texts in every pair were written by Astra, not generated by Qwen.\nThe productive version is correct reasoning that keeps moving forward. It stops mid-thought inside \u0026lt;think\u0026gt; and does not give a boxed final answer. The repetitive version is a short, clean imitation of the loops Astra saw in Qwen\u0026rsquo;s real outputs. For example:\nPrompt: Rafael works 10 hours on Monday and 8 on Tuesday, with 20 hours left to work. He earns $20 per hour. Find his weekly pay.\nProductive: The first two days total 10+8=18 hours. Adding the remaining 20 gives 38 hours. His pay is 38*20=760 dollars. Checking: 200+160+400=760, consistent with the total. The hours and hourly rate have the correct units.\nRepetitive: The first two days total 10+8=18 hours. Adding the remaining 20 gives 38 hours. But is 20 the total? No, it is the remaining time. But is 20 the total? No, it is the remaining time. But is 20 the total? No, it is the remaining time.\nQwen\u0026rsquo;s real loop on this problem kept asking whether the 20 hours were remaining or total; on the arithmetic-sequence problem it repeated \u0026ldquo;Wait, 88−85=3. Wait, 85−82=3.\u0026rdquo; Astra\u0026rsquo;s version of the latter was \u0026ldquo;Now 88-85=3 and 85-82=3.\u0026rdquo; repeated.\nThe contrast is therefore progressing vs. stuck, not correct vs. wrong: the shared prefix is correct in both, and the repetitive version keeps re-asking an already-settled question.\nProblem Source Original model in forward pass 1 Rafael\u0026rsquo;s weekly pay gsm8k_train_1024 Looped, truncated Remainder of $6x^4-14x^3-\\dots$ divided by $2x-6$ math_train_intermediate_algebra_738 Looped, truncated Terms of 88, 85, 82, … before −17 math_train_algebra_808 Looped, truncated Two necklace machines gsm8k_train_1766 Correct (~7,400 tokens), not truncated Only three of the four problems actually looped; the logs do not say why the necklace problem was included. The practice files also contain reference solutions, which Astra had access to when writing the productive versions.\nA caveat from Section 1.4 applies here: the two versions use different words (the repetitive Rafael text says \u0026ldquo;But is 20 the total?\u0026rdquo; three times), so $u$ may partly encode those specific phrases rather than only the looping state. Astra was aware of this and tried a content-matched variant: the difference between the hidden states at the first and the last occurrence of the same repeated sentence. Both sides then have exactly the same tokens, so the only remaining difference is the state \u0026ldquo;this is the $N$-th repetition\u0026rdquo;. The final edit still used the productive-vs-repetitive direction described below.\nThere is a second limitation: with only 4 pairs, the averaging step can cancel very little of $\\epsilon$. A direction estimated from so few samples is a hypothesis, which is why so much of the run (Sections 3.2 and 3.3) went into causal checks.\n2.3 Teacher-forced forward pass: what is recorded Teacher forcing means feeding a fixed, complete text through the model in one forward pass with no sampling. At each position the model predicts the next token given the true preceding text, and the harness returns per-token log-probabilities plus any requested internal states. This allows controlled comparisons between matched texts at low cost.\nThe 8 texts (4 pairs) were run through the original model in the 4th forward pass, returning two kinds of per-token data:\nData Where Shape per text Used for Residual hidden state $h$ 25 positions $L = 0, \\dots, 24$ $T_i \\times 1024$ per position Computing $u$ and choosing layers MLP intermediate activation $a$ Inside each of the 24 MLPs $T_i \\times 3584$ per layer Layer attribution and choosing the 64 neurons Position $L = 0$ is the embedding output; position $L \\ge 1$ is the residual after layer $L-1$ (both its attention/DeltaNet and MLP have been added). So layer 15\u0026rsquo;s output is at position 16, and layer 19\u0026rsquo;s at position 20.\nText lengths ranged from 98 to 146 tokens. The texts are padded to the longest and stacked, giving $H^{(L)} \\in \\mathbb{R}^{8 \\times T_{\\max} \\times 1024}$ per position, with productive texts at even indices and repetitive texts at odd indices. Only the first $T_i$ positions of text $i$ are used.\n2.4 From tokens to $u$ For each residual position $L$:\nAverage the tail of each text. Take the last 36 tokens of text $i$ and average over them: $$\\bar h^{(L)}_i = \\frac{1}{36}\\sum_{t = T_i - 35}^{T_i} h^{(L)}_{i,t} \\in \\mathbb{R}^{1024}$$ Only the tail is used because the prefix is identical within a pair; the difference lives in the continuation, and the end of the repetitive text is in the \u0026ldquo;already looping\u0026rdquo; state. This gives 8 vectors per position. No normalization yet. Subtract within each pair and average over pairs. $$d^{(L)} = \\frac{1}{4}\\sum_{k=1}^{4}\\left(\\bar h^{(L)}_{\\text{rep},k} - \\bar h^{(L)}_{\\text{prod},k}\\right) \\in \\mathbb{R}^{1024}$$ Normalize. $$u^{(L)} = \\frac{d^{(L)}}{\\|d^{(L)}\\|}$$ So each position gets its own unit vector $u^{(L)}$. Astra computed $d^{(L)}$ at every position to see where the contrast grows, but only two directions were used for edits: $u^{(16)}$ for layer 15 and $u^{(20)}$ for layer 19.\nThe shapes through the pipeline:\nStep Operation Shape Input 8 texts (4 pairs: productive + repetitive), length $T_i \\in [98, 146]$ token sequences Forward Teacher-forced pass on the original Qwen, residual at position $L$ $[T_i, 1024]$ per text; padded stack $[8, T_{\\max}, 1024]$ Tail Keep the last 36 tokens of each text $[36, 1024]$ per text Mean Average over tokens $[1024]$ per text, $[8, 1024]$ in total Diff Within-pair difference: repetitive − productive $[4, 1024]$ Average Average over the 4 pairs $d$: $[1024]$ Normalize $u = d / \\lVert d\\rVert$ $u$: $[1024]$, $\\lVert u\\rVert = 1$ $T_i$ is the actual length of text $i$; $T_{\\max} = 146$ is the longest. Padded positions are never used. The same pipeline runs once at $L = 16$ (for layer 15) and once at $L = 20$ (for layer 19).\n2.5 Choosing layers 15 and 19: per-layer attribution Layer $\\ell$ adds to the residual stream:\n$$h^{(\\ell+1)} = h^{(\\ell)} + \\text{attn}_\\ell + m_\\ell$$where $\\text{attn}_\\ell$ is the attention (or DeltaNet) output and $m_\\ell = W_\\ell a_\\ell$ the MLP output. Taking \u0026ldquo;repetitive − productive\u0026rdquo; on both sides gives an exact split of each layer\u0026rsquo;s increment to the contrast:\n$$d^{(\\ell+1)} - d^{(\\ell)} = \\Delta\\text{attn}_\\ell + \\Delta m_\\ell$$With $u_\\ell = d^{(\\ell+1)}/\\|d^{(\\ell+1)}\\|$, three scalars are computed per layer:\nTotal increment: $(d^{(\\ell+1)} - d^{(\\ell)})^\\top u_\\ell$ MLP part: $(W_\\ell\\, \\Delta \\bar a_\\ell)^\\top u_\\ell$, where $\\Delta \\bar a_\\ell \\in \\mathbb{R}^{3584}$ is the \u0026ldquo;repetitive − productive\u0026rdquo; difference of MLP activations, computed the same way as for hidden states (last-36-token mean, averaged over 4 pairs). Because $m = Wa$ is linear, this is exact. Attention part: total increment minus MLP part. Values from the logs (second half of the network):\nLayer $\\ell$ $\\lVert d^{(\\ell+1)}\\rVert$ Total increment MLP part Attention part 13 1.088 0.243 0.093 0.150 14 1.273 0.362 0.150 0.212 15 1.564 0.472 0.263 0.209 16 1.760 0.280 0.123 0.157 17 1.877 0.193 0.049 0.144 18 1.909 0.174 0.078 0.097 19 2.266 0.530 0.250 0.281 20 2.327 0.132 0.105 0.027 21 2.515 0.258 0.164 0.094 22 2.626 0.309 0.110 0.199 The contrast norm grows from about 0.16 at layer 0 to about 2.6, so the looping state is amplified layer by layer. The selection criterion is the MLP part, because the edit targets MLP down_proj; attention contributions cannot be changed by editing the MLP. Layers 15 (0.263) and 19 (0.250) have the two largest MLP parts across all 24 layers (third is layer 21 at 0.164), and also the two largest total increments. The single highest-scoring neuron in the same analysis was neuron 11 of layer 19.\n2.6 First edit: standard directional ablation Attribution only proposed candidates; experiments decided. Astra\u0026rsquo;s sequence was:\nZero out a single neuron (column 11 of $W_{19}$): weak effect. Directional ablation on layer 19 only: still weak. Directional ablation on layers 15 and 19 together: $$W'_{15} = W_{15} - u^{(16)}\\big(u^{(16)\\top} W_{15}\\big), \\qquad W'_{19} = W_{19} - u^{(20)}\\big(u^{(20)\\top} W_{19}\\big)$$ with $W_{15}, W_{19} \\in \\mathbb{R}^{1024 \\times 3584}$ and $u^{(16)}, u^{(20)} \\in \\mathbb{R}^{1024}$. On 128 new practice problems, accuracy went from 8.6% to 32.8%. Astra also probed wider sets (all MLPs in layers 12–22; layers 15, 19 and 22) before settling on 15 and 19, and later confirmed with single-layer isolation that both layers contribute.\nNote how this differs from Arditi et al. (Section 1.5). They project every writer, so $u^\\top h = 0$ everywhere. Astra edits only two down_proj matrices; the embedding, all attention outputs, and the other 22 MLPs can still write along $u$, so the residual stream\u0026rsquo;s $u$ component is not zeroed. What the edit does is stop the two biggest MLP amplifiers from growing the looping signal (and, at α = 6 in Part 3, actively push against it). \u0026ldquo;Cutting the two main amplifiers\u0026rdquo; is a more accurate description than \u0026ldquo;removing the direction from the representation\u0026rdquo;.\nPart 3. Beyond ablation: from α = 1 to the final edit 3.1 Adding a strength parameter α Astra generalized the projection with a strength $\\alpha$:\n$$W' = W - \\alpha\\, u\\,(u^\\top W)$$Arditi et al. have no such parameter; their operation corresponds to $\\alpha = 1$. (Some later open-source \u0026ldquo;abliteration\u0026rdquo; tools added a similar coefficient, usually near 1.) Per column, the component along $u$ becomes $(1 - \\alpha)\\,c_j$:\nα Component along $u$ after edit Meaning 0 $c_j$ No change 1 $0$ Standard directional ablation 2 $-c_j$ Reflection 6 $-5\\,c_j$ Reversed and amplified 5× For $\\alpha \u003e 1$ the edit no longer removes a direction; it pushes against it, which is closer to steering. Unlike activation steering, which adds a fixed vector to every input, this weight edit pushes in proportion to the neurons\u0026rsquo; activations.\nDose-response with all 3584 columns edited:\nα Result 1 128 new problems: 8.6% → 32.8% 2 Another 128 new problems × 2 samples: 8.2% → 43.0% 3 Exploratory: 8.2% → 48.0% Larger α gave larger gains, but also more answers that finished with a wrong result.\n3.2 Mechanism checks with teacher forcing Astra wrote six new contexts (geometry, probability, sequences, modular arithmetic, algebra, calculus), each with a productive and a repetitive continuation, plus controls that close \u0026lt;/think\u0026gt; before the calculation is finished. Comparing per-token log-probabilities before and after the edit, the repetitive continuations became clearly less likely while the productive ones were almost unchanged, and the probability of closing thinking prematurely stayed negligible. This argues against the edit simply forcing early stopping.\nAstra also found copying heads (layer 15 head 7, layer 19 head 5) that attend to the matching position in the previous repetition. Halving their outputs was less selective than the MLP edit, so this route was dropped.\n3.3 Controls Norm-matched random direction: a random direction orthogonal to $u$, with the same Frobenius norm of the change ($\\|W' - W\\|_F$, matched within 0.03%). Result: +1.6pp ± 3.5, i.e. no effect. The direction $u$ matters, not the size of the change. Single-layer isolation: editing only layer 15 or only layer 19 each reduced the likelihood of repetitive text, confirming that both layers contribute. 3.4 Restricting the edit to 64 neurons per layer Each neuron $j$ in layer $\\ell$ (15 or 19) gets a score:\n$$\\text{score}_j = \\underbrace{(u^\\top w_j)}_{c_j} \\times \\underbrace{\\Delta \\bar a_j}_{\\text{activation difference}}$$ $w_j$: column $j$ of down_proj; $u$: the layer\u0026rsquo;s direction ($u^{(16)}$ or $u^{(20)}$). $c_j$: how much neuron $j$ writes along $u$ per unit of activation. $\\Delta \\bar a_j$: neuron $j$\u0026rsquo;s activation difference \u0026ldquo;repetitive − productive\u0026rdquo;, averaged over the last 24 tokens of each text, then over the 4 pairs (same forward-pass data as before). This score is exactly neuron $j$\u0026rsquo;s share of the MLP part from Section 2.5:\n$$u^\\top W\\, \\Delta \\bar a = \\sum_{j=1}^{3584} (u^\\top w_j)\\, \\Delta \\bar a_j = \\sum_{j=1}^{3584} \\text{score}_j$$A positive score means the neuron widens the gap along $u$: either $c_j \u003e 0$ and it fires more during loops, or $c_j \u003c 0$ and it fires less. The 3584 scores are sorted and the top 64 are kept, separately per layer. On the four scoring pairs these 64 account for about 50% (layer 15) and 56% (layer 19) of the total positive score.\nThe sign of $u$ matters here, even though it does not for the projection ($uu^\\top$ is sign-invariant). With a flipped $u$, the top scores would pick neurons that suppress looping. Before scoring, Astra oriented $u$ so that \u0026ldquo;repetitive − productive\u0026rdquo; projects positively onto it, using a separate set of contrast texts.\nThe edit becomes\n$$W' = W - \\alpha\\; u\\,\\big(s \\odot (u^\\top W)\\big)$$where $s \\in \\{0,1\\}^{3584}$ is 1 only for the 64 selected neurons and $\\odot$ is elementwise multiplication. At α = 4, editing only 64 columns performed about the same as editing all columns. Astra then raised α to 6; an exploratory test on 64 problems went from 7/64 to 35/64.\n3.5 Final edit $$W' = W - 6\\; u\\,\\big(s \\odot (u^\\top W)\\big), \\quad \\text{applied to layers 15 and 19}$$ Unselected columns are unchanged. Each selected column becomes $w'_j = w_j - 6\\,c_j\\,u$: its component along $u$ goes from $c_j$ to $-5\\,c_j$, and its orthogonal component is unchanged. 64 columns × 1024 rows × 2 layers = 131,072 changed weights, about 0.016% of the model. Relative Frobenius change $\\|W' - W\\|_F / \\|W\\|_F$: 10.6% (layer 15) and 8.1% (layer 19). The count of changed weights is tiny, but the change is concentrated: the selected columns are rewritten substantially. So the final edit is not strictly an ablation. It is a strong reversal applied to a small set of neurons, built on the directional-ablation framework.\n3.6 Frozen confirmation and results From this point no parameters were changed. On 1,152 practice problems never used before, accuracy went from 10.6% to 60.4%, with 579 problems rescued and 5 made worse. Correct edited answers still had long reasoning (median about 9,200 tokens), another sign the edit did not just shorten thinking.\nThe authors\u0026rsquo; evaluation on held-out benchmarks (thinking on, up to 32,768 tokens):\nQwen3.5-0.8B MATH-500 GSM8K IFEval Hit token limit Base 7.9% 13.8% 21.0% 93.0% Edited (50 forward passes) 48.1% 67.1% 62.3% 33.1% Two further observations from the paper:\nBudget matters. With only 10 or 25 forward passes, looping barely decreased and the gains were much smaller. The 50-pass run had room for the diagnosis, the attribution, and the controls above. Removing an obstacle vs. adding an ability. On Gemma 3 270M, which rarely loops, Astra could not meaningfully improve math scores. The Qwen gain came from removing a failure mode that was hiding ability the model already had; a weight edit like this cannot add ability that is not there. Sources Blog post: https://rikiyat.github.io/blog/weight-space-programming/ Edited model: https://huggingface.co/RikiyaT/Qwen3.5-0.8B-weight-space-programming Full logs: https://github.com/RikiyaT/lm_understand-logs, run exp8_math/runs/p50c/ readable/trajectory.md: all of Astra\u0026rsquo;s code and outputs work_notes/final_report.md, work_notes/experiment_history.md: summaries and the final formula requests/fwd_4.json.claimed: the 8 contrast texts edit/edit.npz, work_notes/frozen_weights_manifest.json: the bit-exact edit and the 64 selected neuron indices per layer Arditi et al. (2024), Refusal in Language Models Is Mediated by a Single Direction. ","permalink":"https://liqq010.github.io/posts/weight-space-programming-qwen-loops/","summary":"How an LLM edited 0.016% of Qwen3.5-0.8B\u0026rsquo;s weights to break its reasoning loops: directional ablation from the refusal-direction work, then pushed further with a strength α and 64 selected neurons per layer.","title":"Directional Ablation, and How Astra Fixed Qwen3.5-0.8B's Reasoning Loops"},{"content":"Paper: A Fix for Softmax and Cross-Entropy\nStandard Softmax and Cross-Entropy can suffer from numerical precision issues that affect training dynamics. This paper identifies two problems and proposes mathematically equivalent but numerically more stable formulations.\n1. Softmax Collapse Problem. For the target class $t$, the Cross-Entropy gradient is\n$$ \\frac{\\partial L}{\\partial z_t} = q_t - 1, \\qquad q = \\operatorname{softmax}(z) $$When the target logit dominates, $q_t$ may round to exactly 1 in floating-point arithmetic, producing a zero gradient even though the true gradient is nonzero. Meanwhile, non-target gradients remain positive, violating the zero-sum gradient property.\nFix. Rewrite Cross-Entropy using Softplus:\n$$ \\begin{aligned} L \u0026= -\\log\\frac{e^{z_t}}{\\sum_j e^{z_j}} \\\\ \u0026= \\log\\Big(1 + \\sum_{j\\ne t} e^{z_j - z_t}\\Big) \\\\ \u0026= \\operatorname{softplus}(a) \\end{aligned} $$where\n$$ a = \\operatorname{logsumexp}_{j\\ne t}(z_j - z_t) $$where $\\operatorname{softplus}(x) = \\log(1+e^x)$, $\\operatorname{logsumexp}_j(x_j) = \\log\\sum_j e^{x_j}$, and $\\sigma$ is the sigmoid.\nThe target gradient becomes\n$$ \\boxed{\\frac{\\partial L}{\\partial z_t} = -\\sigma(a) = q_t - 1} $$This avoids subtracting 1 from a probability extremely close to 1.\nMinimal example (FP32):\nimport torch z = torch.tensor([18.0, 0.0], dtype=torch.float32) q = torch.softmax(z, dim=-1) g_std = q[0] - 1 a = z[1] - z[0] g_fixed = -torch.sigmoid(a) print(g_std) # 0.0 print(g_fixed) # -1.52e-8 The two gradients are mathematically identical, but only the stable formulation preserves the small nonzero gradient.\n2. Softmax Cancellation Problem. In knowledge distillation, let\n$$ p = \\operatorname{softmax}(z_T), \\qquad q = \\operatorname{softmax}(z_S) $$The gradient with respect to student logits is\n$$ \\boxed{g_i = q_i - p_i} $$When student and teacher predictions are nearly identical, subtracting two close probabilities can cause catastrophic cancellation. Even identical logits may produce small nonzero gradients, leading to unnecessary student updates.\nFix. Compute the gradient from logit differences instead of subtracting probabilities.\n$$ g_i = q_i - p_i = p_i\\left(\\frac{q_i}{p_i} - 1\\right) $$Let $\\delta_i = z_{S,i} - z_{T,i}$. From the Softmax definition,\n$$ \\frac{q_i}{p_i} = \\frac{e^{\\delta_i}}{\\sum_j p_j e^{\\delta_j}} $$Then center it:\n$$ d_i = \\delta_i - \\sum_j p_j \\delta_j $$Since Softmax is invariant to a constant shift,\n$$ q_i = \\frac{p_i e^{d_i}}{\\sum_j p_j e^{d_j}} $$Let\n$$ s_i = d_i - \\operatorname{log1p}\\Big(\\sum_j p_j \\operatorname{expm1}(d_j)\\Big) $$where $\\operatorname{log1p}(x) = \\log(1+x)$ and $\\operatorname{expm1}(x) = e^x - 1$ are computed accurately even when $x$ is close to 0.\nThen\n$$ q_i = p_i e^{s_i} $$Therefore,\n$$ \\boxed{g_i = q_i - p_i = p_i \\operatorname{expm1}(s_i)} $$This is mathematically equivalent to the original gradient, but expm1 avoids subtracting nearly identical probabilities.\nTo reduce floating-point violations of the zero-sum gradient property, apply:\n$$ \\boxed{g_i = u_i - q_i \\sum_j u_j, \\qquad u_i = p_i \\operatorname{expm1}(s_i)} $$In exact arithmetic, $\\sum_j u_j = 0$, so the correction does not change the theoretical gradient.\nMinimal example (FP32):\nimport torch import torch.nn.functional as F t = torch.tensor([2., 1., 0.], dtype=torch.float32) s = t.clone().requires_grad_() p = F.softmax(t, dim=-1) loss = -(p * F.log_softmax(s, dim=-1)).sum() g_std = torch.autograd.grad(loss, s)[0] # Stable gradient d = s.detach() - t d = d - (p * d).sum() r = d - torch.log1p((p * torch.expm1(d)).sum()) # s_i in the text u = p * torch.expm1(r) q = F.softmax(s.detach(), dim=-1) g_fixed = u - q * u.sum() print(g_std) # May contain nonzero residuals print(g_fixed) # tensor([0., 0., 0.]) When teacher and student logits are identical, the stable formulation produces zero gradient without subtracting nearly identical Softmax probabilities.\nTakeaway Both fixes preserve the original mathematical objective and gradients:\nSoftmax Collapse: Replace $q_t - 1$ with $-\\sigma(a)$ to avoid precision loss near 1. Softmax Cancellation: Replace $q - p$ with $p \\operatorname{expm1}(s)$ to avoid catastrophic cancellation. Mathematical equivalence does not guarantee numerical equivalence in floating-point arithmetic.\n","permalink":"https://liqq010.github.io/posts/fixing-softmax-cross-entropy/","summary":"Softmax Collapse and Softmax Cancellation: two floating-point pitfalls in Softmax and Cross-Entropy, and mathematically equivalent fixes.","title":"Fixing Softmax and Cross-Entropy: Two Numerical Pitfalls"},{"content":"Introduction joint rotations (SMPL)camera poses (SLAM, 3D reconstruction)Them = 0catm = 1satm = 2onm = 3positional encoding (RoPE) Rotations appear everywhere: the joint rotations of parametric human body models, camera poses in SLAM and 3D reconstruction, and positional encoding in LLMs.\nThere are many ways to represent a 3D rotation: rotation matrices, axis-angle, rotation vectors, Euler angles, quaternions, and the 6D and 9D representations commonly used in neural networks.\nThe post is split into two parts. Part 1 focuses on what is most relevant to neural networks: rotation matrices, RoPE, and the 6D and 9D representations. Part 2 goes further into the other representations: axis-angle, rotation vectors, quaternions and Euler angles, and how they connect to each other.\nRotation Matrices Let $T$ be a rotation. It is a linear transformation.\nAny 3D vector can be written as a combination of the three coordinate axes:\n$$v = (x, y, z) = x\\,e_1 + y\\,e_2 + z\\,e_3$$$$e_1 = (1,0,0), \\quad e_2 = (0,1,0), \\quad e_3 = (0,0,1)$$By linearity,\n$$T(v) = x\\,T(e_1) + y\\,T(e_2) + z\\,T(e_3)$$So once we know where the three axes go, we know where every vector goes. Put the rotated axes into the columns of a matrix:\n$$R = \\begin{bmatrix} T(e_1) \u0026 T(e_2) \u0026 T(e_3) \\end{bmatrix}$$The rule for multiplying a matrix by a vector is exactly \u0026ldquo;use the vector\u0026rsquo;s coordinates as weights and add up the columns\u0026rdquo;:\n$$R\\,v = x\\,(\\text{column } 1) + y\\,(\\text{column } 2) + z\\,(\\text{column } 3) = T(v)$$So $R\\,v$ is the rotated vector.\nThe $j$-th column of a rotation matrix is where the $j$-th axis lands after the rotation. $R\\,v$ adds up the columns, weighted by the coordinates of $v$. 2D rotation matrix where the axes goe₁e₂(cos θ, sin θ)(−sin θ, cos θ)θθcos θ · v + sin θ · (v rotated by 90°)vv rotated by 90°cos θ · vsin θ ·(v rotated by 90°)v′θ Rotating counterclockwise by $\\theta$ sends $e_1$ to $(\\cos\\theta, \\sin\\theta)$ and $e_2$ to $(-\\sin\\theta, \\cos\\theta)$. Putting them in the columns:\n$$R(\\theta) = \\begin{bmatrix} \\cos\\theta \u0026 -\\sin\\theta \\\\ \\sin\\theta \u0026 \\cos\\theta \\end{bmatrix}$$Multiplying out $R(\\theta)\\,v$ for $v = (x, y)$:\n$$R(\\theta)\\begin{bmatrix} x \\\\ y \\end{bmatrix} = \\begin{bmatrix} x\\cos\\theta - y\\sin\\theta \\\\ x\\sin\\theta + y\\cos\\theta \\end{bmatrix} = \\cos\\theta \\begin{bmatrix} x \\\\ y \\end{bmatrix} + \\sin\\theta \\begin{bmatrix} -y \\\\ x \\end{bmatrix}$$$$v \\text{ rotated by } \\theta = \\cos\\theta \\cdot v + \\sin\\theta \\cdot (v \\text{ rotated by } 90^\\circ)$$The result is a weighted sum of two vectors: a $\\cos\\theta$ share of the original $v$, plus a $\\sin\\theta$ share of $v$ rotated by 90°, where $(x, y)$ rotated by 90° is $(-y, x)$.\nRotating the 2D plane with complex numbers Define $i^2 = -1$, and draw the complex number $a + b\\,i$ as the point $(a, b)$ in the plane:\nreal axis (a)imaginary axis (b)1−1i−i3 + 2i → point (3, 2) Multiplying by $i$ = rotating counterclockwise by 90°. Multiply any point $a + b\\,i$ by $i$:\n$$(a + b\\,i) \\times i = a\\,i + b\\,(i \\cdot i) = a\\,i - b = -b + a\\,i$$So the coordinates change as\n$$(a, b) \\xrightarrow{\\ \\times\\, i\\ } (-b, a)$$Rotating by any angle $\\theta$. To rotate by $\\theta$, multiply by the point on the unit circle at angle $\\theta$:\n$$\\text{the point at angle } \\theta \\text{ on the unit circle} = (\\cos\\theta, \\sin\\theta) = \\cos\\theta + i\\sin\\theta$$cos θ + i sin θlength = 1θ1i To see why, multiply any point $a + b\\,i$ by it:\n$$\\begin{aligned} (a + b\\,i)(\\cos\\theta + i\\sin\\theta) \u0026= a\\cos\\theta + a\\sin\\theta\\,i + b\\cos\\theta\\,i + b\\sin\\theta\\,(i \\cdot i) \\\\ \u0026= (a\\cos\\theta - b\\sin\\theta) + (a\\sin\\theta + b\\cos\\theta)\\,i \\end{aligned}$$The coordinates change as\n$$x' = a\\cos\\theta - b\\sin\\theta, \\qquad y' = a\\sin\\theta + b\\cos\\theta$$This is exactly the 2D rotation matrix:\n$$\\begin{bmatrix} x' \\\\ y' \\end{bmatrix} = \\begin{bmatrix} \\cos\\theta \u0026 -\\sin\\theta \\\\ \\sin\\theta \u0026 \\cos\\theta \\end{bmatrix} \\begin{bmatrix} a \\\\ b \\end{bmatrix}$$Conclusion: multiplying by $\\cos\\theta + i\\sin\\theta$ rotates a point counterclockwise by $\\theta$.\nRotation = a $\\cos$ share of \"stay\" + a $\\sin$ share of \"rotate by 90°\". In 2D, \"rotate by 90°\" is multiplying by $i$. Composing rotations = multiplying. Rotate a point $z$ by $\\theta_1$, then by $\\theta_2$:\n$$\\begin{aligned} \u0026z \\times (\\cos\\theta_1 + i\\sin\\theta_1) \\times (\\cos\\theta_2 + i\\sin\\theta_2) \\\\ =\\ \u0026z \\times \\big(\\cos(\\theta_1 + \\theta_2) + i\\sin(\\theta_1 + \\theta_2)\\big) \\end{aligned}$$The angles simply add. Expanding the product in the middle gives exactly the angle-addition formulas for $\\cos$ and $\\sin$.\nEuler\u0026rsquo;s formula.\n$$e^{i\\theta} = \\cos\\theta + i\\sin\\theta$$It simply means \u0026ldquo;the point at angle $\\theta$ on the unit circle\u0026rdquo;.\n\u0026ldquo;Adding the angles\u0026rdquo; is adding the exponents:\n$$e^{i\\theta_1} \\cdot e^{i\\theta_2} = e^{i(\\theta_1 + \\theta_2)}$$RoPE RoPE (rotary position embedding) rotates the query and the key by angles proportional to their positions. The core idea: when two rotated vectors are dotted together, the result naturally contains their relative position.\n$$(R(\\alpha)\\,q)^\\top (R(\\beta)\\,k) = q^\\top R(\\alpha)^\\top R(\\beta)\\,k$$Since $R(\\alpha)^\\top = R(-\\alpha)$ (transposing flips the sign of $\\sin$), and $R(a)\\,R(b) = R(a + b)$ (multiplying rotation matrices adds the angles),\n$$R(\\alpha)^\\top R(\\beta) = R(-\\alpha)\\,R(\\beta) = R(\\beta - \\alpha)$$$$(R(\\alpha)\\,q)^\\top (R(\\beta)\\,k) = q^\\top R(\\beta - \\alpha)\\,k$$The dot product of the two rotated vectors depends only on the angle difference $\\beta - \\alpha$, not on $\\alpha$ and $\\beta$ separately. With $\\alpha = m\\theta$ and $\\beta = n\\theta$ for a query at position $m$ and a key at position $n$, where $\\theta$ is the angle per position, it depends only on $n - m$.\nMatrix form. A pair $(q_0, q_1)$ is rotated by a 2D rotation matrix:\n$$\\begin{bmatrix} \\cos\\theta \u0026 -\\sin\\theta \\\\ \\sin\\theta \u0026 \\cos\\theta \\end{bmatrix} \\begin{bmatrix} q_0 \\\\ q_1 \\end{bmatrix} = \\begin{bmatrix} q_0 \\\\ q_1 \\end{bmatrix} \\cos\\theta + \\begin{bmatrix} -q_1 \\\\ q_0 \\end{bmatrix} \\sin\\theta$$For $q$ and $k$ of dimension $d_{model}$, split them into pairs and rotate each pair. Pair $j$ at position $m$ is rotated by $m\\,\\theta_j$, with $\\theta_j = 1 / 10000^{2j/d_{model}}$: each pair has its own angle per position.\nclass RoPE(nn.Module): def __init__(self, d_model, max_seq_len, theta): super().__init__() # 1 / theta ** (2*i / d_model) inv_freq = 1. / theta ** (torch.arange(0, d_model, 2) / d_model) # d_model / 2 inv_freq = torch.repeat_interleave(inv_freq, 2, dim=-1) # d_model pos = torch.arange(max_seq_len) # L angles = torch.einsum(\u0026#39;l,d-\u0026gt;ld\u0026#39;, pos, inv_freq) # l d_model self.register_buffer(\u0026#39;sin\u0026#39;, angles.sin(), persistent=False) self.register_buffer(\u0026#39;cos\u0026#39;, angles.cos(), persistent=False) def forward(self, x, position_ids): # B H T D sin = self.sin[position_ids][None, None] cos = self.cos[position_ids][None, None] rotate_half = torch.stack([-x[...,1::2], x[...,::2]], dim=-1).flatten(-2) out = cos * x + sin * rotate_half return out Complex form. Treat each pair $(q_0, q_1)$ as the complex number $q_0 + i\\,q_1$. Rotating it by $\\theta$ is multiplying by $\\cos\\theta + i\\sin\\theta = e^{i\\theta}$:\n$$(q_0 + i\\,q_1)(\\cos\\theta + i\\sin\\theta) = (q_0\\cos\\theta - q_1\\sin\\theta) + i\\,(q_0\\sin\\theta + q_1\\cos\\theta)$$the same result as the matrix form. So rotating all pairs is one element-wise complex multiplication.\nclass RoPE_complex(nn.Module): def __init__(self, d_model, max_seq_len, theta): super().__init__() # 1 / theta ** (2*i / d_model) inv_freq = 1. / theta ** (torch.arange(0, d_model, 2) / d_model) # d_model / 2 t = torch.arange(max_seq_len) # L freqs = torch.outer(t, inv_freq) # polar(r, θ) = r·cos(θ) + i·r·sin(θ) = r·e^{iθ} freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # e^{i ang} self.register_buffer(\u0026#39;freqs_cis\u0026#39;, freqs_cis, persistent=False) def forward(self, x, position_ids): B, H, T, d = x.shape fc = self.freqs_cis[position_ids] # [a, b] → a + bi xc = torch.view_as_complex(x.float().reshape(B, H, T, d//2, 2)) out = torch.view_as_real(xc * fc[None, None, :, :]).flatten(-2) return out.type_as(x) 3D rotation matrices Rotations about the $z$, $x$ and $y$ axes:\n$$R_z(\\theta) = \\begin{bmatrix} \\cos\\theta \u0026 -\\sin\\theta \u0026 0 \\\\ \\sin\\theta \u0026 \\cos\\theta \u0026 0 \\\\ 0 \u0026 0 \u0026 1 \\end{bmatrix}$$$$R_x(\\theta) = \\begin{bmatrix} 1 \u0026 0 \u0026 0 \\\\ 0 \u0026 \\cos\\theta \u0026 -\\sin\\theta \\\\ 0 \u0026 \\sin\\theta \u0026 \\cos\\theta \\end{bmatrix}$$$$R_y(\\theta) = \\begin{bmatrix} \\cos\\theta \u0026 0 \u0026 \\sin\\theta \\\\ 0 \u0026 1 \u0026 0 \\\\ -\\sin\\theta \u0026 0 \u0026 \\cos\\theta \\end{bmatrix}$$When rotating about a coordinate axis, that axis stays where it is. Since each column is where an axis lands, the corresponding column is the axis itself: $(0, 0, 1)$ in $R_z$, $(1, 0, 0)$ in $R_x$, $(0, 1, 0)$ in $R_y$. The other two axes rotate in their plane, just like in 2D.\nArbitrary 3D rotations A general 3D rotation is not about a coordinate axis. Take any three orthonormal vectors $u = (x_u, y_u, z_u)$, $v = (x_v, y_v, z_v)$, $w = (x_w, y_w, z_w)$ that form a right-handed set ($w = u \\times v$):\n$$u \\cdot u = v \\cdot v = w \\cdot w = 1, \\qquad u \\cdot v = v \\cdot w = w \\cdot u = 0$$Putting them in the rows gives a rotation:\n$$R_{uvw} = \\begin{bmatrix} x_u \u0026 y_u \u0026 z_u \\\\ x_v \u0026 y_v \u0026 z_v \\\\ x_w \u0026 y_w \u0026 z_w \\end{bmatrix}$$$R_{uvw}\\,u = (u \\cdot u,\\ v \\cdot u,\\ w \\cdot u) = (1, 0, 0)$, so $R_{uvw}$ rotates $u$ to the $x$ axis, $v$ to $y$, and $w$ to $z$. Its transpose has $u, v, w$ as columns and rotates the other way: it sends $x, y, z$ to $u, v, w$ respectively.\nA rotation matrix must be orthogonal with determinant $+1$:\n$$R^\\top R = I, \\qquad \\det R = +1$$Orthogonal: a rotation keeps lengths and angles, so the columns stay unit length and perpendicular to each other. Entry $(i, j)$ of $R^\\top R$ is column $i$ dotted with column $j$, so $R^\\top R = I$.\nThe determinant is the factor by which a transformation scales volume, and its sign tells whether the shape is flipped into its mirror image. A rotation keeps volume and does not flip, so $\\det R = +1$.\nPredicting rotations: 6D and 9D A rotation matrix changes continuously with the rotation, so it is a natural target for a network to predict. The problem is that the numbers a network outputs do not naturally satisfy $R^\\top R = I$ and $\\det R = +1$. We need a correction step that turns any output into a valid rotation matrix. The 6D and 9D representations are two ways to do this.\n6D. Output two columns $a_1, a_2$, orthonormalize them with Gram–Schmidt, and get the third column from the cross product:\n$$b_1 = \\frac{a_1}{\\|a_1\\|}, \\qquad b_2 = \\frac{a_2 - (b_1 \\cdot a_2)\\,b_1}{\\|a_2 - (b_1 \\cdot a_2)\\,b_1\\|}, \\qquad b_3 = b_1 \\times b_2$$$$R = \\begin{bmatrix} b_1 \u0026 b_2 \u0026 b_3 \\end{bmatrix}$$1. network outputa₁a₂not unit length, not perpendicular2. remove the part along b₁a₂(b₁·a₂) b₁b₁b₂subtract, then normalize3. b₃ = b₁ × b₂b₁b₂b₃perpendicular to both 9D. Output the whole $3 \\times 3$ matrix $M$, take its SVD $M = U \\Sigma V^\\top$, and use the nearest rotation:\n$$R = U V^\\top$$If $\\det(U V^\\top) = -1$, flip the sign of the last column of $U$ first, so that $\\det R = +1$.\nThe rotation axis Every 3D rotation matrix has a direction that does not move (Euler\u0026rsquo;s rotation theorem): vectors on the rotation axis $n$ are unchanged by the rotation, $R\\,n = n$. In linear-algebra terms, $n$ is an eigenvector of $R$ with eigenvalue 1. This leads to the next representation: axis-angle and rotation vectors, covered in Part 2. Rodrigues\u0026rsquo; rotation formula, quaternions and Euler angles will also be covered there.\nCitation If you find this post useful, please cite it as:\nLi, Quanquan. \u0026ldquo;Rotation Representations Part 1: Rotation Matrices, RoPE, and 6D/9D.\u0026rdquo; Quanquan Li\u0026rsquo;s blog, October 24, 2026. https://liqq010.github.io/posts/rotation-representations-part1/\nOr in BibTeX:\n@misc{li2026rotationpart1, author = {Li, Quanquan}, title = {Rotation Representations Part 1: Rotation Matrices, RoPE, and 6D/9D}, year = {2026}, month = {October}, day = {24}, howpublished = {\\url{https://liqq010.github.io/posts/rotation-representations-part1/}}, note = {Blog post} } ","permalink":"https://liqq010.github.io/posts/rotation-representations-part1/","summary":"Rotation matrices, complex numbers and RoPE, and the 6D and 9D rotation representations used in neural networks.","title":"Rotation Representations Part 1: Rotation Matrices, RoPE, and 6D/9D"},{"content":"1. A very odd loss curve The epochal sawtooth. Training loss of a 3M-parameter byte-level GPT, Adam, shuffled every epoch. The loss creeps up within each epoch and drops at the next. When a model is trained for many epochs with per-epoch shuffling, the training loss often rises slowly within each epoch and drops abruptly at every epoch boundary. We have seen curves like this before, and they are not a bug; others have reported similar curves, for example when fine-tuning LLMs, and attributed them to the model memorizing examples after seeing them only once [1]. Prior work [2] studied the pattern systematically and named it the Epochal Sawtooth Phenomenon (ESP). They observed it in BERT pretraining, found it strongest with Adam, and showed that it is sensitive to how the data is sampled, to Adam\u0026rsquo;s $\\beta_2$, and to model size.\nTheir explanation centers on the optimizer: Adam\u0026rsquo;s momentum and its per-coordinate normalization by $\\sqrt{v_t}$ make training on other batches slowly raise the loss of a batch not yet visited (Section 2). But it leaves the most striking ablations unexplained: why turning off shuffling weakens the sawtooth, and why sampling with replacement removes it.\nHere we take a different angle, and those unexplained ablations are the clue. Turning off shuffling, sampling with replacement, reversing the order every epoch and changing the batch size all leave the optimizer untouched; they change only the order in which samples are seen. If the order alone can make the sawtooth disappear or amplify it, the key variable should be a quantity determined by the order, not an optimizer statistic. Since neural networks memorize the specific samples they train on and gradually forget them afterwards, a natural candidate is how long ago each sample was last trained on. This leads to a more principled explanation that does not rely on the optimizer: the sawtooth comes from what the training loss measures. Each step\u0026rsquo;s loss is computed on samples of different \u0026ldquo;freshness\u0026rdquo;, and the average freshness, the expected gap since the last visit, follows a sawtooth over the epoch. We derive the sawtooth from the gap, predict the outcome of all four sampling ablations, and verify the mechanism with a small GPT.\n2. Background: the optimizer-based explanation Liu and Ma [2] trace the sawtooth to Adam. Two observations point there: the sawtooth is strongest with Adam and weaker with RMSProp, and it becomes more pronounced as $\\beta_2$ decreases. Recall that, with gradient $g_t$ and learning rate $\\gamma$, Adam updates\n$$ \\begin{aligned} m_t \u0026= \\beta_1 m_{t-1} + (1-\\beta_1)\\, g_t, \\\\ v_t \u0026= \\beta_2 v_{t-1} + (1-\\beta_2)\\, g_t^2, \\\\ \\theta_t \u0026= \\theta_{t-1} - \\gamma\\, \\frac{m_t}{\\sqrt{v_t} + \\epsilon}. \\end{aligned} $$Their argument has two parts, one for each side of the sawtooth.\nThe drop at the start of an epoch. After reshuffling, the last batches of one epoch overlap with the first batches of the next, so the momentum $m_t$ is large at the start of an epoch and the model takes large steps on its first few batches.\nThe rise during the epoch. Let $l_t^b$ be the loss of batch $b$ at step $t$ and $\\Delta\\theta_t = \\theta_t - \\theta_{t-1}$. Before $b$ is visited, to first order\n$$ l_t^b \\approx l_0 + \\sum_{\\tau=0}^{t-1} \\left\\langle \\Delta\\theta_\\tau,\\ \\nabla l_\\tau^b \\right\\rangle, $$where $l_0$ is the loss at the start of the epoch. The gradients of other batches point roughly along $\\nabla l^b$, but the per-coordinate division by $\\sqrt{v_t}$ rotates the update. For $\\beta_2$ not close to 1 the inner product tends to be positive, so the loss of $b$ slowly rises until $b$ is trained on.\nThis explains why a smaller $\\beta_2$ strengthens the sawtooth, but it has two limitations. First, it is not derived from first principles: it relies on empirical fits to how $\\lVert g_t\\rVert$, $\\lVert m_t\\rVert$ and $\\lVert v_t\\rVert$ evolve within an epoch. Second, it does not explain the effect of sampling: the authors acknowledge that their analysis explains neither the effect of shuffling nor that of sampling with replacement.\n3. A new perspective: the gap since last training Consider a dataset of $N$ samples, batch size $B$, and $K = N/B$ steps per epoch, and index the step within an epoch by $q \\in \\{0, 1, \\dots, K-1\\}$. For a sample in the current batch, define its gap $\\Delta$ as the number of steps since it was last trained on. The training loss logged at step $q$ is computed on the current batch before the update, so it is an average of per-sample losses, each at its own gap.\nForgetting curve. Let $\\ell(\\Delta)$ be the expected loss of a sample whose gap is $\\Delta$; we assume\n$$ \\ell(\\Delta) \\ \\text{is increasing in}\\ \\Delta. $$That is, the model retains something specific about each sample it trains on: training on a sample lowers its loss, and training on other samples afterwards partly undoes that. The same effect has been studied as example forgetting [3], and we measure $\\ell(\\Delta)$ directly in Section 5.1.\nHow the gap behaves within an epoch. With an independent shuffle every epoch and no replacement, the sample at step $q$ of this epoch sat at a position $p$ in the previous epoch that is uniform on $\\{0,\\dots,K-1\\}$ and independent of $q$. Its gap is the rest of the previous epoch plus the part of this epoch already done:\n$$ \\Delta = (K - p) + q . $$Averaging over $p$,\n$$ \\mathbb{E}[\\Delta \\mid q] = \\frac{K+1}{2} + q \\approx \\frac{K}{2} + q . $$Within an epoch, the expected gap grows linearly from about $K/2$ to about $3K/2$; at the next epoch it falls back to about $K/2$. Since $\\ell$ is increasing, the training loss rises through the epoch and drops at the boundary. Intuitively, where a sample sat in the previous epoch is random and does not depend on where it sits now. What changes is that every sample drawn later in the epoch has spent more steps away since it was last trained on.\nStrictly, the logged loss is the average of $\\ell(\\Delta)$ over the batch, not $\\ell$ of the average gap. This makes no difference here: given $q$, $\\Delta$ is uniform on $\\{q+1, \\dots, q+K\\}$, so each step shifts the whole distribution of gaps up by one, and the average loss rises with it.\nThis also accounts for the drop at the start of an epoch without any special overlap between consecutive epochs: at $q = 0$ the gaps are smallest on average, so the loss is lowest.\n4. The gap explains the sampling ablations We now check this perspective against the sampling ablations in [2]. For each scheme, we work out the gap $\\Delta$ at step $q$: if it changes with $q$, we predict a sawtooth, otherwise none. We compare each prediction with [2] and with our own runs.\nSetup. A byte-level GPT (4 layers, width 256, 3.3M parameters) is trained on 4,096 sequences of 128 bytes from the TinyStories validation split [4], used here as training data, for 40 epochs, with batch size 32 ($K = 128$) and Adam (learning rate $10^{-3}$, $\\beta=(0.9, 0.999)$), a small analogue of the BERT setting in [2]. The sawtooth is measured as the change from the first 10% to the last 10% of steps of each epoch, averaged over epochs 25–34 and 3 seeds. Code: liqq010/esp.\nThe two figures below are aligned column by column: the first shows the gap under each scheme, the second the training loss of the same model under that scheme. Under shuffling, the expected gap grows within each epoch, and the loss rises by 32% within each epoch.\nGap Δ under each scheme (K = 128). Line: mean; band: range of Δ (for sampling with replacement, whose geometric distribution is unbounded, the 10th–90th percentiles). Training loss under each scheme (epochs 30–34, shared y axis). The loss follows the gap. No shuffling. Every sample keeps the same position in every epoch, $p = q$, so $\\Delta = K$ for all samples at every step. The gap does not depend on $q$, and the sawtooth should disappear. Indeed, [2] finds it nearly gone; the weak residual can come from second-order effects such as samples at different positions having different difficulty. In our runs the loss falls by 6% within each epoch: no sawtooth, only the model improving.\nSampling with replacement. Each step draws $B$ samples independently from the whole dataset. The time since a sample was last drawn is geometric with mean approximately $K$, and its distribution is the same at every step. There is no epoch in any meaningful sense, and the sawtooth should disappear, as [2] observes. In our runs the loss falls by 4% within each epoch.\nReversing the order every epoch (order A, then reversed A, then A, …). The sample at step $q$ sat at position $K-1-q$ in the previous epoch, so\n$$ \\Delta = (K - (K-1-q)) + q = 1 + 2q . $$The gap now runs from 1 to about $2K$, twice the range under shuffling, and the first samples of an epoch were trained on one step ago. The sawtooth should therefore be stronger; how much stronger depends on the shape of $\\ell$. This matches the strong amplification reported in [2]. In our runs the loss rises by 64% within each epoch, twice the 32% under shuffling.\nLarger batch size. With a larger batch, an epoch has fewer steps $K$, so the range of gaps, measured in steps, shrinks. In addition, each step fits any individual sample less. Both effects weaken the sawtooth, as [2] reports. In our runs, batch sizes 16, 32, 64 and 128 give rises of 36%, 32%, 27% and 16%. The ordering also holds when runs are compared at the same training loss, so it is not just that larger batches train more slowly.\nAll four ablations follow directly from this one derivation, while the optimizer-based analysis of [2] leaves two of them, no shuffling and sampling with replacement, unexplained.\n5. Further experiments The sampling ablations test what the gap predicts. The experiments below test the mechanism itself, using the runs of Section 4 and, in Section 5.3, two more model widths. At every step we record, for each sample in the batch, its loss and its gap $\\Delta$, and every 8 steps we evaluate the model on a fixed set of 512 training samples.\n5.1 Measuring the forgetting curve Per-sample loss against the gap Δ (epochs 25–34). Left: shuffling, split by position in the epoch; the three curves nearly coincide. Right: shuffling vs. sampling with replacement; the curves have the same shape. Gaps below about 10 are too rare under shuffling to plot. The assumption of Section 3 can be checked directly. Binning every per-sample loss by its gap gives a smooth, increasing curve, from about 0.12 at $\\Delta \\approx 10$ to about 0.27 at $\\Delta \\approx 200$. Splitting by position in the epoch (left) tests whether the gap is really the variable that matters. Samples from the first, middle and last third of the epoch cover different ranges of $\\Delta$, but where they overlap they follow nearly the same curve, with residual offsets of about 0.01 against a total range of 0.15. The sawtooth is this curve read along a moving window of gaps.\n5.2 The average training loss keeps decreasing Batch loss vs. loss on a fixed set of 512 training samples (shuffling, epochs 30–34). Only the batch loss shows the sawtooth. If the sawtooth comes from which samples are measured, a measurement that does not change its samples should not show it. Evaluating the same model on a fixed set of 512 training samples gives a smooth curve that keeps decreasing within each epoch (−7% from the start to the end of an epoch, against +32% for the batch loss). Forgetting of individual samples is real, and it is what the batch loss picks up as the gaps change; on average, though, the model keeps improving within each epoch.\n5.3 Larger models forget more visibly Sawtooth size over training for three model widths (0.9M, 3.3M and 12.9M parameters; shaded: range over 3 seeds). It grows as the model starts to memorize, and it is larger for wider models. The sampling scheme sets the distribution of gaps; the forgetting curve sets how the loss responds to them, and so how large the sawtooth is. A model that retains more about individual samples has a steeper $\\ell(\\Delta)$ and a bigger sawtooth. Two observations support this. First, the sawtooth grows over training: in the first few epochs the loss falls within each epoch, because general learning dominates, and the rise appears only once the model starts fitting individual sequences. Second, it is larger for larger models: at epoch 15 the rise is +5%, +18% and +23% for widths 128, 256 and 512 (0.05, 0.11 and 0.12 in absolute loss), consistent with the model-size result in [2]. The ordering also holds in absolute terms when widths are compared at the same training loss, so it is not an artifact of measuring the rise relative to a lower baseline.\n6. Discussion What the sawtooth tells you. The sawtooth is not the model degrading within an epoch. It is a sign that the model fits individual training samples, and its amplitude is a rough measure of how much. That is not overfitting by itself; overfitting is a question about held-out data. But in small-data, many-epoch regimes, a sawtooth that suddenly grows is a good reason to look at the validation loss.\nBeyond the gap. The forgetting curve need not depend on the gap alone. In [1], the data is not reshuffled, so every gap equals $K$, yet the loss still rises late in an epoch: there a learning-rate schedule plays the role that the gap plays here, since samples last seen at a high learning rate are remembered better.\nGetting a smooth curve. Sampling with replacement removes the sawtooth. Larger batches or stronger data augmentation flatten the forgetting curve and shrink it. Logging the loss on a fixed subset of training data, as in Section 5.2, gives a smooth curve without changing training.\nWhere Adam comes in: conjectures. Within this picture, the optimizer affects the sawtooth through the forgetting curve $\\ell(\\Delta)$: its steepness sets the size, and its curvature can change the shape; for example, a curve that saturates at large gaps gives a concave rise, as [2] observes for small $\\beta_2$. Apart from a quick $\\beta_2$ sweep, we have not tested the following; they are plausible readings, not results.\nAdam vs. SGD. Adam divides each coordinate\u0026rsquo;s update by $\\sqrt{v_t}$. Coordinates whose gradients are usually small and are driven by only a few samples, such as embeddings of rare tokens or parameters that respond to specific inputs, get their updates amplified. Adam would then fit individual samples more strongly in a single step, giving a steeper forgetting curve. Smaller $\\beta_2$. $v_t$ averages over roughly $1/(1-\\beta_2)$ steps. With a smaller $\\beta_2$, $v_t$ follows the current $g_t^2$ more closely, so the normalization reacts faster to each batch; in the limit $\\beta_2 \\to 0$ the update becomes $m_t/(|g_t|+\\epsilon)$, which is sign descent only when $\\beta_1 \\approx 0$. Each step would then fit the current batch more aggressively and interfere more with other samples, so samples would be memorized and forgotten faster. In our runs, however, this effect is weak: $\\beta_2$ from 0.999 down to 0.8 gives rises of 32–35%, and 0.7 gives 29%. The strong $\\beta_2$ effect in [2] may depend on the setting. Larger weight decay. It limits how far parameters can move to fit an individual sample, which would weaken memorization slightly, matching the small effect seen in [2]. 7. Conclusion The epochal sawtooth comes from what the training loss measures. Each step\u0026rsquo;s loss is taken on samples whose time since the last visit averages $K/2 + q$, a quantity that itself follows a sawtooth over epochs. If a sample\u0026rsquo;s loss grows with that time, this predicts the sawtooth and the outcome of every sampling ablation, and a small GPT confirms each prediction. The sampling scheme determines the distribution of gaps, while the model and optimizer determine how the loss responds to them; together they set the sawtooth\u0026rsquo;s shape and size.\nCitation If you find this post useful, please cite it as:\nLi, Quanquan. \u0026ldquo;Why Training Loss Sawtooths Across Epochs.\u0026rdquo; Quanquan Li\u0026rsquo;s blog, October 15, 2026. https://liqq010.github.io/posts/epochal-sawtooth/\nOr in BibTeX:\n@misc{li2026epochalsawtooth, author = {Li, Quanquan}, title = {Why Training Loss Sawtooths Across Epochs}, year = {2026}, month = {October}, day = {15}, howpublished = {\\url{https://liqq010.github.io/posts/epochal-sawtooth/}}, note = {Blog post} } References Howard \u0026amp; Whitaker. Can LLMs learn from a single example?. fast.ai blog, 2023. Liu \u0026amp; Ma. The Epochal Sawtooth Phenomenon: Unveiling Training Loss Oscillations in Adam and Other Optimizers. Neural Processing Letters, 2025. Code: qiliuchn/training-loss-pattern. Toneva et al. An Empirical Study of Example Forgetting during Deep Neural Network Learning. ICLR 2019. Eldan \u0026amp; Li. TinyStories: How Small Can Language Models Be and Still Speak Coherent English?. 2023. ","permalink":"https://liqq010.github.io/posts/epochal-sawtooth/","summary":"Why does training loss creep up within each epoch and drop at the next? Not the optimizer, but what the loss measures: the expected gap since a sample was last trained on follows a sawtooth, and it explains every sampling ablation.","title":"Why Training Loss Sawtooths Across Epochs"},{"content":"TL;DR One-step generated samples, same initial noise for every method. CT has no guidance, others use classifier-free guidance with scale 3. What I cannot create, I do not understand.\n— Richard Feynman\nKnowing how a method works is not the same as understanding it. Understanding means being able to start from first principles and build it, or derive it, yourself. So I implemented diffusion and flow models from scratch, from the original 1000-step DDPM all the way down to 1-step generators, and put them in one repo: nano_diffusion, a clean, minimal PyTorch implementation of diffusion and flow models.\nWhat\u0026rsquo;s covered DDPM (ε / x₀ / v-prediction) with DDIM, DPM-Solver-2, DPM-Solver++(2M) and UniPC samplers; Rectified Flow; Consistency Models (CT and CD); MeanFlow; DMD2; and LoRA fine-tuning on top. One small DiT (9.79M parameters), one dataset (MNIST), one time convention shared by every few-step method, and one evaluation protocol (same seed, same initial noise). The core of each method is a single file of under 200 lines of PyTorch, and every model trains on a single GPU in 30 minutes.\nThe roadmap The post follows the number of sampling steps down, and each stop answers one question: why can we get away with fewer steps?\nnetwork calls per sample fewer steps → 1000 DDPM denoise a little at a time 5–50 Fast samplers better ODE solver, no retraining 5–50 Rectified Flow straight paths noise → data 1–2 Consistency jump straight to the trajectory end 1–2 MeanFlow learn the average velocity 1 DMD2 match teacher's distribution multi-step: solve an ODE at sampling time one/few-step: learn to take big jumps Setup used everywhere. DiT with 28×28 input, patch size 2, hidden size 256, depth 8, 16 heads (9.79M parameters). MNIST train + test (70k images) scaled to $[-1, 1]$. AdamW, batch size 256 (50 for the small LoRA dataset). All models are class-conditional; with classifier-free guidance drop the label 10% of the time during training.\nPart 1 — The 1000-step Baseline: DDPM Every method in this post is a variation on one idea from DDPM (Ho et al., 2020): destroying an image with noise is trivial, so train a network to undo it, a little at a time.\nAdding noise is easy The forward process gradually mixes an image $x_0$ with Gaussian noise $\\epsilon \\sim \\mathcal{N}(0, I)$ over $T = 1000$ steps. Each step adds a little noise with variance $\\beta_t$, and we can jump straight to any step $t$:\n$$ x_t = \\sqrt{\\bar\\alpha_t}\\, x_0 + \\sqrt{1 - \\bar\\alpha_t}\\, \\epsilon, \\qquad \\bar\\alpha_t = \\prod_{s=1}^{t} (1 - \\beta_s). $$$\\sqrt{\\bar\\alpha_t}$ is how much of the image survives, and $\\sqrt{1-\\bar\\alpha_t}$ is how much noise has been mixed in.\nSweeping $t$ from 0 to 1000 traces the full path from image to noise:\nThe forward process with the linear schedule, sampled every 200 steps. All panels use the same $\\epsilon$. def q_sample(self, x0, t, noise): abar_t = self.abar[t].view(-1, 1, 1, 1) return abar_t.sqrt() * x0 + (1 - abar_t).sqrt() * noise Linear vs cosine schedules The schedule $\\beta_1, \\dots, \\beta_T$ decides how fast the image disappears. The original DDPM uses a linear schedule, $\\beta_t$ from $10^{-4}$ to $0.02$. Improved DDPM (Nichol \u0026amp; Dhariwal, 2021) noticed that this destroys the image too early and proposed a cosine schedule, which defines $\\bar\\alpha_t$ directly as a squared cosine:\n$$ \\bar\\alpha_t = \\frac{f(t)}{f(0)}, \\qquad f(t) = \\cos^2\\!\\left(\\frac{t/T + s}{1 + s} \\cdot \\frac{\\pi}{2}\\right), \\quad s = 0.008. $$ Top: linear. Bottom: cosine. Same image, same noise. With the linear schedule the person is nearly gone by $t = 600$ ($\\sqrt{\\bar\\alpha} = 0.16$), so the last 40% of the steps are spent on what is essentially pure noise. With the cosine schedule the silhouette survives until about $t = 800$. Training: what should the network predict? Training is plain regression. Sample an image, a timestep and some noise, build $x_t$, and ask the network to predict something about it. The original choice is to predict noise $\\epsilon$:\ndef loss(self, model, x0, cond): t = torch.randint(0, self.T, (x0.shape[0],), device=x0.device) noise = torch.randn_like(x0) xt = self.q_sample(x0, t, noise) pred = model(xt, t, cond) if self.pred == \u0026#39;eps\u0026#39;: target = noise elif self.pred == \u0026#39;x0\u0026#39;: target = x0 else: # v = sqrt(abar) * eps - sqrt(1 - abar) * x0 abar_t = self.abar[t].view(-1, 1, 1, 1) target = abar_t.sqrt() * noise - (1 - abar_t).sqrt() * x0 return F.mse_loss(pred, target) But $\\epsilon$ is not the only option. Given $x_t$, the three common targets are interchangeable, since knowing any one of them pins down the other two:\nTarget Network predicts Recovering $x_0$ ε-prediction $\\epsilon$ $x_0 = \\big(x_t - \\sqrt{1-\\bar\\alpha_t}\\,\\hat\\epsilon\\big) / \\sqrt{\\bar\\alpha_t}$ x₀-prediction $x_0$ directly v-prediction (Salimans \u0026amp; Ho, 2022) $v = \\sqrt{\\bar\\alpha_t}\\,\\epsilon - \\sqrt{1-\\bar\\alpha_t}\\,x_0$ $x_0 = \\sqrt{\\bar\\alpha_t}\\,x_t - \\sqrt{1-\\bar\\alpha_t}\\,\\hat v$ They are equivalent in theory. In practice they differ in how the network\u0026rsquo;s error is scaled when it is turned into an image, and that difference will matter a lot once we start taking big steps.\nI trained all three on MNIST with the same settings.\nTraining loss for x₀-, v- and ε-prediction, log scale. Each panel has its own y-axis because the values are not comparable: each curve measures error on a different target. Samples during training. After the first epoch, ε-prediction still produces almost pure noise, while x₀- and v-prediction already show blurry outlines of digits. Sampling: 1000 small steps back To generate, start from pure noise $x_T$ and walk back one step at a time. At each step the network gives an estimate $\\hat x_0$, and the true posterior $q(x_{t-1} \\mid x_t, x_0)$ is a Gaussian with closed-form mean and variance. So we plug in $\\hat x_0$ and sample:\n$$ x_{t-1} = \\underbrace{\\frac{\\sqrt{\\bar\\alpha_{t-1}}\\,\\beta_t}{1-\\bar\\alpha_t}}_{\\text{coef}_1}\\,\\hat x_0 + \\underbrace{\\frac{\\sqrt{1-\\beta_t}\\,(1-\\bar\\alpha_{t-1})}{1-\\bar\\alpha_t}}_{\\text{coef}_2}\\,x_t + \\sigma_t z, \\qquad \\sigma_t^2 = \\frac{1-\\bar\\alpha_{t-1}}{1-\\bar\\alpha_t}\\,\\beta_t. $$@torch.no_grad() def p_sample(self, model, xt, t, cond, cfg_scale=1.0): # any target (eps / x0 / v) -\u0026gt; x0, clipped to [-1, 1] x0, _ = self._pred(model, xt, t, cond, cfg_scale) if t[0] == 0: return x0 coef1 = self._extract(self.posterior_coef1, t, xt.shape) coef2 = self._extract(self.posterior_coef2, t, xt.shape) mean = coef1 * x0 + coef2 * xt var = self._extract(self.posterior_var, t, xt.shape) return mean + var.sqrt() * torch.randn_like(xt) Because every target is converted to $\\hat x_0$ first (and clipped to the valid pixel range), one sampler serves all three models. Here they are with the full 1000 steps:\n1000-step DDPM sampling, guidance scale 3. All three work, but they are not identical. ε-prediction gives the cleanest strokes, while x₀-prediction leaves speckles around the digits and v-prediction sits in between. The real problem with DDPM is the cost: 1000 network calls per sample. Everything that follows is about cutting that number down.\nPart 2 — 1000 → 5 Steps Without Retraining: Fast Samplers Nothing in the DDPM training objective says you have to sample with 1000 steps. The network has learned to estimate $x_0$ from $x_t$ at any noise level. How we walk from noise to data is a separate decision, and we can change it after training. This part keeps the three models from Part 1 exactly as they are and only swaps the sampler.\nDDIM: denoise, then re-noise to a lower level DDIM (Song et al., 2020) makes the observation explicit. At time $t$, use the network to get $\\hat x_0$ and the noise it implies, $\\hat\\epsilon$. Then jump straight to any earlier time $s \u003c t$ by putting them back together with the coefficients of time $s$:\n$$ x_s = \\sqrt{\\bar\\alpha_s}\\,\\hat x_0 + \\sqrt{1-\\bar\\alpha_s}\\,\\hat\\epsilon, \\qquad \\hat\\epsilon = \\frac{x_t - \\sqrt{\\bar\\alpha_t}\\,\\hat x_0}{\\sqrt{1-\\bar\\alpha_t}}. $$In words: denoise all the way, then add back the same noise, but less of it. No fresh randomness is injected, so sampling becomes a deterministic map from the initial noise to an image. That is why every figure in this post can use the same initial noise and be compared cell by cell. (DDIM also has a parameter $\\eta$ that adds noise back in; $\\eta = 1$ recovers DDPM-like sampling. I use $\\eta = 0$ throughout.) Under the hood this is the Euler method applied to the probability-flow ODE, the deterministic ODE whose solutions have the same marginals as the diffusion process.\nOne real DDIM step, 5-step trailing grid. The network turns $x_t$ into $\\hat x_0$ and the noise it implies; the step then recombines them with the coefficients of the lower noise level $s$. The whole step is two lines, and it turns out to be the only building block the fancier samplers need:\n@staticmethod def _ddim_step(x, x0, abar_t, abar_s): \u0026#34;\u0026#34;\u0026#34;Given x0, step deterministically from t to s (eta = 0).\u0026#34;\u0026#34;\u0026#34; eps = (x - abar_t.sqrt() * x0) / (1 - abar_t).sqrt() return abar_s.sqrt() * x0 + (1 - abar_s).sqrt() * eps Higher-order solvers are \u0026ldquo;a better x₀, then a DDIM step\u0026rdquo; DDIM assumes $\\hat x_0$ stays constant over the whole step from $t$ to $s$. That is a first-order approximation, and it is accurate only when the steps are small. DPM-Solver (Lu et al., 2022) showed that if you measure time in half log-SNR,\n$$ \\lambda_t = \\log\\frac{\\sqrt{\\bar\\alpha_t}}{\\sqrt{1-\\bar\\alpha_t}}, $$the probability-flow ODE has an exact solution in which the only unknown is an integral of $\\hat x_0$ (or $\\hat\\epsilon$) over $\\lambda$. Better solvers approximate that integral better. Implemented this way, each one reduces to \u0026ldquo;choose an effective $\\hat x_0$ for this step, then call _ddim_step\u0026rdquo;:\nSampler Network calls per step The $\\hat x_0$ it plugs into the DDIM step DDIM 1 the current estimate $\\hat x_0^{(i)}$ DPM-Solver++(2M) (Lu et al., 2022b) 1 extrapolated from the current and previous estimates UniPC (Zhao et al., 2023) 1 same predictor, plus a corrector that reuses the next step\u0026rsquo;s evaluation DPM-Solver-2 (Lu et al., 2022) 2 the estimate at the midpoint (in $\\lambda$) of the step DPM-Solver++(2M) keeps the previous step\u0026rsquo;s $\\hat x_0$ and extrapolates linearly in $\\lambda$, which costs nothing extra:\nfor i, t in enumerate(ts): x0, _ = self._pred(model, x, t, cond, cfg_scale) if x0_prev is None or i == steps - 1: # first or last step: 1st order x0_eff = x0 else: # r = previous step size / this step size, in lambda r = (lams[i] - lams[i - 1]) / (lams[i + 1] - lams[i]) x0_eff = x0 + (x0 - x0_prev) / (2 * r) x = self._ddim_step(x, x0_eff, abars[i], abars[i + 1]) x0_prev = x0 UniPC uses the same predictor, then notices that the network has to be evaluated at the new point anyway, for the next step. So it uses that evaluation immediately as a corrector: it redoes the current step with the average of the $\\hat x_0$ at both ends (a trapezoidal rule), and hands the same evaluation on to the next step. The correction is free. My version is a simplified, second-order UniC.\nDPM-Solver-2 is the classic single-step midpoint method: take a trial step to the midpoint in $\\lambda$, evaluate the network there, then redo the full step with the midpoint estimate. It is more accurate per step but costs two network calls per step.\nResults Comparison of samplers on the ε-prediction model, guidance scale 3, same initial noise. Rows: sampler. Columns: number of steps (= network calls). Going from 1000 to 5 network calls, a 200× speed-up with no retraining, costs almost nothing visible here. The honest observation is that on MNIST the three samplers are nearly indistinguishable even at 5 steps. Digits are simple, the guided trajectories are fairly straight, and first-order DDIM is already close to converged. The advantage of the higher-order solvers reported in the papers shows up on harder data and at very low step counts. Where things do change on MNIST is below 5 steps, and that is where both pitfalls of this part live. DPM-Solver-2 is not in that figure because it spends two calls per step, so comparing it by step count would be unfair.\nPitfall: 2 steps ≈ 1 step\nMy first DDIM implementation used the obvious timestep grid, torch.linspace(999, 0, steps). With steps = 2 that is $t = [999, 0]$. The first call jumps from pure noise to $t = 0$, where $\\bar\\alpha_0 = 0.9999$, so the image is already finished. The second call, at $t = 0$, has nothing left to do. Half of the budget is wasted, and \u0026ldquo;2 steps\u0026rdquo; gives the same image as 1 step.\nThe fix is the trailing grid from Lin et al., 2023: split $[0, 999]$ into steps equal intervals, call the network only at the noisy end of each, and let the last step land directly on the clean image.\n# steps=2 -\u0026gt; [999, 499] ts = torch.linspace(self.T - 1, 0, steps + 1)[:-1].long() v-prediction, guidance scale 3, DDIM, same initial noise. Each pair compares the two grids: the linspace grid with 2 steps produces the same images as the trailing grid with 1 step, and linspace with 3 steps matches trailing with 2. The final call at $t = 0$ does nothing. A side effect worth knowing: with the trailing grid, DPM-Solver++(2M) and UniPC are exactly identical to DDIM at 1 and 2 steps. The first step has no previous estimate to extrapolate from, and the last step falls back to first order, so with two steps there is no step left where the higher-order correction applies.\nPitfall: ε-prediction cannot take a single step\nAt 1 step, the three models from Part 1 behave completely differently:\n1-step DDIM (one network call at $t = 999$), guidance scale 3. ε-prediction produces noise with a faint digit behind it, while x₀- and v-prediction give blurry but recognizable digits. A likely reason is the conversion to an image:\n$$ \\hat x_0 = \\frac{x_t - \\sqrt{1-\\bar\\alpha_t}\\,\\hat\\epsilon}{\\sqrt{\\bar\\alpha_t}}, \\qquad \\sqrt{\\bar\\alpha_{999}} \\approx 0.0064, $$so at $t = 999$ any error in $\\hat\\epsilon$ is amplified about 157×. With 1000 small steps this never matters, and x₀- and v-prediction avoid the division altogether.\nThe blur of x₀- and v-prediction is not a bug either. From pure noise, the input says nothing about which particular \u0026ldquo;3\u0026rdquo; to draw, and the MSE-optimal answer is the average of all 3s, which is a blur. One network call trained with a regression loss can only give you a mean.\nThis is the real limit of training-free samplers. They make the ODE solve cheaper, but they cannot change what the network was trained to predict. To get a sharp sample in one step, the model has to be trained for it, which is what Part 4 is about. Before that, Part 3 changes the training itself to make the paths from noise to data straighter and easier to integrate.\nPart 3 — Straight Lines: Rectified Flow Part 2 changed only the sampler. This part changes the training: the forward process, the target and the time variable. The result, Rectified Flow (Liu et al., 2022; Lipman et al., 2022), is simpler than DDPM and needs fewer steps. It is also the coordinate system that every method in Part 4 is built on.\nA simpler forward process Drop the noise schedule entirely. Let time run continuously from $t = 0$ (data) to $t = 1$ (noise) and interpolate in a straight line:\n$$ x_t = (1 - t)\\,x_0 + t\\,\\epsilon, \\qquad t \\in [0, 1]. $$Along this line the velocity is constant, and it is what the network learns to predict:\n$$ v = \\frac{dx_t}{dt} = \\epsilon - x_0. $$ $\\epsilon$ $-$ $x_0$ $=$ $v$ The velocity target. The displayed $v$ is min-max rescaled to the pixel range. This is still a Gaussian noising process, just with coefficients $(1 - t, t)$ instead of $(\\sqrt{\\bar\\alpha_t}, \\sqrt{1 - \\bar\\alpha_t})$. Given $x_t$ and $\\hat v$, both other quantities are one line away:\n$$ \\hat x_0 = x_t - t\\,\\hat v, \\qquad \\hat\\epsilon = x_t + (1 - t)\\,\\hat v. $$Keep the first one in mind: $\\hat x_0 = x_t - t\\,\\hat v$ appears in every method of Part 4.\nTraining in five lines def loss(self, model, x0, cond): t = self.sample_t(x0.shape[0], x0.device) x1 = torch.randn_like(x0) # noise tb = t.view(-1, 1, 1, 1) xt = (1 - tb) * x0 + tb * x1 # t is scaled to [0, 1000] for the time embedding v_pred = model(xt, t * self.t_scale, cond) return F.mse_loss(v_pred, x1 - x0) Sampling: just integrate the velocity Generation solves the ODE $dx/dt = v(x, t)$ backwards, from $t = 1$ to $t = 0$. With Euler\u0026rsquo;s method this is a three-line loop:\nx = torch.randn(shape) ts = torch.linspace(1.0, 0.0, steps + 1) for t, t_next in zip(ts[:-1], ts[1:]): x = x + (t_next - t) * v(x, t) # dt \u0026lt; 0: move against v, towards the data Heun\u0026rsquo;s method (Karras et al., 2022) adds a correction: take an Euler step, evaluate the velocity at the end point, and redo the step with the average of the two velocities. That costs 2 network calls per step (1 for the last step), like DPM-Solver-2.\nEuler (top) vs Heun (bottom), guidance scale 3, same initial noise, columns roughly matched by network calls. Heun takes two calls per step except the last. As with the DDPM solvers, the second-order method buys nothing visible here. Euler with 5 steps is already about as good as anything with 50.\nClassifier-free guidance All models so far were trained with the class label dropped 10% of the time (replaced by a \u0026ldquo;null\u0026rdquo; label, index 10). The network therefore learns both a conditional velocity $v(x, t \\mid y)$ and an unconditional one $v(x, t \\mid \\varnothing)$. Classifier-free guidance (Ho \u0026amp; Salimans, 2022) extrapolates away from the unconditional prediction at sampling time:\n$$ \\tilde v = v(x, t \\mid \\varnothing) + w\\,\\big(v(x, t \\mid y) - v(x, t \\mid \\varnothing)\\big), $$with $w = 1$ meaning no guidance. Both branches go through the network as a single batch of twice the size:\nout = model(torch.cat([x, x]), torch.cat([t, t]), torch.cat([cond, null_cond])) v_c, v_u = out.chunk(2) v = v_u + cfg_scale * (v_c - v_u) Top: no guidance ($w = 1$). Bottom: $w = 3$. Euler sampler, same initial noise. Without guidance the samples are more varied but less reliable: thin strokes, a 2 that looks like an R, an 8 that falls apart. With $w = 3$ the digits are thicker, cleaner and more uniform. That is the usual trade of diversity for fidelity. Guidance also doubles the cost of every step, a cost the distillation methods in Part 4 will remove by baking the guided teacher into the student.\nDDPM vs Rectified Flow at few steps The real payoff shows up when the step budget is tiny. Here are DDPM (with v-prediction, which handles big steps far better than ε-prediction) and Rectified Flow side by side at 1–5 network calls:\nGuidance scale 3, same initial noise. Top: DDPM, v-prediction, DDIM. Bottom: Rectified Flow, Euler. The two models also differ in how they sample $t$ during training, so this is not a perfectly controlled comparison. At 5 steps the two are indistinguishable. At 2 and 3 steps Rectified Flow is visibly ahead: its digits are already sharp, while DDPM still shows ghosting and soft edges.\nThe bigger lesson of Parts 2 and 3 is where the limit lies. Better solvers and straighter paths get us from 1000 calls to about 3. To get to one, the model has to be trained to take one big step, which is what Part 4 does.\nPart 4 — The Final Leap: 1-step Generators Parts 2 and 3 ended at the same wall. A network trained by regression to predict $x_0$ (or $v$) from pure noise can only output the average image, because the noise says nothing about which sample to draw. To get a sharp sample from one network call, the training objective itself has to change. The three families in this part change it in three different ways:\ndatanoise Consistency Models Every point on a trajectory should map to the same endpoint, so one call from noise lands on the image.\nzₜzᵣvu MeanFlow Learn the average velocity over an interval instead of the instantaneous one, so one step is exact.\ngeneratorteacher DMD2 Forget trajectories: push the one-step generator's output distribution towards the teacher's.\nConsistency Models: jump straight to the end of the trajectory The idea. The sampling ODE from Part 3 defines a trajectory from every noise sample to an image. Every point $x_t$ on one trajectory should therefore map to the same endpoint $x_0$. A consistency model (Song et al., 2023) learns that map, $f(x_t, t) \\mapsto x_0$, directly. If it succeeds, one call at $t = 1$ produces an image.\nThe parameterization. The map must satisfy $f(x, 0) = x$, since a clean image is its own endpoint. The original paper enforces this with EDM\u0026rsquo;s $c_\\text{skip}$ and $c_\\text{out}$ coefficients. In Rectified Flow coordinates it comes for free: keep the network predicting a velocity and define\n$$ f_\\theta(x_t, t) = x_t - t\\, v_\\theta(x_t, t), $$which equals $x$ exactly at $t = 0$. It is the same $\\hat x_0 = x_t - t\\,\\hat v$ from Part 3, so a Rectified Flow model is already a (bad) consistency model, and a good initialization for one.\nThe loss. Split $[0, 1]$ into $N$ intervals. Take a point $x_t$ and its neighbour $x_{t'}$ one interval closer to the data on the same trajectory. Ask that both map to the same place, with the target computed by the current model under stop-gradient:\n$$ \\mathcal{L} = d\\big(f_\\theta(x_t, t),\\ \\operatorname{sg}[f_\\theta(x_{t'}, t')]\\big). $$At $t' = 0$ the target is the real image, and consistency propagates upward from there. The only question is how to get $x_{t'}$, and the two answers give the two variants.\nConsistency Training (CT): no teacher CT builds $x_{t'}$ from the same $(x_0, \\epsilon)$ pair as $x_t$, which is one Euler step along the conditional velocity $\\epsilon - x_0$. No pretrained model is needed. I followed the recipe of improved CT (Song \u0026amp; Dhariwal, 2023): the target is the current model under stop-gradient (not an EMA copy), the distance is a pseudo-Huber loss, the loss is weighted by $1/\\Delta t$, and $N$ doubles from 10 to 1280 over training so that early steps are large and stable and later ones are small and accurate.\ndef loss_ct(self, model, x0, cond, it): N = self.num_steps(it) # 10 -\u0026gt; 1280, doubling B = x0.shape[0] t = torch.sigmoid(torch.randn(B) * 2.0 - 1.1) # iCT lognormal i = torch.ceil(t * N).clamp(1, N) t, tp = i / N, (i - 1) / N # snap to the grid # (t and tp are broadcast over pixels below) x1 = torch.randn_like(x0) xt = (1 - t) * x0 + t * x1 # same (x0, x1) pair xtp = (1 - tp) * x0 + tp * x1 pred = self.f(model, xt, t, cond) with torch.no_grad(): target = self.f(model, xtp, tp, cond) return (N * self._pseudo_huber(pred, target)).mean() Consistency Distillation (CD): use the Rectified Flow model as teacher CD changes one thing: $x_{t'}$ comes from the Part 3 model taking a single Heun step from $x_t$, so the student learns the teacher\u0026rsquo;s actual trajectories. The teacher runs with guidance $w = 3$, so the guidance gets baked into the student, which then needs no CFG (and no second forward pass) at sampling time. The student starts as a copy of the teacher, and $N = 20$ is fixed.\nwith torch.no_grad(): # one Heun step of the guided teacher dt = tp - t v1 = self._teacher_v(teacher, xt, t, cond, cfg_scale) v2 = self._teacher_v(teacher, xt + dt * v1, tp, cond, cfg_scale) xtp = xt + dt * 0.5 * (v1 + v2) target = self.f(model, xtp, tp, cond) 1-step samples during training. Top: CT from scratch, 100 epochs. Bottom: CD from the Rectified Flow teacher, 50 epochs. The difference is stark. CD produces sharp, guided digits after the very first epoch: it starts from the teacher and only has to learn to skip to the end of trajectories the teacher already knows. CT has to discover the trajectories from data alone. It improves steadily but even after 100 epochs its digits are thinner and rougher, with broken strokes. Part of the gap is guidance, which CT cannot have. CFG mixes velocities, which are derivatives, but a consistency model outputs the endpoint of a trajectory, an integral along the path. The endpoint of the guided ODE is not a mix of the conditional and unconditional endpoints, and getting it requires evaluating the guided velocity along the way, which is exactly what the teacher does in CD.\nMore steps. A consistency model can also sample in several steps: predict $\\hat x_0$, add fresh noise back to a smaller $t$, and predict again. Unlike the ODE samplers, this injects new randomness at every step.\n@torch.no_grad() def sample(self, model, shape, cond, ts=(1.0,)): # (1.0, 0.5) for 2 steps x = torch.randn(shape) for k, t in enumerate(ts): x0 = self.f(model, x, torch.full((shape[0],), t), cond) if k + 1 \u0026lt; len(ts): # re-noise with fresh noise to the next t x = (1 - ts[k + 1]) * x0 + ts[k + 1] * torch.randn_like(x0) return x0 Pitfall: CT collapses when $N$ doubles\nMy first CT run sampled $t$ with Rectified Flow\u0026rsquo;s logit-normal(0, 1). Around epoch 40, right after $N$ switched from 40 to 80, the 1-step samples turned into pure noise. The run recovered over the next ten epochs, and the collapse reproduced in repeated runs.\n1-step CT samples around the collapse. Top: logit-normal(0, 1). Bottom: the lognormal distribution from iCT. Everything else is identical. The likely cause is how rarely the logit-normal visits high noise: only 0.16% of its samples land at $t \\ge 0.95$. That region is exactly where 1-step generation starts, and it sits at the far end of the consistency chain that propagates up from the data. When $N$ doubles, every interval is suddenly half as long and the targets there shift, with almost no training signal to re-anchor them. iCT samples $\\log\\sigma \\sim \\mathcal{N}(-1.1, 2^2)$ with $\\sigma = t/(1-t)$, a much wider distribution that puts 2.2% of the samples at $t \\ge 0.95$, about 14× more. With that change, the run did not collapse.\nLesson: the time distribution decides which parts of the trajectory the model ever learns. The logit-normal that is ideal for Rectified Flow is a poor fit for CT.\nMeanFlow: learn the average velocity The idea. A single Euler step from $t$ to $r$ with Rectified Flow\u0026rsquo;s velocity is wrong because the velocity changes along the path. But there is some velocity that makes one step exact, the average velocity over the interval:\n$$ u(z_t, r, t) = \\frac{1}{t - r} \\int_r^t v(z_\\tau, \\tau)\\, d\\tau, \\qquad z_r = z_t - (t - r)\\, u(z_t, r, t). $$MeanFlow (Geng et al., 2025) trains a network to predict $u$. With $t = 1$ and $r = 0$, one call goes from noise to image, and nothing about it is approximate if $u$ is learned well. When $r = t$, $u$ is just the instantaneous velocity, so ordinary flow matching is a special case.\nThe trick. The integral cannot be computed during training. Multiplying the definition by $(t - r)$ and differentiating with respect to $t$ gives an identity that involves only quantities at time $t$:\n$$ u(z_t, r, t) = v(z_t, t) - (t - r)\\,\\frac{d}{dt} u(z_t, r, t), \\qquad \\frac{d}{dt}u = v\\,\\partial_z u + \\partial_t u. $$So the training target is $u_\\text{tgt} = v - (t - r)\\,\\operatorname{sg}\\big[\\tfrac{d}{dt}u\\big]$, with $v = \\epsilon - x_0$ as in Rectified Flow. The total derivative $\\tfrac{d}{dt}u$ is a Jacobian-vector product of the network with tangent $(v, 0, 1)$ in $(z, r, t)$, and PyTorch computes it in one call:\nwith torch.no_grad(), sdpa_kernel(SDPBackend.MATH): _, dudt = torch.func.jvp( lambda z, r, t: self.u(model, z, r, t, cond), (z, r, t), (v, torch.zeros_like(r), torch.ones_like(t)), ) u_tgt = v - (t - r).view(-1, 1, 1, 1) * dudt u = self.u(model, z, r, t, cond) # normal forward; gradients flow here err = ((u - u_tgt) ** 2).flatten(1).sum(1) The DiT gets one extra input: an embedding of the interval length $h = t - r$, added to the time embedding. Following the paper, $(t, r)$ come from two logit-normal(−0.4, 1) samples, and 75% of the batch uses $r = t$ (plain flow matching).\nMeanFlow can also bake in guidance during training, by replacing the target velocity with a guided one built from the model\u0026rsquo;s own unconditional prediction; ω is that guidance scale.\n1-step MeanFlow samples during training, from scratch, 100 epochs. Top: no guidance ($\\omega = 1$). Bottom: $\\omega = 3$. From scratch and without a teacher, MeanFlow reaches sharp 1-step samples and ends up much cleaner than CT. Training-time guidance has the same effect as CFG had on Rectified Flow: $\\omega = 3$ gives thicker, more uniform, more reliable digits, while $\\omega = 1$ is more varied and makes more mistakes.\nPitfall: t × 1000 silently breaks MeanFlow\nMy first MeanFlow runs reused the Rectified Flow code path, which multiplies $t$ by 1000 before the sinusoidal time embedding (Part 3). The MSE kept rising, and the 1-step samples were blurry and often the wrong class.\nThe cause is the derivative in the target. Rectified Flow never differentiates the network with respect to time, so the scale of $t$ only affects which frequencies the embedding uses. MeanFlow\u0026rsquo;s target contains $\\partial u/\\partial t$, and scaling $t$ by 1000 scales that derivative by up to 1000 through the high-frequency components of the embedding. I measured $\\partial u / \\partial h$ at 14–43, while $|v|$ is about 1.4. The target was dominated by the derivative term, and the model chased it instead of the data.\nThe fix is one line: feed $t$ and $t - r$ in $[0, 1]$, as the official implementation does. Training then converged normally.\nPitfall: the JVP costs over 4× the memory\nTwo separate problems made the first working version need about 53GB at batch size 256:\nPyTorch\u0026rsquo;s default fused attention kernels do not support forward-mode autodiff, so the JVP must run with sdpa_kernel(SDPBackend.MATH). Computing the JVP inside autograd also stores the graph of the tangent computation for the backward pass. But the target is stop-gradient anyway, so none of that is needed. Running the JVP under torch.no_grad() and getting $u$ from a separate ordinary forward pass (with the fast kernel) brings memory down to about 12GB. DMD2: match distributions, not trajectories The idea. CD and MeanFlow still try to reproduce trajectories. Distribution Matching Distillation (Yin et al., 2023; DMD2, Yin et al., 2024) does not care which noise maps to which image. It only asks that the generator\u0026rsquo;s outputs, as a distribution, look like the teacher\u0026rsquo;s. Three networks take part, all of them copies of the Rectified Flow model at the start:\nThe teacher (the \u0026ldquo;real score\u0026rdquo;), frozen, run with guidance $w = 3$. Given a noised sample, it says where the data distribution would denoise it to. The generator $G$, a one-step model: $x = z - v_G(z, t{=}1)$, which is one Euler step from pure noise. At initialization this is exactly the blurry 1-step Rectified Flow sample. The fake score, trained online with flow matching on the generator\u0026rsquo;s outputs. It says where the generator\u0026rsquo;s distribution would denoise the same sample to. noise z Generator G one step x = G(z) real x Teacher frozen · CFG w = 3 Fake score tracks G's outputs Disc. head on fake-score features noised to t noised features DM loss x̂₀ fake − x̂₀ real FM loss flow matching on x GAN loss (D) real vs. generated GAN loss (G) fool the discriminator x̂₀ real x̂₀ fake updates G updates fake score + head frozen The DMD2 setup. Each loss is colored by the parameters it updates: orange losses train the generator, blue losses train the fake score and its discriminator head. The teacher is frozen. The distribution matching gradient. Take a generated $x$, noise it to a random $t \\in [0.02, 0.98]$, and ask both scores for $\\hat x_0$. The difference $\\hat x_0^\\text{fake} - \\hat x_0^\\text{real}$ points from \u0026ldquo;what the generator makes\u0026rdquo; towards \u0026ldquo;what real data looks like\u0026rdquo; around that sample; it is the gradient of $\\mathrm{KL}(p_\\text{fake} \\,\\|\\, p_\\text{real})$ with respect to $x$. Implementing it is a two-line trick: a regression loss whose target is the sample moved one step against the gradient.\n# distribution matching (DM) loss for the generator; x = G(z), with gradients with torch.no_grad(): xt = self._noise(x, t) x0_real = self._x0(teacher, xt, t, cond, self.cfg_scale) x0_fake = self._x0(fake, xt, t, cond) w = (x - x0_real).abs().mean(dim=(1, 2, 3), keepdim=True) grad = (x0_fake - x0_real) / w # normalized per sample loss_dm = 0.5 * F.mse_loss(x, (x - grad).detach()) The DM loss updates only the generator $G$. Both scores are evaluated under no_grad, so the gradient reaches $G$\u0026rsquo;s weights only through $x = G(z)$; the teacher stays frozen, and the fake score is not touched by this loss.\nThe fake score is kept up to date with plain flow matching (FM) on the generator\u0026rsquo;s outputs, exactly the Rectified Flow loss from Part 3 with $x$ in place of real data:\n# flow matching (FM) loss for the fake score; x is detached def loss_fake(self, fake, x, cond): x = x.detach() t = torch.rand(x.shape[0], device=x.device) x1 = torch.randn_like(x) xt = (1 - t.view(-1, 1, 1, 1)) * x + t.view(-1, 1, 1, 1) * x1 return F.mse_loss(fake(xt, t * self.t_scale, cond), x1 - x) The FM loss updates only the fake score. $x$ is detached, so no gradient flows back into $G$. The two losses alternate, with 5 fake-score updates per generator update (below).\nThe GAN losses. The DM loss only pulls the generator towards the teacher, so the generator can never be better than the teacher\u0026rsquo;s view of the data. DMD2 adds a GAN term that compares the generator\u0026rsquo;s outputs with real images directly. It is a diffusion GAN: both real and generated images are first noised to a random $t$, and the discriminator works on the noised versions. The discriminator reuses the fake score network as its backbone: a forward hook grabs the token features of its last DiT block, and a small convolutional head turns them into one logit per image. The DiT itself is unchanged.\ndef _logits(self, fake, head, x, cond): t = torch.rand(x.shape[0], device=x.device) # the forward hook stores the last-block features in self._feat fake(self._noise(x, t), t * self.t_scale, cond) return head(self._feat) # discriminator loss: real images -\u0026gt; real, generator outputs -\u0026gt; fake def loss_gan_d(self, fake, head, x_real, x_fake, c_real, c_fake): l_real = self._logits(fake, head, x_real, c_real) l_fake = self._logits(fake, head, x_fake.detach(), c_fake) return (F.softplus(-l_real) + F.softplus(l_fake)).mean() # generator loss (non-saturating): make the discriminator call G(z) real def loss_gan_g(self, fake, head, x, cond): return F.softplus(-self._logits(fake, head, x, cond)).mean() The discriminator loss updates the fake score network and the head together, with the same optimizer as the FM loss, so the fake network\u0026rsquo;s features learn to tell real from generated. The generated images are detached, so $G$ is not affected. The generator GAN loss updates only $G$: its gradient passes through the head and the fake network to reach $x = G(z)$, and the gradients it leaves on the fake network and the head are cleared before their own update.\nTwo more details from DMD2 complete the recipe:\nTwo time-scale updates (TTUR): the fake score and head take 5 updates for every generator update. The DM gradient is only meaningful if the fake score actually models the generator\u0026rsquo;s current outputs, so it has to keep up with a moving target. No regression loss: the original DMD also regressed the generator onto precomputed teacher samples (noise–image pairs from the full ODE), which is expensive to build and ties the generator to the teacher. DMD2 drops it, and so does this implementation. DMD2 generator, 1 step, during training (20 epochs). At epoch 0 the generator still shows the blur it inherited from the 1-step teacher. By epoch 5 the digits are sharp, thick and guided.\nHow do you know it\u0026rsquo;s working? DMD2\u0026rsquo;s losses are not useful progress bars: the DM loss is a moving target, and the fake score and the discriminator chase a generator that keeps changing. What matters is whether the fake score is good enough for the DM gradient to be trusted, and whether the GAN is balanced.\nI measure the fake score on a fixed probe batch (fixed $z$, labels, $t$ and noise), so the numbers are comparable across training. The figure shows metrics I monitor during training.\nDMD2 diagnostics over 20 epochs. Dashed lines mark reference values: 1 for the lag and the fake/teacher ratio, 0.5 for the discriminator output at equilibrium. fake score lag is the fake score\u0026rsquo;s denoising error on the generator\u0026rsquo;s samples right after a generator update, divided by the same error after its 5 catch-up updates. It asks: does the fake score keep up? Here it starts slightly above 1, while the generator is still moving fast away from its blurry starting point, and settles at 1 within the first epoch. If it stays clearly above 1, raise the TTUR ratio or lower the generator\u0026rsquo;s learning rate. A lag of 1 alone is ambiguous, since a fake score that is not learning at all also gives 1, so read it together with the next signal.\nfake/teacher is the fake score\u0026rsquo;s denoising error on the generator\u0026rsquo;s samples, divided by the teacher\u0026rsquo;s. It asks: does the fake score know the generator\u0026rsquo;s distribution better than the teacher does? Only then is $\\hat x_0^\\text{fake} - \\hat x_0^\\text{real}$ a real signal. Here it stays below 1 throughout. It dips early, when the generator\u0026rsquo;s blurry outputs are unlike real data, and then rises slowly as the generator approaches the data, which is convergence: a perfect generator would give exactly 1. A ratio above 1 means the DM gradient is unreliable: raise the TTUR ratio or check the fake score\u0026rsquo;s learning rate.\np_real and p_fake are the discriminator\u0026rsquo;s average probabilities for real and generated images. They ask: is the GAN balanced? Here the discriminator is ahead but far from saturated. If p_real → 1 and p_fake → 0, the discriminator has won: lower the discriminator loss weight. If both stay near 0.5, it is too weak to matter: raise the discriminator loss weight. The GAN losses in the last panel tell the same story; large swings between them mean the game is unstable.\nThe showdown Here is every method with one network call and the same initial noise, the figure from the top of the post:\nOne network call. DDPM (ε, DDIM) and Rectified Flow were not trained for one step. CT has no guidance; CD and DMD2 distill a guided teacher ($w = 3$); MeanFlow uses training-time guidance ($\\omega = 3$). And with two calls, for the methods that support multi-step sampling:\nTwo network calls. Rectified Flow uses Euler with CFG (4 forward passes in total); CT and CD re-noise between steps; MeanFlow takes two average-velocity steps. Teacher Calls Guidance Training (this repo) 1-step quality DDPM (ε), DDIM — 1 CFG (2 passes) 50 epochs, ~25 min noise Rectified Flow, Euler — 1 CFG (2 passes) 50 epochs, ~25 min blurry average CT none 1–2 none 100 epochs, ~70 min recognizable, rough strokes CD RF, $w = 3$ 1–2 distilled 50 epochs, ~65 min + teacher sharp MeanFlow, $\\omega = 3$ none 1–2 training-time 100 epochs, ~2 h sharp DMD2 RF, $w = 3$ 1 distilled 20 epochs, ~40 min + teacher sharpest, thickest strokes A few things stand out:\nA teacher makes everything easier. CD and DMD2 are sharp within an epoch or two. They inherit both the trajectories and the guidance from a model that already works, and only learn to compress it. Without a teacher, MeanFlow is clearly ahead of CT. Both train from scratch, but MeanFlow\u0026rsquo;s target is exact for any interval, while CT has to propagate consistency up a chain one interval at a time and is sensitive to how it samples $t$. MeanFlow pays for this with a JVP in every step and the longest training time here. More steps help less than you would expect. CT, CD and MeanFlow all look about the same at 1 and 2 steps. Once a model is trained for one step, the second step is mostly a refinement. The cost moves from sampling to training. Every 1-step method here takes longer to train than the 1000-step DDPM. That is the trade: pay once at training time, and every sample afterwards costs a single network call. Part 5 — LoRA Fine-tuning LoRA (Hu et al., 2021) fine-tunes a large model cheaply by freezing its weights and learning only a small low-rank update next to each layer.\nFine-tuning on outlined MNIST I made a synthetic dataset: 100 training images per class, 1000 in total, where each digit is replaced by its outline. The outline is a 3×3 dilation of the digit minus the digit itself, so only a one-pixel ring around each stroke survives.\nThe fine-tuning data. Each pair is an original digit (left) and its outline (right). The base model is the Rectified Flow model from Part 3, trained on normal MNIST. It has never seen an outlined digit. The goal is to make it draw outlines while keeping what it knows: class conditioning, guidance, and the shapes of the digits.\nLoRA in 30 lines Concretely, each frozen weight $W$ gets a low-rank update next to it:\n$$ y = W x + s \\cdot \\frac{\\alpha}{r}\\, B A x, \\qquad A \\in \\mathbb{R}^{r \\times d_\\text{in}},\\ B \\in \\mathbb{R}^{d_\\text{out} \\times r}. $$$B$ starts at zero, so the model is unchanged at the start of training. $s$ is an extra knob that is not trained: it scales the update at inference time.\nclass LoRALinear(nn.Module): def __init__(self, base, r, alpha): super().__init__() self.base, self.factor, self.scale = base, alpha / r, 1.0 self.A = nn.Parameter(torch.empty(r, base.in_features)) self.B = nn.Parameter(torch.zeros(base.out_features, r)) nn.init.kaiming_uniform_(self.A, a=math.sqrt(5)) def forward(self, x): lora = x @ self.A.T @ self.B.T return self.base(x) + self.scale * self.factor * lora apply_lora freezes the model and swaps in LoRALinear for the attention (qkv, proj) and MLP (fc1, fc2) layers of every DiT block. The adaLN modulation layers, which inject the time and class embeddings, are left alone, as is common practice. Only $A$ and $B$ are trained and saved: the adapter for $r = 4$ is about 0.5MB, against 39MB for the full model.\nResults The four fine-tuned models; each title gives the trainable parameters and their share of the model. 50 Euler steps, guidance scale 3, same initial noise everywhere. Every variant learned the style. Even $r = 1$, with a third of a percent of the parameters, draws outlined digits of the right class. More capacity gives cleaner results: $r = 1$ produces doubled, wobbly contours, while $r = 16$ is hard to tell apart from full fine-tuning.\nTraining loss; faint lines are the raw loss, bold lines a moving average. Higher rank fits faster and lower; full fine-tuning is lowest. Samples during fine-tuning. Lower rank adapts more slowly: at epoch 10, full fine-tuning and $r = 16$ already draw outlines, $r = 4$ is halfway, and $r = 1$ still mostly draws filled digits. Lower rank also fits less well, with messier outlines even at epoch 40. Turning the adapter up and down Because the update is scaled by $s$, the same adapter can be applied partially at inference time, without retraining:\nLoRA $r = 4$ applied at strength $s = 0$, $0.5$ and $1$. Guidance scale 3, same initial noise. $s = 0.5$ is a genuine in-between: thin single strokes, neither filled nor outlined. This also gives a free correctness check. At $s = 0$ the output must be the base model\u0026rsquo;s, and it is identical to the base model\u0026rsquo;s samples pixel for pixel.\nLimitations Only MNIST. MNIST is an easy dataset, so the differences between methods, especially at few steps, are probably smaller here than they would be on harder data, at higher resolution, in latent space or with text conditioning. Simplified implementations. UniPC uses a simplified second-order corrector and the consistency models are written in Rectified Flow coordinates rather than EDM\u0026rsquo;s. Try It Yourself Writing the code took a fraction of the time. Most of it went into understanding why something that looked correct did not work, and that turned out to be where the real learning happened. This is what Feynman\u0026rsquo;s line means in practice. Reading a paper gives you the idea. Building it, and finding out exactly why your first version fails, is what gives you the understanding.\nThe code is at github.com/liqq010/nano_diffusion. Everything runs on a single GPU.\nCitation If you find this post useful, please cite it as:\nLi, Quanquan. \u0026ldquo;From 1000 Steps to 1: Building Diffusion and Flow Models from Scratch.\u0026rdquo; Quanquan Li\u0026rsquo;s blog, October 9, 2026. https://liqq010.github.io/posts/nano-diffusion/\nOr in BibTeX:\n@misc{li2026nanodiffusion, author = {Li, Quanquan}, title = {From 1000 Steps to 1: Building Diffusion and Flow Models from Scratch}, year = {2026}, month = {October}, day = {9}, howpublished = {\\url{https://liqq010.github.io/posts/nano-diffusion/}}, note = {Blog post} } References DDPM: Ho, Jain \u0026amp; Abbeel. Denoising Diffusion Probabilistic Models. 2020. Cosine schedule: Nichol \u0026amp; Dhariwal. Improved Denoising Diffusion Probabilistic Models. 2021. v-prediction: Salimans \u0026amp; Ho. Progressive Distillation for Fast Sampling of Diffusion Models. 2022. DDIM: Song, Meng \u0026amp; Ermon. Denoising Diffusion Implicit Models. 2020. DPM-Solver: Lu et al. DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps. 2022. DPM-Solver++: Lu et al. DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models. 2022. UniPC: Zhao et al. UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models. 2023. Trailing grid: Lin et al. Common Diffusion Noise Schedules and Sample Steps are Flawed. 2023. Heun sampler: Karras et al. Elucidating the Design Space of Diffusion-Based Generative Models. 2022. Classifier-free guidance: Ho \u0026amp; Salimans. Classifier-Free Diffusion Guidance. 2022. Rectified Flow: Liu, Gong \u0026amp; Liu. Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. 2022. Flow matching: Lipman et al. Flow Matching for Generative Modeling. 2022. Logit-normal sampling, timestep shift: Esser et al. Scaling Rectified Flow Transformers for High-Resolution Image Synthesis. 2024. Consistency Models: Song et al. Consistency Models. 2023. Improved CT: Song \u0026amp; Dhariwal. Improved Techniques for Training Consistency Models. 2023. MeanFlow: Geng et al. Mean Flows for One-step Generative Modeling. 2025. DMD: Yin et al. One-step Diffusion with Distribution Matching Distillation. 2023. DMD2: Yin et al. Improved Distribution Matching Distillation for Fast Image Synthesis. 2024. DiT: Peebles \u0026amp; Xie. Scalable Diffusion Models with Transformers. 2022. LoRA: Hu et al. LoRA: Low-Rank Adaptation of Large Language Models. 2021. ","permalink":"https://liqq010.github.io/posts/nano-diffusion/","summary":"Clean, minimal PyTorch implementations of diffusion and flow models, from 1000-step DDPM to 1-step generators, on one backbone with one set of conventions.","title":"From 1000 Steps to 1: Building Diffusion and Flow Models from Scratch"},{"content":" I work on generative models. This blog is where I write up things I build to understand them.\n2022 – 2026 Viggle AI Generative Models 2021 – 2022 University of Waterloo Master of Electrical and Computer Engineering 2015 – 2021 SenseTime Computer Vision 2010 – 2015 The Chinese University of Hong Kong Bachelor of Information Engineering ","permalink":"https://liqq010.github.io/about/","summary":"\u003c!-- TODO: 姓名、职位/机构、研究方向，由作者补充 --\u003e\n\u003cp\u003eI work on generative models. This blog is where I write up things I build to understand them.\u003c/p\u003e\n\u003col class=\"timeline\"\u003e\n  \u003cli\u003e\n    \u003ctime\u003e2022 – 2026\u003c/time\u003e\n    \u003cdiv class=\"tl-name\"\u003e\u003ca href=\"https://viggle.ai/\"\u003eViggle AI\u003c/a\u003e\u003c/div\u003e\n    \u003cdiv class=\"tl-desc\"\u003eGenerative Models\u003c/div\u003e\n  \u003c/li\u003e\n  \u003cli\u003e\n    \u003ctime\u003e2021 – 2022\u003c/time\u003e\n    \u003cdiv class=\"tl-name\"\u003eUniversity of Waterloo\u003c/div\u003e\n    \u003cdiv class=\"tl-desc\"\u003eMaster of Electrical and Computer Engineering\u003c/div\u003e\n  \u003c/li\u003e\n  \u003cli\u003e\n    \u003ctime\u003e2015 – 2021\u003c/time\u003e\n    \u003cdiv class=\"tl-name\"\u003eSenseTime\u003c/div\u003e\n    \u003cdiv class=\"tl-desc\"\u003eComputer Vision\u003c/div\u003e\n  \u003c/li\u003e\n  \u003cli\u003e\n    \u003ctime\u003e2010 – 2015\u003c/time\u003e\n    \u003cdiv class=\"tl-name\"\u003eThe Chinese University of Hong Kong\u003c/div\u003e\n    \u003cdiv class=\"tl-desc\"\u003eBachelor of Information Engineering\u003c/div\u003e\n  \u003c/li\u003e\n\u003c/ol\u003e","title":"About"}]