Bigram Subnetworks: Mapping to Next Tokens in Transformer Language Models

  • 2025-04-26 00:34:11
  • Tyler A. Chang, Benjamin K. Bergen
  • 0

Abstract

In Transformer language models, activation vectors transform from currenttoken embeddings to next token predictions as they pass through the model. Toisolate a minimal form of this transformation, we identify language modelsubnetworks that make bigram predictions, naive next token predictions basedonly on the current token. We find that bigram subnetworks can be found infully trained language models up to 1B parameters, and these subnetworks arecritical for model performance even when they consist of less than 0.2% ofmodel parameters. Bigram subnetworks are concentrated in the first TransformerMLP layer, and they overlap significantly with subnetworks trained to optimallyprune a given model. Mechanistically, the bigram subnetworks often recreate apattern from the full models where the first layer induces a sharp change thataligns activations with next token predictions rather than current tokenrepresentations. Our results demonstrate that bigram subnetworks comprise aminimal subset of parameters that are both necessary and sufficient for basicnext token predictions in language models, and they help drive thetransformation from current to next token activations in the residual stream.These subnetworks can lay a foundation for studying more complex language modelcircuits by building up from a minimal circuit.

 

Quick Read (beta)

loading the full paper ...