How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Wu, Jingfeng, Zou, Difan, Chen, Zixiang, Braverman, Vladimir, Gu, Quanquan, Bartlett, Peter L.
Format: Preprint
Published: 2023
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866914715188330496
author Wu, Jingfeng
Zou, Difan
Chen, Zixiang
Braverman, Vladimir
Gu, Quanquan
Bartlett, Peter L.
author_facet Wu, Jingfeng
Zou, Difan
Chen, Zixiang
Braverman, Vladimir
Gu, Quanquan
Bartlett, Peter L.
contents Transformers pretrained on diverse tasks exhibit remarkable in-context learning (ICL) capabilities, enabling them to solve unseen tasks solely based on input contexts without adjusting model parameters. In this paper, we study ICL in one of its simplest setups: pretraining a linearly parameterized single-layer linear attention model for linear regression with a Gaussian prior. We establish a statistical task complexity bound for the attention model pretraining, showing that effective pretraining only requires a small number of independent tasks. Furthermore, we prove that the pretrained model closely matches the Bayes optimal algorithm, i.e., optimally tuned ridge regression, by achieving nearly Bayes optimal risk on unseen tasks under a fixed context length. These theoretical findings complement prior experimental research and shed light on the statistical foundations of ICL.
format Preprint
id arxiv_https___arxiv_org_abs_2310_08391
institution arXiv
publishDate 2023
record_format arxiv
spellingShingle How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?
Wu, Jingfeng
Zou, Difan
Chen, Zixiang
Braverman, Vladimir
Gu, Quanquan
Bartlett, Peter L.
Machine Learning
Transformers pretrained on diverse tasks exhibit remarkable in-context learning (ICL) capabilities, enabling them to solve unseen tasks solely based on input contexts without adjusting model parameters. In this paper, we study ICL in one of its simplest setups: pretraining a linearly parameterized single-layer linear attention model for linear regression with a Gaussian prior. We establish a statistical task complexity bound for the attention model pretraining, showing that effective pretraining only requires a small number of independent tasks. Furthermore, we prove that the pretrained model closely matches the Bayes optimal algorithm, i.e., optimally tuned ridge regression, by achieving nearly Bayes optimal risk on unseen tasks under a fixed context length. These theoretical findings complement prior experimental research and shed light on the statistical foundations of ICL.
title How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?
topic Machine Learning
url https://arxiv.org/abs/2310.08391