Tag: LLM

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

DL0170 Mamba Selective SSM vs Attention

What is Mamba’s Selective State Space Model (S6), and how does its input-dependent gating compare to Transformer self-attention in computational complexity and long-range recall?

Answer

Mamba’s S6 layer is a linear recurrence that carries a fixed-size hidden state $h_t$ per channel and reads it out with a learned projection, exactly like the earlier S4 state space model. The one change that matters is selection: the timestep $\Delta_t$ and the input and output maps $B_t, C_t$ are linear functions of the current token instead of fixed parameters, which turns the discretized decay $\bar{A}_t = \exp(\Delta_t A)$ into an input-dependent forget gate. That single change gives the model the ability to skip filler tokens and write informative ones into the state, but it destroys linear time invariance, so the layer can no longer be evaluated as one long FFT convolution and instead needs a hardware-aware parallel associative scan. On complexity the comparison is lopsided in Mamba’s favor: sequence mixing is $O(L d N)$ instead of $O(L^2 d)$ , and decoding needs a constant $O(d N)$ state rather than a KV cache that grows with context. On recall the comparison reverses, because a bounded state is a lossy summary of the past, so tasks that need verbatim retrieval of an arbitrary earlier span degrade in a way that attention’s exact, if expensive, cache does not.

(1) Selection Is The Whole Idea: $\Delta_t, B_t, C_t$ become projections of $x_t$ , while $A$ stays a learned diagonal matrix whose effective decay $\exp(\Delta_t A)$ is nonetheless input-dependent.
(2) $\Delta_t$ Is A Forget Gate: small $\Delta_t$ gives $\bar{A}_t \approx I$ and the token is ignored, large $\Delta_t$ drives $\bar{A}_t$ toward zero and overwrites the state with the new input.
(3) Time Invariance Is Lost: S4 could precompute one global convolution kernel, Mamba cannot, so training uses a work-efficient associative scan in SRAM with activation recomputation instead of an FFT.
(4) Linear Compute: mixing costs $O(L d N)$ with $N = 16$ in Mamba-1, so the quadratic attention term dominates once $L$ approaches the model width $d$ .
(5) Bounded Decode State: $2 d N$ values per layer regardless of context length, versus $2 L d$ for an MHA KV cache, which is what makes million-token streaming decode cheap.
(6) Recall Is The Trade: fixed capacity caps exact copying and multi-query associative recall, which is precisely why production long-context models are usually hybrids with a few full-attention layers.

Diagram of four input tokens the, key, is, 7Q4 each producing an input-dependent timestep Delta_t shown as 0.04, 1.6, 0.05 and 2.1 with skip or write annotations, feeding a chain of fixed-size state boxes h_1 through h_4 linked by input-dependent decay arrows, and a readout row of outputs y_1 through y_4, with a side box noting the state holds 2dN values for any sequence length

Figure 1: The recurrence is ordinary, the gate is not. Because $\Delta_t$ is computed from the token itself, a filler word produces $\bar{A}_t \approx I$ and passes through without disturbing the state, while a rare identifier produces a large step that writes into the state and forgets older content. The state width never depends on $L$ , which is both the efficiency win and the recall ceiling.

The engineering consequence of dropping time invariance is that the layer becomes memory-bound rather than FLOP-bound. Materializing the expanded states, which have shape $(B, L, 2d, N)$ , would move far more bytes through HBM than the arithmetic justifies, so the reference implementation fuses discretization, scan, and readout into one kernel that keeps states in SRAM and recomputes them during the backward pass. The recall side has a cleaner theoretical story. Attention stores every key and value, so retrieving a specific earlier token is a lookup, whereas a selective SSM must have decided at write time to keep that token, and its $2 d N$ slots bound how many distinct key-value associations can survive. This is why Mamba matches or beats Transformers on language modeling perplexity, audio, and DNA, yet lags on induction-head style copying and needle-in-a-haystack retrieval, and why Mamba-2 raises $N$ from 16 to as much as 256 while hybrids keep a small number of attention layers to do the exact lookups.

Mathematical Formulation:
$h'(t) = A h(t) + B x(t)$
$y(t) = C h(t)$
$\Delta_t = \mathrm{softplus}(W_{\Delta} x_t + b)$
$B_t = W_B x_t, \quad C_t = W_C x_t$
$\bar{A}_t = \exp(\Delta_t A)$
$h_t = \bar{A}_t h_{t-1} + \Delta_t B_t x_t$
$y_t = C_t^{\top} h_t$

Where:

$y_t$ is the layer output for token $t$ and $h_t \in \mathbb{R}^{N}$ is the hidden state held for one channel, so the full layer keeps $2 d N$ values under the standard expansion factor of 2.
$x_t$ is the input activation, $t \in \{1, \ldots, L\}$ indexes the sequence, and $h_0 = 0$ is the required initial condition.
$A$ is a learned diagonal matrix parameterized as $A = -\exp(A_{\log})$ so every eigenvalue is negative and the discrete decay stays stable.
$\Delta_t > 0$ is the input-dependent timestep, and $W_{\Delta}, W_B, W_C$ are the low-rank projections that make the layer selective; $b$ is initialized so that $\Delta_t$ starts in a useful timescale range.
$\bar{A}_t$ is the zero-order-hold discretization of $A$ , and $\Delta_t B_t x_t$ is the simplified input term Mamba uses in place of the exact $\bar{B}_t$ .
$d$ is the model width, $N$ the state dimension (16 in Mamba-1), and $L$ the sequence length; the scan is $O(L d N)$ work with $O(\log L)$ depth.

Decode Memory At 32k Context (values per layer):
$L = 32768, \quad d = 2048, \quad N = 16$
$\mathrm{cache} = 2 L d = 1.34 \times 10^{8}$
$\mathrm{state} = 2 d N = 6.55 \times 10^{4}$
$\mathrm{ratio} = L / N = 2048$

Those are element counts, so in bytes the gap narrows a little when the cache is fp16 and the SSM state is fp32, and it narrows further with grouped-query attention, which divides the cache by the query-to-key-value head ratio. The structural point survives every such adjustment: the attention term is proportional to $L$ and the SSM term is not, so past a few thousand tokens the recurrent model is decoding from a constant working set while the Transformer is streaming a cache that eventually dominates both memory and bandwidth.

Two log-log panels: the left panel plots forward-pass FLOPs per layer against sequence length for a self-attention block with a quadratic term and a Mamba block with a linear term, with a dash-dotted vertical crossover line near two thousand tokens, and the right panel plots decode state in megabytes showing a rising fp16 KV cache line against a flat constant Mamba state line at 0.26 megabytes

Figure 2: Below roughly $L \approx d$ the wider Mamba block, which expands the channel dimension by 2, actually costs more than attention, because both are dominated by their projection matmuls. Above it the $O(L^2 d)$ term takes over, and at 128k tokens the attention layer needs about 20 times the FLOPs and a per-layer cache about 4,000 times larger than the constant SSM state.

Property	Selective SSM (Mamba S6)	Self-attention (MHA)
Sequence-mixing cost	O(L d N), linear in L	O(L^2 d), quadratic in L
Cost per decoded token	O(d N), independent of context	O(L d), grows with context
State carried between tokens	2 d N values per layer, fixed	2 L d values per layer, unbounded
Training parallelism	Associative scan with recomputation; no FFT convolution, since the recurrence is time-varying	Pure matmuls, no sequential dependence at all
Exact retrieval and copying	Lossy; bounded by state capacity, degrades as the number of stored associations grows	Exact within the window; copying long spans is easy
Natural fit	Audio, DNA, streaming, very long inputs summarized rather than quoted	In-context retrieval, many-shot prompts, verbatim citation

August 19, 2026

DL0169 Softmax-1 Off-by-One Attention

What is the Softmax-1 (off-by-one) modification to attention, and why does subtracting one from the denominator improve length generalization and register-token behavior in recent Transformer architectures?

Answer

Softmax-1, also called off-by-one or quiet attention, replaces the attention normalizer $\sum_j e^{z_j}$ with $1 + \sum_j e^{z_j}$ . That extra constant is exactly what you get by appending a virtual key with logit 0 and value vector 0, so the weights over real tokens now sum to $S/(1+S)$ , which is strictly below one, and a head that finds nothing relevant can emit the zero vector. Ordinary softmax has no such option, because its weights are forced to sum to one whether or not any key matches. Heads that want to be a no-op therefore learn to dump their leftover mass on a low-information token, usually the first token or a delimiter, and to inflate that token’s residual to enormous magnitude: the massive activations that dominate activation-quantization error and the artifact tokens that pollute ViT attention maps. The name refers to the denominator being off by one; the operational reading is that one unit of attention budget is subtracted from the real keys and parked in a null option that costs nothing to select.

(1) One Constant, No New Parameters: adding 1 to the denominator is algebraically identical to concatenating a zero-logit, zero-value key, so the change touches one line of the kernel and adds no weights.
(2) Mass Below One Is The Whole Point: the head gains an explicit “abstain” action, and the output norm $\lVert o \rVert$ can go to zero without any token’s value vector having to be zero.
(3) Absolute Instead Of Relative Scores: ordinary softmax is invariant to a constant shift of all logits, so “nothing is similar” is inexpressible; the fixed zero reference makes the logit level itself meaningful.
(4) The Sink Becomes A Mechanism, Not A Token: vanilla models route no-op mass through a specific KV slot that a sliding window may evict, whereas the constant is always present, which is what stabilizes long-context and windowed decoding.
(5) Registers Stop Being Garbage Dumps: once abstention is free, register and background tokens keep their global-storage role without absorbing surplus attention, so their norms stay moderate and attention maps stay readable.
(6) Production Form Is A Learned Sink Logit: replacing the 1 with $e^{s}$ for a per-head learned $s$ strictly generalizes Softmax-1, which is the $s = 0$ special case.

Two-panel grouped bar chart of attention weights over six keys plus a null slot. In panel A every logit is below zero and ordinary softmax still renormalizes to sum one, while softmax-1 keeps only 0.17 of the budget and sends 0.83 to the null slot. In panel B one logit is six and the two methods are nearly identical, with the null slot receiving only a quarter of a percent.

Figure 1: The gate is selective, not a uniform rescaling. When no key matches (panel A), ordinary softmax renormalizes a row of weak scores into a confident-looking distribution, while Softmax-1 keeps 17% of the budget on real keys and abstains with the rest. When a real match exists (panel B), the constant 1 is negligible against $e^{6} \approx 403$ and the two are indistinguishable, so the modification only fires where the head has nothing to say.

Mathematical Formulation:
$\mathrm{softmax}(z)_i = \frac{e^{z_i}}{\sum_{j=1}^{L} e^{z_j}}$
$\mathrm{softmax}_1(z)_i = \frac{e^{z_i}}{1 + \sum_{j=1}^{L} e^{z_j}}$
$S = \sum_{j=1}^{L} e^{z_j}$
$m = \frac{S}{1 + S}$
$m = \sigma(z + \log L)$
$p_i = \frac{e^{z_i}}{e^{s} + \sum_{j=1}^{L} e^{z_j}}$
$o = \sum_{i=1}^{L} p_i v_i$

Where:

$z_i = q^{\top} k_i / \sqrt{d}$ is the scaled dot-product logit for key $i$ , and $L$ is the number of visible keys after masking.
$S$ is the ordinary denominator, so the added 1 is just $e^{0}$ : the virtual key’s exponentiated logit.
$m$ is the total attention mass retained by real keys; the complement $1/(1+S)$ is the abstention share, and it is positive for every finite logit row.
The uniform-logit identity uses $\sigma$ for the logistic function: with all $L$ logits tied at $z$ , the mass is a logistic gate whose half-way point sits at $z = -\log L$ .
$s$ is a learned per-head sink logit replacing the constant, and $s = 0$ recovers Softmax-1 exactly.
$o$ is the head output; because $\sum_i p_i$ can be near zero, $o \approx 0$ is reachable without constraining any $v_i$ .

The length-generalization argument has two halves. First, in ordinary softmax a no-op head’s surplus mass is $1 - m_{\mathrm{rel}}$ , and where it lands depends on how many irrelevant keys exist, so a head calibrated at a 4k training length re-partitions its budget when the context reaches 32k; the model compensates by pinning a token whose residual magnitude was tuned at the training length, and Sun et al. measured such coordinates in the thousands while the median activation sits near 0.1. Second, that sink is a KV slot, not an architectural constant. StreamingLLM showed the consequence directly: evict the first few tokens from a sliding window and perplexity explodes, while pinning four sink tokens keeps generation stable out to millions of tokens. Softmax-1 turns the sink into a term of the denominator that no eviction policy can delete, which is why windowed and extrapolated decoding stop depending on cache bookkeeping. The honest caveat is visible in the uniform-logit identity above: the abstention gate is absolute but its threshold still drifts as $-\log L$ , so pushing a head to stay quiet across a 256x context increase costs about 5.5 nats of logit headroom, which is precisely why deployed variants learn $s$ per head rather than freezing it at zero.

Line chart of total attention mass on real keys versus a uniform attention logit, for context lengths 512, 8192 and 131072. Each softmax-1 curve is a logistic whose half-way point sits at minus log L, marked by dotted vertical lines, while a flat dash-dotted line at mass one shows that ordinary softmax retains the whole budget regardless of logit level or length.

Figure 2: Softmax-1 converts the normalizer into a logistic gate on the absolute logit level, whereas ordinary softmax pins the mass at one for every logit row and every length. The gate’s half-way point moves left by $\log L$ , so a head that abstains comfortably at 512 keys leaks over half its budget at 131k unless it lowers its logits by about 5.5 nats, which is the argument for a learned sink logit rather than a hard-coded 1.

The register story is the same pressure seen in vision. Darcet et al. found that DINOv2, CLIP, and DeiT-III spontaneously repurpose a small fraction of low-information background patches as global scratch space, giving those patches norms roughly an order of magnitude above their neighbours and visibly corrupting attention maps and dense-prediction features; adding explicit register tokens removed the artifacts. Registers supply somewhere to write, while Softmax-1 supplies permission not to write, and the two are complementary rather than competing. Bondarenko et al. reached the identical diagnosis from the quantization side, naming the problem “helping attention heads do nothing” and reporting that suppressing the outliers is what makes per-tensor INT8 activation quantization viable. Reported perplexity gains from the off-by-one change alone are small and setup-dependent, so the case for it rests on outlier suppression and long-context robustness, not on loss curves.

Property	Ordinary softmax	Softmax-1	Learned sink logit	Register tokens
How a head does nothing	Dumps mass on a learned sink token and shrinks its value contribution	Pushes every logit below zero and abstains	Same, with the abstention threshold trained per head	Attends to a dedicated token reserved for scratch space
Weights sum to	Exactly 1	$S/(1+S)$ , strictly below 1	$S/(e^{s}+S)$ , strictly below 1	Exactly 1, including the register slots
Added cost	None	One constant in the denominator, zero parameters	One scalar per head	Extra sequence positions, so extra KV and quadratic prefill
Massive activations	Emerge reliably; block per-tensor INT8 activation quantization	Strongly reduced when trained from scratch	Reduced, and the head can tune how quiet it stays	Confined to registers instead of random patches, not removed
Sliding window safety	Fragile: evicting the sink token collapses perplexity	Safe: the null option is architectural, nothing to pin	Safe, and adapts to the window length during training	Safe only if the registers are never evicted
Retrofit to a trained model	Baseline	Needs continued pretraining; a drop-in swap breaks calibration	Same, though $s$ can be initialized from observed sink mass	Requires retraining or at least adapter tuning of the new tokens

August 19, 2026

DL0168 NoPE: Causal Positional Bias Without Encodings

What is NoPE (No Positional Embeddings), and how does a causal mask alone induce implicit positional bias in decoder-only Transformers without explicit position encodings?

Answer

NoPE is a decoder-only Transformer trained with no position information injected anywhere: no learned or sinusoidal absolute embedding added to the token embeddings, no RoPE rotation inside attention, and no ALiBi distance bias on the logits. The only order-dependent structure left in the network is the causal mask, and that turns out to be enough. A bidirectional self-attention layer without positional encodings is permutation equivariant, so it can only represent a bag of tokens, but causal masking destroys that symmetry because query $i$ attends over exactly $i+1$ keys. The size of each token’s receptive field is therefore itself a strictly monotone function of absolute position, and any head with near-uniform logits converts that count into a readable $1/(i+1)$ signal in the residual stream. Haviv et al. showed that absolute position can be linearly probed out of NoPE hidden states with high accuracy while perplexity stays close to models with explicit encodings, and Kazemnejad et al. proved constructively that a NoPE decoder can represent both absolute and relative position, and that it length-generalizes at least as well as RoPE or ALiBi on small-scale algorithmic tasks.

(1) Definition: NoPE removes every explicit position term and keeps only the lower-triangular attention mask, so position is an emergent property rather than an injected feature.
(2) Broken Permutation Equivariance: without a mask, $f(PX) = P f(X)$ for any permutation $P$ , but the causal mask satisfies $P^\top M P \neq M$ , which is exactly the asymmetry the model exploits.
(3) Counting Is The Mechanism: uniform attention over an $(i+1)$ -token window puts $1/(i+1)$ mass on each visible key, so the output norm encodes absolute position.
(4) Absolute First, Relative Later: layer 1 materializes an absolute code, usually anchored on a dominant BOS sink, and deeper layers subtract two codes to obtain the relative offset $i-j$ .
(5) Resolution Decays Quadratically: the gap between adjacent positions is $O(i^{-2})$ , so the counting code loses discriminative power long before the context window ends.
(6) Implicit Recency Bias: trained NoPE heads develop distance-decaying attention that resembles a learned relative encoding, but with no principled extrapolation knob, which is why long-context NoPE needs attention-temperature scaling or hybrid layers.

The mechanism is easiest to read as counting. Row $i$ of the causal attention matrix has exactly $i+1$ unmasked entries, so a head whose logits are roughly constant across its window spreads $1/(i+1)$ of the probability mass onto each visible key. If the value vectors are dominated by one distinguished token, in practice the BOS token that every position can see, then the output norm at position $i$ is proportional to $1/(i+1)$ : a strictly decreasing, invertible function of absolute position that the next layer can consume as a positional feature. Kazemnejad et al. turn this observation into a theorem, and once an absolute code sits in the residual stream a later layer’s $q_i^\top k_j$ term can compute $i-j$ , recovering relative position too. Real NoPE models show the fingerprints of this construction: very heavy mass on the first token, first-layer heads with near-flat windows, and probes that decode absolute position from early activations.

Three panels: a ten by ten causal attention mask heatmap whose row i has i+1 shaded cells with the visible-key count printed beside each row, the same size bidirectional mask where every row sees all ten keys, and a line chart showing that uniform causal attention yields an output norm of one over i plus one that decreases strictly with position while bidirectional uniform attention yields a flat one over n line carrying no positional information

Figure 1: The causal mask is the position signal. Counting unmasked entries per row gives $1{,}2{,}\ldots{,}n$ under causal masking but a constant $n$ under bidirectional attention, so the same uniform head produces a strictly decreasing code in one case and a flat, uninformative constant in the other.

Mathematical Formulation:
$\alpha_{ij} = \frac{\exp(q_i^\top k_j)}{\sum_{m=0}^{i} \exp(q_i^\top k_m)}$
$\alpha_{ij} = 0 \ \text{for}\ j > i$
$q_i^\top k_j = c \Rightarrow \alpha_{ij} = \tfrac{1}{i+1}$
$o_i = \sum_{j=0}^{i} \alpha_{ij} v_j = \frac{v_0}{i+1}$
$i = \|o_i\|^{-1} - 1$
$f_{\mathrm{bi}}(PX) = P f_{\mathrm{bi}}(X)$
$P^\top M P \neq M$

Where:

$o_i \in \mathbb{R}^{d}$ is the attention output at query position $i$ , and $\alpha_{ij}$ the attention weight it places on key position $j$ .
$q_i, k_j, v_j$ are the query, key, and value vectors, containing no positional term at all under NoPE.
$j, m \in \{0,\ldots,i\}$ index only the unmasked keys; the second line is the causal mask, which is the sole source of order sensitivity.
$c$ is a constant logit, the degenerate case that makes attention uniform over the window; the third and fourth lines assume $v_j = 0$ for $j \geq 1$ and a nonzero anchor value $v_0$ at BOS.
The fifth line inverts the code with $\|v_0\| = 1$ , showing that absolute position is exactly recoverable from a single layer’s output norm.
$P$ is a permutation matrix, $X$ the input sequence, $f_{\mathrm{bi}}$ an unmasked layer, and $M$ the lower-triangular mask; the last two lines state why encoders need explicit position encodings and causal decoders do not.

The same construction explains why NoPE is fragile at long context. Adjacent positions are separated by $\Delta_i = 1/((i+1)(i+2))$ , a relative spacing of only $1/(i+2)$ , so the code that cleanly distinguishes position $i$ from position $i+1$ must resolve relative differences near $8\times10^{-3}$ around token 126 and near $10^{-3}$ around token 1022. Under bf16 activations those differences sit at the edge of machine epsilon, and the second pressure is entropic: uniform attention over a growing window has entropy $H_i \leq \ln(i+1)$ , so a head calibrated to be discriminative at 2K tokens is comparatively diffuse at 32K. Real models do not use the naive counting head alone, but both effects push in the same direction, which is why raw NoPE degrades past its training length while RoPE offers explicit rescaling recipes such as NTK interpolation and YaRN.

Two panels: a log-log plot of the relative gap one over i plus two between adjacent positional codes falling below bf16 machine epsilon around position 126 and below fp16 machine epsilon around position 1022, and a semi-log plot of the maximum attention entropy log of i plus one growing from about 7.6 nats at 2048 tokens to about 10.4 nats at 32768 tokens

Figure 2: Two length pressures on an implicit code. The relative spacing between neighbouring positions shrinks like $1/(i+2)$ , crossing bf16 machine epsilon in the low hundreds of tokens, while the entropy ceiling grows like $\ln(i+1)$ , so the same attention temperature that is sharp at 2K is diffuse at 32K.

Property	NoPE	RoPE	ALiBi
Source of position	Causal mask only; window size acts as an implicit counter	Rotation of queries and keys by angle proportional to index	Fixed per-head linear penalty on the distance i minus j
Works without a mask	No; a bidirectional layer becomes permutation equivariant	Yes; encoders use it directly	Yes, with a symmetric distance bias
Extension knob	Attention temperature or entropy scaling; no standard recipe	Base frequency rescaling, NTK interpolation, YaRN	Slope schedule; extrapolates but truncates effective range
Extra cost	None; consumes capacity and layer-1 heads instead	Two elementwise ops per head; cached keys are pre-rotated	One additive bias term per attention score
Dominant failure mode	Code resolution and softmax sharpness collapse past training length	Unseen rotation phases out of distribution beyond training length	Strong recency bias suppresses genuine long-range retrieval

August 19, 2026

DL0132 LLM Serving Latency: TTFT and ITL

How do you measure and reduce Time to First Token (TTFT), Inter-Token Latency (ITL), and latency variance in a production LLM API such as a streaming chat endpoint served with vLLM?

Answer

TTFT is the wall-clock time from request arrival to the first streamed token and ITL is the gap between consecutive tokens after that, and they are separate metrics because they are produced by two phases with two different hardware bottlenecks. Prefill processes the whole prompt in one or a few forward passes and is compute-bound, so $T_{ttft}$ grows roughly linearly with prompt length; decode generates one token per forward pass and is memory-bandwidth-bound, so $T_{itl}$ is floored by the time to stream the weights and KV cache out of HBM. Measurement must happen at the streaming boundary and be reported as a distribution, namely p50, p95 and p99 merged from histograms and bucketed by input and output length, because a p99 computed over a mix of 200-token and 100k-token prompts is meaningless. Reduction is metric-specific: prefix caching and chunked prefill cut TTFT, FP8 quantization and speculative decoding cut ITL, and the variance work is almost entirely about removing things that can stall the decode loop, namely unchunked prefills, KV cache preemption, and queueing above roughly 70% utilization. The metric that actually decides capacity is goodput, the request rate you can sustain while both the TTFT and ITL SLOs hold at p99, not raw tokens per second.

(1) Two Phases, Two Bottlenecks: prefill is compute-bound and sets TTFT, decode is bandwidth-bound and sets ITL, so a change that helps one often does nothing for the other.
(2) Measure Distributions, Not Means: record one TTFT scalar plus an ITL vector of length $N_{out} - 1$ per request, and merge t-digest or HDR histograms across replicas instead of averaging per-replica p99s.
(3) TTFT Is Mostly Queueing Under Load: split server-side spans into $T_{queue}$ and $T_{prefill}$ , because past about 70% utilization the queue term dominates and no kernel optimization will help.
(4) The ITL Tail Is Interference, Not Steady State: p50 ITL reflects the decode step, while p99 ITL reflects the longest single stall, typically a colliding prefill or a preempted request being recomputed.
(5) Knobs Map To Metrics: prefix caching, chunked prefill and admission control for TTFT; quantization, GQA and speculative decoding for ITL; batch caps and prefill-decode disaggregation for variance.
(6) Optimize Goodput: report the sustained rate under both SLOs simultaneously, since raising batch size always trades TTFT and ITL tails for throughput.

Two timeline rows for one decoding request. In the baseline row, queue and prefill produce the first token, then four decode steps are interrupted by one wide unchunked prefill block that stalls the decode loop, creating a large inter-token gap. In the chunked-prefill row, the same prefill is split into six narrow chunks interleaved with decode steps, so every inter-token gap stays small.

Figure 1: TTFT is queue time plus prefill time; ITL is whatever happens between two decode steps of the same request. A single unchunked prefill scheduled into the shared loop stalls every decoding stream at once, which is the dominant source of p99 ITL. Splitting it into token-budgeted chunks bounds the worst stall to one chunk, at a small throughput cost from re-reading weights more often.

Instrument at the streaming boundary rather than inside the model: timestamp request arrival, the first SSE data event, and every subsequent event, then derive TTFT and the ITL vector client-side while the server emits its own queue, prefill and decode spans for the same request ID. The two views disagree in informative ways, because proxy buffering (an nginx deployment with proxy_buffering on) and a slow client event loop both inflate measured ITL with zero GPU involvement. For load testing, drive the endpoint with realistic input and output length distributions, since synthetic fixed-length prompts hide exactly the length-mixing effects that create the real tail; the vLLM benchmark_serving script and NVIDIA GenAI-Perf both report TTFT and ITL percentiles at a fixed request rate, which is the right shape of experiment. The variance checklist beyond scheduling is mundane but produces most of the surprises in practice: KV cache preemption and recompute when cache headroom runs out, ITL growing with batch size as decode drifts toward compute-bound, CUDA graph misses on uncaptured batch shapes, tensor-parallel all-reduce jitter, autoscaler cold starts that load tens of gigabytes of weights, and speculative decoding turning ITL bimodal because accepted-draft steps emit several tokens at once while rejected steps emit one.

Mathematical Formulation:
$T_{e2e} = T_{ttft} + (N_{out} - 1)\, \overline{T_{itl}}$
$T_{ttft} = T_{queue} + T_{prefill}$
$T_{prefill} \approx \frac{2 N_{in} P}{C_{eff}}$
$T_{itl} \geq \frac{B_{w} + B_{kv}}{BW}$
$T_{itl}^{spec} \approx \frac{T_{step}}{\mathbb{E}[n_{acc}]}$
$G = \lambda \cdot \Pr(T_{ttft} \leq S_{t},\ T_{itl} \leq S_{i})$

Where:

$T_{e2e}$ is end-to-end request latency, $T_{ttft}$ the time to first token, and $\overline{T_{itl}}$ the mean inter-token latency (its reciprocal is the perceived tokens per second of one stream).
$T_{queue}$ is scheduler wait time and $T_{prefill}$ the prompt forward pass; separating them is the single most useful piece of instrumentation, because they need opposite fixes.
$N_{in}$ and $N_{out}$ are prompt and completion token counts, $P$ the parameter count, and $C_{eff}$ the achieved (not peak) FLOP/s; the factor 2 counts one multiply and one add per parameter per token.
$B_{w}$ and $B_{kv}$ are the weight and KV cache bytes read per decode step and $BW$ is per-device HBM bandwidth, so this inequality is the bandwidth floor on ITL that no batching removes.
$T_{step}$ is one verification forward pass and $\mathbb{E}[n_{acc}] \geq 1$ the expected accepted tokens per step under speculative decoding, which divides effective ITL but widens its distribution.
$\lambda$ is offered request rate, $S_{t}$ and $S_{i}$ the TTFT and ITL SLO thresholds, and $G$ the resulting goodput, the only capacity number worth putting in a dashboard.

Plugging in real hardware makes the two bottlenecks concrete. A 70B model in BF16 needs roughly 140 GB of weights, so on eight H100s with tensor parallelism each device reads about 17.5 GB per decode step against roughly 3.35 TB/s of HBM, giving a floor near 5 ms per token before all-reduce and KV reads, which is why a single stream tops out around 100 to 150 tokens per second in practice. The same eight GPUs prefill a 2048-token prompt in about $2 \cdot 2048 \cdot 70$ B FLOPs, roughly 287 TFLOPs, which at an achieved 3.2 PFLOP/s is about 90 ms; a 32k-token prompt is 16 times that, which is exactly why long-context traffic destroys a shared TTFT SLO unless prefixes are cached or prefill is disaggregated. Note that FP8 weights halve $B_w$ and therefore the ITL floor, but barely move TTFT, while prefix caching can cut TTFT by an order of magnitude on a 1500-token shared system prompt and does nothing for ITL.

Two line charts against offered request rate. Left panel shows p50 and p99 TTFT in milliseconds rising sharply near capacity with a dotted 1500 ms SLO line and a vertical marker at the crossing rate. Right panel shows p50 and p99 inter-token latency in milliseconds with a dotted 80 ms SLO line crossed at a lower request rate, marking the binding constraint.

Figure 2: Both p50 curves stay flat while the p99 curves bend upward, because queueing and interference hit the tail first. Here the ITL SLO is crossed at a lower rate than the TTFT SLO, so goodput is ITL-bound and the correct response is to cap batch size or disaggregate, not to add prefill compute.

Technique	Metric it moves	What it costs
Prefix / KV cache reuse	TTFT p50 and p99 on shared system prompts, often 5x to 10x	HBM held by cached blocks, eviction policy tuning, cross-tenant isolation concerns
Chunked prefill	ITL p99 (bounds the worst decode stall to one chunk)	Slightly higher TTFT and a few percent lower throughput from repeated weight reads
Prefill-decode disaggregation	Both tails, by giving each phase its own pool and SLO	A KV cache transfer per request, duplicated weights, needs fast interconnect
Speculative decoding	ITL p50, roughly divided by accepted tokens per step	Bimodal ITL, wasted FLOPs on rejection, gains shrink at large batch sizes
FP8 or INT4 weight quantization	ITL floor, by halving or quartering bytes read per step	Accuracy regression that must be evaluated per task, calibration pipeline
Batch cap plus admission control	Latency variance, by refusing to accept work that will miss the SLO	Lower throughput and visible queueing or 429s, so it needs a priority policy

August 17, 2026

DL0131 DoRA: Weight-Decomposed LoRA

Explain DoRA (Weight-Decomposed Low-Rank Adaptation). How does separating weight magnitude from direction improve LoRA fine-tuning stability?

Answer

DoRA reparameterizes every pretrained weight matrix into a magnitude vector and a directional matrix, and then lets LoRA update only the direction. Concretely it writes $W_0 = m \frac{V}{\|V\|_c}$ , where $\|\cdot\|_c$ is the column-wise L2 norm, the trainable vector $m \in \mathbb{R}^{1 \times k}$ is initialized to $\|W_0\|_c$ , and the low-rank product $BA$ is added to the direction before renormalization. The motivation comes from a weight decomposition analysis: when full fine-tuning updates a layer, the magnitude change $\Delta M$ and the direction change $\Delta D$ are negatively correlated (a large rotation with only slight rescaling, or the reverse), while plain LoRA shows a positive, roughly proportional relation, so it can essentially only push both quantities in the same direction. Giving magnitude its own $k$ parameters removes that coupling, and the renormalization makes the gradient reaching the directional component orthogonal to the current direction, which is the classic weight-normalization conditioning effect: the low-rank branch can rotate the weight without simultaneously inflating its scale. In practice this buys a learning pattern that mirrors full fine-tuning, noticeably better accuracy at low rank (the DoRA paper reports about a 3 to 4 point average gain over LoRA on the eight commonsense-reasoning benchmarks with LLaMA-7B, and matching LoRA at half the rank in its ablations), and zero inference overhead, because $m V' / \|V'\|_c$ collapses back into one merged matrix.

(1) Two Degrees Of Freedom Instead Of One: LoRA has a single additive update $W_0 + BA$ ; DoRA splits the same layer into a scalar-per-column scale and a normalized direction, so scale and rotation are optimized by separate tensors.
(2) Exact Identity At Initialization: with $B = 0$ and $m_0 = \|W_0\|_c$ the adapted layer reproduces $W_0$ bit-for-bit, so training starts from the pretrained function with no loss spike.
(3) Gradient Is Projected, Not Just Scaled: differentiating through the column norm inserts the projector $I - V'V'^{\top} / \|V'\|_c^2$ , which removes the radial component of the directional gradient and bounds its effective step.
(4) Learning Pattern Matches Full Fine-Tuning: measured over training, DoRA reproduces the negative $\Delta M$ versus $\Delta D$ slope of full fine-tuning, whereas LoRA’s slope is positive.
(5) Cheapest Where LoRA Hurts Most: the extra magnitude vector costs $k$ parameters per layer, roughly 6% on top of a rank-8 adapter for a $4096 \times 4096$ projection, and pays off most at $r = 4$ to $r = 8$ .
(6) Training Overhead, Not Serving Overhead: the column norm of $W_0 + BA$ must be recomputed each step, which raises training memory; detaching that norm from the backward graph recovers roughly a quarter of the overhead with negligible accuracy change, and inference is unaffected after merging.

Data-flow diagram of a DoRA layer: the frozen pretrained weight is decomposed into a trainable magnitude vector and a frozen normalized direction, the low-rank product BA is added to the direction, the sum is renormalized column-wise, then scaled by the magnitude vector to produce the merged adapted weight

Figure 1: A DoRA layer. Only the magnitude vector $m$ and the LoRA factors $A, B$ receive gradients; the frozen direction is perturbed by $BA$ , renormalized per column, and rescaled by $m$ . Because the final expression is a single matrix, the adapter still merges away at inference.

The stability argument is easier to see from what LoRA cannot express. In LoRA the only knob is the additive term, so any attempt to rescale a column also rotates it, and any attempt to rotate it also changes its norm; the two effects are welded together by a single rank- $r$ product, which is why its measured $(\Delta D, \Delta M)$ points fall on a line with positive slope. Full fine-tuning has $dk$ free parameters and shows the opposite pattern, with subtle magnitude changes accompanying substantial directional change. DoRA recovers that behavior with $k$ extra parameters, and the practical consequences are the ones interviewers care about: lower sensitivity to the learning rate and to the LoRA scaling factor $\alpha / r$ , because a badly scaled low-rank branch is renormalized away instead of blowing up the weight norm, and much flatter accuracy-versus-rank curves, because a rank-4 direction plus a full-rank magnitude still spans useful updates.

Scatter plot of magnitude change versus direction change per layer for full fine-tuning, LoRA, and DoRA with fitted regression lines; full fine-tuning and DoRA have negative slopes while LoRA has a positive slope

Figure 2: Measured magnitude change $\Delta M$ against direction change $\Delta D$ per adapted module. Full fine-tuning and DoRA trace a negative slope, meaning scale and rotation are traded against each other, while LoRA‘s positive slope shows the two are coupled by its single additive term.

Mathematical Formulation:
$W_0 = m_0 \frac{V_0}{\|V_0\|_c}$
$m_0 = \|W_0\|_c$
$W' = m \frac{W_0 + BA}{\|W_0 + BA\|_c}$
$\nabla_{V'}\mathcal{L} = \frac{m}{C}\left(I - \frac{V'V'^{\top}}{C^2}\right)\nabla_{W'}\mathcal{L}$
$\Delta M_t = \frac{1}{k}\sum_{j} |m_t^{j} - m_0^{j}|$
$\Delta D_t = \frac{1}{k}\sum_{j} \left(1 - \cos(V_t^{j}, W_0^{j})\right)$

Where:

$W'$ is the adapted weight actually used in the forward pass, and $W_0 \in \mathbb{R}^{d \times k}$ is the frozen pretrained weight it starts from.
$m \in \mathbb{R}^{1 \times k}$ is the trainable magnitude vector and $V' = W_0 + BA$ the unnormalized direction, with $V_0 = W_0$ and $m_0$ the initialization that makes $W' = W_0$ at step 0.
$B \in \mathbb{R}^{d \times r}$ (initialized to zero) and $A \in \mathbb{R}^{r \times k}$ are the LoRA factors with rank $r \ll \min(d,k)$ .
$j \in \{1,\ldots,k\}$ indexes output columns, $t$ indexes training checkpoints, and $\|\cdot\|_c$ takes the L2 norm of each column independently, returning a $1 \times k$ row vector.
$C = \|V'\|_c$ is that column norm, and $I - V'V'^{\top}/C^2$ is the orthogonal projector that strips the radial part of the gradient, so the directional update is norm-preserving to first order.
$\Delta M_t$ and $\Delta D_t$ are the magnitude and direction change metrics whose correlation is plotted in Figure 2, with $\Delta D_t \in [0,2]$ because it is one minus a cosine similarity.

Property	DoRA	LoRA	Full fine-tuning
Trainable tensors per layer	A, B, and the magnitude vector m	A and B only	The entire weight matrix
Parameters added	r(d + k) + k, about 6% more than LoRA at r = 8, k = 4096	r(d + k)	dk, plus optimizer state for all of it
Magnitude vs direction coupling	Decoupled; negative correlation like full fine-tuning	Coupled through one additive term; positive correlation	Fully free, the reference behavior
Behavior at very low rank	Degrades slowly; competitive at r = 4 to 8	Degrades quickly below r = 8 on harder tasks	Not applicable
Training and serving cost	Extra column-norm recompute each step; merges cleanly, no inference cost	Cheapest to train; merges cleanly, no inference cost	Full optimizer and gradient memory; one checkpoint per task

August 17, 2026

DL0130 MoE Router Collapse and Auxiliary Loss

What is Router Collapse in Sparse Mixture-of-Experts (MoE) LLMs? Derive the auxiliary load-balancing loss and Router Z-loss used to stabilize MoE training.

Answer

Router collapse is the failure mode in which the learned router of a sparse MoE layer stops using most of its experts and concentrates almost all token assignments on a small subset, so a layer you paid $N$ experts’ worth of memory for behaves like a much smaller model. The cause is a positive feedback loop: an expert that happens to win slightly more tokens early in training receives more gradient signal, becomes genuinely better on those tokens, so the router raises its logit further, while starved experts never see enough tokens to become useful. Under token-choice top-k the visible symptoms are a skewed load histogram, a large dropped-token fraction once the hot experts hit their capacity buffer, and a validation loss that tracks a dense model of the active-parameter size rather than the total-parameter size. Production training stacks suppress it with two extra loss terms rather than with a new routing algorithm: an auxiliary load-balancing loss that pushes router probability mass away from overloaded experts, and a router z-loss that penalizes the squared log-sum-exp of the router logits so the logits cannot grow without bound and destabilize the softmax in low precision. Both are differentiable surrogates for quantities that are not: the actual assignment counts are piecewise constant, and the actual numerical blow-up is a hardware property, so each loss attacks a proxy the gradient can reach.

(1) Collapse Is Self-Reinforcing: a small early advantage in router logits compounds through the gradient the winning expert receives, which is why collapse typically happens in the first few thousand steps and is nearly irreversible afterwards.
(2) The Assignment Is Non-Differentiable: the load fraction $f_i$ comes from a top-k selection and has zero gradient almost everywhere, so the balance loss must pair it with the differentiable mean gate probability $P_i$ .
(3) Balance Loss Is A Normalized Dot Product: $N \sum_i f_i P_i$ equals 1 under a uniform assignment and grows toward $N$ under total collapse, giving a scale-free objective independent of expert count.
(4) Its Gradient Is Load-Proportional: the derivative with respect to $P_i$ is exactly $\alpha N f_i$ , so probability mass is pushed down in proportion to how overloaded each expert already is.
(5) Z-Loss Bounds The Logits: penalizing the squared log-partition function caps $\max_j h_j$ , which matters because the relative round-off in $e^{h}$ grows linearly with $|h|$ in bfloat16 and can flip top-k selection.
(6) Coefficients Are A Quality Trade: $\alpha \approx 10^{-2}$ and $\beta \approx 10^{-3}$ are the common settings, because a large $\alpha$ buys perfect balance by actively fighting the language-modeling objective.

Cycle diagram of router collapse: router logits for one expert edge above the others, top-k dispatches more tokens to it, that expert receives most of the expert gradient, the starved experts stay undertrained and score lower, which feeds back into the logit gap; two green intervention boxes on the right show the router z-loss acting on the logits node and the auxiliary balance loss acting on the dispatch node

Figure 1: Router collapse is a closed loop, not a single bad step: the logit gap, the dispatch skew, and the gradient imbalance each amplify the next. The two stabilizers cut the loop at different points, with the z-loss constraining the logit magnitudes and the balance loss constraining the dispatch distribution.

The derivation of the balance loss starts from what you actually want to penalize, namely the variance of the per-expert token counts, and then asks which part of that quantity carries a gradient. The counts themselves come from a top-k over the router logits, so they are a step function of the parameters and give nothing to backpropagation. The fix used by GShard and simplified by Switch Transformers is to pair the non-differentiable load vector $f$ with the differentiable importance vector $P$ , the mean softmax probability per expert, and minimize their inner product. Treating $f$ as a constant, the objective is linear in $P$ with coefficient $N f_i$ , so each step lowers the router’s probability for exactly the experts that were overloaded in the current batch, and because $f$ is recomputed every step this becomes a self-correcting controller whose fixed point is the uniform assignment. The factor $N$ is a normalization choice: it makes the minimum value 1 regardless of how many experts you have, so the same $\alpha$ transfers from an 8-expert layer to a 256-expert layer. The z-loss has an entirely different motivation: it is a numerical guard, derived from the observation that a logit stored with relative precision $\epsilon$ produces an absolute error of about $|h|\epsilon$ , and exponentiation converts that absolute error into a relative error of the same size, so large logits make the softmax and therefore the selected expert set unreliable.

Mathematical Formulation:
$h_t = W_r x_t, \quad p_t = \mathrm{softmax}(h_t)$
$f_i = \frac{1}{T}\sum_{t=1}^{T} \mathbb{1}[i \in \mathcal{T}_t]$
$P_i = \frac{1}{T}\sum_{t=1}^{T} p_{t,i}$
$\mathcal{L}_{\mathrm{bal}} = N \sum_{i=1}^{N} f_i P_i$
$\frac{\partial \mathcal{L}_{\mathrm{bal}}}{\partial P_i} = N f_i$
$1 \leq \mathcal{L}_{\mathrm{bal}} \leq N$
$\mathcal{L}_{z} = \frac{1}{T}\sum_{t=1}^{T}\left(\log \sum_{j=1}^{N} e^{h_{t,j}}\right)^{2}$
$\max_j h_{t,j} \leq \mathrm{lse}(h_t) \leq \max_j h_{t,j} + \log N$
$\mathcal{L} = \mathcal{L}_{\mathrm{LM}} + \alpha \mathcal{L}_{\mathrm{bal}} + \beta \mathcal{L}_{z}$

Where:

$x_t \in \mathbb{R}^{d}$ is the hidden state of token $t$ , $W_r \in \mathbb{R}^{N \times d}$ the router matrix, and $h_t$ the router logits before any selection.
$T$ is the number of tokens the statistics are aggregated over (micro-batch, device batch, or global batch), $N$ the expert count, and $\mathcal{T}_t$ the set of $k$ experts selected for token $t$ .
$f_i$ is the load, the fraction of tokens dispatched to expert $i$ ; it is piecewise constant in the parameters and therefore contributes no gradient.
$P_i$ is the importance, the mean router probability assigned to expert $i$ ; it is smooth, so all of the balance-loss gradient flows through it and into $W_r$ .
$\mathcal{L}_{\mathrm{bal}} = 1$ exactly when $f_i = P_i = 1/N$ for all $i$ , and approaches $N$ when one expert takes every token, so the value is directly readable as an imbalance factor.
$\mathrm{lse}(h_t) = \log\sum_j e^{h_{t,j}}$ is the log-partition function; squaring it penalizes large logits in either direction and, by the sandwich bound, keeps $\max_j h_{t,j}$ within $\log N$ of a small target.
$\alpha$ and $\beta$ are the balance and z-loss coefficients, commonly $\alpha = 10^{-2}$ and $\beta = 10^{-3}$ ; $\mathcal{L}_{\mathrm{LM}}$ is the ordinary next-token cross-entropy.

Two-panel chart: left panel plots the maximum expert load fraction against training step for an eight-expert layer with balance-loss coefficients zero, one thousandth, and one hundredth, showing full collapse toward one expert without the loss and a curve close to the uniform 0.125 line with alpha one hundredth; right panel plots relative round-off error in exp of a router logit against logit magnitude on a log scale for bfloat16 and float32 storage, with bfloat16 crossing one percent error at a logit magnitude near 2.6

Figure 2: The two losses guard different quantities. Left: without a balance term the maximum load fraction runs from the uniform $1/N = 0.125$ to near 1.0 within a few thousand steps, while $\alpha = 10^{-2}$ holds it close to uniform. Right: the relative round-off in $e^{h}$ grows linearly with $|h|$ , which is why bfloat16 routers become unreliable at moderate logits and the z-loss keeps $\mathrm{lse}(h)$ near $O(1)$ .

Property	Auxiliary balance loss	Router z-loss	Bias-based loss-free balancing
Quantity penalized	The dot product of load and importance, scaled by the expert count	The squared log-sum-exp of the router logits, averaged over tokens	Nothing; a per-expert bias is added to the selection logits only
Failure it prevents	Expert collapse, wasted parameters, and dropped tokens at the capacity buffer	Logit blow-up, low-precision softmax round-off, and loss spikes	Expert collapse, without perturbing the gate weights used in the output
Typical setting	Coefficient 1e-2, aggregated per device batch or per global batch	Coefficient 1e-3, with the router itself computed in float32	Bias update rate around 1e-3, driven by observed per-expert load error
Effect on the LM objective	Adds an interference gradient that trades some quality for balance	Mild regularizer, usually neutral or slightly positive for quality	No interference term at all, since the bias is not part of the output gate
Where it is used	GShard, Switch, Mixtral, OLMoE	ST-MoE and most later open MoE training stacks, including OLMoE	DeepSeek-V3 and its loss-free-balancing follow-ups

August 17, 2026

DL0129 SwiGLU, GELU, and ReLU FFNs

Compare SwiGLU, GeLU, and ReLU activation functions in the Feed-Forward Network (FFN) blocks of LLMs, as used in models like Llama 3 and GPT-3.

Answer

All three sit in the same place, namely the two-layer FFN that follows attention in every transformer block, but they differ in smoothness and in whether the nonlinearity is elementwise or multiplicative. ReLU is $\max(0,x)$ : cheapest, produces exact zeros (roughly 90-95% of FFN units are zero on a typical token), but has a kink at the origin and zero gradient for negative inputs, which gives dead units. GELU is $x\Phi(x)$ : smooth, non-monotonic, keeps a small negative response, and was the default through BERT, GPT-2, and GPT-3. SwiGLU is not an elementwise function at all but a gated layer: the up-projection is split into a gate branch passed through SiLU ( $x\sigma(x)$ ) and a value branch left linear, and the two are multiplied elementwise. That third weight matrix is why practitioners shrink the hidden width by $2/3$ to stay parameter- and FLOP-neutral, and the payoff reported by Shazeer’s GLU-variants ablation is a consistent but modest gain, about 0.05 nats of pretraining log-perplexity on T5’s span-corruption objective at matched budget, which is enough that PaLM, Llama, Mistral, and Qwen2 all adopted it. The cost is one extra matmul, one extra saved activation for the backward pass, loss of exploitable sparsity, and a wider activation range that makes low-bit quantization harder.

(1) Two Matrices Versus Three: a dense FFN is $W_{up}$ then $W_{down}$ ; a gated FFN adds $W_g$ , so parameters go from $2 d\, d_{ff}$ to $3 d\, d_{ff}$ at the same width.
(2) The 2/3 Width Rule: to keep the block budget fixed you set $d_{ff} = \frac{8}{3}d$ instead of $4d$ , which is exactly why Llama 2 7B uses 11008 with $d = 4096$ .
(3) Smoothness Changes Optimization: ReLU’s derivative is a step with no gradient for negative inputs, while GELU and SiLU are $C^\infty$ and pass a small negative signal, so units recover instead of dying.
(4) Gating Is A Multiplicative Interaction: the output is a product of two learned projections of the same input, giving a data-dependent, per-channel scaling that no elementwise function can express.
(5) Sparsity Is A ReLU-Only Asset: exact zeros let inference systems skip entire rows of $W_{down}$ ; GELU and SwiGLU produce values that are small but almost never exactly zero.
(6) SwiGLU Costs Memory Traffic And Range: three matmuls plus an extra stored tensor per layer, and the product of two unbounded projections widens the dynamic range that quantization must cover.

Two-panel line chart: left panel plots ReLU, GELU and SiLU over the range minus four to four, showing ReLU flat at zero for negative inputs while GELU and SiLU dip slightly below zero; right panel plots their derivatives, showing ReLU's step from zero to one at the origin versus the smooth GELU and SiLU derivatives that briefly exceed one and go slightly negative

Figure 1: Shape (left) and derivative (right). ReLU’s derivative is a step function with a kink at the origin and no gradient for negative inputs; GELU and SiLU are smooth, non-monotonic, and dip below zero (SiLU’s minimum is $-0.278$ at $x = -1.28$ ), so a unit pushed into the negative region still receives gradient and can come back.

The historical progression is ReLU → GELU → gated GLU variants, and it is worth being honest that only the first step has a clean story. Replacing ReLU with GELU removes the dead-unit problem and gives a slightly better-conditioned loss surface. Replacing an elementwise nonlinearity with a gate is different in kind: the value branch stays linear, and the gate branch modulates it, so the layer computes a second-order interaction of the input with itself. Shazeer’s paper explicitly declines to explain why this helps, attributing the gain to luck rather than theory, and later analyses mostly note that gating gives the block a cheap multiplicative pathway similar in spirit to the gates in an LSTM. What made SwiGLU the production default is that the gain survives at scale, needs no extra hyperparameters, and is free once you apply the $2/3$ width rule. Llama 3 8B keeps the gated form but widens back out to $d_{ff} = 14336$ with $d = 4096$ , a deliberate choice to spend more parameters in the FFN rather than a parameter-neutral swap.

Mathematical Formulation:
$\mathrm{FFN}_{\mathrm{ReLU}}(x) = \max(0, xW_{up})\,W_{down}$
$\mathrm{GELU}(x) = x\,\Phi(x)$
$\mathrm{SiLU}(x) = x\,\sigma(x)$
$\mathrm{SwiGLU}(x) = \mathrm{SiLU}(xW_g) \odot (xW_u)$
$\mathrm{FFN}_{\mathrm{SwiGLU}}(x) = \mathrm{SwiGLU}(x)\,W_{down}$
$d_{ff} = \frac{2}{3}\cdot 4d = \frac{8}{3}d$

Where:

$x \in \mathbb{R}^{d}$ is the post-norm hidden state for one token, and the FFN output has the same dimension $d$ so it can be added back through the residual connection.
$W_{up}, W_g, W_u \in \mathbb{R}^{d \times d_{ff}}$ are up-projections and $W_{down} \in \mathbb{R}^{d_{ff} \times d}$ the down-projection; the dense block uses $W_{up}$ only, the gated block uses both $W_g$ (gate) and $W_u$ (value).
$\Phi$ is the standard normal CDF and $\sigma$ the logistic sigmoid; GELU is often shipped as the tanh approximation $0.5x(1+\tanh(\sqrt{2/\pi}(x+0.044715x^3)))$ , which is why two frameworks can disagree in the last few digits.
$\odot$ is the elementwise product over the $d_{ff}$ hidden channels, the operation that turns a gate into a multiplicative interaction rather than a pointwise map.
$d_{ff}$ is the FFN hidden width; the last line is the parameter-neutral setting, since $3 d \cdot \frac{8}{3}d = 8d^2 = 2 d \cdot 4d$ , and real implementations round the result to a hardware-friendly multiple such as 256.

Architecture diagram with two stacked left-to-right flows: the top flow is a dense FFN with hidden state, one up-projection, an elementwise ReLU or GELU, a down-projection and output, annotated with two d times d_ff parameters; the bottom flow is a gated SwiGLU FFN where the hidden state branches into a SiLU gate path and a linear value path, the two are combined by an elementwise product, then a down-projection and output, annotated with three d times d_ff parameters and the two-thirds width rule

Figure 2: The structural difference is a branch, not a curve. A dense FFN applies one elementwise function between two matrices; a gated FFN splits the up-projection into a SiLU gate path and a linear value path whose elementwise product feeds $W_{down}$ . The third matrix is what forces the $8d/3$ width if the block budget must stay fixed.

Property	ReLU	GELU	SwiGLU
Form	Elementwise, piecewise linear	Elementwise, smooth, non-monotonic	Gated layer: SiLU branch times linear branch
Matrices per FFN	2	2	3, so width drops to 8d/3 to stay budget-neutral
Gradient behaviour	0 or 1; dead units never recover	Smooth, briefly exceeds 1, small negative region	Smooth, plus a gradient path through the gate itself
Exact-zero sparsity	Yes, typically 90-95% of units, exploitable at inference	No; values are small but nonzero	No; near-zero mass only, so row skipping needs a threshold and loses accuracy
Quantization friendliness	Best: bounded below, no multiplicative amplification	Good	Hardest: the product of two unbounded projections widens outliers into W_down
Representative models	Original Transformer, T5 v1.0, OPT	BERT, GPT-2, GPT-3, ViT	PaLM, Llama 2 and Llama 3, Mistral, Qwen2 (Gemma uses the GELU-gated GeGLU)

August 17, 2026

DL0128 Subword Regularization

What is subword regularization, and why can it help machine translation quality?

Answer

Subword regularization is the practice of training a translation model on multiple stochastically sampled segmentations of the same sentence instead of the one deterministic segmentation a tokenizer normally emits. Any subword vocabulary is ambiguous: the word “smaller” can be a single piece, or small plus er, or sm plus all plus er, and a standard tokenizer always returns the same choice, so the model never sees the alternatives. Two mechanisms implement the sampling: unigram-LM sampling (Kudo 2018, shipped in SentencePiece), which draws a path from the segmentation lattice with probability proportional to $P(\mathbf{x})^{\alpha}$ , and BPE-dropout (Provilkov 2020), which keeps the ordinary BPE merge table but skips each merge with probability $p$ during application. The quality gain has three sources: it is free data augmentation over the same parallel corpus, it trains the embeddings of rare pieces that a deterministic segmenter almost never produces, and it makes the encoder robust to the segmentation shifts caused by typos, morphology, and domain change. Reported gains are roughly 1 BLEU on average, shrinking toward noise on large in-domain high-resource benchmarks and reaching 2 BLEU or more on low-resource pairs and out-of-domain test sets. Sampling is a training-time-only device: decoding uses the deterministic segmentation, or optionally averages scores over $k$ sampled segmentations of the source.

(1) Segmentation Is Ambiguous: a fixed vocabulary admits exponentially many segmentations per sentence, and deterministic tokenization silently commits to one of them in every epoch.
(2) Unigram-LM Sampling: SentencePiece samples from the $l$ -best lattice paths with a temperature $\alpha$ , typically $l = 64$ and $\alpha$ near 0.1 to 0.2, where small $\alpha$ is near-uniform and $\alpha \to \infty$ recovers the Viterbi segmentation.
(3) BPE-Dropout: drop each merge operation independently with $p = 0.1$ on both source and target; the vocabulary and merge table are unchanged, so it is a drop-in change to the tokenizer call.
(4) Approximate Marginalization: the model is pushed to score a sentence pair well under any plausible segmentation, which is closer to the true marginal $P(Y \mid X)$ than a single-path proxy.
(5) Robustness Beats In-Domain Accuracy: the largest wins appear on low-resource, morphologically rich, noisy, and out-of-domain data, exactly where the test-time segmentation drifts from the training-time one.
(6) Costs Are Real: average sequence length grows, tokenization moves into the data loader and cannot be cached, and convergence is slower in wall-clock terms for the same number of updates.

Segmentation lattice for the word smaller: boundary nodes under the characters, with three bracketed paths drawn above them, one spanning the whole word as a single piece, one splitting into small plus er, and one splitting into sm plus all plus er, each annotated with its unigram probability and the sampling weight at alpha equal to 0.2

Figure 1: One word, three legal paths through the segmentation lattice. Deterministic tokenization always returns the top path, so the -er morpheme boundary is never shown to the encoder; sampling at $\alpha = 0.2$ returns the top path only about half the time and spends the rest of the probability mass on the finer decompositions.

The reason the same trick helps two different segmenters is that both are inverting a latent variable the corpus never labels. A unigram language model gives an explicit probability to each path, so its lattice can be sampled cleanly and the temperature $\alpha$ gives a continuous dial from “uniform over the lattice” to “deterministic”. BPE has no probability model at all, only an ordered merge list, so BPE-dropout injects randomness procedurally instead: skipping a merge with probability $p$ produces a shorter piece and therefore a longer, more character-like segmentation, and stacking many independent skips over a long word yields a wide distribution of segmentations. Both settings need a knob that is gentle: too much randomness pushes segmentation toward characters, inflating sequence length, wasting attention on re-learning spelling, and destroying the frequency statistics that made the vocabulary useful in the first place.

Mathematical Formulation:
$P(\mathbf{x}) = \prod_{i=1}^{M} p(x_i)$
$P_{\alpha}(\mathbf{x} \mid X) \propto P(\mathbf{x})^{\alpha}$
$\mathcal{L}(\theta) = \mathbb{E}_{\mathbf{x}, \mathbf{y}}\left[\log P(\mathbf{y} \mid \mathbf{x}; \theta)\right]$
$P(Y \mid X) = \sum_{\mathbf{x}} P(\mathbf{x} \mid X)\, P(Y \mid \mathbf{x})$
$\hat{Y} = \arg\max_{Y} \frac{1}{k} \sum_{j=1}^{k} \log P(Y \mid \mathbf{x}_j)$

Where:

$X$ and $Y$ are the raw source and target sentences, while $\mathbf{x} = (x_1, \ldots, x_M)$ and $\mathbf{y}$ are particular segmentations of them into vocabulary pieces.
$p(x_i)$ is the unigram probability of piece $x_i$ , estimated by EM when the SentencePiece unigram vocabulary is built, and $M$ is the number of pieces on that path.
$\alpha$ is the smoothing exponent of the sampling distribution: $\alpha \to 0$ gives a near-uniform draw over the $l$ -best lattice paths and $\alpha \to \infty$ collapses to the single Viterbi path, so deterministic tokenization is the limiting case.
The expectation in $\mathcal{L}(\theta)$ is taken over $\mathbf{x} \sim P_{\alpha}(\cdot \mid X)$ and $\mathbf{y} \sim P_{\alpha}(\cdot \mid Y)$ , approximated by one fresh sample per example per epoch rather than by an inner loop.
$P(Y \mid X)$ is the true segmentation-marginalized translation probability that the training objective approximates; the sum over $\mathbf{x}$ is intractable, which is exactly why sampling is used.
$k$ is the number of source segmentations used in optional n-best decoding, where $k = 1$ with the Viterbi path is the standard, cheapest inference setting.

Two-panel line chart: left panel shows the entropy in bits of the segmentation sampling distribution decreasing from near-maximum to near zero as alpha grows from 0.02 to 3 for a three-path and a four-path lattice; right panel shows the expected number of pieces per word falling toward one over the same alpha range, both panels marking the default alpha near 0.2

Figure 2: The temperature $\alpha$ controls exactly two things that matter. Left: the entropy of the sampling distribution, which is how much augmentation diversity the model receives per epoch. Right: the expected pieces per word, which is the sequence-length tax you pay for that diversity. Beyond roughly $\alpha = 1.5$ both collapse to the deterministic Viterbi behaviour and the method does nothing.

Property	Deterministic BPE / unigram	Unigram-LM sampling	BPE-dropout
Randomness source	None; one path per sentence forever	Draw from the l-best lattice paths of a unigram LM	Independently skip each merge while applying the merge table
Knob and typical value	Vocabulary size only	l = 64, alpha near 0.1 to 0.2	p = 0.1 on both sides, lower for high-resource pairs
Segmenter changes	n/a	Needs a probabilistic unigram vocabulary, not a merge list	None; reuses an existing BPE model unchanged
Main cost	Rare pieces stay under-trained	Lattice sampling per example; longer average sequences	Re-merging per example; length inflation grows fast with p
Inference	Same path as training	Viterbi path, or average log-scores over k sampled sources	Standard BPE with dropout disabled

August 17, 2026

DL0114 Chain-of-Thought Prompting

What is Chain-of-Thought prompting and why does it improve reasoning?

Answer

Chain-of-Thought (CoT) prompting makes a language model emit intermediate reasoning steps before its final answer, either through few-shot exemplars whose demonstrations contain worked solutions (Wei et al., 2022) or through a zero-shot trigger phrase such as “Let’s think step by step” (Kojima et al., 2022). The gains are large on multi-step problems: on GSM8K, 8-shot CoT lifted PaLM 540B from 17.9% to 56.9%, and the same trigger took GPT-3 (text-davinci-002) from 10.4% to 40.7% with no weight updates. Two mechanisms explain why. The first is serial compute: a transformer with $L$ layers can apply only $O(L)$ sequential operations before it must commit to the next token, so a task needing more sequential steps than the depth provides is simply not representable in one forward pass; every generated token reruns the whole stack, turning the context window into an external scratchpad that raises effective serial depth from $L$ to $L(T+1)$ . Theory backs this up: constant-depth transformers with a polynomially long chain can simulate any polynomial-time computation, while the same model restricted to direct answers cannot. The second mechanism is factorized conditioning: instead of sampling one high-entropy jump from question to answer, the model decomposes the problem into a product of low-entropy conditionals, each conditioned on the already-written steps, which keeps it on the distribution of solution traces seen in pretraining. CoT is not universally free money: it is emergent with scale, and later meta-analysis shows the gains concentrate on math, symbolic, and logical tasks rather than on knowledge or commonsense retrieval.

(1) Two Elicitation Modes: few-shot CoT shows worked examples and controls format tightly, while zero-shot CoT appends one trigger sentence and costs almost no prompt tokens.
(2) Depth Becomes Time: each scratchpad token buys another full pass through the network, so the model trades latency and tokens for sequential computation it structurally lacked.
(3) Easier Conditionals: the chain rewrites one hard prediction as many easy ones, and each step reads all earlier steps as ordinary context.
(4) Emergent, Not Universal: below roughly 10B parameters, CoT often produces fluent but invalid chains and can score below direct prompting.
(5) Chains Are Not Guaranteed Faithful: models shift answers when a prompt carries a bias cue while the written rationale never mentions it, so a chain is an output artifact, not an audit trail.
(6) Cheap Ensembling On Top: because chains are stochastic, sampling $k$ of them and majority-voting the final answers (self-consistency) pushed PaLM 540B on GSM8K to 74.4%.

Two-panel diagram: direct prompting sends the question through one L-layer forward pass to the answer, while chain-of-thought generates T scratchpad tokens, each of which reruns the full L-layer stack before the answer token

Figure 1: Direct prompting caps sequential computation at the network depth $L$ ; CoT reuses the same weights once per generated token, giving $L(T+1)$ sequential steps and letting step $t$ read every earlier step as context.

The costs are as concrete as the benefits. Chains are typically 100 to 400 decode tokens on grade-school math, so a CoT query can cost an order of magnitude more decode passes and time-to-last-token than a direct answer, which matters for anything user-facing. Errors also propagate: the decoder never backtracks, so a slip in the first arithmetic step is carried through every later step, and accuracy behaves roughly like the product of per-step accuracies. Few-shot CoT is additionally prompt-sensitive, since exemplar choice, ordering, and even the formatting of the equations move accuracy by several points. These properties are exactly why the field moved from prompting toward verifying or training the chain rather than just requesting it.

Mathematical Formulation:
$p(a \mid q) = \sum_{z} p(a \mid q, z)\, p(z \mid q)$
$p(z \mid q) = \prod_{t=1}^{T} p(z_t \mid q, z_{1:t-1})$
$D_{\text{direct}} = L$
$D_{\text{CoT}} = L \cdot (T + 1)$
$\hat{a} = \arg\max_{a} \sum_{i=1}^{k} \mathbf{1}[a_i = a]$

Where:

$q$ is the prompt (question plus any exemplars), $a$ the final answer span, and $z = z_{1:T}$ the generated rationale, a latent variable the model writes into its own context.
$t \in \{1, \ldots, T\}$ indexes chain tokens; $T$ is the chain length and $L$ the number of transformer layers.
$D$ counts the sequential layer applications available before the answer is committed, which is the resource CoT actually adds.
$k$ is the number of sampled chains in self-consistency, $a_i$ the answer extracted from sample $i$ , and $\mathbf{1}[\cdot]$ the indicator used for majority voting.
Required condition: the sum over $z$ is intractable, so plain CoT approximates it with a single greedy or sampled $\hat{z}$ , and self-consistency approximates it with $k$ Monte Carlo samples at temperature $T_s > 0$ .

Grouped bar chart of GSM8K solve rate for PaLM at 8B, 62B and 540B parameters under standard prompting and 8-shot chain-of-thought, with the CoT advantage appearing only at 62B and growing at 540B

Figure 2: CoT is an emergent ability: at 8B parameters the chains are fluent but wrong and buy nothing, while at 540B they roughly triple the GSM8K solve rate. Values are approximate figures reported for PaLM by Wei et al. (2022) and Wang et al. (2023).

Property	Zero-shot CoT	Few-shot CoT	CoT + self-consistency
How it is elicited	One trigger sentence appended to the question	A handful of exemplars that show the full reasoning trace	Same CoT prompt sampled k times, majority vote on the extracted answer
Extra prompt tokens	Roughly 10	Hundreds to a few thousand, and it consumes context	Unchanged prompt, but k independent decodes
Reported GSM8K result	10.4% to 40.7% on text-davinci-002	17.9% to 56.9% on PaLM 540B, 8-shot	74.4% on PaLM 540B with 40 sampled chains
Main weakness	Format drift: the model may skip steps or answer immediately	Sensitive to exemplar choice, order, and formatting	k times the decode cost; needs a comparable, extractable final answer
Best fit	Quick baseline and open-ended chat traffic	A fixed task where output format must be stable	High-value math or code queries where accuracy beats latency

August 17, 2026