The Initialization Determines Whether In-Context Learning Is Gradient Descent

TL;DR

This paper investigates how multi-head self-attention in large language models approximates gradient descent in in-context learning, introduces a trainable initial guess to improve performance, and validates these findings through theoretical analysis and experiments.

Contribution

It extends the understanding of in-context learning by analyzing multi-head self-attention with non-zero priors, introduces yq-LSA with trainable initial guesses, and demonstrates improved performance on regression and semantic tasks.

Findings

Multi-head LSA approximates GD under realistic conditions.

Introducing a trainable initial guess improves ICL performance.

Theoretical bounds on the number of attention heads needed.

Abstract

In-context learning (ICL) in large language models (LLMs) is a striking phenomenon, yet its underlying mechanisms remain only partially understood. Previous work connects linear self-attention (LSA) to gradient descent (GD), this connection has primarily been established under simplified conditions with zero-mean Gaussian priors and zero initialization for GD. However, subsequent studies have challenged this simplified view by highlighting its overly restrictive assumptions, demonstrating instead that under conditions such as multi-layer or nonlinear attention, self-attention performs optimization-like inference, akin to but distinct from GD. We investigate how multi-head LSA approximates GD under more realistic conditions specifically when incorporating non-zero Gaussian prior means in linear regression formulations of ICL. We first extend multi-head LSA embedding matrix by introducing an initial estimation of the query, referred to as the initial guess. We prove an upper bound on the number of heads needed for ICL linear regression setup. Our experiments confirm this result and further observe that a performance gap between one-step GD and multi-head LSA persists. To address this gap, we introduce yq-LSA, a simple generalization of single-head LSA with a trainable initial guess yq. We theoretically establish the capabilities of yq-LSA and provide experimental validation on linear regression tasks, thereby extending the theory that bridges ICL and GD. Finally, inspired by our findings in the case of linear regression, we consider widespread LLMs augmented with initial guess capabilities, and show that their performance is improved on a semantic similarity task.