With one attention layer, the model sees context for the first time. This post unfolds the one-layer attention-only model to the end and checks everything it can do.

A one-layer attention-only transformer

Figure 1. A one-layer attention-only transformer. Each head’s output is added to the embedded tokens, which are immediately unembedded into logits.

Path expansion

This model’s output is the sum of two kinds of road. One is the direct path from last post, the built-in autocomplete. The other is the head paths through the attention heads. As a formula, it expands exactly like this.

$$T = \underbrace{W_U W_E}_{\text{direct path}} + \sum_h \underbrace{A^h \otimes (W_U W_{OV}^h W_E)}_{\text{path via head } h}$$

The $\otimes$ in the formula denotes the coupling of two dimensions: the left side (the pattern) handles which position to fetch from, the right side (the matrix) handles how the fetched content transforms. Part 2’s QK/OV split, engraved directly into the formula.

Path expansion of the one-layer model

Figure 2. Expanding the per-layer product yields the sum of the direct path and the head paths. The attention pattern is computed by multiplying token pairs through both sides of $W_{QK}$, then taking a softmax with autoregressive masking.

Each head path is exactly the two tables from part 2: the QK table deciding where to look, and the OV table deciding what an attended token pushes into the output. Both are token-to-token tables, so they can be interpreted from the weights alone, without running the model.

The QK and OV circuits: two independent token-to-token paths

Figure 3. The OV circuit (gold, $W_U W_O W_V W_E$) runs from the source token to the logits; the QK circuit (magenta, $W_E^\top W_Q^\top W_K W_E$) scores source-destination token pairs. Each collapses to a single matrix.

Skip-trigrams

What kind of rule falls out when you actually read the tables? A concrete example. “keep” appeared earlier in the context, and the current token is “in.”

Step one, look it up in the QK table. The query of “in” has been trained to be strongly drawn to an earlier “keep.”

(keep, in)   large   <- when at "in", look at "keep"

Step two, look it up in the OV table. When “keep” is attended to, the score of “mind” has been trained to rise.

(keep -> mind)   large   <- having seen "keep", push "mind"

Chain the two tables and one rule is complete: “keep … in → mind”. A [source] somewhere earlier, the current [destination], and an [out] to push. Because the source and destination may sit far apart, the paper calls this three-token pattern a skip-trigram. Trained heads turn out to hold such rules in bulk, and that is the entire capability of a one-layer model: a skip-trigram table stacked on top of the autocomplete table.

Examples of skip-trigrams found in trained heads

Figure 4. Large entries of the expanded QK/OV matrices. Reading, for each source token, the destination tokens it attends from and the out tokens it boosts gives a skip-trigram table directly. The fourth row is the model learning LaTeX; the fifth, HTML escape sequences.

Not every head attends by content. Some heads are found to attend mostly by position, for example to the vicinity of the immediately preceding token. They look like supporting cast for now, but in the next post they become the induction head’s partner, under the name previous-token head.

Copying heads and eigenvalues

The most common skip-trigram by far is copying: rules like “Tom … T → Tom” that push a token seen earlier back out. This implements the statistic that words already in a document tend to recur, and the paper calls it a primitive form of in-context learning. The division of labor is familiar by now: the OV table is set up to raise the attended token’s own score, and the QK table only looks back from places where that token would be plausible next. Copy, but only where autocomplete permits.

There is an entertaining special case too. Tokenizers usually merge the leading space into a word, so a rare word appearing without a space splits into two tokens (" Ralph" is one token; “Ralph” becomes “R” + “alph”). Some heads partially specialize in this case: seeing the fragment “R”, they attend back to a possible " Ralph" earlier and predict “alph”. A very special case of a one-layer model mimicking the induction heads we meet next post.

Heads copying words that split into two tokens without a leading space

Figure 5. Entries handling copying for words that split without a leading space. With the full word (" Ralph") as source and the fragment (“R”) as destination, the head outputs the remaining piece (“alph”).

How do you tell a copying head at a glance? The intuition is the table’s diagonal. A rule like “see Tom, push Tom” means the diagonal cells of the OV table are large. Rather than inspecting the diagonal of a 50,000 x 50,000 table cell by cell, though, the paper measures this with one summary statistic: eigenvalues. An eigenvalue is the factor by which a matrix stretches vectors along a particular direction; in $Mv = \lambda v$, a positive $\lambda$ means the direction $v$ comes back scaled up, not flipped. So OV eigenvalues leaning positive means that feeding in a token, or a combination of tokens, raises those very tokens’ scores: the table is copying overall. By this measure, 10 of the 12 heads in the analyzed model show clear copying. These seemingly enormous matrices are in fact low-rank, of rank $d_{head}$ (64 or 128), which is what makes the computation feasible. One qualifier attaches: a copying matrix must have positive eigenvalues, but the converse does not hold, so positive eigenvalues are strong evidence of “copying on average,” not a dispositive proof. This summary statistic keeps being used in the later papers.

Heads by fraction of positive eigenvalues: 10 of 12 sit at the copying extreme

Figure 6. Heads histogrammed by the fraction of their eigenvalues that are positive. Ten of twelve pile up at the right end, the copying extreme.

The bugs the factored structure produces

Say one head learned both “keep … in → mind” and “keep … at → bay”. A head’s rules are always the product of two tables. In the QK table, the cells (keep, in) and (keep, at) are large; in the OV table, (keep → mind) and (keep → bay) are large. But the product structure has no cell that could tie the combinations together, “mind only after in, bay only after at.” So probability leaks onto the crossed combination “keep … in → bay” as well. It is like being able to write a three-variable function $f(a, b, c)$ only in the form $f_1(a, b) \cdot f_2(a, c)$: the interaction the three tokens make together cannot be expressed freely. The model’s errors are not random noise but logical consequences of its structure: a good example of the predictive power reverse engineering buys.

Skip-trigram bugs produced by the factored structure

Figure 7. Limited expressivity creates bugs that look strange from the outside. A head that learned the correct skip-trigrams cannot help also boosting the crossed combinations (highlighted).

Do we fully understand it?

The paper answers this question carefully. In one sense, yes. Path expansion has given every parameter meaning in context, and no algorithmic mystery remains. But this is understanding in the sense of understanding the weights of a giant linear regression, or knowing how to query a large database. With a 50,000-token vocabulary, a single expanded OV matrix has about 2.5 billion entries. Without further work on summarizing, no human can hold that in their head. In the paper’s phrase, the one-layer model was revealed to be a compressed Chinese room, and we are now left holding a giant pile of cards.

A few technical reservations attach as well. As always with linear models over correlated variables, two heads can substitute for each other’s roles, so a zero weight need not mean “doesn’t do it,” and there is no standard way to normalize QK scores across different queries. The paper’s self-assessment: the contract has been made readable, but not yet read in full.

Takeaways

  • A one-layer model expands completely into the direct path (autocomplete) plus head paths (skip-trigrams).
  • A skip-trigram is one QK-table cell chained to one OV-table cell, and the flagship case is copying, summarized by positive OV eigenvalues (strong evidence, not dispositive proof).
  • The very fact that rules are a product of two tables predicts the bug of leaking crossed combinations.
  • The algorithmic mystery is gone, but 2.5 billion entries cannot be held in a head without summarization. What remains is a summarization problem.

The next post is the climax of the paper. With two layers, heads compose with heads, and induction heads appear.


Source: the One-Layer Attention-Only Transformers section of A Mathematical Framework for Transformer Circuits. All figures in this post are from the original paper.