Towards Understanding How Transformers Learn In-context Through a Representation Learning Lens

TL;DR

This paper dives into the training process of this dual model in Transformers from a representation learning standpoint and derives a generalization error bound related to the quantity of demonstration tokens.

Abstract

Pre-trained large language models based on Transformers have demonstrated remarkable in-context learning (ICL) abilities. With just a few demonstration examples, the models can implement new tasks without any parameter updates. However, it is still an open question to understand the mechanism of ICL. In this paper, we attempt to explore the ICL process in Transformers through a lens of representation learning. Initially, leveraging kernel methods, we figure out a dual model for one softmax attention layer. The ICL inference process of the attention layer aligns with the training procedure of its dual model, generating token representation predictions that are equivalent to the dual model's test outputs. We delve into the training process of this dual model from a representation learning standpoint and further derive a generalization error bound related to the quantity of demonstration tokens. Subsequently, we extend our theoretical conclusions to more complicated scenarios, including one Transformer layer and multiple attention layers. Furthermore, drawing inspiration from existing representation learning methods especially contrastive learning, we propose potential modifications for the attention layer. Finally, experiments are designed to support our findings.

Figures (14)

• Figure 1: The ICL output ${\bm{h}}'_{N+1}$ of one softmax attention layer is equivalent to the test prediction $\hat{{\bm{y}}}_{test}$ of its trained dual model $f({\bm{x}}) = \widehat{{\bm{W}}}\phi({\bm{x}})$. The training data and test input can be obtained by linear transformations of demonstration and query tokens, respectively.
• Figure 2: Left Part: The representation learning process for the ICL inference by one attention layer. Remaining Part: Comparison of the ICL Representation Learning Process (Center Left), Contrastive Learning without Negative Samples (Center Right), and Contrastive Kernel Learning (Right).
• Figure 3: The equivalence between ICL of one softmax attention layer and gradient descent, along with analysis on different model modifications. Left Part:$\| \hat{{\bm{y}}}_{test} - {\bm{h}}'_{T+1} \|_{2}$ as the gradient descent proceeds under setting $N = 15$; Remaining Part: the performance for regularized models (Center Left), augmented models (Center Right) and negative models (Right) with different settings.
• Figure 4: The representation learning process for the ICL inference by one Transformer layer.
• Figure 5: Illustrating the ICL inference process of multiple softmax attention layers from the perspective of dual models. The layer-wise process of ICL can be viewed as a gradual gradient descent on the dual model sequence. The datasets used for each gradient descent, including training data and test input, are obtained from the outputs of the previous dual model before and after training.

Theorems & Definitions (14)

• Theorem 3.1
• Theorem 3.2
• Theorem A.1
• proof
• Theorem B.1
• proof
• Theorem B.2
• proof
• Theorem C.1
• proof