Transformer Key-Value Memories Are Nearly as Interpretable as Sparse Autoencoders

SAE

SAE decomposes and reconstructs the neuron activations, typically after the FF layer (residual stream). That is, let \(\bm{x}_\mathrm{FF_{out}} \in \mathbb{R}^{d_\mathrm{model}}\) be neuron activations after the FF layer, and \(d_{\mathrm{SAE}}\) denotes the dimension of SAE features; SAE module computes as follows:

with \(\bm{W}_{\mathrm{enc}} \in \mathbb{R}^{d_\mathrm{model} \times d_\mathrm{SAE}}\), \(\bm{W}_{\mathrm{dec}} \in \mathbb{R}^{d_\mathrm{SAE} \times d_\mathrm{model}}\), \(\bm{b}_{\mathrm{enc}} \in \mathbb{R}^{d_\mathrm{SAE}}\), and \(\bm{b}_{\mathrm{dec}} \in \mathbb{R}^{d_\mathrm{model}}\) in the SAE module. \(\text{ReLU}(\cdot): \mathbb{R}^d\to \mathbb{R}^d\) denotes an element-wise ReLU activation. Each activation dimension is treated as a potentially interpretable feature, and the matrix maps each feature dimension to its feature vector in the representation space. This module is trained so that the activations are as sparse as possible with a sparsity loss to disentangle the potentially polysemantic input neurons.

Transcoder

Notably, as an alternative to SAEs and perhaps the closest attempt to this study, Transcoders have recently been proposed. This approximates the original FF by training a sparse MLP as a proxy to predict FF output \(\bm{x}_\mathrm{FF_{out}}\) from FF input \(\bm{x}_\mathrm{FF_{in}}\), and its internal activations (\(\in \mathbb{R}^{d_\mathrm{TC}}\)) are evaluated in the same way as the standard SAE. Still, their work did not clearly evaluate how interpretable the original FF's internal activations are, and this work complements this overlooked question.

FF-KVs

Feed-forward layers in Transformers once project the FF input \(\bm{x}_\mathrm{FF_{in}} \in \mathbb{R}^{d_\mathrm{model}}\) to \(d_\mathrm{FF}\)-dimensional representation (\(d_\mathrm{model} < d_\mathrm{FF}\)), applies an element-wise non-linear activation \(\bm{\phi}(\cdot): \mathbb{R}^d\to \mathbb{R}^d\), such as SwiGLU, and projects it back, as follows:

where \(\bm{W}_K \in \mathbb{R}^{d_\mathrm{model} \times d_\mathrm{FF}}\) and \(\bm{W}_V \in \mathbb{R}^{d_\mathrm{FF} \times d_\mathrm{model}}\) are learnable weight matrices, and \(\bm{b}_K \in \mathbb{R}^{d_\mathrm{FF}}\), \(\bm{b}_V \in \mathbb{R}^{d_\mathrm{model}}\) are learnable biases. \(d_{\mathrm{FF}}\) is typically set as \(4d_{\mathrm{model}}\). One interpretation of the FF layer is a knowledge retrieval module; that is, the module first creates keys (activations) from an input \(\bm{x}_\mathrm{FF_{in}}\) and then aggregates their associated values (feature vectors).

TopK FF-KV

To encourage the alignment with SAE research, we also introduce sparsity to activations in FFs by applying a top-\(k\) activation function to the key vector. This keeps only the \(k\) neurons with the \(k\) largest activations in each inference, zeroing out the activation for the rest.

Normalized FF-KV

The feature vectors of SAE are typically normalized, whereas those in FF are not. If a particular feature vector \(\bm{W}_{V[i,:]}\) has a large norm, the magnitude of its corresponding activation may be underestimated. To handle this potential concern, we normalize each row of \(\bm{W}_V\), and the discounted vector norm is weighted to activations.

Here, \(\tilde{\bm{W}}_V = \text{diag}(\bm{s})^{-1} \bm{W}_V\), where \(\text{diag}(\cdot)\) expands a vector \(\mathbb{R}^d\) to a diagonal matrix \(\mathbb{R}^{d\times d}\).