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

Large language models (LLMs) exhibit remarkable in-context learning capabilities, but standard softmax attention mechanisms suffer from quadratic computational and memory complexity, limiting their use in long-context scenarios. This paper investigates the theoretical foundations of linear transformers—which offer linear complexity—to understand how they achieve few-shot generalization without parameter updates.

The authors frame in-context learning as an operator learning problem within a domain generalization framework. They model the transformer's input as a two-staged sampling process: first, sampling a meta-distribution of tasks, and second, sampling context-augmented inputs from these tasks. By representing context information as a kernel embedding, they analyze the approximation and generalization abilities of linear transformers, specifically examining how they interact with latent feature spaces.

The study proves that linear transformers effectively learn a mapping from context distributions to response functions. A key result is the derivation of dimension-independent convergence rates for this learning process. The authors identify a fast spectral decay phenomenon in the attention weight matrices of pretrained LLMs, which they prove helps linear transformers mitigate the negative effects of distribution shifts. This theoretical framework provides a new perspective for designing activation functions and loss objectives to linearize pretrained softmax-based LLMs.

This work bridges the gap between the empirical success of efficient transformer variants (like Mamba or RetNet) and formal learning theory. By showing that linear transformers can mimic the behavior of standard softmax attention while maintaining computational efficiency, the paper provides a rigorous basis for scaling LLMs to handle massive context windows, such as entire codebases or long-form documents, without the prohibitive costs of quadratic attention.