An Adversarial Example for Direct Logit Attribution: Memory Management in GELU -4 L (2024)

Abstract: Prior work suggests that language models manage the limited bandwidth of the residual stream through a “memory management” mechanism, where certain attention heads and MLP layers clear residual stream directions set by earlier layers. Our study provides concrete evidence for this erasure phenomenon in a 4-layer transformer, identifying heads that consistently remove the output of earlier heads. We further demonstrate that direct logit attribution (DLA), a common technique for interpreting the output of intermediate transformer layers, can show misleading results by not accounting for erasure.

## introduction understanding the internal mechanisms of language models is an increasingly urgent scientific and practical challenge @xcite . for instance, we lack a clear explanation of the interaction between internal components, such as attention heads and mlps. @xcite referred to residual stream dimensions as memory or bandwidth that components use to communicate with each other. memory management @xcite observe that there are much more computational dimensions (such as neurons and attention head result dimensions) than residual stream dimensions, thus we should expect residual stream bandwidth to be in high demand. the authors speculated that some model components perform a memory management role, clearing residual stream dimensions set by earlier components to free some of this bandwidth. direct logit attribution (dla) is a technique for interpreting the output activations of model components in vocabulary space @xcite @xcite . in particular, dla applies the unembedding matrix to model internal activations, effectively skipping further computation of downstream components. this method implicitly assumes continuity of the residual stream, meaning a direction written to the stream stays conserved throughout the forward pass. however, the continuity assumption would not hold if some components erase residual directions set by earlier ones. overall, our main contributions are as follows: ## methods we characterize erasure as 3 steps during a forward pass of a model: (1) a writing component adds its output to the residual stream. (2) subsequent components read this information to perform their function. ## conclusion in this paper, we presented a concrete example of memory management in a 4-layer transformer model. it is important to note that our study focused on a single model and a specific attention head. further research is needed to determine the extent to which these phenomena generalize across different model components and model sizes. our findings also highlight the need for caution when using dla, as in the presence of the erasure phenomenon, these results can be misleading. to mitigate this, we advocate for testing effects across varied prompts, particularly those with different correct next token completions, as averaging over many prompts can cancel out spurious results. moreover, we recommend complementing dla with activation patching to measure both direct and indirect effects of model components.