Language Models “Grok” to Copy (2025)

Abstract: We examine the pre-training dynamics of language models, focusing on their ability to copy text from preceding context—a fundamental skill for various LLM applications, including in-context learning (ICL) and retrieval-augmented generation (RAG). We propose a novel perspective that Transformer-based language models develop copying abilities similarly to grokking, which refers to sudden generalization on test set long after the model fit to the training set. Our experiments yield three arguments: (1) The pre-training loss decreases rapidly, while the context copying ability of models initially lags and then abruptly saturates. (2) The speed of developing copying ability is independent of the number of tokens trained, similarly to how grokking speed is unaffected by dataset size as long as the data distribution is preserved. (3) Induction heads, the attention heads responsible for copying, form from shallow to deep layers during training, mirroring the development of circuits in deeper layers during grokking. We contend that the connection between grokking and context copying can provide valuable insights for more effective language model training, ultimately improving in-context performance. For example, we demonstrated that techniques that enhance grokking, such as regularization, either accelerate or enhance the development of context copying.

## introduction large language models (llms) can learn, retrieve, and reason from input context, facilitating various applications such as in-context learning @xcite and retrieval-augmented generation @xcite . despite these achievements, several shortcomings have been reported regarding llms' in-context capacities. for instance, the order of icl demonstrations matters @xcite and llms' awareness of different contextual positions fluctuates @xcite . we believe that studying the mechanisms behind the development of in-context capabilities during pre-training offers valuable insights for enhancing llms from a novel perspective. in this paper, we examine the pre-training dynamics of language models, focusing specifically on their context copying capabilities. these capabilities are crucial for various llm applications, including icl and rag. for example, @xcite interpret icl as a process that entails copying and then fuzzy pattern completion. similarly, rag exhibits this characteristic, as it requires the in-context retrieval of key information, which is then copied (or integrated with additional paraphrasing and reasoning) as the output. this paper presents empirical evidence demonstrating that transformer-based language models @xcite develop context copying capabilities in a manner akin to "grokking" @xcite . grokking refers to the abrupt improvement in test set generalization long after models have overfit. our experimental method is summarized as follows: we trained 12-layer llama models @xcite using 40 billion tokens and saved checkpoints at regular intervals. to evaluate context copying, we presented the models with an input context comprising multiple random token subsequences, each beginning with a unique prefix, and let them complete one of the prefixes presented in the context. the accuracy of these completions served as a measure of the models' context copying abilities. by analyzing the evolution of context copying accuracy and the development of circuits (i.e., the subnetworks responsible for completing the specific task) across the saved checkpoints, we argue there is a potential connection between grokking and the development of context copying capabilities, as outlined in the following arguments: argument 1: grokked context copying. we observe that context copying accuracy shows a sudden increase long after the training loss stabilizes, akin to "grokking" on the test set when neural networks trained on small training sets. ## general setup model architecture and hyper-parameters. we train small llama models @xcite on a subset of the redpajama dataset (computer, 2023), comprising 40 billion tokens, with the task 94071d6780d7ba717e\n 5015906d28e186d2bb\n …… 8bc19970dd5b8e7350\n 94071d6780d7 ## random token subsequences the prefix of the 1st subsequence of next-token prediction. our model has 162m parameters (12 layers, each with 12 attention heads; the hidden state dimension is 768, and the intermediate dimension of mlp layers is 3,072.) the context length is 1,024 tokens. we use the llama tokenizer with a vocabulary of 32,000 tokens. unless otherwise specified, the following hyperparameters are used: the adamw optimizer @xcite with (β 1 , β 2 ) = (0.9, 0.999), a learning rate of 0.1, 2000 warmup steps, and the norm clip value of 1. our training is conducted on 8 a100 gpus, with a batch size of 64 per gpu. ## application viewing the development of context copying as a special grokking inspires us to examine the impact of regularization, as it enhances grokking @xcite . we train models using (1) 10% attention dropout and (2) weight decay (λ = 0.1). figure 7 shows that their positive impact: with dropout, the model groks to copy earlier; both techniques improve the accuracy compared to the vanilla model. ## discussions we sincerely appreciate the anonymous reviewers for their valuable feedback. in this section, we address key points raised in their reviews, which may also be of interest to a broader audience. 1. our motivation for using copying tasks to measure in-context ability. induction heads, the key components responsible for in-context learning, are known to perform "copy and paste," as described by @xcite . in essence, induction heads "complete the pattern" by copying and extending sequences that have occurred previously. this behavior motivates our exploration of copying, which are foundational to understanding in-context abilities. moreover, the copying task employed in this study has proven effective in previous research on rag @xcite and in-context abilities (chen et al., 2024). to demonstrate that copying on the training data has reached saturation, we measured the "icl score" proposed by @xcite , which tracks the development of in-context abilities. our results show that after approximately 4,000 training steps (about 1.85 billion tokens), the icl score stabilizes at -0.5 nats. since testing accuracy continues to improve well after this saturation point, we infer that once copying accuracy is "grokked," reductions in training loss primarily stem from improvements in other abilities, rather than further progress in in-context copying. ## conclusions this paper introduces a novel perspective that the development of context copying is a special grokking. it holds the potential to provide meaningful insights that can be applied to language models, as we did in section 4. we hope a better understanding of grokking in future works provide more insights for developing stronger language models.