Probing Layers › The Anatomy of a Transformer's Layers

Before diving into probing methodology, it helps to understand what each layer in a transformer is doing. A transformer encoder like BERT consists of a stack of $L$ identical blocks. Given an input sequence of $n$ tokens, each block turns a matrix of hidden states $\mathbf{H}^{(l)} \in \mathbb{R}^{n \times d}$ into a new matrix $\mathbf{H}^{(l+1)}$ of the same shape.

Each layer applies two sub-operations:

• Multi-head self-attention: Each token attends to all other tokens and gathers contextually weighted information.
• Position-wise feed-forward network (FFN): An MLP applied independently to each token's representation, often interpreted as a key-value memory store.

The transformation is:

$$\mathbf{H}^{(l+1)} = \text{LayerNorm}\!\left(\mathbf{H}^{(l)} + \text{FFN}\!\left(\text{LayerNorm}\!\left(\mathbf{H}^{(l)} + \text{Attn}\!\left(\mathbf{H}^{(l)}\right)\right)\right)\right)$$

where:

• $\mathbf{H}^{(l)}$: the hidden state matrix at layer $l$, shape $n \times d$
• $\text{Attn}$: the multi-head self-attention operation
• $\text{FFN}$: the feed-forward sub-network
• $\text{LayerNorm}$: layer normalization applied before each sub-operation (in the pre-norm variant)

Layer 0 is a special case: it is the embedding layer, which maps each token to a $d$-dimensional vector using a learned lookup table plus a positional encoding. This layer has no attention mechanism and no contextual mixing; each token's representation depends solely on its identity (and position), not its neighbors. Every subsequent layer introduces context by allowing tokens to communicate through attention.

The architecture's design has a direct consequence for probing: because each layer builds on the output of the previous one, representations accumulate information progressively. The embedding layer knows only the token identity. Layer 1 can see immediate neighbors through attention. By layer 6, a token's representation has been shaped by six rounds of contextual integration, with the attention patterns in each round potentially attending to very different parts of the sequence.

The residual connections deserve special mention here. In the formula above, each sub-operation adds its output back to the input: $\mathbf{H}^{(l)} + \text{Attn}(\mathbf{H}^{(l)})$. This means information from every preceding layer is directly available at every subsequent layer without needing to be re-encoded. A feature computed at layer 2 can flow through the residual stream untouched all the way to layer 12 if no subsequent layer overwrites it. This has important consequences for interpreting probing results, which we will return to in the limitations section.

The FFN sub-network, while often described as a simple two-layer MLP, has been given a more interesting interpretation in recent years. Geva et al. (2021) showed that FFN layers act as key-value memories: the first weight matrix stores patterns (keys) while the second stores factual content (values). When a key pattern matches the current token representation, the corresponding value is retrieved and added to the representation. This means upper-layer FFNs may be storing and retrieving factual knowledge (like "Paris is the capital of France"), while lower-layer FFNs may be storing and retrieving linguistic patterns (like "a token that looks like a past-tense verb ending in -ed"). Layer probing can detect both types of stored information, though it cannot easily distinguish memory retrieval from attention-based contextual integration.

This progression from context-free to context-rich representations is what layer probing exploits. By training probes at each layer, you can pinpoint the transition from "this token is the word bank" to "this token is the word bank used as a geographical feature."