Handout #14 — LLM Architecture Choices and Dimensions

Session 13 ended with the runtime question on slide 13. This continuation starts with the RMSNorm evidence, then develops the deferred architecture choices: what a gated FFN computes, which dependencies a parallel block removes, and how a shared embedding learns through two roles. It then connects FFN width, attention-head dimensions, and depth/width choices to shape constraints and systems budgets.

Warm-up: shapes, sharing, and dataflow in 5–7 minutes

Use the prerequisites from #11–#13 before learning the new variants. Predict first, write your reasoning, then open an answer. The scalar puzzle deliberately omits normalization so only the dependency difference remains.

a = nn.Parameter(torch.tensor([1., 2.]))
b = nn.Parameter(a.detach().clone())
c = a
with torch.no_grad():
    a.add_(1)

1From a runtime argument to empirical evidence

#13 §2.4 explains why small FLOP share can coexist with meaningful runtime cost. Slide 14 now supplies a measured comparison: the displayed baseline and RMSNorm runs have the same rounded parameter/operation totals, but different throughput and final loss. A rounded count cannot reveal kernel launches, reductions, memory traffic, or all small operation differences.

Keep two separate choices from slides 15–16 in mind: placement (pre-norm versus post-norm) and statistics (LN versus RMSNorm). Pre-norm describes where normalization sits; it does not specify mean subtraction. RMSNorm controls a token’s scale using its second moment; it does not itself specify the residual layout. Model-family examples show common combinations, not a theorem that every architecture must use one combination.

2Gated FFNs: more projections, a fairer comparison

2.1 A smooth activation is not a random procedure

ReLU is $\max(0,z)$; GELU smoothly attenuates inputs and can produce small negative outputs. At $z=-1$, ReLU is 0 but GELU is approximately −0.158655. GELU's Gaussian interpretation does not make an ordinary GELU forward pass stochastic: identical inputs give identical outputs. Libraries may also offer a tanh approximation, so check the chosen implementation when comparing numbers.

Look near zero, then far to the left. The small negative lobe is easy to miss if you look only at large positive values, where the curves nearly coincide. The numeric value at $z=-1$ above makes that difference measurable. The plot describes a deterministic activation; it does not show random masks or a training-quality comparison.

2.2 Trace both views of one token

Use row vectors for FFN equations. A bias-free dense FFN is $f(xW_1)W_2$, with $W_1\in\mathbb R^{D\times F}$ and $W_2\in\mathbb R^{F\times D}$. A gated FFN has two input projections:

For SwiGLU, $f(g)=\operatorname{SiLU}(g)=g\operatorname{sigmoid}(g)$; for GeGLU, $f$ is GELU. SiLU names the nonlinearity; SwiGLU names the gated construction. Replacing GELU with SiLU in a two-matrix dense FFN does not add the second input projection or the elementwise gate. The multiplication is elementwise, not a dot product. Each hidden coordinate's response depends on two learned views of the same input token. All tokens use the same matrices.

These are the hidden products before the down-projection, not vocabulary probabilities. The gate can be negative, and SiLU(2) is about 1.762, so “volume knob” does not mean a probability constrained to $[0,1]$. Nor is the gate shared across tokens: matrices are shared, but their outputs depend on each token.

class SwiGLU(nn.Module):
    def __init__(self, d, hidden):
        super().__init__()
        self.gate = nn.Linear(d, hidden, bias=False)
        self.up = nn.Linear(d, hidden, bias=False)
        self.down = nn.Linear(hidden, d, bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

2.3 Match parameters before judging the gate

Ignoring biases, dense FFN parameters are $2DF_{\rm dense}$; gated FFN parameters are $3DF_{\rm gated}$. Equating these gives:

For $D=768$, dense width 3072 and gated width 2048 each use 4,718,592 parameters. Keeping gated width 3072 gives 7,077,888: 50% more FFN parameters. A quality gain in that comparison could partly be a budget gain. Forward projection arithmetic is also matched when the two different hidden widths obey this equation: approximately $2BT$ times the matrix parameter count, counting a multiply and add as two FLOPs. Activation operations and saved intermediates still differ.

Real implementations often round hidden widths to hardware-friendly multiples. Report the actual width and actual count; “roughly two thirds” is not an exact equality after rounding. PaLM’s bias-removal stability observation (noted on slide 18) is an empirical result in its setting, separate from the small parameter saving. Removing the dense biases saves only $F+D$ parameters: at $D=768,F=3072$, that is 3,840 versus 4.72 million matrix entries.

3Serial versus parallel: find the missing dependency

Write both formulations with explicit intermediates. $A$ is attention and $M$ the FFN:

In serial, the FFN sees attention's update from this block. In parallel, both branches read the original normalized stream. Attention information can reach a later block's FFN, but there is no attention-to-FFN edge inside the current block. Changing this edge generally changes the function, even if the parameter values are identical.

3.1 A small counterexample to “same computation, new schedule”

To isolate the dependency, use $x=[1,2,3,4]$, affine-free LN with epsilon $10^{-6}$, $A(z)=[z_0,0,0,0]$, and $M(z)=z$. These toy branches are not an attention implementation. The same functions give different outputs; the lab computes both:

Parallel inputs permit concatenating attention QKV and FFN input-projection weights into a larger projection of $n$, then splitting outputs. This does not remove the constituent matrix entries. It can reduce launches and improve utilization. Python statements running one after another are still sequential; actual fusion/concurrency requires a suitable implementation. Some parallel designs use separate branch norms, so a shared norm is an explicit choice here.

4Embedding tying is a constraint on learning

4.1 Shapes and actual sharing

An embedding matrix $E\in\mathbb R^{V\times D}$ maps token ID $i$ to row $E_i$. For final row-vector hidden state $h\in\mathbb R^D$, tied logits are $hE^\top\in\mathbb R^V$. In mathematical notation the output matrix is transposed; in PyTorch both nn.Embedding(V,D).weight and nn.Linear(D,V).weight are stored as $[V,D]$. Assign the same Parameter without transposing it.

Follow the two uses of the green parameters. At the input, a token ID selects a row. At the output, a hidden vector is compared with every vocabulary row to form logits. The red arrow denotes one shared parameter table; it does not make lookup and output projection the same operation, and it does not imply that softmax is removed.

self.embed = nn.Embedding(vocab, d)
self.head = nn.Linear(d, vocab, bias=False)
self.head.weight = self.embed.weight
# Tie during construction, before creating the optimizer.
assert self.head.weight is self.embed.weight

Copying values using copy_ or wrapping a detached clone in a new Parameter leaves two independently trainable objects. They can start equal and diverge on the first step. Use model.parameters() to enumerate the registered shared parameter once; do not concatenate overlapping module parameter lists into the optimizer. Replacing a parameter after constructing the optimizer can leave the optimizer tracking an obsolete object.

Tying is a modeling constraint, not merely compression: input and output roles must agree on one representation. Separate weights allow more independent choices. Tying also requires compatible hidden widths (or an extra adapter). A separate output bias is possible, but omitted here.

4.2 Why a token absent from the input can still update

For one prediction, let $p_j$ be the softmax probability of vocabulary item $j$ and $t$ its target ID. Holding $h$ fixed, the output-side derivative is:

The embedding-lookup route contributes only to rows looked up as inputs. The output route has a coefficient for every vocabulary row under full softmax, including rows absent from the input. Under tying, autograd sums both routes into the same gradient buffer:

“Dense signal” is not a guarantee every row has a nonzero or helpful update: zero hidden states, numerical underflow, cancellation, or modified objectives can change that. In a single ordinary full-softmax prediction with nonzero $h$, an absent, non-target word gets coefficient $p_j>0$. Holding $h$ fixed, the output-side gradient descent update lowers that word’s logit. With multiple predictions or a changing hidden state, individual logit changes need not all have this sign. It learns from being an alternative prediction, not from being observed as input.

4.3 Count savings without promising fewer logits

Untied input/output matrices contain $2VD$ parameters; tied matrices contain $VD$, saving $VD$. For $V=32{,}000,D=768$, this saves 24,576,000 parameters, or 46.875 MiB of two-byte weights. Under an illustrative training layout with a two-byte weight, two-byte gradient, and two float32 Adam moments, saving 12 bytes per parameter gives 281.25 MiB. A float32 master copy would add another 93.75 MiB. Actual optimizer, sharding, and dtype policies determine the real memory.

The logits still require a projection into $V$ classes. Tying does not remove the $[B,T,V]$ output, vocabulary softmax, or their arithmetic. The saved parameter fraction depends on depth and width of the rest of the model. Press & Wolf (2017) motivate sharing and study its language-modeling benefits; equal roles are an inductive bias, not a mathematical necessity for every model.

5Dimensions: conventions, constraints, and budgets

Slides 39–49 ask how to choose the sizes inside the architecture. Separate three kinds of statement: a shape constraint makes an operation legal, a budget equation counts its cost, and an empirical convention describes choices that have worked. Only the first is mandatory for the mathematics; a particular software interface may impose additional restrictions.

5.1 FFN ratios are budget choices, not shape requirements

The familiar dense choice $F=4D$ gives $2DF=8D^2$ matrix parameters per FFN. A bias-free gated FFN with $F=8D/3$ has the same count, as derived in §2.3. There is no requirement that the intermediate width be exactly either value: the down-projection can map any positive hidden width back to $D$. Hardware-friendly rounding and quality/compute tradeoffs help determine the actual setting.

The FFN sweeps on slides 41 and 49 suggest a broad useful region in their experimental setting, not a universally optimal multiplier. When a sweep holds non-embedding parameter count fixed, changing $F/D$ also requires changing other dimensions: it is not a comparison that merely adds hidden units to an otherwise identical model. Slides 42–43 show how the gated default follows from matching the projection budget. Slide 49 also summarizes a wider spread of actual gated ratios; rounding and different allocations of capacity can move a model away from exactly $8/3$.

5.2 Optional reading · Head count need not divide the residual width

For this subsection use ordinary MHA, with $H$ heads and the same per-head width $d_h$ for queries, keys, and values. Define the total attention width $A=H d_h$. We may choose $A$ independently of residual width $D$:

The widespread convention $A=D$ gives $d_h=D/H$, so that convention requires divisibility. The general construction requires $A=H d_h$, not $D=H d_h$. Queries and keys must have matching feature widths for their dot products; values can even use a different width if the concatenation/output projection is adjusted. We keep their widths equal here to isolate one choice.

This flexibility has a cost: the four bias-free projection matrices have $4DA$ parameters rather than necessarily $4D^2$. At fixed $A$, increasing $H$ shrinks $d_h$ without changing this projection count, but a materialized score tensor still has $BH T^2$ entries. At fixed $H$, increasing $d_h$ increases $A$, projection cost, and ordinary MHA KV-cache width. Equal parameter counts alone do not establish equal runtime or intermediate memory.

5.3 Optional reading · Equal block parameters can hide different systems costs

Let $L$ be the number of decoder blocks. Assume standard MHA with $A=D$ and a dense $F=4D$ FFN, and ignore biases and norms. Attention contributes $4D^2$ parameters and the FFN $8D^2$, so:

The same leading count holds for a budget-matched gated FFN. Vocabulary weights, positional parameters, and final normalization are extra. In particular, changing $D$ changes the $VD$ embedding count, so equal block counts do not imply equal total model counts.

These are counting examples, not recommended configurations. Width here is the hidden feature dimension, not context length. The deeper model has four times as many sequential block dependencies and smaller matrix operations. The wider model may offer more useful parallel work within a block; actual utilization depends on shapes and hardware. Pipeline parallelism can overlap work from different microbatches, but does not remove a single sequence’s layer dependencies.

Recall #12’s ordinary MHA KV-cache count, $2BLTDs$ bytes for $s$ bytes per cache element. With $B,T,s$ fixed, the two example models have cache sizes proportional to $12\cdot768$ and $48\cdot384$: the deeper one needs twice as much cache. Their full-sequence attention score/value products also scale with $LD$ at fixed $B,T$, while the main projection/FFN arithmetic scales with $LD^2$. Even their total arithmetic need not be equal.

Read the basin, not just one preferred ratio. Each curve compares shapes at approximately fixed non-embedding parameter count. Slides 46–49 highlight ratios around 100–200, but the plotted basin and the paper’s discussion are broader. The vertical axis in this slide crop has no numerical labels: use it to interpret the qualitative trend, not to read exact losses. These experiments motivate a range of candidate shapes; they do not prove a universal optimum or predict which candidate runs fastest on your hardware.

5.4 A configuration audit before training

The architecture recap on slide 49 combine conventions with choices for which consensus is weaker. Record the actual FFN width, head count and width, residual width, depth, normalization, activation, layout, and tying setting. Derive the shapes and count the real Parameters, then measure quality, throughput, and memory under the stated budget. Copying a successful configuration is a starting hypothesis; a legal shape, a familiar ratio, and an empirical quality result answer different questions.

6Optional lab: checks before training

These are component checks, not evidence of better language-model quality or throughput. Make each deliberate mutation separately, observe the failed assertion, and restore it before the next experiment. Session 13’s normalization checks remain in #13’s optional lab.

7Optional code reading: revisit nanoGPT with new questions

7.1 What does the second normalization read?

model.py · Block.forward · verbatim excerpt

        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))

The second line reads the updated $x$: this is the serial pre-norm layout in §3. Find the final normalization before lm_head; pre-norm does not imply an unnormalized output stream. Then inspect the MLP's two linear layers and GELU: it is a dense FFN, not a gated one.

        n = self.ln_1(x)
        x = x + self.attn(n) + self.mlp(n)

7.2 Find the assignment that constrains the optimizer

model.py · GPT.__init__ · excerpt, trailing source comment omitted

        self.transformer.wte.weight = self.lm_head.weight

The assignment direction differs from §4's teaching model, but both modules end up referencing one Parameter. In a separate checkout, replace the assignment's right-hand side with nn.Parameter(self.lm_head.weight.detach().clone()). Use GPTConfig(vocab_size=5, n_embd=12, n_head=3, n_layer=1, block_size=4, dropout=0.0, bias=False). The untied version has exactly 60 more parameters, and the identity assertion fails. Count with sum(p.numel() for p in model.parameters()), keeping every other setting fixed.

There is an initialization trap: the clone initially has equal values, but nanoGPT later calls self.apply(self._init_weights). Its two separate Parameters are then initialized independently, so they are generally unequal when construction finishes. To isolate “equal values, different objects,” clone after constructing the model and before creating its optimizer. The local lab does this with $V=5,D=3$: 15 versus 30 unique parameters.

Further edit: replace its dense FFN with the lab's bias-free SwiGLU. At $D=12$, compare hidden widths 48 (dense) and 32 (gated); both have 1,152 matrix entries when the dense biases are off. Keep causal attention intact and repeat prefix invariance after the substitution. This tests correctness of integration, not whether the new activation improves training.

nanoGPT excerpts: Copyright (c) 2022 Andrej Karpathy, MIT License. License and permission notice. The proposed edits are original exercises.

7.3 Find an implementation assumption about heads

Optional code extension of bonus slides 44–45 and §5.2.

model.py · CausalSelfAttention.__init__ · verbatim excerpt

        assert config.n_embd % config.n_head == 0

Locate the QKV output width and the reshape that computes head width from n_embd // n_head. Explain why removing this assertion alone cannot implement §5.2’s $D=10,H=3,d_h=4$ example. Then read FlexibleAttention in the local lab: its projections use the independent attention width $A$, and its output projection maps $A$ back to $D$. In the Q/K/V reshape, replace self.head_dim with x.size(-1) // self.heads and confirm that the reshape fails.

8Practice: diagnose the claim, then design the check

Questions A–D cover the core reading. E–F are optional practice for bonus slides 44–48 (§5.2–§5.3).