Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Do Attention Heads Compete or Cooperate during Counting?

Pál Zsámboki, Ádám Fraknói, Máté Gedeon, András Kornai, Zsolt Zombori

TL;DR

This work uses mechanistic interpretability to analyze how a small, single-layer transformer solves a counting task. It shows that attention heads behave as a pseudo-ensemble for the counting semantics, but the output layer must aggregate their representations non-uniformly to satisfy the syntactic EOS constraint. Through three metrics (l-acc, ROC AUC, s-acc), attention patterns, and attention-intervention experiments, the authors demonstrate that the key factor is how heads attend to tokens (notably balancing $w_{01}$ and suppressing $w_{02}$) and that model performance largely aligns with a linear combination of head outputs. The findings illuminate the interplay between semantic and syntactic requirements in transformers and highlight the value of mechanistic interpretability for understanding head-level contributions in even simple algorithmic tasks.

Abstract

We present an in-depth mechanistic interpretability analysis of training small transformers on an elementary task, counting, which is a crucial deductive step in many algorithms. In particular, we investigate the collaboration/competition among the attention heads: we ask whether the attention heads behave as a pseudo-ensemble, all solving the same subtask, or they perform different subtasks, meaning that they can only solve the original task in conjunction. Our work presents evidence that on the semantics of the counting task, attention heads behave as a pseudo-ensemble, but their outputs need to be aggregated in a non-uniform manner in order to create an encoding that conforms to the syntax. Our source code will be available upon publication.

Do Attention Heads Compete or Cooperate during Counting?

TL;DR

This work uses mechanistic interpretability to analyze how a small, single-layer transformer solves a counting task. It shows that attention heads behave as a pseudo-ensemble for the counting semantics, but the output layer must aggregate their representations non-uniformly to satisfy the syntactic EOS constraint. Through three metrics (l-acc, ROC AUC, s-acc), attention patterns, and attention-intervention experiments, the authors demonstrate that the key factor is how heads attend to tokens (notably balancing and suppressing ) and that model performance largely aligns with a linear combination of head outputs. The findings illuminate the interplay between semantic and syntactic requirements in transformers and highlight the value of mechanistic interpretability for understanding head-level contributions in even simple algorithmic tasks.

Abstract

We present an in-depth mechanistic interpretability analysis of training small transformers on an elementary task, counting, which is a crucial deductive step in many algorithms. In particular, we investigate the collaboration/competition among the attention heads: we ask whether the attention heads behave as a pseudo-ensemble, all solving the same subtask, or they perform different subtasks, meaning that they can only solve the original task in conjunction. Our work presents evidence that on the semantics of the counting task, attention heads behave as a pseudo-ensemble, but their outputs need to be aggregated in a non-uniform manner in order to create an encoding that conforms to the syntax. Our source code will be available upon publication.

Paper Structure

This paper contains 22 sections, 10 equations, 15 figures, 1 table.

Table of Contents

  1. Introduction
  2. Related Work
  3. Problem Statement
  4. Understanding Attention
  5. A Minimal Solution
  6. Understanding Heads
  7. Three Metrics: l-acc, ROC AUC, and s-acc
  8. Understanding Individual Heads
  9. The Role of Attention
  10. Interventions
  11. Head Evolution During Training
  12. Interaction among Heads
  13. Head Cooperation
  14. The Syntactic Task: Predicting the '[EOS]' Token
  15. Discussion and Conclusion
  16. ...and 7 more sections

Figures (15)

  • Figure 1: Next token logit distributions for '4', '5', and '[EOS]', separated by current and true next tokens.
  • Figure 2: Average model accuracy on the Count01 task, along with standard deviation based on 10 random seeds.
  • Figure 3: Single and dual head (a) l-acc, (b) ROC AUC, and (c) s-acc on the test set.
  • Figure 4: Distribution of s-acc, ROC AUC, and l-acc values of all subsets of heads of a given size.
  • Figure 5: Each dot corresponds to one of the 16 trained heads. Color indicates linear separation accuracy. The x axis is $w_{01}$, the weight ratio of '0' and '1' tokens. The y axis is $w_{02}$ the weight ratio of '0' and '2' tokens.
  • ...and 10 more figures

Theorems & Definitions (1)

  • Definition 3.1: Count01 language