Transformer Architecture

From token IDs to a next-token distribution

Let B be batch size, T sequence length, d hidden width and V vocabulary size. An embedding lookup maps integer IDs of shape (B,T) to vectors of shape (B,T,d). A bigram table predicts from the current token alone. Contextual attention lets the model use a longer prefix.

In a conventional head, Q,K,V have per-head width p. Attention computes $\operatorname{softmax}\left(\frac{QK^\top}{\sqrt p}+M\right)V$. The additive mask $M$ excludes future keys before softmax. Concatenated head outputs are projected back to d. GQA and MLA change the key/value representation; see Attention →.

Training shifts the target by one token. For [A,B,C,D], input [A,B,C] predicts [B,C,D]. All three positions can be trained together because the causal mask prevents target leakage. Generation appends one sampled token and repeats; a KV cache avoids recomputing earlier keys and values.

# Shape pseudocode; attention applies the causal mask internally.
x = embedding(ids)                    # B,T,d
for block in blocks:
    x = x + block.attention(block.norm1(x))
    x = x + block.ffn(block.norm2(x))
logits = final_norm(x) @ output_weight.T  # B,T,V
loss = cross_entropy(logits, next_token_targets)

Without positional information, unrestricted self-attention is permutation-equivariant: permuting inputs permutes outputs. A causal mask already breaks unrestricted permutation symmetry. Position mechanisms add explicit order or relative-distance structure; see Positional Encodings →. Equivariance is not the same as invariance.

Why keep an identity path?

Without skips, backpropagation multiplies sublayer Jacobians. Many contractions can shrink gradients; many expansions can increase them. With a residual, a layer can start near the identity instead of learning that identity through every weight matrix.

The identity term helps but does not guarantee stable gradients. Products of I+J can still contract or expand. A branch can cancel a direction. Initialization, normalization, optimizer settings and depth remain coupled.

Unrolling gives $x_L=x_0+\sum_{\ell=0}^{L-1}F_\ell(x_\ell)$. This is an exact sum of the actual updates, not a sum of independent random variables. A state can grow while each normalized branch input remains bounded.

Choose which statistics to share

Before choosing a normalization method, specify the axes. BatchNorm (2015) shares feature statistics across examples. LayerNorm (2016) computes statistics within one example. In a Transformer, LayerNorm usually operates over one token's d features. RMSNorm (2019) keeps the same token-local axis but omits centering.

BatchNorm uses the corresponding batch mean and variance per feature, then applies gain and bias. In a CNN, its reduction often also includes spatial positions. Training generally updates running statistics; inference uses those stored statistics. Forward variance and running-variance estimators can differ by implementation.

The original BatchNorm motivation was to reduce changes in layer-input distributions during training. Later work studied smoother optimization landscapes as an explanation of its benefits. These are explanations to test, not guarantees for every model. BatchNorm and optimization smoothness ↗

Small batches can make BatchNorm statistics noisy. Synchronizing them across workers adds communication. Pooling sequence positions requires careful masking and causality. This is not proof that BatchNorm cannot work in language models.

LayerNorm and RMSNorm do not depend on other batch examples for their statistics. Their identical train/test normalization rules do not mean the whole network is context-independent. Attention still couples token representations.

“Internal covariate shift” was BatchNorm's original motivation, not a settled explanation of all its benefits. Avoid assuming a normalization change always permits a higher learning rate. RMSNorm's paper reports model-dependent running-time reductions of 7–64%; that is not a guaranteed modern-kernel or end-to-end speedup. RMSNorm report ↗

Centering and scaling are different operations

Take one vector and add the same offset to every feature. LayerNorm discards that offset. RMSNorm does not. Both methods divide by a statistic computed from all four features.

Positive common rescaling approximately cancels in both methods, subject to ε. Common shifts cancel only in LayerNorm. Learned gain and bias can change the normalized mean and variance afterward. For a constant vector, LayerNorm produces zero before its affine transform; RMSNorm can retain a nonzero constant direction.

These operations couple feature derivatives. For normalized outputs before learned gain and bias, define s and r as the corresponding divisors.

Replacing normalization therefore changes optimization, not only runtime. A trainable tanh scale can control saturation without recreating coordinate coupling.

Pre-Norm and Post-Norm change the gradient path

Pre-Norm controls what the branch reads, not the magnitude of the accumulated state. As that state grows, a similarly sized update can become a smaller relative change. This can reduce effective depth in some regimes. Post-Norm controls the resulting state's scale but changes backward propagation.

Neither is universally stronger or more stable. DeepNorm couples a modified Post-Norm skip with initialization. Its form is $\operatorname{LN}(\alpha x+F(x))$, not simply a small branch multiplier. For an M-block decoder-only configuration, the report uses $\alpha=(2M)^{1/4}$ and selected Xavier gains $\beta=(8M)^{-1/4}$; Q/K gains remain 1. Do not transplant this recipe into a different block without checking its derivation. DeepNet Figure 2 ↗

Keep depth and width scales separate

Suppose a linear projection sums d independent unit-variance inputs with independent zero-mean weights. Weight variance $\frac1d$ keeps each output variance near 1. This is a width calculation. Nonlinearities, weight sharing and correlated inputs change its assumptions.

A branch multiplier of order $\frac1{\sqrt L}$ controls accumulated energy under the uncorrelated model. Coherent accumulation calls for a different analysis. This generic scaling argument is not the DeepNorm formula.

BERT-style 0.02 initialization is a model recipe, not a universal width law. At $d=768$, $\frac1{\sqrt d}\approx0.036$. This numerical comparison alone does not explain why a particular training run succeeds. Biases, nonlinearities, normalization placement and warmup matter.

Likewise, $\sqrt p$ attention scaling follows from a specific independent-coordinate dot-product variance calculation. Real queries and keys can be correlated. The derivation and its limits belong in Attention →.

What does an activation change?

Composing only affine layers remains affine. A nonlinearity lets a feedforward branch build a richer token-wise transformation. Its value and derivative shape which signals and gradients pass through.

A ReLU unit can receive zero branch gradients when its preactivations remain negative. This “dying ReLU” state is not necessarily permanent: upstream changes can move its inputs. Leaky ReLU changes that local derivative; it does not guarantee global optimization success.

Sigmoid and tanh still serve as output links and gates. Saturating tails do not make them obsolete. A sigmoid output maps one logit to a binary-class probability; stable training uses a logits-based loss. Binary cross-entropy → Hardtanh replaces smooth saturation with clipping; quantization or bounded-state applications need their own accuracy and gradient checks.

Softmax couples several logits and normalizes them to a probability vector; it is not an independent scalar activation. Sparsemax and entmax use different geometries and matched losses; see Losses and Objectives →. No curve establishes a universally best activation or a universal variance guarantee.

From an activation to a gated feedforward layer

A gated linear unit (GLU) multiplies one linear branch by a sigmoid-transformed branch. SwiGLU changes that gate to SiLU. The gate scales values; it is not a sampled probability that a neuron fires. SwiGLU gates can be negative or exceed 1.

An ordinary two-projection FFN has about 2dm parameters. A three-projection gated FFN has about 3dm. Matching this parameter budget requires reducing gated m to about two-thirds of the ordinary m. Equal hidden ratios are not equal budgets.

MoE makes the feedforward transformation conditional on expert routing. It does not remove the gate or residual design choices. See Mixture of Experts → for routing, balancing and expert layout.

A residual stream also mixes information across depth

The ordinary skip does more than help gradients. Its running sum gives every preceding update coefficient 1. Individual contributions can cancel. Once only the sum remains, a later layer cannot retrieve one contribution separately.

Each layer learns a pseudo-query w_l. Its keys depend on the token's earlier states. Therefore a fixed learned query can still produce token-dependent weights. Keys are RMS-normalized for scoring; values remain the original source vectors.

Token attention chooses positions. Depth attention chooses earlier sublayer outputs. Full Attention Residuals (AttnRes, 2026) retains every source. For L sublayers and width d, depth retrieval adds O(L²d) arithmetic and O(Ld) history per token. Counting attention and MLP separately avoids a factor-of-two ambiguity.

The proposing team's target is Pre-Norm dilution. A convex depth mixture avoids an unbounded all-ones sum of fixed sources. The sources themselves can still grow. This is an architecture change, not a post-hoc patch for an existing checkpoint.

Read a source, not just the accumulated sum

A zero pseudo-query initially produces uniform weights, not the ordinary unit-weight sum. Normalization changes the read's scale as well as its source preference. Key normalization reduces magnitude-based score domination; learned gains and source norms still need monitoring.

Compress history into block summaries

Full history is costly to retain and communicate. Block AttnRes sums sublayer outputs within a block. A later layer attends to the embedding, completed block summaries and the current block's partial sum.

With N blocks of roughly $S=\frac LN$ sublayers, history storage becomes $O(Nd)$ per token. $N=L$ approaches Full AttnRes. $N=1$ retains the embedding separately and is close to ordinary accumulation, but is not identical to it.

A block boundary changes which sources are stored, not the meaning of an attention versus MLP operation. Compute a sublayer's read before updating its partial sum. At a boundary, append the finished summary once and begin a fresh partial sum. Do not count the same output in both sources.

The report's ablations find about eight blocks recover most Full AttnRes gains across its tested scales. A few distant summaries outperform retaining only nearby individual outputs in those comparisons. Block count is an empirical architecture/system choice, not a universal constant.

Make depth retrieval practical

Without recomputation, training already retains many activations. With recomputation or pipeline parallelism, additional history changes the storage and transfer pattern. The report caches summaries on physical pipeline ranks and transfers newly completed blocks rather than retransmitting the full history.

For P physical and V virtual stages, its peak transition transfer falls from O(PV) to O(P) chunks under that schedule. This is a transfer-count comparison, not a complete byte or latency model. Measured Block AttnRes training overhead is negligible without pipeline parallelism and under 4% with it on the reported stack.

At inference, fixed pseudo-queries permit a two-phase schedule. First batch all S queries against completed blocks. Then compute the sequential intra-block partial source and merge it with the precomputed result.

For $L=128, N=8, S=16$, estimated residual I/O per token per sublayer is $3d$ elements for ordinary residuals, $5.5d$ for Block AttnRes, $24d$ for scheduled Full AttnRes and $34d$ for four-stream mHC. These are residual-mechanism estimates, not whole-model traffic.

The report measures under 2% inference-latency overhead on its typical workloads. Long-context prefill adds $NTd$ history elements: its 128K, eight-block example is about 15 GB before sharding. Sequence sharding across P devices reduces that cache to $N\frac TP d$ elements per device, with integration into reduce-scatter/all-gather. See Distributed Training and Parallelism →.

What the Attention Residuals experiments establish

The proxy study uses MoE models with 194M–528M activated non-embedding parameters, 38.7B–119B tokens, 12–17 attention layers plus equally many MLP layers, and 8K context. Its fitted curves describe proxy training loss, not serving latency.

At 5.6 PFLOP/s-days, the fitted Block loss is 1.692 versus baseline 1.714. The inferred 1.25× compute advantage is a curve comparison, not a universal wall-clock saving. At the largest proxy point, measured validation losses are baseline 1.719, Block 1.693 and Full 1.692.

The input-dependent query improves this loss but requires a d×d projection and sequential reads. The simpler pseudo-query supports batched inference. The multi-head result suggests shared depth choices are useful in this setting; it does not prove channel-specific retrieval is never useful.

The larger controlled experiment is Kimi Linear: 48B total, 3B active parameters; 27 Transformer blocks counted as 54 sublayers; eight of 256 routed experts plus one shared expert. Six sublayers per AttnRes block produce nine block summaries plus the embedding. Training uses Muon, WSD, an 8M-token global batch, 1T pretraining tokens at 4K, about 400B high-quality mid-training tokens and context extension to 32K.

The method matches or improves every listed score; MMLU-Pro ties. Larger reasoning/code gains are consistent with the team's compositional-retrieval hypothesis, not proof of that mechanism. This is not a controlled post-training or agent comparison.

Learned weights are mainly local, with long skips and persistent embedding reads. Attention sublayers use broader depth histories; MLP reads are more concentrated near recent sources. Some sources become depth-wise attention sinks. Block summaries may regularize as well as compress, but that causal interpretation remains a hypothesis.

The report interprets recurrent multi-stream mechanisms as structured depth mixers, unlike dense cross-layer access. Stronger residual baselines matter. No independent replication is established here. Its largest controlled model has 3B active parameters; fitted compute and custom-kernel overheads must not be generalized to larger or unoptimized deployments. Deeper/narrower proxy preferences also increase serial inference depth. Attention Residuals report ↗ · Implementation ↗

Gated Residual: spend extra capacity on the read

A single residual sum gives every new sublayer the same accumulated state. Widening that stream creates separate paths, but each path also costs memory traffic. Qwen separates two questions: does width help, and which read/write machinery justifies its cost?

A simple four-branch baseline reads a learned scalar mixture and writes to one branch in round-robin order. It already improves loss. Hyper-Connections adds data-dependent reads, writes and branch mixing. Gated Residual (GR) keeps richer reads and writes but removes the explicit branch-mixing matrix.

The gate builds on GatedNorm, which supplies an explicit feature-dependent rescaling after normalization:

Normalize each branch independently, with its own gain. Predict gates from all branches rather than pooling first. Each read gate lies between zero and one; each write strength lies between zero and two. Read weights are not a softmax distribution and need not sum to one. GR replaces the sublayer pre-normalization; adding another ordinary Pre-Norm would change the architecture.

The dynamic variant reduces loss only 0.002 relative to static mHC, but improves mean accuracy by 1.98 percentage points. Width and data dependence therefore look different under loss and downstream evaluation. GR and dynamic mHC are close at this scale; the argument for GR also includes memory traffic and stability, not a decisive quality margin alone.

Channel-specific reads help; channel-specific writes add almost nothing. Branch mixing adds no significant benefit after improving the read/write path. Removing that operator saves one full residual-state read per block. Standard random initialization suffices in the final GR configuration; a learned static term adds no useful gain. Each attention and MLP sublayer has its own GR module.

What do the extra paths preserve?

In a matched 20-layer probe, the authors decompose writes into later reads at the gates taken by the actual forward pass. Across five GR checkpoints, one branch carries long-range paths; the others remain more local. The branch identity is exchangeable, not assigned in advance.

The normalized contribution of layer-0 GDN to layer-15 attention rises from 0.020 to 0.138. Typical long-range skips are 10.9 layers versus 3.4–3.9 for other branches. These contribution magnitudes are descriptive, not causal attribution or probabilities that a specific fact survives. The average skip changes little; GR redistributes influence toward selected short and long paths.

Reading only the two highest-gated branches looks nearly harmless during pretraining but degrades post-training quality. They reject that shortcut. FP8 residual storage instead halves this state's bytes relative to BF16 with almost no reported quality loss. This does not halve all activation memory or establish FP8 arithmetic throughout the model. Read and write are separately fused. Gates under stress → · Residual traffic →.

Qwen3.8-Next architecture report, §2.2, Tables 5–6 and Figure 7, pp. 10–14 (26 August 2026) ↗.

One state, many vocabulary decisions

Each token competes through its output embedding w. The context network can be deep and nonlinear. The final head still compresses every vocabulary comparison into d shared directions. Numerical stabilization subtracts each row's maximum before exponentiation; that preserves probabilities but does not increase expressivity.

The output softmax is distinct from attention's softmax. One selects vocabulary tokens, the other selects source positions. Cross-entropy measures prediction error against targets; see Losses and Objectives →.

Sharing embeddings changes the parameter budget

Input embeddings and an ordinary output matrix both have V×d parameters. Tying uses the same matrix for both. A tied head projects h against the rows used for token lookup; it does not skip the vocabulary projection.

At V=50,000 and d=2,048, tying saves 102.4M parameters. That can be substantial in a small model. In a larger backbone, it is a smaller fraction of the budget.

Compare two distinct experiments. Keeping the backbone fixed isolates the extra freedom and extra parameters of untying. Keeping total parameters fixed requires reallocating width, depth or another part. Those comparisons answer different questions.

Tying can regularize or reduce quality in a particular setting. Untying can improve freedom without removing the output-rank limit. Large-model parameter fractions alone do not establish a quality benefit.

Why softmax cannot recover every distinction

A bias-free linear head factors logits through width d. Log normalization subtracts one scalar per context, so it adds at most one rank-one term.

The sharper test removes these row constants. Choose a reference token r. The log-odds matrix $D_{c,x}=\log\frac{p(x\mid c)}{p(r\mid c)}$ equals $h_c^\top(w_x-w_r)$. Its rank is at most d. The probability matrix P itself can have higher rank; that is not a contradiction.

This 2×2 log-odds matrix has determinant $(\log 2)^2$ and rank 2. A bias-free $d=1$ head cannot represent it. Increasing Transformer depth without increasing the final factorization width cannot repair that exact limitation.

A row shift changes rank by at most one. Thus $d<\operatorname{rank}(A)-1$ is a sufficient obstruction. Output bias adds a shared vocabulary term of rank at most one. It changes these precise bounds; the two-context example above is deliberately bias-free.

Mix distributions after normalization

A mixture of contexts combines hidden states first: $h'=\sum_k\pi_k h_k$, then $\operatorname{softmax}(h'W^\top)$. This is still one d-dimensional state. A mixture of softmaxes (MoS) normalizes each component separately, then adds probability vectors.

MoS is not MoE. A routed FFN changes how h is produced. MoS changes how the final distribution is formed. A model can use either, both or neither.

Separate algebraic rank from useful probability error

Yang et al.'s ICLR 2018 work evaluates RNN language models, not contemporary large Transformers. On Penn Treebank with 10,000 tokens, reported empirical log-probability ranks are 400 for a width-400 softmax, 280 for width-280 MoC and 9,981 for width-280 MoS. Numerical rank depends on the threshold.

Additional components initially improve rank and likelihood, but can stop helping or overfit. The original settings report roughly 2–3× slower training with up to 15 components. These experiments establish a useful counterexample, not a recommendation that every LLM replace its head. Softmax bottleneck: theory and PTB evidence ↗

High exact rank can be dominated by tiny singular values. The practical question is probability error in the discarded directions. The true language distribution's high rank is a hypothesis, not an observable theorem. Computing vocabulary logits can also be expensive; that computational bottleneck is separate from representational rank.

For a diagnostic, collect diverse contexts and a fixed token subset. Compute stable log probabilities, subtract one reference-token column, and inspect the singular-value spectrum. Report precision, threshold and matrix dimensions. Too few contexts or tokens cap measured rank. Better rank alone does not guarantee better cross-entropy.

A gradient boundary is not a teacher update

A self-supervised matching loss can compare two augmented views of the same input. The student predicts the detached target representation. Detachment changes the backward path, not either view's forward value.

A detached target branch can still share parameters with the student. Gradients through the student path update those shared parameters. Detaching one tensor does not freeze every use of its weights.

An exponential moving average (EMA) teacher is a separate update rule.

BYOL and DINO use momentum targets. SimSiam uses stop-gradient without a momentum teacher. Stop-gradient does not inherently require storing two parameter copies. SimSiam ↗

For a plain matching loss, a constant output can still minimize disagreement. Stop-gradient alone is not a general anti-collapse proof. Predictors, normalization, centering, variance objectives or negatives change the full system. Self-Supervised Vision Learning connects DINO, BYOL, SimSiam, LeJEPA and variance regularization to their complete anti-collapse objectives.

Reconstruct activations instead of storing every layer

Ordinary x+F(x) is not automatically invertible. Reversible blocks impose an additive coupling structure. Split channels into x₁,x₂, then update them in sequence. This Feistel-like additive coupling is a reversible data path, not an encryption scheme.

Under the ideal reversible-stack assumptions, saved boundary activations can become O(1) in depth rather than O(L). Parameters, gradients, optimizer states, temporary branch activations and non-reversible boundaries still consume memory.

Backward reconstructs an input, recomputes the local forward and differentiates that local computation. Ordinary eager autograd may still retain the whole forward graph unless a custom backward/storage strategy is used. Reconstruction is not a gradient approximation in exact arithmetic.

The original RevNet operation count estimates 4N versus ordinary 3N: about 33% extra. With roughly equal GPU forward/backward costs, the paper discusses about 50% instead. Neither is a universal measured overhead. Checkpointing makes a related recomputation/storage tradeoff without requiring this architecture. RevNet: coupling and recomputation ↗

Floating-point subtraction can accumulate reconstruction error. Recompute dropout with the same random state. Store or checkpoint non-reversible operations. Verify reconstructed values and parameter gradients against a stored-activation baseline before claiming memory savings.

Count the actual architecture

Use conventional attention geometry $d=hp$, with h query heads, $h_{\mathrm{kv}}$ key/value heads and p head width. GQA usually requires h divisible by $h_{\mathrm{kv}}$. Some architectures use a different attention-inner width; the following formulas must then change.

Depth scales this parameter estimate linearly. Width generally contributes quadratically when head geometry and FFN ratio scale with it. KV cache instead depends directly on key/value heads and context length. Local-window eviction changes cache only if the serving implementation actually evicts those states.

A 4,096-wide model with 32 heads cannot use $p=64$ under $d=hp$. For $p=64$, use 64 query heads or explicitly change attention-inner width. The 14,336 FFN width is $3.5d$: its gated cost is $10.5d^2$, not the ordinary $4d$ FFN's $8d^2$.

Head sizes 64 or 128 and GQA ratios 2,4,8 are candidate configurations, not quality or hardware guarantees. Profile kernels and compare trained quality. Use a dense baseline before attributing benefits to sparse layouts, shared experts or dense/MoE interleaving.

Parameter count is not wall-clock cost. MoE has distinct total and active parameters; looped models can reuse parameters while increasing compute. MTP, multimodal frontends, state-space mixers and nonuniform layer widths change accounting. Do not force MLA, OpenELM or Mamba into these conventional formulas.

Match data, batch size, optimizer, token budget and evaluation when comparing shapes. Critical batch size and FP8 behavior require their own training analysis. A depth-before-width heuristic must be tested against serial latency and the residual mechanism, not treated as a universal allocation law.

Turn a budget into testable designs

Choose deployment memory, latency and context requirements first. Then choose width and depth together, validate head geometry, set KV sharing, and tune feedforward capacity. These are experiment axes, not an optimal recipe.

Count unique parameters for storage, executed active parameters for arithmetic, and repeated passes for total compute. Dense interleaving can make shared experts less necessary in a tested design; it does not prove shared experts are useless. Dropless routing still needs a declared capacity and imbalance policy.

Loop boundaries can be candidates for early exit when training exposes those states to output objectives. Evaluate stopping quality and latency rather than assuming every reused block is output-ready. Recurrent and Looped Models → · Routing and expert layout →

Increasing batch size beyond the useful critical regime can reduce returns. The critical regime depends on the data, loss and optimizer. Compare more steps, different data or architecture changes at matched compute. Optimization → · Scaling Laws →

Measure widths, precision formats and kernel support on the actual stack. Nonuniform layer widths, latent-compression models, sparse multimodal frontends and state-space mixers require their own accounting. Precision → · Local/global access → · Performance accounting →

Checks before scaling the model

Start with tensor contracts: residual additions must match shapes; heads must match the declared inner width; training targets must be shifted. Confirm causal masks with a future-token perturbation test.

Then measure hidden-state RMS, branch-output RMS and gradient RMS by depth. Change normalization placement and initialization together only with an explicit rationale. A bounded branch input is not proof of a bounded residual stream.

Check tied weight identity, vocabulary shape and parameter counts. Distinguish rank limitation from numerical softmax stability. For reversible execution, compare values and gradients and measure peak memory, not only the number of saved tensors.

Finally measure complete training and serving cost. A lower loss at matched FLOPs can lose its advantage through communication, memory traffic or deeper serial execution. Use the parallelism reader to connect architecture to placement and communication.

Sources and connected notes

BatchNorm (2015) ↗ · LayerNorm (2016) ↗ · RMSNorm (2019) ↗ · DeepNorm (2022) ↗ · Attention Residuals (2026) ↗

Softmax bottleneck (2017 / ICLR 2018) ↗ · Monotonic nonlinear logits ↗ · Sigsoftmax ↗ · Reversible networks ↗ · SimSiam ↗

Smol training playbook ↗ · Initialization and parameterization ↗ · BERT’s 0.02 initialization ↗ · DeepNorm discussion ↗ · Pre-Norm versus Post-Norm ↗ · Tanh versus normalization ↗

Attention Residuals ↗ · Kimi Linear ↗.