Author: admin

DL0180 Convex vs Non-Convex Optimization Landscapes

Explain the difference between convex and non-convex optimization landscapes in high-dimensional deep learning, and why non-convexity in deep networks is empirically more tractable than classical theory predicts.

Answer

A convex objective is shaped like a single bowl: every local minimum is a global minimum, the stationarity condition $\nabla f(x) = 0$ is sufficient for optimality, and convergence rates hold from any initialization. Deep networks violate convexity by construction, because composing affine maps with nonlinearities produces a loss that is invariant under permuting hidden units within a layer, so the surface already carries $\prod_l h_l!$ equivalent copies of every solution and no convex reparameterization of the weights exists. Classical theory then warns that gradient descent could stall in a poor local minimum, but in high dimension the geometry is different: a local minimum requires all $d$ Hessian eigenvalues to be positive at once, which random-matrix and spin-glass arguments make exponentially improbable except very close to the loss floor, so the overwhelming majority of critical points are saddles that gradient noise escapes. Overparameterized networks add a second effect: along the training trajectory the loss satisfies a local Polyak-Lojasiewicz condition, which gives linear convergence to a near-zero-loss solution without any convexity assumption. The solutions SGD reaches are also not isolated, since distinct runs are joined by low-loss curved paths, so the practical obstructions are ill-conditioning, plateaus, and degenerate saddles rather than the isolated bad basins convexity theory teaches you to fear.

(1) What Convexity Buys: local equals global, a checkable optimality certificate, duality gaps, and initialization-independent rates such as $O(1/k)$ for gradient descent on smooth convex losses.
(2) Deep Nets Are Non-Convex By Construction: nested nonlinearities plus permutation and rescaling symmetries mean the loss is a highly multimodal function of the weights even when it is convex in the network output.
(3) Saddles Dominate Critical Points: the fraction of negative Hessian eigenvalues, the index, is almost never zero away from the loss floor, so most flat regions are passable rather than terminal.
(4) Band Structure Of Critical Values: the expected loss of a critical point grows roughly linearly with its index, so high-loss critical points come with escape directions and low-loss ones are what remains.
(5) Overparameterization Tames The Trajectory: in the wide-network regime the loss is nearly a convex function of the output and satisfies a PL inequality locally, giving linear convergence to global training loss near zero.
(6) The Low-Loss Set Is Connected: mode connectivity shows independent solutions are linked by low-loss curves, and permutation alignment removes most of the naive linear-interpolation barrier.
(7) The Real Enemies Are Conditioning And Degeneracy: plateaus with near-zero curvature in every direction, saturated units, and Hessian condition numbers in the thousands cost far more wall-clock time than bad minima.

Two contour maps side by side. Left: a single elliptical bowl with three gradient descent trajectories from different initializations all converging to the same central minimum. Right: a rugged multi-basin surface where three trajectories from different initializations settle into three different basins of comparable depth, with one saddle point marked between two basins.

Figure 1: The structural difference in two dimensions. On the convex bowl every trajectory reaches the same point, so the answer is a property of the objective. On the rugged surface the endpoint depends on initialization, yet the basins reached have comparable depth, which is the two-dimensional cartoon of what actually happens in $d \approx 10^{8}$ : many solutions, not many bad solutions.

The intuition that non-convexity means “gradient descent gets trapped” comes from low-dimensional pictures, where a basin is bounded by walls in both directions. In $d$ dimensions a critical point is a minimum only if the Hessian is positive definite in every one of $d$ directions, and a single negative eigenvalue is enough to make it a saddle with a downhill escape route. Under a Gaussian-random-field or spin-glass model of the loss, the probability that a randomly encountered critical point has index zero decays like $e^{-c d}$ , and the critical values organize into a band: the higher the loss, the larger the expected index. Measured Hessian spectra at converged deep-network solutions match the qualitative picture, with a bulk of near-zero eigenvalues (the flat, degenerate directions created by symmetry and overparameterization), a handful of large positive outliers roughly matching the number of classes, and a thin slightly negative tail. Theory for deep linear networks makes the same point exactly: every local minimum is a global minimum, and the non-convexity manifests as saddles rather than as spurious basins.

Left panel: scatter plot of loss value versus index fraction for simulated critical points, forming an upward band so that critical points with zero or near-zero index cluster at the lowest loss values, with a shaded strip marking the index-zero region. Right panel: log-scale histogram of Hessian eigenvalues at a converged solution showing a very tall bulk concentrated at zero, a thin negative tail, and a dozen isolated large positive outliers.

Figure 2: Left, the band structure: critical points with a high fraction of negative curvature directions sit at high loss, and only points near the loss floor are close to index zero, so “escapable” and “bad” are the same set. Right, a representative Hessian spectrum: thousands of near-zero eigenvalues, a few large outliers that set the effective step-size limit, and a small negative tail that keeps the point technically a saddle.

Mathematical Formulation:
$f(\lambda x + (1-\lambda) y) \leq \lambda f(x) + (1-\lambda) f(y)$
$\nabla^2 f(x) \succeq 0$
$\alpha(x) = k_{-}(x) / d$
$\mathbb{E}[f \mid \alpha] \approx f_0 + c\,\alpha$
$P(\alpha = 0) \sim e^{-c d}$

Where:

$f$ is the training loss as a function of the parameter vector, $x, y$ are two parameter settings, and $\lambda \in [0,1]$ interpolates between them; the first line is the definition of convexity and the second its twice-differentiable equivalent.
$\nabla^2 f(x) \succeq 0$ means the Hessian is positive semidefinite everywhere, which is exactly the condition a deep network’s loss fails to satisfy.
$k_{-}(x)$ counts negative Hessian eigenvalues at a critical point, $d$ is the parameter count, and $\alpha(x)$ is the index fraction; $\alpha = 0$ is a local minimum and any $\alpha > 0$ is a saddle.
$f_0$ is the loss floor and $c$ a positive constant, so the fourth line states the band structure: expected critical value increases with index.
The last line is the random-matrix scaling for an index-zero critical point of a Gaussian random field, which is why bad local minima are measure-zero in practice for large $d$ .

Why Non-Convexity Is Still Solvable:
$\|\nabla f(x)\|^{2} \geq 2\mu\,(f(x) - f^{*})$
$f(x_k) - f^{*} \leq (1 - \mu/L)^{k}\,(f(x_0) - f^{*})$
$\|\nabla f(x)\| \leq \epsilon$
$\lambda_{\min}(\nabla^2 f(x)) \geq -\sqrt{\epsilon}$
$T = O(\epsilon^{-2}\log^{4}(d/\epsilon))$

Where:

The first line is the Polyak-Lojasiewicz (PL) inequality with constant $\mu > 0$ and optimal value $f^{*}$ : wherever the loss is far from optimal, the gradient is large.
PL plus $L$ -smoothness yields the second line, a linear convergence rate with no convexity assumption, and wide networks provably satisfy a local PL condition in a neighborhood of their initialization.
The next two lines define an $\epsilon$ –second-order stationary point: small gradient and no strongly negative curvature, which rules out strict saddles rather than merely first-order stalls.
$T$ is the iteration count in which perturbed gradient descent reaches such a point, with only polylogarithmic dependence on the dimension $d$ ; SGD’s minibatch noise plays the same role for free.

The last piece of the empirical story is that the solutions are not isolated points in separate valleys. Naively interpolating linearly between two independently trained networks produces a large loss barrier, which looks like evidence for distinct basins, but that barrier is largely an artifact of the permutation symmetry: after matching the hidden units of one network to the other, the linear path becomes far flatter, and a low-loss quadratic Bezier path can be found between essentially any pair of solutions. Functionally the minima behave like a single connected low-loss manifold, modulo relabeling, which is why techniques that assume a connected solution set (weight averaging, ensembling along a path, cyclical learning rates) work at all. In production this reframes debugging: a run stuck at high loss is almost never in a bad minimum, it is in a plateau caused by dead units, a saturated nonlinearity, a bad initialization scale, or a step size fighting the largest Hessian eigenvalue.

Line chart of training loss along a path between two independently trained solutions as the interpolation coefficient goes from zero to one. The naive linear interpolation curve rises to a tall barrier in the middle, the permutation-aligned linear interpolation curve rises only slightly, and the curved Bezier path stays essentially flat at the endpoint loss, with a double-headed arrow marking the barrier height.

Figure 3: The apparent wall between two solutions is mostly symmetry, not geometry. Naive linear interpolation crosses a large barrier; permutation alignment flattens most of it, and a learned curved path stays at the endpoint loss the whole way, evidence that the low-loss set is connected rather than a collection of isolated basins.

Property	Convex objective (logistic regression, SVM, LASSO)	Non-convex deep network loss
Stationary points	One connected set of global minima; no saddles	Exponentially many critical points, mostly saddles; global minima replicated by permutation symmetry
Optimality certificate	Zero gradient (or a duality gap bound) proves global optimality	None available; you report training loss and validation metrics, not optimality
Effect of initialization	Affects speed only; the solution is unique up to degeneracy	Selects which solution you land in; scale of initialization can decide whether training starts at all
Role of dimension	Higher $d$ costs compute and risks overfitting	Higher $d$ helps: index-zero critical points become rare and a local PL condition appears
What stalls training	Condition number of the Hessian, non-smooth regularizers	Degenerate plateaus, dead or saturated units, exploding curvature at the step-size limit
Reproducibility	Bitwise-comparable solutions across seeds and solvers	Different weights every seed; only the function learned is comparable

August 20, 2026

DL0179 Grounding DINO vs Faster R-CNN

How does Grounding DINO’s open-vocabulary object detection, which conditions on free-form text, differ from classic Faster R-CNN’s closed-category detection in architecture, training data, and deployment flexibility?

Answer

Faster R-CNN is a closed-set detector: the label space is baked into a final linear layer of shape $(K+1) \times d$ , so the model can only ever emit one of the $K$ categories it was trained on plus background. Grounding DINO keeps the same output contract (boxes with scores) but deletes that fixed head and replaces it with a region-text similarity. A prompt such as “dog . traffic cone .” is encoded by BERT, its token features are fused with image features in three places (the neck feature enhancer, the language-guided query selection stage, and the cross-modality decoder), and each of the 900 decoder queries is scored by a dot product against every text token instead of by a softmax over classes. Because the class list moved out of the weight matrix and into the input, adding a category becomes a string edit rather than an annotation, retraining, and redeployment cycle. The bill arrives in training data (detection plus grounding plus caption-derived boxes instead of one boxed dataset) and in inference cost (a transformer detector with a text encoder instead of a ResNet with a two-layer head).

(1) Fixed Head vs Contrastive Alignment: Faster R-CNN classifies with $\mathrm{softmax}(W f_i + b)$ over $K+1$ rows; Grounding DINO scores each query against each text token and applies a sigmoid per pair, so categories do not compete for a shared probability mass.
(2) Text Is An Input, Not A Label Set: the prompt is runtime configuration, which means the same weights detect “forklift” today and “spilled pallet” tomorrow with no gradient step.
(3) Fusion Happens Early And Often: a late-fusion design that only compares final features would leave the proposal stage text-blind, so Grounding DINO injects language into the encoder, into query selection, and into the decoder.
(4) Detector Family Differs Too: anchors, RoIAlign, and NMS are replaced by DETR-style one-to-one Hungarian matching with learned queries, which removes the anchor and NMS hyperparameters but slows convergence.
(5) Training Data Is The Real Difference: COCO’s 80 categories over ~118k images versus a mixture of Objects365, GoldG, and caption-mined pseudo boxes covering tens of thousands of phrase types.
(6) Flexibility Costs Latency And Calibration: open-vocabulary scores are per-phrase and poorly comparable across phrases, so thresholds must be tuned per prompt rather than set once.

Side-by-side vertical pipelines: left panel shows Faster R-CNN with image only, ResNet-50 plus FPN, Region Proposal Network producing about a thousand proposals, RoIAlign with a two-FC box head, and a softmax over K plus one fixed classes followed by NMS; right panel shows Grounding DINO with two input lanes for image and text prompt, a Swin backbone and a BERT text encoder, a feature enhancer combining deformable self-attention with bi-directional image-text cross-attention, language-guided query selection producing 900 cross-modality queries, and a cross-modality decoder emitting region-text contrastive logits whose label space is the prompt

Figure 1: The vocabulary lives in a different place. On the left it is a row of $W$ ; on the right it is a string that enters the network beside the pixels, so language influences which regions are proposed at all, not only how a finished proposal is labeled.

The training recipe follows from that architecture. Faster R-CNN needs one homogeneous boxed dataset, and every category must appear with exhaustive box annotation, which is why closed-set benchmarks stall around a few hundred classes. Grounding DINO is trained on a mixture of supervision grades: fully annotated detection data (COCO, Objects365 with 365 categories), human phrase grounding data (GoldG, built from Flickr30k Entities and Visual Genome), and caption data whose boxes are pseudo-labeled by a teacher in the GLIP lineage. The reformulation that makes this legal is treating detection as grounding: a detection dataset is just a caption of concatenated category names, so a single per-token alignment loss (focal loss on region-token logits) consumes all three grades plus the usual $L_1$ and GIoU box terms. The payoff is that a Swin-T model reaches about 48.4 AP zero-shot on COCO without seeing a single COCO image, and roughly 27 AP on LVIS minival where the long tail is exactly what a fixed 80-way head cannot express.

Mathematical Formulation:
$p_i = \mathrm{softmax}(W f_i + b)$
$W \in \mathbb{R}^{(K+1) \times d}$
$s_{ij} = q_i^{\top} t_j$
$\hat{p}_{ij} = \sigma(s_{ij})$
$s_i(P) = \max_{j \in P} s_{ij}$

Where:

$p_i$ is the closed-set posterior for RoI $i$ over $K+1$ outcomes, and $f_i \in \mathbb{R}^{d}$ is its pooled RoIAlign feature; the row count of $W$ is the vocabulary, which is why the label space is a weight-shape decision.
$q_i \in \mathbb{R}^{d}$ is the $i$ -th decoder query (Grounding DINO uses $N_q = 900$ ) and $t_j \in \mathbb{R}^{d}$ is the $j$ -th projected text token feature; both are projected into one shared embedding space, so the dot product is directly the logit.
$\sigma$ is the sigmoid, so every region-token pair is an independent binary decision; nothing forces the scores of a query to sum to one across the prompt.
$P$ is the set of sub-word indices belonging to one phrase, and $s_i(P)$ is the phrase-level score obtained by taking the maximum over that phrase’s tokens.
Index ranges are $i \in \{1,\ldots,N_q\}$ and $j \in \{1,\ldots,L_t\}$ with $L_t \leq 256$ BERT sub-word tokens, which is the hard cap on how large a prompt vocabulary can be in one forward pass.

Two-panel figure: left panel is a grayscale bar chart of a Faster R-CNN softmax over person, car, dog, chair, tv, and background probabilities that sum to one, annotated that traffic cone has no column so mass is forced onto the nearest in-vocabulary class; right panel is a grayscale heatmap of sigmoid alignment scores between four decoder queries and the seven prompt sub-word tokens of dog . traffic cone . leash ., with high values where the cone query meets traffic and cone, the dog query meets dog, and the leash query meets leash, and near-zero values for the pavement query

Figure 2: Two different score semantics. The closed-set head must spend all probability mass inside its vocabulary, so an unseen object is mislabeled with confidence; the contrastive head gives each query an independent score per token, and a new phrase adds a column rather than a retrained row of $W$ .

Deployment flexibility is therefore real but not free. The prompt format matters: phrases are period-separated and Grounding DINO uses sub-sentence masking so unrelated category names do not attend to each other, yet cramming hundreds of categories into 256 tokens still degrades both accuracy and score calibration, and long or rare names fragment into sub-words whose max-pooled score behaves differently from a short common noun. Because $\hat{p}_{ij}$ is not normalized across phrases, a single global confidence threshold that is right for “person” is usually wrong for “loose cable”, so production systems keep a per-phrase threshold table. The common industrial pattern is not to serve the open-vocabulary model at all: use it plus a segmenter as an auto-labeler to bootstrap a dataset, then train or distill a fast closed-set detector for the frames-per-second and cost envelope the product actually needs.

Property	Faster R-CNN (closed set)	Grounding DINO (open vocabulary)
Label space	Rows of the classifier weight matrix, fixed at training time	Tokens of the prompt, chosen per request
Localization mechanism	Anchors, RPN proposals, RoIAlign, NMS at inference	900 learned queries, language-guided query selection, one-to-one Hungarian matching, no NMS
Classification	Softmax cross-entropy over $K+1$ classes	Focal loss on per-token region-text logits, sigmoid per pair
Training data	One exhaustively boxed dataset (COCO: 80 classes, ~118k images)	Detection (Objects365, 365 classes) plus grounding (GoldG) plus caption-mined pseudo boxes
Adding one category	Annotate, grow the head by about 5.1k parameters, retrain, revalidate, redeploy	Edit the prompt string, zero new parameters, no retraining
Reported accuracy	About 40 box AP on COCO with R50-FPN, undefined outside its 80 classes	48.4 AP zero-shot COCO and 57.2 AP fine-tuned with Swin-T; 52.5 AP zero-shot with Swin-L
Inference cost	Convolutional backbone plus a two-FC head, edge-deployable, easy to quantize	Transformer detector plus a text encoder; prompt features are cacheable when the vocabulary is fixed
Dominant failure mode	Confidently mislabels unseen objects as the nearest known class	Prompt-sensitive, per-phrase thresholds, degradation past the 256-token prompt budget

August 19, 2026

DL0178 SAM: Promptable Segmentation

How does SAM (Segment Anything Model) use promptable Vision Transformer encoders for zero-shot object segmentation, and what makes its prompt interface generalizable to unseen object categories?

Answer

SAM splits segmentation into one expensive image step and a very cheap prompt step. A ViT-H image encoder (MAE-pretrained, about 632M parameters) takes a 1024×1024 image, converts it into $64 \times 64$ patch tokens, and emits a single image embedding of shape $64 \times 64 \times 256$ that is computed once and cached. A tiny prompt encoder maps whatever the user supplies (clicks, a box, a coarse mask, text) into 256-dimensional tokens, and a two-block two-way cross-attention mask decoder reads the cached embedding together with those tokens to produce three mask logits plus three predicted IoU scores in roughly 50 ms. Generalization to unseen categories comes from the task definition rather than from any category machinery: the objective is “return a valid mask for this prompt”, there is no classification head and no fixed label set, so the model only ever learns a geometry-to-region mapping, supervised by 1.1B masks over 11M images (SA-1B) that a model-in-the-loop data engine collected without class labels. The one thing a category-free interface cannot resolve is ambiguity, and SAM handles it by predicting three nested candidates and back-propagating only the lowest-loss one.

(1) Asymmetric Compute: the ViT encoder runs once per image (about 0.15 s on a GPU) while the decoder runs once per prompt (about 50 ms, light enough for a browser), which is what makes interactive clicking feasible.
(2) Prompts Become Tokens: points and boxes are Fourier positional encodings of their coordinates plus a learned type embedding; a coarse mask prompt is instead embedded by a small convolutional stack and added elementwise to the image embedding.
(3) Two-Way Attention Decoder: prompt tokens attend to image features and image features attend back to the tokens, so a handful of tokens can reshape a dense 4096-token feature map at negligible cost.
(4) Dynamic Mask Head: each mask token is turned by an MLP into the weights of a linear classifier that is dotted with the 4x-upscaled embedding, so the mask is produced by a per-prompt hypernetwork rather than by fixed per-class filters.
(5) Ambiguity-Aware Output: three masks (subpart, part, whole) with a min-over-masks loss and an IoU-prediction token for ranking, which prevents the averaging that destroys single-output models on ambiguous clicks.
(6) Composable Interface: because the prompt is only geometry, any upstream module can drive SAM (a detector’s box, a CLIP text embedding, or a regular grid of points for fully automatic mask generation) with no retraining.

Architecture diagram showing an input image passing through a ViT-H encoder into a cached 64 by 64 by 256 image embedding on the top row, prompts passing through a small prompt encoder into prompt tokens plus three mask tokens and one IoU token on the bottom row, a dashed path where a dense mask prompt is convolved and added to the image embedding, both paths feeding a lightweight two-block two-way cross-attention mask decoder, which outputs three mask logits and three predicted IoU scores

Figure 1: The whole design is an amortization argument. Everything expensive is a function of the image alone and is cached; everything that depends on the prompt is a few million parameters of decoder. Note the absence of any classification branch: the output is a region and a quality score, never a category.

The reason this interface transfers to categories that never appeared in training is that the prompt carries no semantics at all. A click is a coordinate, a box is two coordinates, and the encoder was pretrained to represent generic visual structure, so segmenting a species of coral or an industrial part reuses exactly the same computation as segmenting a cat. What must be learned instead is the notion of “object-ness at the scale the prompt implies”, and SA-1B provides that at massive scale precisely because its data engine was label-free: annotators clicked regions, then a partially trained SAM proposed masks that annotators only corrected, then the final stage generated masks automatically from a $32 \times 32$ point grid with NMS and stability filtering. Ambiguity is the residual difficulty. A single click on a shirt pocket is consistent with the pocket, the shirt, and the person, so a single-output model trained with an averaged loss learns a blurred compromise. SAM emits three masks and only the one with the lowest loss receives gradient, which forces the three slots to specialize by scale.

Three side-by-side stylized figures of a person wearing a shirt with a pocket, each marked by the same star-shaped point prompt on the pocket; the first panel highlights only the pocket as the subpart mask, the second highlights the shirt as the part mask, the third highlights the entire person as the whole mask, each labeled with a different predicted IoU score

Figure 2: One prompt, three valid answers. The min-over-masks loss back-propagates only through whichever slot best matches the ground-truth region, so slot 1 drifts toward subparts, slot 2 toward parts, and slot 3 toward whole objects. The IoU-prediction token then supplies the ranking the interface needs when no human is in the loop.

Mathematical Formulation:
$F = E_{\mathrm{img}}(I) + \mathrm{conv}(m)$
$p_k = \mathrm{PE}(x_k, y_k) + t_{\mathrm{type}(k)}$
$P = (p_1, \ldots, p_{N_p}, o_1, o_2, o_3, o_{\mathrm{iou}})$
$(Z, H) = \mathrm{Dec}(F, P)$
$M_j = \langle \mathrm{MLP}_j(H), \mathrm{Up}_4(Z) \rangle$
$\mathcal{L}_j = 20\,\ell_{\mathrm{focal}}(M_j) + \ell_{\mathrm{dice}}(M_j)$
$\mathcal{L} = \min_{j \in \{1,2,3\}} \mathcal{L}_j$

Where:

$F \in \mathbb{R}^{64 \times 64 \times 256}$ is the cached image embedding, $E_{\mathrm{img}}$ the ViT-H encoder applied to the 1024×1024 image $I$ , and $m$ an optional coarse mask prompt whose convolutional embedding is added elementwise.
$p_k$ is the token for sparse prompt $k$ , built from a random Fourier positional encoding $\mathrm{PE}$ of its coordinates plus a learned embedding $t_{\mathrm{type}}$ for foreground point, background point, box corner, or text.
$P$ is the decoder’s token sequence: $N_p$ prompt tokens plus three learned mask tokens $o_j$ and one IoU token $o_{\mathrm{iou}}$ .
$Z$ is the updated image feature map and $H$ the updated tokens after two blocks of two-way cross-attention; $\mathrm{Up}_4$ is a pair of transposed convolutions giving a $256 \times 256$ map.
$M_j$ is mask $j$ ‘s logit map, formed by dotting the dynamically predicted weights $\mathrm{MLP}_j(H)$ with every spatial location of $\mathrm{Up}_4(Z)$ , then bilinearly resized to 1024×1024.
$\ell_{\mathrm{focal}}$ and $\ell_{\mathrm{dice}}$ are combined in a 20:1 ratio, and only the argmin slot receives gradient; the IoU head is trained separately with an MSE loss against the realized IoU.

Latency Accounting For $K$ Prompts On One Image:
$C_{\mathrm{enc}} = O(N^2 d), \quad N = 4096$
$C_{\mathrm{dec}} = O(N_t N d)$
$T(K) = t_{\mathrm{enc}} + K \, t_{\mathrm{dec}}$
$T(32) = 150\ \mathrm{ms} + 32 \times 50\ \mathrm{ms}$
$T(32) = 1750\ \mathrm{ms}$

The quadratic term belongs entirely to the encoder over $N = 4096$ tokens, while the decoder is linear in the $N_t \approx 8$ tokens it carries, which is why 32 clicks cost 1.75 s with caching but 6.4 s if the ViT is re-run each time. Automatic mask generation exploits the same asymmetry: a $32 \times 32$ grid of 1,024 point prompts, run in batches against one cached embedding, yields a full mask inventory for the image without ever re-encoding it.

Line chart of total latency in milliseconds versus number of prompts on a single image, comparing a caching design that pays 150 milliseconds once plus 50 milliseconds per prompt against a design that re-runs the encoder for every prompt at 200 milliseconds each, with an annotation marking 1.75 seconds versus 6.4 seconds at 32 prompts

Figure 3: Why the encoder/decoder split is the architecture, not an optimization. Caching turns the marginal cost of a prompt from 200 ms into 50 ms, and the gap widens with every additional click, which is exactly the regime an interactive annotation tool or a 1,024-point automatic sweep lives in.

Property	Point clicks	Box	Coarse mask	Text
Encoding path	Sparse: Fourier positional encoding plus fg/bg type embedding	Sparse: two corner tokens with corner-specific embeddings	Dense: conv stack added elementwise to the image embedding	Sparse: a CLIP text embedding used as one token
What it pins down	Location only, so scale stays ambiguous	Location and extent, which resolves most ambiguity	Full spatial support, used for iterative refinement	Semantics with no localization
Typical source	Human clicks, or a 32×32 grid for automatic generation	An upstream detector such as Grounding DINO	SAM’s own previous output in a refinement loop	A user query in the exploratory prototype
Main failure mode	Wrong scale slot chosen when no human ranks the three masks	Loose or overlapping boxes leak into neighboring instances	Errors in the input mask are reinforced across iterations	Weakest of the four; not a released production interface

August 19, 2026

DL0177 Discrete vs Continuous Visual Tokens

What is the difference between discrete visual tokenization (e.g., VQ-VAE, VQ-GAN) and continuous vision embeddings (e.g., ViT patch outputs), and when is each preferred for generation versus understanding tasks?

Answer

Both paths begin identically, with a convolutional or ViT encoder turning the image into a grid of $d$ -dimensional vectors. The only structural difference is whether a quantizer follows. A discrete tokenizer (VQ-VAE, VQ-GAN) snaps each vector to its nearest entry in a learned codebook of $K$ vectors and keeps only the integer index, so an image becomes a string of ids over a finite vocabulary that a jointly trained decoder can invert back to pixels. A continuous representation (ViT patch outputs, CLIP or SigLIP features, KL-VAE latents) keeps the float vector, which preserves far more information but has no finite support, so no softmax and no cross-entropy can be defined over it. That single difference decides the downstream interface: discrete ids plug into next-token or masked-token prediction with exactly the machinery used for text, while continuous features must be projected into an LLM or denoised by a diffusion model. As a default, understanding prefers continuous features because quantization discards the high-frequency detail that OCR and fine-grained VQA depend on, while generation historically preferred discrete tokens because a categorical likelihood is easy to train and easy to sample.

(1) Only The Quantizer Differs: the encoder, the patch grid, and the spatial downsampling factor can be identical; adding a nearest-code lookup converts a float grid into an id grid.
(2) Information Budget: a discrete token carries $\log_2 K$ bits (10 to 18 in practice), while a continuous patch vector carries roughly $d \times 16$ bits of activation, three orders of magnitude more.
(3) Gradient Path: $\arg\min$ is non-differentiable, so VQ needs a straight-through estimator plus codebook and commitment losses, whereas continuous encoders train by plain backpropagation.
(4) Reconstruction Cost: heavy compression makes plain L2 reconstruction blurry, which is why VQ-GAN adds perceptual and patch-GAN losses to keep 16x-downsampled decodes sharp.
(5) Downstream Interface: discrete gives one softmax vocabulary shared with text; continuous gives features for a projector, cross-attention, or a latent diffusion denoiser.
(6) Task Split: continuous features dominate VLM understanding benchmarks; discrete tokens dominate when the goal is a single unified next-token model that also emits pixels.

Two horizontal pipelines: the upper understanding path runs input image to ViT patch encoder to 256 continuous vectors in R^1024 to a linear projector to an LLM emitting text, with no quantizer and no pixel decoder; the lower generation path runs input image to a CNN or ViT encoder to a quantizer that picks the nearest of K codes, to a 16 by 16 grid of integer ids, to a transformer with a softmax over K, to a decoder producing pixels

Figure 1: The same encoder, two endings. Deleting the quantizer leaves continuous features that a projector feeds to an LLM; inserting it buys a finite vocabulary and a pixel decoder at the price of $\log_2 K$ bits per token. Note that the understanding path has no decoder at all, which is why an understanding-only encoder is never required to be invertible.

The practical difficulty of discrete tokenization is that the codebook must be learned through a non-differentiable lookup. The straight-through estimator simply copies the decoder gradient past the quantizer, which is a biased estimate that works only if the encoder output stays close to its assigned code, hence the commitment loss. The characteristic failure is codebook collapse: a few entries win most assignments, the rest receive no gradient and die, and effective vocabulary size stops tracking nominal $K$ . Standard mitigations are EMA codebook updates, low-dimensional $\ell_2$ -normalized codes, dead-code re-initialization, and an entropy bonus on the assignment distribution. Continuous encoders have none of this machinery, but they also cannot be sampled from, since there is no distribution over $\mathbb{R}^{d}$ that a softmax can express, which is precisely why continuous-latent generation requires a diffusion or flow model rather than a token classifier.

Mathematical Formulation:
$z = E(x)$
$k = \arg\min_{j} \lVert z - e_j \rVert_2$
$z_q = e_k$
$\hat{x} = D(z_q)$
$\mathcal{L}_{\mathrm{com}} = \beta \lVert z - \mathrm{sg}(e_k) \rVert_2^2$

Where:

$x$ is the input image, $E$ the encoder, and $z \in \mathbb{R}^{d}$ one continuous patch vector from the encoder grid; keeping $z$ and stopping here is the continuous path.
$e_j$ for $j \in \{1,\ldots,K\}$ are the learned codebook vectors, $k$ is the selected index (the actual token), and $K$ is the vocabulary size.
$z_q$ is the quantized vector fed to the decoder $D$ , and the residual $z - z_q$ is information the model can never recover.
$\mathrm{sg}(\cdot)$ is the stop-gradient operator and $\beta$ (typically $0.25$ ) weights the commitment loss that pulls encoder outputs toward their assigned codes.
The full VQ-GAN objective adds a reconstruction term, an LPIPS perceptual term, and a patch-discriminator term to $\mathcal{L}_{\mathrm{com}}$ ; only the reconstruction term survives in a plain VQ-VAE.

Bit Budget For One 256×256 Image:
$B_{\mathrm{disc}} = 256 \times 14 = 3584$
$B_{\mathrm{cont}} = 256 \times 1024 \times 16$
$B_{\mathrm{cont}} = 4194304$
$B_{\mathrm{cont}} / B_{\mathrm{disc}} \approx 1170$

All four numbers are in bits. A 16x-downsampling tokenizer with $K = 16384$ compresses the image to 256 ids of 14 bits each, about 448 bytes, while a ViT-L/14 tower keeps 256 patch vectors of 1024 bf16 activations, about 512 KiB. That ratio is the whole argument: it is why a discrete sequence is short enough to model autoregressively alongside text, and equally why an OCR-heavy or chart-reading task should not be routed through it.

Left panel shows a two-dimensional scatter of continuous encoder outputs partitioned into nine square cells by dashed boundaries, with a black X codebook entry at each cell center and one highlighted red point joined by an arrow to its nearest code, labelled quantization error. Right panel plots reconstruction FID against bits per token from 10 to 18, with a plain VQ curve that stalls near 5 and rises after 14 bits, a lookup-free or FSQ curve that keeps falling toward 1.2, and a dashed horizontal line marking the continuous KL-VAE floor near 0.74

Figure 2: Left, quantization is a Voronoi partition of the latent space: the id names the cell, and the offset inside the cell is thrown away. Right, approximate published reconstruction results show that a plain VQ codebook stops improving past about $2^{14}$ entries because of codebook collapse, whereas lookup-free and FSQ-style quantizers keep scaling toward the continuous-latent floor.

Property	Discrete (VQ-VAE / VQ-GAN)	Continuous (ViT / KL-VAE)
What a token is	An integer index into a learned codebook of K vectors	An unconstrained d-dimensional float vector
Information per token	log2 K bits, typically 10 to 18	About d x 16 bits, typically 4k to 16k
Gradient path	Non-differentiable argmin, needs a straight-through estimator plus commitment loss	Plain end-to-end backpropagation
Known instability	Codebook collapse and dead codes, with usage often below 10 percent at large K	Latent scale drift without a KL term or normalization
Generation interface	Cross-entropy over K, autoregressive or masked sampling, one vocabulary shared with text	Latent diffusion or flow matching, or a small per-token diffusion head
Understanding quality	Weaker OCR, charts, and fine-grained recognition after the bottleneck	Default choice: CLIP or SigLIP features feed the projector in most VLMs
Pixel decoder	Required, trained jointly with the codebook	Only if the task emits pixels; understanding-only towers have none

August 19, 2026

DL0176 Online vs Offline DPO

How does Online DPO, which samples new responses from the current policy during training, differ from Offline DPO on a fixed preference dataset, and why does the online variant generalize better to out-of-distribution prompts?

Answer

Both variants minimize the same Bradley-Terry log-loss over the implicit reward $\hat{r}_\theta = \beta \log (\pi_\theta / \pi_{\mathrm{ref}})$ , so the difference is not the objective but the distribution the training pairs are drawn from. Offline DPO consumes a dataset collected once from some behavior policy, which makes it a purely supervised procedure: four log-probability forward passes per pair, no decoding, and reference log-probs that can be cached before training starts. Online DPO decodes two fresh completions from the current policy at every step, has them ranked by a preference oracle (reward model, LLM judge, or human), and applies the DPO gradient to that on-policy pair, so the data distribution moves as the model moves. The generalization gap follows directly from this. DPO’s derivation only pins the implicit reward down where the data has support, and a frozen dataset’s support stops moving while the policy keeps drifting, so on unfamiliar prompts the model is being shaped by an extrapolated reward that no label ever corrected. On-policy sampling forces the coverage mismatch back toward 1 by construction, so the KL-regularized objective is enforced exactly at the responses the model will actually emit.

(1) Same Loss, Different Sampler: the gradient formula is identical, but offline draws $(y_w, y_l)$ from a fixed $\mu$ while online draws them from $\pi_\theta$ and labels them on the fly.
(2) An Oracle Becomes Mandatory: online DPO needs a ranker in the loop, so the quality ceiling shifts from the dataset to the reward model or judge.
(3) Coverage Versus Exploration: offline guarantees scale with a concentrability coefficient over the whole response space, while on-policy data only needs local coverage around the current policy.
(4) Off-Policy Drift Is Cumulative: every gradient step moves $\pi_\theta$ further from $\mu$ , so the last epoch of offline training is the most off-policy and the least trustworthy.
(5) Cost Is Decode-Bound: an online step adds two autoregressive generations plus an oracle call, typically 5x to 10x the wall-clock of an offline step at the same batch size.
(6) Hybrids Dominate In Practice: iterative or batched online DPO regenerates the preference set every few thousand steps, recovering most of the on-policy benefit at a fraction of the sampling overhead.

Mechanically, offline DPO is a four-forward-pass classification problem. You score the chosen and rejected completions under the policy and under the frozen reference, take the difference of differences, and push it through a log-sigmoid. Nothing in that loop ever asks what the model would say today. Online DPO inserts two extra stages before the loss, decode → rank → update, and the decode stage is what costs money: generating two 512-token completions is roughly a thousand sequential memory-bandwidth-bound forward steps, whereas the loss itself is four parallel prefills. The payoff is that the pair being contrasted is a pair the model genuinely produced, so the gradient always removes probability mass from an error the model is currently making rather than from an error some other model made months ago.

Two training loops compared. The top row shows offline DPO as a straight left-to-right chain from a frozen preference dataset to a stored pair to the four log-probability forward passes to a gradient step, with no arrow returning to the data. The bottom row shows online DPO as a cycle from a prompt to sampling two completions from the current policy to a preference oracle to the DPO loss to a gradient step, with a feedback arrow returning to the sampling stage.

Figure 1: The structural difference is one arrow. Offline DPO is an open chain whose data never changes, so the mismatch between the dataset and the policy grows monotonically; online DPO closes the loop, which is what keeps the preference signal on-policy at the price of two decodes and one oracle call per prompt.

Mathematical Formulation:
$\hat{r}_\theta(x,y) = \beta \log \frac{\pi_\theta(y \mid x)}{\pi_{\mathrm{ref}}(y \mid x)}$
$\Delta = \hat{r}_\theta(x,y_w) - \hat{r}_\theta(x,y_l)$
$\mathcal{L}(\theta) = -\mathbb{E}\left[\log \sigma(\Delta)\right]$
$\nabla_\theta \mathcal{L} = -\beta\, \sigma(-\Delta)\, g_\theta$
$\text{offline: } (y_w,y_l) \sim \mu(\cdot \mid x)$
$\text{online: } (y_1,y_2) \sim \pi_\theta(\cdot \mid x)$

Where:

$\hat{r}_\theta$ is the implicit reward that DPO optimizes in place of an explicit reward model, and $\beta$ is the inverse KL penalty controlling how far $\pi_\theta$ may move from $\pi_{\mathrm{ref}}$ .
$x$ is the prompt, $y_w$ the preferred completion, $y_l$ the rejected one, and $\Delta$ their implicit-reward margin.
$\sigma$ is the logistic function, so $\sigma(-\Delta)$ is the per-pair gradient weight: pairs the model already ranks correctly contribute almost nothing.
$g_\theta = \nabla_\theta \log \pi_\theta(y_w \mid x) - \nabla_\theta \log \pi_\theta(y_l \mid x)$ is the contrastive direction.
$\mu$ is the fixed behavior policy that produced the offline corpus, typically an SFT checkpoint or a different and often stronger model.

Coverage Requirement:
$C_\mu = \max_{y} \frac{\pi^{*}(y \mid x)}{\mu(y \mid x)}$
$\mu = \pi_\theta \Rightarrow C_\mu \approx 1$

The offline sample complexity carries a factor of $C_\mu$ , the worst-case density ratio between the target policy and the data-collecting policy. On out-of-distribution prompts, $\mu$ assigns near-zero mass to the region the trained policy now occupies, so $C_\mu$ explodes and the bound says nothing. This is visible empirically as likelihood displacement: DPO frequently drives down the log-probability of the chosen response as well as the rejected one, and the displaced mass lands on unlabeled completions that the loss never inspects. On-policy sampling closes exactly that hole, and the DeepMind analysis of the online-offline gap points the same way: offline algorithms remain excellent at classifying preferences while degrading at generating preferred text, and scaling offline data does not close the gap.

Two density panels over a one-dimensional projection of response space. The left panel shows a fixed grey offline data support centred at zero, the policy at step zero overlapping it, and the trained policy shifted to the right so most of its mass sits in a shaded region with no labelled pairs. The right panel shows three sampling distributions at successive training steps that move rightwards together with the policy, so the labelled region always coincides with the policy mass.

Figure 2: Offline training grades the policy where the data is, not where the policy is. Once $\pi_\theta$ has drifted, most of its probability mass lies in the shaded region with no labeled pairs, and the implicit reward there is pure extrapolation. Online sampling drags the labeled region along with the policy, which is why the coverage ratio stays near 1 even as the model changes.

Property	Offline DPO	Online DPO
Pair source	Frozen corpus from a behavior policy, collected once	Two fresh decodes from the current policy per prompt
Labeler in the loop	None at training time; labels are pre-collected	Required: reward model, LLM judge, or human
Cost per step	Four log-prob prefills; reference scores cacheable	Two autoregressive decodes plus an oracle pass, roughly 5x to 10x
Coverage needed	Global: the corpus must cover the responses the final policy will emit	Local: only around the current policy, so the ratio stays near 1
Dominant failure	Likelihood displacement and reward extrapolation off the data support	Reward hacking of the oracle once samples leave its training distribution
Best fit	Human-labeled or teacher-generated data, tight compute, one-shot alignment	A trustworthy proxy oracle, broad prompt distribution, long training runs

August 19, 2026

DL0175 RLAIF and Constitutional AI

What is RLAIF (Reinforcement Learning from AI Feedback), and how does replacing human preference labelers with a constitution-guided critic model like Anthropic’s Constitutional AI change the scaling and alignment properties of RLHF?

Answer

RLAIF keeps every moving part of RLHF and swaps only the label source. Instead of paying annotators to rank two candidate responses, a frozen critic (feedback) model is shown the prompt, both candidates, and a short natural-language principle, and it answers a multiple-choice question whose option log-probabilities become the preference label. Constitutional AI is the best-known instantiation and runs in two stages: a supervised stage where the model repeatedly critiques and revises its own harmful answers against principles sampled from a written constitution, producing the SL-CAI checkpoint, and a reinforcement stage where AI-generated harmlessness comparisons train a preference model that then drives PPO, producing RL-CAI. The scaling change is blunt: a comparison that costs roughly a dollar and minutes of human attention becomes two forward passes, so label volume becomes a compute decision rather than a hiring decision, and labels arrive as calibrated soft probabilities rather than single bits. The alignment change is subtler and more interesting, because the optimization target moves from an implicit, unauditable aggregate of annotator taste into a short document you can read, diff, and version, and because self-critique lets a model be harmless without collapsing into evasive refusals.

(1) Only The Labeler Changes: the preference-model loss, the KL-regularized PPO objective, and the reward-hacking dynamics are identical to RLHF, so RLAIF is not a new algorithm but a new data supply.
(2) Two Stages In Constitutional AI: critique → revision → SL-CAI supervised finetuning, then AI-labeled comparisons → preference model → RL-CAI, with SL-CAI serving as both policy initialization and KL reference.
(3) One Principle Per Judgment: a single principle is sampled from the constitution for each critique and each comparison, which keeps prompts short and spreads coverage across the document instead of demanding the model weigh 16 rules at once.
(4) Soft Labels, Not Bits: the critic emits $p(A \succ B)$ from option token logits, which carries more information per comparison than a human click and lets you distil graded preference strength.
(5) Auditable Objective: a behavior change becomes a text edit to the constitution plus a rerun, rather than a new annotation guideline, a new labeler cohort, and a multi-week collection cycle.
(6) Bounded By The Critic: AI feedback inherits the critic’s position bias, self-preference, and capability ceiling, so it is strongest on axes where judging is easier than generating and weakest on novel or expert domains.

Two-lane pipeline diagram: the top lane runs left to right from a helpful-only model given a red-team prompt, to self-critique against a sampled principle, to revision repeated k times, to supervised finetuning that yields the SL-CAI model; the bottom lane runs right to left from SL-CAI sampling a response pair, to the feedback model choosing A or B given one principle, to a preference model trained on AI labels, to PPO with a KL penalty toward SL-CAI yielding RL-CAI; a constitution box between the two lanes sends arrows up to the critique step and down to the feedback model

Figure 1: The constitution is the only human-written artifact in the loop, and it touches the pipeline twice: once to steer self-critique and revision in the supervised stage, once to condition the critic’s pairwise judgments in the RL stage. Everything downstream of the label, the preference model and the KL-penalized policy update, is unchanged RLHF.

Three implementation details do most of the real work. First, order debiasing: language models judging “(A) or (B)” have a strong position preference, so each pair is scored twice with the options swapped and the two probabilities are averaged, which is cheap because compute is the only cost. Second, chain-of-thought judging with clamped probabilities: letting the critic reason before answering improves agreement with humans but produces near-deterministic 0.0 or 1.0 preferences, so the CAI work clamps the resulting probabilities into a narrow band around 50 percent to keep the preference model from training on overconfident targets. Third, a deliberately mixed label supply: RL-CAI does not eliminate humans, it reassigns them, training the preference model on human helpfulness comparisons alongside AI harmlessness comparisons, because harmlessness is the axis where labeling is unpleasant, high-variance, and easy to specify in writing. Empirically, the reported effect is not merely cheaper labels but a better trade-off, with RL-CAI models both more harmless and less evasive than an RLHF baseline, while the RLAIF study on summarization found AI-labeled and human-labeled policies statistically indistinguishable to human raters despite the critic agreeing with humans only about 78 percent of the time.

Mathematical Formulation:
$y^{(k+1)} \sim \pi_0(\cdot \mid x, y^{(k)}, c_k)$
$p_{\mathrm{AI}}(A \succ B) = \sigma(z_A - z_B)$
$\bar{p} = \frac{1}{2}\left(p_{AB} + 1 - p_{BA}\right)$
$\mathcal{L}_{\mathrm{PM}} = -\log \sigma\!\left(r_\phi(x,y_w) - r_\phi(x,y_l)\right)$
$J(\theta) = \mathbb{E}\left[r_\phi(x,y)\right] - \beta\,\mathrm{KL}\!\left(\pi_\theta \,\|\, \pi_{\mathrm{ref}}\right)$

Where:

$y^{(k)}$ is the response after $k$ revisions of the original answer to red-team prompt $x$ , and $\pi_0$ is the helpful-only model doing both the critique and the rewrite.
$c_k \sim \mathcal{C}$ is one principle sampled uniformly from the constitution $\mathcal{C}$ , typically around 16 short written rules.
$z_A$ and $z_B$ are the critic’s logits for the option tokens of the multiple-choice question, and $\sigma$ is the logistic function, so the label is a probability rather than a bit.
$p_{AB}$ and $p_{BA}$ are the probabilities assigned to the same candidate under the two presentation orders, and $\bar{p}$ is the order-debiased label actually stored.
$r_\phi$ is the preference model, $y_w$ and $y_l$ the winning and losing responses under $\bar{p}$ , and $\mathcal{L}_{\mathrm{PM}}$ the standard Bradley-Terry loss, unchanged from RLHF.
$\pi_\theta$ is the policy being optimized, $\pi_{\mathrm{ref}}$ is the frozen SL-CAI checkpoint, and $\beta > 0$ sets how far the policy may drift before the KL penalty dominates the reward.

Label Volume At A Fixed 100k USD Budget:
$N_{H} \approx 1.3 \times 10^{5}$
$N_{\mathrm{AI}} \approx 5.0 \times 10^{7}$
$N_{\mathrm{AI}} / N_{H} \approx 375$

The numbers assume a fully loaded cost near 0.75 USD per human comparison against roughly 0.002 USD for two short critic forward passes, which is why the same budget buys about 375 times more comparisons. Exact rates move with model size, context length, and vendor pricing, but the two orders of magnitude are robust, and they change what is even worth attempting: per-principle coverage sweeps, full relabeling after a constitution edit, and comparison sets larger than any human dataset ever collected. The binding constraint stops being annotation throughput and becomes critic quality plus reward-model overoptimization.

Left panel: log-log chart of labeling cost in dollars versus number of preference comparisons, with a steep line for human labels at 0.75 dollars each and a much lower line for AI labels at 0.002 dollars each, a horizontal line marking a 100000 dollar budget, and two marked crossing points at 133000 human labels and 50 million AI labels. Right panel: schematic scatter of harmlessness against helpfulness showing an evasive refusal model, a helpful-only RLHF model, an HH-RLHF model and an SL-CAI model lying on a dashed trade-off frontier, with RL-CAI plotted above and to the right of that frontier

Figure 2: Two different claims. On the left, the cost curve is the scaling argument, and at a fixed budget it is worth about 375 times more comparisons. On the right, the frontier shift is the alignment argument: human feedback tends to trade helpfulness against harmlessness, while self-critique plus AI harmlessness labels moves the operating point off that line instead of sliding along it. Right-panel positions are schematic illustrations of the reported qualitative result, not measured Elo scores.

Property	Classic RLHF	Generic RLAIF	Constitutional AI (RL-CAI)
Label source	Paid annotators following a private rubric	An off-the-shelf judge model with a prompt	Critic conditioned on one sampled principle from a published constitution
Cost and throughput	Roughly 1 USD per comparison, days to weeks per batch	Roughly 0.002 USD, limited only by inference capacity	Same as RLAIF, doubled for order swapping and chain-of-thought judging
Label form	Binary choice or coarse Likert rating	Soft probability from option logits	Soft probability, order-averaged and clamped near 50 percent
Auditability of the target	Low, values live in guidelines and labeler judgment	Medium, values live in an ad hoc judge prompt	High, values are a short versioned document you can diff
Human role	Label every comparison	Validate the judge on a held-out sample	Write and revise the constitution, still label helpfulness
Dominant failure mode	Annotator noise, sycophancy, length bias, slow iteration	Position and self-preference bias, silent capability ceiling	Behaviors no principle covers, and policy hacking of a critic that shares its blind spots

August 19, 2026

DL0174 Wide vs Deep LLM Architectures

Under a fixed parameter budget, what are the trade-offs between wider (more hidden dimensions, fewer layers) and deeper (more layers, narrower) Transformer architectures for LLMs, in terms of expressivity, training stability, and inference latency?

Answer

A decoder layer costs about $4d^{2}$ parameters in the attention projections and about $8d^{2}$ in the MLP, so a fixed budget pins the product $L d^{2}$ and the only remaining freedom is the aspect ratio $\rho = d/L$ . Cutting depth by 4x buys only 2x width, which is why the two axes are so asymmetric in practice. Depth buys serial computation: each layer is one more composition step, and multi-hop reasoning and compositional generalization improve with layers in a way that width cannot replicate. Width buys parallel capacity and hardware efficiency: larger GEMMs, higher arithmetic intensity, clean tensor-parallel sharding, and a smaller KV cache per token because the cache grows linearly in $L$ . Empirically the pretraining loss is remarkably insensitive to shape across the middle of the range, so the decision is usually settled by systems constraints rather than by quality. Depth is the axis that costs you latency and training stability, and width is the axis that costs you head redundancy, embedding budget, and eventual depth starvation.

(1) The Budget Is Quadratic In Width: the iso-parameter constraint is $L d^{2} = \mathrm{const}$ , so width is an expensive axis and depth is a cheap one. Doubling $d$ forces you to delete three quarters of the layers.
(2) Depth Is Serial Expressivity: a layer is one composition step, and tasks requiring $k$ chained inferences need enough layers to host them. Depth-efficiency saturates, though, with the useful depth for a given width growing only logarithmically in $d$ .
(3) Width Is Capacity And Throughput: knowledge storage tracks total parameters rather than shape, but wide layers give larger matrix multiplies, better MFU, and more heads per layer for parallel feature lookup.
(4) Depth Destabilizes Training: residual variance accumulates over $L$ branches, so deep stacks need $1/\sqrt{2L}$ output-projection init, Pre-LN or DeepNet-style residual scaling, and they show more loss spikes and attention rank collapse.
(5) Depth Is A Latency Multiplier: autoregressive decode pays a fixed per-layer serial cost (kernel launches plus two tensor-parallel all-reduces), so $L\tau$ adds directly to time per token even though total FLOPs are unchanged.
(6) Small Budgets Flip The Answer: the embedding matrix costs $Vd$ , which at a 128k vocabulary swamps a sub-billion model, so on-device work favors deep and thin shapes.

Log-log plot of model width d against number of layers L along the constant-parameter curve L times d squared equals 6.04e8, with three marked shapes: wide and shallow at 16 layers and width 6144 with aspect ratio 384, balanced at 36 layers and width 4096 with aspect ratio 114, and deep and thin at 64 layers and width 3072 with aspect ratio 48, plus region labels above the curve for exceeding the budget and below for budget unspent

Figure 1: Every shape on the curve costs the same parameters. The curve is a straight line in log-log space with slope $-1/2$ , which is the whole trade-off in one picture: width is quadratically expensive, so moving from 64 layers to 16 layers pays for only a 2x wider residual stream. Real production shapes cluster in the middle band, roughly $\rho$ between 60 and 160.

The stability side is where depth stops being free. In a Pre-LN residual stack the output is a sum of $L$ branch contributions, so without correction the variance entering the final norm grows roughly linearly in depth, gradients through the early layers get squeezed, and the effective learning rate per layer drifts as you add layers. The standard fixes are scaling the residual output projections by $1/\sqrt{2L}$ at init, DeepNet-style residual weighting for very deep stacks, and QK-norm to stop attention logits from drifting. Deep stacks also suffer token-uniformity (rank) collapse, where repeated pure attention drives representations toward a rank-one subspace and the MLPs plus residual path are what hold it off. Width has the opposite profile. Hyperparameter transfer across width is well understood through muP, so you can tune a small proxy model and reuse the learning rate, whereas transfer across depth needed its own theory. The failure mode of extreme width is not divergence but waste: many heads become redundant, the per-layer computation is over-provisioned relative to the number of serial steps available, and the model has no place to put a long chain of dependent inferences.

Mathematical Formulation:
$N \approx 12 L d^{2}$
$\rho = d / L$
$C_{\mathrm{kv}} = 2 L n_{\mathrm{kv}} d_{h} b$
$t_{\mathrm{dec}} \approx W / B + L \tau$

Where:

$N$ is the non-embedding parameter count, from $4d^{2}$ of attention projections plus $8d^{2}$ of MLP weights per layer (a SwiGLU MLP with hidden width $8d/3$ matches the same count).
$L$ is depth and $d$ the residual width; a fixed budget holds $L d^{2}$ constant, so 4x fewer layers buys 2x width.
$\rho$ is the aspect ratio, the single knob this question is really about. GPT-3 175B sits at $12288/96 = 128$ .
$C_{\mathrm{kv}}$ is KV-cache bytes per token, with $n_{\mathrm{kv}}$ key-value heads, head dimension $d_{h}$ , and $b$ bytes per element. Under GQA with a fixed head count it is proportional to $L$ alone.
$t_{\mathrm{dec}}$ is per-token decode latency, $W$ the weight bytes read per step, $B$ the achievable HBM bandwidth, and $\tau$ the fixed serial cost per layer from kernel launches and collective synchronization.
Because $W$ is set by the budget and not by the shape, the $W/B$ term is identical across shapes and only the $L\tau$ term differs.

Iso-Parameter Shapes At About 7.2B Non-Embedding Parameters:
$64 \times 3072^{2} \approx 6.04 \times 10^{8}$
$36 \times 4096^{2} \approx 6.04 \times 10^{8}$
$16 \times 6144^{2} \approx 6.04 \times 10^{8}$
$N \approx 12 \times 6.04 \times 10^{8} \approx 7.2 \times 10^{9}$

These three shapes are interchangeable on the training-cost sheet and nearly interchangeable on validation loss, yet they behave very differently in a serving stack. With GQA at 8 key-value heads and $d_{h} = 128$ in fp16, each layer stores 4 KiB per token, so the 64-layer shape carries 4x the KV cache of the 16-layer shape and supports a proportionally smaller batch at the same context length. On the latency side, a plausible tensor-parallel-8 deployment has roughly $\tau \approx 40$ microseconds per layer once you count two all-reduces plus launch overhead, which turns depth into 2.56 ms of pure serial cost at 64 layers against 0.64 ms at 16 layers, on top of an identical weight-read term.

Two panel chart: left panel shows stacked bars of per-token decode latency for the three iso-parameter shapes, each with an identical 0.90 millisecond weight-read term plus a serial per-layer term of 2.56, 1.44 and 0.64 milliseconds, giving totals of 3.46, 2.34 and 1.54 milliseconds and single-stream rates of 289, 427 and 649 tokens per second; right panel shows KV cache per 8192-token sequence of 2.0, 1.125 and 0.5 gibibytes for 64, 36 and 16 layers

Figure 2: Same parameters, same FLOPs, 2.2x difference in single-stream decode latency. The green weight-read term is fixed by the budget, so all of the spread comes from the $L\tau$ serial term. The right panel shows the second penalty of depth: KV cache is linear in $L$ , so the deep shape gives up 4x in memory per sequence and therefore in achievable batch size.

Property	Deep and thin (L=64, d=3072)	Balanced (L=36, d=4096)	Wide and shallow (L=16, d=6144)
Aspect ratio	48	114	384
Serial steps per token	64 composition steps, best compositional depth	36, enough for most multi-hop patterns	16, depth-starved on chained inference
GEMM shape	Many small matmuls, lower MFU in decode	Good balance on current accelerators	Few large matmuls, highest arithmetic intensity
KV cache per 8k sequence	2.00 GiB	1.125 GiB	0.50 GiB
Decode latency (TP=8 estimate)	3.46 ms per token, 289 tok/s	2.34 ms per token, 427 tok/s	1.54 ms per token, 649 tok/s
Parallelism fit	Pipeline stages are easy, but bubbles hurt small batches	Both TP and PP viable	Tensor-parallel friendly, hard to pipeline
Stability risk	Residual variance growth, loss spikes, rank collapse	Standard recipe suffices	Benign, muP transfers cleanly across width
Tied embedding cost (V=128k)	0.39B parameters	0.52B parameters	0.79B parameters

August 19, 2026

DL0173 Continuous vs Static Batching

How does continuous batching (iteration-level scheduling) in serving systems like vLLM and TGI differ from static request batching, and why does it improve GPU utilization under heterogeneous request lengths?

Answer

Static batching schedules at request granularity: the server collects up to $B$ requests, launches one generation loop, and that loop runs until the longest sequence in the group finishes. Every shorter sequence keeps holding its slot, its KV memory, and its lane in every matmul while emitting nothing useful, and a request that arrives one iteration after launch waits for the entire group to drain. Continuous batching, introduced as iteration-level scheduling in Orca and now the default in vLLM and TGI, moves the scheduling decision inside the loop: after every single forward pass the scheduler retires finished sequences, frees their KV blocks, and admits queued requests into the next pass. The batch is re-formed at each of the thousands of decode steps a request lives through, so the number of sequences actually producing a token stays pinned near the memory limit instead of decaying toward one. Because decoding is memory-bandwidth bound, that sustained batch size converts almost directly into tokens per second, which is why the win grows with the variance of output lengths.

(1) Scheduling Granularity: static batching makes one admission decision per batch, continuous batching makes one per forward pass, which is the entire conceptual difference.
(2) No Head-Of-Line Blocking: a finished sequence is evicted and a waiting request admitted on the next iteration, so queueing delay stops scaling with the longest generation in the current group.
(3) Padding Disappears: Orca’s selective batching batches the position-independent linear layers over flattened tokens and runs attention per sequence with variable-length kernels, so no pad tokens are ever computed.
(4) Paged KV Cache Makes Admission Cheap: PagedAttention allocates KV in fixed blocks rather than reserving a contiguous max-length buffer, so freed blocks immediately become admission capacity for a new prompt.
(5) Decode Is Bandwidth Bound: one decode step reads all model weights once regardless of $B$ , so a sustained batch amortizes that read across more tokens and lifts arithmetic intensity roughly linearly.
(6) The Cost Is Prefill Interference: admitting a long prompt injects a compute-heavy prefill into the loop and stalls every decoding sequence, producing inter-token latency spikes that chunked prefill was designed to remove.

Two stacked Gantt charts of four GPU batch slots across decode iterations for the same eight requests. The top chart shows static batching: the first four requests start together, three of them finish early and leave hatched idle slots until the longest request finishes at iteration 16, at which point the second group of four is admitted and finishes at iteration 23. The bottom chart shows continuous batching: as soon as a short request finishes, a queued request is admitted into the freed slot, so all eight requests complete by iteration 16 with far fewer idle cells.

Figure 1: The same eight requests and the same 50 sequence-iterations of useful decode work, scheduled two ways. Static batching spreads them over 23 iterations at 54% slot occupancy because the group cannot retire until its longest member does; continuous batching backfills every freed slot on the next iteration and finishes in 16 at 78% occupancy. The hatched cells are the entire cost of request-level scheduling: paid compute and reserved KV memory that produce no tokens.

The mechanics that make iteration-level scheduling possible are as important as the policy. A naive implementation would need all sequences in a batch to sit at the same generation position so the whole thing is one dense padded tensor, which is exactly why static batching pads. Selective batching breaks the batch apart: the QKV projections, MLP, and output head act on tokens independently, so they run over a flattened ragged tensor, while attention is dispatched per sequence with its own context length. PagedAttention then removes the second obstacle, memory fragmentation, by storing KV in fixed-size blocks (typically 16 tokens) that need not be contiguous, so a sequence grows block by block and a new arrival can be admitted whenever a handful of blocks are free. With both pieces in place, admission is bounded by free KV blocks rather than by a pre-declared batch shape, and the scheduler’s job becomes a per-iteration packing problem over a memory budget.

Why this shows up as GPU utilization is a roofline argument rather than a scheduling one. Generating one token for one sequence requires reading every weight from HBM, roughly 16 GB for an 8B model in fp16, but only about $2P$ FLOPs of arithmetic, so a batch of one runs at a tiny fraction of peak FLOPs and the GPU is idle waiting on memory. Adding sequences to the same forward pass reuses that single weight read for more tokens, so throughput climbs steeply until the KV-cache reads and finally the matmuls take over. Static batching’s effective batch size decays as its short members retire, so it spends most of its time in the low-intensity regime; continuous batching holds the batch near $B_{\max}$ and stays in the high-intensity regime. Reported end-to-end gains follow from this: Orca measured up to 36.9x throughput over a static FasterTransformer baseline at matched latency, and vLLM measured a further 2x to 4x from paged memory alone.

Line chart of aggregate decode throughput in tokens per second versus the number of sequences in the decode batch, from one to one hundred twenty-eight, for an eight-billion-parameter fp16 model with a one-thousand-token context on a two-terabyte-per-second GPU. The curve rises steeply and then bends as KV cache reads grow, with two marked points: an average effective batch of seventeen for static batching at about eighteen hundred tokens per second, and a sustained batch of thirty-two for continuous batching at about thirty-two hundred tokens per second.

Figure 2: Decode throughput against sustained batch size for an 8B fp16 model with 1K-token contexts on a 2 TB/s GPU. The curve is bandwidth-bound everywhere in this range, so throughput rises almost linearly at small $B$ and bends only as KV reads start to rival the 16 GB weight read. Continuous batching does not move the curve; it moves the operating point, from the average effective batch a draining static group achieves to the memory-limited maximum.

Mathematical Formulation:
$U_{\mathrm{static}} = \frac{\sum_{i=1}^{B} L_i}{B \, L_{\max}}$
$b_{\mathrm{kv}} = 2 \, n_l \, h_{kv} \, d_h \, s$
$t_{\mathrm{step}}(B) \approx \frac{2P + B \bar{L} b_{\mathrm{kv}}}{\mathrm{BW}}$
$\lambda(B) = B \, / \, t_{\mathrm{step}}(B)$
$B_{\max} = M_{\mathrm{free}} \, / \, (\bar{L} \, b_{\mathrm{kv}})$

Where:

$U_{\mathrm{static}}$ is the fraction of batch slot-iterations that produce a token under static batching, with $L_i$ the output length of request $i$ and $L_{\max} = \max_i L_i$ ; it is exactly 1 only when all lengths are equal.
$b_{\mathrm{kv}}$ is KV bytes per token, where $n_l$ is layers, $h_{kv}$ key/value heads (small under GQA), $d_h$ head dimension, $s$ bytes per element, and the leading 2 counts K and V.
$P$ is the parameter count, so $2P$ bytes is the fp16 weight read paid once per iteration regardless of $B$ , and $\mathrm{BW}$ is achievable HBM bandwidth.
$\bar{L}$ is the mean context length in the running batch and $B$ the number of sequences decoding in one forward pass.
$\lambda(B)$ is aggregate decode throughput in tokens per second; it is concave in $B$ because the KV term grows with $B$ while the weight term does not.
$M_{\mathrm{free}}$ is HBM left after weights and activations, so $B_{\max}$ is the admission ceiling the scheduler targets; the formula assumes $\bar{L}$ stays bounded, which it does not for sequences still growing.

Throughput At Two Operating Points (8B fp16, 1K context, 2 TB/s):
$\bar{L} \, b_{\mathrm{kv}} = 1024 \times 128\ \mathrm{KB} = 0.13\ \mathrm{GB}$
$t_{\mathrm{step}}(17) = 18.2 / 2000 = 9.1\ \mathrm{ms}$
$t_{\mathrm{step}}(32) = 20.2 / 2000 = 10.1\ \mathrm{ms}$
$\lambda(17) \approx 1870$
$\lambda(32) \approx 3170$

The byte totals are the 16 GB weight read plus $B \times 0.13$ GB of KV, divided by 2000 GB/s. Nearly doubling the sustained batch costs only 11% more time per step and yields 1.7x the tokens per second, which is the whole economic case for iteration-level scheduling: the extra sequences ride along in memory traffic that was already being paid.

Property	Static (request-level) batching	Continuous (iteration-level) batching	Continuous + chunked prefill
Scheduling unit	One whole generation loop per group	One forward pass	One forward pass with a token budget split across prefill and decode
Wait for a new arrival	Until the longest sequence in the current group finishes	One iteration, if KV blocks are free	One iteration, and its prefill is spread over several
Wasted compute	Pad tokens plus idle slots, growing with length variance	Near zero: ragged attention, no padding	Near zero, with better SM occupancy on decode-only steps
KV memory model	Contiguous buffer reserved for max length per slot	Paged blocks allocated on demand, freed on retirement	Same paged blocks, filled incrementally during prefill
Inter-token latency	Stable within a group, terrible queueing before it	Spikes whenever a long prompt is admitted	Bounded by the chunk size, at slightly higher TTFT
Dominant failure mode	Throughput collapse under heavy-tailed output lengths	KV exhaustion causing preemption and recompute thrash	Chunk size mistuned, trading TTFT against throughput

August 19, 2026

DL0172 Attention Collapse in Ultra-Deep Transformers

What actually causes attention collapse in ultra-deep Transformers with 100+ layers, uniform attention scores or rank degeneration, and how do register tokens, coupled attention, and SkipNet-style gated skipping address it?

Answer

Both phenomena are real, but they are not the same object and only one of them is the root cause. Rank degeneration is a property of the token matrix: after enough layers every row of $X^{(l)}$ converges to the same vector, so the representation is effectively rank-1 and no head can distinguish positions any more. Uniform attention scores are one route to that state but not a requirement, because every softmax attention map is row-stochastic, and a product of row-stochastic matrices contracts any mean-zero component of the token set by the second-largest eigenvalue at each step, even when each individual map looks sharp and interpretable. Dong et al. proved the sharper version of this: pure self-attention without residual connections or MLPs loses rank doubly exponentially in depth, so the residual stream, the FFN nonlinearity, and the normalization placement are the load-bearing defenses that let 100-layer stacks exist at all. A third, opposite pathology is attention entropy collapse, where rows become nearly one-hot as the spectral norm of the query-key map grows; that destroys training stability rather than smoothing representations, and it is the failure that spectral reparameterization targets. Register tokens, cross-layer attention coupling, and gated layer skipping each attack a different factor of the product that drives the representation toward rank-1.

(1) Two Different Objects: entropy of the attention rows measures how a layer mixes, while the rank of $X^{(l)}$ measures what survives after mixing. A model can have healthy per-layer entropy and still be rank-collapsed at layer 90.
(2) Doubly Exponential Rank Loss: for attention-only stacks the distance to rank-1 obeys a cubic recursion, so it is numerically zero within a handful of layers, not after 100.
(3) Non-Uniform Attention Still Collapses: because $A^{(l)}\mathbf{1} = \mathbf{1}$ , the all-ones vector is always a fixed point, and repeated application contracts everything else toward the Perron eigenvector.
(4) The Opposite Failure Mode: entropy collapse to near one-hot rows correlates with exploding query-key spectral norm and produces loss spikes; $\sigma$ Reparam bounds it by spectral-normalizing the logit map with a learned scalar.
(5) Register Tokens: a handful of learnable non-content tokens (typically 4 to 16) give softmax a place to dump probability mass, so content rows are not forced to spread uniformly and the emergent high-norm artifact tokens disappear.
(6) Coupled Attention: adding the previous layer’s raw logits to the current layer’s logits means the effective attention is no longer a fresh independent stochastic matrix per layer, which preserves diversity and gives gradients a direct path across depth.
(7) Gated Skipping: a SkipNet-style gate that drops a block reduces the effective mixing depth below the nominal $L$ , shortening the product of contraction factors and cutting sequential latency at the same time.

The reason a 100-layer Transformer is not already dead is that a residual block computes $X + \mathrm{Attn}(X)$ rather than $\mathrm{Attn}(X)$ . In path-decomposition terms, the identity path carries the full-rank input straight through, and only paths that traverse many attention modules are strongly contracted. The FFN nonlinearity adds a second defense by increasing the Lipschitz constant of the layer map away from a pure average, which is why the empirical decay of the distance to rank-1 in a working model looks geometric with a factor close to 1 instead of doubly exponential. The remaining problem at extreme depth is that these defenses only slow the contraction: with a per-layer factor of 0.95, a 120-layer stack still retains only about 0.2% of the initial token spread, which shows up as flat similarity matrices, near-duplicate hidden states in the last third of the network, and layers whose removal barely changes the loss.

Two panel chart. Left panel plots distance from rank-1 on a log scale against layer index for four settings: pure attention plunging off the chart by layer five, pre-LN residual only decaying geometrically to about 1e-4 at layer 120, residual plus MLP decaying to about 0.03, and a stabilized configuration staying above 0.5. Right panel plots attention entropy normalized by log N against layer index, showing one curve rising toward the uniform limit of 1.0, one curve falling toward zero labeled entropy collapse, and a stabilized curve staying inside a shaded healthy band.

Figure 1: The two collapse modes are measured on different axes. Left: rank degeneration is catastrophic for attention-only stacks, already below $10^{-8}$ by layer 5, and merely slow once residuals and FFNs are present. Right: entropy can fail in either direction, drifting up toward the uniform limit $\log N$ (over-smoothing) or down toward zero (entropy collapse and loss spikes), and a healthy deep model must stay in the band between them.

Mathematical Formulation:
$\mathrm{res}(X) = X - \mathbf{1}x^{\top}$
$\|\mathrm{res}(X^{l+1})\| \leq c\,\|\mathrm{res}(X^{l})\|^{3}$
$c^{1/2}\|\mathrm{res}(X^{L})\| \leq \left(c^{1/2}\|\mathrm{res}(X^{0})\|\right)^{3^{L}}$
$A^{(l)}\mathbf{1} = \mathbf{1}$
$\|A^{(l)}u\| \leq \lambda_2^{(l)}\|u\|$
$H(a_i) = -\sum_{j=1}^{N} a_{ij}\log a_{ij}$
$0 \leq H(a_i) \leq \log N$

Where:

$X^{(l)} \in \mathbb{R}^{N \times d}$ is the token matrix at layer $l$ , with $N$ tokens of width $d$ , and $L$ is the total depth.
$\mathrm{res}(X)$ is the distance to the nearest rank-1 matrix whose rows are all equal; $\mathbf{1}$ is the all-ones vector and $x$ the common row it would collapse to.
$c$ collects the head geometry, roughly $4\gamma\beta/\sqrt{d_{qk}}$ , where $\gamma$ and $\beta$ bound the value and query-key weight norms and $d_{qk}$ is the head dimension.
The cubic recursion compounds into a doubly exponential bound with exponent $3^{L}$ , which is why attention-only depth is hopeless while residual depth is merely expensive.
$A^{(l)} \in \mathbb{R}^{N \times N}$ is the row-stochastic attention map, so $\mathbf{1}$ is always its eigenvector with eigenvalue 1, and $u$ is any mean-zero deviation across tokens, $\mathbf{1}^{\top}u = 0$ .
$\lambda_2^{(l)}$ is the second-largest eigenvalue modulus of $A^{(l)}$ , strictly less than 1 whenever all entries are positive; the surviving spread after $L$ layers scales like $\prod_l \lambda_2^{(l)}$ .
$H(a_i)$ is the entropy of attention row $i$ ; the upper bound $\log N$ is the uniform row (maximal mixing) and 0 is the one-hot row (entropy collapse).

Collapse Speed, Two Regimes:
$0.9^{3^{5}} = 0.9^{243} \approx 7 \times 10^{-12}$
$0.95^{120} \approx 2.1 \times 10^{-3}$

The first line is the attention-only regime: five layers are enough to destroy the representation. The second is the realistic regime for a 120-layer residual stack with a mild per-layer contraction, and it is the number that motivates the three interventions. Register tokens change the geometry of each individual $A^{(l)}$ so that content rows keep structure instead of hedging uniformly; the attention sink observed in decoder-only LLMs, where the first token absorbs a large share of the mass, is the same phenomenon arising without being designed. Coupled attention changes the product itself, since adding the previous layer’s logits makes consecutive maps correlated rather than independent draws. Gated skipping changes the number of factors in the product, and it also attacks a separate ultra-deep pathology: under Pre-LN the output variance grows with depth, so late blocks approach the identity and contribute almost nothing, which means paying their latency buys no capacity.

Left to right architecture diagram of a deep Transformer segment. An input sequence box containing N content tokens plus R register tokens feeds block l minus 1, then block l, then a skip gate, then block l plus one, then an output box. A dashed arc above the stack shows attention logits from block l minus 1 being added to block l as coupled attention, an annotation above the input box explains that registers give softmax a non-content place to dump attention mass, and a routed path below the stack shows the gate bypassing block l plus one when its gate value is zero, reducing effective mixing depth.

Figure 2: Three interventions at three different levels of the same product. Registers reshape each individual attention map, coupled attention correlates consecutive maps so the stack stops re-averaging from scratch, and the skip gate removes factors entirely by lowering the effective depth. The residual stream and FFN remain the baseline defense underneath all three.

Property	Register tokens	Coupled attention	SkipNet-style gating
What it changes	The geometry of each single attention map	The correlation between consecutive maps	The number of maps in the product
Mechanism	Extra learnable non-content tokens absorb attention mass, so content rows need not spread uniformly	Pre-softmax logits of layer l minus 1 are added to layer l, giving a residual path through attention itself	A learned gate executes or bypasses a block per input, so effective depth is data dependent
Cost	Sequence grows to N plus R, so prefill cost grows quadratically in that length; registers are discarded at the head	Must retain the previous layer’s logit tensor, which is $O(hN^2)$ activation memory per block	Gate parameters plus non-differentiable routing, usually trained with a straight-through or RL estimator
Where it fails	Cannot be bolted on after pretraining, and too many registers waste context without adding capacity	Incompatible with kernels that never materialize the score matrix, so it fights FlashAttention-style fusion	Ragged per-example depth hurts batched throughput, and gates can collapse to always-on or always-off
Diagnostic it fixes	High-norm artifact tokens and noisy attention maps	Entropy drifting toward the uniform limit in late layers	Late blocks that behave as the identity and can be pruned for free

August 19, 2026

DL0171 Jamba Hybrid SSM-Transformer Architecture

Explain Jamba and hybrid SSM-Transformer architectures. Why blend Mamba blocks with attention blocks instead of using either alone?

Answer

Jamba (AI21 Labs, 2024) is a decoder-only language model whose 32 layers are built from a repeating 8-layer block that contains seven Mamba mixers and one attention mixer, with a mixture-of-experts feed-forward on every other layer (16 experts, top-2 routing, 52B total and 12B active parameters). Every layer still writes into the same residual stream, so the only thing that changes from layer to layer is which token-mixing operator runs: a selective state space recurrence that is linear in sequence length and keeps a fixed-size state, or full self-attention that is quadratic in prefill and keeps a growing KV cache. The blend exists because the two operators fail in opposite directions. A pure Mamba stack of this size does not reliably form induction heads, so it degrades at few-shot in-context learning, format copying, and verbatim retrieval, while a pure Transformer of the same depth carries a KV cache eight times larger and loses long-context throughput. Four attention layers out of thirty-two are enough to restore the copy and in-context behaviour, and the remaining twenty-eight Mamba layers deliver the memory and speed: at a 256K-token context the KV cache is 4 GB instead of 32 GB, and long-context throughput is roughly 3x that of a comparable all-attention MoE model.

(1) Two Mixers, One Residual Stream: Mamba and attention layers are interchangeable drop-ins at the token-mixing position, so no fusion or adapter machinery is needed to combine them.
(2) A 1:7 Attention Ratio: one attention layer per eight-layer block was chosen after ablations showed 1:3 and 1:7 score alike, so the cheaper ratio wins.
(3) MoE Buys Capacity Without FLOPs: replacing the MLP on every second layer with 16 experts and top-2 routing raises total parameters to 52B while keeping 12B active per token.
(4) Attention Supplies Exact Recall: the few attention layers are the only components that can address an arbitrary earlier token exactly, which is what induction heads and few-shot copying need.
(5) Mamba Supplies Position: the recurrence is inherently ordered, so Jamba ships with no explicit positional encoding and RoPE gave no measurable gain.
(6) The Payoff Is Memory, Not Perplexity: a 256K context fits in 4 GB of KV cache, letting a 52B model serve 140K tokens on a single 80 GB GPU in int8.

A ribbon of 32 layer cells split into four identical eight-layer blocks, with the fourth cell of every block shaded as an attention mixer and the other twenty-eight shaded as Mamba mixers, above an expanded view of one block showing each layer as a mixer cell plus a feed-forward cell where MLP alternates with a 16-expert MoE

Figure 1: The whole architecture is one repeated block. Attention appears at layer 4 of every 8, giving 4 attention layers out of 32, and the MoE feed-forward alternates with a dense MLP so that capacity grows without raising the per-token FLOP count.

The reason neither operator survives alone is a difference in what they can store. A selective SSM compresses the entire prefix into a state of fixed size, so its recall is lossy and content-addressed by whatever the gating learned to keep. Theoretical and empirical work on copying shows that a fixed-state recurrent model needs state size proportional to the string it must reproduce, whereas attention copies with a constant number of heads. That is exactly the gap seen in Jamba’s ablations: a pure Mamba model trained on the same 250B tokens tracks the Transformer on log-probability benchmarks yet collapses on few-shot tasks where the model must imitate the label format shown in the prompt, because it never develops induction heads. Inserting attention into one layer in eight repairs this, and the repaired model then inherits Mamba’s cost profile for the other seven eighths of its depth. Ordering matters too, since spreading the attention layers evenly through the stack lets every group of Mamba layers be followed by an exact-lookup step rather than concentrating all lookup capability at one depth.

Mathematical Formulation:
$h_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t$
$y_t = C_t h_t + D x_t$
$\bar{A}_t = \exp(\Delta_t A)$
$C_{\mathrm{ssm}} = O(L d N)$
$C_{\mathrm{attn}} = O(L^2 d)$
$M_{kv} = 2 b L n_a h_{kv} d_h$

Where:

$h_t$ is the SSM state at step $t$ , $x_t$ the layer input and $y_t$ the layer output; $D$ is the skip term.
$\bar{A}_t, \bar{B}_t, C_t$ are the discretized transition, input, and output matrices. In a selective SSM they depend on the current token, which is what lets the layer decide what to keep and what to forget.
$\Delta_t > 0$ is the input-dependent step size; a large $\Delta_t$ overwrites the state with the new token, a small one carries the old state forward.
$L$ is sequence length, $d$ the model width, and $N$ the state dimension (16 in Mamba), so SSM cost is linear in $L$ while attention prefill is quadratic.
$b$ is bytes per element (2 in fp16), $n_a$ the number of attention layers only, $h_{kv}$ the KV heads under GQA, and $d_h$ the head dimension; the leading 2 counts keys and values.
The hybrid changes exactly one factor in $M_{kv}$ , namely $n_a$ , which drops from 32 to 4.

KV Cache At A 256K Context (fp16, 8 KV heads, $d_h = 128$ ):
$M_{\mathrm{hybrid}} = 16\ \mathrm{KB} \times 262144$
$M_{\mathrm{hybrid}} = 4\ \mathrm{GB}$
$M_{\mathrm{full}} = 128\ \mathrm{KB} \times 262144$
$M_{\mathrm{full}} = 32\ \mathrm{GB}$

The 28 Mamba layers contribute a per-sequence state of only a few megabytes that does not grow with $L$ at all, so the entire cache curve of the hybrid is set by its four attention layers. That is what turns long context from a memory problem into a compute problem: batch size at 256K stops being limited by cache residency, and decode throughput at 128K measures roughly 3x a comparable all-attention MoE because far fewer bytes move per generated token.

Log-log chart of KV cache memory in gigabytes versus context length in tokens, with a steep line for a 32-layer pure Transformer at 128 KB per token, a parallel line eight times lower for the Jamba hybrid at 16 KB per token, a flat line near seven megabytes for a pure Mamba model with fixed state, a horizontal marker for a single 80 GB GPU, and an annotation at 256K tokens showing 4 GB versus 32 GB

Figure 2: Only the attention layers create a cache that grows with context. Cutting them from 32 to 4 shifts the whole line down by a constant factor of 8x, while a pure SSM stack is flat because its state size is independent of $L$ . The hybrid buys most of the SSM memory profile at the price of four layers.

Property	Pure Transformer	Pure Mamba	Hybrid (1 attention per 8)
Prefill cost	Quadratic in every layer	Linear via hardware-aware parallel scan	Linear in 28 layers, quadratic in 4
Per-sequence memory at 256K	32 GB of KV cache	A few MB of fixed state	4 GB, an 8x reduction
Exact copy and few-shot format following	Strong, induction heads emerge reliably	Weak at scale, degrades on in-context label imitation	Matches the Transformer with 4 attention layers
Position information	Requires RoPE or ALiBi	Implicit in the recurrence	No explicit encoding needed, RoPE gave no gain
Serving ecosystem	Mature: paged attention, prefix caching, speculative decoding	Custom scan kernels, thin tooling	Needs both paths, plus state handling in the scheduler

August 19, 2026