SourceUnderstanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions
The authors freeze all GPT-2 XL weights and train only new input/output embedding layers on a nearest-neighbor prediction task. The frozen model achieves near-perfect accuracy on the 80 points whose labels are determined by the nearest-neighbor rule. By computing an 'nn-score' for each attention head, they identify 8-10 heads that attend to the label of the nearest-neighbor input, closely implementing the NN algorithm. A randomly initialized transformer of the same size cannot do this, indicating the mechanism was learned during pretraining on language data.