Decomposing Representation Space into Interpretable Subspaces with Unsupervised Learning

TL;DR

The method, neighbor distance minimization (NDM), learns non-basis-aligned subspaces in an unsupervised manner and offers a new perspective on understanding model internals and building circuits.

Abstract

Understanding internal representations of neural models is a core interest of mechanistic interpretability. Due to its large dimensionality, the representation space can encode various aspects about inputs. To what extent are different aspects organized and encoded in separate subspaces? Is it possible to find these ``natural'' subspaces in a purely unsupervised way? Somewhat surprisingly, we can indeed achieve this and find interpretable subspaces by a seemingly unrelated training objective. Our method, neighbor distance minimization (NDM), learns non-basis-aligned subspaces in an unsupervised manner. Qualitative analysis shows subspaces are interpretable in many cases, and encoded information in obtained subspaces tends to share the same abstract concept across different inputs, making such subspaces similar to ``variables'' used by the model. We also conduct quantitative experiments using known circuits in GPT-2; results show a strong connection between subspaces and circuit variables. We also provide evidence showing scalability to 2B models by finding separate subspaces mediating context and parametric knowledge routing. Viewed more broadly, our findings offer a new perspective on understanding model internals and building circuits.

Figures (35)

• Figure 1:
• Figure 2:
• Figure 3:
• Figure 5: Illustration of our subspace decomposition. On the right, the interpretation of the column of token "a" is obtained by inspecting Figure \ref{['fig:preimage-examples-main']}. Likewise, each row from bottom to top corresponds to Figure \ref{['fig:preimage-examples-space2']}, \ref{['fig:preimage-examples-space7']}, \ref{['fig:preimage-examples-space6']}, \ref{['fig:preimage-examples-space8']}, respectively.
• Figure 6: Preimages of the subspace activation corresponding to the same token "a" in 4 different subspaces of the residual stream $\boldsymbol{h}^{4,post}$ in GPT-2 Small. Tokens highlighted with red background are those where the activations are taken from. Tokens highlighted with dashed boxes reflect the encoded information. These activations are projected into the subspaces and the first column "sim" is the similarity between this projection in the preimage and in the query activation. The second column "pos" is the position indices of the highlighted tokens. We are showing only the most recent part of the context and as well as 5 future tokens, the actual input sequences are much longer. Interpretation of encoded information: (a) current token "a". (b) the prior token "under". (c) current position, around 118. (d) the topic, the overall context is about licence/permission. In sum, we can see in different subspaces of the same layer, different aspects of the current context are encoded.