Transcoders

Introduction

In addition to SAEs (Sparse Autoencoders) and probes, transcoders are another approach for constructing interpretable features from a model. Their central idea is to replace the original Transformer MLP sublayer with a new representantion that is wider and sparser.

Let \(x\) denote an input text. The initial representation for any \(x \in D\) (where \(D\) is the dataset) is denoted as:

\[ x^{0, t} \in \mathbb{R}^{d_{model}} \]

Here, the superscript \(t\) indexes the token position.

Token Updates in a Transformer

Recall how token states are updated in a Transformer. If we rewrite this process schematically, we obtain:

• \(x_{pre}\) — token representation before the attention mechanism,
• \(x_{mid}\) — representation after attention,
• \(x_{post}\) — the final representation after one block.

In this notation, token updates through the Transformer layers proceed as follows:

1. Attention update:

\[ x_{mid}^{(l,t)} = x_{pre}^{(l,t)} + \sum_h attn^{(l,h)} \big(x_{pre}^{(l,t)} ; x_{pre}^{(l,1:t)}\big) \]

1. MLP update (residual connection):

\[ x_{pre}^{(l+1,t)} = x_{mid}^{(l,t)} + MLP^{(l)}(x_{mid}^{(l,t)}) \]

• Formula 1 reflects how the attention sublayer updates the hidden state of token \(t\).
• Formula 2 reflects how the MLP sublayer updates the hidden state of token \(t\).
• Note that in Formula 2, the output of the MLP is added back via a residual connection.

The combination of formulas (1) and (2) leads to a key theoretical result:
any hidden state of an arbitrary token \(t\) can be decomposed into the outputs of all preceding layers.

This result is directly related to the concept of the residual stream —
defined as the sum of outputs from all previous layers together with the input embedding.
Equivalently, it is the hidden state of a token at the current layer.

Source: A Mathematical Framework for Transformer Circuits, Anthropic

What Does a Transcoder Look Like?

At this point, we have all the tools to understand how a transcoder works.
As mentioned earlier, a transcoder is a trainable representation.
It consists of two sequential transformations:

\[ z_{TC}(x) = ReLU(W_{enc} x + b_{enc}) \]

\[ TC(x) = W_{dec} z_{TC}(x) + b_{dec} \]

Here \(W_{enc}, W_{dec}\) are trainable feature matrices, obtained by minimizing the following loss function:

\[ L_{TC}(x) = \|MLP(x) - TC(x)\|_2^2 + \lambda_1 \|z_{TC}(x)\|_1 \]

• The first term enforces accuracy of approximation to the original MLP.
• The second term enforces sparsity of activations.

Interpretable Features

As a result, each interpretable feature \(i\) is associated with two vectors:

• \(f^{i,l}_{enc}\) — the \(i\)-th row of the encoder matrix \(W_{enc}\),
• \(f^{l,i}_{dec}\) — the \(i\)-th column of the decoder matrix \(W_{dec}\).

In words:
a transcoder is a sparse approximation of the MLP.

Relation to Transformer Updates

Recall the embedding state update in the Transformer:

\[ x^{l+1, t}_{pre} = x^{l, t}_{mid} + MLP^l(x^{l, t}_{mid}) = x^{l, t}_{mid} + MLP^l\Big(x^{l, t}_{pre} + \sum_h attn^{l, h}(x^{l, t}_{pre}, x^{l, 1:t}_{pre})\Big) \]

The transcoder, during training, approximates the MLP part:

\[ MLP^l\Big(x^{l, t}_{pre} + \sum_h attn^{l, h}(x^{l, t}_{pre}, x^{l, 1:t}_{pre})\Big) \]

Thus, the embedding update can be rewritten approximately as:

\[ x^{l+1, t}_{pre} \approx x^{l, t}_{mid} + TC^l(x) \]

Circuits Between Layers

The goal of transcoders is not only to decompose the MLP into interpretable features,
but also to construct circuits between these features.

By definition, for two transcoders at layers \(l\) and \(l' > l\),
the contribution of feature \(i\) at layer \(l\) to the activation of feature \(i'\) at layer \(l'\) is:

\[ z_{TC}^{(l,i)}(x_{mid}^{(l,t)}) \cdot \big(f_{dec}^{(l,i)} \cdot f_{enc}^{(l',i')}\big) \]

where:

• \(f^{i, l'}_{enc}\) — the \(i\)-th row of the encoder matrix \(W_{enc}\) at layer \(l'\),
• \(f^{l,i}_{dec}\) — the \(i\)-th column of the decoder matrix \(W_{dec}\) at layer \(l\),
• \(z_{TC}^{(l,i)}(x_{mid}^{(l,t)})\) — a scalar reflecting the activation of feature \(i\) at layer \(l\) for token \(t\).

The product captures the interaction mechanism of features:
an activation at one layer is multiplied by the projection of the previous layer’s decoder onto the next layer’s encoder.

Why This Contribution Measure Arise

The embedding state update can be rewritten as:

\[ x^{l+1, t}_{pre} \approx x^{l, t}_{mid} + TC^l(x). \]

For describing feature \(i\), we have two vectors:

• \(f^{i, l}_{enc}\) — the \(i\)-th row of the encoder matrix \(W_{enc}\),
• \(f^{l, i}_{dec}\) — the \(i\)-th column of the decoder matrix \(W_{dec}\).

The activation of feature \(i\) in the encoder is given by:

\[ ReLU(f^{l, i}_{enc} \cdot x_{mid}^{l, t} + b_{enc}^{l,i}), \]

where the bias \(b_{enc}^{l,i}\) can be ignored.

First conclusion:
the main contribution to the activation of feature \(f^{l, i}_{enc}\) depends only on \(x^{l,t}\).

Because of residual connections, each \(x^{l, t}\) in turn can be decomposed into the sum of all previous components.

Since the following decompositions hold:

Attention block (first sub-block in a standard Transformer encoder):

\[ x_{mid}^{(l,t)} = x_{pre}^{(l,t)} + \sum_h attn^{(l,h)}(x_{pre}^{(l,t)} ; x_{pre}^{(l,1:t)}) \]

MLP block (second sub-block):

\[ x_{pre}^{(l+1,t)} = x_{mid}^{(l,t)} + MLP^{(l)}(x_{mid}^{(l,t)}) \]

Particular Case

For \(l=2\):

\[ x_{mid}^{2,t} = x_{pre}^{(2,t)} + \sum_h attn^{(2,h)}(x_{pre}^{(2,t)} ; x_{pre}^{(2,1:t)}) \]

but since

\[ x_{pre}^{2,t} = x_{mid}^{1,t} + MLP^1(x_{mid}^{1,t}), \]

we have

\[ x_{mid}^{1,t} + MLP^{(1)}(x_{mid}^{1,t}) + \sum_h attn^{(2,h)}(x_{pre}^{(2,t)} ; x_{pre}^{(2,1:t)}) \]

and since

\[ x_{mid}^{1,t} = x_{pre}^{1,t} + \sum_h attn^{(1,h)}(x_{pre}^{(1,t)} ; x_{pre}^{(1,1:t)}), \]

it follows:

\[ = x_{pre}^{1,t} + MLP^{(1)}\Big(x_{mid}^{(1,t)} + \sum_h attn^{(1,h)}(x_{pre}^{(1,t)} ; x_{pre}^{(1,1:t)})\Big) + \sum_h attn^{(2,h)}(x_{pre}^{(2,t)} ; x_{pre}^{(2,1:t)}) \]

\[ = x_{pre}^{1,t} + TC^1(x) + \sum_h attn^{(2,h)}(x_{pre}^{(2,t)} ; x_{pre}^{(2,1:t)}) \]

Interpretation

From this decomposition it is clear:
the influence of the previous MLP on feature \(i'\) in the next layer \(l' > l\) is expressed as

\[ f^{l', i'} \cdot MLP^l(x_{mid}^{l,t}) \;\;\approx\;\; f^{l', i'} \cdot TC(x_{mid}^{l,t}). \]

Thus, the contribution measure arises naturally from the residual decomposition of Transformer updates:
each feature’s influence in the next layer is captured through its projection via the transcoder. By definition, the transcoder at layer \(l\) can be decomposed into individual features:

\[ TC^l(x_{mid}^{l,t}) = \sum_{j \;\text{features}} z_{TC}^{l, j}(x_{mid}^{l,t}) \cdot f^{l, j}_{dec} + b^{l, j}_{dec} \]

As before, the bias term can be ignored.

Thus, from the earlier relation:

\[ f^{l', i'} \cdot MLP^l(x_{mid}^{l,t}) \;\;\approx\;\; f^{l', i'} \cdot TC(x_{mid}^{l,t}), \]

we obtain the interaction between features:

\[ z_{TC}^{(l, i)}(x_{mid}^{(l, t)}) \cdot \big(f_{dec}^{(l, i)} \cdot f_{enc}^{(l', i')}\big). \]

Here we use the encoder feature of the next layer (\(f_{enc}^{(l', i')}\)),
because this feature is the input component to the subsequent chain,
which enables feature-to-feature interactions across layers.

Adding the Contribution of Attention Heads

Up to this point, we have traced connections between MLP layers using transcoders.
However, since each new token representation also depends on the contribution of attention heads,
we must incorporate these contributions in order to build complete feature-level circuits.

How attention heads form the output vector in the attention mechanism (source)

Suppose we have a transcoder with feature \(i'\) at layer \(l'\). Let \(h\) denote an attention head in a previous layer \(l < l'\).

We want to answer:

What is the contribution of source token \(s\) at layer \(l\) in head \(h\) to the activation of feature \(i'\) at layer \(l'\) for target token \(t\)?

In other words:
how does a specific token in a specific attention head at the previous layer influence our fixed feature in the later layer?

Attention mechanism (source)

This can be expressed as:

\[ score^{(l,h)}(x^{(l,t)}_{pre}, x^{(l,s)}_{pre}) \cdot \Big(((W_{OV}^{(l,h)})^T f_{enc}^{l',i'}) \cdot x^{(l,s)}_{pre}\Big) \]

Why This Contribution Measure Arise

The output of layer \(l\) from the attention sub-block for token \(t\) is:

\[ \sum_h attn^{(l,h)}(x_{pre}^{(l,t)} ; x^{(l,1:t)}_{pre}) \]

Each head can be further decomposed into a sum over source tokens \(s\):

\[ attn^{(l,h)}(x_{pre}^{(l,t)} ; x^{(l,1:t)}_{pre}) = \sum_{s \;\text{tokens}} score^{(l,h)}(x_{pre}^{(l,t)}, x_{pre}^{(l,s)}) \; W_{OV}^{(l,h)} x_{pre}^{(l,s)} \]

Where does \(W_{OV}^{(l,h)}\) come from?

From the definition of the OV-circuit:

If the head output is

\[ z_h = softmax\Big(\frac{q \cdot k^T}{\sqrt{d_k}}\Big) v, \]

then the final output vector is:

\[ concat(z_0, z_1, \dots, z_H) \cdot W^O = \big((W^O)^T \cdot \begin{bmatrix} z_0 \\ z_1 \\ \vdots \\ z_H \end{bmatrix}\big)^T \]

Since \(W^O = [W^O_{h_0}, W^O_{h_1}, \dots, W^O_{h_H}]\),
we can equivalently express the contribution of head \(h\) as:

\[ z = \sum_h W^O_{h} z_h, \]

which corresponds to:

\[ W_{OV}^{(l,h)} x_{pre}^{(l,s)}. \]

Final Expression

The contribution of head \(h\) at layer \(l\) to feature \(i'\) at layer \(l'\) is:

\[ f^{i', l'}_{enc} \cdot attn^{(l,h)}(x_{pre}^{(l,t)} ; x^{(l,1:t)}_{pre}) = \sum_{s \;\text{tokens}} score^{(l,h)}(x_{pre}^{(l,t)}, x_{pre}^{(l,s)}) \cdot ((W_{OV}^{(l,h)})^T f_{enc}^{l',i'})^T x_{pre}^{(l,s)}. \]

Therefore, for each source token \(s\), the contribution to the activation of feature \(i'\) is:

\[ score^{(l,h)}(x_{pre}^{(l,t)}, x_{pre}^{(l,s)}) \cdot \big((W_{OV}^{(l,h)})^T f_{enc}^{l',i'}\big)^T x_{pre}^{(l,s)}. \]