Nathan Godey, Yoav Artzi
9 min
Modern Transformer language models are adept at copying information from their context, a skill that is vital given that a large percentage of training data consists of recurring tokens. Despite this, the standard architecture for next-token prediction—the LM head—is a single, low-rank linear projection. This setup creates a softmax bottleneck that not only limits the expressivity of the model but also destructively compresses the training signal during backpropagation. To address this, the authors introduce Copy-Paste Language Models (CPLMs), which augment the standard LM head with a lightweight, learned pointer mechanism.
CPLMs add a dedicated Copy-Paste (CP) head alongside the standard LM head. The CP head calculates a probability distribution over past token positions using query and key projections, allowing the model to "copy" an earlier token directly rather than generating it from scratch via the LM head. A gated mixture, controlled by a special <copy> token and a sink mechanism, allows the model to intelligently choose between the standard generation path and the copy path. Because the CP head provides an alternative, higher-rank path for gradients to reach the backbone, it mitigates the information loss usually associated with the LM head during the training process.
CPLMs consistently outperform standard vanilla models across a variety of scales (135M to 1.7B parameters) and backbone architectures, including full-attention, sliding-window, and hybrid models. Notably, sliding-window CPLMs are able to achieve performance levels comparable to full-attention vanilla models while utilizing significantly less memory for the KV cache. Furthermore, the authors demonstrate that CPLMs achieve faster convergence and offer improved zero-shot performance on standard benchmarks. They perform exceptionally well in long-context retrieval tasks, such as finding specific "needles" in long documents, where they outperform vanilla models by wide margins.
This research demonstrates that the inherent architecture of the final output layer in modern LLMs is a significant bottleneck for training efficiency and recall. By showing that an explicit, learned copy mechanism effectively widens the gradient channel, this work provides a scalable, low-overhead method to improve model performance and efficiency without increasing compute or memory footprint substantially.
Transformers are intuitively considered to be efficient at copying from their inputs. We actually show that the ability of LLMs to copy in context can be improved through explicit modeling. We propose Copy-Paste Language Models (CPLMs), which add a lightweight Copy-Paste head to the conventional language modeling (LM) head. CPLMs improve training efficiency, downstream performance, and offer new operating points in long-context and retrieval setups. We validate our approach by pretraining models from 135M to 1.7B parameters, showing improvements across all scales and backbones with negligible training and inference overheads. We explain the success of CPLMs by showing that they partially mitigate the LM head gradient bottleneck through an alternative gradient path.
Sam: A name might appear several times. Does the model have to select one occurrence as the correct answer?
Alex: No. Probability assigned to different occurrences of the same token is added together. The usual vocabulary prediction and the copying prediction are then mixed. A special copy option in the ordinary output layer controls how much weight copying receives.
Sam: What if it opens that route but nothing suitable is actually in the earlier text?
Alex: There's a fallback called the sink. The copying component can select it and return weight to the ordinary predictor. The whole system learns through the standard next-token objective, without labels telling it when or where to copy.
Sam: Pointer networks already selected input positions, and earlier models mixed pointing with vocabulary prediction. What is genuinely new here?
Alex: The authors explicitly say inventing copying isn't their contribution. Their claim is that jointly training this mechanism improves modern language-model pretraining across architectures. And their main explanation concerns the learning signal, not just an extra option at prediction time.
Sam: Let's separate performance from that explanation. How broad was the training comparison?
Alex: They trained several model sizes, topping out at one point seven billion parameters. Copying and ordinary models used identical training hyperparameters. They tested full attention, locally restricted attention, mixtures of those, and a recurrent architecture that compresses earlier information into a running state.
Sam: The locally restricted models gain access they didn't have before. Could that alone account for the improvements?
Alex: It can't explain all of them. Models with full attention also improved, despite already accessing every earlier position. The authors report better prediction loss across the tested sizes and architectures, with smaller but generally positive gains on standard zero-shot evaluations.
Sam: What is the clearest training-efficiency result? Better loss isn't necessarily less time spent training.
Alex: In the largest sliding-window comparison, the copying model reached the ordinary model's final loss twice as fast. Sliding-window attention restricts the backbone to recent text. Full-attention and mixed models also reached their comparison targets sooner, but their speedups were smaller.
Sam: And what does that added access cost? Keeping another representation of every earlier token still uses memory.
Alex: It does, and the reported memory comparison includes that copying cache. At the training context length, the sliding-window copying model matched the full-attention baseline using twelve times less memory. It stores a copying key for older tokens rather than keeping full attention records throughout the network.
Sam: So the gain comes from replacing expensive memory, not making memory disappear. What about computation?
Alex: In the largest trained configuration, measured training throughput dropped by about three percent. That's small compared with the reported reduction in training needed to reach the baseline loss. But the extra copying work isn't free.
Sam: Give me a concrete test where that cheaper memory actually retrieves something beyond the local window.
Alex: They hide a seven-digit magic number inside repeated neutral filler sentences, then ask for it. The sliding-window copying model retrieves it almost perfectly across the tested lengths and insertion depths. Its ordinary counterpart mostly succeeds only when the number remains inside its recent window.
Sam: That's clean retrieval, but fairly artificial. Does the benefit survive questions and generated answers?
Alex: They report average improvements on longer-context benchmarks and after supervised fine-tuning on instruction data. Math evaluations show consistent gains, and coding and instruction-following results generally improve too. They also test question answering with documents hidden from the backbone, accessible only through the copying memory.
Sam: That last setup sounds relevant to retrieval systems. Does hiding the documents destroy performance?
Alex: Not in their tested setup. The copying keys retain much of the visible-document accuracy while using substantially less cache memory. Documents are encoded independently, so those keys can also be reused across questions. That's a specific question-answering experiment, not a demonstration across every retrieval application.
Sam: Now the causal story. What makes the authors think better training signals matter more than simply adding a useful output predictor?
Alex: They stop the copying gradient from reaching the backbone, while still training the copying component itself. That version retains only about a quarter of the full improvement. So most of the benefit disappears when copying can read the backbone's representations but cannot shape them.
Sam: That implicates the backward route. Do they also test a setting where the ordinary output layer shouldn't have the proposed bottleneck?
Alex: They train models predicting bytes instead of their usual text tokens. The vocabulary is then smaller than the internal representation, removing the rank restriction at issue. Copying stays active, but brings no improvement. Together, these interventions support their explanation, although both use a small full-attention setup and a single seed.
Sam: What should listeners resist generalizing from these results?
Alex: Learning experiments stop at the model size we mentioned. Larger architectures appear only in timing measurements, not trained performance tests. Copying also doesn't repair every long-context failure: full-attention models still fail beyond certain lengths. And this route only helps predict tokens already present in the context.
Sam: Given those boundaries, who should read the full document, and where should they begin?
Alex: Pretraining researchers should start with the bottleneck analysis, then Appendix H's stopped-gradient and byte-level tests. Retrieval researchers should read Efficient Retrieval, especially the hidden-document comparison. For everyone else, remember: access to earlier text is not the same as an easy path to predicting it.
Sam: Thanks for listening.