RESEARCH NOTESThe Next Word · 1

Anatomy of a Language Model: When Everything Comes Down to the Next Word

How do you find a book in a library with no librarian?

What if an AI assistant’s “memory” isn’t all stored in the model’s weights?

It may remember your name, your preferences, even what you talked about last week. Some of that information may be in the conversation itself; some may have been saved elsewhere and supplied to the model again. Keeping a conversation going and updating a model’s weights are two different things.

The starting point is ordinary inference with fixed weights. A saved record might arrive with the input, much like a briefing before a meeting: “This is Mr Zhang. You discussed renovations last time. He likes a minimalist style.” The model can use that briefing without changing its parameters. Further training, or a mechanism specifically designed to update memory, raises a different set of questions.

Imagine a brilliant person who carries a notebook everywhere. They know languages, can prove theorems, and write poems that move you. To pick up yesterday’s conversation, though, they may need to turn back a few pages.

Here is the stranger part. Many autoregressive language models generate text by repeatedly predicting the next token. We often say “the next word”, but a token can also be part of a word, a character, or a punctuation mark.

It is a little like tapping a cane along a street. At each step, the question is: where next?

How does a training objective like that lead to solving problems, debugging complicated code, or responding to someone’s anxiety at three in the morning?

I want to take you inside this library without a librarian: to see how these abilities connect to a seemingly simple objective, and what they might tell us about memory, meaning, and thought.

Chapter 1: The Curse of Forgetting↗

1.1 Why do sequence models forget?↗

Imagine reading a long novel with only a small notebook for summaries. By the later chapters, those summaries are what keep you connected to the beginning.

A unidirectional recurrent neural network, or RNN, works a little like this. It updates a hidden state one step at a time, carrying information forward. A detail from much earlier in the sequence may have to survive many transformations before it can affect the current decision.

Think of a river. Information travels downstream, from one point to the next. To work out whether a later “it” refers to an earlier “cat” or “dog”, the model may need to carry that connection a long way.

Long short-term memory networks, or LSTMs, add gates to this river. They learn what to retain and what to let go. They ease the difficulty of learning long-range dependencies, but do not make every long-sequence problem disappear.

“Forgetting” is a metaphor here. The model does not lose a fixed percentage of information on every page. The question is whether it can carry distant details effectively, given its representations and the conditions under which it was trained.

1.2 What if information could take a shorter route?↗

In 2017, Vaswani and colleagues put attention at the centre of sequence modelling in Attention Is All You Need.

If a piece of information has to pass through many intermediate states to reach another position, could we shorten that path?

Let positions in the sequence connect directly, instead of always passing their messages along the same river.

That is our way into the library.

Chapter 2: A Library Without a Librarian↗

2.1 From the river to the library↗

Forget the river for a moment. Let me try another story.

Imagine a vast library with no librarian and no fixed shelf numbers. The books float in the air.

You walk in looking for something about the French Revolution. In an ordinary library, you would consult the catalogue and head for a particular shelf. Here, you begin by calling out:

“Who knows about the French Revolution?”

Each book compares its index with your question. A book on Napoleon glows brightly. Another on Marie Antoinette lights up too. A quantum physics textbook gives off only a faint glow.

You gather that light into a cloud of information drawn from different pages.

The analogy gives us three roles in an attention calculation. A query expresses what is being sought; a key is used for matching; a value supplies the content to be gathered.

2.2 One attention calculation↗

Scaled dot-product attention can be written as:

Attention⁡(Q,K,V)=softmax⁡ ⁣(QKTdk)V\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\!\left(\frac{QK^{\mathsf T}}{\sqrt{d_k}}\right)V

Here, QQ, KK, and VV are matrices of queries, keys, and values, and dkd_k is the key dimension. QKTQK^{\mathsf T} gives the matching scores. Scaling and softmax turn them into weights, which are then used to combine the value vectors.

Back in the library: a book that is more relevant to the question contributes more of its content.

This shortens the route information must travel. There is an important boundary, though. In autoregressive text generation, a causal mask blocks future positions. The model can use the current and earlier positions; it cannot peek ahead at the answer.

2.3 A shorter path is not an unlimited memory↗

Recall the RNN’s river. For information on page 1 to affect page 100, it may have to pass through many state updates.

Within the positions it is allowed to attend to, self-attention can connect the two directly. Information can take a shortcut instead of passing through every intermediate state.

Chapter 3: The Transformer’s Five Departments↗

Before we look more closely, let’s draw a map of the library. For this part, it also helps to imagine a company:

Table 1. Five Transformer components and their explanatory analogies
Component Analogy Main role
Embedding layer Reception Map tokens to vectors
Positional encoding ID numbers Provide information about sequence positions
Multi-head attention Meetings across departments Combine information in context
Feed-forward network (FFN) Specialists Apply a nonlinear transformation at each position
Residual connections and normalisation Working procedures Support information flow and stable training

We talk a great deal about attention, but FFNs can also account for a large share of the parameters. A meeting determines what information needs to be exchanged; the departments still have work to do with it.

3.1 Reception: the embedding layer↗

Before entering the library, each token needs a pass: it is mapped to a list of numbers—a vector.

Imagine words as people. The embedding layer draws a portrait of each one, using many numerical features. Relationships between those portraits develop during training. Some words end up near one another; some differences can be represented as directions.

A famous analogy from early word-vector research is:

eking−eman+ewoman≈equeen\mathbf{e}_{\mathrm{king}}-\mathbf{e}_{\mathrm{man}}+\mathbf{e}_{\mathrm{woman}}\approx\mathbf{e}_{\mathrm{queen}}

This illustrates an approximate relationship in a particular kind of word-vector space. It does not mean that every language model literally performs this arithmetic.

The embedding dimension, dmodeld_{\mathrm{model}}, sets the width of the representation. Greater width changes the model’s capacity and its parameter and compute costs. It does not increase capability in a fixed proportion.

3.2 Check-in: positional encoding↗

The floating library has another problem. How does the model know what comes first?

“The cat chases the dog” and “the dog chases the cat” contain the same words, but mean different things. Attention needs access to positional information to take that order into account.

The original Transformer constructs positional encodings with sine and cosine functions. Think of a timestamp built from several frequencies. Another approach is to learn an embedding for each position during training.

Rotary position embeddings, or RoPE, bring position into rotations of queries and keys, allowing their inner products to depend on relative position.

I like the image this suggests: two people turning on a dance floor. What matters is not just how far each has turned, but the relation between them.

Using RoPE does not automatically make a model reliable at arbitrary sequence lengths. Extending the context still calls for specific design choices and tests.

3.3 The reading room: multi-head self-attention↗

You call out “the French Revolution”, and the relevant books respond. What happens if you ask several different questions at once?

One voice asks about politics. Another asks for a timeline. A third asks for the people involved.

Each gathers different material. Their results are then combined into a representation that draws on several perspectives.

That is the intuition behind multiple heads. Heads can learn different patterns of attention, though each need not correspond to a stable, neatly named human concept.

We will return to the trade-off between heads and depth in Chapter 4. For now, keep in mind that heads are not free: their dimensions, the total representation width, and compute costs all enter the picture.

3.4 The study: feed-forward networks↗

Attention is the meeting room, where information from different positions comes together. The feed-forward network, or FFN, transforms the resulting representations further.

In a common standard design, the FFN expands a vector to a higher dimension, applies a nonlinearity, then projects it back. The original Transformer’s base configuration goes from 512 dimensions to 2048 and back to 512. That fourfold expansion is a design choice, not the only possible answer.

Geva and colleagues offer a concrete way to think about FFNs in Transformer Feed-Forward Layers Are Key-Value Memories: some key vectors relate to textual patterns, while value vectors influence the distribution over output words.

This gives the “book storage” analogy some substance. It does not establish that most knowledge resides exclusively in FFNs. How knowledge is distributed and used also depends on cooperation between components.

A small intermediate dimension may constrain capacity; a large one costs more parameters and compute. Gated designs also change the number of matrices and their dimensions. Comparing expansion ratios alone cannot tell us which design works better.

3.5 The corridors: residual connections and normalisation↗

Stack attention and FFNs repeatedly, and both information and training signals must cross many layers. How can a layer change a representation without having to rebuild everything from scratch?

Residual connections provide a simple form:

y=x+F(x)y=x+F(x)

The layer learns a transformation that can be added to its input. This gives information and gradients a more direct route through the network, helping with the training of deep models.

Normalisation adjusts the scale of activations. LayerNorm uses the mean and variance; RMSNorm scales by the root mean square without subtracting the mean. Where these operations sit in the network varies between architectures.

Chapter 4: Scaling Up—and Deciding What to Scale↗

Now that we know the departments, where should we put a fixed compute budget?

4.1 Depth and width↗

A deeper model offers more stages of computation. A wider model offers different representational capacity. Choosing between them requires looking at the task, training, and total parameter count together.

Within the settings they studied, Kaplan and colleagues’ Scaling Laws for Neural Language Models found that model size, dataset size, and training compute had more pronounced effects on loss than changing the depth-to-width ratio alone.

That does not make depth useless, or imply that doubling width doubles capability. It suggests that a seemingly crucial architectural detail may not always be the main driver of performance. The experimental setting matters.

Self-attention shortens the path between distant positions. Depth can still support successive computations and transformations. Reaching distant information in one layer does not mean that all the reasoning has already been done.

4.2 Attention and FFNs↗

Apple’s One Wide Feedforward Is All You Need explores a different allocation of resources. In an encoder–decoder Transformer for machine translation, the researchers remove the FFNs from the decoder layers and share a single FFN across the encoder.

This cuts parameters at some cost to accuracy. They then widen the shared FFN and, in their experiments, improve both accuracy and latency relative to the original Transformer Big.

These are results for a particular architecture and task. They do not establish that every large language model should have only one FFN.

The question remains interesting: does the model need many repeated processing modules, or fewer modules configured more effectively?

4.3 Heads and layers↗

Leaner Transformers: More Heads, Less Depth tries trading additional attention heads for fewer layers.

One ViT-B configuration in the paper’s appendix uses 7 layers, 16 heads, and a narrower MLP. It reaches 80.4% top-1 accuracy on ImageNet-1k with about 40.1M parameters. The corresponding original configuration has 12 layers, 12 heads, 86.6M parameters, and 80.1% accuracy. This is a combined architectural change; its result cannot be attributed simply to adding heads.

The authors examine the conditioning of the attention block and study how additional heads relate to optimisation stability under specified theoretical assumptions and experimental conditions.

The useful idea is that other components may compensate for reduced depth. Whether that trade-off works still needs to be tested in the model and task we care about.

4.4 The trade-offs↗

Table 2. Potential benefits and trade-offs in allocating model capacity
Choice Potential benefit Cost or limitation to consider
More layers More stages of computation and representation change More sequential computation; potentially longer training and inference latency
Larger hidden dimension A wider representation space More parameters, memory, and compute
More attention heads Different routes for combining information; better optimisation in some settings At fixed total width, each head becomes narrower; benefits depend on the configuration
Larger FFN More capacity for nonlinear transformations at each position More compute; possible redundancy

That is also what makes resource allocation interesting under a limited budget. Sometimes there is more to learn from placing capacity carefully than from making every component bigger.

Epilogue: Both Sides of the Mirror↗

Let’s return to the library.

Details of yesterday’s conversation may sit in a register by the door. Connections between pages have been written into the model’s parameters through training. Both can help it answer, but they are different ways of “remembering”.

You walk in and ask a question. Relevant pages begin to glow, some brightly, some faintly. Their light gathers into a cloud of information that becomes an answer.

The answer may be wrong. It may bring together things whose connection has never been established. Sometimes, though, it offers an opening: a relation you had not thought to look for.

The library is also a mirror.

When we use models to study human language and cognition, we are asking more than what the machine does. Where does it resemble us, and where does it differ? Asking whether it “understands” also makes us say more clearly what we mean by understanding.

When its memory depends on external records, we turn back to human memory. When predicting text produces complex representations, we ask how prediction, representation, and thought relate to one another.

Perhaps part of the interest lies in the questions these models make us ask again.

Memory and forgetting. Understanding and performance. Symbols and meaning. What allows matter to think—whether it is organised into carbon-based neurons or silicon circuitry?

We walk through a maze of our own making, looking for exits and answers. Perhaps learning how to ask for directions is one thing we find there.

Attention is a beginning. What does a machine attend to? What do we discover when we attend to what it attends to?

Outside, it is getting dark. Somewhere in a data centre, a model is waiting for its next input. Whether it understands its answer still calls for clearer questions and evidence.

But when you ask, the relevant pages begin to glow.