How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression

TL;DR

This paper investigates how trained transformers use multi-head attention in in-context learning for sparse linear regression, revealing layer-specific patterns and providing theoretical explanations for their mechanisms.

Contribution

It offers a comprehensive analysis of multi-head attention utilization in trained transformers, combining experimental observations with theoretical insights for sparse linear regression.

Findings

Multi-heads are essential in the first layer but often only one is used in subsequent layers.

The first layer preprocesses data, and later layers perform simple optimization steps.

Preprocess-then-optimize approach outperforms naive gradient descent and ridge regression.

Abstract

Despite the remarkable success of transformer-based models in various real-world tasks, their underlying mechanisms remain poorly understood. Recent studies have suggested that transformers can implement gradient descent as an in-context learner for linear regression problems and have developed various theoretical analyses accordingly. However, these works mostly focus on the expressive power of transformers by designing specific parameter constructions, lacking a comprehensive understanding of their inherent working mechanisms post-training. In this study, we consider a sparse linear regression problem and investigate how a trained multi-head transformer performs in-context learning. We experimentally discover that the utilization of multi-heads exhibits different patterns across layers: multiple heads are utilized and essential in the first layer, while usually only a single head is sufficient for subsequent layers. We provide a theoretical explanation for this observation: the first layer preprocesses the context data, and the following layers execute simple optimization steps based on the preprocessed context. Moreover, we demonstrate that such a preprocess-then-optimize algorithm can significantly outperform naive gradient descent and ridge regression algorithms. Further experimental results support our explanations. Our findings offer insights into the benefits of multi-head attention and contribute to understanding the more intricate mechanisms hidden within trained transformers.