<?xml version="1.0" encoding="UTF-8"?>
<rss  xmlns:atom="http://www.w3.org/2005/Atom" 
      xmlns:media="http://search.yahoo.com/mrss/" 
      xmlns:content="http://purl.org/rss/1.0/modules/content/" 
      xmlns:dc="http://purl.org/dc/elements/1.1/" 
      version="2.0">
<channel>
<title>Vaibhav Nakrani</title>
<link>https://nakranivaibhav.github.io/</link>
<atom:link href="https://nakranivaibhav.github.io/index.xml" rel="self" type="application/rss+xml"/>
<description>ML and Agentic AI, explained without the hand-waving. Notes on diffusion, agents, computer vision, and whatever I&#39;m currently figuring out.</description>
<generator>quarto-1.8.25</generator>
<lastBuildDate>Fri, 31 Jul 2026 18:30:00 GMT</lastBuildDate>
<item>
  <title>einsum, one index at a time</title>
  <dc:creator>Vaibhav Nakrani</dc:creator>
  <link>https://nakranivaibhav.github.io/posts/einsum-one-index-at-a-time.html</link>
  <description><![CDATA[ 




<p>Every PyTorch codebase I read eventually hits a line like this:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb1" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb1-1">attn <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> torch.einsum(<span class="st" style="color: #20794D;
background-color: null;
font-style: inherit;">'bhqd,bhkd-&gt;bhqk'</span>, q, k)</span></code></pre></div></div>
<p>I could always decode these strings mechanically. Match the letters, remember that repeated letters get summed, move on. But decoding is not the same as reading. I wanted <code>'ij,ja-&gt;ia'</code> to be a sentence I understand, not an incantation I pattern-match.</p>
<p>It turned out the whole system rests on one rule, and everything else is that rule applied to progressively bigger objects. This post builds it up in that order: one vector, two vectors, a matrix, a batch, three tensors, and finally the attention string above. Prerequisites: vectors, dot products, and basic matrix multiplication. Nothing else.</p>
<section id="the-tempting-model-a-lookup-table" class="level2">
<h2 class="anchored" data-anchor-id="the-tempting-model-a-lookup-table">The tempting model: a lookup table</h2>
<p>The way most people hold einsum in their head is as a table of memorized entries:</p>
<ul>
<li><code>'ij,jk-&gt;ik'</code> is matrix multiply</li>
<li><code>'ij-&gt;ji'</code> is transpose</li>
<li><code>'ii-&gt;'</code> is trace</li>
<li><code>'i,i-&gt;'</code> is dot product</li>
</ul>
<p>The table has no entry for <code>'bhqd,bhkd-&gt;bhqk'</code>, and adding more entries doesn’t help, because the strings are not arbitrary codes. Each one is a compressed equation, and the compression scheme is a physics convention from 1916 called Einstein summation. Learn the scheme once and every string, including ones you have never seen, becomes legible.</p>
</section>
<section id="step-1-a-dot-product-and-one-rule" class="level2">
<h2 class="anchored" data-anchor-id="step-1-a-dot-product-and-one-rule">Step 1: a dot product, and one rule</h2>
<p>Start with the smallest thing that involves a sum. The dot product of two vectors:</p>
<p><img src="https://latex.codecogs.com/png.latex?%5Cmathbf%7Bv%7D%5Ccdot%5Cmathbf%7Bw%7D%20=%20v_1%20w_1%20+%20v_2%20w_2%20+%20v_3%20w_3%20=%20%5Csum_%7Bj=1%7D%5E%7B3%7D%20v_j%5C,%20w_j"></p>
<p>The index <img src="https://latex.codecogs.com/png.latex?j"> is doing the walking: it visits every slot, multiplies the paired entries, and the sigma adds them up.</p>
<p>Einstein’s convention is a single observation about equations like this: whenever an index appears <strong>exactly twice</strong> in a term, it is always being summed over, so the sigma carries no information.</p>
<p><img src="https://latex.codecogs.com/png.latex?%5Cmathbf%7Bv%7D%5Ccdot%5Cmathbf%7Bw%7D%20=%20v_j%5C,%20w_j"></p>
<p>That is the entire convention. An index written twice means “sum over it, silently.” Everything that follows is this rule meeting larger objects.</p>
</section>
<section id="step-2-dummy-indices-and-free-indices" class="level2">
<h2 class="anchored" data-anchor-id="step-2-dummy-indices-and-free-indices">Step 2: dummy indices and free indices</h2>
<p>The dropped sigma creates two kinds of index, and telling them apart is the entire skill. Both have exact counterparts in code you already write, so let’s build the distinction there.</p>
<p>A <strong>dummy index</strong> is a loop variable inside a sum. Here is <img src="https://latex.codecogs.com/png.latex?v_j%20w_j"> as the loop it abbreviates:</p>
<div id="ad0acbe4" class="cell" data-execution_count="3">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb2" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb2-1"><span class="im" style="color: #00769E;
background-color: null;
font-style: inherit;">import</span> numpy <span class="im" style="color: #00769E;
background-color: null;
font-style: inherit;">as</span> np</span>
<span id="cb2-2"></span>
<span id="cb2-3">v <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.array([<span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">1.0</span>, <span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">2.0</span>, <span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">3.0</span>])</span>
<span id="cb2-4">w <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.array([<span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">4.0</span>, <span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">5.0</span>, <span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">6.0</span>])</span>
<span id="cb2-5"></span>
<span id="cb2-6">total <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> <span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">0.0</span></span>
<span id="cb2-7"><span class="cf" style="color: #003B4F;
background-color: null;
font-weight: bold;
font-style: inherit;">for</span> j <span class="kw" style="color: #003B4F;
background-color: null;
font-weight: bold;
font-style: inherit;">in</span> <span class="bu" style="color: null;
background-color: null;
font-style: inherit;">range</span>(<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>):          <span class="co" style="color: #5E5E5E;
background-color: null;
font-style: inherit;"># j lives only inside this loop</span></span>
<span id="cb2-8">    total <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">+=</span> v[j] <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">*</span> w[j]</span>
<span id="cb2-9"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(total)                 <span class="co" style="color: #5E5E5E;
background-color: null;
font-style: inherit;"># the answer has no j in it</span></span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>32.0</code></pre>
</div>
</div>
<p>Everything you know about the variable <code>j</code> transfers to the index <img src="https://latex.codecogs.com/png.latex?j">. It exists only inside the computation. It visits every value. And once the loop finishes, it is gone: the answer is <code>32.0</code>, not something indexed by <img src="https://latex.codecogs.com/png.latex?j">. Its name is irrelevant too. Rename <code>j</code> to <code>k</code> and it is the same computation, which is why <img src="https://latex.codecogs.com/png.latex?v_j%20w_j"> and <img src="https://latex.codecogs.com/png.latex?v_k%20w_k"> denote the same number. An index that appears twice is a loop variable, and spotting one should trigger the same reflex: <em>this will be summed away and won’t exist in the output.</em></p>
<p>A <strong>free index</strong> is the opposite: it survives into the output and labels its slots. In code, it is the variable that indexes the result you are building. Take matrix-vector multiplication, <img src="https://latex.codecogs.com/png.latex?u_i%20=%20A_%7Bij%7D%20v_j">:</p>
<div id="854359d4" class="cell" data-execution_count="4">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb4" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb4-1">A_small <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.arange(<span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">9.0</span>).reshape(<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>, <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>)</span>
<span id="cb4-2"></span>
<span id="cb4-3">u <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> [<span class="bu" style="color: null;
background-color: null;
font-style: inherit;">sum</span>(A_small[i, j] <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">*</span> v[j] <span class="cf" style="color: #003B4F;
background-color: null;
font-weight: bold;
font-style: inherit;">for</span> j <span class="kw" style="color: #003B4F;
background-color: null;
font-weight: bold;
font-style: inherit;">in</span> <span class="bu" style="color: null;
background-color: null;
font-style: inherit;">range</span>(<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>))   <span class="co" style="color: #5E5E5E;
background-color: null;
font-style: inherit;"># j: loop variable, consumed</span></span>
<span id="cb4-4">     <span class="cf" style="color: #003B4F;
background-color: null;
font-weight: bold;
font-style: inherit;">for</span> i <span class="kw" style="color: #003B4F;
background-color: null;
font-weight: bold;
font-style: inherit;">in</span> <span class="bu" style="color: null;
background-color: null;
font-style: inherit;">range</span>(<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>)]                             <span class="co" style="color: #5E5E5E;
background-color: null;
font-style: inherit;"># i: labels the output slot</span></span>
<span id="cb4-5"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(u)</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>[np.float64(8.0), np.float64(26.0), np.float64(44.0)]</code></pre>
</div>
</div>
<p>Two nested roles, visually separated. The inner <code>j</code> is a dummy again, consumed by <code>sum</code>. The outer <code>i</code> is free: the output is a list <em>indexed by</em> <code>i</code>, three numbers, one per value it takes. A free index appears exactly once in each term of the equation, and the number of free indices tells you the output’s number of dimensions before you compute anything. Zero free indices, scalar. One, vector. Two, matrix.</p>
<p>One consistency law binds every valid equation: <strong>the free indices on the left must exactly match the free indices on the right.</strong> In loop terms: both sides must be built by the same outer <code>for</code> loops. A dummy on the right is fine, it is consumed by its own <code>sum</code>. A free index on only one side is like writing <code>u[i] = ...</code> where the right side loops over <code>m</code> instead: the program is malformed before it runs. This law catches most index mistakes before any arithmetic happens.</p>
<p>The analogy has one limit worth naming. A literal loop nest fixes an execution order, but the einsum equation does not: the sums can be regrouped and reordered freely, which is exactly the freedom optimizers exploit later. The loop picture is for <em>reading</em> the equation, not for how the machine runs it.</p>
<p>From here on, whenever an equation appears, I’ll mention the roles the same way: which letter is the loop variable, which letters shape the output.</p>
</section>
<section id="step-3-the-strings-are-the-equations" class="level2">
<h2 class="anchored" data-anchor-id="step-3-the-strings-are-the-equations">Step 3: the strings are the equations</h2>
<p><code>np.einsum</code> takes the equation you would write on paper and asks you to type it, letters only:</p>
<ul>
<li>input tensors are described by their indices, comma-separated</li>
<li>the arrow shows the free indices that survive into the output</li>
<li>any index missing after the arrow is a dummy: summed</li>
</ul>
<p>The dot product <img src="https://latex.codecogs.com/png.latex?v_j%20w_j"> over two vectors <img src="https://latex.codecogs.com/png.latex?v"> and <img src="https://latex.codecogs.com/png.latex?w"> has no free index, so nothing follows the arrow:</p>
<div id="beb20292" class="cell" data-execution_count="5">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb6" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb6-1"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(np.einsum(<span class="st" style="color: #20794D;
background-color: null;
font-style: inherit;">'j,j-&gt;'</span>, v, w))</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>32.0</code></pre>
</div>
</div>
<p>Now flip one decision. Keep both indices free instead of matching them, <img src="https://latex.codecogs.com/png.latex?M_%7Bij%7D%20=%20v_i%20w_j">, and no sum happens at all. Every component of <img src="https://latex.codecogs.com/png.latex?%5Cmathbf%7Bv%7D"> meets every component of <img src="https://latex.codecogs.com/png.latex?%5Cmathbf%7Bw%7D">:</p>
<div id="4f0e5392" class="cell" data-execution_count="6">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb8" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb8-1"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(np.einsum(<span class="st" style="color: #20794D;
background-color: null;
font-style: inherit;">'i,j-&gt;ij'</span>, v, w))</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>[[ 4.  5.  6.]
 [ 8. 10. 12.]
 [12. 15. 18.]]</code></pre>
</div>
</div>
<p>This is the outer product, a 3×3 grid, and the string told you its shape before running anything: two free indices in, two dimensions out. Compare the two strings. <code>'j,j-&gt;'</code> and <code>'i,j-&gt;ij'</code> differ only in whether the letters match, and that single choice is the difference between collapsing to a scalar and expanding to a matrix.</p>
<p>Two more one-liners fall out. Transposing is just reordering free indices, no sum anywhere:</p>
<div id="e31ea76a" class="cell" data-execution_count="7">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb10" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb10-1">M <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.einsum(<span class="st" style="color: #20794D;
background-color: null;
font-style: inherit;">'i,j-&gt;ij'</span>, v, w)</span>
<span id="cb10-2"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(M.shape)</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>(3, 3)</code></pre>
</div>
</div>
<p>And repeating an index within a <em>single</em> tensor walks its diagonal. <img src="https://latex.codecogs.com/png.latex?A_%7Bii%7D"> means sum the entries where row equals column, the trace:</p>
<div id="9ccb0e35" class="cell" data-execution_count="8">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb12" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb12-1"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(np.einsum(<span class="st" style="color: #20794D;
background-color: null;
font-style: inherit;">'ii-&gt;'</span>, M))</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>32.0</code></pre>
</div>
</div>
</section>
<section id="step-4-a-matrix-acts-on-a-vector" class="level2">
<h2 class="anchored" data-anchor-id="step-4-a-matrix-acts-on-a-vector">Step 4: a matrix acts on a vector</h2>
<p>A matrix-vector product, written per component, is a dot product between row <img src="https://latex.codecogs.com/png.latex?i"> and the vector:</p>
<p><img src="https://latex.codecogs.com/png.latex?u_i%20=%20A_%7Bij%7D%5C,%20v_j"></p>
<p>This is exactly the loop nest from Step 2, now in notation: <img src="https://latex.codecogs.com/png.latex?j"> is the loop variable, consumed by its sum; <img src="https://latex.codecogs.com/png.latex?i"> names the output slot. One free index, so the output is a vector. The string is a direct transcription:</p>
<div id="ac80046a" class="cell" data-execution_count="9">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb14" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb14-1">A_demo <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.arange(<span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">9.0</span>).reshape(<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>, <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>)</span>
<span id="cb14-2">dotted <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.einsum(<span class="st" style="color: #20794D;
background-color: null;
font-style: inherit;">'ij,j-&gt;i'</span>, A_demo, v)</span>
<span id="cb14-3"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(dotted, dotted.shape)</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>[ 8. 26. 44.] (3,)</code></pre>
</div>
</div>
<p>Matrix-matrix multiplication is the same sentence with one more free index along for the ride, <img src="https://latex.codecogs.com/png.latex?C_%7Bik%7D%20=%20A_%7Bij%7DB_%7Bjk%7D">: contract <img src="https://latex.codecogs.com/png.latex?A">’s columns against <img src="https://latex.codecogs.com/png.latex?B">’s rows over the shared <img src="https://latex.codecogs.com/png.latex?j">, keep <img src="https://latex.codecogs.com/png.latex?i"> and <img src="https://latex.codecogs.com/png.latex?k"> to shape the result.</p>
<div id="e4484d06" class="cell" data-execution_count="10">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb16" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb16-1">B_demo <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.ones((<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>, <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>))</span>
<span id="cb16-2">matmul <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.einsum(<span class="st" style="color: #20794D;
background-color: null;
font-style: inherit;">'ij,jk-&gt;ik'</span>, A_demo, B_demo)</span>
<span id="cb16-3"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(matmul, matmul.shape)</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>[[ 3.  3.  3.]
 [12. 12. 12.]
 [21. 21. 21.]] (3, 3)</code></pre>
</div>
</div>
</section>
<section id="step-5-where-a-real-matrix-comes-from" class="level2">
<h2 class="anchored" data-anchor-id="step-5-where-a-real-matrix-comes-from">Step 5: where a real matrix comes from</h2>
<p>So far every matrix has been handed to us. The notation gets genuinely powerful when an equation <em>produces</em> a matrix you didn’t know was there. Here is the example that made this click for me: lets sayP&gt; sunlight with direction <img src="https://latex.codecogs.com/png.latex?%5Chat%7B%5Cmathbf%7Bs%7D%7D"> casts a point <img src="https://latex.codecogs.com/png.latex?%5Cmathbf%7Br%7D"> onto the ground (the plane where the third coordinate is zero). Sliding along the ray until it hits the ground gives the shadow point</p>
<p><img src="https://latex.codecogs.com/png.latex?%5Cmathbf%7Br%7D'%20=%20%5Cmathbf%7Br%7D%20-%20%5Chat%7B%5Cmathbf%7Bs%7D%7D%5C,%5Cfrac%7B%5Cmathbf%7Br%7D%5Ccdot%5Chat%7B%5Cmathbf%7Be%7D%7D_3%7D%7Bs_3%7D"></p>
<p>where <img src="https://latex.codecogs.com/png.latex?%5Chat%7B%5Cmathbf%7Be%7D%7D_3%20=%20(0,0,1)"> picks out height and <img src="https://latex.codecogs.com/png.latex?s_3"> is the sun direction’s third component. The right side is linear in <img src="https://latex.codecogs.com/png.latex?%5Cmathbf%7Br%7D">, so a matrix is hiding in it. Index notation digs it out in three moves, and to keep each move honest we’ll drag a dummy input through all of them: the sun at <img src="https://latex.codecogs.com/png.latex?%5Chat%7B%5Cmathbf%7Bs%7D%7D%20=%20(4/13,%20-3/13,%20-12/13)">, so <img src="https://latex.codecogs.com/png.latex?s_3%20=%20-12/13">, and a point at <img src="https://latex.codecogs.com/png.latex?%5Cmathbf%7Br%7D%20=%20(6,%202,%203)">, three units above the ground.</p>
<div id="84e9ba8c" class="cell" data-execution_count="11">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb18" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb18-1">s <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.array([<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">4</span><span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">/</span><span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">13</span>, <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">-</span><span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span><span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">/</span><span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">13</span>, <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">-</span><span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">12</span><span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">/</span><span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">13</span>])</span>
<span id="cb18-2">r <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.array([<span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">6.0</span>, <span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">2.0</span>, <span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">3.0</span>])</span>
<span id="cb18-3">e3 <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.array([<span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">0.0</span>, <span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">0.0</span>, <span class="fl" style="color: #AD0000;
background-color: null;
font-style: inherit;">1.0</span>])</span></code></pre></div></div>
</div>
<p><strong>Move 1: components.</strong> Write the <img src="https://latex.codecogs.com/png.latex?i">-th component of each side, with the dot product in convention form:</p>
<p><img src="https://latex.codecogs.com/png.latex?r'_i%20=%20r_i%20-%20s_i%5C,%20%5B%5Chat%7B%5Cmathbf%7Be%7D%7D_3%5D_j%5C,%20r_j%20/%20s_3"></p>
<p>Roles: <img src="https://latex.codecogs.com/png.latex?j"> is the loop variable, <img src="https://latex.codecogs.com/png.latex?i"> shapes the output, present once in every term on both sides. Well-formed.</p>
<p>Now watch <img src="https://latex.codecogs.com/png.latex?j"> actually loop and vanish. The piece <img src="https://latex.codecogs.com/png.latex?%5B%5Chat%7B%5Cmathbf%7Be%7D%7D_3%5D_j%5C,%20r_j"> expands to <img src="https://latex.codecogs.com/png.latex?0%5Ccdot%206%20+%200%5Ccdot%202%20+%201%5Ccdot%203"> i.e.&nbsp;a dot product between the two vectors, the basis <img src="https://latex.codecogs.com/png.latex?%5Chat%7B%5Cmathbf%7Be%7D%7D_3"> vector and the point <img src="https://latex.codecogs.com/png.latex?%5Cmathbf%7Br%7D">, which is exactly the height of the point. Check it:</p>
<div id="cd94c176" class="cell" data-execution_count="12">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb19" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb19-1"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(<span class="bu" style="color: null;
background-color: null;
font-style: inherit;">sum</span>(e3[j] <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">*</span> r[j] <span class="cf" style="color: #003B4F;
background-color: null;
font-weight: bold;
font-style: inherit;">for</span> j <span class="kw" style="color: #003B4F;
background-color: null;
font-weight: bold;
font-style: inherit;">in</span> <span class="bu" style="color: null;
background-color: null;
font-style: inherit;">range</span>(<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>)))</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>3.0</code></pre>
</div>
</div>
<p>It collapses to <img src="https://latex.codecogs.com/png.latex?3">, the height of the point, with no <img src="https://latex.codecogs.com/png.latex?j"> anywhere in the answer. The equation for the first component, <img src="https://latex.codecogs.com/png.latex?i=1">, is then plain arithmetic: <img src="https://latex.codecogs.com/png.latex?r'_1%20=%206%20-%20(4/13)%5Ccdot%203/(-12/13)%20=%206%20+%201%20=%207">.</p>
<div id="18b8a0ca" class="cell" data-execution_count="13">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb21" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb21-1"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(r[<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">0</span>] <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">-</span> s[<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">0</span>] <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">*</span> <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span> <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">/</span> s[<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">2</span>])</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>7.0</code></pre>
</div>
</div>
<p>So the first shadow coordinate is <img src="https://latex.codecogs.com/png.latex?7"> before any matrix exists. The remaining moves don’t change this number; they reorganize the computation so a matrix becomes visible.</p>
<p><strong>Move 2: rewrite both terms using the identity matrix.</strong> Here is where we left off:</p>
<p><img src="https://latex.codecogs.com/png.latex?r'_i%20=%20r_i%20-%20s_i%5C,%20%5B%5Chat%7B%5Cmathbf%7Be%7D%7D_3%5D_j%5C,%20r_j%20/%20s_3"></p>
<p>The goal of this move is simple: make both terms look alike, so the next move can factor them. Two small rewrites do it.</p>
<p>First, <img src="https://latex.codecogs.com/png.latex?%5B%5Chat%7B%5Cmathbf%7Be%7D%7D_3%5D_j"> is the list <img src="https://latex.codecogs.com/png.latex?(0,0,1)"> exactly row 3 of the identity matrix. So we can write it as <img src="https://latex.codecogs.com/png.latex?I_%7B3j%7D">:</p>
<div id="0c5dcf58" class="cell" data-execution_count="14">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb23" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb23-1">I <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.eye(<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>)</span>
<span id="cb23-2"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(I[<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">2</span>])</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>[0. 0. 1.]</code></pre>
</div>
</div>
<p>Second, the bare <img src="https://latex.codecogs.com/png.latex?r_i"> can be written as <img src="https://latex.codecogs.com/png.latex?I_%7Bij%7Dr_j">. Row <img src="https://latex.codecogs.com/png.latex?i"> of the identity has a single 1 in slot <img src="https://latex.codecogs.com/png.latex?i">, so dotting it with <img src="https://latex.codecogs.com/png.latex?%5Cmathbf%7Br%7D"> just returns component <img src="https://latex.codecogs.com/png.latex?i">. Check for <img src="https://latex.codecogs.com/png.latex?i=2">: row <img src="https://latex.codecogs.com/png.latex?(0,1,0)"> against <img src="https://latex.codecogs.com/png.latex?(6,2,3)"> gives <img src="https://latex.codecogs.com/png.latex?2"> same for the remaining values meaning <img src="https://latex.codecogs.com/png.latex?r_i%20=%20I_%7Bij%7D%20r_j"> is valid for all <img src="https://latex.codecogs.com/png.latex?i">:</p>
<div id="e084d4e1" class="cell" data-execution_count="15">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb25" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb25-1"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(<span class="bu" style="color: null;
background-color: null;
font-style: inherit;">sum</span>(I[<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">1</span>, j] <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">*</span> r[j] <span class="cf" style="color: #003B4F;
background-color: null;
font-weight: bold;
font-style: inherit;">for</span> j <span class="kw" style="color: #003B4F;
background-color: null;
font-weight: bold;
font-style: inherit;">in</span> <span class="bu" style="color: null;
background-color: null;
font-style: inherit;">range</span>(<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>)))</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>2.0</code></pre>
</div>
</div>
<p>Why dress up a term that was already simple? Because factoring needs both terms to end in the same <img src="https://latex.codecogs.com/png.latex?r_j">, and <img src="https://latex.codecogs.com/png.latex?r_i"> alone has no <img src="https://latex.codecogs.com/png.latex?j"> to offer.</p>
<p><strong>Move 3: factor out <img src="https://latex.codecogs.com/png.latex?r_j">.</strong> With both rewrites in place, every term ends in <img src="https://latex.codecogs.com/png.latex?r_j">, and it pulls out:</p>
<p><img src="https://latex.codecogs.com/png.latex?r'_i%20=%20I_%7Bij%7D%5C,%20r_j%20-%20%5Cfrac%7Bs_i%5C,%20I_%7B3j%7D%7D%7Bs_3%7D%5C,%20r_j%20=%20%5Cleft(I_%7Bij%7D%20-%20%5Cfrac%7Bs_i%5C,%20I_%7B3j%7D%7D%7Bs_3%7D%5Cright)%20r_j"></p>
<p>The bracket has two free indices, <img src="https://latex.codecogs.com/png.latex?i"> and <img src="https://latex.codecogs.com/png.latex?j"> so it is a matrix. This is the matrix-vector pattern <img src="https://latex.codecogs.com/png.latex?%5BA%5Cmathbf%7Br%7D%5D_i%20=%20A_%7Bij%7Dr_j"> from Step 4, with</p>
<p><img src="https://latex.codecogs.com/png.latex?A_%7Bij%7D%20=%20I_%7Bij%7D%20-%20%5Cfrac%7Bs_i%5C,%20I_%7B3j%7D%7D%7Bs_3%7D"></p>
<p>Check with the dummy input: row <img src="https://latex.codecogs.com/png.latex?i=1"> of the bracket is <img src="https://latex.codecogs.com/png.latex?(1,%200,%201/3)">, and dotting it with <img src="https://latex.codecogs.com/png.latex?(6,2,3)"> should reproduce the <img src="https://latex.codecogs.com/png.latex?7"> from Move 1:</p>
<div id="425541e3" class="cell" data-execution_count="16">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb27" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb27-1">A <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> I <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">-</span> np.outer(s, I[<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">2</span>]) <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">/</span> s[<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">2</span>]   <span class="co" style="color: #5E5E5E;
background-color: null;
font-style: inherit;"># A_ij = I_ij - s_i I_3j / s3</span></span>
<span id="cb27-2"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(A)</span>
<span id="cb27-3"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(<span class="bu" style="color: null;
background-color: null;
font-style: inherit;">sum</span>(A[<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">0</span>, j] <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">*</span> r[j] <span class="cf" style="color: #003B4F;
background-color: null;
font-weight: bold;
font-style: inherit;">for</span> j <span class="kw" style="color: #003B4F;
background-color: null;
font-weight: bold;
font-style: inherit;">in</span> <span class="bu" style="color: null;
background-color: null;
font-style: inherit;">range</span>(<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>)))</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>[[ 1.          0.          0.33333333]
 [ 0.          1.         -0.25      ]
 [ 0.          0.          0.        ]]
7.0</code></pre>
</div>
</div>
<p>Same number, now produced by a matrix row. Note <code>np.outer(s, I[2])</code> is exactly the <img src="https://latex.codecogs.com/png.latex?s_i%20I_%7B3j%7D"> term: two free indices, no shared letter the outer product from Step 3 in the wild.</p>
<p>The matrix itself is readable. <img src="https://latex.codecogs.com/png.latex?I_%7B3j%7D"> is zero unless <img src="https://latex.codecogs.com/png.latex?j=3">, so only column 3 differs from the identity: all the shadow action lives in one column. And at <img src="https://latex.codecogs.com/png.latex?i=3"> the two terms cancel, so row 3 is all zeros: every output has height zero the 3D-to-2D collapse written as a dead row.</p>
<p>Applying the full matrix finishes the example:</p>
<div id="d8d20007" class="cell" data-execution_count="17">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb29" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb29-1"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(A <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">@</span> r)</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>[7.   1.25 0.  ]</code></pre>
</div>
</div>
<p>The shadow lands at <img src="https://latex.codecogs.com/png.latex?(7,%5C%205/4,%5C%200)">: shifted <img src="https://latex.codecogs.com/png.latex?+x"> because the sun direction has positive <img src="https://latex.codecogs.com/png.latex?x">, shifted down in <img src="https://latex.codecogs.com/png.latex?y"> because it has negative <img src="https://latex.codecogs.com/png.latex?y">, and flat on the ground.</p>
<p>One bonus check. A shadow of a shadow is itself, so <img src="https://latex.codecogs.com/png.latex?A%5E2%20=%20A">:</p>
<div id="64840ef5" class="cell" data-execution_count="18">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb31" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb31-1"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(np.allclose(A <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">@</span> A, A))</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>True</code></pre>
</div>
</div>
<p>Maps with this property are called projections, and this one is oblique: unless the sun is directly overhead, points travel diagonally to the plane rather than dropping straight down.</p>
</section>
<section id="step-6-the-batch-dimension-appears" class="level2">
<h2 class="anchored" data-anchor-id="step-6-the-batch-dimension-appears">Step 6: the batch dimension appears</h2>
<p>Now shadow four points at once. Stack them as columns of a matrix <img src="https://latex.codecogs.com/png.latex?R"> and write the transformation of the whole batch:</p>
<p><img src="https://latex.codecogs.com/png.latex?R'_%7Bia%7D%20=%20A_%7Bij%7D%5C,%20R_%7Bja%7D"></p>
<p>Read the roles. <img src="https://latex.codecogs.com/png.latex?j">: appears twice, the loop variable, doing the actual contraction. <img src="https://latex.codecogs.com/png.latex?i">: free, output component. <img src="https://latex.codecogs.com/png.latex?a">: free, and it means <em>which vector you are on</em>. The crucial observation is what <img src="https://latex.codecogs.com/png.latex?a"> does not do. It never enters the sum. Nothing in the equation lets column 3 influence column 2. Batching works not by accident but because the batch index is free on both sides, and the notation makes that independence visible rather than asking you to trust it.</p>
<div id="a84d8c95" class="cell" data-execution_count="19">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb33" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb33-1">R <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.array([[<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">5</span>, <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">-</span><span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">1</span>, <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">-</span><span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>,  <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">7</span>],</span>
<span id="cb33-2">              [<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">4</span>, <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">-</span><span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">4</span>,  <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">1</span>, <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">-</span><span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">2</span>],</span>
<span id="cb33-3">              [<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">9</span>,  <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>,  <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">0</span>, <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">12</span>]], dtype<span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">float</span>)</span>
<span id="cb33-4"></span>
<span id="cb33-5">R_batch <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.einsum(<span class="st" style="color: #20794D;
background-color: null;
font-style: inherit;">'ij,ja-&gt;ia'</span>, A, R)</span>
<span id="cb33-6"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(R_batch)</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>[[ 8.   -0.   -3.   11.  ]
 [ 1.75 -4.75  1.   -5.  ]
 [ 0.    0.    0.    0.  ]]</code></pre>
</div>
</div>
<p>The second check confirms it: the einsum equals four independent matrix-vector products laid side by side. This is the same mechanism by which a batch dimension in a deep learning framework broadcasts through a matmul without mixing examples.</p>
<p>Before scaling up further, a checkpoint. Write your answer before expanding.</p>
<div class="callout callout-style-default callout-note callout-titled" title="Checkpoint: what does `np.einsum('ij,jk,ka->ia', A, A, R)` compute, and what should it equal here?">
<div class="callout-header d-flex align-content-center collapsed" data-bs-toggle="collapse" data-bs-target=".callout-1-contents" aria-controls="callout-1" aria-expanded="false" aria-label="Toggle callout">
<div class="callout-icon-container">
<i class="callout-icon"></i>
</div>
<div class="callout-title-container flex-fill">
<span class="screen-reader-only">Note</span>Checkpoint: what does <code>np.einsum('ij,jk,ka-&gt;ia', A, A, R)</code> compute, and what should it equal here?
</div>
<div class="callout-btn-toggle d-inline-block border-0 py-1 ps-1 pe-0 float-end"><i class="callout-toggle"></i></div>
</div>
<div id="callout-1" class="callout-1-contents callout-collapse collapse">
<div class="callout-body-container callout-body">
<p>Read the string one index at a time. <img src="https://latex.codecogs.com/png.latex?j"> appears in <code>ij</code> and <code>jk</code> a dummy, contracting the first <img src="https://latex.codecogs.com/png.latex?A"> against the second. <img src="https://latex.codecogs.com/png.latex?k"> appears in <code>jk</code> and <code>ka</code> another dummy, contracting that result against <img src="https://latex.codecogs.com/png.latex?R">. What survives is <img src="https://latex.codecogs.com/png.latex?i"> (output component) and <img src="https://latex.codecogs.com/png.latex?a"> (which vector in the batch), so the output is <code>ia</code>, the same shape as one batched matmul:</p>
<p><img src="https://latex.codecogs.com/png.latex?R''_%7Bia%7D%20=%20A_%7Bij%7D%5C,%20A_%7Bjk%7D%5C,%20R_%7Bka%7D"></p>
<p>In words: apply <img src="https://latex.codecogs.com/png.latex?A"> twice to the whole batch, examples still kept independent by the free <img src="https://latex.codecogs.com/png.latex?a">.</p>
<div id="5107320a" class="cell" data-execution_count="20">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb35" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb35-1"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(np.einsum(<span class="st" style="color: #20794D;
background-color: null;
font-style: inherit;">'ij,jk,ka-&gt;ia'</span>, A, A, R))</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>[[ 8.   -0.   -3.   11.  ]
 [ 1.75 -4.75  1.   -5.  ]
 [ 0.    0.    0.    0.  ]]</code></pre>
</div>
</div>
</div>
</div>
</div>
</section>
<section id="step-7-reading-the-attention-string" class="level2">
<h2 class="anchored" data-anchor-id="step-7-reading-the-attention-string">Step 7: reading the attention string</h2>
<p>The opening line can now be read as a sentence:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb37" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb37-1">attn <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> torch.einsum(<span class="st" style="color: #20794D;
background-color: null;
font-style: inherit;">'bhqd,bhkd-&gt;bhqk'</span>, q, k)</span></code></pre></div></div>
<p>Hunt for the loop variable first: which letter is missing after the arrow? Only <img src="https://latex.codecogs.com/png.latex?d">, the feature dimension. So <img src="https://latex.codecogs.com/png.latex?d"> is the one that loops and vanishes: queries and keys are contracted along features, a dot product per pair, exactly the <code>'j,j-&gt;'</code> from Step 1 wearing more clothes.</p>
<p>Then account for the free indices. <img src="https://latex.codecogs.com/png.latex?b"> (batch) and <img src="https://latex.codecogs.com/png.latex?h"> (head) appear in both inputs and the output in matching positions: they ride along untouched, exactly like <img src="https://latex.codecogs.com/png.latex?a"> did in Step 6, two batch dimensions instead of one. <img src="https://latex.codecogs.com/png.latex?q"> comes only from the first tensor and <img src="https://latex.codecogs.com/png.latex?k"> only from the second, and both survive: every query meets every key, the outer-product pattern from Step 3. The output is a <img src="https://latex.codecogs.com/png.latex?q%20%5Ctimes%20k"> grid of similarity scores for each batch and head.</p>
<div id="47841593" class="cell" data-execution_count="21">
<div class="code-copy-outer-scaffold"><div class="sourceCode cell-code" id="cb38" style="background: #f1f3f5;"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb38-1">rng <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.random.default_rng(<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">0</span>)</span>
<span id="cb38-2">q <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> rng.standard_normal((<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">2</span>, <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>, <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">4</span>, <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">8</span>))   <span class="co" style="color: #5E5E5E;
background-color: null;
font-style: inherit;"># batch=2, heads=3, queries=4, features=8</span></span>
<span id="cb38-3">k <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> rng.standard_normal((<span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">2</span>, <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">3</span>, <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">5</span>, <span class="dv" style="color: #AD0000;
background-color: null;
font-style: inherit;">8</span>))   <span class="co" style="color: #5E5E5E;
background-color: null;
font-style: inherit;"># batch=2, heads=3, keys=5, features=8</span></span>
<span id="cb38-4"></span>
<span id="cb38-5">scores <span class="op" style="color: #5E5E5E;
background-color: null;
font-style: inherit;">=</span> np.einsum(<span class="st" style="color: #20794D;
background-color: null;
font-style: inherit;">'bhqd,bhkd-&gt;bhqk'</span>, q, k)</span>
<span id="cb38-6"><span class="bu" style="color: null;
background-color: null;
font-style: inherit;">print</span>(scores.shape)</span></code></pre></div></div>
<div class="cell-output cell-output-stdout">
<pre><code>(2, 3, 4, 5)</code></pre>
</div>
</div>
<p>The shape check alone shows the grammar working: the string predicted <code>(2, 3, 4, 5)</code> before the computation ran, and the einsum matches the conventional transpose-and-matmul spelling while stating its intent far more directly.</p>
</section>
<section id="where-the-mental-model-stops" class="level2">
<h2 class="anchored" data-anchor-id="where-the-mental-model-stops">Where the mental model stops</h2>
<p>First, reading an einsum tells you <em>what</em> is computed, not <em>how fast</em>. With three or more tensors, as in the checkpoint above, there is a choice of which pair to contract first, and for large tensors that choice can change the cost by orders of magnitude. NumPy’s <code>optimize=True</code> flag and the <a href="https://dgasmith.github.io/opt_einsum/">opt_einsum</a> library exist to search for a good contraction order. The notation is about meaning; scheduling is a separate problem.</p>
<p>Second, not everything is an einsum. The convention expresses sums of products over matched indices: contractions, traces, transposes, outer products, and their batched combinations. Elementwise nonlinearities, softmax, max-pooling, anything with a comparison or a division by a computed sum, live outside it. When you see <code>einsum</code> in a codebase, you are looking at the multilinear skeleton of the computation, with the nonlinear flesh applied elsewhere. In the attention example, the softmax and the value-weighting that follow the score computation are separate lines for exactly this reason.</p>


</section>

 ]]></description>
  <category>machine-learning</category>
  <guid>https://nakranivaibhav.github.io/posts/einsum-one-index-at-a-time.html</guid>
  <pubDate>Fri, 31 Jul 2026 18:30:00 GMT</pubDate>
</item>
<item>
  <title>NVFP4 Is a Scaling System, Not Just a 4-Bit Number</title>
  <dc:creator>Vaibhav Nakrani</dc:creator>
  <link>https://nakranivaibhav.github.io/posts/nvfp4-scaling-system.html</link>
  <description><![CDATA[ 




<style>
.nv-lab {
  border: 1px solid #cbd5e1;
  border-radius: .8rem;
  background: #f8fafc;
  margin: 1.25rem 0;
  padding: 1rem 1.1rem;
}
.nv-lab h3 { margin-top: 0; font-size: 1.05rem; }
.nv-lab-note { color: #475569; font-size: .92rem; margin-bottom: 1rem; }
.nv-controls { display: grid; grid-template-columns: repeat(3, minmax(10rem, 1fr)); gap: 1rem; }
.nv-control { display: flex; flex-direction: column; gap: .4rem; }
.nv-control label { font-weight: 650; color: #1e293b; }
.nv-control select, .nv-control input { width: 100%; }
.nv-control output { font-variant-numeric: tabular-nums; color: #1d4ed8; font-weight: 700; }
.nv-equation {
  margin: 1rem 0 .75rem;
  padding: .9rem;
  border-radius: .6rem;
  background: #eff6ff;
  color: #1e3a8a;
  text-align: center;
  font: 700 1.1rem/1.5 ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
}
.nv-scope-grid, .nv-metric-grid { display: grid; grid-template-columns: repeat(3, 1fr); gap: .65rem; }
.nv-scope, .nv-metric { border-radius: .55rem; background: white; border: 1px solid #dbe3ec; padding: .7rem; }
.nv-scope strong, .nv-metric strong { display: block; color: #0f172a; }
.nv-scope span, .nv-metric span { color: #475569; font-size: .88rem; }
.nv-codebook { display: flex; flex-wrap: wrap; gap: .35rem; margin-top: .75rem; }
.nv-code { background: #e2e8f0; border-radius: 99px; padding: .18rem .5rem; font: .85rem ui-monospace, SFMono-Regular, Menlo, Consolas, monospace; }
.nv-code.active { background: #2563eb; color: white; }
.nv-stack { display: flex; height: 1.5rem; overflow: hidden; border-radius: .35rem; margin: 1rem 0 .45rem; background: #e2e8f0; }
.nv-stack-payload { background: #2563eb; }
.nv-stack-local { background: #f59e0b; min-width: 2px; }
.nv-stack-global { background: #7c3aed; min-width: 2px; }
.nv-legend { display: flex; flex-wrap: wrap; gap: .8rem; font-size: .86rem; color: #475569; }
.nv-swatch { display: inline-block; width: .75rem; height: .75rem; border-radius: .18rem; margin-right: .25rem; vertical-align: -.05rem; }
.nv-payload { background: #2563eb; } .nv-local { background: #f59e0b; } .nv-global { background: #7c3aed; }
.nv-result { margin: .9rem 0 0; color: #334155; }
@media (max-width: 760px) {
  .nv-controls, .nv-scope-grid, .nv-metric-grid { grid-template-columns: 1fr; }
}
</style>
<p>Suppose a model tensor is stored in NVFP4. It is tempting to picture every entry as a miniature, self-contained FP16 value: four bits go in, one approximate real number comes out, and the rest is a hardware detail.</p>
<p>That picture is incomplete. The four bits only select from a tiny set of normalized values. To reconstruct a value in the tensor’s actual range, NVFP4 also needs a scale shared by a small block and another scale shared by the whole tensor.</p>
<p>The question I want to resolve is therefore narrower than “Is FP4 good?” It is: <strong>what job do the payload, block scale, and tensor scale each perform, and when can that representation help a real workload?</strong></p>
<p>I will assume familiarity with tensors, ordinary floating point, quantization, and matrix multiplication. Everything specific to FP4 and NVIDIA’s implementation starts from scratch.</p>
<section id="four-bits-do-not-provide-many-choices" class="level2">
<h2 class="anchored" data-anchor-id="four-bits-do-not-provide-many-choices">Four bits do not provide many choices</h2>
<p>A floating-point layout is often named by how many exponent and mantissa bits it contains. <strong>E2M1</strong> has one sign bit, two exponent bits, and one mantissa bit. Together, they make a four-bit payload, so E2M1 is one kind of FP4.</p>
<p>Ignoring the sign, its finite magnitudes are:</p>
<p><code>0, 0.5, 1, 1.5, 2, 3, 4, 6</code></p>
<p>The sign bit mirrors the nonzero values below zero. This is a very small <strong>codebook</strong>, meaning the set of values the bit patterns can select. NVIDIA documents this E2M1 payload as the element format used by NVFP4 (<a href="https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/features/low_precision_training/nvfp4/nvfp4.html">Transformer Engine documentation</a>).</p>
<p>Now consider a block whose useful values cluster near <code>0.01</code>, and another whose values cluster near <code>100</code>. A bare E2M1 codebook cannot closely cover both neighborhoods. Its largest finite magnitude is <code>6</code>, and the gaps between available values are large relative to many tensors.</p>
<p>Quantization solves this by treating the codebook values as numbers in a local unit rather than final tensor values. A scale defines that unit. The payload says “four units”; the scale says how large one unit is here.</p>
</section>
<section id="the-smallest-example-the-same-payload-two-real-values" class="level2">
<h2 class="anchored" data-anchor-id="the-smallest-example-the-same-payload-two-real-values">The smallest example: the same payload, two real values</h2>
<p>NVFP4 reconstructs a value with three factors:</p>
<p><code>x = x_e2m1 × s_block × s_global</code></p>
<p>Their scopes are more important than their names:</p>
<ul>
<li><code>x_e2m1</code> is the four-bit code stored for one element.</li>
<li><code>s_block</code> is a scale shared by a local block of 16 elements.</li>
<li><code>s_global</code> is a scale shared by the whole tensor.</li>
</ul>
<p>Take two 16-value blocks and fix the tensor-wide scale at <code>s_global = 0.25</code>.</p>
<pre class="text"><code>Tensor-wide scale: s_global = 0.25

Block A: s_block = 1.0       Block B: s_block = 1.5
         payload = 4                  payload = 4
              |                            |
              v                            v
      4 × 1.0 × 0.25 = 1.0        4 × 1.5 × 0.25 = 1.5</code></pre>
<p>Both elements contain the same E2M1 payload. Yet one reconstructs to <code>1.0</code> and the other to <code>1.5</code>, because each block uses a different local unit. The payload is not a standalone measurement; it is a code interpreted within two shared coordinate systems.</p>
<p>The explorer below lets you vary those three factors separately. Choose an exact E2M1 payload, then move the block scale or tensor scale. The block control changes the unit for one toy 16-value neighborhood; the tensor control changes the unit shared by every block. The blue result is the reconstructed value, not a quantization-error or model-quality measurement.</p>
<section class="nv-lab" id="nv-scale-lab" aria-labelledby="nv-scale-title">
  <h3 id="nv-scale-title" class="anchored">Toy explorer: one stored code, two scaling levels</h3>
  <p class="nv-lab-note"><strong>Toy calculation, not benchmark data.</strong> The controls use exact codebook values and simplified scale choices to expose the reconstruction rule.</p>
  <div class="nv-controls">
    <div class="nv-control">
      <label for="nv-payload-input">Stored E2M1 value</label>
      <select id="nv-payload-input">
        <option value="-6">−6</option><option value="-4">−4</option><option value="-3">−3</option>
        <option value="-2">−2</option><option value="-1.5">−1.5</option><option value="-1">−1</option>
        <option value="-0.5">−0.5</option><option value="0">0</option><option value="0.5">0.5</option>
        <option value="1">1</option><option value="1.5">1.5</option><option value="2">2</option>
        <option value="3">3</option><option value="4" selected="">4</option><option value="6">6</option>
      </select>
      <span>Selects one normalized codebook value.</span>
    </div>
    <div class="nv-control">
      <label for="nv-block-input">Block scale: <output id="nv-block-value">1.50</output></label>
      <input id="nv-block-input" type="range" min="0.25" max="2" step="0.25" value="1.5">
      <span>Shared by this toy block of 16 values.</span>
    </div>
    <div class="nv-control">
      <label for="nv-global-input">Tensor scale: <output id="nv-global-value">0.25</output></label>
      <input id="nv-global-input" type="range" min="0.05" max="1" step="0.05" value="0.25">
      <span>Shared by every block in the toy tensor.</span>
    </div>
  </div>
  <div class="nv-codebook" id="nv-codebook" aria-label="Signed E2M1 codebook values"></div>
  <div class="nv-equation" id="nv-equation" aria-live="polite">4 × 1.50 × 0.25 = 1.50</div>
  <div class="nv-scope-grid">
    <div class="nv-scope"><strong>Payload</strong><span>Chooses a normalized E2M1 value for one element.</span></div>
    <div class="nv-scope"><strong>Block scale</strong><span>Changes the local unit for one 16-value neighborhood.</span></div>
    <div class="nv-scope"><strong>Tensor scale</strong><span>Places all local units within one tensor-wide range.</span></div>
  </div>
  <p class="nv-result" id="nv-scale-reading">With the defaults, the stored code 4 becomes 1.5 in real tensor units.</p>
</section>
<script>
(() => {
  const payloadInput = document.getElementById('nv-payload-input');
  const blockInput = document.getElementById('nv-block-input');
  const globalInput = document.getElementById('nv-global-input');
  const blockValue = document.getElementById('nv-block-value');
  const globalValue = document.getElementById('nv-global-value');
  const equation = document.getElementById('nv-equation');
  const reading = document.getElementById('nv-scale-reading');
  const codebook = document.getElementById('nv-codebook');
  const values = [-6, -4, -3, -2, -1.5, -1, -0.5, 0, 0.5, 1, 1.5, 2, 3, 4, 6];
  const clean = value => {
    const rounded = Math.round(value * 10000) / 10000;
    return Number.isInteger(rounded) ? rounded.toFixed(0) : rounded.toString();
  };
  values.forEach(value => {
    const chip = document.createElement('span');
    chip.className = 'nv-code';
    chip.dataset.value = value;
    chip.textContent = clean(value);
    codebook.appendChild(chip);
  });
  const update = () => {
    const payload = Number(payloadInput.value);
    const block = Number(blockInput.value);
    const global = Number(globalInput.value);
    const result = payload * block * global;
    blockValue.value = block.toFixed(2);
    globalValue.value = global.toFixed(2);
    equation.textContent = `${clean(payload)} × ${block.toFixed(2)} × ${global.toFixed(2)} = ${clean(result)}`;
    reading.textContent = `The four-bit selection is still ${clean(payload)}. Under this block's local unit and the tensor-wide unit, it reconstructs to ${clean(result)}.`;
    codebook.querySelectorAll('.nv-code').forEach(chip => chip.classList.toggle('active', Number(chip.dataset.value) === payload));
  };
  [payloadInput, blockInput, globalInput].forEach(input => input.addEventListener('input', update));
  update();
})();
</script>
<p>Try leaving the payload at <code>4</code> and changing only the block scale from <code>1.0</code> to <code>1.5</code>. The selected four-bit value does not move, but the reconstructed result does. Changing the tensor scale then moves every local unit together. This is the division of labor that the phrase “two-level scaling” can otherwise hide.</p>
<p>Why two levels instead of one? The block scale adapts the codebook to a small neighborhood, where values are more likely to have a similar range. The global scale places those local scales within the range that their own storage format can represent. Transformer Engine describes the same hierarchy and derives the global scale from the tensor’s maximum absolute value before deriving local scales from block maxima (<a href="https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/features/low_precision_training/nvfp4/nvfp4.html">Transformer Engine documentation</a>).</p>
<p>The format names can now be attached without hiding their roles:</p>
<ul>
<li>The element payload is <strong>E2M1</strong>, a four-bit floating-point code.</li>
<li>Each group of 16 payloads shares an <strong>E4M3</strong> scale. E4M3 is an eight-bit floating-point format with four exponent and three mantissa bits, so its scale can take fractional values between powers of two.</li>
<li>The tensor shares an <strong>FP32</strong> global scale, stored in ordinary 32-bit floating point.</li>
</ul>
<p>This hierarchy is what <strong>two-level scaling</strong> means in NVFP4: a local scale per block, followed by a global scale per tensor. <strong>Per-block scaling</strong> adapts to variation among small groups. <strong>Per-tensor scaling</strong> provides one range-setting factor for the entire tensor. Neither phrase says that every operation or every piece of model state is stored at that precision.</p>
</section>
<section id="bit-is-a-payload-count-not-a-memory-measurement" class="level2">
<h2 class="anchored" data-anchor-id="bit-is-a-payload-count-not-a-memory-measurement">“4-bit” is a payload count, not a memory measurement</h2>
<p>Once the local scale is visible, a common storage shortcut breaks.</p>
<p>One block contains 16 four-bit payloads and one eight-bit local scale:</p>
<p><code>16 × 4 + 8 = 72 bits</code></p>
<p>Amortized across the block:</p>
<p><code>72 / 16 = 4.5 bits per value</code></p>
<p>So the basic payload-plus-local-scale representation is already 4.5 bits per value, not four. The tensor’s FP32 global scale adds a small amortized cost whose size depends on the number of elements in that tensor.</p>
<p>The next explorer makes that amortization visible. Move the slider to change the number of 16-value blocks in one toy tensor. It counts four payload bits per value, one eight-bit local scale per block, and one 32-bit global scale for the tensor. It intentionally excludes padding, alignment, duplicate layouts, high-precision components, and runtime buffers, so its result is a representation calculation rather than an application-memory forecast.</p>
<section class="nv-lab" id="nv-storage-lab" aria-labelledby="nv-storage-title">
  <h3 id="nv-storage-title" class="anchored">Toy explorer: payload bits plus scale metadata</h3>
  <p class="nv-lab-note"><strong>Defined assumptions:</strong> 16 values per block, 4 payload bits per value, one 8-bit block scale, and one 32-bit tensor scale. No padding or runtime state is included.</p>
  <div class="nv-control">
    <label for="nv-block-count">Tensor size: <output id="nv-tensor-values">1,024 values</output></label>
    <input id="nv-block-count" type="range" min="0" max="10" step="1" value="6">
    <span>Each step doubles the number of 16-value blocks.</span>
  </div>
  <div class="nv-stack" aria-label="Share of counted bits by component">
    <div class="nv-stack-payload" id="nv-stack-payload" title="E2M1 payload bits"></div>
    <div class="nv-stack-local" id="nv-stack-local" title="E4M3 local-scale bits"></div>
    <div class="nv-stack-global" id="nv-stack-global" title="FP32 global-scale bits"></div>
  </div>
  <div class="nv-legend">
    <span><span class="nv-swatch nv-payload"></span>E2M1 payloads</span>
    <span><span class="nv-swatch nv-local"></span>E4M3 block scales</span>
    <span><span class="nv-swatch nv-global"></span>FP32 tensor scale</span>
  </div>
  <div class="nv-metric-grid">
    <div class="nv-metric"><strong id="nv-payload-bits">4,096 bits</strong><span>payload</span></div>
    <div class="nv-metric"><strong id="nv-scale-bits">544 bits</strong><span>all scale metadata</span></div>
    <div class="nv-metric"><strong id="nv-bits-per-value">4.531 bits/value</strong><span>under these assumptions</span></div>
  </div>
  <p class="nv-result" id="nv-storage-reading">At 1,024 values, the single global scale is already thinly amortized. The local scale still contributes 0.5 bit per value.</p>
</section>
<script>
(() => {
  const blockCountInput = document.getElementById('nv-block-count');
  const tensorValues = document.getElementById('nv-tensor-values');
  const payloadBitsOut = document.getElementById('nv-payload-bits');
  const scaleBitsOut = document.getElementById('nv-scale-bits');
  const bitsPerValueOut = document.getElementById('nv-bits-per-value');
  const payloadBar = document.getElementById('nv-stack-payload');
  const localBar = document.getElementById('nv-stack-local');
  const globalBar = document.getElementById('nv-stack-global');
  const reading = document.getElementById('nv-storage-reading');
  const number = value => value.toLocaleString('en-US');
  const update = () => {
    const blocks = 2 ** Number(blockCountInput.value);
    const values = blocks * 16;
    const payloadBits = values * 4;
    const localBits = blocks * 8;
    const globalBits = 32;
    const scaleBits = localBits + globalBits;
    const totalBits = payloadBits + scaleBits;
    const bitsPerValue = totalBits / values;
    tensorValues.value = `${number(values)} values (${number(blocks)} ${blocks === 1 ? 'block' : 'blocks'})`;
    payloadBitsOut.textContent = `${number(payloadBits)} bits`;
    scaleBitsOut.textContent = `${number(scaleBits)} bits`;
    bitsPerValueOut.textContent = `${bitsPerValue.toFixed(3)} bits/value`;
    payloadBar.style.width = `${100 * payloadBits / totalBits}%`;
    localBar.style.width = `${100 * localBits / totalBits}%`;
    globalBar.style.width = `${100 * globalBits / totalBits}%`;
    reading.textContent = `At ${number(values)} values, the local scales contribute 0.5 bit per value. The one tensor scale contributes ${(globalBits / values).toFixed(3)} more, giving ${bitsPerValue.toFixed(3)} bits/value before any uncounted system overhead.`;
  };
  blockCountInput.addEventListener('input', update);
  update();
})();
</script>
<p>At the smallest setting, the one global scale is expensive because only 16 values share it. As the toy tensor grows, that one-time cost fades, while the local-scale cost remains fixed at <code>8 / 16 = 0.5</code> bit per value. The output approaches 4.5 bits per value but never turns the calculation into a full-memory estimate.</p>
<p>Even 4.5 bits is not a full runtime-memory prediction. Transformer Engine documents padded and swizzled scale layouts and a TN-only NVFP4 GEMM layout; training may also keep rowwise and columnwise representations of weights (<a href="https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/features/low_precision_training/nvfp4/nvfp4.html">Transformer Engine documentation</a>). Alignment, padding, higher-precision layers, activations, optimizer state, KV cache, workspaces, and framework allocations all remain outside the small calculation.</p>
<p>Comparing payloads alone, four bits are one quarter of a BF16 payload and one half of an FP8 payload. Those ratios are useful for understanding the upper-level opportunity. They do not establish a fourfold reduction in total application memory.</p>
</section>
<section id="generic-fp4-mxfp4-and-nvfp4" class="level2">
<h2 class="anchored" data-anchor-id="generic-fp4-mxfp4-and-nvfp4">Generic FP4, MXFP4, and NVFP4</h2>
<p>“FP4” is not a complete deployment specification. It can identify a four-bit element format such as E2M1 without telling us how a tensor is scaled. Two systems can store the same E2M1 payload and reconstruct it differently.</p>
<p>MXFP4 makes the comparison concrete. <strong>MXFP4</strong>, the FP4 member of the OCP microscaling family, also uses E2M1 payloads. It shares one <strong>E8M0</strong> scale across 32 values. E8M0 has eight exponent bits and no mantissa, so its scale is restricted to powers of two. NVFP4 instead shares an E4M3 fractional scale across 16 values and adds the FP32 scale for the tensor (<a href="https://developer.nvidia.com/blog/introducing-nvfp4-for-efficient-and-accurate-low-precision-inference/">NVIDIA’s NVFP4 technical overview</a>).</p>
<table class="caption-top table">
<colgroup>
<col style="width: 25%">
<col style="width: 25%">
<col style="width: 25%">
<col style="width: 25%">
</colgroup>
<thead>
<tr class="header">
<th>Property</th>
<th>Generic FP4 / E2M1</th>
<th>MXFP4</th>
<th>NVFP4</th>
</tr>
</thead>
<tbody>
<tr class="odd">
<td>Element payload</td>
<td>May be E2M1</td>
<td>E2M1</td>
<td>E2M1</td>
</tr>
<tr class="even">
<td>Local scale</td>
<td>Not specified by the name</td>
<td>E8M0, powers of two</td>
<td>E4M3, fractional FP8</td>
</tr>
<tr class="odd">
<td>Values per local scale</td>
<td>Not specified</td>
<td>32</td>
<td>16</td>
</tr>
<tr class="even">
<td>Second scaling level</td>
<td>Not specified</td>
<td>None in the MX definition</td>
<td>One FP32 scale per tensor</td>
</tr>
</tbody>
</table>
<p>The table is a comparison of representations, not model quality. NVFP4 gets two ways to fit a local distribution more closely: its scale covers half as many elements, and that scale can sit between powers of two. NVIDIA illustrates lower quantization error from these choices, but a representation-error example is not an end-to-end accuracy result (<a href="https://developer.nvidia.com/blog/introducing-nvfp4-for-efficient-and-accurate-low-precision-inference/">NVIDIA’s NVFP4 technical overview</a>). Smaller blocks also mean more scale metadata and handling.</p>
<p>The better mental model is now precise: generic FP4 names a payload class; MXFP4 and NVFP4 specify different systems around an E2M1 payload.</p>
</section>
<section id="a-format-is-not-a-training-recipe" class="level2">
<h2 class="anchored" data-anchor-id="a-format-is-not-a-training-recipe">A format is not a training recipe</h2>
<p>If scaling were the whole story, low-precision training would amount to casting tensors and calling the same matrix multiplications. NVIDIA’s documented NVFP4 training path contains several additional mechanisms.</p>
<p><strong>Two-dimensional weight scaling</strong> changes the grouping used for weights. Transformer Engine uses 16-by-16 tiles by default so rowwise and columnwise weight representations can support the forward and backward matrix multiplications. Activations and gradients use one-dimensional 16-value blocks.</p>
<p><strong>Stochastic rounding</strong> chooses probabilistically between neighboring representable values instead of always taking the nearest one. For a gradient between two E2M1 values, the probabilities are chosen so the expected rounded value equals the original value. Transformer Engine enables this for gradient quantization. “Unbiased in expectation” does not mean an individual rounding is exact, and the added randomness has reproducibility implications.</p>
<p>A <strong>Random Hadamard Transform</strong> is an orthogonal mixing operation applied to selected operands before quantization. Its purpose here is to spread isolated large values across coordinates so that one outlier is less likely to set an unhelpful scale for an entire block. Transformer Engine currently documents this path for BF16 inputs and gradients used in selected weight-gradient operations (<a href="https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/features/low_precision_training/nvfp4/nvfp4.html">Transformer Engine documentation</a>).</p>
<p>The broader recipe also retains higher precision for sensitive layers or operations. In other words, “NVFP4 training” is mixed-precision training built around NVFP4 matrix multiplications, not a claim that the whole training state and every computation use four bits.</p>
<p>NVIDIA reports pretraining a 12-billion-parameter model for 10 trillion tokens with loss and downstream accuracy comparable to an FP8 baseline using this combined method (<a href="https://arxiv.org/abs/2509.25149">NVFP4 training paper</a>). This is evidence for the model, data, schedule, hardware, and recipe in that experiment. It is not evidence that naive E2M1 conversion works for arbitrary architectures, or that every model will match FP8.</p>
</section>
<section id="hardware-support-still-needs-a-full-sentence" class="level2">
<h2 class="anchored" data-anchor-id="hardware-support-still-needs-a-full-sentence">Hardware support still needs a full sentence</h2>
<p>NVFP4 is designed for native block-scaled matrix operations on NVIDIA Blackwell Tensor Cores. But “Blackwell supports NVFP4” leaves out which operation, device, software path, and tensor are involved.</p>
<p>Transformer Engine 2.16 documents NVFP4 training on compute capabilities SM100 and SM103, and inference on SM100 or later. It exposes the recipe to PyTorch and JAX through <code>NVFP4BlockScaling</code>, while also documenting layout and distributed-training constraints (<a href="https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/features/low_precision_training/nvfp4/nvfp4.html">Transformer Engine documentation</a>). These are versioned software conditions, not permanent properties of every Blackwell deployment.</p>
<p>Inference has at least two separate quantization targets:</p>
<ol type="1">
<li><strong>Weights and activations</strong> are operands used by model computations.</li>
<li>The <strong>KV cache</strong> stores attention keys and values from earlier tokens so autoregressive decoding does not recompute them.</li>
</ol>
<p>TensorRT-LLM’s quantization matrix, inspected in the source package after its 2026-07-14 update, listed NVFP4 weights and activations on Blackwell SM100/103 and SM120. Its NVFP4 KV-cache support was narrower, limited to SM100/103, and the documented KV-cache combination required FP8 weights and activations (<a href="https://nvidia.github.io/TensorRT-LLM/latest/features/quantization.html">TensorRT-LLM quantization guide</a>). Because this support changes with releases, the live matrix should be checked for the exact TensorRT-LLM version before deployment.</p>
<p>Training GEMMs, weight-and-activation inference, and KV-cache compression are therefore three different decisions. Each needs an exact compute capability, framework version, model and layout, calibration or training path, and quality test.</p>
</section>
<section id="where-speed-and-accuracy-claims-can-fail" class="level2">
<h2 class="anchored" data-anchor-id="where-speed-and-accuracy-claims-can-fail">Where speed and accuracy claims can fail</h2>
<p>Native Tensor Core execution creates a credible path to higher matrix throughput and lower operand bandwidth. End-to-end performance improves only if supported matrix multiplications dominate the workload.</p>
<p>Quantization and scale handling consume work. Attention, communication, unsupported operators, memory bandwidth, small tensor shapes, and irregular layouts may dominate instead. A faster GEMM can simply move the bottleneck elsewhere. Peak hardware throughput is not a substitute for measuring prefill latency, decode latency, generation throughput, or training step time at the intended shapes and batch sizes.</p>
<p>Accuracy needs a separate boundary. Finer local scaling can reduce the numerical error introduced while representing a tensor. Model quality also depends on where quantization is applied, how calibration is performed, whether training adapts to it, which computations stay at higher precision, and how errors propagate through the network. Lower error in a toy tensor does not prove equal language-model accuracy, much less equal behavior on rare or sensitive cases.</p>
<p>The strongest large-scale result cited above was produced by NVIDIA for a particular pretraining experiment. Independent behavior on a target model, task distribution, context length, and deployment stack remains unknown until it is measured.</p>
</section>
<section id="a-decision-test-for-an-nvfp4-claim" class="level2">
<h2 class="anchored" data-anchor-id="a-decision-test-for-an-nvfp4-claim">A decision test for an NVFP4 claim</h2>
<p>When a benchmark or deployment proposal says “NVFP4,” I would unpack it in this order:</p>
<ol type="1">
<li><strong>Name the quantized object.</strong> Is it weights, activations, training GEMM operands, or the KV cache?</li>
<li><strong>Count the representation.</strong> Include local and global scales, padding, duplicate layouts, and any higher-precision components instead of counting only E2M1 payload bits.</li>
<li><strong>Verify the supported path.</strong> Record the compute capability, software versions, model family, tensor layout, and allowed precision combination.</li>
<li><strong>Describe how the artifact was made.</strong> Distinguish a supported pre-quantized checkpoint, calibrated post-training quantization, and training with the full mixed-precision recipe.</li>
<li><strong>Validate quality on the intended workload.</strong> Compare with BF16 or FP8 on held-out tasks, long contexts, outlier-heavy inputs, and rare behaviors that matter.</li>
<li><strong>Measure the system boundary.</strong> Report checkpoint size, peak memory, KV-cache size where relevant, end-to-end latency or step time, throughput, and quality under the same conditions.</li>
</ol>
<p>If those questions are unanswered, “4-bit” tells us the width of the element payload and little more.</p>
<p>NVFP4’s useful idea is not merely squeezing a floating-point value into four bits. It is giving the same tiny E2M1 codebook different local units, then putting those units into a tensor-wide range. That hierarchy can represent tensors more closely than a coarser power-of-two scheme, and Blackwell hardware can execute supported operations natively. The remaining memory, accuracy, and speed claims belong to the whole mixed-precision system and the workload running through it, not to the four payload bits alone.</p>


</section>

 ]]></description>
  <category>machine-learning</category>
  <guid>https://nakranivaibhav.github.io/posts/nvfp4-scaling-system.html</guid>
  <pubDate>Fri, 17 Jul 2026 18:30:00 GMT</pubDate>
</item>
<item>
  <title>A Visual Guide to Causal, Bidirectional, and Hybrid Inference in Video Generation</title>
  <dc:creator>Vaibhav Nakrani</dc:creator>
  <link>https://nakranivaibhav.github.io/posts/causal-bidirectional-hybrid-video-inference.html</link>
  <description><![CDATA[ 




<style>
.schedule { display:grid; grid-template-columns:6.5rem repeat(4,minmax(4.7rem,1fr)); gap:.42rem; margin:1.2rem 0 .65rem; }
.schedule > div { min-height:3.2rem; padding:.55rem .35rem; border:1px solid #cbd5e1; border-radius:.45rem; display:flex; align-items:center; justify-content:center; text-align:center; font-size:.9rem; }
.schedule .axis { border:0; justify-content:flex-end; color:#475569; font-weight:600; }
.schedule .head { min-height:auto; border:0; font-weight:700; color:#334155; }
.unknown { background:#f1f5f9; color:#64748b; }
.rough { background:#dbeafe; border-color:#60a5fa !important; color:#1e3a8a; }
.mid { background:#bfdbfe; border-color:#3b82f6 !important; color:#1e3a8a; }
.ready { background:#2563eb; border-color:#1d4ed8 !important; color:white; }
.current { outline:3px solid #f59e0b; outline-offset:1px; }
.legend { display:flex; flex-wrap:wrap; gap:1rem; margin:.4rem 0 1.1rem; color:#475569; font-size:.9rem; }
.swatch { display:inline-block; width:.8rem; height:.8rem; margin-right:.3rem; border-radius:.18rem; vertical-align:-.05rem; }
.checkpoint { border-left:4px solid #f59e0b; background:#fffbeb; padding:.8rem 1rem; margin:1.2rem 0; }
</style>
<p>“Causal” and “bidirectional” sound like labels attached permanently to a video model. One model moves left to right; another sees the whole clip. Pick an architecture, inherit its behavior.</p>
<p>That picture is useful, but incomplete. A generator also needs a policy for deciding <strong>which parts of a noisy video may exchange information, and when</strong>. Once we look at that policy, causal and bidirectional generation become endpoints of a larger scheduling problem.</p>
<p>I want to build that idea with four chunks of an imaginary video. No benchmark is hiding in the diagrams. They show information availability, not speed or quality measurements.</p>
<section id="start-with-what-each-position-is-allowed-to-know" class="level2">
<h2 class="anchored" data-anchor-id="start-with-what-each-position-is-allowed-to-know">Start with what each position is allowed to know</h2>
<p>Suppose we split a video into four consecutive chunks: A, B, C, and D. A <strong>chunk</strong> is simply a group of neighboring video positions. It might contain one frame, several frames, or latent video tokens rather than pixels.</p>
<p>In causal generation, a position may use an available prefix: the current chunk and earlier chunks, but not later chunks that have not been generated. When producing C, information from A and B can flow forward; D cannot flow backward.</p>
<p>In bidirectional generation, positions across the clip may condition on one another while the clip is being refined. A rough state in A can influence D, and D can influence A. “Bidirectional” describes information flow, not playback direction. The finished video still plays from beginning to end.</p>
<p>The tempting model is therefore:</p>
<ul>
<li>causal means <code>A → B → C → D</code>;</li>
<li>bidirectional means <code>{A, B, C, D}</code> are refined together.</li>
</ul>
<p>Nothing is wrong with those endpoints. The mistake is assuming there can be no useful schedule between them.</p>
</section>
<section id="there-are-two-kinds-of-time-on-the-page" class="level2">
<h2 class="anchored" data-anchor-id="there-are-two-kinds-of-time-on-the-page">There are two kinds of time on the page</h2>
<p>Diffusion adds a second clock. <strong>Video time</strong> tells us whether a chunk appears early or late in the clip. <strong>Denoising time</strong> tells us how far a noisy representation has been refined toward a clean one.</p>
<p>The grid below keeps these clocks separate. Columns run through the video. Rows move from a rough, noisy state toward a ready state. A row is not another moment in the story; it is another stage in the computation.</p>
<div class="schedule">
<div>

</div>
<div class="head">
A<br>earliest
</div>
<div class="head">
B
</div>
<div class="head">
C
</div>
<div class="head">
D<br>latest
</div>
<div class="axis">
rough state
</div>
<div class="rough">
A · rough
</div>
<div class="rough">
B · rough
</div>
<div class="rough">
C · rough
</div>
<div class="rough">
D · rough
</div>
<div class="axis">
more refined
</div>
<div class="mid">
A · refining
</div>
<div class="mid">
B · refining
</div>
<div class="mid">
C · refining
</div>
<div class="mid">
D · refining
</div>
<div class="axis">
ready state
</div>
<div class="ready">
A · ready
</div>
<div class="ready">
B · ready
</div>
<div class="ready">
C · ready
</div>
<div class="ready">
D · ready
</div>
</div>
<div class="legend">
<p><span><span class="swatch rough"></span>rough</span> <span><span class="swatch mid"></span>partly refined</span> <span><span class="swatch ready"></span>ready</span></p>
</div>
<p>This separation matters because a schedule can change along either axis. It can group more or fewer neighboring video positions, and it can change those groups as denoising progresses.</p>
<div class="checkpoint">
<p><strong>Pause and predict:</strong> must D be completely unknown until A is ready? A strictly causal schedule says yes. A hybrid schedule does not have to.</p>
</div>
</section>
<section id="the-rigid-endpoints" class="level2">
<h2 class="anchored" data-anchor-id="the-rigid-endpoints">The rigid endpoints</h2>
<p>The tabs show two idealized reveal orders. Step through them from left to right. “Available” means the scheduler allows that region to participate; it does not mean the model has rendered a viewable frame at that instant.</p>
<div class="tabset-margin-container"></div><div class="panel-tabset">
<ul class="nav nav-tabs"><li class="nav-item"><a class="nav-link active" id="tabset-1-1-tab" data-bs-toggle="tab" data-bs-target="#tabset-1-1" aria-controls="tabset-1-1" aria-selected="true" href="">Bidirectional</a></li><li class="nav-item"><a class="nav-link" id="tabset-1-2-tab" data-bs-toggle="tab" data-bs-target="#tabset-1-2" aria-controls="tabset-1-2" aria-selected="false" href="">Causal · after A</a></li><li class="nav-item"><a class="nav-link" id="tabset-1-3-tab" data-bs-toggle="tab" data-bs-target="#tabset-1-3" aria-controls="tabset-1-3" aria-selected="false" href="">Causal · after B</a></li></ul>
<div class="tab-content">
<div id="tabset-1-1" class="tab-pane active" aria-labelledby="tabset-1-1-tab">
<div class="schedule">
<div>

</div>
<div class="head">
A
</div>
<div class="head">
B
</div>
<div class="head">
C
</div>
<div class="head">
D
</div>
<div class="axis">
1 · rough
</div>
<div class="rough current">
available
</div>
<div class="rough current">
available
</div>
<div class="rough current">
available
</div>
<div class="rough current">
available
</div>
<div class="axis">
2 · refine
</div>
<div class="mid">
together
</div>
<div class="mid">
together
</div>
<div class="mid">
together
</div>
<div class="mid">
together
</div>
<div class="axis">
3 · ready
</div>
<div class="ready">
ready
</div>
<div class="ready">
ready
</div>
<div class="ready">
ready
</div>
<div class="ready">
ready
</div>
</div>
<p>All four chunks enter the computation together and can help establish a global plan. The price of this simple picture is that the full working set must be handled together; it also does not naturally emit a completed prefix while later chunks remain absent.</p>
</div>
<div id="tabset-1-2" class="tab-pane" aria-labelledby="tabset-1-2-tab">
<div class="schedule">
<div>

</div>
<div class="head">
A
</div>
<div class="head">
B
</div>
<div class="head">
C
</div>
<div class="head">
D
</div>
<div class="axis">
rough
</div>
<div class="rough">
available
</div>
<div class="unknown">
not available
</div>
<div class="unknown">
not available
</div>
<div class="unknown">
not available
</div>
<div class="axis">
refine
</div>
<div class="mid">
refining
</div>
<div class="unknown">
not available
</div>
<div class="unknown">
not available
</div>
<div class="unknown">
not available
</div>
<div class="axis">
ready
</div>
<div class="ready current">
A ready
</div>
<div class="unknown">
not available
</div>
<div class="unknown">
not available
</div>
<div class="unknown">
not available
</div>
</div>
<p>A is completed without future chunks. This makes a finished prefix possible, but A cannot use evidence that will appear only in C or D.</p>
</div>
<div id="tabset-1-3" class="tab-pane" aria-labelledby="tabset-1-3-tab">
<div class="schedule">
<div>

</div>
<div class="head">
A
</div>
<div class="head">
B
</div>
<div class="head">
C
</div>
<div class="head">
D
</div>
<div class="axis">
rough
</div>
<div class="ready">
past context
</div>
<div class="rough">
available
</div>
<div class="unknown">
not available
</div>
<div class="unknown">
not available
</div>
<div class="axis">
refine
</div>
<div class="ready">
past context
</div>
<div class="mid">
refining
</div>
<div class="unknown">
not available
</div>
<div class="unknown">
not available
</div>
<div class="axis">
ready
</div>
<div class="ready">
A ready
</div>
<div class="ready current">
B ready
</div>
<div class="unknown">
not available
</div>
<div class="unknown">
not available
</div>
</div>
<p>B may use A as context. C and D still cannot help. Repeating the pattern eventually completes the clip from left to right.</p>
</div>
</div>
</div>
<p>These diagrams deliberately exaggerate the contrast. Real systems differ in caching, attention masks, chunk size, and what counts as “ready.” The grid gives us a vocabulary for asking about those details instead of treating the labels as complete specifications.</p>
</section>
<section id="a-schedule-can-live-between-the-endpoints" class="level2">
<h2 class="anchored" data-anchor-id="a-schedule-can-live-between-the-endpoints">A schedule can live between the endpoints</h2>
<p><a href="https://arxiv.org/abs/2607.03509">Flex-Forcing</a> makes chunking flexible over both the temporal dimension and denoising steps. The same video diffusion model can operate under bidirectional, autoregressive, and hybrid generation regimes. Its <a href="https://research.nvidia.com/labs/genair/flex-forcing/">project page</a> describes larger chunks as closer to global bidirectional planning and smaller chunks as closer to autoregressive streaming.</p>
<p>Here is a conceptual four-chunk hybrid. It is an explanatory sketch, not the paper’s exact tensor layout or a reproduced experimental configuration.</p>
<div class="tabset-margin-container"></div><div class="panel-tabset">
<ul class="nav nav-tabs"><li class="nav-item"><a class="nav-link active" id="tabset-2-1-tab" data-bs-toggle="tab" data-bs-target="#tabset-2-1" aria-controls="tabset-2-1" aria-selected="true" href="">1 · rough plan</a></li><li class="nav-item"><a class="nav-link" id="tabset-2-2-tab" data-bs-toggle="tab" data-bs-target="#tabset-2-2" aria-controls="tabset-2-2" aria-selected="false" href="">2 · finish A</a></li><li class="nav-item"><a class="nav-link" id="tabset-2-3-tab" data-bs-toggle="tab" data-bs-target="#tabset-2-3" aria-controls="tabset-2-3" aria-selected="false" href="">3 · finish B</a></li><li class="nav-item"><a class="nav-link" id="tabset-2-4-tab" data-bs-toggle="tab" data-bs-target="#tabset-2-4" aria-controls="tabset-2-4" aria-selected="false" href="">4 · complete</a></li></ul>
<div class="tab-content">
<div id="tabset-2-1" class="tab-pane active" aria-labelledby="tabset-2-1-tab">
<div class="schedule">
<div>

</div>
<div class="head">
A
</div>
<div class="head">
B
</div>
<div class="head">
C
</div>
<div class="head">
D
</div>
<div class="axis">
rough
</div>
<div class="rough current">
global hint
</div>
<div class="rough current">
global hint
</div>
<div class="rough current">
global hint
</div>
<div class="rough current">
global hint
</div>
<div class="axis">
refine
</div>
<div class="unknown">
pending
</div>
<div class="unknown">
pending
</div>
<div class="unknown">
pending
</div>
<div class="unknown">
pending
</div>
<div class="axis">
ready
</div>
<div class="unknown">
pending
</div>
<div class="unknown">
pending
</div>
<div class="unknown">
pending
</div>
<div class="unknown">
pending
</div>
</div>
<p>All chunks receive a rough global context. Later positions are not finished, but they are not entirely absent either. This is the bidirectional part of the mental model.</p>
</div>
<div id="tabset-2-2" class="tab-pane" aria-labelledby="tabset-2-2-tab">
<div class="schedule">
<div>

</div>
<div class="head">
A
</div>
<div class="head">
B
</div>
<div class="head">
C
</div>
<div class="head">
D
</div>
<div class="axis">
rough
</div>
<div class="rough">
global hint
</div>
<div class="rough">
global hint
</div>
<div class="rough">
global hint
</div>
<div class="rough">
global hint
</div>
<div class="axis">
refine
</div>
<div class="mid">
refined
</div>
<div class="rough">
planned
</div>
<div class="rough">
planned
</div>
<div class="rough">
planned
</div>
<div class="axis">
ready
</div>
<div class="ready current">
A ready
</div>
<div class="unknown">
pending
</div>
<div class="unknown">
pending
</div>
<div class="unknown">
pending
</div>
</div>
<p>A can now be completed while the rough plan for B–D remains available. Fine synthesis moves forward even though coarse planning did not.</p>
</div>
<div id="tabset-2-3" class="tab-pane" aria-labelledby="tabset-2-3-tab">
<div class="schedule">
<div>

</div>
<div class="head">
A
</div>
<div class="head">
B
</div>
<div class="head">
C
</div>
<div class="head">
D
</div>
<div class="axis">
rough
</div>
<div class="rough">
global hint
</div>
<div class="rough">
global hint
</div>
<div class="rough">
global hint
</div>
<div class="rough">
global hint
</div>
<div class="axis">
refine
</div>
<div class="ready">
past context
</div>
<div class="mid">
refined
</div>
<div class="rough">
planned
</div>
<div class="rough">
planned
</div>
<div class="axis">
ready
</div>
<div class="ready">
A ready
</div>
<div class="ready current">
B ready
</div>
<div class="unknown">
pending
</div>
<div class="unknown">
pending
</div>
</div>
<p>B uses the completed prefix and the wider plan. Call this “autoregressive synthesis within the plan”: left-to-right completion—not that every internal operation must be strictly causal.</p>
</div>
<div id="tabset-2-4" class="tab-pane" aria-labelledby="tabset-2-4-tab">
<div class="schedule">
<div>

</div>
<div class="head">
A
</div>
<div class="head">
B
</div>
<div class="head">
C
</div>
<div class="head">
D
</div>
<div class="axis">
rough
</div>
<div class="rough">
global hint
</div>
<div class="rough">
global hint
</div>
<div class="rough">
global hint
</div>
<div class="rough">
global hint
</div>
<div class="axis">
refine
</div>
<div class="ready">
refined
</div>
<div class="ready">
refined
</div>
<div class="ready">
refined
</div>
<div class="ready">
refined
</div>
<div class="axis">
ready
</div>
<div class="ready">
A ready
</div>
<div class="ready">
B ready
</div>
<div class="ready">
C ready
</div>
<div class="ready current">
D ready
</div>
</div>
<p>The result is still an ordinary video. What changed was the reveal order of intermediate information.</p>
</div>
</div>
</div>
<p>The paper’s claim is broader than this toy schedule. Flex-Forcing defines flexible chunking jointly over video positions and denoising steps, supports chunk sizes chosen for different device budgets, and permits any-order, any-timestep autoregressive generation without a strict causal constraint. The authors also report better video quality, long-video stability, and faster inference than rigid-schedule baselines across their experiments (<a href="https://arxiv.org/abs/2607.03509">paper abstract</a>; <a href="https://research.nvidia.com/labs/genair/flex-forcing/">project results</a>). Those are results for the evaluated configurations, not a guarantee that every hybrid schedule improves every metric.</p>
<p>The project page offers one concrete example of why the schedule itself matters. For a five-second video represented by 21 latent frames, it reports that a coarse-to-fine chunk configuration <code>[9, 6, 6]</code> lies on its measured quality–speed frontier and outperforms uniform <code>[7, 7, 7]</code> and its Self-Forcing baseline on the displayed measures. The numbers describe chunk configurations in that experiment; they are not denoising-step counts or a universal recipe.</p>
</section>
<section id="the-policy-changes-the-questions-we-should-ask" class="level2">
<h2 class="anchored" data-anchor-id="the-policy-changes-the-questions-we-should-ask">The policy changes the questions we should ask</h2>
<p>Once generation becomes a schedule, “Is this model causal?” is too coarse. More useful questions are:</p>
<ul>
<li>Which video positions are grouped into a chunk?</li>
<li>Can rough future positions influence the current chunk?</li>
<li>Does chunk size change along the denoising trajectory?</li>
<li>When is the first usable prefix available?</li>
<li>Which intermediate states must remain in memory?</li>
<li>Can a later chunk be edited without rebuilding the whole sequence?</li>
</ul>
<p>Different answers can trade global planning, streaming behavior, memory use, and fine-grained control. The direction of those trade-offs is plausible from the information-flow pattern, but their size requires matched experiments on the same model, outputs, hardware, and quality target. A grid cannot produce a latency number.</p>
<p>Flex-Forcing also demonstrates why “inference policy” is a useful phrase rather than just a metaphor. Its project page describes editing at low-level refinement steps while keeping high-level planning fixed, and allowing a chunk to attend to both past and future tokens. The schedule determines which information can move and which parts can change.</p>
</section>
<section id="lingbot-exposes-a-different-kind-of-inference-policy" class="level2">
<h2 class="anchored" data-anchor-id="lingbot-exposes-a-different-kind-of-inference-policy">LingBot exposes a different kind of inference policy</h2>
<p>LingBot-Video is useful here as a contrast, not as another Flex-Forcing implementation. Its public <a href="https://github.com/robbyant/lingbot-video">repository</a> and <a href="https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b">model card</a> expose a user-facing workflow with prompt rewriting, base video generation, an optional refiner, and multiple serving paths. The documented MoE text-to-video command uses a 40-step base stage followed by an 8-step refiner stage.</p>
<p>That sequence is also an inference policy: it decides how a prompt is prepared and how generation is divided across two denoising stages. But it operates along different axes from Flex-Forcing’s joint temporal-and-denoising chunk schedule. <code>40 + 8</code> does not tell us whether information flows causally or bidirectionally across video chunks, and the public LingBot materials cited here do not claim Flex-Forcing’s mechanism.</p>
<p>The comparison therefore stops at a shared lesson: a model’s weights and architecture do not completely describe the behavior a user receives. Scheduling, refinement, conditioning, and serving choices are part of the generation path. The two systems expose different choices, so their counters cannot be placed on one speed scale without a matched evaluation.</p>
</section>
<section id="a-compact-way-to-explain-it-back" class="level2">
<h2 class="anchored" data-anchor-id="a-compact-way-to-explain-it-back">A compact way to explain it back</h2>
<p>Try drawing a grid with video position on the horizontal axis and denoising progress on the vertical axis.</p>
<p>A rigid bidirectional schedule refines a large region together. A rigid causal schedule completes regions from left to right, leaving future regions unavailable. A flexible schedule can make rough information available across a wide region, then complete finer regions in an autoregressive order. Flex-Forcing turns that two-axis chunking pattern into a training and inference framework for one model.</p>
<p>This mental model establishes <strong>how</strong> causal and bidirectional behavior can become scheduling choices. It does not establish the best schedule for a deployment, the memory required by an implementation, or a speed comparison with LingBot. Those answers need released implementations or matched measurements.</p>
<p>The next time a video model is called “causal” or “bidirectional,” the useful follow-up is not “which one is better?” It is: <strong>what information becomes available at each point in video time and denoising time?</strong></p>


</section>

 ]]></description>
  <category>machine-learning</category>
  <guid>https://nakranivaibhav.github.io/posts/causal-bidirectional-hybrid-video-inference.html</guid>
  <pubDate>Fri, 10 Jul 2026 18:30:00 GMT</pubDate>
</item>
<item>
  <title>Diffusion Without the Hand-Waving</title>
  <dc:creator>Vaibhav Nakrani</dc:creator>
  <link>https://nakranivaibhav.github.io/posts/diffusion-without-the-hand-waving.html</link>
  <description><![CDATA[ 




<p>The cleanest way to understand diffusion is to build it up one step at a time. Here’s the version that finally made it click.</p>
<p>Start with the goal. You want to turn random noise into a realistic image.</p>
<p>Now imagine you had a magic function. You show it any image and it tells you how realistic that image is. A clear photo scores high. Pure static scores low.</p>
<p>The useful part isn’t the score itself. It’s the gradient of that score with respect to the pixels. That gradient tells you exactly which pixels to change, and in which direction, to make the image look a little more real. Follow it over and over and noise slowly turns into a real image.</p>
<p>The problem is nobody can hand you that function.</p>
<p>The trick that makes the whole field work. You can learn that gradient without ever building the function. Take a real image. Add a known amount of noise to it. Train a network to predict the noise you added. That predicted noise is, give or take a scaling factor, the gradient you wanted. Subtract a bit of it and the image moves toward “more realistic.”</p>
<p>The network does not learn to clean up images. It learns which direction points toward real data. Denoising is just what that looks like from the outside. The proper name for that direction is the score, but you can keep thinking of it as a gradient.</p>
<p>Training is boring on purpose. Take images, add a random amount of noise to each, ask the model to predict the noise, penalize it with mean squared error. Repeat millions of times. You also feed in how much noise was added, so the model knows how messy the picture currently is. ( the magical timestep )</p>
<p>Generation just runs it backward. Start from pure noise. Ask the model which direction is more realistic. Take a small step that way. Ask again. Step again. Repeat until an image appears. The steps are small because the model’s guess is only approximate, especially early on when the input looks nothing like real data.</p>
<p>That stepping strategy is the sampler. DDPM, DDIM, Euler, DPM-Solver. They teach the model nothing new. They’re just different ways of walking downhill. The right mental model is optimizers. SGD versus Adam. Same loss surface, different path, different speed. A good sampler gets you there in 20 steps instead of 1000.</p>
<p>So far this was on raw pixels. In practice you don’t work on pixels, and you also want to control what gets generated. Two extra pieces handle that, and both are trained ahead of time and then frozen.</p>
<p>The VAE is just a compressor. It’s trained separately, before any diffusion, with one job: squash an image into a small latent and decode it back faithfully. It never sees noise and knows nothing about diffusion. Once trained, you freeze it. All the diffusion now happens on these small latents instead of big images, which is what makes it cheap.</p>
<div class="quarto-figure quarto-figure-center">
<figure class="figure">
<p><img src="https://nakranivaibhav.github.io/assets/images/diffusion-without-the-hand-waving/im4.png" class="img-fluid figure-img"></p>
<figcaption>vae</figcaption>
</figure>
</div>
<p>The text encoder is the text half of a CLIP model, also trained earlier and frozen. Its job is to turn a caption into an embedding. That embedding is how the prompt steers generation. It enters the U-Net through cross-attention, so the gradient now points toward realistic and matches this caption.</p>
<p>Here’s the actual training flow, start to finish: - Take an (image, caption) pair. - Caption goes through the frozen text encoder. Out comes a text embedding. - Image goes through the frozen VAE encoder. Out comes a clean latent. - Pick a random timestep. Blend the clean latent with random noise according to that timestep. Low timestep is mostly clean latent, high timestep is mostly noise. This is a one-shot formula, you do not iterate here. - That noisy latent, the timestep, and the text embedding go into the U-Net. - The U-Net predicts the noise that was added. - Loss is MSE between the real noise and the predicted noise. Only the U-Net updates. VAE and text encoder stay frozen.</p>
<div class="quarto-figure quarto-figure-center">
<figure class="figure">
<p><img src="https://nakranivaibhav.github.io/assets/images/diffusion-without-the-hand-waving/im5.png" class="img-fluid figure-img"></p>
<figcaption>diffusion training process</figcaption>
</figure>
</div>
<p>Step 4 is the thing people get stuck on. You do not add noise to the image. You encode the image to a latent first, then add noise to the latent with a formula. The clean latent is the target, the noise is what the model predicts.</p>
<p>Generation is the same pipeline in reverse:</p>
<ul>
<li>Caption through the text encoder, get the embedding.</li>
<li>Start from a pure noise latent.</li>
<li>Each step, the U-Net predicts the noise from the current latent, the timestep, and the text embedding. The sampler uses that to step to a slightly cleaner latent. Repeat.</li>
<li>You end with a clean latent. Run it through the VAE decoder to get the final image.</li>
</ul>
<div class="quarto-figure quarto-figure-center">
<figure class="figure">
<p><img src="https://nakranivaibhav.github.io/assets/images/diffusion-without-the-hand-waving/im2.png" class="img-fluid figure-img"></p>
<figcaption>forward pass diffusion</figcaption>
</figure>
</div>
<p>And the network itself is swappable. Replace the U-Net with something else and you get modern DiT backbones. Same training flow, same everything else.</p>
<p>So the whole thing is still three pieces: A direction that says which way is more realistic. A network that learns it by predicting noise. A sampler that walks that direction from noise to image.</p>
<p>The VAE and text encoder are just the frozen helpers that make it cheap and controllable. Everything else is detail.</p>
<p>Optional: both the vae and the text encoder can be fine-tuned when tuning these models. text encoder tuning is more beneficial generally than vae from my expreiments.</p>
<p><em>Inspired by <a href="https://course.fast.ai/Lessons/lesson9.html">fast.ai Practical Deep Learning, Lesson 9</a>. If you want the long version, the lesson has the full lecture video walking through all of this end to end. Grab a coffee, it’s worth the couple of hours.</em></p>



 ]]></description>
  <category>machine-learning</category>
  <guid>https://nakranivaibhav.github.io/posts/diffusion-without-the-hand-waving.html</guid>
  <pubDate>Sat, 27 Jun 2026 18:30:00 GMT</pubDate>
  <media:content url="https://nakranivaibhav.github.io/assets/images/diffusion-without-the-hand-waving/im1.png" medium="image" type="image/png" height="108" width="144"/>
</item>
</channel>
</rss>
