Learning Less Is More: Premature Upper-Layer Attention Specialization Hurts Language Model Pretraining

2026-05-11

AI Summary

This work identifies and addresses a critical yet previously overlooked issue in the pretraining of GPT-style language models: premature specialization, wherein upper-layer attention heads rapidly converge to sharp, rigid patterns before lower-layer representations have stabilized, thereby degrading overall performance. The study demonstrates that this phenomenon is a key mechanism underlying pretraining inefficacy and proposes mitigating it by reducing the learning rates of query and key projections in upper layers during early training stages, thereby delaying attention from locking onto immature residual features. Through hierarchical learning rate scheduling, path decomposition, and ablation studies, the approach is validated on both LLaMA and GPT architectures, showing that adjusting only upper-layer learning rates significantly lowers perplexity and improves downstream task accuracy. Additionally, the work reveals that gated feed-forward networks inherently alleviate this issue by suppressing residual writes, uncovering a deeper interplay between architectural design and optimization dynamics.

Abstract

A causal-decoder block is hierarchical: lower layers build the residual basis that upper layers attend over. We identify a failure mode in GPT pretraining: upper layers commit to sharp attention patterns before lower-layer features stabilize. We call this premature upper-layer attention specialization. Temporarily slowing only upper-layer Q/K projections during early training improves final perplexity and downstream accuracy without altering other parameters; it prevents upper attention from collapsing onto an immature residual basis. In LLaMA-style blocks, the same intervention is nearly unnecessary. Through ablations, we isolate multiplicative gated FFNs (not RMSNorm or bias removal) as the component that suppresses the upstream residual writes driving the failure. A pathwise analysis unifies both findings: the learning-rate intervention reduces a step-size factor, while gated FFNs reduce a residual-energy factor on the same growth pathway. Our results identify upper-layer Q/K timing as a concrete interaction point between decoder architecture and optimization.

Problem

• premature attention specialization
• language model pretraining
• upper-layer attention
• residual basis instability
• decoder architecture

Innovation

• premature attention specialization
• layer-wise learning rate scheduling
• gated FFN
• residual basis stability
• decoder architecture optimization