Ghost in the Kernel: In-Context Learning with Efficient Transformers via Domain Generalization

2026-07-01

AI Summary

This work addresses the quadratic computational and memory costs of conventional Softmax attention in in-context learning and the lack of theoretical understanding of feature mappings in existing linear Transformers. From the perspective of domain generalization, the authors model in-context learning as a mapping from context distributions to response functions and analyze the approximation and generalization capabilities of linear Transformers through a two-stage sampling framework. They establish the first theoretical foundation for linear Transformers in in-context learning, revealing a trade-off between data distribution and feature regularity, and derive dimension-independent generalization convergence rates. This study opens a new pathway for efficiently fine-tuning large language models, achieving strong in-context learning performance while preserving linear complexity.

Abstract

Transformer-based large models have demonstrated remarkable generalization abilities across different tasks by leveraging a context-aware attention module for in-context learning. With richer context, transformers adapt more effectively to the current use case without any parameter updates. However, the quadratic computational and memory complexity with respect to context length significantly slows data processing in softmax transformers. Linear transformers were proposed to address this issue by reducing the complexity to linear dependence on context length, but the design and understanding of the feature mapping in linear attention, from a theoretical viewpoint, remain unclear. In this paper, we investigate the approximation and generalization abilities of linear transformers under a two-staged sampling process from domain generalization. We show that linear transformers perform in-context learning as learning a mapping from context distributions to response functions. A dimension-independent convergence rate is obtained for our generalization analysis, which also exhibits the tradeoff between the regularities of data distributions and latent features. Guided by our theoretical framework, we propose a new perspective on activation and loss design for linearizing pretrained softmax large language models.

Problem

• linear transformers
• in-context learning
• domain generalization
• feature mapping
• generalization

Innovation

• linear transformers
• in-context learning
• domain generalization
• feature mapping
• computational efficiency