[{"content":"Open any tutorial on speech processing and you will meet the spectrogram in the first five minutes — usually with a hand-wave: \u0026ldquo;we apply the Fourier transform in sliding windows\u0026rdquo;. Open a math textbook, and you will find the Fourier series in its full rigor — with no hint of why an ML engineer should care. The road between these two points is almost never walked end to end: every explanation starts somewhere in the middle. This series of three posts walks the whole road: from air pressure and a guitar string, through sampling and quantization, through the Fourier series and the DFT, to the spectrogram — the picture that speech models actually consume. In this first part we build the core machinery: we follow the sound wave into the computer and end with the DFT derived — honestly, delta functions and all — from the Fourier series itself.\nAlong the way we will meet several different creatures that all answer to the name \u0026ldquo;Fourier\u0026rdquo;: the Fourier series, the Fourier transform, its discrete-time cousin, and the discrete Fourier transform. They form a tidy 2×2 family — and untangling their family relationships is a story of its own, which gets a separate post (and more): for now, one picture as a teaser.\n(a teaser of a separate post to come — in this series we walk the bottom row: from the Fourier series to the discrete Fourier transform)\nWhat is sound The air around us is filled with molecules. Pull a guitar string — it creates a vibration that travels through them as alternating zones of compression and rarefaction: a pressure wave. A microphone senses exactly this: variations of pressure relative to the atmospheric baseline. Plot that pressure against time — and you get the most honest picture of sound there is: the waveform.\n💡 Sound is a mechanical wave — it needs a medium. In the vacuum of space there is nothing to compress, which is why we can see the Sun but cannot hear it.\nAnalog → digital The microphone\u0026rsquo;s output is an analog signal — and the word is more literal than it sounds: the voltage on the microphone\u0026rsquo;s wire is an analogue of the air pressure, one physical quantity tracing the shape of another. Nothing has been measured yet; the signal has merely changed its carrier — from pressure to voltage — and is still continuous in time and in value.\nA computer, however, needs numbers, and the Analog-to-Digital Converter (ADC) produces them by discretizing along both axes:\nsampling — measure the signal at regular moments, $f_s$ times per second (time discretization); quantization — round each measurement to the nearest level of a fixed grid (amplitude discretization). The result — a sequence of integers at a fixed rate — is Pulse-Code Modulation (PCM), the format inside every .wav file. Two numbers fully describe the grid: the sampling rate (e.g. 44.1 kHz) and the bit depth (e.g. 16 bits).\nWhy the waveform is not enough The waveform is honest but unhelpful. One second of CD-quality audio is 44,100 numbers — and no individual number tells you anything about what you hear. Is this a male or a female voice? Which note is the guitar playing? Is there a hum from the power line polluting the recording? Staring at pressure values will not answer any of these questions.\nFourier\u0026rsquo;s idea flips the axis. A complex signal can be decomposed into a sum of simple oscillations — the way a chord can be decomposed into individual notes. Instead of asking \u0026ldquo;what is the pressure at each moment of time?\u0026rdquo;, we ask \u0026ldquo;how much of each frequency does this signal contain?\u0026rdquo; The answer lives on the frequency axis rather than the time axis, and it is called the spectrum.\nIf decomposing sound into frequencies feels like an arbitrary idea, think of a piano. Pressing a key produces an oscillation at a known frequency — the keyboard is literally a frequency axis, laid out left to right. Playing music is easy in this direction: choose which keys to press, how hard, and when — frequencies in, melody out. Fourier analysis asks for the reverse: given only the recorded sound, can we recover which keys were pressed and how hard? Most of this post is the machinery that turns this \u0026ldquo;can we?\u0026rdquo; into \u0026ldquo;here is how\u0026rdquo;.\nTwo things make this representation valuable in practice. First, it is compact: a long sequence of time samples collapses into a handful of meaningful frequency components. The wiggly curve below takes hundreds of numbers to store — and just three frequency components to describe:\nSecond, it is actionable: frequencies can be inspected and edited. Say a recording picked up the 50 Hz hum of the power line. In the time domain the hum is smeared over every sample and there is nothing to grab; in the frequency domain it is one column. Transform, erase that column, transform back — the melody survives untouched, the hum is gone:\nThe Fourier series Let us make the \u0026ldquo;sum of simple oscillations\u0026rdquo; idea precise. We start with the version that history started with, too: a periodic function.\nAny periodic function $x(t)$ with period $P$ that is absolutely integrable over $\\left[-\\frac{P}{2}, \\frac{P}{2}\\right]$ can be represented as a Fourier series:\n$$ x(t) = a_0 + \\sum_{n = 1}^\\infty \\left( a_n \\cos\\!\\left( 2 \\pi \\frac{n}{P} t \\right) + b_n \\sin\\!\\left( 2 \\pi \\frac{n}{P} t \\right) \\right) $$with the coefficients\n$$ \\begin{gathered} a_0 = \\frac{1}{P} \\int_{-P/2}^{P/2} x(t)\\, dt, \\\\ a_n = \\frac{2}{P} \\int_{-P/2}^{P/2} x(t) \\cos\\!\\left(2 \\pi \\frac{n}{P} t \\right) dt, \\qquad b_n = \\frac{2}{P} \\int_{-P/2}^{P/2} x(t) \\sin\\!\\left(2 \\pi \\frac{n}{P} t \\right) dt. \\end{gathered} $$The building blocks are sines and cosines whose frequencies are integer multiples of $\\frac{1}{P}$ — the harmonics of the base frequency. Nothing else is allowed: only oscillations that fit a whole number of times into the period. This restriction is easy to read past, and much of what follows grows out of it — so let us stare at it once, properly. A periodic function repeats: whatever happens on $[0, P]$ must glue seamlessly to its own copy on $[P, 2P]$. A harmonic with a whole number of oscillations arrives at the seam exactly where it started, so the copies join smoothly. An oscillation with a fractional count arrives somewhere else — and the glued copies tear:\nA pair $(a_n, b_n)$ at the same frequency is really one oscillation in disguise. Using the cosine-of-difference formula, the pair collapses into a single cosine with an amplitude and a shift:\n$$ \\begin{gathered} x(t) = a_0 + \\sum_{n = 1}^\\infty A_n \\cos\\!\\left( 2 \\pi \\frac{n}{P} t - \\phi_n\\right), \\\\ A_n = \\sqrt{a_n^2 + b_n^2}, \\qquad \\phi_n = \\operatorname{atan2}(b_n, a_n). \\end{gathered} $$Here $A_n$ is the magnitude of the $n$-th harmonic, $\\frac{n}{P}$ its frequency, and $\\phi_n$ its phase. This form is the one to keep in mind: a periodic signal is a recipe — this much of this frequency, shifted by this much.\n💡 Why sinusoids, of all things? Mathematics knows plenty of ways to decompose a function — why not, say, a Taylor series, with polynomials as the building blocks? Several reasons stack on top of each other. First, the shape of the data: sound is locally quasi-periodic — a guitar note, a vowel — and periodic building blocks describe such signals with a handful of coefficients, while a polynomial cannot even be periodic (it must run off to infinity) and needs ever more terms for every extra period. Second, and deeper: the physical world plays along. Put a speaker in one corner of a room, play a pure 440 Hz tone through it, and record with a microphone in the opposite corner. The walls reflect the sound, echoes pile on top of each other — yet the recording is still a 440 Hz tone: louder or quieter, shifted in time, but at the same frequency. The reason is one line of trigonometry: echoes are delayed, scaled copies, and a sum of sinusoids of one frequency — whatever their amplitudes and shifts — is again a sinusoid of that frequency.\n(Can a room be an anti-carpet and make a frequency louder? It cannot add energy — but it can concentrate it: when the copies arrive in phase, they add up constructively. That is resonance, and it is exactly why your voice blossoms in a tiled bathroom — at the room\u0026rsquo;s resonant frequencies, the echoes conspire in your favor.)\nThe same holds for a vibrating string, a microphone membrane, an electronic filter: none of them can invent new frequencies. (A distortion pedal can — precisely because it is not linear; that is what its \u0026ldquo;dirty\u0026rdquo; sound is made of.) That is why sound is built out of sinusoids in the first place, and the quasi-periodicity of the previous argument is not a lucky accident but physics. (Why this happens — and why it earns sinusoids the grand title of eigenfunctions of linear time-invariant systems — is a story of its own; for a very accessible account, see chapter 5 of The Scientist and Engineer\u0026rsquo;s Guide to DSP.) Third, stability. To compare fairly, fix the reference point — the moment you press \u0026ldquo;record\u0026rdquo; — and delay the signal past it by $\\tau$. The Fourier description barely notices: delaying the signal turns each harmonic $A_n \\cos(2\\pi \\frac{n}{P} t - \\phi_n)$ into $A_n \\cos(2\\pi \\frac{n}{P} (t - \\tau) - \\phi_n)$ — which is the same cosine with the same magnitude $A_n$, only its phase nudged to $\\phi_n + 2\\pi \\frac{n}{P} \\tau$. The Taylor description — the derivatives at the reference point — has no such luck: every new coefficient becomes a mixture of all the old higher-order ones. The cleanest example: $\\sin t$ and $\\cos t$ are one signal shifted by a quarter period, yet one has only odd-degree terms and the other only even-degree ones. A note sounds the same whenever you play it, and the magnitude spectrum agrees; the Taylor coefficients do not. (It is the room argument again, in disguise: a time shift leaves every sinusoid being itself, just rotated — while it smears each monomial $t^k$ across all the degrees below it.)\nThe exponential form One more rewrite, and the notation becomes so compact that every later formula in this series will use it. Euler\u0026rsquo;s formula,\n$$ e^{i t} = \\cos t + i \\sin t \\quad\\Longleftrightarrow\\quad \\cos t = \\tfrac{1}{2} \\left(e^{it} + e^{-it} \\right), $$lets us split every cosine into two complex exponentials — one rotating \u0026ldquo;forward\u0026rdquo; and one \u0026ldquo;backward\u0026rdquo;:\n$$ \\cos\\!\\left( 2 \\pi \\tfrac{n}{P} t - \\phi_n \\right) = \\tfrac{1}{2} e^{-i \\phi_n} e^{2 \\pi i \\frac{n}{P} t} + \\tfrac{1}{2} e^{i \\phi_n} e^{-2 \\pi i \\frac{n}{P} t}. $$Absorbing the magnitudes and phases into complex coefficients, the whole series collapses into a single sum:\n$$ x(t) = \\sum_{n = -\\infty}^\\infty c_n e^{2 \\pi i \\frac{n}{P} t}, \\qquad c_n = \\frac{1}{P} \\int_{-P/2}^{P/2} x(t)\\, e^{-2 \\pi i \\frac{n}{P} t}\\, dt. $$The set of coefficients $\\{c_n\\}$ is called the spectrum of the signal — the same word we met informally above, now with an exact meaning. Each $c_n$ is one complex number that stores both the magnitude and the phase of the $n$-th harmonic: $|c_n| = A_n / 2$ and $\\arg c_n = -\\phi_n$ for $n \\ge 1$.\n💡 Wait, negative frequencies? The sum now runs over all integers $n$, including negative ones — that is the price of the compact notation: each real oscillation split into a forward- and a backward-rotating exponential. For a real-valued signal the two halves are not independent: $c_{-n} = \\overline{c_n}$, so the negative-frequency half of the spectrum is a mirror image carrying no new information. Remember this — the same symmetry will resurface in the DFT and explain why half of its coefficients can be thrown away.\nFrom the series to the DFT This is the part most explanations skip. Textbooks stop at the Fourier series for nice continuous functions; engineering tutorials start from the DFT formula, presented as an axiom. But the road between them is short, honest, and worth walking — and it starts where honesty demands: by asking what a sampled signal even is as a mathematical object.\nAn honest model of a discrete signal Remember the Analog-to-Digital Converter from the digitization section, back before all the mathematics? It is about to become the protagonist again:\nThis is what it left us with: $N$ numbers $x(0), x(T), \\dots, x\\big((N-1)T\\big)$, measured every $T$ seconds. Can we recover a spectrum from these points? A spectrum means Fourier coefficients, and coefficients are integrals — so before computing anything, we owe the integral a well-definedness check. For a bounded function, the Riemann integral exists precisely when the function is continuous almost everywhere: its discontinuities must form a set of measure zero — this is Lebesgue\u0026rsquo;s integrability criterion.\nOur samples are not yet a function of continuous time, so let us complete them in the most straightforward way imaginable: keep the measured values at the grid points, put zero everywhere else,\n$$ \\tilde{x}(t) = \\begin{cases} x(nT), \u0026 t = nT, \\quad n = 0, \\dots, N-1, \\\\ 0, \u0026 t \\in [0, NT], \\; t \\neq nT. \\end{cases} $$\nDoes $\\tilde{x}$ pass the check? It is bounded; and it is discontinuous only at the grid points — finitely many on one period, and still just countably many after the periodic extension that the series insists on. A countable set has measure zero — continuous almost everywhere, check. The coefficients are well-defined, and we may integrate with a clear conscience:\n$$ c_k = \\frac{1}{P} \\int_{0}^{P} \\tilde{x}(t)\\, e^{-2 \\pi i \\frac{k}{P} t}\\, dt \\equiv 0 \\quad \\text{for every } k. $$Every single coefficient is zero — our signal has vanished from the mathematics. To see why, watch the Riemann sums converge: each spike gets trapped in a rectangle of finite height and ever-shrinking width, so its contribution — height times width — dies together with the mesh. Countably many spikes stand against a continuum of zeros, and the zeros win:\nThe integral honestly reports that $\\tilde{x}$ is almost everywhere indistinguishable from the zero function. The verdict is not against Fourier; it is against our model: \u0026ldquo;a value at a point and zero elsewhere\u0026rdquo; is the wrong mathematical object for a sample.\nSo what is a sample, really? Think about how the measurement is actually made. No instrument reads a value at an instant — a measurement takes some time $\\tau$: around every grid point $t = nT$ the device opens its gate, and for $\\tau$ seconds the signal pours in:\nWhat single number should the device report for its window? It saw not one value but a continuum of them — everything the signal did between $nT - \\tau/2$ and $nT + \\tau/2$. The natural answer is the average. And what is the average of a continuum of values? For a handful of numbers the average is \u0026ldquo;add them up, divide by how many\u0026rdquo;; for a continuum, the sum becomes an integral and the count becomes the length of the window:\n$$ \\hat{x}(nT) = \\frac{1}{\\tau} \\int_{nT - \\tau/2}^{nT + \\tau/2} x(t)\\, dt. $$It is useful to rewrite this average as an integral against a kernel. Let $r_\\tau(t)$ be the rectangular pulse of width $\\tau$ and unit height centered at zero:\nWith the scaled pulse as the kernel, the average becomes\n$$ \\hat{x}(nT) = \\int_{-\\infty}^{\\infty} x(t)\\, \\tfrac{1}{\\tau} r_\\tau(t - nT)\\, dt. $$The kernel $\\frac{1}{\\tau} r_\\tau$ cuts a column of width $\\tau$ out from under the graph of $x$ and reports its area, divided by the width — the average height of the graph inside the window. Remember this shape: it is about to have a famous limit. For finite $\\tau$ this is an estimate with an error: the signal keeps changing inside the window.\nNow improve the instrument. As $\\tau$ shrinks, the kernel becomes a rectangle ever narrower and ever taller — width $\\tau$, height $\\frac{1}{\\tau}$, area always exactly $1$ — and the estimate sharpens:\nThe limit of this process,\n$$ \\delta(t) = \\lim_{\\tau \\to 0} \\tfrac{1}{\\tau} r_\\tau(t), $$is the Dirac impulse: an \u0026ldquo;infinitely narrow, infinitely tall\u0026rdquo; spike of unit area. It is not a function in the classical sense — it is defined by what it does inside an integral, namely the sifting property, the $\\tau \\to 0$ limit of our averaging:\n$$ \\int_{-\\infty}^{\\infty} x(t)\\, \\delta(t - t_0)\\, dt = x(t_0) $$— and any integration limits that enclose $t_0$ work just as well, since the spike carries all of its area at the single point $t_0$.\nProof of the sifting property (a physicist's proof: we swap limits and integrals without asking permission) Step 0: a notation. For two real square-integrable signals (the space $L^2$ — where the Cauchy–Schwarz inequality guarantees the integral below is finite), their scalar (inner) product is\n$$ \\langle f(t), g(t) \\rangle = \\int_{-\\infty}^{\\infty} f(t)\\, g(t)\\, dt $$— the continuous cousin of the dot product of vectors: multiply the two signals pointwise, then add everything up (with the sum, as usual by now, becoming an integral; complex signals conjugate the second factor). In this notation, the sifting property reads $\\langle x(t), \\delta(t - t_0) \\rangle = x(t_0)$.\nStep 1: the spike at zero. Substitute the definition of $\\delta$ as the limit of our rectangles and move the limit outside the integral:\n$$ \\begin{aligned} \\langle x(t), \\delta(t) \\rangle \u0026= \\int_{-\\infty}^{\\infty} x(t)\\, \\delta(t)\\, dt = \\int_{-\\infty}^{\\infty} x(t) \\lim_{\\tau \\to 0} \\tfrac{1}{\\tau} r_\\tau(t)\\, dt \\\\ \u0026= \\lim_{\\tau \\to 0} \\int_{-\\infty}^{\\infty} x(t)\\, \\tfrac{1}{\\tau} r_\\tau(t)\\, dt. \\end{aligned} $$Step 2: the integral as a Riemann sum. Note what Step 1 bought us: a Riemann sum needs the values of the integrand at the grid points, and the original integrand contained $\\delta(t)$ — which has no value at zero to sample. After the trade, for each fixed $\\tau$ the integrand $x(t)\\, \\frac{1}{\\tau} r_\\tau(t)$ is an ordinary bounded function, and sampling it is legal. So write the integral as the limit of rectangle areas, choosing the mesh width to be the same $\\tau$ as in the pulse (a shortcut of the genre: the mesh refinement and the pulse shrinkage merge into a single limit):\n$$ \\int_{-\\infty}^{\\infty} f(t)\\, dt = \\lim_{\\tau \\to 0} \\sum_{n=-\\infty}^{\\infty} f(n\\tau)\\, \\tau. $$\nApplying this to our integrand, the $\\tau$ of the mesh cancels the $\\frac{1}{\\tau}$ of the kernel:\n$$ \\begin{aligned} \\langle x(t), \\delta(t) \\rangle \u0026= \\lim_{\\tau \\to 0} \\sum_{n=-\\infty}^{\\infty} x(n\\tau)\\, \\tfrac{1}{\\tau} r_\\tau(n\\tau)\\, \\tau \\\\ \u0026= \\lim_{\\tau \\to 0} \\sum_{n=-\\infty}^{\\infty} x(n\\tau)\\, r_\\tau(n\\tau). \\end{aligned} $$To be clear, this is where the genre\u0026rsquo;s quiet cheat actually lives. Honestly there are two independent limits — a mesh $h \\to 0$ that defines the integral at each fixed $\\tau$, and only then $\\tau \\to 0$ for the pulse. Setting $h = \\tau$ walks a single diagonal path through a double limit, and diagonals are not automatically legal. Here the diagonal happens to be right: with the mesh equal to the pulse width, the sum quietly replaces the average of $x$ over the window by its value at the window\u0026rsquo;s midpoint, and the error of that replacement dies with $\\tau$ by continuity — which is exactly what the Bonus section below proves properly.\nStep 3: one term survives. The pulse $r_\\tau$ is zero outside its window of width $\\tau$ around zero — so of all the grid points $n\\tau$, only $n = 0$ lands inside. The infinite sum collapses to a single term:\n$$ \\langle x(t), \\delta(t) \\rangle = \\lim_{\\tau \\to 0} x(0) \\underbrace{r_\\tau(0)}_{=\\,1} = x(0). $$Step 4: the shifted spike. For $\\delta(t - t_0)$, change variables $\\xi = t - t_0$ (so $t = \\xi + t_0$, $dt = d\\xi$, and the infinite limits stay infinite):\n$$ \\int_{-\\infty}^{\\infty} x(t)\\, \\delta(t - t_0)\\, dt = \\int_{-\\infty}^{\\infty} x(\\xi + t_0)\\, \\delta(\\xi)\\, d\\xi = x(t_0) $$by the case we just proved — the spike always reports the value of $x$ at the point where it stands. $\\blacksquare$\nBonus: can the limit swap in Step 1 be saved by a theorem? A natural hope: perhaps uniform convergence, or Lebesgue\u0026rsquo;s dominated convergence, justifies moving the limit outside the integral? Here is the surprise — no theorem can, because the equality they would justify is false. The pointwise limit of $\\frac{1}{\\tau} r_\\tau(t)$ is zero at every $t$ except the single point $t = 0$, where it blows up to $\\infty$. Strictly speaking, \u0026ldquo;the integral of the limit\u0026rdquo; therefore does not even exist: the limit function is unbounded, and the proper Riemann integral is defined only for bounded functions — the very criterion we checked earlier. And if we repair the lone point (assign any finite value at $t = 0$, or integrate in the Lebesgue sense, where a single point has measure zero), the integral of a function that is zero almost everywhere is $0$. Either way — nonexistent or zero — the left-hand side is certainly not $x(0)$, and any theorem permitting the swap would be proving nonsense. Accordingly, every hypothesis fails on purpose. The convergence is not uniform — uniformity would demand $\\sup_t \\big| \\frac{1}{\\tau} r_\\tau(t) \\big| \\to 0$, the whole graph fitting into an ever-thinner tube around zero, whereas our supremum is $\\frac{1}{\\tau}$ and blows up instead. And no integrable dominating function exists — at a fixed $t$, the largest value of $\\frac{1}{\\tau} r_\\tau(t)$ over all $\\tau$ is $\\frac{1}{2|t|}$, and $\\frac{1}{2|t|}$ is not integrable near zero.\nThe honest classical route runs in the opposite direction: swap the roles. Since \u0026ldquo;integrate the limit\u0026rdquo; is meaningless, we declare that the expression $\\langle x(t), \\delta(t) \\rangle$ shall mean \u0026ldquo;the limit of the integrals\u0026rdquo;:\n$$ \\langle x(t), \\delta(t) \\rangle \\;\\overset{\\text{def}}{=}\\; \\lim_{\\tau \\to 0} \\int_{-\\infty}^{\\infty} x(t)\\, \\tfrac{1}{\\tau} r_\\tau(t)\\, dt. $$At each fixed $\\tau$ everything inside is classical — a bounded integrand, an ordinary integral — and only after integrating do we let $\\tau \\to 0$. This limit we can compute directly. For a fixed $\\tau$ there is no delta anywhere in sight:\n$$ \\int_{-\\infty}^{\\infty} x(t)\\, \\tfrac{1}{\\tau} r_\\tau(t)\\, dt = \\frac{1}{\\tau} \\int_{-\\tau/2}^{\\tau/2} x(t)\\, dt $$— the plain average of $x$ over the window. By the mean value theorem for integrals, for continuous $x$ this average equals $x(\\xi_\\tau)$ for some point $\\xi_\\tau$ inside the window; as $\\tau \\to 0$ the window collapses, $\\xi_\\tau \\to 0$, and continuity gives $x(\\xi_\\tau) \\to x(0)$. Fully rigorous — for continuous signals. Families like $\\frac{1}{\\tau} r_\\tau$ are called approximate identities, or mollifiers: they are how analysis makes the delta respectable without ever letting it exist as a function.\nA closing caveat in the same physicist\u0026rsquo;s spirit: the delta itself is famously not square-integrable, so $\\langle x, \\delta \\rangle$ is not a true $L^2$ inner product — the angle brackets are a convenient notation for the pairing our formulas suggest. The bonus above makes our particular computation honest for continuous $x$; making the whole delta calculus systematic — for far rougher objects than continuous signals — is the job of distribution theory.\nThe perfect instrument, then, measures $x(nT)$ by integrating $x$ against $\\delta(t - nT)$. Place one impulse at every grid point — the infinite train of shifted deltas is called the Dirac comb (the Russian DSP literature knows it as the lattice function):\n$$ \\text{Ш}_T(t) = \\sum_{n=-\\infty}^{\\infty} \\delta(t - nT), $$and the honest model of a discrete signal is the analog signal multiplied by the comb (we follow the construction from ru.dsplib.org, archived):\n$$ x_d(t) = x(t) \\cdot \\text{Ш}_T(t) = \\sum_{n} x(t)\\, \\delta(t - nT). $$\nNote that nothing here is approximate anymore: the finite-$\\tau$ estimate $\\hat{x}$ with its error stayed behind in the limit. This is the exact mathematical model, the one the rest of the derivation stands on.\nDid it fix the integration? Let us check — integrate $x_d$ in a small neighborhood of a sampling point $t_0 = kT$, taking limits $kT \\pm \\tau$ with $\\tau \\lt T$ so that exactly one tooth of the comb falls inside:\n$$ \\begin{aligned} \\int_{kT - \\tau}^{kT + \\tau} x(t) \\left( \\sum_{n} \\delta(t - nT) \\right) dt \u0026= \\int_{kT - \\tau}^{kT + \\tau} x(t)\\, \\delta(t - kT)\\, dt \\\\ \u0026= x(kT), \\end{aligned} $$where the first equality holds because all the deltas except $n = k$ sit outside the integration limits. The integral no longer returns zero — it returns the sample. The sampling points survive integration, which is exactly what the naive model could not do.\n💡 Yes, that symbol is a Cyrillic letter. The comb is traditionally denoted by Ш — \u0026ldquo;sha\u0026rdquo; — and Western literature adopted both the symbol and the name: the Shah function. It is quite possibly the only Cyrillic letter in standard mathematical notation, chosen for the obvious reason: the letter looks like the comb.\n💡 A units check (a detail dsplib is careful about, and most sources skip): $\\delta(t)$ has dimension $1/\\text{time}$ — its area over time is the dimensionless $1$. So if $x(t)$ is in volts, the model $x_d(t)$ is in volts per second: it is a density, not a value — in the same sense as a probability density, whose pointwise values are not probabilities (they can even exceed $1$); only its integrals are. Here too the pointwise values of $x_d$ are useless — zero between the grid points, \u0026ldquo;infinite\u0026rdquo; on them — and the sample lives in the impulse\u0026rsquo;s area, not its height. The volts come back when you integrate — as we just saw. Keep this in mind whenever a stray $T$ or $\\frac{1}{T}$ appears in sampling formulas — it is usually this density speaking.\nNow bring in the series Time to cash in. Remember where this whole detour started: we had $N$ samples out of the ADC and wanted their spectrum, but on the naive model the Fourier coefficient formula returned $c_k \\equiv 0$ — the integrals looked straight through the isolated points. That is what sent us rebuilding the model. Now we hold $x_d$, whose integrals do see the samples — so this is attempt number two at the very same computation.\nThe Fourier series has one more requirement, though: a periodic function. Our recording is time-limited — so extend it, gluing copies of the $N$-sample stretch end to end. The smallest period that works is $P = NT$, the duration of the recording itself:\nFeed the periodically extended comb into the Fourier coefficient formula. Over one period $[0, NT)$ exactly $N$ teeth of the comb fall inside, so under the integral sits a finite sum — and swapping a finite sum with an integral is plain linearity, no anxieties about interchanging limits this time. After the swap each summand is a sifting integral, and it evaluates the exponential at its grid point $t = nT$:\n$$ \\begin{aligned} c_k \u0026= \\frac{1}{NT} \\int_{0}^{NT} x_d(t)\\, e^{-2 \\pi i \\frac{k}{NT} t}\\, dt \\\\ \u0026= \\frac{1}{NT} \\sum_{n=0}^{N-1} \\int_{0}^{NT} x(t)\\, \\delta(t - nT)\\, e^{-2 \\pi i \\frac{k}{NT} t}\\, dt \\\\ \u0026= \\frac{1}{NT} \\sum_{n=0}^{N-1} x(nT)\\, e^{-2 \\pi i \\frac{k}{NT} \\cdot nT} = \\frac{1}{NT} \\sum_{n=0}^{N-1} x(nT)\\, e^{-2 \\pi i \\frac{k n}{N}}. \\end{aligned} $$The dreaded integral has collapsed into a finite sum. And look at what happened in the exponent: the sampling period $T$ cancelled out. The basis functions no longer care about seconds — only about the two integers $k$ and $n$. The continuous world has quietly left the stage.\nOnly N distinct coefficients The formula above is valid for any integer $k$ — but try shifting $k$ by $N$:\n$$ e^{-2 \\pi i \\frac{(k+N) n}{N}} = e^{-2 \\pi i \\frac{k n}{N}} \\underbrace{e^{-2 \\pi i n}}_{=\\,1} = e^{-2 \\pi i \\frac{k n}{N}}, $$so $c_{k+N} = c_k$: the coefficients repeat with period $N$. Of the infinitely many harmonics the Fourier series offered us, only $N$ are genuinely distinct. $N$ numbers in, $N$ numbers out — the books balance.\nAnd notice that both laws from the Four Shades table at the top of the post have just fired, each with its own culprit. Our spectrum is discrete — a list of coefficients $c_k$ rather than a function of a continuous frequency — because we periodized the signal: that is the Fourier series\u0026rsquo; side of the bargain, baked in from the start. And the spectrum is periodic because the signal is sampled: look back at the computation — the whole proof of $c_{k+N} = c_k$ hangs on $e^{-2 \\pi i n} = 1$, which is true only because $n$ is an integer, i.e. because the signal lives on a grid. Periodic in time ⇒ discrete in frequency; discrete in time ⇒ periodic in frequency. Not analogies — theorems, and we walked into both bottom-up.\nThe Discrete Fourier Transform One last cosmetic step: new notation to mark the new attitude. The basis functions have forgotten about $T$, so let the signal forget it too — treat the input as a plain list of numbers and write\n$$ x[n] \\overset{\\text{def}}{=} x(nT), $$where the square brackets signal \u0026ldquo;the $n$-th entry of a list\u0026rdquo;, not \u0026ldquo;the value at a moment of time\u0026rdquo;. Drop the overall $\\frac{1}{NT}$ factor as well (a convention we will revisit in a second) and call what remains $X[k]$ — the Discrete Fourier Transform:\n$$ X[k] = \\sum_{n=0}^{N-1} x[n]\\, e^{-2 \\pi i \\frac{k n}{N}}, \\qquad k = 0, 1, \\dots, N-1, $$and its inverse, which reassembles the samples from the spectrum:\n$$ x[n] = \\frac{1}{N} \\sum_{k=0}^{N-1} X[k]\\, e^{2 \\pi i \\frac{k n}{N}}, \\qquad n = 0, 1, \\dots, N-1. $$ 💡 Where did the $\\frac{1}{N}$ go? Between the forward and the inverse transform, a total factor of $\\frac{1}{N}$ must appear somewhere — but the two formulas only constrain the product of their scale factors. The common engineering convention puts all of it into the inverse (as above), because the forward transform is computed far more often; a symmetric convention with $\\frac{1}{\\sqrt{N}}$ on both sides also exists and makes the transform unitary. Libraries differ — always check before comparing numbers.\nThe pair is an exact, lossless round trip: $N$ complex coefficients fully encode $N$ samples. And computing it is cheap: the naive sum costs $O(N^2)$ operations, but the Fast Fourier Transform computes exactly the same $X[k]$ in $O(N \\log N)$ — the algorithmic miracle that makes everything downstream (including every spectrogram ever displayed) practical. This video is a beautiful walkthrough of the idea.\nSo here is what this part promised: an honest model of sampling, a periodic extension, and the Fourier series itself handed us the DFT — no axioms required. The machinery for \u0026ldquo;which piano keys were pressed?\u0026rdquo; is built.\nIn Part 2 we put the machinery to work and learn to read its output: what the $N$ complex numbers $X[k]$ actually say about the sound, which physical frequencies they correspond to, why half of them mirror the other half, and what limits the frequency resolution. In Part 3 the DFT starts sliding along the signal — windows, the short-time Fourier transform, and finally the spectrogram itself.\n","permalink":"https://jen1995.github.io/posts/fourier-series-to-spectrogram-part-1/","summary":"Part 1 of 3: from air pressure and a guitar string, through sampling and the Fourier series — to the DFT, derived honestly, delta functions and all.","title":"From the Fourier Series to the Spectrogram, Part 1: From Sound to the DFT"},{"content":"Part 1 ended with the machinery built. The Discrete Fourier Transform takes the $N$ samples of a recording and returns $N$ complex numbers:\n$$ X[k] = \\sum_{n=0}^{N-1} x[n]\\, e^{-2 \\pi i \\frac{k n}{N}}, \\qquad k = 0, 1, \\dots, N-1. $$$N$ numbers in, $N$ numbers out — and every trace of physical time gone: the formula sees only the two integers $k$ and $n$. That was a feature during the derivation, but it leaves us unable to answer the simplest practical question. Remember the hum? Back in Part 1, arguing that the frequency view is worth having, we took a recording polluted by the 50 Hz buzz of the power line — hopeless to fix in the time domain, where the hum is smeared over every sample — and fixed it in the frequency domain, where the hum is a single column: transform, erase that column, transform back, and the melody survives while the buzz is gone. A fine trick — except that now, with the transform actually in our hands, try to perform it. Which column? Which $k$ is 50 Hz? This part is about learning to read the DFT\u0026rsquo;s output: matching indices to physical frequencies, seeing what the resolution of that matching costs, and discovering along the way why half of the output is a mirror image of the other half.\nWhat does X[k] measure? Look at the DFT formula through the lens of the inner product from Part 1\u0026rsquo;s sifting-property proof. There we needed it for continuous signals: multiply pointwise, then integrate — the integral standing in for \u0026ldquo;add everything up\u0026rdquo; over a continuum. For finite lists of samples the same recipe is simpler still: multiply pointwise and literally add. That is the ordinary dot product of vectors — the very case the notation was borrowed from in the first place — plus one refinement for complex signals: the second factor gets conjugated (that is what the overline over $w_k[n]$ denotes):\n$$ X[k] = \\sum_{n=0}^{N-1} x[n]\\, \\overline{w_k[n]} = \\langle x, w_k \\rangle, \\qquad w_k[n] = e^{2 \\pi i \\frac{k n}{N}}. $$Each coefficient is the inner product of the signal with one basis oscillation $w_k$ — a number that measures how much the signal resembles that oscillation, exactly like a dot product measures how much one vector leans along another. As $n$ runs through the $N$ samples, the exponent of $w_k$ grows to $2 \\pi i k$: the basis oscillation makes exactly $k$ full turns across the recording — Part 1\u0026rsquo;s \u0026ldquo;only oscillations that fit a whole number of times into the period\u0026rdquo; rule, wearing its discrete clothes. Here is the whole cast at once, drawn over one window of $N = 8$ samples:\nThis picture is worth pausing on, because it shows what dictates the basis: nothing but $N$ itself. Fix the window — $N$ samples, $NT$ seconds — and the slowest nonzero probe is forced: one full turn over exactly those samples, frequency $1/(NT)$. The rest of the vocabulary follows automatically, because every other probe is an integer multiple of that fundamental — two turns, three turns, up the ladder. Choose the window and the basis comes with it; the DFT then offers its $N$ probe frequencies — zero turns, one turn, two turns, … — and reports the signal\u0026rsquo;s resemblance to each.\n(A detail worth a second look: in the picture, the curves run past the last sample, all the way to $t = NT$. That is not sloppiness — the window genuinely lasts $NT$ seconds: $N$ intervals of $T$ each, with a sample at the start of every interval. The last sample therefore sits at $n = N - 1$, while $t = NT$ is where the next period would begin — and a whole-turn probe arrives there exactly as it started. That endpoint belongs to the copy, not to this window; if we also sampled at $t = NT$, we would be counting the same point twice.)\nThese $N$ oscillations deserve a name, because they are more than measuring sticks: look back at the inverse DFT formula from Part 1 — it is $x[n] = \\frac{1}{N} \\sum_k X[k]\\, w_k[n]$, the signal reassembled as a weighted sum of these very oscillations. So the set $w_0, \\dots, w_{N-1}$ is everything the DFT can say: the forward transform measures how much of each word the signal contains, the inverse composes the signal back out of the words. We will call this set the DFT\u0026rsquo;s vocabulary — $N$ words, and nothing in between.\nHere is what the probes look like in the flesh — the beginning of the vocabulary, drawn over the same sampled signal for four different recording lengths:\nIn every panel the bold orange curve is the $k = 1$ probe — the lowest oscillating frequency the basis has; the fainter curves behind it are $k = 2$ and $k = 3$, twice and three times faster. One full turn per recording, whatever the recording turns out to be — and that is worth staring at while $N$ is small. At $N = 4$ even the slowest word in the vocabulary races through its full turn before the signal has done anything at all: every probe the DFT owns is far too fast for this signal, and a good description is simply not on offer. Only as the recording grows does the vocabulary reach down to where the signal actually lives — by $N = 25$ the lowest probe finally oscillates at almost the signal\u0026rsquo;s own pace.\n(Keep one thing in mind while looking at these curves: they are drawn at phase zero, every one starting from its peak. In the true approximation each probe will additionally come with its own phase shift — the weights the DFT assigns are complex numbers, and a complex weight slides its probe in time. How exactly that works — and how peak-started cosines manage to assemble a zero-started signal — is a couple of pictures away.)\nAnd what does the vocabulary say, once every word gets its weight? Sum all $N$ probes with the coefficients the DFT assigns them — that is exactly the inverse DFT — and draw the sum as a continuous curve. This is the signal the DFT actually believes in:\nTwo things to see here. Inside the window, the green model passes through every blue sample exactly — $N$ numbers in, $N$ numbers out, the books balance as always.\n(Why exactly, and not merely very closely? That has a beautiful answer, but it deserves its own stretch of road — and gets one: the appendix at the end of the post, for whenever the linear-algebra mood strikes.)\n💡 Wait — how do cosines add up to a sine? Every probe we drew starts at its peak, yet the green model starts wherever the signal starts — at zero, even. The resolving detail is that the weights $X[k]$ are complex: each carries a magnitude and a phase, and a complex weight shifts its probe in time. Pair each $k$ with its mirror twin $N - k$, and their joint contribution to the model is $\\tfrac{2 |X[k]|}{N} \\cos\\!\\big( 2 \\pi \\tfrac{k}{NT}\\, t + \\arg X[k] \\big)$ — a cosine nudged to its own starting point, Part 1\u0026rsquo;s $A \\cos(\\dots - \\phi)$ story again. The pictures draw every probe at phase zero; the sum deploys each one shifted, with the $k = 0$ word — the plain mean — soaking up any constant offset. A sine is merely a cosine whose phase is nudged by a quarter turn.\nOutside the window, the model does the only thing a sum of whole-turn oscillations can do: repeats, with the window\u0026rsquo;s own period. The DFT never models your signal as it is — it models a periodic world assembled from your window, Part 1\u0026rsquo;s glued copies meeting us yet again. At $N = 4$ that world has almost nothing to do with the real signal; by $N = 25$ it is a faithful model of the window — and still a pure invention everywhere else. (This picture was suggested by a friend of the blog — thank you!)\nZoom out, and the same experiment shows when the invention comes true. Everything depends on how the window relates to the signal\u0026rsquo;s own period:\nThe middle panel is the special one. When the window holds exactly one period of the signal (or any whole number of them), the glued copies reproduce the signal — the model is correct not only inside the window but everywhere, forever. That is the DFT at its happiest: the signal\u0026rsquo;s frequency coincides with one of the probes. In the other two panels the window holds half a period and one and a half: the model still passes through every sample it saw, but its periodic continuation has nothing to do with the real signal — the top one never even goes negative, the bottom one continues in counter-phase — and at every seam the curve kinks. Those kinks have a price in the spectrum, and Part 3 charges it under the name spectral leakage.\nTurn the \u0026ldquo;whole turns\u0026rdquo; rule around, and it becomes a restriction important enough to put in bold: whole numbers of turns are all the DFT has. Its basis contains the constant ($k = 0$) and the oscillations that fit a whole number of times into the recording — nothing else. A tone that completes, say, two and a half turns over our $N$ samples is simply not in the vocabulary: no single $X[k]$ is \u0026ldquo;its\u0026rdquo; coefficient. Real recordings contain such tones all the time, of course, and the DFT must express them somehow — smearing them across the whole-turn vocabulary it does have. The consequences of that smearing (it goes by the name spectral leakage) will matter a great deal when we build the spectrogram in Part 3; for now, keep in mind that the DFT\u0026rsquo;s world is quantized to whole turns.\n💡 The $k = 0$ probe makes zero turns: $w_0[n] \\equiv 1$, and $X[0] = \\sum_n x[n]$ is just $N$ times the average of the signal. Audio engineers call it the DC component (from \u0026ldquo;direct current\u0026rdquo; — the electrical origin shows). For sound it is normally near zero: pressure oscillates around the atmospheric baseline, and the microphone measures only the deviation.\nFrom the index to hertz So $X[k]$ measures the content of \u0026ldquo;$k$ turns per recording\u0026rdquo;. To turn that into hertz, bring back the physical time that the DFT so pointedly forgot. The $n$-th sample was taken at $t = nT$, and the whole recording lasts $NT$ seconds — so \u0026ldquo;$k$ full turns per recording\u0026rdquo; is the physical frequency\n$$ f_k = \\frac{k}{NT} = k \\frac{f_s}{N}, $$where $f_s = 1/T$ is the sampling rate. The probe frequencies are not arbitrary: they form a uniform grid with step\n$$ \\Delta f = \\frac{f_s}{N} = \\frac{1}{NT}, $$called the frequency resolution — no probe exists between $f_k$ and $f_{k+1}$, so the DFT simply cannot distinguish frequencies closer together than $\\Delta f$.\nNow, a question worth pausing on. The sampling rate is not really ours to choose — it is a property of the microphone and the ADC, fixed in hardware. The one knob we do control is $N$: how many samples we feed the transform. What exactly does that knob turn? Look again at the probe picture in the previous section — it holds the answer.\nThe basis rides the window, not the clock: more samples at the same rate means a longer recording, and the one turn of $w_1$ spreads over it. Write $w_1$ out and walk it back from the index notation to physical time, exactly the way we did for a general $k$:\n$$ w_1[n] = e^{2 \\pi i \\frac{n}{N}} = e^{2 \\pi i \\frac{1}{NT} \\cdot nT} $$— the right-hand form is a sinusoid of physical frequency $\\frac{1}{NT}$, caught at the moments $t = nT$. The $N$ sits in the denominator of the frequency: lengthen the recording, and the slowest probe slows down with it — $\\Delta f = \\frac{1}{NT}$ drops. And since every other probe sits at a multiple of it, the whole frequency grid tightens:\nThat is what the knob turns. Look at the second formula for $\\Delta f$ again: $NT$ is simply the duration of the recording, so\n$\\Delta f = \\dfrac{1}{\\text{duration}}$ — frequency resolution is one over the listening time.\nTo tell two tones one hertz apart, you must listen for at least a second — only then does the slower tone fall a full turn behind and the difference become visible. There is no way around this trade, only a choice along it: this very tension — resolution in frequency versus locality in time — will return as the central design decision of the spectrogram in Part 3.\nA worked example in unforgiving numbers: at $f_s = 8000$ Hz, taking $N = 100$ samples (an eighth of a second) gives $\\Delta f = 8000 / 100 = 80$ Hz. That grid is too coarse to tell middle C (262 Hz) from the B just below it (247 Hz): both land in bin $k = 3$. To separate them you need $\\Delta f$ around their 15 Hz gap — that is $N \\approx 530$ samples, a fifteenth of a second. Music transcription from spectra is possible, but the resolution bill must be paid first. (A caveat for the careful: $\\Delta f$ equal to the gap is the bare threshold — at that margin the two peaks only just stand apart, as the companion notebook at the end of this post shows; comfortable separation wants a finer grid still.)\nThe spectrum, at last We can now do properly what Part 1 could only preview: plot the magnitudes $|X[k]|$ against their physical frequencies $f_k$. This picture is the (magnitude) spectrum of the recording:\nA pure sine shows up as a sharp spike at its frequency — plus a curious twin at the far end of the axis (explained below, in The mirror) — and a mix of three sines as three spikes with the right heights (and three twins), each component recoverable at a glance. This is the \u0026ldquo;hundreds of numbers → three meaningful ones\u0026rdquo; compression promised at the very start of Part 1 — delivered.\nWhat about the other half of the complex number? Write $X[k]$ in polar form — every complex number is a length times a direction:\n$$ X[k] = |X[k]|\\, e^{i \\varphi_k}, \\qquad \\varphi_k = \\arg X[k]. $$The length $|X[k]|$ is the magnitude we have just been plotting; the angle $\\varphi_k$ is the phase, and it stores where in its cycle the $k$-th oscillation starts — Part 1\u0026rsquo;s $\\phi_n$, one per harmonic. Much of speech processing works with magnitudes alone: what a vowel sounds like, which note was played, whether the hum is there — all of it lives in the magnitudes. The phase becomes essential the moment you need to rebuild the waveform — shifting every component back to its proper starting point — which is why speech synthesis and enhancement systems must treat it with care while a speech recognizer can throw it away.\nThe mirror Look at the spectra again: the single 4 Hz sine lights up its own bin and a twin at 96 Hz, and in the three-sine mix the whole right half of the axis mirrors the left. This is not an artifact of the example — it is a theorem about every real-valued signal, and it takes four lines to prove. Compute the coefficient at index $N - k$ (the overline here and below is complex conjugation, $\\overline{a + bi} = a - bi$ — the flip of the imaginary part\u0026rsquo;s sign):\n$$ \\begin{aligned} X[N-k] \u0026= \\sum_{n=0}^{N-1} x[n]\\, e^{-2 \\pi i \\frac{(N-k) n}{N}} \\\\ \u0026= \\sum_{n=0}^{N-1} x[n]\\, e^{-2 \\pi i n}\\, e^{2 \\pi i \\frac{k n}{N}} \\\\ \u0026= \\sum_{n=0}^{N-1} x[n]\\, e^{2 \\pi i \\frac{k n}{N}} \\\\ \u0026= \\overline{\\sum_{n=0}^{N-1} x[n]\\, e^{-2 \\pi i \\frac{k n}{N}}} = \\overline{X[k]}. \\end{aligned} $$One move per line. First, split the exponent: $\\frac{N-k}{N} = 1 - \\frac{k}{N}$, so the exponential factors into $e^{-2 \\pi i n} \\cdot e^{2 \\pi i \\frac{k n}{N}}$. Second, $e^{-2 \\pi i n} = 1$ because $n$ is an integer — the same \u0026ldquo;signal lives on a grid\u0026rdquo; card we played to prove $c_{k+N} = c_k$ in Part 1. Third, recognize a conjugate: flipping the sign of the exponent conjugates each exponential, and the samples $x[n]$ are real — conjugation passes through them untouched — so the conjugation bar slides over the entire sum, and the sum under the bar is exactly the DFT formula for $X[k]$.\nThe second half of the spectrum is therefore the complex conjugate of the first, read backwards — and conjugation does not change magnitudes: $|X[N-k]| = |X[k]|$. Every spike below the middle has a twin above it, and the magnitude spectrum of any real signal is symmetric about $k = N/2$.\nWe have met this mirror before. Part 1\u0026rsquo;s exponential form split every real oscillation into a forward- and a backward-rotating exponential, with $c_{-n} = \\overline{c_n}$ — and we promised the symmetry would resurface in the DFT. Here it is: by the periodicity $c_{k+N} = c_k$ that we proved when deriving the DFT, the coefficient $c_{-k}$ is $c_{N-k}$. The upper half of the DFT output is precisely the negative-frequency half of the spectrum, wrapped around by periodicity into the range $k = N/2, \\dots, N-1$. The twin spike at \u0026ldquo;96 Hz\u0026rdquo; is really the spike at $-4$ Hz, filed under an alias.\n💡 Do the books still balance? If half of the output mirrors the other half, doesn\u0026rsquo;t the DFT return only $N/2$ numbers\u0026rsquo; worth of information for $N$ numbers of input? Count the real degrees of freedom (say $N$ is even). $X[0]$ is a sum of real samples — real, one number. $X[N/2]$ pairs with itself in the mirror ($N - N/2 = N/2$), forcing $X[N/2] = \\overline{X[N/2]}$ — real, one number. The remaining $N - 2$ coefficients come in conjugate pairs, each pair carrying one independent complex number — two real ones. Total: $1 + 1 + \\frac{N-2}{2} \\cdot 2 = N$ real numbers. Exactly the information content of $N$ real samples — the books balance to the cent, as they must for an invertible transform.\nThe Nyquist frequency The mirror axis itself deserves a name. The probe at the fold, $k = N/2$, is the oscillation\n$$ w_{N/2}[n] = e^{2 \\pi i \\frac{(N/2) n}{N}} = e^{\\pi i n} = (-1)^n $$— the sequence $+1, -1, +1, -1, \\dots$, flipping sign at every single sample. No oscillation representable on the grid can flip faster: there is nothing between the samples to flip in. Its physical frequency comes from the bin-to-hertz formula $f_k = k \\frac{f_s}{N}$ of the resolution section, with $k = N/2$ plugged in:\n$$ f_{N/2} = \\frac{N}{2} \\cdot \\frac{f_s}{N} = \\frac{f_s}{2}, $$is the Nyquist frequency — the ceiling of what a given sampling rate can represent, sitting always at exactly half of it. Draw this axis onto the spectra from earlier, and the symmetry becomes something you can fold with your eyes:\nEverything above the ceiling and below $f_s$ is the mirror land we mapped in the previous section — present in the numbers, but carrying nothing new for a real signal. So here is the full geography of the DFT\u0026rsquo;s frequency axis on one picture:\nThe genuinely informative range runs from $0$ to $f_s/2$ — which finally explains a number from Part 1: CD audio samples at 44.1 kHz because its ceiling must clear the ~20 kHz limit of human hearing, with a little engineering margin on top.\nIn fact, the ceiling is the visible edge of one of the most celebrated results in all of signal processing — and we finally know enough to state it properly:\nThe sampling theorem (Kotelnikov, 1933; independently Shannon, 1949; the frequency bears Nyquist\u0026rsquo;s name — this theorem was discovered by everyone). A continuous signal containing no frequencies higher than $f$ Hz is completely determined by its samples taken $f_s = 2f$ times per second.\nRead the claim slowly, because it is startling. Between two neighboring samples, a continuous signal could seemingly wiggle any way it pleases — and the theorem says a band-limited one cannot: the samples pin down the entire continuous curve, exactly, nothing lost. This is the license behind everything we have done since Part 1 — the reason a list of numbers can honestly stand in for a sound wave, provided the wave had nothing above $f_s/2$ to begin with.\nAnd if it did? The ceiling is a statement about what the grid can represent — but the analog world does not consult our grid. Frequencies above $f_s/2$ do not politely disappear at sampling; they fold back into the visible range under false names — aliasing, the theorem\u0026rsquo;s dark twin. The proof of the theorem, the folded world of aliasing, and even the fine print hiding in the statement above (sticklers: a sinusoid at exactly $f$ needs care) are a story for later in this Fourier journey — the theorem with three names (Kotelnikov, Shannon, Nyquist) will get a post of its own, and by then the machinery we keep building will have turned its proof into a single picture. Consider it teased.\nOnward We can now read a spectrum: find any frequency\u0026rsquo;s bin, trust the left half, ignore the mirror, and budget the resolution by the length of the recording. So let us read a real one. Here is a phrase of live speech — the title of this series, spoken aloud — and the magnitude spectrum of the entire recording:\nNotice the asymmetry in what the two pictures know. In the waveform you can practically count the words — bursts of energy separated by silences — but no frequencies are visible. The spectrum knows all the frequencies: the tall peaks on the left are the voice and its harmonics. But it is one column of numbers for the whole phrase — every phoneme\u0026rsquo;s frequencies stacked into the same bins, and when is gone entirely. A chord is a fine thing to summarize with one spectrum; a sentence is not: \u0026ldquo;which frequencies appear\u0026rdquo; is not the same as \u0026ldquo;which frequencies appear when\u0026rdquo;. In Part 3 we make the Fourier view local — slide a window along the recording, pay the resolution bill we just learned about at every stop, and stack the results into the picture this series is named after: the spectrogram.\nEvery computation of this part can be rerun and poked at: a ready-made notebook lives in this blog\u0026rsquo;s repository and opens in Colab in one click. It also covers the corner cases the post glossed over: the sine that vanishes at exactly the Nyquist frequency, what happens to the Nyquist bin when $N$ is odd, and the factor-of-2 bookkeeping of one-sided amplitude spectra (the same machinery scipy applies inside its one-sided routines).\nAppendix: the probes are a basis Back in What does X[k] measure?, the green model hit every sample dead on — and that was promised to be no accident. This interlude pays the debt — a stretch of honest linear algebra, ending with a reunion with an old classic.\nFirst, write the model down. As a function of continuous time, the $k$-th probe is $e^{2 \\pi i \\frac{k}{NT} t}$ — the $k$-th grid frequency — so the weighted probe sum is\n$$ \\tilde{x}(t) = \\frac{1}{N} \\sum_{k=0}^{N-1} X[k]\\, e^{2 \\pi i \\frac{k}{NT} t}. $$Now put a sample instant $t = nT$ into it. The $T$ cancels in the exponent, and\n$$ \\tilde{x}(nT) = \\frac{1}{N} \\sum_{k=0}^{N-1} X[k]\\, e^{2 \\pi i \\frac{k n}{N}} = x[n] $$— the right-hand side is literally the inverse DFT formula from Part 1, and its output is the original samples. So at the grid instants the model has no freedom at all: it is contractually obliged to return $x[n]$.\nThat pushes the question one step deeper: why does the inverse formula reproduce the samples exactly? To see it, switch to vector language — and start slow. The inverse DFT is one equation per sample — write them all out, one under another:\n$$ \\begin{aligned} x[0] \u0026= \\tfrac{X[0]}{N}\\, w_0[0] + \\tfrac{X[1]}{N}\\, w_1[0] + \\dots + \\tfrac{X[N-1]}{N}\\, w_{N-1}[0] \\\\ x[1] \u0026= \\tfrac{X[0]}{N}\\, w_0[1] + \\tfrac{X[1]}{N}\\, w_1[1] + \\dots + \\tfrac{X[N-1]}{N}\\, w_{N-1}[1] \\\\ \u0026\\;\\;\\vdots \\\\ x[N-1] \u0026= \\tfrac{X[0]}{N}\\, w_0[N-1] + \\dots + \\tfrac{X[N-1]}{N}\\, w_{N-1}[N-1]. \\end{aligned} $$Now read this system by columns instead of rows. Every column carries one and the same scalar $\\tfrac{X[k]}{N}$ — and the stack of numbers it multiplies, $w_k[0], w_k[1], \\dots, w_k[N-1]$, is the $k$-th probe written out top to bottom. So the $N$ equations are really one equation between columns:\n$$ \\begin{pmatrix} x[0] \\\\ x[1] \\\\ \\vdots \\\\ x[N-1] \\end{pmatrix} = \\frac{X[0]}{N} \\begin{pmatrix} w_0[0] \\\\ w_0[1] \\\\ \\vdots \\\\ w_0[N-1] \\end{pmatrix} + \\dots + \\frac{X[N-1]}{N} \\begin{pmatrix} w_{N-1}[0] \\\\ w_{N-1}[1] \\\\ \\vdots \\\\ w_{N-1}[N-1] \\end{pmatrix}. $$Name the columns — $\\mathbf{x}$ on the left, $\\mathbf{w}_k$ for the $k$-th one on the right — and the whole system collapses into one line, with the forward transform alongside it. The second formula is nothing new: it is the inner-product form of the DFT from the very first section, now wearing vector notation:\n$$ \\mathbf{x} = \\sum_{k=0}^{N-1} \\frac{X[k]}{N}\\, \\mathbf{w}_k, \\qquad X[k] = \\langle \\mathbf{x}, \\mathbf{w}_k \\rangle $$So the claim to be proved now reads: measure the resemblances $X[k] = \\langle \\mathbf{x}, \\mathbf{w}_k \\rangle$, use them as weights — and the weighted probes reassemble $\\mathbf{x}$ itself. One property of the probes does all the work here: as vectors, they are orthogonal to one another. Compute their inner product:\n$$ \\langle w_j, w_k \\rangle = \\sum_{n=0}^{N-1} e^{2 \\pi i \\frac{(j - k) n}{N}}. $$For $j = k$ every term is $1$ and the sum is $N$. For two different probes it is a geometric series with ratio $q = e^{2 \\pi i (j-k)/N}$ — a ratio that is not $1$ itself, yet satisfies $q^N = 1$ — so the sum, $\\frac{q^N - 1}{q - 1}$, is exactly zero.\nThat settles everything. Orthogonal vectors are linearly independent, so the $N$ probes form a genuine basis of the $N$-dimensional space of sample lists — the word we have been using all along, now earned — and every $\\mathbf{x}$ is a weighted sum of probes in exactly one way. And for an orthogonal basis, the weights are given by the textbook formula: the inner product with the basis vector, divided by that vector\u0026rsquo;s squared length. That is precisely our pair of formulas — with squared length $\\langle \\mathbf{w}_k, \\mathbf{w}_k \\rangle = N$, which is where the $\\frac{1}{N}$ of the inverse DFT has been hiding all along. (Part 1\u0026rsquo;s convention inset told you where that factor is put; this is why it exists. And one step further, for the matrix-minded: stack the probes as the columns of an $N \\times N$ matrix $W$; orthogonality reads $W^{*} W = N I$, so $W / \\sqrt{N}$ is unitary — the DFT is, up to scale, a rotation of $\\mathbb{C}^N$.)\nOne more name for the same matrix Look at the entries of $W$ once more: $(W)_{nk} = w_k[n] = \\omega^{nk}$ with $\\omega = e^{2 \\pi i / N}$, so the $n$-th row is $1, \\omega^n, \\omega^{2n}, \\dots$ — the successive powers of a single number. A matrix whose rows are geometric progressions of their own \u0026ldquo;nodes\u0026rdquo;, $V_{jk} = z_j^{\\,k}$, has a classical name: the Vandermonde matrix. Written out in full, with one node per row:\n$$ V = \\begin{pmatrix} 1 \u0026 z_0 \u0026 z_0^2 \u0026 \\cdots \u0026 z_0^{N-1} \\\\ 1 \u0026 z_1 \u0026 z_1^2 \u0026 \\cdots \u0026 z_1^{N-1} \\\\ \\vdots \u0026 \\vdots \u0026 \\vdots \u0026 \u0026 \\vdots \\\\ 1 \u0026 z_{N-1} \u0026 z_{N-1}^2 \u0026 \\cdots \u0026 z_{N-1}^{N-1} \\end{pmatrix}. $$It is the matrix of polynomial interpolation — solving $V \\mathbf{c} = \\mathbf{y}$ means finding a polynomial with coefficients $c_k$ that passes through the points $(z_j, y_j)$ — and its famous determinant,\n$$ \\det V = \\prod_{0 \\,\\le\\, i \\,\\lt\\, j \\,\\le\\, N-1} (z_j - z_i), $$says it is invertible exactly when all the nodes are distinct.\nThe DFT matrix, then, is the Vandermonde matrix with its nodes placed at the $N$-th roots of unity, $z_j = \\omega^j$ — watch the rows fill in with powers of $\\omega$:\n$$ W = \\begin{pmatrix} 1 \u0026 1 \u0026 1 \u0026 \\cdots \u0026 1 \\\\ 1 \u0026 \\omega \u0026 \\omega^2 \u0026 \\cdots \u0026 \\omega^{N-1} \\\\ 1 \u0026 \\omega^2 \u0026 \\omega^4 \u0026 \\cdots \u0026 \\omega^{2(N-1)} \\\\ \\vdots \u0026 \\vdots \u0026 \\vdots \u0026 \u0026 \\vdots \\\\ 1 \u0026 \\omega^{N-1} \u0026 \\omega^{2(N-1)} \u0026 \\cdots \u0026 \\omega^{(N-1)^2} \\end{pmatrix}. $$Nodes as distinct, and as symmetric, as $N$ points can be — and that special placement is what upgrades \u0026ldquo;invertible\u0026rdquo; to \u0026ldquo;unitary up to scale\u0026rdquo;. Two consequences are worth savoring.\nFirst, our sample-hitting model turns out to be a polynomial interpolation problem, posed the most natural way: find the polynomial of degree $N-1$,\n$$ q(z) = c_0 + c_1 z + \\dots + c_{N-1} z^{N-1}, $$whose values at the $N$ roots of unity are exactly our samples — one node per sample:\n$$ q(\\omega^n) = \\sum_{k=0}^{N-1} c_k\\, \\omega^{k n} = x[n], \\qquad n = 0, \\dots, N-1. $$Stare at this system of $N$ conditions: the matrix multiplying the unknown coefficients has entries $\\omega^{nk}$ — it is exactly our $W$. Finding the polynomial through the samples means solving $W \\mathbf{c} = \\mathbf{x}$ — written out in full:\n$$ \\underbrace{\\begin{pmatrix} 1 \u0026 1 \u0026 \\cdots \u0026 1 \\\\ 1 \u0026 \\omega \u0026 \\cdots \u0026 \\omega^{N-1} \\\\ \\vdots \u0026 \\vdots \u0026 \u0026 \\vdots \\\\ 1 \u0026 \\omega^{N-1} \u0026 \\cdots \u0026 \\omega^{(N-1)^2} \\end{pmatrix}}_{W} \\underbrace{\\begin{pmatrix} c_0 \\\\ c_1 \\\\ \\vdots \\\\ c_{N-1} \\end{pmatrix}}_{\\mathbf{c}} = \\underbrace{\\begin{pmatrix} x[0] \\\\ x[1] \\\\ \\vdots \\\\ x[N-1] \\end{pmatrix}}_{\\mathbf{x}}. $$And this system we can simply solve, because the inverse matrix costs nothing: a few paragraphs ago orthogonality gave us $W^{*} W = N I$, that is,\n$$ W^{-1} = \\frac{1}{N} W^{*}, \\qquad\\text{so}\\qquad \\mathbf{c} = \\frac{1}{N} W^{*} \\mathbf{x}. $$Now look at what the product $W^{*} \\mathbf{x}$ computes, row by row. Written out, the conjugate transpose is simply $W$ with the sign of every power flipped — conjugation negates the exponents, and transposition changes nothing because $W$ is symmetric:\n$$ W^{*} = \\begin{pmatrix} 1 \u0026 1 \u0026 1 \u0026 \\cdots \u0026 1 \\\\ 1 \u0026 \\omega^{-1} \u0026 \\omega^{-2} \u0026 \\cdots \u0026 \\omega^{-(N-1)} \\\\ 1 \u0026 \\omega^{-2} \u0026 \\omega^{-4} \u0026 \\cdots \u0026 \\omega^{-2(N-1)} \\\\ \\vdots \u0026 \\vdots \u0026 \\vdots \u0026 \u0026 \\vdots \\\\ 1 \u0026 \\omega^{-(N-1)} \u0026 \\omega^{-2(N-1)} \u0026 \\cdots \u0026 \\omega^{-(N-1)^2} \\end{pmatrix}. $$The $k$-th row holds the conjugated probe $\\overline{w_k[n]} = \\omega^{-k n}$, so the $k$-th entry of $W^{*} \\mathbf{x}$ is $\\sum_n x[n]\\, \\omega^{-k n}$ — the forward DFT, term for term. The solution of the interpolation system turns out to be our old coefficients:\n$$ c_k = \\frac{X[k]}{N}. $$To interpolate through the samples, run a forward DFT — that is the entire algorithm. (Interpolation normally costs solving a linear system; this particular system\u0026rsquo;s matrix inverts by mere conjugation, because its columns are orthogonal.)\nAnd the interpolating polynomial is an old friend. Substitute $z = e^{2 \\pi i t / (NT)}$: as $t$ runs through the window, $z$ walks once around the unit circle, visiting the node $\\omega^n$ exactly at the sample instant $t = nT$. Under this substitution — with the solution $c_k = X[k]/N$ plugged in on the second line —\n$$ \\begin{aligned} q\\!\\left( e^{2 \\pi i t/(NT)} \\right) \u0026= \\sum_{k=0}^{N-1} c_k\\, e^{2 \\pi i \\frac{k}{NT} t} \\\\ \u0026= \\frac{1}{N} \\sum_{k=0}^{N-1} X[k]\\, e^{2 \\pi i \\frac{k}{NT} t} = \\tilde{x}(t) \\end{aligned} $$— $q$ turns into $\\tilde{x}(t)$, the green model from the pictures. The curve that hit every sample was this interpolating polynomial all along, traced along the circle:\nA fine-print footnote: why the drawn curve is real (it leans on the mirror symmetry) Traced literally with the frequencies $0, \\dots, N-1$, the curve $q(e^{2 \\pi i t/(NT)})$ is complex-valued between the samples. The drawn curve lets the mirror twins act as a conjugate pair instead: for $0 \\lt k \\lt N/2$ the coefficients $X[k]$ and $X[N-k] = \\overline{X[k]}$ keep their full values, but the twin plays its negative frequency $k - N$ rather than $N - k$ — and the pair\u0026rsquo;s two terms, now complex conjugates of each other, sum to the real quantity $2 \\operatorname{Re}\\!\\left( X[k]\\, e^{2 \\pi i k t/(NT)} \\right)$. No averaging happens there; the only coefficient literally split in half is the lone Nyquist one when $N$ is even — half of $X[N/2]$ goes to $+f_s/2$, half to $-f_s/2$, and the halves sum to a real cosine. At the sample instants every one of these choices is indistinguishable, because $e^{2 \\pi i (k-N) n/N} = e^{2 \\pi i k n/N}$ — the periodicity $c_{k+N} = c_k$ yet again.\nSecond, this coefficients-to-values dictionary, cheap in both directions, is exactly how the FFT multiplies polynomials fast — convert both factors to their values at the roots of unity, multiply the values pointwise, convert the product back to coefficients — the trick at the heart of big-integer arithmetic.\n","permalink":"https://jen1995.github.io/posts/fourier-series-to-spectrogram-part-2/","summary":"Part 2 of 3: learning to read the DFT\u0026rsquo;s output — which bin is which frequency, why resolution is one over duration, the mirror symmetry, and the Nyquist frequency.","title":"From the Fourier Series to the Spectrogram, Part 2: Reading the DFT"},{"content":"Part 2 ended on an asymmetry. The waveform of a spoken phrase knows when — you can count the words in it — but no frequencies. Its spectrum knows what — every frequency of every phoneme — but stacked into one column of numbers with no notion of time. A chord survives such a summary; a sentence does not. Speech is non-stationary: its spectrum changes many times per second, and the whole point of listening is to follow the changes.\nThis closing part builds the tool that keeps both answers at once. The plan writes itself: if one spectrum for the whole recording is too coarse, take many — cut the signal into short pieces and transform each piece separately. Everything else in this post is the honest accounting for that one idea: what the cutting costs (spectral leakage), how to soften the blow (windows), what the assembled picture is (the spectrogram), which trade-off it can never escape (time versus frequency), and one final compression pass borrowed from your own ear (the mel scale).\nCut the signal into frames Slide a short window along the recording and take the DFT of each position separately. Each windowed stretch is called an (acoustic) frame, and two numbers govern the slicing: the window length $W$ — how many samples each frame holds — and the hop length $H$ — how far the window advances between frames. With $H \\lt W$ the frames overlap, and every sample gets seen by several of them. The overlap is not decoration: sounds do not schedule themselves to fit our slicing, and an event that falls on a frame boundary would otherwise be chopped in half — overlapping frames guarantee that every moment is also seen whole, near the middle of some frame. (A second, sneakier reason will surface when we meet window functions in a couple of sections.)\nFormally, this is the Short-Time Fourier Transform (STFT): the DFT of Part 2, applied to the $m$-th frame,\n$$ \\mathrm{STFT}[m, k] = \\sum_{n=0}^{W-1} x[m H + n]\\; w[n]\\; e^{-2 \\pi i \\frac{k n}{W}}, $$Unpack the ingredients. $W$ and $H$ are the window and hop lengths from the figure, both in samples; $x[mH + n]$ walks through the $m$-th frame — its first sample sits $m$ hops from the start of the recording. Note the two jobs $W$ does: it is the number of samples summed and the size of the DFT in the exponent — each frame gets the full Part 2 treatment as if it were an entire recording of length $W$. Finally, $w[n]$ is a window function whose job we are about to discover — for now, imagine $w[n] \\equiv 1$, a plain rectangular cutout. The result is indexed by two integers: $k$ still means \u0026ldquo;which frequency\u0026rdquo;, exactly as in Part 2, and the new index $m$ means \u0026ldquo;which moment\u0026rdquo;. This is the whole idea; the rest of the post is fine print. But in signal processing, as we have learned twice already, the fine print is where the theorems live.\nThe price of cutting: spectral leakage Here is the debt from Part 2 coming due. Remember how the DFT was derived in Part 1: we took $N$ samples and glued copies of them end to end — the Fourier series demanded a periodic function, so we manufactured one. That construction is still inside the machine: the DFT of a frame is honestly analyzing the infinite periodic signal made of glued copies of that frame.\nPart 1 also showed what gluing demands: whatever happens on the frame must join seamlessly to its own copy. A frequency that completes a whole number of turns per frame arrives at the seam exactly where it started — its copies glue smoothly, and the DFT gives it one clean bin. But our window now lands wherever it lands: no real oscillation consults the frame boundaries. A tone that completes, say, $8.5$ turns arrives at the seam mid-swing — and the glued copies tear:\nThe DFT does not complain; it reports what it sees. And what it sees is a signal with a jump at every seam — and jumps, as the square wave showed us back in Part 1, take a whole choir of frequencies to build. So the energy of one pure tone leaks across the spectrum: a tall pair of bins where the tone roughly is, and a skirt of nonzero magnitudes everywhere else. This is spectral leakage — Part 2\u0026rsquo;s \u0026ldquo;not in the vocabulary\u0026rdquo; effect, finally caught in the act. (You have already met it once: in the Part 2 notebook, the strange dip between the two nearly-resolved notes was the leakage skirts of two off-grid tones interfering.)\nWindows: taper, don\u0026rsquo;t chop The tear happens at the frame\u0026rsquo;s edges — so treat the edges. Instead of cutting the signal out with a rectangular cookie-cutter, multiply the frame by a window function that rises smoothly from zero and falls smoothly back: whatever the signal was doing, the windowed frame now starts and ends at zero, and the glued copies meet without a jump. The standard first choice is the Hann window,\n$$ w[n] = \\tfrac{1}{2} \\left( 1 - \\cos \\tfrac{2 \\pi n}{W - 1} \\right). $$\nLook at the right panel (magnitudes in dB — decibels, a logarithmic scale; we will justify it in a moment). The rectangular cut leaves a skirt decaying so slowly that a strong tone can drown quiet neighbors three octaves away; the Hann window pushes the skirt down by tens of dB. The price, visible at the very top: the central peak becomes about twice as wide — a windowed tone occupies two-ish bins even when perfectly on-grid. Windowing trades a little blur near the true frequency for enormous cleanliness far from it. (There is a whole zoo of windows — Hamming, Blackman, Kaiser — each choosing this trade slightly differently; the Hann window is the workhorse default in speech.)\nAnd here is the promised second reason for overlapping frames: a windowed frame is nearly deaf at its own edges — the taper multiplies the edge samples by almost zero. If the frames merely touched ($H = W$), the signal near every boundary would go essentially unheard. With enough overlap, what one frame tapers away sits at full volume near the middle of a neighboring frame — and the classic choice $H = W/2$ does something almost magical: the shifted Hann windows sum to exactly one. No magic, just a half-period shift flipping the sign of the cosine:\n$$ \\begin{aligned} w[n] + w\\!\\left[n + \\tfrac{W}{2}\\right] \u0026= \\tfrac{1}{2}\\left(1 - \\cos\\tfrac{2\\pi n}{W}\\right) + \\tfrac{1}{2}\\left(1 + \\cos\\tfrac{2\\pi n}{W}\\right) \\\\ \u0026= 1. \\end{aligned} $$\nEvery sample gets exactly its fair share of attention — while at a lazier hop the sum ripples, and the signal near the seams is systematically underheard. (A technicality for the careful: exact constancy holds for the periodic variant of the Hann window, with $W$ rather than $W - 1$ in the denominator — a one-sample difference that matters only when you need to reconstruct the signal from its frames.)\nThe spectrogram Now assemble. Compute the windowed DFT of frame $0$, frame $1$, frame $2$, …; keep only the informative half of each spectrum (Part 2\u0026rsquo;s mirror); take magnitudes; stand each result upright as a column and line the columns up in time order. The resulting matrix — frequency up the side, time along the bottom, magnitude as brightness — is the spectrogram:\nThe same phrase as in Part 2 — but where the whole-recording spectrum collapsed every syllable into one anonymous forest of peaks, the spectrogram lays the sentence out like a musical score. Read it top to bottom and left to right:\nthe bright horizontal striations in the lower half are the harmonics of the voice — the Fourier series of Part 1, alive and visible, one row per harmonic; their spacing is the pitch of the speaker; the tall unstructured columns reaching high frequencies are the hissy consonants — \u0026ldquo;s\u0026rdquo;, \u0026ldquo;f\u0026rdquo;, \u0026ldquo;sh\u0026rdquo; — noise-like sounds with energy smeared across the spectrum; the dark vertical gaps are the silences between words: the when that Part 2\u0026rsquo;s spectrum had lost is back on the horizontal axis. Two familiar knobs decide the geometry of this picture, and both are Part 2 veterans. The sampling rate sets the ceiling: rows run from $0$ to $f_s/2$, the Nyquist frequency — nothing above the ceiling exists on the grid. The window length sets the resolution: each column is a $W$-sample DFT, so its rows are spaced $\\Delta f = f_s / W$ apart — the number of rows is the window length (divided by two). And one new knob: the hop length sets the frame rate of the movie — how many columns per second of signal.\n💡 Why dB? Spectrogram magnitudes are conventionally shown as $20 \\log_{10}$ of the magnitude — decibels. Two reasons. First, the dynamic range: the loud harmonics and the quiet fricative hiss differ by factors of thousands; on a linear brightness scale everything but the harmonics would be black. Second, the ear: loudness perception is itself roughly logarithmic — the Weber–Fechner law: what the senses register is the relative change of a stimulus, not the absolute one — so equal steps in dB feel like equal steps of loudness. The picture is drawn the way it is heard.\n💡 Typical numbers for speech: window $\\approx 25$ ms, hop $\\approx 10$ ms — a hundred columns per second, each with a $\\Delta f = 40$ Hz grid. Why 25 ms? Short enough that speech barely changes within one frame (quasi-stationarity), long enough to resolve the harmonics of a typical voice. The next section is about why you cannot have both at once. And the third number, the sampling rate: consumer audio\u0026rsquo;s 44.1 kHz exists precisely so that the ceiling $f_s/2$ from Part 2 — here 22.05 kHz — clears the upper limit of the human hearing range, about 20 kHz. Speech systems, whose signal of interest ends far lower, often settle for 16 kHz: an 8 kHz ceiling covers the voice comfortably at a third of the samples.\nThe trade-off you cannot escape Part 2 proved a hard law: the DFT\u0026rsquo;s probe frequencies form a grid whose step is set by nothing but the length of the recording,\n$$ \\Delta f = \\frac{f_s}{N} = \\frac{1}{NT} = \\frac{1}{\\text{duration}} $$— frequency resolution is one over the listening time. In the STFT the $N$ of that formula is the window length $W$: each column listens for $W$ samples — $W/f_s$ seconds — and no longer. So the law becomes a genuine dilemma. A short window pins events precisely in time but smears them in frequency; a long window resolves frequencies finely but averages away the timing. Watch the law act on a chirp — a tone whose frequency climbs steadily (the name is honest: it is the sound of a bird call or a slide whistle). In the time domain a chirp looks like this — constant amplitude, ever-shrinking period:\nOur actual test chirp climbs from $200$ to $3600$ Hz over two seconds — thousands of oscillations, far too many to draw sample by sample; which is precisely the kind of signal you need a spectrogram to see. Here it is under two different window lengths:\nFirst, note what the two pictures share and what they do not. The vertical axis is the same on both: it runs from $0$ to $f_s/2$ — the Nyquist ceiling from Part 2, the grid\u0026rsquo;s edge that the Kotelnikov–Shannon–Nyquist sampling theorem is about — and that range is fixed by the sampling rate alone. What differs is the grid packed inside it: the left picture has $129$ rows spaced $31.2$ Hz apart, the right one $1025$ rows spaced $3.9$ Hz apart — eight times finer, on paper. The bottom row of the figure makes the grids visible: it re-plots the same green patch — $500$ Hz by a quarter of a second — from each picture, pixel for pixel. On the left the patch is spanned by $17$ chunky rows and $32$ columns; on the right, by $129$ fine rows but only four columns. The long window\u0026rsquo;s grid is finer along frequency and far coarser along time — and every pixel of a spectrogram is exactly one (frame, bin) cell of the STFT, so the pixels are the grid. (Where does the column count come from? One column per hop, and we hop by a quarter of the window — $8$ ms on the left, $64$ ms on the right. A smaller hop would draw more columns, but they would be near-duplicates: neighboring frames share three quarters of their samples. The honest time resolution is set by the window length itself; the hop only chooses how densely you sample it.)\nNow the arithmetic of what each column actually sees. The chirp climbs at $(3600 - 200)/2 = 1700$ Hz per second. During one $32$ ms frame it sweeps through $1700 \\cdot 0.032 \\approx 54$ Hz — a couple of bins of the left grid: the frame is nearly a constant tone, and the line comes out about as thin as $\\Delta f$ allows. During one $256$ ms frame the chirp sweeps through $1700 \\cdot 0.256 \\approx 435$ Hz — more than a hundred bins of the right grid. And the column reports the truth: the frame genuinely contained all of those frequencies, so the fine grid faithfully resolves… a smear $435$ Hz wide. The sharper $\\Delta f$ bought a worse picture. Here is the lesson: $\\Delta f$ measures the fineness of the frequency axis — not the thinness of what lands on it. A tone\u0026rsquo;s image is one bin thin only when the tone\u0026rsquo;s frequency stays inside one bin — changes by less than $\\Delta f$ — over the whole window. Drift further, and the frame genuinely contains every frequency the tone visited along the way: a band of them, which no grid, however fine, can render thinner than it really is. Check this against the numbers: the left window catches $54$ Hz of drift on a grid of $31$ Hz bins — about two bins thick, nearly as thin as the axis allows; the right window catches $435$ Hz of drift on $3.9$ Hz bins — a hundred-bin-wide band, rendered in loving detail. This is also what quasi-stationarity — the word from the typical-numbers inset — actually buys: if the signal\u0026rsquo;s frequencies barely move during one window, its image stays sharp. Speech manages that within $25$ ms; a chirp, by definition, manages it at no window length.\nThere is no window length that wins both ways — only a choice matched to the signal. This is not an engineering shortcoming but mathematics: time locality and frequency locality are fundamentally at odds (the same tension that in quantum mechanics goes by the name uncertainty principle — a story for another day). The speech-processing compromise from the previous section — 25 ms — is simply where the trade sits comfortably for human voices.\nThe mel scale: compress like the ear does The spectrogram is already a fine input for a machine. But it spends its rows wastefully — from the listener\u0026rsquo;s point of view. Your ear does not weigh all frequencies equally: the cochlea resolves low frequencies finely and high frequencies coarsely. Neighboring keys at the bottom of a piano differ by a couple of hertz and you hear the step clearly; at the top of the keyboard the same one-semitone step is a couple of hundred hertz — and sounds no bigger. Perceptually, frequency is closer to logarithmic than linear.\nThe mel scale (from melody) is a practical approximation of that perception, fitted to listening experiments:\n$$ m = 2595 \\, \\log_{10}\\!\\left(1 + \\frac{f}{700}\\right). $$Equal steps in mel are meant to sound like equal steps of pitch — which means the scale walks slowly through the low frequencies (where the ear discriminates finely) and takes ever-larger strides up high:\nTo convert a spectrogram, build a mel filter bank: a few dozen triangular filters, evenly spaced in mel — therefore narrow and dense at low frequencies, wide and sparse at high ones (bottom panel above). Now stack the triangles into a matrix $F$, one filter per row: each row is one triangle from the picture above, written out as its weights over all the frequency bins. With $80$ filters — the de-facto standard in speech synthesis; recognition systems often get by with $40$ — over the $257$ frequency rows of our spectrogram, $F$ is an $80 \\times 257$ matrix:\nThe whole conversion is then a single matrix multiplication. One column of the spectrogram is a vector of $257$ magnitudes; multiplying by $F$ takes $80$ weighted sums of it — one sum per filter, each collecting the bins under its triangle. And multiplying $F$ by all the columns at once handles the entire recording in one stroke:\n$$ M = F \\cdot S, \\qquad (80 \\times 257) \\cdot (257 \\times T) = 80 \\times T, $$where $T$ is the number of frames. A logarithm on top — the same dB story as before — and the result is the mel-spectrogram, every column squeezed from $257$ linear-frequency numbers into $80$ perceptually spaced ones:\nHundreds of linear-frequency rows become a few dozen mel bands, spending their budget where the ear spends its attention — the harmonics-rich bottom gets most of the rows, the hissy top is summarized coarsely. This picture — compact, perceptually weighted, still laid out in time — is the actual input of most speech recognition and synthesis models: when a paper says \u0026ldquo;we feed the audio to the network\u0026rdquo;, this matrix is almost always what is being fed.\nOne last piece of accounting closes the story: what did the whole pipeline buy, in raw numbers? For our recording:\nrepresentation shape numbers total the time axis waveform $52\\,225$ samples $52\\,225$ $22\\,050$ values per second spectrogram $S$ $257 \\times 405$ $104\\,085$ $\\approx 170$ columns per second mel-spectrogram $M$ $80 \\times 405$ $32\\,400$ $\\approx 170$ columns per second Two honest surprises in this table. First, the spectrogram is not a compression: overlapping frames see every sample about four times, so $S$ holds twice as many numbers as the waveform it came from. What changed is the organization: the time axis became $128$ times coarser — one column per hop instead of one value per sample — and in exchange each column spells out explicitly what the samples only implied: which frequencies are present at that moment. Second, the genuine shrinkage arrives only with the mel step: $M$ is about $60\\%$ of the raw waveform and less than a third of $S$ — and yet, as the pictures show, still perfectly legible: the harmonics, the fricative bursts, the silences between words all survived. Fewer numbers, arranged so that the structure shows — that, in one line, is what audio feature extraction is for.\nThe road, walked This is where the series set out to arrive, so let us look back at the whole road. A pressure wave shook a microphone, and an ADC turned the voltage into $N$ numbers (Part 1). To do mathematics with those numbers we built them an honest home — deltas, sifting, the comb — and the Fourier series itself handed us the DFT (Part 1). We learned to read the DFT\u0026rsquo;s output: which bin is which hertz, why resolution is one over duration, why half the spectrum is a mirror (Part 2). And today we made the Fourier view local — frames, windows against the leakage, the spectrogram, and the mel compression that matches the ear (Part 3). From air molecules to the input tensor of a speech model, with no step taken on faith.\nThe journey of this series continues past the trilogy. The Fourier transform proper — continuous time, continuous frequency, the one shade we never defined — and the family bridges between all four shades are next. Then the theorem with three names gets its promised proof (one picture, as vowed in Part 2). And the finale reads the human voice itself: F0, pitch and the cepstrum. The four shades await.\nEverything in this part runs in a ready-made notebook that opens in Colab in one click: the STFT in three lines of numpy, leakage measured in percent, the Hann sum-to-one property verified to machine precision (periodic vs symmetric variant included), a spectrogram and a mel bank built from scratch — and, as the finale, the signal reassembled exactly from its complex STFT by overlap-add, with a note on why the magnitude-only spectrogram cannot be inverted so easily (that is what vocoders are for).\n","permalink":"https://jen1995.github.io/posts/fourier-series-to-spectrogram-part-3/","summary":"Part 3 of 3: cut the signal into frames, pay for the cut with leakage, patch it with windows, stack the columns — the spectrogram, and its mel-compressed cousin that speech models actually consume.","title":"From the Fourier Series to the Spectrogram, Part 3: From the DFT to the Spectrogram"},{"content":"Part 0 — Introduction: seq2seq, language models, attention Natural language processing spans many different tasks: text classification (say, spam vs. not spam), sentiment analysis, entity extraction, question answering, summarization, text generation and others. Among them, machine translation stands apart.\nTranslation was the area where a broad audience first noticed a dramatic jump in quality: systems suddenly learned to translate coherent texts so well that the output often needed only light editing — and sometimes almost none. But, more importantly, it quickly became clear that the ideas and methods born in translation actually solve a much more general problem: transforming one message into another while preserving its meaning. In this form, \u0026ldquo;translation\u0026rdquo; generalizes far beyond natural languages — to summarization (text → text), code generation (text → code), code explanation (code → text), and even multimodal transformations (text → image, text → audio, image → image, and so on).\nMuch of this breakthrough is tied to the Transformer architecture (Vaswani et al., 2017), which made training on large amounts of data substantially more efficient and gave models a new level of expressiveness. Yet the Transformer did not appear out of nowhere: it builds on several key ideas, which we will briefly recall in this introduction — the sequence-to-sequence paradigm, language models, and the attention mechanism, in its classical formulation usually associated with the work of Bahdanau and co-authors (Bahdanau et al., 2014).\nFormal setup (Sequence to Sequence) Formally, the translation task can be described as follows: we have an input sequence $x=(x_1,\\dots,x_m)$ and an output sequence $y=(y_1,\\dots,y_n)$ (their lengths may differ). We want to find the translation $y$ that is most probable given the input $x$:\n$$ y^* = \\arg\\max_y\\, p(y\\mid x). $$One standard way to model such tasks is the encoder-decoder architecture. It has two main parts:\nthe encoder reads the source sequence and builds its internal representation; the decoder uses this representation to generate the target sequence. For example, if the source sentence is «Я видел кота на мате», the encoder transforms it into a vector representation capturing the \u0026ldquo;overall meaning\u0026rdquo; of the input sentence, and the decoder generates the translation from this representation: \u0026ldquo;I saw a cat on a mat\u0026rdquo;.\nDiagram after Lena Voita\u0026rsquo;s NLP Course\nWe will see different models below, but they all rely on the same general scheme: first encode the input, then decode the output from this representation.\nLanguage models A language model is a function that can estimate how \u0026ldquo;plausible\u0026rdquo; a string of text is:\nthe input of a language model is $x_1,\\dots,x_n$ — the tokens of a sequence of length $n$, the output is $p(x_1,\\dots,x_n)$ — the probability of the whole sequence among all possible sequences of length $n$. Estimating the probability of an entire sentence directly — as a single indivisible object — is very hard: there are too many possible sentences, and for most of them we will never have enough observations. So, instead of treating a sentence as one \u0026ldquo;atomic\u0026rdquo; unit, let us decompose its probability into probabilities of smaller fragments.\nFor example, take the sentence \u0026ldquo;I saw a cat on a mat\u0026rdquo; and read it word by word. At each step we estimate the probability of the next token given all the context seen so far. Previous computations are not wasted: when a new word arrives, we simply refine the probability of the whole sequence by multiplying it by the probability of the new token given the already-known prefix.\n$$ \\begin{aligned} p(\\text{I saw a cat}) \u0026= p(\\text{I}) \\cdot p(\\text{saw}\\mid\\text{I}) \\cdot p(\\text{a}\\mid\\text{I saw}) \\\\ \u0026\\quad \\cdot p(\\text{cat}\\mid\\text{I saw a}). \\end{aligned} $$Formally, this is just the chain rule from probability theory:\n$$ \\begin{aligned} p(x_1,\\dots,x_n) \u0026= p(x_1)\\,p(x_2\\mid x_1)\\,p(x_3\\mid x_1,x_2)\\,\\dots\\,p(x_n\\mid x_1,\\dots,x_{n-1}). \\end{aligned} $$That is,\n$$ p(x_1,\\dots,x_n)=\\prod_{t=1}^{n} p(x_t\\mid x_1,\\dots,x_{t-1}). $$One usually introduces the shorthand $x_{\\lt t}= (x_1,\\dots,x_{t-1})$ and writes compactly:\n$$ p(x_1,\\dots,x_n)=\\prod_{t=1}^{n} p(x_t\\mid x_{\\lt t}). $$Once we have a language model, we can use it to generate text. This is done one token at a time: at each step the model predicts the probability distribution of the next token given the preceding context, and we then pick or sample a token from this distribution.\nTraining and the cross-entropy loss Neural language models are trained to predict the probability distribution of the next token given the preceding context. Initially the model predicts a uniform distribution over all tokens, but during training it learns to assign more and more probability to the correct next token.\nFormally, let $y_1, \\dots, y_n$ be a training sequence of tokens. Then at step $t$ the model predicts the distribution\n$$ p^{(t)} = p(* \\mid y_1, \\dots, y_{t-1}). $$The target distribution at this step is $p^* = \\mathrm{one\\text{-}hot}(y_t)$: we want the model to assign probability 1 to the correct token $y_t$ and probability 0 to all other tokens.\nAt this point we may notice that what we are looking at is nothing other than a classification problem: the classes are the different tokens, and the target class is the correct token at the current step. The standard loss function for such a problem is cross-entropy. For the target distribution $p^*$ and the predicted distribution $p$ it is written as:\n$$ \\mathrm{Loss}(p^*, p) = -p^* \\log(p) = -\\sum_{i=1}^{|V|} p_i^* \\log(p_i). $$Here $|V|$ is the vocabulary size, i.e. the number of possible tokens over which the model spreads its probability.\nSince only one component $p_i^*$ of the one-hot distribution is nonzero — the one corresponding to the correct token $y_t$ — the expression simplifies:\n$$ \\mathrm{Loss}(p^*, p) = -\\log(p_{y_t}) = -\\log p(y_t \\mid y_{\\lt t}). $$That is, at each step we minimize the negative log-probability of the correct next token. The higher the probability the model assigns to the correct token, the lower the loss.\nDiagram after Lena Voita\u0026rsquo;s NLP Course\nFor the whole sequence, the loss is obtained by summing over all steps:\n$$ \\mathcal{L} = -\\sum_{t=1}^{n} \\log p(y_t \\mid y_{\\lt t}). $$This is exactly the quantity that is usually minimized when training a language model.\nConditional language models In the section on language models we learned to estimate the unconditional probability $p(y)$ of a token sequence $y=(y_1, y_2, \\dots, y_n)$. The step from such language models to sequence-to-sequence models amounts to replacing the unconditional probability $p(y)$ with $p(y \\mid x)$: the probability of the output sequence $y$ given the source sequence $x$.\nThis is why sequence-to-sequence tasks can be viewed as conditional language modeling (CLM). Such models work almost like ordinary language models, but additionally receive information about the source $x$.\nFor an ordinary language model:\n$$ P(y_1, y_2, \\dots, y_n) = \\prod_{t=1}^{n} p(y_t \\mid y_{\\lt t}). $$For a conditional language model:\n$$ P(y_1, y_2, \\dots, y_n \\mid x) = \\prod_{t=1}^{n} p(y_t \\mid y_{\\lt t}, x). $$Here $x$ is the condition — the source information the model relies on during generation.\nImportantly, conditional language modeling is not only a way to solve sequence-to-sequence tasks. In a more general sense, $x$ does not have to be a sequence of tokens at all. For example, in image captioning, $x$ is an image and $y$ is a text description of that image.\nSince the only difference from ordinary language models is the presence of the source $x$, modeling and training are organized very similarly. At a high level, training and generation look like this:\nfeed the network the source and the already-generated tokens of the target sequence; obtain a vector representation of the context — both the source and the previous target history; from this representation, predict the probability distribution of the next token. Diagram after Lena Voita\u0026rsquo;s NLP Course\nBahdanau attention Before turning to the mechanism itself, let us take a closer look at the encoder-decoder architecture described at the beginning. It has a bottleneck: the entire meaning of the input sentence — however long it may be — has to be packed into a single fixed-size vector. While sentences are short, this is tolerable, but the longer the input, the more information is lost in the compression. Worse still, the decoder sees the same frozen representation at every step, although intuitively it needs different parts of the input at different moments: when generating the word \u0026ldquo;cat\u0026rdquo;, it would like to look at «кота» rather than at the whole sentence.\nThis is the problem solved by the attention mechanism, proposed in the paper Neural Machine Translation by Jointly Learning to Align and Translate (Bahdanau et al., 2014). The core idea is to let the model \u0026ldquo;focus\u0026rdquo; on different parts of the input sequence at different decoding steps.\nDiagram after Lena Voita\u0026rsquo;s NLP Course\nAt each decoder step, the attention mechanism:\nreceives the decoder state $h_t$ and all encoder states $s_1, s_2, \\dots, s_m$;\ncomputes attention scores.\nFor each encoder state $s_k$, the attention mechanism estimates its \u0026ldquo;relevance\u0026rdquo; to the current decoder state $h_t$. Formally, this is done by an attention function, which takes one decoder state and one encoder state and returns a scalar value $\\mathrm{score}(h_t, s_k)$;\ncomputes the attention weights: a probability distribution over the input tokens — that is, a softmax of the attention scores;\ncomputes the output: a weighted sum of the encoder states with the attention weights.\nNow, at each step, the decoder receives not one compressed representation of the whole input, but its own context, tailored to the current step. The bottleneck is gone: the input length is no longer limited by the capacity of a single vector.\nWhat\u0026rsquo;s next: from attention to the Transformer Let us look back at the path we have taken. We started with the encoder-decoder architecture and saw its bottleneck — the entire input is squeezed into a single vector. Attention removed this bottleneck: the decoder itself chooses where to look at each step.\nBut one problem has not gone anywhere. Both the encoder and the decoder are still recurrent networks: the state $s_k$ cannot be computed until $s_{k-1}$ is ready. A sentence of $n$ tokens means $n$ sequential steps that cannot be parallelized. Yet modern GPUs are remarkably good at executing huge matrix operations in parallel — while an RNN forces them to idle, processing tokens strictly one at a time.\nNow let us look at attention from this angle. The scores $\\mathrm{score}(h_t, s_k)$ for different $k$ do not depend on each other — they can all be computed at once. What does that look like concretely? For definiteness, take the simplest variant of the score — the dot product $\\mathrm{score}(h_t, s_k) = h_t^\\top s_k$ (in Bahdanau\u0026rsquo;s work the score is a small neural network, but the idea is the same). Stack all encoder states into a matrix $S$ of size $m \\times d$, one row per token — then all $m$ scores for the current decoder step are obtained by a single matrix-vector product: $S h_t$. Moreover, if all decoder states are known at once — stack them too, into a matrix $H$ of size $n \\times d$ — then the entire score table, $n \\times m$ numbers, is computed by a single matrix product $H S^\\top$. Remember this construction: in the Transformer it will become the main character under the name $QK^\\top$. Admittedly, as long as the decoder is an RNN, the states $h_t$ are still born one at a time, and the full power of this trick remains untapped. But attention itself requires no sequential computation.\nDiagram after Lena Voita\u0026rsquo;s NLP Course\nA caveat on the word \u0026ldquo;steps\u0026rdquo;: a Transformer stage — one layer — does more raw arithmetic than an RNN step, and that amount grows with the sentence length. What stays constant is the number of stages that must run one after another: within a layer no score waits for any other score, everything is computed in parallel, and the number of layers does not depend on the sentence length. FLOPs and sequential steps are different currencies — Part 1 makes this precise.\nHence a daring thought: if attention is so good — maybe throw out the RNN entirely and keep only attention? That is exactly what the authors called the paper that started the Transformer: \u0026ldquo;Attention Is All You Need\u0026rdquo;. How it works, what has to be added to the architecture once recurrence is removed, and why it changed the entire field — that is what the rest of this post is about.\nPart 1 — Why the Transformer at all Part 0 left off at an unsolved problem. The attention mechanism removed the bottleneck: the decoder is no longer forced to reconstruct the entire meaning of the input from a single vector. But both the encoder and the decoder remained recurrent networks, and recurrence has a price of its own. In this part we will figure out what that price is — and why the authors of \u0026ldquo;Attention Is All You Need\u0026rdquo; concluded that the RNN can be dropped altogether.\nTwo problems of RNNs An RNN processes a sequence strictly one token at a time. The state at step $k$ is computed from the previous state and the current input:\n$$ s_k = f(s_{k-1}, x_k). $$Both problems grow out of this formula.\nProblem 1: $O(n)$ sequential steps. Until $s_{k-1}$ is ready, we cannot start computing $s_k$. A sentence of $n$ tokens means $n$ computations that are forced to run one after another. GPUs are built exactly the other way around: thousands of cores tuned for large matrix operations, where everything is computed at once. Each RNN step is a small computation that, on top of that, waits for the previous one; the cores sit idle. We can parallelize across examples in a batch — but not across positions within a sequence. When training on millions of sentence pairs, this becomes the main brake.\nProblem 2: the long path between distant tokens. For the model to learn a dependency between tokens at positions $i$ and $j$, the signal (and, during training, the gradient) has to travel $|i-j|$ steps of recurrence. At every step it gets multiplied by yet another Jacobian — and over dozens of steps it manages to vanish or explode. LSTMs and GRUs alleviate this disease but do not cure it: the path still has length $O(n)$.\n💡 Why does path length matter? Take the sentence: \u0026ldquo;The cat, which had already eaten five fish and still looked hungry, was sitting on the mat.\u0026rdquo; The number of the verb \u0026ldquo;was\u0026rdquo; is determined by the word \u0026ldquo;cat\u0026rdquo; — with a dozen and a half tokens in between. The longer the path a signal must travel from one word to another, the harder it is to learn such a dependency: the useful gradient gets diluted along the way.\nAttention has already solved half of it Notice: the attention from Part 0 is precisely a solution to problem 2, at least half of one. The decoder gets direct access to any encoder state in a single step: the path from \u0026ldquo;cat\u0026rdquo; to «кота» has length $O(1)$, no matter how many tokens separate them.\nBut attention was an add-on to the RNN: recurrent networks still sit in both the encoder and the decoder. Problem 1 has not gone anywhere, and within the encoder distant tokens still communicate through a long chain of states.\nAnd here the authors of the paper asked a question that seems obvious in retrospect but was audacious at the time: if attention transfers information so well — maybe the RNN is not needed at all? Maybe attention is all we need? That is exactly what they called the paper: \u0026ldquo;Attention Is All You Need\u0026rdquo; (Vaswani et al., 2017).\nSelf-attention: every token looks at everyone Suppose we removed the RNN. But the attention from Part 0 operated on top of encoder states built by an RNN. If there is no RNN, where do context-aware token representations come from?\nThe answer is self-attention: the same attention mechanism applied to a sequence with itself — and it is exactly what replaces recurrence. Previously, contextual token representations were built by an RNN passing information along a chain of states; now every token gathers information from all tokens of the same sequence directly (including itself), weighting them by relevance. The word \u0026ldquo;sitting\u0026rdquo; can directly ask: \u0026ldquo;who is the subject here?\u0026rdquo; — and receive a large weight on the token \u0026ldquo;cat\u0026rdquo;.\nAnimation after Lena Voita\u0026rsquo;s NLP Course\nRecall the matrix trick from the end of Part 0: all scores were computed by a single product $H S^\\top$, where $H$ holds the decoder states and $S$ the encoder states. In self-attention both matrices are one and the same sequence $X$ of size $n \\times d$: the everyone-with-everyone score table is, roughly speaking, $X X^\\top$. One matrix multiplication — and all $n^2$ scores are ready. Nobody waits for anybody: the representations of all tokens are updated simultaneously.\nWhat this gives us with respect to the two problems:\nsequential steps — $O(1)$: there is no chain inside a layer; the number of layers in the model is fixed and does not depend on $n$; path between any two tokens — $O(1)$: a single self-attention layer connects any pair directly. These advantages come at a price, and we pay it twice. First, the attention table has size $n \\times n$ — the complexity became quadratic in the sequence length (we will return to this more than once). Second, a weighted sum does not depend on the order of its summands: shuffle the tokens — and self-attention will output the same vectors, only permuted. The RNN knew the token order \u0026ldquo;for free\u0026rdquo;, by construction; self-attention does not see the order at all, and it will have to be brought back into the model by a separate mechanism — positional encoding, which we will discuss in Part 2.\nTwo closing remarks. First, the classical attention from Part 0 does not disappear: in the Transformer it survives under the name cross-attention — the decoder, as in Bahdanau\u0026rsquo;s model, looks at the encoder, only the scheme is slightly generalized. The mechanism is the same in both cases; the only difference is who asks and whom: in cross-attention the decoder queries the encoder, in self-attention the sequence queries itself. Second, on naming: the attention variant where scores are computed via dot products — with additional scaling — is called scaled dot-product attention; what exactly gets scaled and why we actually need three different projections of the input ($Q$, $K$, $V$) will be sorted out in Part 2.\nA back-of-the-envelope count: self-attention vs RNN Let us compare, honestly, one self-attention layer with one RNN layer. A full FLOPs count of all blocks comes in Part 2; here — an order-of-magnitude estimate.\nA self-attention layer. There are $n^2$ token pairs, the score of each pair is a dot product of vectors of dimension $d$, i.e. $d$ multiplications. Total $O(n^2 \\cdot d)$. The weighted summation of values costs the same. All of it is a couple of large matrix multiplications: $O(1)$ sequential operations.\nAn RNN layer. There are $n$ steps, and at each step a state of dimension $d$ is multiplied by a $d \\times d$ matrix: $O(d^2)$ per step. Total $O(n \\cdot d^2)$, and all $n$ steps are strictly sequential.\nLayer Complexity per layer Sequential operations Path length between tokens Self-attention $O(n^2 \\cdot d)$ $O(1)$ $O(1)$ RNN $O(n \\cdot d^2)$ $O(n)$ $O(n)$ Look closely at the first column: self-attention is quadratic in the sequence length $n$, while the RNN is quadratic in the model dimension $d$. An amusing consequence: at a typical $d = 512$ and an ordinary sentence length ($n$ of a few dozen tokens), self-attention is actually cheaper in raw FLOPs.\n💡 So where is the win, if not in FLOPs? In the second and third columns. FLOPs can be multiplied away with GPU cores, but sequential steps cannot: no hardware purchase can beat the law \u0026ldquo;first $s_{k-1}$, then $s_k$\u0026rdquo;. Self-attention turns sequence processing into a few large matrix multiplications — exactly what a GPU does best. Plus the $O(1)$ path between any tokens: long-range dependencies are learned without the gradient decaying along the way.\nThe architecture from a bird\u0026rsquo;s-eye view So, here is what the authors assembled the model from, once left without an RNN.\nAfter Figure 1 of Vaswani et al., 2017 and Lena Voita\u0026rsquo;s NLP Course\nEmbeddings + positional encoding. Tokens are turned into vectors of dimension $d$, and information about position is mixed in — that very fix for self-attention\u0026rsquo;s blindness to order. The encoder is a stack of $N$ identical layers (in the paper $N = 6$). Each layer: self-attention (tokens exchange information) + a small fully connected network FFN (each token is processed independently), plus a few auxiliary mechanisms that stabilize the training of a deep stack of layers — we will cover those separately in Part 2. The decoder is also $N$ layers, but each has three blocks: self-attention over the already-generated tokens (with a mask — no peeking into the future), cross-attention over the encoder output (the direct heir of Bahdanau\u0026rsquo;s attention from Part 0), and the same FFN. The output layer: a linear projection $d \\to |V|$ and a softmax — the probability distribution of the next token. Everything here is exactly as in the language models of Part 0: the same cross-entropy, the same principle of generating one token at a time. Nothing recurrent: every block is either attention or a per-token operation, and everything is computed with matrix multiplications over the whole sequence at once.\nThis is the plan for the next part: we will go through the list bottom-up and assemble every block with our own hands — intuition, formula, PyTorch code, a numpy reference and an honest FLOPs count. We will start with embeddings and positional encoding, and the central place will be taken by scaled dot-product attention.\nPart 2 — The building blocks, with code In Part 1 we looked at the architecture from a bird\u0026rsquo;s-eye view. Now let us come down to earth and assemble every block by hand. The format for every building block is the same: intuition → formula → PyTorch code → a reference implementation in pure numpy for cross-checking → an honest complexity count for that block.\nAll the code of this part can be run without assembling it cell by cell: a ready-made notebook lives in this blog\u0026rsquo;s repository and opens in Colab in one click; no GPU needed.\nWhy a numpy reference? PyTorch modules are convenient, but they hide the details behind library calls. Implementing the same computations \u0026ldquo;by hand\u0026rdquo; in numpy leaves no room for misunderstanding: if two independent pieces of code produce the same numbers, we really do understand what happens inside. We will check every block via np.allclose.\nNotation for the whole series: $n$ — sequence length, $d$ — model dimension, $h$ — number of attention heads, $d_{ff}$ — inner dimension of the FFN, $|V|$ — vocabulary size.\nOne thing we leave out of scope: tokenization. We assume the text has already been turned into a sequence of integer token ids — how exactly text is cut into tokens (usually with the BPE algorithm) is well described, for example, in Lena Voita\u0026rsquo;s course chapter and in Sennrich et al., 2016.\nSet up the environment:\nimport numpy as np import torch import torch.nn as nn import torch.nn.functional as F torch.manual_seed(0) np.random.seed(0) n, d, h, d_ff, vocab_size = 6, 16, 4, 64, 100 # toy sizes for the checks Tokens → embeddings Intuition. The model cannot work with token ids directly — it needs vectors. The embedding table is simply a matrix $E$ of size $|V| \\times d$: one row per vocabulary token. \u0026ldquo;Turning a token into a vector\u0026rdquo; means taking the row of the matrix with that token\u0026rsquo;s id.\nFormula. For a sequence of ids $t_1, \\dots, t_n$:\n$$ X = (E_{t_1}, \\dots, E_{t_n}) \\in \\mathbb{R}^{n \\times d}. $$Code. In PyTorch this is nn.Embedding, in numpy — plain indexing:\nemb = nn.Embedding(vocab_size, d) token_ids = torch.randint(0, vocab_size, (n,)) x_torch = emb(token_ids) # (n, d) E = emb.weight.detach().numpy() # (|V|, d) x_np = E[token_ids.numpy()] # the same operation by hand assert np.allclose(x_torch.detach().numpy(), x_np) print(\u0026#34;embeddings: ok\u0026#34;, x_torch.shape) Complexity. No multiplications: $n$ lookups into the table\u0026rsquo;s rows, an $n \\times d$ matrix as output — that is, $O(n \\cdot d)$ in the size of the result. The main cost of embeddings is not compute but parameters: $|V| \\cdot d$ numbers, a noticeable share of the model when the vocabulary is large.\nA small detail from the original paper: before entering the model, the embeddings are multiplied by $\\sqrt{d}$ — so that their scale is not drowned out by the positional encoding we are about to add to them. We will account for this when assembling the full model in Part 3.\nPositional encoding Intuition. In Part 1 we established: self-attention does not see token order — a weighted sum does not care in what order its summands come. For language this is a catastrophe (\u0026ldquo;the cat ate the mouse\u0026rdquo; and \u0026ldquo;the mouse ate the cat\u0026rdquo; are different events), so position information must be put back into the model explicitly. The original paper\u0026rsquo;s solution: add to each token\u0026rsquo;s embedding a vector that depends only on the token\u0026rsquo;s position in the sequence.\nFormula. The positional vector for position $pos$ is composed of sines and cosines of different frequencies:\n$$ PE_{(pos,\\, 2i)} = \\sin\\!\\left(\\frac{pos}{10000^{2i/d}}\\right), \\qquad PE_{(pos,\\, 2i+1)} = \\cos\\!\\left(\\frac{pos}{10000^{2i/d}}\\right), $$where $i$ runs over pairs of coordinates. Each coordinate pair is its own \u0026ldquo;clock hand\u0026rdquo;: the first pairs rotate fast (change from token to token), the last ones — very slowly. A set of $d/2$ such hands encodes the position unambiguously, just as a clock with a second, minute and hour hand encodes the time.\nCode.\ndef positional_encoding(n, d): pos = torch.arange(n).unsqueeze(1) # (n, 1) i = torch.arange(0, d, 2) # (d/2,) angles = pos / 10000 ** (i / d) # (n, d/2) pe = torch.zeros(n, d) pe[:, 0::2] = torch.sin(angles) pe[:, 1::2] = torch.cos(angles) return pe def positional_encoding_np(n, d): pos = np.arange(n)[:, None] i = np.arange(0, d, 2) angles = pos / 10000 ** (i / d) pe = np.zeros((n, d)) pe[:, 0::2] = np.sin(angles) pe[:, 1::2] = np.cos(angles) return pe pe = positional_encoding(n, d) assert np.allclose(pe.numpy(), positional_encoding_np(n, d), atol=1e-6) print(\u0026#34;positional encoding: ok\u0026#34;, pe.shape) It is worth seeing this matrix with your own eyes once:\nimport matplotlib.pyplot as plt pe_big = positional_encoding(100, 128) plt.figure(figsize=(8, 4)) plt.imshow(pe_big.numpy(), aspect=\u0026#34;auto\u0026#34;, cmap=\u0026#34;RdBu\u0026#34;) plt.xlabel(\u0026#34;vector coordinate\u0026#34;) plt.ylabel(\u0026#34;token position\u0026#34;) plt.colorbar(label=\u0026#34;value\u0026#34;) plt.title(\u0026#34;Positional encoding: fast frequencies on the left, slow on the right\u0026#34;) plt.tight_layout() plt.show() Each column is one coordinate of the vector, each row is a position in the sequence. The left coordinates oscillate fast and distinguish neighboring tokens; the right ones change slowly and encode the \u0026ldquo;coarse scale\u0026rdquo; of the position.\nComplexity. Adding a positional vector to each of the $n$ vectors of size $d$ — $O(n \\cdot d)$. The positional vectors themselves are not trained and are computed once.\nScaled dot-product attention The central block of the whole architecture.\nIntuition. In Part 0, attention worked like this: there is a query (the decoder state), there is a set of candidates (the encoder states); we compute each candidate\u0026rsquo;s relevance to the query and take a weighted sum. Let us generalize this into three roles a vector can play:\nquery $q$ — \u0026ldquo;what I am looking for\u0026rdquo;; key $k$ — \u0026ldquo;by what feature I can be found\u0026rdquo;; value $v$ — \u0026ldquo;what I will hand over if I am chosen\u0026rdquo;. The analogy is a library search: the query is compared against the catalog cards (keys), and based on the comparison a mixture of books (values) is returned. Importantly, the \u0026ldquo;card\u0026rdquo; and the \u0026ldquo;book\u0026rdquo; are different objects: a token can be found by one feature and hand over entirely different information.\nDiagram after Lena Voita\u0026rsquo;s NLP Course\nIn the diagram, each vector receives its three roles through three matrices $W^Q$, $W^K$, $W^V$ — these learnable projections will appear a bit below, in multi-head attention; the attention operation itself, which we are about to write, takes ready-made $Q$, $K$, $V$.\nFormula. We stack the queries, keys and values into matrices $Q$ ($n_q \\times d_k$), $K$, $V$ ($n_k \\times d_k$ and $n_k \\times d_v$):\n$$ \\mathrm{Attention}(Q, K, V) = \\mathrm{softmax}\\!\\left(\\frac{QK^\\top}{\\sqrt{d_k}}\\right)V. $$Here $QK^\\top$ is the score table \u0026ldquo;every query against every key\u0026rdquo;, familiar from Parts 0–1; the softmax turns each row into a distribution of weights, and the multiplication by $V$ takes weighted sums of the values.\n💡 Why $\\sqrt{d_k}$, and not $d_k$ or nothing? The dot product of two random vectors with $d_k$ independent components of zero mean and unit variance has variance $d_k$ (a careful derivation via the variance of a product of independent variables) — that is, the typical spread of the scores grows as $\\sqrt{d_k}$. One might think the softmax cares not about scale but only about the relative magnitudes of the scores. Let us write out the formula and check: for a score vector $s = (s_1, \\dots, s_n)$, the weight of the $i$-th element is\n$$p_i = \\mathrm{softmax}(s)_i = \\frac{e^{s_i}}{\\sum_{k=1}^{n} e^{s_k}}.$$For a shift the intuition is correct: add a constant $c$ to all scores — the numerator and every term of the denominator get multiplied by $e^c$, the common factor cancels, the weights do not change. But multiplying the scores by a constant does not cancel: it stretches the differences between scores, which is what the softmax genuinely depends on — dividing $p_i$ by $p_j$, we see $p_i / p_j = e^{s_i - s_j}$. Variance $d_k$ means the typical difference between two scores is of order $\\sqrt{d_k}$ in absolute value; at $d_k = 512$ that is tens, and already a logit difference of 20 gives a weight ratio of $e^{20} \\approx 5 \\cdot 10^8$. The sign of the difference does not matter. To see this, divide the numerator and denominator of the largest score\u0026rsquo;s weight $s_{\\max}$ by $e^{s_{\\max}}$:\n$$p_{\\max} = \\frac{e^{s_{\\max}}}{\\sum_{k} e^{s_k}} = \\frac{1}{1 + \\sum_{k \\neq \\max} e^{s_k - s_{\\max}}}.$$All the exponents $s_k - s_{\\max}$ are negative, and if they are on the order of tens in absolute value, every term in the denominator is $\\sim e^{-20} \\approx 2 \\cdot 10^{-9}$ — that is, $p_{\\max} \\approx 1$: the largest score takes almost all the weight. And there always is a largest one. That is the problem: at the start the weights are random, so attention \u0026ldquo;locks onto\u0026rdquo; a random token, and the gradients through the rest are practically zero — training can hardly fix an unlucky choice. Dividing by $\\sqrt{d_k}$ returns the differences to a scale of order one. And why not divide by $d_k$? Then typical differences shrink to $1/\\sqrt{d_k}$. The exponential is no longer scary on small arguments — it is linear there, $e^x \\approx 1 + x$ — so the weight ratios are $p_i/p_j \\approx 1 \\pm 1/\\sqrt{d_k}$: at $d_k = 512$ all weights are almost equal, differing by mere percents. The softmax becomes almost uniform, and attention does not choose — it just averages all the values. Note that this failure is milder than the previous one: the gradients here are alive, but the model would have to learn projections with larger norms, pulling the differences back to a scale of one against the imposed division. Dividing by $\\sqrt{d_k}$ is the golden mean: differences of order 1, a softmax that is selective but not locked, and trainable in both directions. Let us check the variance experimentally:\nd_ks = 2 ** np.arange(2, 13) # 4 ... 4096 var_raw, var_scaled = [], [] for d_k in d_ks: q = np.random.randn(2000, d_k) k = np.random.randn(2000, d_k) scores = (q * k).sum(axis=1) var_raw.append(scores.var()) var_scaled.append((scores / np.sqrt(d_k)).var()) plt.figure(figsize=(7, 3.5)) plt.loglog(d_ks, d_ks, \u0026#34;--\u0026#34;, color=\u0026#34;gray\u0026#34;, label=\u0026#34;theory: $\\\\mathrm{var} = d_k$\u0026#34;) plt.loglog(d_ks, var_raw, \u0026#34;o-\u0026#34;, label=\u0026#34;var($q \\\\cdot k$)\u0026#34;) plt.loglog(d_ks, var_scaled, \u0026#34;s-\u0026#34;, label=\u0026#34;var($q \\\\cdot k \\\\,/\\\\, \\\\sqrt{d_k}$)\u0026#34;) plt.xlabel(\u0026#34;$d_k$\u0026#34;) plt.ylabel(\u0026#34;score variance\u0026#34;) plt.grid(alpha=0.3) plt.legend() plt.tight_layout() plt.show() The variance of the \u0026ldquo;raw\u0026rdquo; scores falls exactly on the line $\\mathrm{var} = d_k$; the scaled ones stay around 1 at any dimension.\nCode. The operation itself is four lines:\ndef scaled_dot_product_attention(q, k, v, mask=None): scores = q @ k.transpose(-2, -1) / q.shape[-1] ** 0.5 # (..., n_q, n_k) if mask is not None: scores = scores.masked_fill(~mask, float(\u0026#34;-inf\u0026#34;)) weights = F.softmax(scores, dim=-1) # (..., n_q, n_k) return weights @ v, weights Numpy reference. We write the softmax ourselves (subtracting the maximum for numerical stability):\ndef softmax_np(x, axis=-1): x = x - x.max(axis=axis, keepdims=True) e = np.exp(x) return e / e.sum(axis=axis, keepdims=True) def attention_np(q, k, v, mask=None): scores = q @ np.swapaxes(k, -2, -1) / np.sqrt(q.shape[-1]) if mask is not None: scores = np.where(mask, scores, -np.inf) weights = softmax_np(scores) return weights @ v, weights q, k, v = torch.randn(n, d), torch.randn(n, d), torch.randn(n, d) out_torch, w_torch = scaled_dot_product_attention(q, k, v) out_np, w_np = attention_np(q.numpy(), k.numpy(), v.numpy()) assert np.allclose(out_torch.numpy(), out_np, atol=1e-6) assert np.allclose(w_torch.numpy(), w_np, atol=1e-6) print(\u0026#34;scaled dot-product attention: ok\u0026#34;, out_torch.shape) Complexity (with $n_q = n_k = n$, $d_k = d_v = d$):\n$QK^\\top$: multiplying $(n \\times d)$ by $(d \\times n)$ — $O(n^2 \\cdot d)$; softmax: $n$ elements in each of $n$ rows — $O(n^2)$; the weighted sum $\\mathrm{weights} \\cdot V$: $(n \\times n)$ by $(n \\times d)$ — $O(n^2 \\cdot d)$. Total $O(n^2 \\cdot d)$ — that very quadraticity in sequence length from Part 1. Note: this block has zero parameters — all the learnable weights appear one level up, in multi-head attention.\nMulti-head attention Intuition. One attention — one \u0026ldquo;view\u0026rdquo; of the sequence: one weight table per pair of tokens. But tokens are related in several ways at once: syntactically, semantically, by coreference (\u0026ldquo;he\u0026rdquo; → \u0026ldquo;the cat\u0026rdquo;). The idea of multi-head: run $h$ small attentions in parallel, each in its own subspace of dimension $d/h$ — and let each head learn its own type of relations.\nAnimation after Lena Voita\u0026rsquo;s NLP Course\nFormula. The input $X$ ($n \\times d$) first passes through three learnable linear projections:\n$$ Q = XW^Q, \\quad K = XW^K, \\quad V = XW^V, \\qquad W^Q, W^K, W^V \\in \\mathbb{R}^{d \\times d}. $$Then $Q, K, V$ are sliced along the coordinates into $h$ heads of size $n \\times d/h$; each head computes ordinary scaled dot-product attention (with the scale $\\sqrt{d/h}$ — the head\u0026rsquo;s dimension!), the results are glued back into $n \\times d$ and pass through the output projection $W^O \\in \\mathbb{R}^{d \\times d}$:\n$$ \\mathrm{head}_i = \\mathrm{Attention}(Q_i, K_i, V_i), \\qquad \\mathrm{MultiHead}(X) = \\mathrm{Concat}(\\mathrm{head}_1, \\dots, \\mathrm{head}_h)\\,W^O. $$And here, by the way, is the answer promised in Part 1 to why we need three different projections: they are what creates the query/key/value roles — without them a token would be forced to \u0026ldquo;search\u0026rdquo; and \u0026ldquo;be found\u0026rdquo; with one and the same vector.\n💡 Why do we divide by $\\sqrt{d/h}$ in the softmax, and not by $\\sqrt{d}$? Because the dot products are computed inside a head, between vectors of dimension $d/h$ — so the score variance also grows as $d/h$, and that is what we must normalize by. The general principle: divide by the square root of the dimension of the vectors actually being multiplied.\nCode.\nclass MultiHeadAttention(nn.Module): def __init__(self, d, h): super().__init__() assert d % h == 0 self.d, self.h, self.d_head = d, h, d // h self.w_q = nn.Linear(d, d, bias=False) self.w_k = nn.Linear(d, d, bias=False) self.w_v = nn.Linear(d, d, bias=False) self.w_o = nn.Linear(d, d, bias=False) def split_heads(self, x): # (n, d) -\u0026gt; (h, n, d/h) n = x.shape[0] return x.view(n, self.h, self.d_head).transpose(0, 1) def forward(self, x_q, x_k, x_v, mask=None): q = self.split_heads(self.w_q(x_q)) k = self.split_heads(self.w_k(x_k)) v = self.split_heads(self.w_v(x_v)) out, weights = scaled_dot_product_attention(q, k, v, mask) n = out.shape[1] concat = out.transpose(0, 1).reshape(n, self.d) # glue the heads return self.w_o(concat), weights For self-attention all three inputs are the same sequence: mha(x, x, x). For the cross-attention of Part 1, the queries will come from the decoder, and the keys and values from the encoder: mha(dec, enc, enc) — the code is one and the same.\n💡 Slicing into heads is free. view and transpose in PyTorch (like reshape in numpy) do not copy or move numbers — they only change the tensor\u0026rsquo;s metadata: shapes and memory strides. \u0026ldquo;Sliced into $h$ heads and glued back\u0026rdquo; sounds like work, but in fact it is a reinterpretation of the same memory.\nNumpy reference. We take the weights from the PyTorch module and repeat all the steps by hand (nn.Linear stores its matrix transposed, hence the .T):\ndef multi_head_attention_np(x, w_q, w_k, w_v, w_o, h): n, d = x.shape d_head = d // h def split(m): # (n, d) -\u0026gt; (h, n, d/h) return m.reshape(n, h, d_head).transpose(1, 0, 2) q, k, v = split(x @ w_q), split(x @ w_k), split(x @ w_v) out, _ = attention_np(q, k, v) concat = out.transpose(1, 0, 2).reshape(n, d) return concat @ w_o mha = MultiHeadAttention(d, h) x = torch.randn(n, d) out_torch, _ = mha(x, x, x) weights_np = [m.weight.detach().numpy().T for m in (mha.w_q, mha.w_k, mha.w_v, mha.w_o)] out_np = multi_head_attention_np(x.numpy(), *weights_np, h) assert np.allclose(out_torch.detach().numpy(), out_np, atol=1e-6) print(\u0026#34;multi-head attention: ok\u0026#34;, out_torch.shape) Complexity, line by line. Let us walk down the forward pass:\nStep Operation Complexity Projections $XW^Q, XW^K, XW^V$ $(n \\times d) \\cdot (d \\times d)$, three times $O(n \\cdot d^2)$ Slicing into heads metadata $O(1)$* $Q_i K_i^\\top$ per head $(n \\times d/h) \\cdot (d/h \\times n)$ $O(n^2 \\cdot d/h)$ …over all $h$ heads $O(n^2 \\cdot d)$ softmax per head $n \\times n$ elements $O(n^2)$, over all heads $O(n^2 \\cdot h)$ $\\mathrm{weights}_i \\cdot V_i$ over all heads $O(n^2 \\cdot d)$ Gluing the heads metadata + an $n \\times d$ copy $O(n \\cdot d)$ Output projection $W^O$ $(n \\times d) \\cdot (d \\times d)$ $O(n \\cdot d^2)$ * — when physically contiguous memory is required (as before reshape after transpose), the copy costs $O(n \\cdot d)$; it does not affect the total.\nTotal: $O(n^2 \\cdot d + n \\cdot d^2)$. The first term is the attention itself (quadratic in length), the second is the projections (quadratic in model dimension). Which one dominates depends on the ratio of $n$ and $d$; a detailed bottleneck analysis comes in the summary in Part 3. Note: slicing into $h$ heads does not change the total complexity compared to one big head — the work is the same, just divided into independent chunks.\nMasks Attention as we wrote it allows every token to look at all tokens. That is not always permissible — in the architecture overview in Part 1 we already mentioned that self-attention in the decoder works with a mask that forbids peeking into the future. Masks will start working for real in Part 3, when we assemble the decoder and the training — but the mechanics of the prohibitions belong to the attention block itself (the mask parameter is already in our implementation), so let us sort it out here, and at assembly time apply a ready-made part with one line.\nA mask is a boolean matrix where True = \u0026ldquo;may look\u0026rdquo;. To forbidden positions we assign the score $-\\infty$ before the softmax: after exponentiation they turn into an honest zero weight, and the distribution over the allowed positions stays valid (sums to 1).\nThere are two masks in the Transformer:\nThe padding mask. Sequences in a batch have different lengths, and the short ones are padded with the special token \u0026lt;pad\u0026gt; up to a common length. Looking at \u0026lt;pad\u0026gt; is pointless — it is not text but packing material.\nThe causal mask (decoder). The decoder is trained to predict the next token, and the token at position $t$ has no right to see positions $\u003e t$ — otherwise the prediction task turns into peeking at the answer. The mask is lower-triangular:\ndef causal_mask(n): return torch.tril(torch.ones(n, n, dtype=torch.bool)) Let us draw it for a familiar sentence:\nfrom matplotlib.colors import ListedColormap tokens = [\u0026#34;I\u0026#34;, \u0026#34;saw\u0026#34;, \u0026#34;a\u0026#34;, \u0026#34;cat\u0026#34;, \u0026#34;on\u0026#34;, \u0026#34;a\u0026#34;, \u0026#34;mat\u0026#34;] m = causal_mask(len(tokens)).int().numpy() plt.figure(figsize=(5, 4.5)) plt.imshow(m, cmap=ListedColormap([\u0026#34;#f0f0f0\u0026#34;, \u0026#34;#aec7e8\u0026#34;])) plt.xticks(range(len(tokens)), tokens, rotation=45) plt.yticks(range(len(tokens)), tokens) plt.xlabel(\u0026#34;who is looked at (keys)\u0026#34;) plt.ylabel(\u0026#34;who is looking (queries)\u0026#34;) for i in range(len(tokens)): for j in range(len(tokens)): plt.text(j, i, \u0026#34;✓\u0026#34; if m[i, j] else \u0026#34;✗\u0026#34;, ha=\u0026#34;center\u0026#34;, va=\u0026#34;center\u0026#34;) plt.title(\u0026#34;Causal mask: only the prefix is visible\u0026#34;) plt.tight_layout() plt.show() A row is the \u0026ldquo;observer\u0026rdquo; token, a column is the token it wants to look at. The token \u0026ldquo;cat\u0026rdquo; sees itself and everything before it, but not \u0026ldquo;on\u0026rdquo;, \u0026ldquo;a\u0026rdquo; or \u0026ldquo;mat\u0026rdquo;: at generation time those tokens do not exist yet, and training must proceed under the same conditions.\nLet us check that the mask really works: the weights above the diagonal must become zeros.\nq, k, v = torch.randn(n, d), torch.randn(n, d), torch.randn(n, d) _, w_masked = scaled_dot_product_attention(q, k, v, mask=causal_mask(n)) assert np.allclose(np.triu(w_masked.numpy(), k=1), 0.0) # zeros above the diagonal assert np.allclose(w_masked.numpy().sum(axis=-1), 1.0, atol=1e-6) # rows are distributions print(\u0026#34;causal mask: ok\u0026#34;) Complexity. Applying the mask is an elementwise operation on the score table: $O(n^2)$.\nFeed-forward network (FFN) Intuition. Attention is responsible for the exchange of information between tokens, but by itself it is an almost linear operation (the only nonlinearity is the softmax over the weights). The FFN is the \u0026ldquo;digestion\u0026rdquo; of what was gathered: two linear layers with a nonlinearity in between, applied to each token independently. There is no interaction between positions here — all the communication has already happened in attention.\nFormula.\n$$ \\mathrm{FFN}(x) = \\mathrm{ReLU}(xW_1 + b_1)\\,W_2 + b_2, \\qquad W_1 \\in \\mathbb{R}^{d \\times d_{ff}},\\; W_2 \\in \\mathbb{R}^{d_{ff} \\times d}. $$Note: there is no nonlinearity after the second layer. In the original, $d_{ff} = 4d$ — the network first \u0026ldquo;expands\u0026rdquo; the representation fourfold, then compresses it back.\nCode.\nclass FeedForward(nn.Module): def __init__(self, d, d_ff): super().__init__() self.lin1 = nn.Linear(d, d_ff) self.lin2 = nn.Linear(d_ff, d) def forward(self, x): return self.lin2(F.relu(self.lin1(x))) def feed_forward_np(x, w1, b1, w2, b2): return np.maximum(x @ w1 + b1, 0.0) @ w2 + b2 ffn = FeedForward(d, d_ff) x = torch.randn(n, d) out_torch = ffn(x) out_np = feed_forward_np( x.numpy(), ffn.lin1.weight.detach().numpy().T, ffn.lin1.bias.detach().numpy(), ffn.lin2.weight.detach().numpy().T, ffn.lin2.bias.detach().numpy(), ) assert np.allclose(out_torch.detach().numpy(), out_np, atol=1e-6) print(\u0026#34;FFN: ok\u0026#34;, out_torch.shape) Complexity. First layer: $(n \\times d) \\cdot (d \\times d_{ff})$ — $O(n \\cdot d \\cdot d_{ff})$, plus ReLU $O(n \\cdot d_{ff})$. Second layer — again $O(n \\cdot d \\cdot d_{ff})$. Total $O(n \\cdot d \\cdot d_{ff})$; at $d_{ff} = 4d$ this is $O(n \\cdot d^2)$ with a constant of 8 — in FLOPs the FFN is usually more expensive than the attention projections.\nAdd \u0026amp; Norm Intuition. These are those \u0026ldquo;auxiliary mechanisms\u0026rdquo; from the overview in Part 1 — they do not process information, they help a deep stack of layers train at all. There are two of them, and they are applied around every block (attention or FFN):\nThe residual connection: the block\u0026rsquo;s output is added to its input, $x + \\mathrm{Block}(x)$. The block learns not a \u0026ldquo;new representation from scratch\u0026rdquo; but a correction to the current one; and the gradient gains a short bypass through all the layers — the same short-path logic as with attention, only along the depth. LayerNorm: each vector is normalized over its own $d$ coordinates — subtract the mean, divide by the standard deviation, then apply a learnable scale $\\gamma$ and shift $\\beta$. The normalization is per token: for an $n \\times d$ input, $n$ independent normalizations are performed; every token has its own mean and variance, and tokens do not interact through LayerNorm in any way. This keeps the scale of activations stable from layer to layer. Formula.\n$$ \\mathrm{AddNorm}(x) = \\mathrm{LayerNorm}(x + \\mathrm{Block}(x)), \\qquad \\mathrm{LayerNorm}(z) = \\gamma \\odot \\frac{z - \\mu(z)}{\\sqrt{\\sigma^2(z) + \\varepsilon}} + \\beta. $$Code.\nclass AddNorm(nn.Module): def __init__(self, d): super().__init__() self.norm = nn.LayerNorm(d) def forward(self, x, block_out): return self.norm(x + block_out) def add_norm_np(x, block_out, gamma, beta, eps=1e-5): z = x + block_out mu = z.mean(axis=-1, keepdims=True) var = z.var(axis=-1, keepdims=True) return gamma * (z - mu) / np.sqrt(var + eps) + beta add_norm = AddNorm(d) x, sub = torch.randn(n, d), torch.randn(n, d) out_torch = add_norm(x, sub) out_np = add_norm_np( x.numpy(), sub.numpy(), add_norm.norm.weight.detach().numpy(), add_norm.norm.bias.detach().numpy(), ) assert np.allclose(out_torch.detach().numpy(), out_np, atol=1e-5) print(\u0026#34;Add \u0026amp; Norm: ok\u0026#34;, out_torch.shape) Complexity. Everything is elementwise or a reduction over $d$: $O(n \\cdot d)$ — free compared to attention and the FFN.\n💡 Pre-LN vs post-LN. We described the variant from the original paper: normalization after the residual sum (post-LN). Modern models more often put the LayerNorm before the block (pre-LN): $x + \\mathrm{Block}(\\mathrm{LayerNorm}(x))$ — deep pre-LN models train noticeably more stably and do not require a careful learning-rate warmup; a systematic comparison of the two variants is in Xiong et al., 2020. For our toy transformer the difference is not fundamental, and we stay faithful to the original.\n💡 Why LayerNorm and not BatchNorm? First, the difference itself. Take the input tensor $(b, n, d)$: $b$ examples, $n$ positions, $d$ coordinates per token vector. LayerNorm averages over a single axis — the coordinates: each vector, i.e. the representation of one token, is normalized on its own. That is $b \\cdot n$ independent normalizations with $d$ summands in each; tokens do not interact, and the batch plays no role at all. BatchNorm averages over the two remaining axes at once — examples and positions: for the $j$-th coordinate it computes one shared mean and variance over all $b \\cdot n$ tokens of the batch (one statistic per coordinate — not a separate one per example!). That is $d$ normalizations with $b \\cdot n$ summands in each, and through them every token depends on its random batch neighbors. Both conventions are easy to verify against dniku/dl-norms — minimal from-scratch implementations of the norm layers, tested to match PyTorch\u0026rsquo;s: BatchNorm averages over \u0026ldquo;all dims except C\u0026rdquo;, while LayerNorm averages over the trailing normalized_shape dims — for the transformer input that is (d,), i.e. all dims except batch and position, and it is precisely the \u0026ldquo;except position\u0026rdquo; part that makes it per-token.\nNow, why is BatchNorm a poor fit here? The dependence on batch neighbors has a well-known cost, not specific to text at all: at inference the batch statistics are replaced by running averages accumulated during training, so the behavior at training and at inference slightly diverges. Computer vision lives with this cost successfully: there the statistics are stable from batch to batch, and the batch noise even helps as regularization. With sequences it is worse. Lengths differ — and although padding is easy to filter out of the statistics, the number of real tokens still jumps from batch to batch; moreover, the activation distributions depend heavily on which tokens landed in the batch — rare words produce outliers. As a result, batch statistics in a transformer are so noisy that the running averages become a poor estimate, and the train/inference mismatch starts to genuinely hurt quality — this is exactly what is shown experimentally in Shen et al., 2020. Finally, in autoregressive generation the batch may consist of one example and one new token — there is simply nothing to compute statistics from. LayerNorm is free of all this because it is per-token: identical at training and inference, at any sequence length and any batch size.\nThere is one more variant on the axis menu, sitting between the two: a single normalization over the whole sequence — all $n \\cdot d$ numbers of one example. Note that the arguments above do not touch it: it has no batch statistics, hence no running averages and no train/inference mismatch. It fails for reasons of its own. First, its statistics would include future tokens — for the decoder that is an answer leak bypassing the causal mask, and during generation the statistics (and the representations of already-generated tokens) would change with every new token. Second, it does not solve the actual problem: the scale is evened out only on average over the sentence, while an individual outlier token with a huge norm remains an outlier — whereas it is precisely individual vectors that saturate the softmax in attention. Per-token normalization fixes each vector individually and opens no extra channels of interaction between tokens.\nRecap: all the blocks and their complexity All the building blocks are ready and checked against the reference. A complexity summary:\nBlock Complexity Parameters Embeddings $O(n \\cdot d)$ (lookup) $\\lvert V \\rvert \\cdot d$ Positional encoding $O(n \\cdot d)$ — Multi-head attention $O(n^2 \\cdot d + n \\cdot d^2)$ $4d^2$ Masks $O(n^2)$ — FFN $O(n \\cdot d \\cdot d_{ff})$ $2 d \\cdot d_{ff} + d_{ff} + d$ Add \u0026amp; Norm $O(n \\cdot d)$ $2d$ In the next part we will assemble an encoder layer and a decoder layer out of these blocks, stack them into a full model, train it, and count what it all costs — in FLOPs and, no less importantly, in memory.\nPart 3 — Assembly, training, inference The building blocks of Part 2 are ready and checked against the references. In this part we will assemble a full model out of them, train it on a toy task (for real, on a CPU, in a couple of minutes), learn to generate answers — and pay off the main debt of the series: count the complexity of the Transformer as a whole, including memory and inference.\nAll the code of this part — including the training — can be run end to end: a ready-made notebook lives in this blog\u0026rsquo;s repository and opens in Colab in one click; no GPU needed.\nThe blocks of Part 2 — now with batches In Part 2 we worked with a single sequence for clarity: the input had shape $(n, d)$. For training we need to process a pack of examples at once, so all tensors gain a batch dimension: $(b, n, d)$. The good news: almost nothing needs to change — all the operations are written with @ and work with any number of leading dimensions; only split_heads has changed.\nimport numpy as np import torch import torch.nn as nn import torch.nn.functional as F torch.manual_seed(0) np.random.seed(0) def scaled_dot_product_attention(q, k, v, mask=None): scores = q @ k.transpose(-2, -1) / q.shape[-1] ** 0.5 if mask is not None: scores = scores.masked_fill(~mask, float(\u0026#34;-inf\u0026#34;)) weights = F.softmax(scores, dim=-1) return weights @ v, weights class MultiHeadAttention(nn.Module): def __init__(self, d, h): super().__init__() assert d % h == 0 self.d, self.h, self.d_head = d, h, d // h self.w_q = nn.Linear(d, d, bias=False) self.w_k = nn.Linear(d, d, bias=False) self.w_v = nn.Linear(d, d, bias=False) self.w_o = nn.Linear(d, d, bias=False) def split_heads(self, x): # (b, n, d) -\u0026gt; (b, h, n, d/h) b, n, _ = x.shape return x.view(b, n, self.h, self.d_head).transpose(1, 2) def forward(self, x_q, x_k, x_v, mask=None): q = self.split_heads(self.w_q(x_q)) k = self.split_heads(self.w_k(x_k)) v = self.split_heads(self.w_v(x_v)) out, _ = scaled_dot_product_attention(q, k, v, mask) b, _, n, _ = out.shape return self.w_o(out.transpose(1, 2).reshape(b, n, self.d)) class FeedForward(nn.Module): def __init__(self, d, d_ff): super().__init__() self.lin1 = nn.Linear(d, d_ff) self.lin2 = nn.Linear(d_ff, d) def forward(self, x): return self.lin2(F.relu(self.lin1(x))) def positional_encoding(n, d): pos = torch.arange(n).unsqueeze(1) i = torch.arange(0, d, 2) angles = pos / 10000 ** (i / d) pe = torch.zeros(n, d) pe[:, 0::2] = torch.sin(angles) pe[:, 1::2] = torch.cos(angles) return pe def causal_mask(n): return torch.tril(torch.ones(n, n, dtype=torch.bool)) A mask of shape $(n, n)$ is automatically broadcast over all batches and heads of the score tensor $(b, h, n, n)$ — another place where PyTorch\u0026rsquo;s numpy conventions save code.\nThe encoder layer The recipe from the overview in Part 1: self-attention (tokens exchange information), then the FFN (each token digests what was gathered), and each of the two blocks is wrapped in a residual connection with LayerNorm.\nclass EncoderLayer(nn.Module): def __init__(self, d, h, d_ff): super().__init__() self.attn = MultiHeadAttention(d, h) self.ffn = FeedForward(d, d_ff) self.norm1 = nn.LayerNorm(d) self.norm2 = nn.LayerNorm(d) def forward(self, x, mask=None): x = self.norm1(x + self.attn(x, x, x, mask)) x = self.norm2(x + self.ffn(x)) return x Let us check the whole layer against a numpy reference assembled from the verified parts of Part 2 (repeated here once more, unchanged):\ndef softmax_np(x, axis=-1): x = x - x.max(axis=axis, keepdims=True) e = np.exp(x) return e / e.sum(axis=axis, keepdims=True) def attention_np(q, k, v, mask=None): scores = q @ np.swapaxes(k, -2, -1) / np.sqrt(q.shape[-1]) if mask is not None: scores = np.where(mask, scores, -np.inf) weights = softmax_np(scores) return weights @ v, weights def multi_head_attention_np(x, w_q, w_k, w_v, w_o, h): n, d = x.shape d_head = d // h def split(m): return m.reshape(n, h, d_head).transpose(1, 0, 2) q, k, v = split(x @ w_q), split(x @ w_k), split(x @ w_v) out, _ = attention_np(q, k, v) concat = out.transpose(1, 0, 2).reshape(n, d) return concat @ w_o def feed_forward_np(x, w1, b1, w2, b2): return np.maximum(x @ w1 + b1, 0.0) @ w2 + b2 def add_norm_np(x, block_out, gamma, beta, eps=1e-5): z = x + block_out mu = z.mean(axis=-1, keepdims=True) var = z.var(axis=-1, keepdims=True) return gamma * (z - mu) / np.sqrt(var + eps) + beta def encoder_layer_np(x, layer, h): p = lambda t: t.detach().numpy() a = layer.attn attn = multi_head_attention_np(x, p(a.w_q.weight).T, p(a.w_k.weight).T, p(a.w_v.weight).T, p(a.w_o.weight).T, h) x = add_norm_np(x, attn, p(layer.norm1.weight), p(layer.norm1.bias)) ffn = feed_forward_np(x, p(layer.ffn.lin1.weight).T, p(layer.ffn.lin1.bias), p(layer.ffn.lin2.weight).T, p(layer.ffn.lin2.bias)) return add_norm_np(x, ffn, p(layer.norm2.weight), p(layer.norm2.bias)) d, h, d_ff, n = 64, 4, 256, 10 layer = EncoderLayer(d, h, d_ff) x = torch.randn(1, n, d) assert np.allclose(layer(x).detach().numpy()[0], encoder_layer_np(x[0].numpy(), layer, h), atol=1e-5) print(\u0026#34;encoder layer: ok\u0026#34;) Layer complexity. Attention $O(n^2 \\cdot d + n \\cdot d^2)$ + FFN $O(n \\cdot d \\cdot d_{ff})$ + Add \u0026amp; Norm $O(n \\cdot d)$. At the standard $d_{ff} = 4d$ the whole thing remains $O(n^2 \\cdot d + n \\cdot d^2)$ — the attention formula absorbs the rest.\nThe decoder layer The decoder layer has three blocks, and all three are familiar:\nmasked self-attention — the target-sequence tokens look at their own prefix (the causal mask from Part 2); cross-attention — the heir of Bahdanau\u0026rsquo;s attention from Part 0: queries from the decoder, keys and values from the encoder. The same MultiHeadAttention class, just different inputs; FFN — unchanged. class DecoderLayer(nn.Module): def __init__(self, d, h, d_ff): super().__init__() self.self_attn = MultiHeadAttention(d, h) self.cross_attn = MultiHeadAttention(d, h) self.ffn = FeedForward(d, d_ff) self.norm1 = nn.LayerNorm(d) self.norm2 = nn.LayerNorm(d) self.norm3 = nn.LayerNorm(d) def forward(self, y, enc_out, mask): y = self.norm1(y + self.self_attn(y, y, y, mask)) y = self.norm2(y + self.cross_attn(y, enc_out, enc_out)) y = self.norm3(y + self.ffn(y)) return y The cross-check is organized just like for the encoder — all the parts have already been verified, so we omit the reference code. The complexity is the same up to a constant: cross-attention adds another $O(n^2 \\cdot d + n \\cdot d^2)$ (more precisely $O(n_{tgt} \\cdot n_{src} \\cdot d + \\dots)$ — the score table is now \u0026ldquo;target × source\u0026rdquo;).\nThe full model We assemble everything according to the scheme from Part 1 — here it is again as a reminder:\nAfter Figure 1 of Vaswani et al., 2017 and Lena Voita\u0026rsquo;s NLP Course\nEmbeddings (multiplied by $\\sqrt{d}$ — that debt from Part 2) + positional encoding, a stack of encoder layers, a stack of decoder layers, an output projection into the vocabulary size — we now have code for every element of the scheme.\nclass Transformer(nn.Module): def __init__(self, vocab_size, d, h, d_ff, n_layers, max_len=512): super().__init__() self.d = d self.emb_src = nn.Embedding(vocab_size, d) self.emb_tgt = nn.Embedding(vocab_size, d) self.register_buffer(\u0026#34;pe\u0026#34;, positional_encoding(max_len, d)) self.enc_layers = nn.ModuleList(EncoderLayer(d, h, d_ff) for _ in range(n_layers)) self.dec_layers = nn.ModuleList(DecoderLayer(d, h, d_ff) for _ in range(n_layers)) self.out_proj = nn.Linear(d, vocab_size) def encode(self, src): x = self.emb_src(src) * self.d ** 0.5 + self.pe[: src.shape[1]] for layer in self.enc_layers: x = layer(x) return x def decode(self, tgt, enc_out): mask = causal_mask(tgt.shape[1]) y = self.emb_tgt(tgt) * self.d ** 0.5 + self.pe[: tgt.shape[1]] for layer in self.dec_layers: y = layer(y, enc_out, mask) return self.out_proj(y) def forward(self, src, tgt): return self.decode(tgt, self.encode(src)) The output layer deserves a separate look at its complexity: the projection $(n \\times d) \\cdot (d \\times |V|)$ costs $O(n \\cdot d \\cdot |V|)$. This is an often underestimated expense item. A quick estimate for machine translation: $|V| = 32{,}000$, $d = 512$. The output projection is $d \\cdot |V| \\approx 16.4$ million multiplications per token. And one layer (attention + FFN with $d_{ff} = 4d$) is roughly $10 d^2 + 2nd \\approx 2.7$ million per token at $n = 50$; six layers — about 16 million. That is, the output projection costs roughly as much as all six encoder layers put together.\nHow many parameters? Let us count by the formulas and check against the fact:\nembeddings: $2 \\lvert V \\rvert d$ (source and target); an encoder layer: $4d^2$ (MHA) + $2 d \\cdot d_{ff} + d_{ff} + d$ (FFN with bias) + $2 \\cdot 2d$ (two LayerNorms); a decoder layer: $8d^2$ (two MHAs) + FFN + $3 \\cdot 2d$; the output projection: $d\\lvert V \\rvert + \\lvert V \\rvert$. vocab_size, n_layers = 11, 2 model = Transformer(vocab_size, d, h, d_ff, n_layers) ffn_p = 2 * d * d_ff + d_ff + d enc_p = 4 * d * d + ffn_p + 2 * 2 * d dec_p = 8 * d * d + ffn_p + 3 * 2 * d formula = 2 * vocab_size * d + n_layers * (enc_p + dec_p) + d * vocab_size + vocab_size fact = sum(p.numel() for p in model.parameters()) assert formula == fact print(f\u0026#34;parameters: {fact:,}\u0026#34;) Note: positional encoding is absent from the list — it is not trained (we put it into register_buffer, not into nn.Parameter).\nTraining: string reversal To check whether our assembly \u0026ldquo;comes alive\u0026rdquo;, we do not need a big dataset. Let us take a task that has everything a seq2seq model needs — an input, an output, a dependency of every output token on the input — but trains in minutes on a CPU: reversing a sequence of symbols. Input \u0026ldquo;3 7 1 5\u0026rdquo; → output \u0026ldquo;5 1 7 3\u0026rdquo;. We reserve token 0 for \u0026lt;bos\u0026gt; (begin of sequence — the decoder starts generation from it).\nseq_len, n_symbols = 10, 10 # vocabulary: 0 = \u0026lt;bos\u0026gt;, 1..10 — \u0026#34;letters\u0026#34; def make_batch(batch_size): src = torch.randint(1, n_symbols + 1, (batch_size, seq_len)) tgt = src.flip(1) bos = torch.zeros(batch_size, 1, dtype=torch.long) dec_in = torch.cat([bos, tgt[:, :-1]], dim=1) # decoder input shifted by 1 return src, dec_in, tgt The main idea of the training is visible right here in the code — teacher forcing. The decoder learns to predict the $t$-th token from a prefix of the correct tokens $y_{\\lt t}$ (not from its own past predictions): we feed the target shifted one position to the right (with \u0026lt;bos\u0026gt; at the front) as input, and require the unshifted target as the prediction. Thanks to the causal mask, all positions are trained simultaneously, in one pass — that very parallelization over the sequence for which the whole thing was started in Part 1.\nThe loss function is the cross-entropy from Part 0, applied at every position:\nmodel = Transformer(vocab_size=n_symbols + 1, d=64, h=4, d_ff=256, n_layers=2) opt = torch.optim.Adam(model.parameters(), lr=5e-4) losses = [] for step in range(2000): src, dec_in, tgt = make_batch(128) logits = model(src, dec_in) # (b, n, |V|) loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), tgt.reshape(-1)) opt.zero_grad() loss.backward() opt.step() losses.append(loss.item()) print(f\u0026#34;loss: {losses[0]:.3f} → {losses[-1]:.4f}\u0026#34;) import matplotlib.pyplot as plt plt.figure(figsize=(7, 3.5)) plt.semilogy(losses) plt.xlabel(\u0026#34;training step\u0026#34;) plt.ylabel(\u0026#34;cross-entropy (log scale)\u0026#34;) plt.title(\u0026#34;Training on the string reversal task\u0026#34;) plt.grid(alpha=0.3) plt.tight_layout() plt.show() The rare spikes on the curve are not a bug: with Adam on small models, occasional gradient outliers happen, after which the loss returns to its trajectory within dozens of steps. In large models such spikes are fought with learning-rate warmup and gradient clipping; for our toy task it is enough that training recovers on its own.\nIn the original paper, label smoothing is added to the cross-entropy (the one-hot target distribution is slightly smeared over the vocabulary, $\\varepsilon = 0.1$) — a regularization useful on real data; our task does not need it.\nInference: generation token by token The trained model predicts the next token — but how do we get a whole sequence? The same way the language model generated in Part 0: one token at a time. We start with \u0026lt;bos\u0026gt;, predict the first token, append it to the decoder input, predict the second — and so on. The simplest strategy is greedy decoding: at every step take the most probable token.\n@torch.no_grad() def greedy_decode(model, src, out_len): enc_out = model.encode(src) # encode the source once ys = torch.zeros(src.shape[0], 1, dtype=torch.long) # \u0026lt;bos\u0026gt; for _ in range(out_len): logits = model.decode(ys, enc_out) next_tok = logits[:, -1].argmax(-1, keepdim=True) ys = torch.cat([ys, next_tok], dim=1) return ys[:, 1:] src, _, tgt = make_batch(1000) pred = greedy_decode(model, src, seq_len) exact = (pred == tgt).all(dim=1).float().mean() print(f\u0026#34;accuracy (exact match): {exact:.1%}\u0026#34;) accuracy (exact match): 99.5% Let us look at the predictions with our own eyes. Numeric tokens are a convention; let us display symbols 1–10 as latin letters a–j, so the reversal is literally visible:\ndef to_str(tokens): return \u0026#34;\u0026#34;.join(chr(ord(\u0026#34;a\u0026#34;) + t - 1) for t in tokens.tolist()) for i in range(5): ok = \u0026#34;✓\u0026#34; if (pred[i] == tgt[i]).all() else \u0026#34;✗\u0026#34; print(f\u0026#34;{to_str(src[i])} → {to_str(pred[i])} {ok}\u0026#34;) wrong = (pred != tgt).any(dim=1).nonzero().flatten() if len(wrong) \u0026gt; 0: i = wrong[0] print(f\u0026#34;\\nerror example: {to_str(src[i])} → {to_str(pred[i])}\u0026#34;) print(f\u0026#34;correct: {\u0026#39; \u0026#39; * (seq_len + 3)}{to_str(tgt[i])}\u0026#34;) gjcfhghhfg → gfhhghfcjg ✓ hgfcacbgah → hagbcacfgh ✓ gehbahfgja → ajgfhabheg ✓ ieidhjgfag → gafgjhdiei ✓ jjbbifcije → ejicfibbjj ✓ error example: bhchhicehh → hhecihhcbh correct: hhecihhchb The character of the rare errors is telling in itself: they happen on strings with many repeats of the same symbol (five \u0026ldquo;h\u0026quot;s in the example above). This is no coincidence: identical tokens have identical embeddings, and attention can distinguish them only by the positional component — that very positional encoding from Part 2. Copying in reverse order is, at its core, pure positional work, and it is what fails first.\nGreedy is not the only strategy. A greedy choice is irreversible: one mistake at the beginning ruins the whole tail, even if the model \u0026ldquo;knew\u0026rdquo; a better option. Beam search softens this: at every step not one but $k$ best prefixes are kept (by total log-likelihood), and the final answer is chosen among $k$ complete hypotheses. In translation, typical $k = 4\\text{–}8$ give a noticeable quality boost; for our toy task, where the model is confident in every token, greedy is enough.\nMemory at training time: the hidden cost of quadraticity We have been counting FLOPs all series long, but training has a second currency — memory. For the backward pass we must store the activations of all intermediate layers: gradients are computed by the chain rule from the output back to the input, and every layer needs its forward-pass inputs for that.\nThe hungriest activation is the attention weight table: $n \\times n$ for every head of every layer. A quick estimate for a modest-by-today\u0026rsquo;s-standards config:\nn, h_, N_ = 4096, 8, 12 # context length, heads, layers attn_bytes = N_ * h_ * n * n * 4 # float32 print(f\u0026#34;attention matrices: {attn_bytes / 2**30:.1f} GiB per example\u0026#34;) attention matrices: 6.0 GiB per example Six gigabytes — just for the attention tables, just for one example in the batch, and we have not yet counted the FFN activations or the gradients with the optimizer state. The FLOPs, meanwhile, are perfectly manageable — it is memory we hit first: it grows as $O(N \\cdot h \\cdot n^2)$, and doubling the context length quadruples the bill. That is why long context has historically been a memory problem, not an arithmetic one.\n💡 Modern attention implementations (e.g. FlashAttention) do not materialize the $n \\times n$ table in full: the scores are computed in blocks and immediately folded into the output, and on the backward pass the missing parts are recomputed. The quadraticity of FLOPs stays; the quadraticity of memory goes away. This is beyond our series, but it is useful to know that \u0026ldquo;quadratic memory\u0026rdquo; is not a verdict.\nKV-cache: do not recompute the past Let us look closely at our greedy_decode. At step $t$ we call model.decode on the entire prefix of $t$ tokens: we recompute Q, K, V for tokens that have not changed, and rebuild the $t \\times t$ attention table. The cost of step $t$ is $O(t^2 \\cdot d)$, and of the whole generation of $T$ tokens — $O(T^3 \\cdot d)$. Wasteful: from step to step, exactly one token in the prefix changes.\nThe key observation: for its attention, the new token needs its own query $q_t$ — and the keys and values of all the previous tokens, which we already computed at past steps. So they can be cached: we store the accumulated $K$ and $V$ (a pair per layer), at each step compute the projections only for the one new token and append them to the cache. That is the KV-cache. The cost of step $t$ drops to $O(t \\cdot d)$: one query against $t$ keys. The whole generation of $T$ tokens gets cheaper from $\\sum_t O(t^2 d) = O(T^3 \\cdot d)$ to $\\sum_t O(t \\cdot d) = O(T^2 \\cdot d)$ — an order of $T$ better; the quadraticity remains (the last token still looks at all its predecessors), but the cubic wastefulness is gone.\nLet us verify, at the level of a single attention layer, that incremental generation with a cache gives exactly the same result as full recomputation with a causal mask:\nmha = MultiHeadAttention(d, h) x = torch.randn(1, 20, d) with torch.no_grad(): full = mha(x, x, x, causal_mask(20)) # full recomputation k_cache, v_cache, outs = None, None, [] for t in range(x.shape[1]): x_t = x[:, t : t + 1] # only the new token q = mha.split_heads(mha.w_q(x_t)) k_new = mha.split_heads(mha.w_k(x_t)) v_new = mha.split_heads(mha.w_v(x_t)) k_cache = k_new if k_cache is None else torch.cat([k_cache, k_new], dim=2) v_cache = v_new if v_cache is None else torch.cat([v_cache, v_new], dim=2) out, _ = scaled_dot_product_attention(q, k_cache, v_cache) outs.append(mha.w_o(out.transpose(1, 2).reshape(1, 1, mha.d))) incremental = torch.cat(outs, dim=1) assert torch.allclose(full, incremental, atol=1e-5) print(\u0026#34;KV-cache: matches full recomputation\u0026#34;) Notice: the causal mask was not needed in the incremental variant. The cache is the past — the new token physically cannot look into the future, because the future keys do not exist yet. The mask is only needed during training, when all positions are computed simultaneously.\nWe pay for the speedup with memory: the cache holds $2 \\cdot N \\cdot t \\cdot d$ numbers and grows linearly with the generation length — in long dialogues with large models it is precisely the KV-cache that occupies a notable share of GPU memory. This trade — recomputation for storage — is an engineering classic, and in Part 2 we already saw its mirror image: FlashAttention trades storage for recomputation.\n📦 Transformer complexity: the whole series in one place The promised summary — all the answers in one place.\nFormulas by block (derived in Part 2):\nBlock Complexity Embeddings $O(n \\cdot d)$ Positional encoding $O(n \\cdot d)$ Multi-head attention $O(n^2 \\cdot d + n \\cdot d^2)$ FFN $O(n \\cdot d \\cdot d_{ff})$, at $d_{ff}=4d$ — $O(n \\cdot d^2)$ Add \u0026amp; Norm $O(n \\cdot d)$ Output projection $O(n \\cdot d \\cdot \\lvert V \\rvert)$ Total for a model of $N$ layers:\n$$ \\underbrace{O\\big(N \\cdot (n^2 d + n d^2)\\big)}_{\\text{layers}} \\;+\\; \\underbrace{O(n \\cdot d \\cdot |V|)}_{\\text{output}}. $$Comparison with the RNN (derived in Part 1):\nFLOPs per layer Sequential operations Path between tokens Self-attention $O(n^2 \\cdot d)$ $O(1)$ $O(1)$ RNN $O(n \\cdot d^2)$ $O(n)$ $O(n)$ The Transformer\u0026rsquo;s win is not in FLOPs (there may even be more of them) but in the two right columns: parallelism and short paths for the gradient.\nWhere is the bottleneck? Look at $O(n^2 d + n d^2)$ and compare $n$ with $d$:\n$n \\gg d$ (long contexts): $n^2 d$ dominates — attention. Hence the whole industry of \u0026ldquo;cheap\u0026rdquo; attention. $n \\approx d$: both terms are $O(n^3)$, no single bottleneck. $n \\ll d$ (typical translation: $n \\sim 50$, $d = 512\\text{–}1024$): $n d^2$ dominates — the linear projections and the FFN, not attention at all. This is worth spelling out because it contradicts a popular myth: on short sequences the quadraticity of attention is practically invisible; the bulk of the FLOPs goes into ordinary matrix multiplications. \u0026ldquo;Complexity\u0026rdquo; is three different questions:\nTraining FLOPs: the forward pass is $O(N(n^2 d + n d^2) + n d |V|)$ per example; the backward is roughly twice the forward; in total, training costs about three forwards. Training memory: the activations for the backward pass, the worst being the attention tables $O(N \\cdot h \\cdot n^2)$. Quadratic memory hits the hardware sooner than quadratic FLOPs. Inference cost: naive generation — $O(t^2 d)$ per token; with a KV-cache — $O(t \\cdot d)$ per token at the price of a cache of size $O(N \\cdot t \\cdot d)$, growing with the length of the dialogue. Series finale We have walked the whole path: from seq2seq and language models (Part 0), through the motivation of \u0026ldquo;why throw out the RNN\u0026rdquo; (Part 1) and the hand-assembly of every block with numpy checks (Part 2) — to a working model that trained, generates, and fits the formulas we derived (Part 3). The Transformer has stopped being a black box: it is a dozen small parts, each of which fits into 10–20 lines of code and one complexity formula.\nWhere to go next if you want to dig deeper: BPE tokenization, decoder-only architectures (GPT), efficient attention (FlashAttention and relatives), quantization, and everything that makes inference cheap. But those are other stories.\nAcknowledgments. The diagrams and animations in this series are redrawn after the wonderful illustrations in Lena Voita\u0026rsquo;s magical NLP Course For You — a course that inspired much of this series in the first place. If after these posts you want a broader and deeper dive into NLP, that is the place to go. Thank you, Lena! 💛\n","permalink":"https://jen1995.github.io/posts/transformers/","summary":"The full series in one post: why attention replaced recurrence; every building block of the Transformer — intuition, formula, PyTorch code, a numpy reference and an honest FLOPs count; then the assembled model, trained on a toy task, with the memory costs and the KV-cache explained.","title":"Transformers from Scratch"},{"content":" This post is an English translation of my chapter on VAE from the Machine Learning Handbook by Yandex School of Data Analysis (originally in Russian). The text follows the latest revised version of the chapter; the figures are the original English-language versions.\nIntroduction There is a fairly broad area of machine learning devoted to training generative models. Their goal is to learn the distribution from which the objects of the training set could have been sampled.\nOnce trained, a generative model can sample new objects from the learned distribution — objects that do not belong to the original data. Most often this is associated with generating new images: from pictures of handwritten digits to face swapping in videos.\nThe model this post is about is called the variational autoencoder, or VAE. It belongs to the family of generative models. Here is a quick roadmap of what lies ahead.\nThe sections \u0026ldquo;Problem setup\u0026rdquo; and \u0026ldquo;Training a VAE\u0026rdquo; describe how a VAE is built and trained in its classical form. These two sections are enough to get a general understanding of VAE.\nThe section \u0026ldquo;Key papers\u0026rdquo; is not required for a first understanding, but may be interesting to those who want to learn about the foundational works that grew out of the VAE idea.\nProblem setup Imagine that we need to draw a horse. How would we go about it?\nWe would probably first sketch the overall silhouette of the horse, its size and pose, and then start adding details: the mane, the tail, the hooves, picking a coat color, and so on. It seems that while learning to draw, we learn to identify a basic set of factors that matter most for generating a new image — overall silhouette, size, color and the like — and during drawing we simply plug in particular values of those factors.\nAt the same time, identical combinations of the same factors can lead to different pictures — after all, you most likely cannot draw something exactly the same way twice.\nLet us try to formalize the process described above. Suppose we have a dataset $D$ living in a high-dimensional space of raw data $X^N$ — the objects we wish to generate — and a lower-dimensional space $Z^M$ of hidden (latent) variables, which encode the hidden factors in the data. The generative process then consists of two consecutive stages (see the picture below):\nSample $z\\in{Z^M}$ from the distribution $p(z)$\nSample $x\\in{X^N}$ from the distribution $p(x\\mid z)$\nImage source\nThat is, thinking in terms of drawing pictures of horses, we first mentally sample some $z$ (size, shape, color), then fill in all the necessary details — that is, we sample from the distribution $p(x\\mid z)$ — and in the end we hope that the result will resemble a horse.\nThus, building a generative model in our case means being able to sample, via the two-stage process described above, objects that are close to the objects of the training set $D$.\nMore formally, we would like our model to maximize the likelihood $p(x)$ of the elements of the training set $D$ under the described generation procedure:\n$$p(x)=\\int_{Z^M}{p(x\\mid z)p(z)dz}\\to{\\max} $$Assume that the joint distribution $p(x,z)$ is parameterized by some parameter $\\theta\\in{\\Theta}$ and is expressed by a function continuous in $\\theta$ for every fixed $x$ and $z$:\n$$p_{\\theta}(x,z)=p(x,z\\mid\\theta)\\in{C(\\Theta)} $$Then\n$$p_{\\theta}(x,z)=p_{\\theta}(x\\mid z)p_{\\theta}(z), $$and we can write down the following optimization problem:\n$$p_{\\theta}(x)=\\int_{Z^M}{p_{\\theta}(x\\mid{z})p_{\\theta}(z)dz}\\to\\max_{\\theta\\in\\Theta}\\tag{1} $$Solving it, we obtain our generative model.\nRemark 1. After the analogy with learning to draw, one might mistakenly conclude that latent variables always carry some nicely interpretable meaning. In practice this does not have to be the case: the latent variables we end up finding may or may not have a simple interpretation. The explanations above were primarily meant to illustrate the notion of \u0026ldquo;latent variables.\u0026rdquo;\nRemark 2. It might seem that $p(x)$ is somehow already known to us, and then it is unclear why we need all these complications with latent variables and integrals. Indeed, we can actually build a statistical estimate $\\hat{p}(x)$ from the data $D$ and even try to generate new data with such models (as done, for example, here). But statistical methods come with various limitations, the most serious of which appears to be the curse of dimensionality: the more dimensions your data has, the more diverse examples you need to build an adequate estimate $\\hat{p}(x)$. We will talk about the curse of dimensionality in a bit more detail later.\nRemark 3. Another natural question: why introduce latent variables at all, model the joint distribution $p(x,z)$, and define the target distribution $p(x)$ as the marginalization of $p(x,z)$ over $z$? Why should this approach work in the first place? The answer is that even with relatively simple expressions for $p(z)$ and $p(x \\mid z)$, one can describe a rather complex distribution $p(x)$ — which is illustrated quite vividly in the example below.\nExample: mixture of Gaussians Imagine you have a table with a finite number of rows, where the $k$-th row contains two numbers — the mean $\\mu_k$ and the variance $\\sigma_k^2$ of a normal distribution. Suppose a discrete distribution $p(z)$ is defined over the row indices of this table, such that:\n$$p(z=k) = \\lambda_k $$Say we sampled an index $k$, took the distribution parameters from the corresponding row, and sampled an object $x$ with those parameters. The distribution from which $x$ was obtained equals:\n$$p(x \\mid z=k) = \\mathcal{N}(\\mu_k, \\sigma_k^2) $$The distribution $p(x)$ is obtained by marginalizing the joint distribution $p(x,z)$ over $z$:\n$$\\begin{aligned} p(x) \u0026= \\sum_{k=1}^K p(x,z=k) = \\sum_{k=1}^K p(x \\mid z=k) p(z=k) \\\\ \u0026= \\sum_{k=1}^K \\lambda_k \\mathcal{N}(\\mu_k, \\sigma_k^2) \\end{aligned}$$It turns out that $p(x)$ is described by a mixture of Gaussians and has a more complex form than $p(z)$ and $p(x \\mid z)$:\nImage source\nClearly, the more Gaussians in our sum, the more complex the shape of $p(x)$ can be. So, with simple $p(z)$ and $p(x \\mid z)$, we can model complex multimodal distributions. Now imagine that the prior distribution $p(z)$ takes not discrete but continuous values. Consider, for example, the following case:\n$$p(z) = \\mathcal{N}(0,1) $$$$p(x \\mid z) = \\mathcal{N}(\\mu(z), \\sigma^2(z)) $$The distribution $p(x)$, analogously to the discrete-prior case, is obtained by integrating $p(x,z)$ over $z$ and is, in a sense, an \u0026ldquo;infinite\u0026rdquo; mixture of Gaussians:\nImage source\nThat settles it. In the next section we continue with the optimization problem we started discussing above — stay tuned!\nTraining a VAE Before trying to solve optimization problem $(1)$, let us think about how we could even compute such an integral. The first thing that comes to mind is to approximate it with the Monte Carlo method:\n$$\\begin{aligned} p_\\theta(x) \u0026= \\int\\limits_{Z^M} p_\\theta(x \\mid z) p_\\theta(z) dz = \\mathbb{E}_{z \\sim p_\\theta(z)} [p_\\theta(x \\mid z)] \\\\ \u0026\\approx \\frac{1}{K} \\sum_k p_\\theta(x \\mid z_k), \\end{aligned}$$where in the last step we use samples $z_k \\sim p_\\theta(z)$. However, if $z \\in Z^M$ and $M$ is large enough, we run into the curse of dimensionality — the number of samples needed to cover $Z^M$ well grows exponentially with $M$:\nImage source\nIs there a way to somehow reduce the number of samples needed to compute $(1)$? As it happens, it is often the case that far from all possible $z$ map into elements of $D$, and the contribution of most $z$ to the estimate of $p_\\theta(x \\mid z)$ is practically zero. This suggests that for each $x$ we could benefit from knowing the distribution $q(z \\mid x)$ of those $z$ that are preimages of $x$. We may assume that the distribution $q$ is parameterized by some family of parameters $\\Phi$:\n$$q(z \\mid x) = q_\\phi(z \\mid x), \\quad \\phi \\in \\Phi $$Knowing the distribution $q_\\phi(z \\mid x)$, we could sample only from it rather than from the whole $p_\\theta(z)$, and if the distribution $q$ turns out to be good enough, the number of required samples will drop significantly.\nWe will discuss how to build $q_\\phi$ later. For now, note that the processes of sampling from the distributions $q_\\phi(z \\mid x)$ and $p_\\theta(x \\mid z)$ are mutually inverse: the first maps dataset elements into a subset of the latent space $Z^M$, i.e. acts as an encoder, while the second maps latent variables into a subset of $X^N$, i.e. acts as a decoder:\nImage source\nSince both of these distributions will take part in training the VAE, an analogy arises between VAE and autoencoder models, which have a similar structure.\nDeriving the loss function We now have everything ready to write down the general form of the loss function for training a variational autoencoder. Recall that we train the model by maximizing the likelihood $p_\\theta(x)$ with respect to $\\theta$. For convenience, we switch to the log-likelihood:\n$$\\log p_\\theta(x) = \\log \\int_{Z^M} p_\\theta(x \\mid z) p_\\theta(z) dz \\to \\max_{\\theta \\in \\Theta} $$Optimizing this expression directly is hard because of the curse of dimensionality discussed in the previous section. To defeat it, we would like to replace sampling from the prior distribution $p_\\theta(z)$ with sampling from $q_\\phi(z \\mid x)$, which requires a certain trick. For any $q_\\phi(z \\mid x)$ that is nonzero for all $z \\in Z^M$, we can write out the following chain of equalities:\n$$\\log p_\\theta(x) = \\mathbb{E}_{q_\\phi(z \\mid x)} [\\log p_\\theta(x)] = $$$$= \\mathbb{E}_{q_\\phi(z \\mid x)} \\left[ \\log \\left( \\frac{p_\\theta(x,z)}{p_\\theta(z \\mid x)} \\right) \\right] = $$$$= \\mathbb{E}_{q_\\phi(z \\mid x)} \\left[ \\log \\left( \\frac{p_\\theta(x,z)}{q_\\phi(z \\mid x)} \\frac{q_\\phi(z \\mid x)}{p_\\theta(z \\mid x)} \\right) \\right] = $$$$\\begin{aligned} ={} \u0026 \\underbrace{\\mathbb{E}_{q_\\phi(z \\mid x)} \\left[ \\log \\left( \\frac{p_\\theta(x,z)}{q_\\phi(z \\mid x)} \\right) \\right]}_{\\mathcal{L}_{\\theta,\\phi}(x) \\text{ (ELBO)}} \\\\ \u0026 + \\underbrace{\\mathbb{E}_{q_\\phi(z \\mid x)} \\left[ \\log \\left( \\frac{q_\\phi(z \\mid x)}{p_\\theta(z \\mid x)} \\right) \\right]}_{D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z \\mid x))} \\end{aligned}$$The second term in the last equality is the $KL$ divergence between $q_\\phi(z \\mid x)$ and $p_\\theta(z \\mid x)$, which, as is well known, is nonnegative:\n$$D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z \\mid x)) \\ge 0 $$And the first term is the evidence lower bound (ELBO):\n$$\\mathcal{L}_{\\theta,\\phi}(x) = \\mathbb{E}_{q_\\phi(z \\mid x)} [\\log p_\\theta(x,z) - \\log q_\\phi(z \\mid x)] = $$$$= \\underbrace{\\mathbb{E}_{q_\\phi(z \\mid x)} [\\log p_\\theta(x \\mid z)]}_{\\text{reconstruction loss}} - \\underbrace{D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z))}_{\\text{regularization term}} $$The first term in the last step is called the reconstruction loss, since it measures how well the decoder reconstructs the object $x$ from its latent representation $z$. The second one plays the role of a regularization term and pushes the distribution produced by the encoder to be closer to the prior distribution.\nSince the $KL$ divergence is nonnegative, the ELBO is a lower bound on the log-likelihood of the data:\n$$\\mathcal{L}_{\\theta,\\phi}(x) = \\log p_\\theta(x) - D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z \\mid x)) \\le \\log p_\\theta(x) $$Let us take a closer look at the equalities we have written down.\nThe function $\\mathcal{L}_{\\theta,\\phi}$ can be optimized by gradient descent (SGD), once we choose a convenient form for $p_\\theta(x \\mid z)$, $q_\\phi(z \\mid x)$ and $p_\\theta(z)$. By maximizing $\\mathcal{L}_{\\theta,\\phi}$, we increase $\\log p_\\theta(x)$, thereby improving our generative model. We will discuss ELBO optimization with SGD in detail in the next section.\nBy maximizing $\\mathcal{L}_{\\theta,\\phi}$, we simultaneously minimize $D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z \\mid x))$. The distribution $p_\\theta(z \\mid x)$ estimates which $z$ the object $x$ could have been generated from, and it is not known to us in advance. But if we pick a sufficiently large model for $q_\\phi(z \\mid x)$, then during optimization $q_\\phi(z \\mid x)$ can get very close to $p_\\theta(z \\mid x)$, and then we will be directly optimizing $\\log p_\\theta(x)$. As a pleasant bonus, to estimate the distribution of preimages of $x$ we can use $q_\\phi(z \\mid x)$ instead of the intractable $p_\\theta(z \\mid x)$. That is, $q_\\phi$, which we introduced in the derivation as an arbitrary distribution, will indeed play the role of the model\u0026rsquo;s encoder.\nAn alternative derivation of the ELBO In the reasoning above, introducing $q_\\phi(z \\mid x)$ might have seemed rather formal. So here we present another approach to deriving the expression for the ELBO, which may feel more natural. It consists of successively applying a technique called importance sampling and Jensen\u0026rsquo;s inequality.\nImage source\nIn many practical problems, a situation arises where we want to compute $\\mu = \\mathbb{E}[f(X)]$, but $f(x)$ is close to zero outside some region $A$, while the probability of landing in this important region is very small: $P(X \\in A) \\approx 0$.\nThe set $A$ may either have too small a measure, or sit in the tail of the distribution of the random variable $X$. Ordinary Monte Carlo sampling may generate almost no examples that fall into the set $A$. Problems of this kind are quite common in high-energy physics, Bayesian inference, forecasting of natural hazards, and many other areas.\nA rather intuitive solution is to try to artificially increase the share of important examples among all the rest. This can be done by using a distribution that gives more weight to examples from the important region. Hence the name of the method — importance sampling.\nSo, suppose our task is to compute the expectation $\\mu = \\mathbb{E}_p[f(x)] = \\int_{\\mathcal{D}} f(x) p(x) dx$, where\n$p$ is a probability density on the set $\\mathcal{D} \\in \\mathbb{R}^d$,\n$f$ is some integrable function.\nLet $q$ be a probability density function, defined and positive on $\\mathcal{D}$, that allows us to sample examples from some narrow subset of interest. Our task is to switch from sampling from $p$ to sampling from $q$ when estimating $\\mu$. Since the mean $\\mathbb{E}_q[f(x)]$ is, generally speaking, not equal to $\\mu$, we write the following:\n$$\\begin{aligned} \\mu = \\mathbb{E}_p[f(x)] \u0026= \\int_{\\mathcal{D}} f(x) p(x) dx = \\int_{\\mathcal{D}} \\frac{f(x)p(x)}{q(x)} q(x) dx \\\\ \u0026= \\mathbb{E}_q \\left[ \\frac{f(x)p(x)}{q(x)} \\right] \\end{aligned}$$The original density $p$ is called the nominal distribution, and the density $q$ the importance distribution. The likelihood ratio $\\frac{p(x)}{q(x)}$ compensates for the bias introduced when switching from $p$ to $q$.\nLet us also recall the statement of Jensen\u0026rsquo;s inequality for random variables: if $\\xi$ is a random variable with finite expectation and $g(x)$ is a convex function, then:\n$$\\mathbb{E}[g(\\xi)] \\ge g(\\mathbb{E}[\\xi]) $$Now back to the original problem. Again, for any $q_\\phi(z \\mid x)$ that is nonzero for all $z \\in Z^M$, we can write:\n$$\\log p_\\theta(x) = \\log \\int_{Z^M} p_\\theta(x \\mid z) p_\\theta(z) dz = $$$$= \\log \\mathbb{E}_{p_\\theta(z)} [p_\\theta(x \\mid z)] = $$$$= \\log \\mathbb{E}_{q_\\phi(z \\mid x)} \\left[ \\frac{p_\\theta(x \\mid z) p_\\theta(z)}{q_\\phi(z \\mid x)} \\right] \\ge $$$$\\ge \\mathbb{E}_{q_\\phi(z \\mid x)} \\log \\left[ \\frac{p_\\theta(x \\mid z) p_\\theta(z)}{q_\\phi(z \\mid x)} \\right] = $$$$= \\mathbb{E}_{q_\\phi(z \\mid x)} \\log [p_\\theta(x \\mid z)] - \\mathbb{E}_{q_\\phi(z \\mid x)} \\log \\left[ \\frac{q_\\phi(z \\mid x)}{p_\\theta(z)} \\right] = $$$$= \\mathbb{E}_{q_\\phi(z \\mid x)} \\log [p_\\theta(x \\mid z)] - D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z)) $$As a result of these manipulations we have, as you can see, once again obtained the expression for the ELBO. In the third step we applied importance sampling, and in the fourth — Jensen\u0026rsquo;s inequality for $g(x)=\\log(x)$ (note that for the concave $\\log$ the inequality flips).\nThe downside of this approach is that, unlike the previous method, it does not let us write out an explicit formula for the gap between $\\log p_\\theta(x)$ and the ELBO:\n$$\\log p_\\theta(x) - \\text{ELBO} = D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z \\mid x)) $$On the other hand, this derivation follows naturally from more general methods, without requiring artificial tricks.\nTraining a VAE with gradient descent An important property of the ELBO is that it can be optimized by gradient descent with respect to the parameters $\\phi$ and $\\theta$. If the objects of the dataset $D$ are independent and identically distributed, then $\\mathcal{L}_{\\theta,\\phi}(\\mathcal{D})$ is written as a sum (or mean) of the values $\\mathcal{L}_{\\theta,\\phi}(x)$ over the objects $x \\in D$:\n$$\\mathcal{L}_{\\theta,\\phi}(\\mathcal{D}) = \\sum_{x \\in \\mathcal{D}} \\mathcal{L}_{\\theta,\\phi}(x) $$The values $\\mathcal{L}_{\\theta,\\phi}(x)$ and their gradients $\\nabla \\mathcal{L}_{\\theta,\\phi}(x)$ cannot be computed exactly in the general case, but we can obtain unbiased estimates of them, which will let us use stochastic gradient descent.\nAn estimate for the gradient with respect to the parameters $\\theta$ is easy to get:\n$$\\nabla_\\theta \\mathcal{L}_{\\theta,\\phi}(x) = \\nabla_\\theta \\mathbb{E}_{q_\\phi(z \\mid x)} [\\log p_\\theta(x,z) - \\log q_\\phi(z \\mid x)] = $$$$= \\mathbb{E}_{q_\\phi(z \\mid x)} [\\nabla_\\theta (\\log p_\\theta(x,z) - \\log q_\\phi(z \\mid x))] = $$$$= \\mathbb{E}_{q_\\phi(z \\mid x)} [\\nabla_\\theta \\log p_\\theta(x,z)] \\approx $$$$\\approx \\frac{1}{K} \\sum_k \\nabla_\\theta \\log p_\\theta(x, z_k), $$where in the last line $z_k \\sim q_\\phi(z \\mid x)$. However, an estimate for the gradient with respect to the parameters $\\phi$ is harder to obtain, because they also take part in the sampling:\n$$\\nabla_\\phi \\mathcal{L}_{\\theta,\\phi}(x) = \\nabla_\\phi \\mathbb{E}_{q_\\phi(z \\mid x)} [\\log p_\\theta(x,z) - \\log q_\\phi(z \\mid x)] \\neq $$$$\\neq \\mathbb{E}_{q_\\phi(z \\mid x)} [\\nabla_\\phi (\\log p_\\theta(x,z) - \\log q_\\phi(z \\mid x))] $$In the general case this problem is unsolvable. However, some distributions admit the reparameterization trick: representing the variable $z$ as an invertible differentiable function of random noise, the parameters $\\phi$, and the variable $x \\in D$:\n$$z = g(\\varepsilon, \\phi, x) $$Here the distribution $\\varepsilon \\sim p_\\varepsilon$ does not depend on $\\phi$ or $x$. For example, let $\\varepsilon \\sim \\mathcal{N}(0, I)$. Then $g$ can have the following form:\n$$z = g(\\varepsilon, \\phi, x) = \\mu_\\phi(x) + \\varepsilon \\cdot \\sigma_\\phi(x) \\sim \\mathcal{N}(\\mu_\\phi(x), \\sigma_\\phi^2(x)) $$After this substitution we can obtain an estimate of the gradient with respect to $\\phi$:\n$$\\nabla_\\phi \\mathcal{L}_{\\theta,\\phi}(x) = \\nabla_\\phi \\mathbb{E}_{q_\\phi(z \\mid x)} [\\log p_\\theta(x,z) - \\log q_\\phi(z \\mid x)] = $$$$= \\nabla_\\phi \\mathbb{E}_{p_\\varepsilon} [\\log p_\\theta(x, g(\\varepsilon, \\phi, x)) - \\log q_\\phi(g(\\varepsilon, \\phi, x) \\mid x)] = $$$$= \\mathbb{E}_{p_\\varepsilon} [\\nabla_\\phi (\\log p_\\theta(x, g(\\varepsilon, \\phi, x)) - \\log q_\\phi(g(\\varepsilon, \\phi, x) \\mid x))] \\approx $$$$\\approx \\frac{1}{K} \\sum_k \\nabla_\\phi (\\log p_\\theta(x, g(\\varepsilon_k, \\phi, x)) - \\log q_\\phi(g(\\varepsilon_k, \\phi, x) \\mid x)), $$where in the last line $\\varepsilon_k \\sim p_\\varepsilon$. The reparameterization trick is well illustrated by the following picture:\nImage source\nHere $f$ is the loss function. The values of $f$ are the same in both diagrams, but in the left one the gradients with respect to $\\phi$ cannot be computed, since we cannot differentiate through the random variable $z$.\nIn the right diagram, however, the source of randomness moves into the input data thanks to reparameterization, and the gradients are computed with respect to deterministic variables. We thus arrive at a setup typical for optimization with SGD: there we approximate the gradient of the loss function over random batches of input data, and here the role of random batches is played jointly by batches of the variables $x$ and of the random variables $\\varepsilon$.\nBesides the normal distribution, there are quite a few examples of distributions that admit reparameterization. They can be found in the VAE paper, in the section \u0026ldquo;The reparameterization trick\u0026rdquo;. Most VAE implementations, however, use the normal distribution.\nIn the end, the rough algorithm for training a VAE looks like this:\ndataset = np.array(...) epsilon = RandomDistribution(...) # Encoder q_phi(z|x) — a neural network with parameters phi encoder = Encoder() # Decoder p_theta(x|z) — a neural network with parameters theta decoder = Decoder() for step in range(max_steps): # Sample a batch of input data and of random noise batch_x = sample_batch(dataset) batch_noise = sample_batch(epsilon) # Compute the parameters of q(z | x) with the encoder latent_distribution_parameters = encoder(batch_x) # Apply reparameterization (sample from q(z | x)) z = reparameterize(latent_distribution_parameters, batch_noise) # The decoder outputs the parameters of the output distribution output_distribution_parameters = decoder(z) # Compute the ELBO and update the model parameters L = -ELBO( latent_distribution_parameters, output_distribution_parameters, batch_x ) L.backward() It is worth emphasizing that the decoder outputs precisely the parameters of the output distribution, not a particular sample from it. For example, if you model output images with a normal distribution $\\mathcal{N}(\\mu(z), \\sigma^2(z))$, the decoder will predict some $\\hat{\\mu}(z)$ and $\\hat{\\sigma}(z)$, which, together with the parameters of the latent distribution (the encoder\u0026rsquo;s output), are fed into the ELBO.\nTo generate an actual picture at inference time, you either honestly sample from $\\mathcal{N}(\\hat{\\mu}(z), \\hat{\\sigma}^2(z))$, or, as is often done, simply take the mean $\\hat{\\mu}(z)$ as the output image. In general, the exact way inference is carried out depends on the type of the output distribution used.\nChoosing the distributions It is time to give examples of concrete $p_\\theta(x \\mid z)$, $q_\\phi(z \\mid x)$ and $p_\\theta(z)$ with which a VAE can be built. To begin with, assume that $p_\\theta(z)$ can be set to the standard normal distribution:\n$$p_\\theta(z) = \\mathcal{N}(0, I) $$Note that in this case the prior distribution of $z$ has no dependence on the parameters $\\theta$.\nThe distribution $p_\\theta(x \\mid z)$ depends on the distribution your data comes from. If your data has a continuous distribution, then $p_\\theta(x \\mid z)$ can be set, for example, to a Gaussian:\n$$p_\\theta(x \\mid z) = \\mathcal{N}(f_\\theta(z), \\sigma^2) $$The mean vector in this example is given by a function $f$ of the variables $\\theta$ and $z$, and the covariance matrix is a constant diagonal matrix. The function $f$ can be defined by a neural network with parameters $\\theta$. If desired, the covariance matrix can also be given by some function rather than restricted to constant matrices. If your data is discrete, a categorical distribution may be suitable:\n$$p_\\theta(x \\mid z) = \\operatorname{Categorical}(f_\\theta(z)), $$in which the probability vector $f_\\theta(z) = (p_1, \\dots, p_n)$ is the network output after applying $\\text{softmax}$. If you have binary data, you can use a Bernoulli distribution:\n$$p_\\theta(x \\mid z) = \\operatorname{Bernoulli}(f_\\theta(z)), $$where $f_\\theta(z) = p$ is the neural network output after applying a sigmoid.\nThe distribution $q_\\phi(z \\mid x)$ can in principle be anything, but in the simplest case it is a Gaussian with a diagonal covariance matrix:\n$$q_\\phi(z \\mid x) = \\mathcal{N}(\\mu_\\phi(x), \\sigma_\\phi^2(x)) $$Such a distribution, in particular, admits the reparameterization trick discussed above. If we choose $z$ to be two-dimensional, the distributions defined by $q$ can be nicely visualized:\nImage source\nNow recall how the ELBO is defined:\n$$\\mathcal{L}_{\\theta,\\phi}(x) = \\mathbb{E}_{q_\\phi(z \\mid x)} [\\log p_\\theta(x \\mid z)] - D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z)) $$Let us compute it for the distributions given above.\nWe start with $D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z))$. The $KL$ divergence between the distributions $\\mathcal{N}(\\mu, \\Sigma)$ and $\\mathcal{N}(0, I)$ equals:\n$$D_{KL}(\\mathcal{N}(\\mu, \\Sigma) \\| \\mathcal{N}(0, I)) = \\frac{1}{2} (\\mu^T \\mu + \\operatorname{tr} \\Sigma - M - \\log(\\det \\Sigma)), $$where $M$ is the dimension of these distributions. A derivation of this relation can be found here. In our case $\\mu_\\phi(x) = (\\mu_1, \\dots, \\mu_M)$, $\\sigma_\\phi^2(x) = \\operatorname{diag}(\\sigma_1^2, \\dots, \\sigma_M^2)$ and\n$$D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z)) = D_{KL}(\\mathcal{N}(\\mu_\\phi(x), \\sigma_\\phi^2(x)) \\| \\mathcal{N}(0, I)) = $$$$= \\frac{1}{2} \\sum_{j=1}^M (\\sigma_j^2 + \\mu_j^2 - 1 - \\ln \\sigma_j^2) $$Then the ELBO is computed as:\n$$\\mathcal{L}_{\\theta,\\phi}(x) = \\mathbb{E}_{q_\\phi(z \\mid x)} [\\log p_\\theta(x \\mid z)] - D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z)) = $$$$= \\mathbb{E}_{\\mathcal{N}(\\mu_\\phi(x), \\sigma_\\phi^2(x))} [\\log p_\\theta(x \\mid z)] - \\frac{1}{2} \\sum_{j=1}^M (\\sigma_j^2 + \\mu_j^2 - 1 - \\ln \\sigma_j^2) \\approx $$$$\\approx \\frac{1}{K} \\sum_{k=1}^K \\log p_\\theta(x \\mid z_k) + \\frac{1}{2} \\sum_{j=1}^M (1 + \\ln \\sigma_j^2 - \\mu_j^2 - \\sigma_j^2), $$where $z_k \\sim \\mathcal{N}(\\mu_\\phi(x), \\sigma_\\phi^2(x))$. As mentioned in the paper by the authors of VAE, in section 2.3, the number of samples $K$ can be set to one, provided the batch size is large enough (e.g. 100).\nIf you choose a Bernoulli $p_\\theta(x \\mid z)$, then\n$$\\log p_\\theta(x \\mid z) = \\sum_{j=1}^D \\log p_\\theta(x_j \\mid z) = \\sum_{j=1}^D \\log \\operatorname{Bernoulli}(x_j, p_j) = $$$$= \\sum_{j=1}^D x_j \\log p_j + (1 - x_j) \\log (1 - p_j) $$If a Gaussian $\\mathcal{N}(f_\\theta(z), \\sigma^2)$, then\n$$\\begin{aligned} \\log p_\\theta(x \\mid z) \u0026= \\sum_{j=1}^D \\log p_\\theta(x_j \\mid z) \\\\ \u0026= \\sum_{j=1}^D \\log \\left( \\frac{1}{\\sqrt{2\\pi\\sigma^2}} \\exp \\left( -\\frac{(x_j - f_{\\theta,j}(z))^2}{2\\sigma^2} \\right) \\right) = \\end{aligned}$$$$= -\\frac{D}{2} \\log 2\\pi - D \\log \\sigma - \\frac{1}{2\\sigma^2} \\sum_{j=1}^D (x_j - f_{\\theta,j}(z))^2 $$An example implementation of training and using a VAE on the MNIST dataset can be found in Keras and in PyTorch.\nInference with a trained model Once we have trained a VAE, we can generate new samples simply by feeding $z \\sim \\mathcal{N}(0, I)$ into the decoder:\nImage source\nThe encoder is not needed for generating new samples. However, we may want to estimate $p(x) = \\int p(x \\mid z) p(z) dz$ for $x$ from a test set, to understand how likely the model is to generate $x$. To estimate the integral we need to sample some number of $z$, and if we take samples from $z \\sim \\mathcal{N}(0, I)$, the estimate may converge poorly. But we can again use the ELBO as a lower bound on $\\log p(x)$ and estimate it instead, sampling from the distribution $q_\\phi(z \\mid x)$. Such an estimate converges faster and gives a rough idea of how well the model handles a particular example $x$.\nIt is also interesting to look at how the codes of the training examples are distributed in the latent space. This is what the distribution of latent codes of MNIST digits may look like for a trained VAE with a two-dimensional latent space:\nImage source\nDifferent types of digits are shown in different colors (the correspondence between digits and colors is shown on the scale at the side). One can see that the model distinguishes zeros and ones best of all, and eights and threes worst of all. It is worth noting, of course, that the latent space was chosen to be two-dimensional for visualization purposes; with a higher dimension the model could learn to distinguish the digits better.\nFor a two-dimensional latent space there is another interesting way to visualize the structure of the manifold learned by the VAE. One can take a uniform grid on the unit square and map it into the latent space by applying the inverse CDF of the normal distribution.\nWhy this works The nodes of a uniform grid $u_{ij}$ can, to some approximation, be treated as samples from a uniform distribution: $u_{ij} \\sim \\text{Uniform}([0,1])$. Therefore the samples $\\Phi^{-1}(u_{ij})$ approximately follow the normal distribution:\n$$\\mathbb P(\\Phi^{-1}(u_{ij}) \\le t) = \\mathbb P(u_{ij} \\le \\Phi(t)) = \\Phi(t) $$ The resulting samples can be fed into the decoder to see which pictures correspond to the grid nodes:\nLeft: the learned Frey Face manifold; right: the learned MNIST manifold. Image source\nShown here are examples generated for the Frey Face and MNIST datasets (both available here). This visualization lets us see the smooth transition of latent codes of some objects into the codes of others, as well as the relative arrangement of the latent codes.\nFor MNIST we again see, in particular, that the model has placed the codes of zeros and ones far apart, while the codes of threes and eights are very close. It is also fun to observe the smooth transition from sixes to zeros and from sevens to ones. For Frey Face, the happy faces sit far from the sad ones, and along the main diagonal of the square one can trace a gradual transition from a serious face to a smiling one.\nIt is also interesting to look at how the quality of the generated digits changes depending on the dimension of the latent space (the pictures show random samples from the model):\nLatent space dimension, left to right: 2, 5, 10, 20. Image source\nA noticeable transition is visible between dimensions 2 and 5; further increase of the dimension has almost no significant effect.\nConditional VAE (CVAE) Sometimes we may want to generate not just an arbitrary object from the dataset, but one belonging to a particular group or class. Earlier we wrote out an equation for $\\log p_\\theta(x)$:\n$$\\begin{aligned} \\log p_\\theta(x) ={} \u0026 \\mathbb{E}_{q_\\phi(z \\mid x)} [\\log p_\\theta(x \\mid z)] \\\\ \u0026 - D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z)) + D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z \\mid x)) \\end{aligned}$$We can make all the distributions participating in this equation conditional on a variable $y$:\n$$\\begin{aligned} \\log p_\\theta(x \\mid y) ={} \u0026 \\mathbb{E}_{q_\\phi(z \\mid x, y)} [\\log p_\\theta(x \\mid z, y)] \\\\ \u0026 - D_{KL}(q_\\phi(z \\mid x, y) \\| p_\\theta(z \\mid y)) + D_{KL}(q_\\phi(z \\mid x, y) \\| p_\\theta(z \\mid x, y)) \\end{aligned}$$The variable $y$ can be the label of the object $x$, or an entirely arbitrary tensor characterizing $x$ in some way. Instead of a single $p_\\theta(z)$ shared by all $x$ in the training set, there is now a separate prior distribution $p_\\theta(z \\mid y)$ for each value of $y$.\nThe variable $y$ can take both discrete and continuous values. It can even be, for example, half of an image that the model is asked to complete. Just in case, let us stress that training a CVAE is not the same as training several independent VAEs, since the CVAE weights are shared across all classes.\nAt the implementation level this is quite simple: you just concatenate the inputs of the encoder and decoder with the tensor corresponding to $y$. If $y$ takes categorical values, it is often useful to first encode them as one-hot vectors. The algorithm looks roughly like this:\ndataset, labels = np.array(...), np.array(...) epsilon = RandomDistribution(...) # Encoder q_phi(z|x) — a neural network with parameters phi encoder = Encoder() # Decoder p_theta(x|z) — a neural network with parameters theta decoder = Decoder() for step in range(max_steps): # Sample a batch of input data, labels and random noise batch_x = sample_batch(dataset) batch_y = sample_batch(labels) batch_noise = sample_batch(epsilon) # Feed the encoder the concatenation of inputs and labels encoder_input = concatenate([batch_x, batch_y]) # Compute the parameters of the distribution of z with the encoder latent_distribution_parameters = encoder(encoder_input) # Apply reparameterization z = reparameterize(latent_distribution_parameters, batch_noise) # Concatenate the resulting random vector with the labels decoder_input = concatenate([z, batch_y]) # The decoder outputs the output image output_distribution_parameters = decoder(decoder_input) # Compute the ELBO and update the parameters L = -ELBO( latent_distribution_parameters, output_distribution_parameters, batch_x ) L.backward() A CVAE implementation in PyTorch and TensorFlow can be found, for example, here.\nIf we visualize the distribution of latent codes for MNIST digits obtained after conditioning the model on the digit class, we see something like this:\nImage source\nWe see an unintelligible mixture of points instead of the distinct clusters that the plain VAE produced. The point is that instead of trying to place all digits in a single space $p(z) \\sim \\mathcal{N}(0, I)$, the model uses a separate latent space $p(z \\mid y) \\sim \\mathcal{N}(0, I)$ for each digit:\nImage source\nImage source\nIn each picture, the right part shows the prior distributions for the digits 6 and 7, and the left part visualizes the structure of the learned manifolds for these digits, built the same way as the analogous visualization for the VAE. The quality of the images of each individual digit improves noticeably:\nImage source\nThe variability of the generated digits has also grown noticeably, and the model can imitate digits written in different handwriting styles.\nKey papers Beyond the standard description of how VAE works, we will discuss several works building on the VAE idea. Although these papers came out in 2017–2021, we cover them not as \u0026ldquo;recent results\u0026rdquo; but as the foundation on which modern generative models are built: the idea of discrete latent codes from VQ-VAE became the basis for tokenizing non-textual data — images, audio and video, — and DALL-E was the first demonstration that a large language model can be trained on top of such tokens.\nVQ-VAE and VQ-VAE-2 The models VQ-VAE and VQ-VAE-2 are interesting in that they employ discrete distributions as priors. In which situations can discrete distributions be more suitable than continuous ones? For example, when we deal with tokens in NLP tasks or phonemes in speech processing. Images, too, could be encoded by a set of integers: one number could encode the object type, another its color, a third the background color, and so on:\nImage source\nMoreover, there exist quite powerful algorithms (for example, the Transformer) designed to work with discrete data. Learning good discrete representations makes it possible to use such algorithms effectively for, say, image generation.\nLooking ahead: this very idea — compressing a continuous signal into a sequence of discrete codes — turned out to be one of the most fruitful in modern deep learning. The neural audio codecs SoundStream and EnCodec are VQ-VAEs with multi-level residual quantization; the voice cloning model VALL-E operates on top of EnCodec tokens. Visual tokenizers (VQGAN, MAGVIT) encode images and videos as discrete tokens for autoregressive generative models, and multimodal LLMs that generate images (for example, Meta\u0026rsquo;s Chameleon) treat pictures as sequences of the same kind of tokens as text. The idea of structuring latent representations is developing in other directions as well — for example, the nested, \u0026ldquo;matryoshka-like\u0026rdquo; variable-size embeddings of Matryoshka Representation Learning.\nVQ-VAE The authors of VQ-VAE introduce a discrete latent space in the form of $K$ real-valued vectors $e_1, \\dots, e_K$ of dimension $D$. The vectors of this space are called code vectors, or codes. The figure below shows a rough scheme of training the proposed model.\nImage source\nThe encoder takes an image $x$ as input and outputs a tensor $z_e(x)$. In the figure this tensor has shape $M \\times M \\times D$: the last dimension coincides with the length of the code vectors, and $M \\times M$ is the spatial dimension of the CNN output (for simplicity we do not write out the batch dimension explicitly).\nEach of the $M \\times M$ vectors of $z_e(x)$ is mapped to the code vector nearest to it in $L_2$ distance. After this procedure the tensor $z_e(x)$ turns into a tensor $z_q(x)$ consisting of $M \\times M$ code vectors. The decoder receives the tensor $z_q(x)$ as input and maps it back to the original image. For speech and text the authors used a two-dimensional tensor $z_e(x)$ instead of a three-dimensional one.\nThe output distribution of the encoder $q(z \\mid x)$ is defined here as follows:\n$$q(z = k \\mid x) = \\begin{cases} 1, \u0026 k = \\arg\\min_j \\| z_e(x) - e_j \\|_2, \\\\ 0, \u0026 \\text{otherwise} \\end{cases} $$During training, a uniform distribution $p(z)=\\frac 1K$ is used as the prior over the latent space, so the term $D_{KL}(q(z \\mid x) \\| p(z))$ turns out to be constant and equal to $\\log K$:\n$$\\begin{aligned} D_{KL}(q(z \\mid x) \\| p(z)) \u0026= -\\sum_{k=1}^K q(z = k \\mid x) \\log \\left( \\frac{p(z)}{q(z = k \\mid x)} \\right) \\\\ \u0026= \\log K \\end{aligned}$$At the points where $q(z = k \\mid x) = 0$, the next-to-last expression is extended by zero by continuity. Thus, the ELBO for such distributions takes the form\n$$ELBO(x) = \\mathbb{E}_{q(z \\mid x)} [\\log p_\\theta(x \\mid z_e(x))] - D_{KL}(q(z \\mid x) \\| p(z)) = \\log p_\\theta(x \\mid z_q(x)) - \\log K, $$where $\\theta$ are the decoder parameters. During optimization the $\\log K$ term can be ignored. The mapping of the encoder output to code vectors is not differentiable, so the following trick is used during training: on the backward pass, the gradient is copied directly from the decoder to the encoder, skipping the layer that maps encoder outputs to code vectors.\nThis trick is very close to the technique known as the straight-through estimator: on the backward pass the gradient flows through a non-differentiable operation as if it were not there. This technique was first proposed in this paper (and a simple explanation of it can be found here). Using the straight-through estimator, however, does not allow training the code vectors themselves, since no gradients are computed for them. Therefore the loss function for training the model consists of three components:\n$$\\mathcal L = \\log p(x \\mid z_q(x)) + \\| \\operatorname{sg}[z_e(x)] - z_q(x) \\|_2^2 + \\beta \\| z_e(x) - \\operatorname{sg}[z_q(x)] \\|_2^2 $$Here $\\operatorname{sg}[\\cdot]$ denotes the stop-gradient operator: no gradients flow through its argument.\nIn the paper the loss is written somewhat differently:\n$$\\mathcal L = \\log p(x \\mid z_q(x)) + \\| \\operatorname{sg}[z_e(x)] - e \\|_2^2 + \\beta \\| z_e(x) - \\operatorname{sg}[e] \\|_2^2 $$This notation seems somewhat confusing, for two reasons:\nThe letter $e$ in the subscript of $z_e(x)$ is only meant to indicate that this is the encoder output, not the existence of a connection between the code vectors $e$ and the encoder parameters. But the latter is quite easy to assume by mistake.\nSubtracting $e$ means subtracting not all elements of the codebook from the corresponding position of the tensor $z_e(x)$, but only the nearest neighbor of the element of $z_e(x)$ at that position. That is, in effect, subtracting $e$ in this notation is equivalent to subtracting $z_q(x)$. This is not clarified in the paper, but can be seen in the official implementation.\nThe first term is the ELBO up to a constant. The second term is responsible for moving the code vectors toward the encoder outputs. To prevent a situation where the encoder outputs keep dragging the code vectors around via the second loss component while themselves producing, at every iteration, vectors far from the current code vectors, a third term is added. It makes the encoder strive to output vectors close to the code vectors, and its importance is regulated by the coefficient $\\beta$.\nHowever, during training we lost the regularization term $D_{KL}(q(z \\mid x) \\| p(z))$, because of which the encoder distribution was under no obligation to approximate the prior and remained a narrow subset of it. As a result, when sampling from the uniform categorical distribution, we will most likely get plain noise instead of nice pictures:\nLeft to right: test data, their reconstructions, samples from the uniform prior. Image source\nIn a bit more detail When training a plain VAE, we minimize the distance between the prior distribution and the distribution produced by the encoder via the regularization term $D_{KL}(q(z \\mid x) \\| p(z))$.\nThanks to it, for example, the two-dimensional latent codes of MNIST digits approximately arrange themselves into a ball — the normal prior. And if each digit is given its own latent space (by conditioning on the digit class), the conditional prior for each digit is very close to normal.\nIn the case of VQ-VAE we cannot force the distribution predicted by the encoder to be the uniform categorical one; we simply get some categorical distribution with unknown parameterization. This is reminiscent of the situation with a plain autoencoder: it also maps input images into a latent space, but we cannot sample from that space.\nTo fix this problem, the authors propose to learn, with an additional model, the prior distribution $p(z)$ of those latent variables that the model learned to generate during training. Since any code representation can be flattened into a sequence, and the number of codes is finite and fixed in advance, this task is close to training a language model.\nIndeed, there we must predict the next word from the available vocabulary given the sequence of preceding words of a sentence, and in our case — predict the next latent code given the input sequence of discrete latent codes.\nFor images, the authors proposed to model the prior distribution of latent codes with PixelCNN. The details of this model\u0026rsquo;s architecture and training can be found in the original paper; here we describe only the general idea.\nPixelCNN generates the pixels of an image sequentially, moving from the top-left corner to the bottom-right. It traverses all rows one by one from top to bottom, and within each row moves left to right:\nImage source\nFor color images, the channels (R, G, B) are also modeled sequentially: when generating, channel B depends on R and G, and G only on R. When predicting the value of each next pixel, the model uses the values of already-generated neighbors from some surrounding square. To prevent the model from reading pixels that come after the currently predicted one, a special mask is used, an example of which is shown in the right part of the figure.\nIn the case of VQ-VAE, PixelCNN is trained not on pixels but on latent codes. Sampling from the learned prior looks much better than attempts to sample from the uniform one:\nLeft: samples from the uniform prior; right: samples from the prior learned by PixelCNN. Image source\nFor audio, the authors use WaveNet instead of PixelCNN. When training the prior models, one can also feed in class labels, so that later one can sample from those classes (the same principle as for CVAE).\nThe results of reconstructing ImageNet images with VQ-VAE look quite good (by reconstruction we mean the output of the full model consisting of the encoder and decoder):\nLeft: original ImageNet images (128×128); right: their VQ-VAE reconstructions with a 32×32 latent space and a codebook of 512 code vectors. Image source\nAnd this is what sampling from a VQ-VAE with a PixelCNN-learned prior looks like:\n128×128 samples from a VQ-VAE with a PixelCNN prior trained on ImageNet. Left to right: kit fox, gray whale, brown bear, admiral butterfly, coral reef, alp, microwave, pickup. Image source\nVQ-VAE-2 The VQ-VAE-2 model is not used today by itself, but it is interesting as an example of how the VQ-VAE idea evolved — above all, through a hierarchy of discrete latent spaces.\nMore about VQ-VAE-2 The VQ-VAE-2 model is an extension of VQ-VAE. It shows a significant leap in the quality of generated images:\nClass-conditional 256×256 samples from a two-level VQ-VAE-2 trained on ImageNet. Image source\nWhat is impressive is that the picture shows the result of sampling from the distribution learned by the model, not the result of reconstruction. The first key difference between VQ-VAE and VQ-VAE-2 is the use of hierarchical latent variables:\nImage source\nBefore moving on to the architecture description, a small disclaimer: whenever the text below says \u0026ldquo;a tensor of size $M \\times M$\u0026rdquo;, it means the tensor has shape $(B,M,M,C)$, where the first component corresponds to batches and the last one to channels.\nThe picture shows an example of a two-level architecture (although there may be more levels). Each level has its own encoder, decoder and set of code vectors (of a common dimension $D$ for all levels). Denote the bottom and top encoders by $Enc_\\text{bottom}$ and $Enc_\\text{top}$, and the decoders by $Dec_\\text{bottom}$ and $Dec_\\text{top}$.\n$Enc_\\text{bottom}$ takes a three-channel image of size $256 \\times 256$ pixels, maps it to a tensor of size $64 \\times 64$ and passes it to $Enc_\\text{top}$. $Enc_\\text{top}$ outputs a tensor of size $32 \\times 32$, which is then mapped to a tensor of code vectors $z_\\text{top}$ (quantized)\n$z_\\text{top}$ is fed into $Dec_\\text{top}$, then the outputs of $Enc_\\text{bottom}$ and $Dec_\\text{top}$ are concatenated and quantized into $z_\\text{bottom}$\n$z_\\text{top}$ and $z_\\text{bottom}$ are concatenated and fed into $Dec_\\text{bottom}$, which maps them to the original image\nThe model is trained with almost the same loss as VQ-VAE. For VQ-VAE it had the form:\n$$\\mathcal L = \\log p(x \\mid z_q(x)) + \\| \\operatorname{sg}[z_e(x)] - z_q(x) \\|_2^2 + \\beta \\| z_e(x) - \\operatorname{sg}[z_q(x)] \\|_2^2 $$For VQ-VAE-2 the first and third terms keep their form, while the second term is replaced by updating the code vectors $e_i$ with an exponential moving average. Let $E(x)^{(t)}$ be the encoder output at step $t$, flattened into a two-dimensional tensor whose last dimension equals the dimension $D$ of the code vectors.\nLet $\\{ E_{i,1}^{(t)}, \\dots, E_{i,n_i^{(t)}}^{(t)} \\}$ be the set of $n_i^{(t)}$ vectors for which, at step $t$, the nearest code vector was $e_i^{(t-1)}$. Then $e_i$ is updated at step $t$ by the following formulas:\n$$e_i^{(t)} = \\frac{m_i^{(t)}}{N_i^{(t)}} $$$$m_i^{(t)} = m_i^{(t-1)} \\cdot \\gamma + \\sum_j^{n_i^{(t)}} E(x)_{i,j}^{(t)} (1 - \\gamma) $$$$N_i^{(t)} = N_i^{(t-1)} \\cdot \\gamma + n_i^{(t)} (1 - \\gamma) $$Here $\\gamma$ is some real-valued parameter.\nJust as for VQ-VAE, the prior for VQ-VAE-2 is learned separately after the main model has been trained, but in the case of VQ-VAE-2 it has a hierarchical structure. The picture shows an example of such a distribution for a two-level architecture:\nImage source\nA separate PixelCNN model is trained for each level: one on the code vectors of the first level, the other on the code vectors of the first and second levels. Both models also take as input the label of the class from which an image should be sampled.\nSampling from the final model goes as follows:\nvectors $e_\\text{top}$ are sampled from the top distribution\nvectors $e_\\text{bottom}$ are sampled from the bottom distribution conditioned on the vectors $e_\\text{top}$\nthe decoder takes the vectors $e_\\text{top}$ and $e_\\text{bottom}$ as input and outputs the final picture\nSampling results from a two-level VQ-VAE-2 trained on ImageNet:\nClass-conditional samples from VQ-VAE-2. Classes by row, top to bottom: sea anemone, brain coral, slug, goldfinch, flamingo, redshank, Pekinese, papillon, drake, spotted salamander. Image source\nAnd here are the sampling results from a three-level VQ-VAE-2 trained on FFHQ:\nSamples from a three-level model trained on FFHQ 1024×1024: the model maintains global dependencies (eye color, facial symmetry) while also covering rare modes of the dataset — for example, green hair. Image source\nDALL-E Another work that largely shaped the development of generative models is DALL-E by OpenAI. They trained a model with 12 billion parameters that generates pictures from their text descriptions. For training, the authors collected a dataset of 250 million image–caption pairs. Here are some examples of this model at work:\nPrompt: \u0026ldquo;an armchair in the shape of an avocado\u0026rdquo;. Image source\nPrompt: \u0026ldquo;a store front that has the word \u0026lsquo;openai\u0026rsquo; written on it\u0026rdquo;. Image source\nMore examples of generations for various text descriptions can be found in OpenAI\u0026rsquo;s blog post about DALL-E.\nPrompt: \u0026ldquo;a painting of a capybara sitting in a field at sunrise\u0026rdquo; in different styles: painting, pop art, cubism, surrealism, Van Gogh, Monet, pencil drawing, charcoal, crayons and chalk. Image source\nConceptually, DALL-E builds on the results of VQ-VAE: first, code vectors for pictures are learned, and then a Transformer is trained to model the joint prior distribution of texts and code vectors. (I\u0026rsquo;m preparing a separate series of posts about transformers — stay tuned.)\nIn essence, DALL-E is the first demonstration that the \u0026ldquo;textual\u0026rdquo; pretraining recipe (autoregressive next-token prediction) also works for non-textual data, provided it is first discretized with a VQ-VAE-style encoder. This recipe quickly spread beyond image generation: for example, the voice cloning system Tortoise TTS is modeled after DALL-E — its author explicitly described the project as \u0026ldquo;DALL-E for speech synthesis\u0026rdquo;.\nDALL-E employs an architecture based on the decoder part of the original Transformer, so it is also worth reading up on the GPT-2 model, which works in a similar way.\nTraining proceeds in two stages:\nFirst, a discretized VAE (dVAE) is trained, with an encoder compressing $256 \\times 256$ RGB images into a tensor of $32 \\times 32 = 1024$ code vectors. This training stage strongly resembles VQ-VAE, but instead of adding extra loss terms for the code vectors, the DALL-E authors use the Gumbel relaxation — a trick that makes honest differentiation with respect to the encoder parameters possible. We will discuss dVAE training in more detail below.\nThen a Transformer is trained (more precisely, only the decoder part of the original Transformer architecture), whose task is to learn the joint distribution of pictures and their text descriptions. It takes as input the concatenation of the embeddings of the text tokens and the code vectors of the pictures, and learns to predict the continuation of each input sequence. Some details of the Transformer training will also be covered below.\nInference with the trained model goes like this: the embeddings of a picture\u0026rsquo;s text description are fed into the Transformer, which autoregressively predicts the code vectors of a picture matching that description; the resulting code vectors are then passed through the dVAE decoder.\ndVAE The dVAE is trained by maximizing the ELBO for pictures $x$ and their discrete latent representations $z$:\n$$\\ln p_\\theta(x) \\ge \\mathbb{E}_{q_\\phi(z \\mid x)} [\\log p_\\theta(x \\mid z)] - \\beta \\, D_{KL}(q_\\phi(z \\mid x) \\| p(z)), $$where $\\phi$ and $\\theta$ are the parameters of the encoder and decoder of the discretized VAE, and $p(z)$ is the uniform categorical distribution over the code vectors. Notice the extra coefficient $\\beta$, which in a standard VAE always equals 1. The DALL-E authors introduced the additional parameter $\\beta$ following the results of the paper on $\\beta$-VAE. Unlike the original paper, though, in their experiments the value of $\\beta$ is gradually decreased during training.\nThe dVAE encoder maps $256 \\times 256$ images into a tensor $z_e(x)$ of shape $32 \\times 32 \\times 8192$, where $8192$ is the number of code vectors. That is, to each of the $32 \\times 32$ positions the encoder assigns a categorical distribution over the $8192$ code vectors, parameterized by the output logits.\nTo obtain the tensor $z_q(x)$ of code vectors, one could first apply $\\text{softmax}$ to the distributions at each of the $32 \\times 32$ positions, and then assign to each position the code vector whose index has the maximum probability (take the $\\arg\\max$ for that position).\nHowever, the $\\arg\\max$ operation is not differentiable — and besides, in the VAE framework the decoder should receive a sample from the distribution predicted by the encoder, and taking the $\\arg\\max$ at each position is not sampling from the predicted distribution.\nWe will therefore need a couple of tricks that will let us simultaneously:\napproximate sampling from the $\\text{softmax}$\nmake the sampling differentiable\nGumbel-Max Trick and Gumbel-Softmax The first trick is known as the Gumbel-Max Trick. Imagine we have network output logits $x_1, \\dots, x_k$, and we want to use them to obtain a sample from the categorical distribution — that is, to predict a class stochastically. For this we usually apply $\\text{softmax}$ to the logits to obtain probabilities $\\pi_i$:\n$$\\pi_i = \\frac{\\exp x_i}{\\sum_j \\exp x_j}, $$and then sample a class from the resulting categorical distribution $\\{\\pi_1, \\dots, \\pi_k\\}$. It turns out that these two steps are equivalent to the following procedure:\nsample numbers $g_1, \\dots, g_k$ from the standard Gumbel distribution,\nadd the sample $g_i$ to each logit $x_i$,\npick the class $j$ such that $j = \\arg\\max_i (x_i + g_i)$.\nWhy this is indeed the case can be read here. But the Gumbel-Max Trick alone will not help us — the operation still is not differentiable. So we need one more trick, proposed almost simultaneously in two papers (first and second) and named Gumbel-Softmax in one of them.\nTo describe this trick, note that the result of the $\\arg\\max$ operation is the index of some class $j$. Such an index can be described by one-hot encoding — a vector of length $k$ in which all elements are zero except the $j$-th, which equals one.\nGumbel-Softmax consists of doing the following instead of taking the $\\arg\\max$ at the last step of the Gumbel-Max Trick:\ncompute $y_i = \\frac{\\exp((x_i + g_i)/\\tau)}{\\sum_{j=1}^k \\exp((x_j + g_j)/\\tau)}$, $i = 1, \\dots, k$, — an approximation of the one-hot vector via the $\\text{softmax}$ activation with temperature;\nsum the code vectors $e_i$ with weights $y_i$: $z = \\sum_i y_i e_i$;\noutput the vector $z$ as the latent vector for the given position\nStrictly speaking, the DALL-E authors did not specify how the output vector $z$ is aggregated from the code vectors and the $y_i$, but this is the approach taken in the PyTorch implementation of DALL-E.\nAs $\\tau \\to 0$, sampling from the distribution $\\frac{\\exp((x_i + g_i)/\\tau)}{\\sum_{j=1}^k \\exp((x_j + g_j)/\\tau)}$ tends to the $\\arg\\max$, and during dVAE training the authors gradually decreased the value of $\\tau$. In the following picture, the left part shows the plain Gumbel-Max Trick, and the right part its differentiable variant:\nImage source\nThus, training the code vectors of the dVAE requires no extra loss terms on top of the ELBO, nor copying gradients from the decoder to the encoder (as was the case in VQ-VAE).\nMoreover, it is worth noting that $D_{KL}(q_\\phi(z \\mid x) \\| p(z))$ in this case does not degenerate into a constant, but genuinely acts as a regularizer. Let us spell out how it works: the encoder outputs a $32 \\times 32$ grid of positions and at each of them predicts its own categorical distribution over the $8192$ code vectors, while the prior $p(z)$ takes the code at each position to be equally probable among all options and independent of the other positions. The regularizer pulls the encoder\u0026rsquo;s prediction at each position separately toward this uniform distribution.\nImportantly, such regularization says nothing about the joint structure of all $1024$ codes of a picture: a real image corresponds to a strongly correlated set of codes, not at all to an independent set of equally probable ones. So one cannot sample pictures directly from the dVAE (drawing the code at each position independently) — the joint distribution of the codes has to be modeled separately, and that is exactly what the Transformer does in the second stage. Essentially, this is the same role that PixelCNN played in modeling the prior over codes for VQ-VAE.\nFinally, the non-constant $KL$ has one more practical advantage — it encourages high codebook utilization. By pulling the distribution at each position toward the uniform one, the regularizer prevents the model from collapsing to the use of just a handful of codes (a problem known as codebook collapse). In hard VQ-VAE there is no such regularizer, so codebook collapse is fought there with separate techniques — for example, updating the code vectors with an exponential moving average (EMA; this technique is covered above in the section on VQ-VAE-2) or re-initializing \u0026ldquo;dead\u0026rdquo; codes that have stopped being selected.\nWhy a vanilla VAE blurs pictures and VQ-VAE does not In the models we are discussing there are two completely different independence assumptions, and they are easy to confuse.\nThe first is the conditional independence of the object\u0026rsquo;s components given the latent representation: $p_\\theta(x \\mid z) = \\prod_j p_\\theta(x_j \\mid z)$. It is precisely what allowed us to write the reconstruction loss as a sum over pixels (see the section \u0026ldquo;Choosing the distributions\u0026rdquo;; for DALL-E — the Logit-Laplace distribution below).\nThe second is the independence of the latent variables themselves across positions, i.e. the factorized form of the prior $p(z)$ discussed above.\nIn a vanilla VAE the latent vector is single and small, so it is unable to explain all the correlations between pixels. And the decoder, by construction, cannot fill them in: it generates pixels independently. As a result, the model is forced to average over all admissible variants of the fine details — hence the characteristic blurriness of VAE samples.\nWhen we move to a grid of $32 \\times 32$ latents, the situation flips. The latent representation is now rich enough to explain almost the entire structure of the image, so the conditional independence of pixels given $z$ becomes a good approximation — and the reconstructions come out sharp (which we saw in the VQ-VAE pictures). But the correlations have not disappeared: they have moved into the distribution of the codes themselves, which is no longer factorized.\nIn other words, the modeling burden shifts from the decoder to the prior. Before, the dependencies between pixels had to be explained by the decoder — and it could not; now the dependencies between codes have to be explained by $p(z)$ — and a factorized $p(z)$ cannot handle that either. That is why it is replaced by an autoregressive model: PixelCNN in VQ-VAE, the Transformer in DALL-E.\nThe Logit-Laplace distribution One more trick in dVAE training concerns the output distribution $p_\\theta(x \\mid z)$. The DALL-E authors noticed a problem arising with the commonly chosen Laplace and Gaussian distributions for $p_\\theta(x \\mid z)$: both are defined on the entire real line, whereas pixels take values from a bounded interval. Thus, part of the density is \u0026ldquo;lost\u0026rdquo; during modeling, ending up outside the feasible range of pixel values.\nTo fix this problem, the authors propose to use a distribution they called \u0026ldquo;Logit-Laplace\u0026rdquo;. Its density is defined on the interval $(0,1)$ and is expressed by the following formula:\n$$f(x \\mid \\mu, b) = \\frac{1}{2b x(1-x)} \\exp\\left( -\\frac{|\\operatorname{logit}(x) - \\mu|}{b} \\right), $$$$\\operatorname{logit}(x) = \\log \\frac{x}{1-x} $$This density corresponds to a random variable obtained by applying the sigmoid to a Laplace-distributed random variable. The expression for the Logit-Laplace distribution can be derived from the standard formula for the density of a random variable obtained by applying a monotone differentiable function to another random variable (see the formula, for example, here). The logarithm of this density is substituted into the ELBO in place of $\\ln p_\\theta(x \\mid z)$.\nThe decoder outputs 6 tensors: the first three correspond to $\\mu$ for the RGB channels, the remaining three correspond to $\\ln b$, and these 6 tensors are used to compute the loss. Before being fed to the encoder, the image values are normalized by the function $\\phi: [0,255] \\to (\\varepsilon, 1-\\varepsilon)$:\n$$\\phi: x \\mapsto \\frac{1-2\\varepsilon}{255} x + \\varepsilon $$This way the authors ensure that the decoder models values from $(\\varepsilon, 1-\\varepsilon)$, which mitigates the computational problems associated with dividing by $x(1-x)$ in the density formula. At inference time, the reconstruction $\\hat x$ of a picture $x$ is computed by the formula:\n$$\\hat{x} = \\phi^{-1}(\\operatorname{sigmoid}(\\mu)), $$where $\\mu$ is the first three tensors of the decoder output. The outputs corresponding to $\\ln b$ are not used here.\nThe prior over texts and images In the second stage, the authors freeze the parameters $\\phi$ and $\\theta$ and model the joint distribution of pictures and their text descriptions with a Sparse Transformer with 12 billion parameters. As input it receives the concatenation of a picture\u0026rsquo;s text description and its code vectors. A picture is represented by 1024 code vectors obtained from the encoder $q_\\phi$, and when sampling code sequences the plain $\\arg\\max$ is used, without adding noise from the Gumbel distribution.\nThe text description is tokenized with the BPE procedure (see the section on BPE here), and each token is assigned a vector of real numbers representing it (an embedding). At most 256 tokens are used to represent the text, and the vocabulary size is 16,384 tokens.\nThe Transformer\u0026rsquo;s task during training is to predict, for each initial segment of the input sequence, the token that follows it. This can be either a text token or an image code vector. Since the code vectors of a picture always come after the text tokens, when generating code vectors the attention mechanism also attends to all the preceding text tokens.\nFurthermore, the attention mask for the code vectors takes into account that they are originally arranged not linearly one after another, but on a rectangular grid. The paper presents several variants of geometric patterns used for the attention mask over code vectors.\nThe loss is a weighted sum of the cross-entropy for the text tokens and the cross-entropy for the picture code vectors, with weights $\\frac 18$ and $\\frac 78$ respectively (image generation is given higher priority, hence the larger weight for its loss).\nOf course, training a huge Transformer is anything but easy, and a substantial part of the paper is devoted to the tricks the authors applied to train such a large model.\nInference At inference time, the tokens of a picture\u0026rsquo;s text description are fed into the model, and based on them the model autoregressively predicts the code vectors:\nImage source\nThe picture\u0026rsquo;s code vectors are fed into the dVAE decoder, which maps them into the final picture:\nImage source\nTo improve prediction quality, the authors first generate 512 pictures for each text description, and then pick the best picture among the predictions. Different sets of code vectors for the same text can be obtained, for example, by randomly picking a code vector at each generation step according to the distribution predicted by the Transformer. The ranking of the resulting 512 pictures is done with CLIP — a large neural network trained without supervision on a large amount of data to model the joint distribution of pictures and texts.\nConclusion So, in this post we discussed how the VAE works in its classical form — with a continuous distribution of latent variables — and covered the works based on the idea of using discrete distributions in VAEs.\nOf course, the various modifications of VAE are not limited to swapping continuous latent variables for discrete ones. There are many other possible directions for improving the model: hierarchical latent distributions (which we saw, by the way, in the context of VQ-VAE-2), loss functions other than the ELBO, various shapes of latent spaces, adversarial training, and much more.\nA good list of papers on VAE modifications can be found here. Among the works developing hierarchical distributions, NVAE is worth noting — there is a good video review of it by Yannic Kilcher. It deserves a separate mention that the ideas of VAE underlie latent diffusion models (such as Stable Diffusion): in them, diffusion happens not in pixel space but in the latent space of a trained autoencoder.\nThis concludes our story about VAE. Hopefully, it gave you a general picture both of the original ideas from which the VAE model grew and of the most interesting results connected with it.\nSelf-check questions 1. Why introduce hidden (latent) variables when building a generative model? Why not just estimate $p(x)$ from the data directly?\nAnswer First, even very simple $p(z)$ and $p(x \\mid z)$ can, after marginalization over $z$, yield a very complex, multimodal $p(x)$ — for example, a discrete $p(z)$ with Gaussian $p(x \\mid z)$ gives a mixture of Gaussians, and a continuous $p(z)$ an \u0026ldquo;infinite\u0026rdquo; mixture. In other words, latent variables are a way to describe a complex distribution through simple building blocks.\nSecond, direct statistical density estimation $\\hat{p}(x)$ runs into the curse of dimensionality: the higher the dimension of the data, the exponentially more examples are needed for an adequate estimate.\n2. Why can\u0026rsquo;t the likelihood $p_\\theta(x) = \\mathbb{E}_{z \\sim p_\\theta(z)}[p_\\theta(x \\mid z)]$ simply be estimated by Monte Carlo, sampling $z$ from the prior?\nAnswer The number of samples needed to cover the latent space $Z^M$ well grows exponentially with the dimension $M$ (the curse of dimensionality). Meanwhile, the contribution of the overwhelming majority of $z$ to the estimate is practically zero: only a small part of the latent space maps to objects resembling the elements of the dataset. Hence the idea (akin to importance sampling): introduce a distribution $q_\\phi(z \\mid x)$ over the \u0026ldquo;preimages\u0026rdquo; of the object $x$ and sample only from it.\n3. What is the ELBO, and why is maximizing it a reasonable substitute for maximizing $\\log p_\\theta(x)$?\nAnswer The ELBO (evidence lower bound) is the functional\n$$\\mathcal{L}_{\\theta,\\phi}(x) = \\mathbb{E}_{q_\\phi(z \\mid x)} [\\log p_\\theta(x,z) - \\log q_\\phi(z \\mid x)]$$The log-likelihood decomposes into the sum\n$$\\log p_\\theta(x) = \\mathcal{L}_{\\theta,\\phi}(x) + D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z \\mid x)),$$and since the $KL$ divergence is nonnegative, the ELBO is a lower bound on $\\log p_\\theta(x)$, with the gap between them equal to $D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z \\mid x))$ — the distance between the encoder\u0026rsquo;s distribution and the true (intractable) posterior. Therefore, by maximizing the ELBO we simultaneously increase $\\log p_\\theta(x)$ and bring $q_\\phi(z \\mid x)$ closer to $p_\\theta(z \\mid x)$.\n4. Which two terms does the ELBO consist of, and what is each of them responsible for?\nAnswer $$\\mathcal{L}_{\\theta,\\phi}(x) = \\underbrace{\\mathbb{E}_{q_\\phi(z \\mid x)} [\\log p_\\theta(x \\mid z)]}_{\\text{reconstruction loss}} - \\underbrace{D_{KL}(q_\\phi(z \\mid x) \\| p_\\theta(z))}_{\\text{regularization term}}$$The first term is the reconstruction loss: it measures how well the decoder reconstructs the object $x$ from its latent representation. The second is the regularization term: it pushes the distribution produced by the encoder toward the prior $p_\\theta(z)$. It is thanks to this term that one can sample from a trained VAE by feeding $z \\sim p(z)$ into the decoder.\n5. Why is the reparameterization trick needed when training a VAE? Why can an unbiased estimate of the ELBO gradient with respect to $\\theta$ be obtained without it, but not with respect to $\\phi$?\nAnswer The ELBO is an expectation over $z \\sim q_\\phi(z \\mid x)$. The parameters $\\theta$ do not participate in this distribution, so $\\nabla_\\theta$ can simply be moved under the expectation sign and estimated from samples. But $\\phi$ parameterizes the very distribution over which the expectation is taken, so $\\nabla_\\phi$ cannot be moved under the expectation.\nThe reparameterization trick represents $z$ as a deterministic differentiable function of the parameters and independent noise: for example, $z = \\mu_\\phi(x) + \\varepsilon \\cdot \\sigma_\\phi(x)$, where $\\varepsilon \\sim \\mathcal{N}(0, I)$. The source of randomness moves into the \u0026ldquo;input data\u0026rdquo;, the expectation is now taken over $p_\\varepsilon$, which does not depend on $\\phi$, and gradients flow through deterministic variables — resulting in a setup typical for SGD.\n6. Which distributions are usually chosen for $p_\\theta(z)$, $q_\\phi(z \\mid x)$ and $p_\\theta(x \\mid z)$ in a classical VAE?\nAnswer Prior: $p_\\theta(z) = \\mathcal{N}(0, I)$. Encoder: $q_\\phi(z \\mid x) = \\mathcal{N}(\\mu_\\phi(x), \\sigma_\\phi^2(x))$ with a diagonal covariance matrix — such a distribution admits reparameterization, and its $KL$ divergence with $\\mathcal{N}(0, I)$ has a closed form: $\\frac{1}{2} \\sum_j (\\sigma_j^2 + \\mu_j^2 - 1 - \\ln \\sigma_j^2)$. Decoder: depends on the nature of the data — Gaussian $\\mathcal{N}(f_\\theta(z), \\sigma^2)$ for continuous data (then the reconstruction loss turns into MSE up to constants), Bernoulli for binary, categorical for discrete. 7. How do you generate a new object with a trained VAE? Is the encoder needed for this?\nAnswer It suffices to sample $z \\sim \\mathcal{N}(0, I)$ and feed it into the decoder — the encoder is not needed for generation. It comes in handy, though, if we want to estimate $\\log p(x)$ for a particular object $x$ (say, a test one): the ELBO-based estimate with sampling from $q_\\phi(z \\mid x)$ converges much faster than direct Monte Carlo estimation with sampling from the prior.\n8. How does CVAE differ from a plain VAE? Why is it not the same as training a separate VAE per class?\nAnswer In a CVAE all distributions are conditioned on an additional variable $y$ (a label or an arbitrary tensor characterizing the object): $q_\\phi(z \\mid x, y)$, $p_\\theta(x \\mid z, y)$, $p_\\theta(z \\mid y)$. In practice this comes down to concatenating $y$ (e.g. as a one-hot vector) with the inputs of the encoder and decoder. For each value of $y$ the model effectively gets its own prior $p_\\theta(z \\mid y)$.\nAt the same time, the network weights are shared across all classes — unlike a collection of independent VAEs — so the model can reuse patterns learned across classes, and $y$ can also be continuous.\n9. How is the latent space of VQ-VAE organized? What happens to the $KL$ term of the ELBO during training, and how do gradients pass through the quantization operation?\nAnswer The latent space of VQ-VAE is a codebook of $K$ learnable code vectors $e_1, \\dots, e_K$. Each vector of the encoder output $z_e(x)$ is replaced by the nearest code vector in $L_2$ distance. The encoder distribution $q(z \\mid x)$ is then degenerate (one-hot), and the prior during training is uniform, so $D_{KL}(q(z \\mid x) \\| p(z)) = \\log K$ — a constant that can be ignored during optimization.\nQuantization is non-differentiable, so on the backward pass the gradient is copied from the decoder to the encoder \u0026ldquo;through\u0026rdquo; it — this is the straight-through estimator applied to the quantization operation. But this way the code vectors receive no gradients, so two more terms are added to the loss:\n$$\\mathcal L = \\log p(x \\mid z_q(x)) + \\| \\operatorname{sg}[z_e(x)] - z_q(x) \\|_2^2 + \\beta \\| z_e(x) - \\operatorname{sg}[z_q(x)] \\|_2^2$$The second term moves the code vectors toward the encoder outputs; the third makes the encoder produce vectors close to the code vectors.\n10. Why does sampling from the uniform prior of a trained VQ-VAE yield noise instead of meaningful pictures, and how do the authors solve this problem?\nAnswer The regularizing $KL$ term degenerated into a constant during VQ-VAE training, so nothing forced the distribution of latent codes to approach the uniform one — it remained an unknown \u0026ldquo;narrow subset\u0026rdquo; (a situation resembling a plain autoencoder, from whose latent space one cannot sample). A sample from the uniform distribution almost surely misses the region of actually used codes.\nThe solution is to learn the prior $p(z)$ with a separate autoregressive model over the discrete codes (which is close to training a language model): PixelCNN for images, WaveNet for audio. Samples from the learned prior look incomparably better.\n11. In a plain (continuous) VAE one can sample directly from the prior $\\mathcal N(0,I)$, while VQ-VAE required a separate prior model (PixelCNN). Why?\nAnswer In a continuous VAE the regularizing $KL$ term does not degenerate: it penalizes $q_\\phi(z\\mid x)$ for deviating from $\\mathcal N(0,I)$ for every $x$, so the aggregate distribution of the codes is trained to resemble the prior. Then sampling $z\\sim\\mathcal N(0,I)$ and passing it through the decoder yields meaningful objects. In VQ-VAE the $KL$ collapsed into a constant — the distribution of the codes is unregularized, arbitrary and unknown, so a sample from the uniform gives noise.\nBut there is a deeper reason, unrelated to the degeneration of the $KL$. A plain VAE typically encodes an object into a single latent vector, whose entire structure is \u0026ldquo;deciphered\u0026rdquo; by the decoder — and a simple $\\mathcal N(0,I)$ suffices for it. VQ-VAE, in order to preserve details, keeps a grid of many latents, and to reconstruct a coherent picture the codes in different cells must be correlated (neighboring cells describe neighboring parts of the image). A factorized prior — a product of independent distributions over positions, $p(z)=\\prod_i p(z_i)$, — cannot express such dependence in principle: it will only fit the marginals at each position, and an independent sample from it will produce an incoherent mosaic.\nThis can also be seen formally. If we fit a factorized $p(z)=\\prod_i p_i(z_i)$ by maximum likelihood to the true distribution of codes $q(z)$, the objective decouples into per-position cross-entropies, each depending only on its own marginal $q_i$; the optimum is reached at $p_i = q_i$, and the irreducible gap equals $\\sum_i H(q_i) - H(q)$ — the total correlation between positions. The objective simply has no term that would penalize incorrectly learned dependencies.\nWhat fixes this is precisely a non-factorized, autoregressive prior $p(z)=\\prod_i p(z_i \\mid z_{\\lt i})$ (PixelCNN, and in DALL-E — the Transformer), which does capture the dependencies between positions.\nIncidentally, it is useful to notice where the correlations come from — and first to spell out what exactly is contrasted with what. The whole grid of codes is a single latent variable $z = (z_1, \\ldots, z_n)$: one sample from a categorical distribution over the (huge) space of all code combinations, and the dependencies we are discussing are dependencies between the components of this one $z$. For a fixed $x$, the dVAE encoder predicts a factorized distribution:\n$$q_\\phi(z \\mid x) = \\prod_i q_\\phi(z_i \\mid x),$$but the dataset-aggregated distribution of codes $q(z) = \\mathbb{E}_x\\, q_\\phi(z \\mid x)$ is a mixture of such factorized distributions over all $x$, and a mixture of independent distributions is no longer independent:\n$$q(z) \\ne \\prod_i q(z_i)$$A toy example makes this immediate. Let there be two positions, two codes, and two kinds of images in equal proportion — cats and dogs. Suppose the trained encoder maps every cat to the pair $(0,0)$ and every dog to $(1,1)$ (for a fixed $x$ the distribution is degenerate, hence trivially independent). Now aggregate over the dataset: the marginals are $q(z_1{=}0) = q(z_2{=}0) = \\tfrac12$, so if the positions were independent, the combination $(0,1)$ would have probability $\\tfrac14$ — yet in fact $q(0,1) = 0$: \u0026ldquo;cat ears + dog tail\u0026rdquo; never occurs in the data. The correlation comes not from the encoder (which is independent given $x$), but from the mixing over the data: the code at position 1 reveals which image was encoded, and that, in turn, shifts the distribution of the code at position 2. A factorized prior would happily sample such chimeras — exactly the incoherent mosaic mentioned above; these dependencies are what the prior model has to capture.\nThis also shows that discreteness as such is not the point: a continuous VAE with a latent grid would face the same issue and would also require a separate prior model — in fact, latent diffusion models (Stable Diffusion) do exactly that: they learn a distribution over the continuous latent grid of an autoencoder.\n12. When deriving the ELBO we assumed that the components of an object are conditionally independent given the latent representation: $p_\\theta(x \\mid z) = \\prod_j p_\\theta(x_j \\mid z)$. But in VQ-VAE and dVAE the codes at different positions are, on the contrary, strongly correlated. Doesn\u0026rsquo;t this violate the original assumption and render the ELBO incorrect?\nAnswer It does not: two different independence assumptions are being conflated here.\nConditional independence of the object\u0026rsquo;s components given the latent, $p_\\theta(x \\mid z) = \\prod_j p_\\theta(x_j \\mid z)$ — this is what turns the reconstruction loss into a sum over pixels. In VQ-VAE and dVAE it still holds: the decoder still outputs an independent distribution for each pixel (in DALL-E it is the Logit-Laplace).\nIndependence of the latent variables across positions, i.e. the factorized form of $p(z)$. It is this assumption that turns out to be inadequate — and the correlation of the codes contradicts precisely it, not the first assumption.\nMoreover, the first assumption only becomes more plausible with the move to a grid of latents: a rich latent representation explains almost the entire structure of the picture, so the decoder hardly needs to \u0026ldquo;fill in\u0026rdquo; correlated details. Hence the sharp reconstructions of VQ-VAE — in contrast to the blurry samples of a vanilla VAE, where a small latent could not explain the pixel correlations, and the decoder by construction had no right to add them (see the box \u0026ldquo;Why a vanilla VAE blurs pictures and VQ-VAE does not\u0026rdquo;).\nThe ELBO itself remains a valid lower bound on $\\log p_\\theta(x)$ — but for the model we defined. If the prior is factorized and uniform, the ELBO honestly estimates the likelihood of the model \u0026ldquo;independent uniform codes $\\to$ decoder\u0026rdquo;. There is no error in the derivation: it is the model family that is inadequate, not the bound. Such a model is simply a poor generator — it would produce a mosaic.\nFinally, training the prior in the second stage is not a departure from the ELBO either. For the two-stage model $p(x) = \\sum_z p_\\psi(z) p_\\theta(x \\mid z)$, with the encoder and decoder frozen,\n$$\\begin{aligned} \\log p(x) \\ge{} \u0026 \\mathbb{E}_{q_\\phi(z \\mid x)} [\\log p_\\theta(x \\mid z)] + \\mathbb{E}_{q_\\phi(z \\mid x)} [\\log p_\\psi(z)] \\\\ \u0026 - \\mathbb{E}_{q_\\phi(z \\mid x)} [\\log q_\\phi(z \\mid x)], \\end{aligned}$$and maximization over $\\psi$ reduces to $\\mathbb{E}_{q_\\phi(z \\mid x)}[\\log p_\\psi(z)] \\to \\max$ — that is, exactly to training PixelCNN or the Transformer by maximum likelihood on the codes produced by the encoder. So the two-stage scheme is coordinate ascent on a valid ELBO of the final model.\n13. What do the Gumbel-Max Trick and Gumbel-Softmax do, and why are they needed in the dVAE (DALL-E)?\nAnswer The Gumbel-Max Trick is a way to sample from a categorical distribution given by logits $x_1, \\dots, x_k$: add independent samples $g_i$ from the Gumbel distribution to the logits and take $\\arg\\max_i (x_i + g_i)$. This is equivalent to sampling from $\\text{softmax}(x)$, but is still non-differentiable because of the $\\arg\\max$.\nGumbel-Softmax replaces the $\\arg\\max$ with a temperature $\\text{softmax}$: $y_i = \\frac{\\exp((x_i + g_i)/\\tau)}{\\sum_j \\exp((x_j + g_j)/\\tau)}$; the latent vector is assembled as a weighted sum of code vectors $z = \\sum_i y_i e_i$. As $\\tau \\to 0$ the procedure tends to honest sampling with the $\\arg\\max$, which is why the temperature is gradually lowered during dVAE training.\nAs a result, sampling of discrete codes becomes differentiable: the dVAE is trained by honest ELBO optimization — without gradient copying or extra loss terms, which VQ-VAE needed. And the $KL$ term here does not degenerate into a constant and genuinely works as a regularizer.\nA caveat is in order, though: the dVAE prior is factorized over positions. The encoder outputs a $32 \\times 32$ grid of positions, each with its own categorical distribution over the $8192$ code vectors; the prior $p(z)$ takes the code index at each position to be equally probable among all $8192$ options and independent of the other positions. The $D_{KL}$ regularizer pulls the encoder\u0026rsquo;s distribution at each position separately toward this prior, but says nothing about the joint structure of all $32 \\times 32$ codes of a picture (much less about their connection with the text). And a meaningful picture corresponds to a strongly correlated set of codes — a negligible fraction of the space of all combinations. So sampling pictures \u0026ldquo;from the dVAE alone\u0026rdquo; (drawing the code at each position independently) is still impossible: the joint distribution of codes and text is learned in the second stage by the Transformer — just as PixelCNN played this role for VQ-VAE.\n14. Gumbel-Softmax allows training a dVAE by honest ELBO optimization, but in practice discrete tokenizers more often use hard VQ (straight-through + commitment loss). Does this mean that careful regularization of the token prior is not that important?\nAnswer To a large extent, yes. Since the joint distribution of the codes is learned by a separate second-stage model (the Transformer) anyway, how uniform the marginal of an individual token is has almost no effect on the final generation quality: the prior model will learn whatever token statistics the tokenizer produces — just as a language model works perfectly well with the fact that real tokens are not equally probable. So the \u0026ldquo;honest non-constant $KL$\u0026rdquo; is not the decisive practical advantage of Gumbel-Softmax.\nThe value of this approach is rather in the cleanliness of the formulation: an honest ELBO, differentiable sampling, no straight-through hack and no separate terms for the code vectors. Hard VQ is simpler (no temperature schedule and no bias from the relaxation) and works no worse, which is why most modern tokenizers use exactly it.\nOne caveat: pulling toward the uniform distribution is not entirely useless — it encourages high codebook utilization and prevents codebook collapse. But this property matters for training the tokenizer itself well, not for the ability to sample directly.\n15. Describe the two training stages of DALL-E and the inference procedure.\nAnswer Stage 1: a discretized VAE (dVAE) is trained, compressing a $256 \\times 256$ picture into $32 \\times 32 = 1024$ discrete codes from a codebook of size 8192; differentiability of the sampling is provided by the Gumbel relaxation (Gumbel-Softmax).\nStage 2: the dVAE is frozen, and a decoder-Transformer is trained on the concatenation of the caption\u0026rsquo;s text tokens and the picture\u0026rsquo;s code vectors — it learns the joint distribution of texts and pictures by predicting the continuation of the sequence.\nInference: the text description is fed into the Transformer, which autoregressively samples the picture\u0026rsquo;s codes (predicts the distribution of the next code, samples from it and feeds the result back as input), after which the resulting codes are passed through the dVAE decoder, which produces the final picture.\nPractice The hands-on notebook accompanying this post lives right in this blog\u0026rsquo;s repository: notebooks/vae.ipynb. You can open it in Colab — no GPU needed, everything runs on a CPU in a few minutes.\nIn the notebook we build a VAE on MNIST from scratch and reproduce with our own hands all the key pictures of this chapter:\nthe encoder, reparameterization, a Bernoulli decoder and the ELBO as the loss function; reconstructions and samples from the prior $\\mathcal{N}(0, I)$; the latent space map for $M = 2$ — those very clusters of digits; the learned manifold: a uniform grid passed through $\\Phi^{-1}$ and the decoder; the effect of the latent space dimension on sample quality; CVAE: conditioning on the class, generating a chosen digit, and per-class manifolds. The exercises at the end of the notebook are worth doing separately. The most important one is to train the model with zero weight on the $KL$ term: the reconstructions get better, while the samples from the prior turn into garbage. This is exactly the situation we analyzed for VQ-VAE, and seeing it with your own eyes is more useful than reading about it.\n","permalink":"https://jen1995.github.io/posts/vae/","summary":"A deep dive into variational autoencoders: the ELBO and its derivation, the reparameterization trick, CVAE — and the discrete-latent line of work (VQ-VAE, VQ-VAE-2, DALL-E) that grew out of it. With self-check questions and a hands-on notebook.","title":"Variational Autoencoder (VAE)"},{"content":"Hi! I\u0026rsquo;m Eugenia Elistratova. Cohomology Zero is a blog where I write about machine learning and about math for its own sake.\nThe name comes from «Группы и теория гомотопий (трэш трейлер)» — a legendary Russian-language trash trailer for a homotopy theory course. We unapologetically nerd out about math here, patiently waiting for the environment to answer: cohomology — zero?..\nCurrent and upcoming series:\nTransformers from scratch — a step-by-step walk from RNNs with attention to the full Transformer architecture. Generative models — VAEs and friends, with the math worked out carefully. Mathematics — a bit of everything, coming soon. You can find me on GitHub, reach me on Telegram: @evg3307, or write me an email: evg3307@yandex.ru.\n","permalink":"https://jen1995.github.io/about/","summary":"About me","title":"About"}]