Search papers, labs, and topics across Lattice.
This study investigates the impact of the language model (LM) head's projection on optimization efficiency by conducting a backward-only intervention that reduces the rank of the gradient sent into the Transformer while maintaining ordinary logits. The findings reveal that while reducing backward rank leads to an increase in validation loss, the loss increase is significantly smaller compared to using a factorized forward head, suggesting that the LM head's geometric compression does not necessarily constitute a harmful bottleneck. Ultimately, the results challenge the notion that the LM head's projection is detrimental to optimization, highlighting the complexity of its role in model performance.
Reducing the rank of gradients sent to the Transformer can increase validation loss, but surprisingly, it鈥檚 less detrimental than using a factorized forward head.
The language-model head maps a hidden state of width D to a vocabulary of size V, so its transpose can return at most D independent directions to the Transformer. Godey and Artzi argue that this severe projection is a harmful optimization bottleneck. We separate the geometry from the causal claim. Our backward-only intervention keeps the ordinary logits and the exact LM-head parameter update while reducing only the rank of the gradient sent into the Transformer. Across five paired seeds on byte-level and BPE-8192 WikiText-2 models, reducing backward rank increases validation loss. An equally ranked factorized forward head, however, increases loss substantially more. At half rank in the larger model, the backward-only loss increase is 0.0586 (95% CI [0.0167, 0.1005]), while the factorized forward head increases loss by 0.1795 ([0.1547, 0.2042]). The vocabulary-space residual also contributes to the ordinary LM-head update, and removing that contribution is harmful. Additional controls show that repeated-token failures are confounded by the number of independently sampled symbols, that adding never-target output classes does not impair learning, and that projection diagnostics do not reliably predict progress in our runs. Tested auxiliary feedback routes do not beat tuned backpropagation. These results confirm strong geometric compression but do not establish that it is a harmful optimization bottleneck.
