Learning Linear Regression with Low-Rank Tasks In-Context

Kaito Takanami, Takashi Takahashi, Yoshiyuki Kabashima

Abstract: In-context learning (ICL) is a key building block of modern large language models, yet its theoretical mechanisms remain poorly understood. It is particularly mysterious how ICL operates in real-world applications where tasks have a structure. In this work, we address this problem by analyzing a linear attention model trained on low-rank regression tasks. Within this setting, we precisely characterize the distribution of predictions and the generalization error in the high-dimensional limit. Moreover, we find that statistical fluctuations in finite pre-training data induce an implicit regularization. Finally, we identify a sharp phase transition of the generalization error governed by task structure. These results provide a framework for understanding how transformers learn to learn the task structure.

Code Dataset Promise: Yes

Code Dataset Url: https://github.com/taka255/icl-replica-analysis

Format Confirmation: I agree that I have read and followed the formatting instructions for the camera ready version.

Submission Number: 2100