Interpreting Affine Recurrence Learning in GPT-style Transformers

TL;DR

This paper investigates how GPT-style transformers internally perform affine recurrence learning, revealing a layered copying and refinement mechanism that enhances understanding of in-context learning and recursive task processing.

Contribution

It introduces a mechanistic interpretability analysis of transformers trained on affine recurrences, uncovering a layered process of copying and refinement within the model.

Findings

Initial sequence estimates are made via copying in the zeroth layer.

Refinement occurs through negative similarity heads in the second layer.

Insights contribute to understanding recursive tasks in transformers.

Abstract

Understanding the internal mechanisms of GPT-style transformers, particularly their capacity to perform in-context learning (ICL), is critical for advancing AI alignment and interpretability. In-context learning allows transformers to generalize during inference without modifying their weights, yet the precise operations driving this capability remain largely opaque. This paper presents an investigation into the mechanistic interpretability of these transformers, focusing specifically on their ability to learn and predict affine recurrences as an ICL task. To address this, we trained a custom three-layer transformer to predict affine recurrences and analyzed the model's internal operations using both empirical and theoretical approaches. Our findings reveal that the model forms an initial estimate of the target sequence using a copying mechanism in the zeroth layer, which is subsequently refined through negative similarity heads in the second layer. These insights contribute to a deeper understanding of transformer behaviors in recursive tasks and offer potential avenues for improving AI alignment through mechanistic interpretability. Finally, we discuss the implications of our results for future work, including extensions to higher-dimensional recurrences and the exploration of polynomial sequences.