The Architecture of Intelligence: A Deep Autopsy of the Transformer

The Sequential Cage — Why RNNs Were a Dead End

Before we can appreciate what the Transformer is, we need to feel the weight of what it replaced. The year is 2017. The dominant paradigm for processing language—for any sequential data, really—is the Recurrent Neural Network. The LSTM variant, introduced by Hochreiter and Schmidhuber in 1997 and refined through the early 2010s, is the workhorse. It's clever. It carries a "hidden state" forward through a sequence, updating it token by token. Its gating mechanisms—the input gate, forget gate, output gate—are designed to selectively remember and selectively erase. On paper, it's elegant. In practice, it is a machine for producing catastrophic forgetting.

Here is the central tragedy of the RNN: it must process tokens one at a time, in strict order. This is not a software limitation. It is structural. The hidden state $h_t$ at time $t$ depends on $h_{t-1}$. You cannot compute $h_5$ until you have $h_4$. The computation graph is a chain. This makes parallelisation across the length of a sequence fundamentally impossible. When you train on modern GPU hardware—hardware that is extraordinarily good at doing thousands of operations simultaneously—the RNN makes you throw that hardware's strength away and work serially. It's like hiring a hundred architects and then insisting they draft blueprints in a single room, one at a time, each forced to wait for the person before them to finish.

Then there's the gradient problem. When you backpropagate through a chain of 500 time steps, the gradient must travel that entire chain to update weights at the beginning. It either vanishes—shrinking exponentially because repeated multiplication of values less than one collapses to zero—or it explodes. The LSTM's gates mitigate this but do not eliminate it. Long-range dependencies—understanding that the pronoun at position 300 refers to a noun at position 12—remain genuinely difficult. Models tend to rely disproportionately on recent tokens, making them constitutionally bad at long contexts.

The field knew this was a problem. Bahdanau's attention mechanism in 2014 was already a workaround—it let the decoder, at each output step, look back at all encoder states simultaneously, creating a weighted summary. It was brilliant. But it was an add-on, a patch applied to the sequential model, not a replacement. The question that Vaswani et al. dared to ask in 2017 was: what if attention isn't a supplement to recurrence, but its entire replacement? What if recurrence itself is the disease?

Anatomy of the Paper That Started Everything

"Attention Is All You Need" was submitted to NeurIPS 2017. The title is a manifesto. It doesn't say "attention helps" or "attention augments"—it says attention is sufficient. The paper's eight authors—Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser, and Polosukhin—were Google Brain and Google Research researchers. Within the ML community, the paper spread like a slow fire that suddenly became an inferno. The model they proposed, the Transformer, achieved state-of-the-art on English-to-German and English-to-French translation, trained in a fraction of the time of the best RNN-based models. The BLEU scores weren't just better—they were achieved with dramatically fewer GPU-hours because the model parallelised across the entire sequence simultaneously.

The architecture, at its highest level, is an Encoder-Decoder model. The encoder reads the input sequence and produces a rich, contextualised representation of it. The decoder then uses that representation, along with its own previously generated output, to produce the next token autoregressively. But inside each encoder and decoder block, the mechanism is entirely attention-based. No convolutions. No recurrence. No LSTM gates. Just attention and feed-forward layers, stacked repeatedly. Let's dissect each component.

Embeddings & the Geometry of Meaning

Before any attention can happen, tokens need to exist in a form the model can reason about. Language—words, subwords, characters—is discrete. Neural networks are continuous. You need a bridge. That bridge is the embedding layer.

Every token in the vocabulary is assigned a learned vector in $\mathbb{R}^{d_{model}}$. In the original paper, $d_{model} = 512$. This means every word exists as a point in a 512-dimensional space. The remarkable thing that Mikolov et al. discovered with Word2Vec in 2013, and that remains true in Transformers, is that this space is not random. It is geometrically meaningful. The vector for "king" minus the vector for "man" plus the vector for "woman" is very close to the vector for "queen." Syntactic and semantic relationships become directions in this high-dimensional manifold.

In the Transformer, these embeddings are learned jointly with the rest of the model. The embedding matrix $W_E \in \mathbb{R}^{|V| \times d_{model}}$ is essentially a lookup table indexed by token ID. The paper also notes that the same weight matrix is shared between the encoder embedding, decoder embedding, and the pre-softmax linear transformation—a weight-tying trick that reduces the parameter count and, empirically, improves generalisation.

The model then scales these embeddings by $\sqrt{d_{model}}$ before adding positional encodings. The rationale: the embedding weights are initialised and trained to have magnitude roughly 1, but positional encodings have magnitude 1 by construction (they're sines and cosines). Scaling up the embeddings prevents the positional signal from dominating the content signal in the early stages of training.

Positional Encoding — Injecting the Sense of Order

Here is the first deep problem with abolishing recurrence: you lose the model's inherent sense of position. In an RNN, position is implicit—the fact that token $t_5$ is processed after $t_4$ is baked into the sequential structure. Remove that structure, and every token is equidistant from every other token in the model's perception. "The dog bit the man" and "The man bit the dog" would produce identical representations for every word, which is obviously catastrophic.

The Transformer's solution is elegant in its simplicity: add a positional signal directly to the token embedding. The question is: what signal? You could use learned positional embeddings—give every position a trainable vector. The paper actually tried this and found it worked about as well as the fixed approach. They chose the fixed sinusoidal encoding, and the reasoning reveals something deep about what they were optimising for.

The sinusoidal encoding uses sine and cosine functions at different frequencies—one function per dimension pair of the $d_{model}$-dimensional space:

Why sinusoids? Because relative position can be expressed as a linear function of absolute position under this scheme. Notice that for any fixed offset $k$, the encoding at position $pos + k$ can be expressed as a linear transformation of the encoding at position $pos$. This means the model can learn to attend by relative distance using simple linear attention patterns. A head can learn to "look two tokens to the left" by learning the appropriate linear transformation. With learned embeddings, you'd need separate representations for each absolute position, making generalisation to longer sequences at inference time impossible.

The interleaving of sine and cosine at different frequencies also means that each position gets a unique "fingerprint." The lowest-frequency pair changes so slowly across positions that it signals coarse-grained location (beginning, middle, end of sequence). The highest-frequency pair oscillates every two tokens, providing fine-grained positional discrimination. It's a binary counting system in continuous space.

Scaled Dot-Product Attention — The Engine Itself

We have arrived at the mechanism that defines the entire paradigm. Attention, in the abstract, is a function that takes a set of queries, a set of keys, and a set of values, and returns a weighted combination of the values where the weights are determined by the compatibility between queries and keys. If that sounds abstract, let's make it viscerally concrete.

Imagine you are at a library, and you have a research question in mind—your Query. Every book in the library has a spine with keywords—these are the Keys. The actual content of the book is the Value. You scan the shelves, and for each book, you compare your query against its keywords. Books with highly relevant keywords get more of your attention. You then synthesise what you find, weighted by that relevance. Self-attention is this process, where every token in a sequence simultaneously acts as a querying reader and a book on the shelf.

More precisely, for a sequence of $n$ tokens, each represented as a $d_{model}$-dimensional vector, we compute three projections using learned weight matrices:

Here $X \in \mathbb{R}^{n \times d_{model}}$ is the matrix of input token embeddings. The key insight is that these projections are learned. The model learns, through training, how to project each token into a query space (what to look for) and a key space (what to advertise). These are not the same projections, which is crucial—the grammar of what a token offers as a key may be entirely different from the grammar of what it seeks as a query.

The attention score between query $i$ and key $j$ is the dot product $q_i \cdot k_j$. If $q_i$ and $k_j$ point in similar directions in this learned space, their dot product is large, meaning the model has learned that token $i$ should attend strongly to token $j$. We then compute these scores for all pairs simultaneously:

The $\frac{1}{\sqrt{d_k}}$ scaling factor deserves serious attention. As $d_k$ grows large, the dot products grow in magnitude (their variance scales with $d_k$). Large dot products push the softmax into regions of extremely small gradients—the function saturates, and learning stalls. Dividing by $\sqrt{d_k}$ keeps the dot products in a regime where the softmax's gradient remains useful. This is not a trivial detail; it is the difference between training converging and not.

After scaling and softmax, we have an $n \times n$ matrix of attention weights—a map of how much every token should attend to every other token. This matrix is then used to compute a weighted sum of the value vectors. The output for token $i$ is a blended representation of all tokens' values, weighted by how relevant each token is to token $i$. The model is, in a very real sense, dynamically routing information across the sequence in a single parallel operation.

The Computational Cost — and Why It Matters

The $n \times n$ matrix tells you immediately what the fundamental bottleneck is. Attention has $O(n^2)$ complexity in both time and memory. For a sequence of 512 tokens, that's 262,144 attention scores per head. For 2,048 tokens—GPT-3's context—that's 4,194,304. For 1,000,000 tokens, which modern models are beginning to handle, that's a trillion entries. This is why the innovation around attention efficiency has become one of the most productive areas of ML research in the years since 2017. But we'll get to that.

# Scaled Dot-Product Attention — the mathematical core
import torch
import torch.nn.functional as F
import math

def scaled_dot_product_attention(
    query:  torch.Tensor,   # (batch, heads, seq_len, d_k)
    key:    torch.Tensor,   # (batch, heads, seq_len, d_k)
    value:  torch.Tensor,   # (batch, heads, seq_len, d_v)
    mask:   torch.Tensor = None  # (batch, 1, seq_len, seq_len)
) -> tuple[torch.Tensor, torch.Tensor]:

    d_k = query.size(-1)

    # Core operation: QKᵀ / √d_k
    # scores.shape → (batch, heads, seq_len, seq_len)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    # Causal masking: prevent attending to future tokens
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # Softmax over the key dimension → attention weights
    attn_weights = F.softmax(scores, dim=-1)

    # Weighted aggregation of values
    output = torch.matmul(attn_weights, value)

    return output, attn_weights

Multi-Head Attention — The Ensemble Within a Single Layer

A single attention head is powerful but limited. It has one set of $W^Q, W^K, W^V$ matrices and therefore one "lens" through which to view the relationships between tokens. The question the paper asks is: what if different types of relationships need different lenses? Syntactic dependencies (which noun is the subject of which verb?) look fundamentally different from semantic relationships (which adjective modifies which noun?) or coreference (which pronoun refers to which entity?). A single head cannot simultaneously specialise in all of these.

Multi-Head Attention is the solution. Instead of one large attention operation in $\mathbb{R}^{d_{model}}$, we run $h$ parallel attention operations, each in a lower-dimensional subspace $\mathbb{R}^{d_k}$ where $d_k = d_{model} / h$. In the original paper, $h = 8$ and $d_{model} = 512$, so $d_k = d_v = 64$. Each head has its own $W_i^Q, W_i^K, W_i^V$ matrices. They are trained simultaneously but independently, and they are free to learn entirely different projection strategies.

The outputs of all $h$ heads are concatenated, giving a vector of dimension $h \cdot d_v = d_{model}$, and then projected through a learned output matrix $W^O \in \mathbb{R}^{h \cdot d_v \times d_{model}}$. This final projection mixes information from all heads into a coherent unified representation.

Interpretability research has since shown that attention heads do indeed specialise. Some heads reliably attend to syntactically related words regardless of semantic content. Others track long-range semantic similarity. Some heads in decoder models appear to function as "induction heads"—pattern matchers that complete sequences by analogy. The mechanistic interpretability community has spent years probing these heads, and the picture that emerges is of a surprisingly modular computational structure emerging from the training dynamics, not from architectural constraints. The model discovers the utility of specialisation on its own.

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"

        self.d_model    = d_model
        self.num_heads  = num_heads
        self.d_k        = d_model // num_heads   # per-head dimension

        # Separate linear projections for Q, K, V and output
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

        self.dropout = nn.Dropout(p=dropout)
        self.scale   = math.sqrt(self.d_k)

    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
        # x: (B, T, d_model) → (B, heads, T, d_k)
        B, T, _ = x.shape
        x = x.view(B, T, self.num_heads, self.d_k)
        return x.transpose(1, 2)

    def forward(
        self,
        query:  torch.Tensor,
        key:    torch.Tensor,
        value:  torch.Tensor,
        mask:   torch.Tensor = None
    ) -> torch.Tensor:
        B = query.size(0)

        # 1. Project inputs to Q, K, V and split across heads
        Q = self._split_heads(self.W_q(query))   # (B, h, T, d_k)
        K = self._split_heads(self.W_k(key))
        V = self._split_heads(self.W_v(value))

        # 2. Scaled dot-product attention (per head)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        attn   = self.dropout(torch.softmax(scores, dim=-1))
        output = torch.matmul(attn, V)    # (B, h, T, d_k)

        # 3. Concatenate heads and project back to d_model
        output = output.transpose(1, 2).contiguous()  # (B, T, h, d_k)
        output = output.view(B, -1, self.d_model)        # (B, T, d_model)
        return self.W_o(output)                           # (B, T, d_model)

The Feed-Forward Network — and Why It Exists at All

After the multi-head attention sub-layer, every Transformer block applies a position-wise feed-forward network. This is applied identically and independently to each token's representation. It's two linear layers with a ReLU (or later, GELU) activation between them:

In the original paper, the inner dimension is 2048—four times $d_{model}$. This means the FFN first expands the representation into a higher-dimensional space, applies a non-linearity, and then projects back down. The question every student asks when they first encounter this is: why? We just computed all this contextualised attention. Why run each token through an independent MLP that doesn't even let tokens interact?

The answer, which has become clearer through mechanistic interpretability work, is that the FFN appears to be the model's memory store. Attention handles routing—deciding which information should flow where. The FFN handles computation and recall—applying learned transformations to the now-contextualised representations. Research by Geva et al. (2020) suggests that the FFN layers function somewhat like key-value memories: the first weight matrix's rows act as keys that are activated by specific input patterns, and the second matrix's columns are the corresponding value vectors that modify the output. When the model "knows" a fact, that fact tends to be stored in the FFN weights, not the attention matrices.

This also explains why scaling the FFN dimension is so important for model capability. The capacity to store factual associations scales roughly with the number of parameters in the FFN layers, which is $2 \times d_{model} \times d_{ff} \times \text{num\_layers}$. For GPT-3's 175 billion parameters, a large fraction lives in these FFN matrices—this is not wasted capacity, it is the model's knowledge base.

Residual Streams & Layer Normalisation — The Infrastructure of Deep Learning

Between every sub-layer—attention and FFN—the Transformer applies two operations that seem mundane but are, in reality, load-bearing infrastructure: a residual connection and layer normalisation. Without them, training a 6-layer Transformer would be extremely difficult. Training a 96-layer GPT-3 would be essentially impossible.

The residual connection is deceptively powerful. By adding the input $x$ directly to the sub-layer's output, we create a gradient superhighway. When backpropagating through many layers, the gradient can flow directly through the residuals without passing through any non-linearities, dramatically mitigating vanishing gradients. He et al. demonstrated this with ResNets in 2016 for vision, and it generalises completely. You can think of the residual stream as an "information bus" that persists throughout the entire depth of the network, with each layer reading from it and writing back to it.

Layer Normalisation, introduced by Ba et al. in 2016, normalises the activations across the feature dimension (as opposed to Batch Norm, which normalises across the batch dimension). For each token's representation vector $x \in \mathbb{R}^{d_{model}}$:

The learned scale $\gamma$ and shift $\beta$ allow the model to undo the normalisation if necessary. The key benefit is stability: it prevents the internal activations from growing unboundedly across layers, which would cause numerical instability and chaotic gradient behaviour during training. The original paper applies LayerNorm after the sub-layer (Post-LN). Most modern implementations use Pre-LN (normalising the input before the sub-layer), which has been found to train more stably without requiring learning rate warmup schedules.

The Complete Encoder-Decoder Stack

With these components defined, we can now describe the full architecture of the original Transformer. The encoder is a stack of $N=6$ identical layers. Each layer has two sub-layers: (1) multi-head self-attention and (2) the position-wise FFN, each wrapped with a residual connection and layer norm. The encoder processes the entire input sequence in parallel and produces a set of contextualised representations—one per input token—of dimension $d_{model}$.

The decoder is also a stack of $N=6$ layers, but with three sub-layers per layer. The first is masked multi-head self-attention: the decoder can attend to its own previous outputs, but with a causal mask preventing it from seeing future tokens (since at inference time, we don't have those future tokens yet). The second is cross-attention: the queries come from the decoder's own representations, but the keys and values come from the encoder's output. This is how the decoder "reads" the input—it attends to relevant parts of the encoded source. The third is the FFN, same as in the encoder.

The decoder's output at each position is passed through a final linear layer (with weights shared with the embedding matrix) and softmax to produce a probability distribution over the vocabulary. During training, we feed the ground-truth output tokens shifted by one position (teacher forcing). During inference, we generate tokens autoregressively—one token at a time, feeding each generated token back as input for the next step.

Training — Loss, Masking, and the Cost of Understanding

The Transformer is trained using the standard cross-entropy loss for sequence-to-sequence models. Given a target sequence of length $m$ and the model's predicted probability distribution at each step, the loss is:

The model is optimised using Adam with $\beta_1 = 0.9$, $\beta_2 = 0.98$, and $\epsilon = 10^{-9}$. The learning rate schedule follows a custom warmup-then-decay scheme:

This ramps the learning rate linearly for the first 4,000 steps, then decays it as the inverse square root of the step number. The warmup is critical: at the start of training, the parameters are random and the gradients are unreliable, so a large learning rate would cause chaotic, destructive updates. The warmup period gives the model time to settle into a useful regime before we start optimising aggressively.

The paper also uses three regularisation techniques. Dropout (rate 0.1) is applied to the output of each sub-layer before the residual addition, as well as to the attention weights themselves. Label smoothing (value 0.1) replaces the hard one-hot target distribution with a softer distribution that assigns probability $0.1 / (V-1)$ to all non-target tokens—this prevents the model from becoming overconfident and generally improves calibration. Residual dropout, as mentioned, helps regularise the residual stream.

The Great Bifurcation — BERT vs GPT and the Architecture Wars

The original Transformer was designed for translation: a specific task with a clear input-output structure. The field quickly realised that the architecture's power was far more general, and within a year, two entirely different research programmes emerged from it, each taking half of the architecture and doing something remarkable with it.

BERT: Taking the Encoder

In October 2018, Google released BERT—Bidirectional Encoder Representations from Transformers. Devlin et al. took only the encoder stack and trained it with two pre-training objectives. The first was Masked Language Modelling (MLM): randomly mask 15% of tokens in the input and train the model to predict the masked tokens from context. This is critically different from the causal (left-to-right) training of RNN-based models—BERT can attend to both left and right context simultaneously, making it bidirectional. The second was Next Sentence Prediction (NSP): train the model to predict whether two sentences are adjacent in the original document.

The result was a model that could be fine-tuned on downstream tasks with remarkable efficiency. BERT-Large had 340 million parameters and achieved state-of-the-art on eleven NLP benchmarks including GLUE, SQuAD, and MultiNLI. The key insight was that the encoder's bidirectional representations, pre-trained on massive unlabelled corpora, are so rich that even a linear classifier sitting on top of the frozen BERT representations can solve complex NLP tasks. Pre-training followed by fine-tuning became the dominant paradigm overnight.

GPT: Taking the Decoder

OpenAI's approach, which would prove to be the more generative—and arguably the more transformative—research direction, took only the decoder stack. GPT-1 (2018) demonstrated that a decoder trained with the simple causal language modelling objective (predict the next token) on large corpora could be fine-tuned for a wide range of downstream tasks. GPT-2 (2019) scaled this to 1.5 billion parameters and demonstrated, controversially, that at sufficient scale the model could perform tasks zero-shot—without any fine-tuning—simply by being prompted with a description of the task in natural language.

GPT-3 (2020) scaled to 175 billion parameters and demonstrated something genuinely astonishing: in-context learning. The model, seeing only a handful of examples in its context window, could perform novel tasks it had never been explicitly trained on. This was not retrieval from a memory—it was genuine generalisation happening in the forward pass. The model was, in some sense, running a tiny gradient descent in its attention patterns without updating any weights.

Scaling Laws & the Emergent Phenomena

In 2020, Kaplan et al. at OpenAI published one of the most consequential empirical papers in recent ML history: "Scaling Laws for Neural Language Models." The finding was startling in its regularity: model performance (measured as loss on a held-out test set) improves as a smooth power law of three variables—model parameter count $N$, training data size $D$, and compute budget $C$:

The critical implication is that these laws hold across multiple orders of magnitude, with no sign of hitting a wall. You don't need new architectures; you just need more parameters, more data, and more compute, and the model gets predictably better. This gave the field something it had rarely had: a roadmap. If you know your compute budget, you can calculate, before training, roughly how good your model will be.

Hoffmann et al.'s 2022 "Chinchilla" paper refined this, arguing that Kaplan's scaling laws had been suboptimally applied: models had been trained with too many parameters relative to their training data. The optimal policy is to scale parameters and tokens proportionally. A 70 billion parameter model trained on 1.4 trillion tokens (the Chinchilla optimal ratio) outperforms a 280 billion parameter model trained on the same compute but with fewer tokens. This recalibrated how the entire field approaches model training.

More philosophically interesting than the scaling laws themselves are the emergent capabilities that appear as models scale. These are abilities that simply do not exist at smaller scales and appear, somewhat suddenly, once a threshold is crossed. Multi-step arithmetic, chain-of-thought reasoning, code generation, word analogies in low-resource languages—these capabilities are not present in 1-billion-parameter models and are strongly present in 100-billion-parameter models. The discontinuity of emergence is one of the most actively debated questions in modern AI research: are these genuinely phase transitions in capability, or are they measurement artefacts that would look smooth with better metrics?

Modern Advances — FlashAttention, RoPE, MoE, and the Architecture's Future

The core Transformer architecture of 2017 remains recognisable in 2026, but the engineering and theoretical refinements layered on top of it are substantial. Let's survey the most consequential.

FlashAttention — Rethinking the Memory Hierarchy

Tri Dao et al.'s FlashAttention (2022, refined in 2023 and 2024) is not a change to the mathematical definition of attention—the output is identical. It is a reimplementation of attention that is aware of the GPU's memory hierarchy. Standard attention materialises the full $n \times n$ attention matrix in HBM (High-Bandwidth Memory), the large but slower memory pool on the GPU. For long sequences, this matrix is enormous, and the time spent reading and writing it to HBM dominates the computation time. FlashAttention uses a tiling strategy to compute attention in blocks that fit in SRAM (the small, fast on-chip memory), never materialising the full $n \times n$ matrix. The result is 2-4x faster attention with significantly lower memory usage, with no approximation. This is what made 100K+ token context windows practically feasible.

Rotary Position Embeddings (RoPE)

RoPE, introduced by Su et al. (2021) and adopted by LLaMA, PaLM, and most modern open models, encodes positional information not by adding a fixed vector to the token embedding, but by rotating the query and key vectors in complex space before computing their dot product. For a pair of positions $m$ and $n$, the inner product $\langle q_m, k_n \rangle$ depends only on their relative offset $m - n$, not their absolute positions. This gives RoPE the extrapolation properties of sinusoidal encodings while being learnable and generally achieving better empirical performance. Extended RoPE variants have further enabled long-context extrapolation far beyond the training sequence length.

Grouped Query Attention (GQA) & Multi-Query Attention (MQA)

During autoregressive inference, the model must cache the key and value matrices for all previous tokens—the KV cache. This grows linearly with sequence length and model size, and for large models with large batch sizes, it becomes a significant memory bottleneck. Multi-Query Attention (Shazeer, 2019) uses a single set of $K$ and $V$ heads shared across all query heads, dramatically reducing the KV cache. Grouped Query Attention (Ainslie et al., 2023) is a middle ground—groups of query heads share a single $K$/$V$ pair. Llama 2 and 3 use GQA; it provides most of the memory savings of MQA with less quality degradation.

Mixture of Experts (MoE)

MoE replaces the single FFN in each layer with a collection of $E$ expert FFNs, and a lightweight router network that selects the top-$k$ experts (typically $k=2$ or $k=4$) for each token. This allows the model to have a very large total parameter count (more knowledge capacity) while only activating a fraction of parameters per token (constant compute cost). Mixtral 8x7B, for example, has 47B total parameters but activates only ~13B per token. The sparsity means you get the knowledge of a massive model at the inference cost of a smaller one.

class MixtureOfExperts(nn.Module):
    def __init__(self, d_model, num_experts=8, top_k=2, d_ff=2048):
        super().__init__()
        self.num_experts = num_experts
        self.top_k       = top_k

        # A pool of expert FFNs
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.GELU(),
                nn.Linear(d_ff, d_model)
            ) for _ in range(num_experts)
        ])

        # Router: maps each token to a probability over experts
        self.router = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape

        # Router logits → select top-k experts per token
        logits       = self.router(x)                        # (B, T, E)
        weights, ids = torch.topk(logits.softmax(-1), self.top_k, dim=-1)
        weights      = weights / weights.sum(-1, keepdim=True)  # renormalise

        # Aggregate expert outputs (simplified; production uses scatter ops)
        output = torch.zeros_like(x)
        for k in range(self.top_k):
            for e_idx in range(self.num_experts):
                mask = (ids[..., k] == e_idx)              # tokens routed to expert e
                if mask.any():
                    expert_in  = x[mask]
                    expert_out = self.experts[e_idx](expert_in)
                    output[mask] += weights[mask, k:k+1] * expert_out

        return output

The World Model Hypothesis

As the field has scaled Transformers and probed their internals, a more fundamental question has emerged: what are these models actually learning? The naive answer is "statistical correlations in text." The more intriguing hypothesis—supported by models that can solve physics problems, prove theorems, and navigate simulated 3D environments—is that Transformers are learning compressed internal models of the domains they're trained on. When a large language model correctly solves a novel spatial reasoning problem it has never seen, it is not retrieving a memorised solution. Something more abstract is happening. The Transformer, in learning to model text, may be learning to model the reality that text describes.