“Learning to Grok: Emergence of In-Context Learning and Skill Composition in Modular Arithmetic Tasks”, Tianyu He, Darshil Doshi, Aritra Das, Andrey Gromov2024-06-04 (Transformer, grokking (NN), meta-learning)⁠:

Large language models can solve tasks that were not present in the training set. This capability is believed to be due to in-context learning and skill composition.

In this work, we study the emergence of in-context learning and skill composition in a collection of modular arithmetic tasks. Specifically, we consider a finite collection of linear modular functions z = a x + b y [mod] p labeled by the vector (a, b) ∈ ℤ_p². We use some of these tasks for pre-training and the rest for out-of-distribution testing.

We empirically show that a GPT-2-style transformer exhibits a transition from in-distribution to out-of-distribution generalization as the number of pre-training tasks increases. We find that the smallest model capable of out-of-distribution generalization requires two transformer blocks, while for deeper models, the out-of-distribution generalization phase is transient, necessitating early stopping.

Finally, we perform an interpretability study of the pre-trained models, revealing the highly structured representations in both phases; and discuss the learnt algorithm.