Pretrain-Test Task Alignment Governs Generalization in In-Context Learning
In-context learning (ICL) is a central capability of Transformer models, but the structures in data that enable its emergence and govern its robustness remain poorly understood. In this work, we study how the structure of pretraining tasks governs generalization in ICL. Using a solvable model for ICL of linear regression by linear attention, we derive an exact expression for ICL generalization error in high dimensions under arbitrary pretraining-testing task covariance mismatch. This leads to a new alignment measure that quantifies how much information about the pretraining task distribution is useful for inference at test time. We show that this measure directly predicts ICL performance not only in the solvable model but also in nonlinear Transformers. Our analysis further reveals a tradeoff between specialization and generalization in ICL: depending on task distribution alignment, increasing pretraining task diversity can either improve or harm test performance. Together, these results identify train-test task alignment as a key determinant of generalization in ICL.
Code (0)
등록된 구현이 없습니다.
Similar Papers 제목 키워드 기반
An Effective Gram Matrix Characterizes Generalization in Deep Networks
We derive a differential equation that governs the evolution of the generalization gap when a deep network is trained by gradient descent. This differential equation is controlled by two quantities, a contraction factor …
image-classificationImage ClassificationSmall Vocabularies, Big Gains: Pretraining and Tokenization in Time Series Models
Tokenization and transfer learning are two critical components in building state of the art time series foundation models for forecasting. In this work, we systematically study the effect of tokenizer design, specificall…
Representation LearningTransfer LearningIs Large-Scale Pretraining the Secret to Good Domain Generalization?
Multi-Source Domain Generalization (DG) is the task of training on multiple source domains and achieving high classification performance on unseen target domains. Recent methods combine robust features from web-scale pre…
Domain GeneralizationDomain Generalization using Pretrained Models without Fine-tuning
Fine-tuning pretrained models is a common practice in domain generalization (DG) tasks. However, fine-tuning is usually computationally expensive due to the ever-growing size of pretrained models. More importantly, it ma…
Domain GeneralizationEnsemble LearningEnhancing Compositional Generalization via Compositional Feature Alignment
Real-world applications of machine learning models often confront data distribution shifts, wherein discrepancies exist between the training and test data distributions. In the common multi-domain multi-class setup, as t…