nanoGPT Speedrun Techniques: How Delayed Untying Solves the Sparse Gradient Problem in Embedding Layers

Introduction: A Technical Puzzle from the nanoGPT Speedrun Community

In the deep learning community, nanoGPT training speedruns have been a topic of significant interest. nanoGPT is a minimalist GPT training codebase released by Andrej Karpathy in 2023, with core code of only about 300 lines, designed to help researchers and learners understand and reproduce GPT-2 level language model training from scratch. A speedrun community quickly formed around it, with participants competing to reach a specific validation loss in the shortest wall-clock time on NVIDIA A100 or H100 GPUs. This competition has spawned numerous engineering optimization techniques spanning from low-level CUDA kernel optimizations to high-level training strategy adjustments, making it a cutting-edge testing ground for deep learning engineering practices.

Researchers continuously push the limits of training high-quality language models in minimal time with minimal compute. A recent quiz question on Reddit about nanoGPT speedruns revealed a clever and profound engineering technique — one that touches on a long-debated core question in language model training: Should the input embedding and the output projection layer (lm_head) share weights or remain separate?

In a standard Transformer language model, the input embedding matrix has shape [vocab_size, d_model], mapping discrete token IDs to continuous vectors. The output projection layer (lm_head) has shape [d_model, vocab_size] — essentially the transpose of the embedding matrix — responsible for projecting the final hidden states back into vocabulary space to compute the probability distribution over the next token. The two are mathematically dual operations, and this symmetry is the theoretical foundation that makes weight tying possible.

The question itself is highly instructive, guiding us to think about how training dynamics change across different phases.

The Contradiction Between Sparse and Dense Gradients

To understand this question, we first need to clarify the phenomenon described in the prompt. In language model training, the input embedding matrix stores the vector representation of each token. The key issue lies in the distribution characteristics of gradients.

Early Training: Embedding Gradients Are Sparse

As the question states, during early training, embedding layer gradients are sparse: only tokens that actually appear in the current batch receive gradient updates. This means that low-frequency, rare tokens go long periods without effective gradient signals, causing their embedding vectors to update very slowly or even stagnate.

From an implementation perspective, in frameworks like PyTorch, the forward pass of an embedding layer (nn.Embedding) is essentially a lookup operation (index_select), and its backward pass produces gradients only on the indexed rows. For a GPT-2 vocabulary of size 50257, if only 2000 distinct tokens appear in a batch, then 96% of embedding rows have zero gradients per update. This stands in stark contrast to fully connected layers where every parameter participates in computation and receives gradients. The momentum estimates in the Adam optimizer also become biased for sparsely-updated tokens, causing their effective learning rate to be far lower than expected.

In contrast, the output projection layer (lm_head) carries much denser gradients — because it participates in computing logits over the entire vocabulary, and the softmax operation causes all token output weights to receive gradients. Specifically, when computing cross-entropy loss, the softmax function normalizes across all vocab_size logits. According to the softmax gradient formula, for the correct label position i, the gradient is p_i - 1; for all other positions j, the gradient is p_j. This means even if a token never appears as a target in the current batch, as long as its logit value is non-zero (i.e., softmax assigns non-zero probability), the corresponding weight column in lm_head will receive a gradient. This is the fundamental reason why output layer gradients are "dense" — every forward pass updates all output weights corresponding to the entire vocabulary.

Therefore, sharing lm_head's dense gradients is a more stable way to drive insufficiently trained token vectors to move.

Late Training: The Two Need Different Geometric Structures

However, an inherent contradiction exists here. As training progresses, the input embeddings and output logits actually require different representational geometric structures. Input embeddings need to encode semantic information about tokens for model comprehension, while the output layer needs discriminative structure that can correctly distinguish next-token probability distributions. Forcing both to permanently share the same matrix ultimately limits the model's expressive capacity.

This creates a classic trade-off: early sharing benefits stability, while late separation benefits expressiveness.

Technical Analysis of the Four Candidate Answers

Each of the four options corresponds to a real technical approach from either academia or engineering practice, and each deserves individual analysis.

Option A: Weight Tying

This is the most classic approach, proposed almost simultaneously and independently by Inan et al. (2016) and Press & Wolf (2017). Its theoretical motivation comes from an intuition: if two semantically similar words (like "car" and "automobile") should be close in input embedding space, then they should also be interchangeable candidates during output prediction, so the output weights should reflect this similarity as well. Empirical studies show that weight tying typically improves performance on small-to-medium models (like GPT-2 124M) and reduces embedding-related parameters by about 30%, but its benefits diminish at very large scales — partly because of the diverging geometric structure needs of input/output.

It keeps embed and lm_head sharing the same matrix throughout the entire training process. The advantage is significant parameter reduction and faster convergence; the disadvantage is exactly what was mentioned — it cannot satisfy the different geometric structure needs that emerge later. This is the baseline that the nanoGPT speedrun seeks to "surpass," not the technique it employs.

Option C: 75x Embedding Learning Rate

This option proposes keeping the matrices untied while dramatically amplifying the embedding layer's learning rate (e.g., 75x) to compensate for the sparse update problem. The logic is: since sparse tokens get fewer update opportunities, make them "travel further" when they do get updated. This is indeed a viable engineering approach, and similar differential learning rate strategies have appeared in some implementations, but it is not the correct answer here. An excessively high learning rate can also cause unstable oscillations in the embedding space, especially for rarely-appearing tokens where a single large step might push them to an unreasonable position.

Option D: Multi-token Prediction

This is a hot research direction from teams at Meta and elsewhere, providing richer training signals by predicting the next k tokens. In theory, this could give rare tokens more gradient opportunities, but it changes the training objective itself and represents a more macro-level architectural adjustment, operating on a different plane from the embed/lm_head coupling problem this question focuses on.

Option B: Delayed Untying — The Correct Answer

The correct answer is B: Delayed Untying. The core idea of this technique is: during the first 2/3 of training, tie the embed and lm_head weights together; upon reaching a certain checkpoint, copy the weights and optimizer states, then let both train independently as separate parameters.

This design elegantly leverages the advantages of both shared and separate modes:

• Early shared phase: Leverages lm_head's dense gradients to stably drive sparse embedding updates, ensuring all tokens (including rare ones) receive reasonable initialization and directional guidance;
• Late untied phase: After untying, the input embeddings and output projection layer each evolve toward the geometric structure best suited to their respective tasks, unlocking the model's expressive potential.

A notable detail: the question specifically emphasizes "copying optimizer states" — a point easily overlooked in implementation but critically important. If only weights are copied while optimizer momentum and other states are reset, it causes training instability at the moment of untying. Modern deep learning universally uses Adam or AdamW optimizers, which maintain two states for each parameter: first-order momentum (exponential moving average of gradients) and second-order momentum (exponential moving average of squared gradients). When delayed untying occurs, if only the weight matrix is copied while the new parameter's optimizer states start from zero, Adam goes through a bias correction phase where the effective learning rate slowly climbs from near zero, causing hundreds of steps of training efficiency loss. More seriously, the sudden disappearance of momentum is equivalent to "forgetting" the parameter's prior gradient history direction, potentially causing a sudden spike in training loss. Therefore, fully copying both first-order and second-order momentum is a necessary condition for ensuring a smooth transition, preserving optimizer states so that both independent trajectories can smoothly diverge from the shared point.

The Phased Training Philosophy Behind Delayed Untying

Behind the delayed untying technique lies an engineering philosophy of "phased training." It acknowledges that models have different needs at different stages of training, and therefore adopts dynamically adjusted strategies rather than fixed configurations from start to finish.

This approach is increasingly common in modern large model training. Phased training thinking has become a standard paradigm in the large model era. Typical examples include: progressive sequence length growth used by models like GPT-4 (training efficiently on short sequences first, then switching to long sequences later); dynamic batch size scheduling inspired by the Chinchilla paper (using small batches early for more update steps, then large batches later to reduce gradient noise); and the recently popular μP (maximal update parameterization) strategy of applying different learning rate scaling to different layers. Learning rate scheduling (warmup + decay), gradual batch size increases, curriculum learning, and similar techniques are all essentially applying different optimization strategies at different training phases. The common philosophy of these methods is: acknowledging that the geometric properties of the loss surface change dynamically during training, and fixed hyperparameter configurations cannot simultaneously satisfy the needs of all phases.

Delayed untying can be seen as a specific instance of applying this philosophy to parameter sharing structure.

For nanoGPT speedruns pursuing maximum efficiency, this kind of fine-grained optimization is precisely what shortens training time and improves final quality. The significance of speedruns lies not just in showing off skills, but in forcing researchers to deeply understand every detail of training dynamics, distilling reusable engineering wisdom.

Conclusion

This seemingly simple multiple-choice question actually connects several core deep learning concepts: weight tying, gradient sparsity, optimizer state management, and phased training. The delayed untying technique tells us that in language model training, there is no permanently optimal structure — only optimal strategies that adjust dynamically with training. For developers who wish to deeply understand model training mechanisms, the nanoGPT speedrun and its underlying collection of techniques is a treasure trove worth exploring repeatedly.

Key Takeaways