Latest / Embodied AI 101 / One Learning Rate Doesn't Fit All: Layerwise Spectral Scheduling for Transformers

Transcript

One Learning Rate Doesn't Fit All: Layerwise Spectral Scheduling for Transformers · Embodied AI 101

• One Learning Rate Doesn’t Fit All: Layerwise Spectral Scheduling for Transformers. Transformer-based large language models have reshaped AI, but training them remains a costly endeavor. Conventionally we pick a single global learning rate (LR) schedule for all layers, adjusting it over time with warmup and decay. This “one LR fits all” strategy is simple, but it ignores a key fact: different Transformer layers often behave quite differently during training. In a recent work, Di He et al. (2026) show that paying attention to layerwise differences can pay off. They introduce Layerwise Learning Rate (LLR), a method that assigns each layer its own LR based on the spectral properties of its weights. Strikingly, this yields substantial speedups: up to 1.5× faster training on models like LLaMA and GPT variants, with even better final accuracy. In other words, by letting each layer learn at its own pace, we can train transformers faster and reach higher quality. In this episode we unpack why standard practice fails and how LLR fixes it. We’ll delve into the theory of heavy-tailed weight spectra for deep nets, explain how LLR measures each layer’s training state, and see how those measurements map to per-layer LRs. Then we’ll survey the experimental results: speedups and accuracy gains, across architectures and optimizers. Finally, we’ll place this work in context of related techniques (layerwise decay in fine-tuning, adaptive optimizers, etc.) and discuss how one might adopt LLR in practice. Why Might Transformers Learn at Different Rates?. Transformers are deep stack of layers, each consisting of attention and feedforward sub-layers, plus embeddings at input/output. Intuitively, not all layers are created equal. Early layers close to the input tend to learn very general features (embedding words into vector spaces, coarse grammar), while deeper layers capture more task-specific patterns. In many networks (especially vision or language models), practitioners often treat different layers differently. For example, in fine-tuning a pretrained BERT, it is common to use layerwise learning rate decay, giving the lower (earlier) layers smaller learning rates so they retain general linguistic knowledge, while the top layers adapt more aggressively. As a survey notes, “lower layers receive smaller learning rates to preserve general features, while upper layers and the classifier adapt more aggressively”. Beyond transfer learning, there is evidence that even during from-scratch training, layers progress unevenly. Gradient norms and weight updates in transformers can vary by depth. Some layers may converge or saturate faster, while others lag behind. From an optimization standpoint, a single LR that is good for one layer may be suboptimal for another. If the global LR is too small, we waste capacity of “slow-learning” layers; too large, and “fast-learning” layers might overshoot. Despite this, the standard practice in LLM training has been uniform LR schedules. The new work argues this is putting all layers in a “one-size-fits-all” straitjacket, possibly slowing down overall convergence. To make this concrete, Di He et al. turned to tools from deep learning theory: the idea of heavy-tailed weight spectra. Over the last few years, research in random matrix theory has found that well-trained neural network layers often have heavy-tailed (power-law) distributions of eigenvalues in their weight correlation matrices. In plain terms, if you look at the squared singular values of a layer’s weights, they often follow a power-law decay rather than a sharp cutoff. In such a regime, one can define a “tail exponent” α: roughly speaking, smaller α means the tail of large eigenvalues is fatter (heavier tail), indicating a more “self-regularized” or well-structured weight matrix. Early in training or in under-parameterized settings, the distribution might not show a heavy tail – it might look more like noise plus a cut-off. But as training progresses and the network learns rich features, layers tend to develop heavier tails. Mahoney & Martin (2019) famously analyzed lots of DNNs and found exactly this pattern: smaller or older networks had spectra resembling classical “Marchenko–Pastur plus spike” (like traditional ridge regularization), whereas modern deep nets showed truly heavy-tailed spectra – a phenomenon they called Heavy-Tailed Self-Regularization. In other words, each layer’s weight matrix has an empirical spectral density (ESD) whose tail can be fit by a power-law (p(\lambda)\propto\lambda^{-\alpha}). The exponent α is a way to quantify the layer’s “training progress”: large α (thin tail) implies the layer hasn’t fully developed its complexity, while small α (thick tail) implies more structure and specialization. This connection between spectral tails and training suggests a neat idea: if one layer’s tail exponent is high (not heavy-tailed), maybe that layer is “behind” and could use a larger learning rate to catch up. Conversely, a layer with an already heavy-tailed spectrum should perhaps train more cautiously with a smaller LR. That is exactly the intuition behind LLR. By continually measuring the spectral shape of each layer’s weight matrix during training, we get a proxy for how far along each layer is. LLR leverages this to assign per-layer learning rates in an adaptive way, rather than a single global LR. Heavy-Tailed Spectra: A Primer. Before diving into LLR itself, it’s worth pausing to explain the heavy-tailed spectrum idea a bit more. Consider a single linear weight matrix (W) from a trained layer. Its covariance or correlation matrix (X = W^T W) is positive semi-definite, so it has eigenvalues (\lambda_1,\dots,\lambda_n). The empirical spectral density (ESD) is simply a distribution that puts mass (1/n) at each eigenvalue. What researchers observed is that for many hidden-layer weight matrices in state-of-the-art networks, the upper tail of this ESD follows a power-law (p(\lambda)\propto\lambda^{-\alpha}). The exponent α often falls in a range like 2–5 for large nets; smaller α means a “heavier tail” (a slower decay of large eigenvalues). Why does this happen? Intuitively, having a few very large eigenvalues (heavy tail) means the layer has captured strong directions of variation from the data, while having many small eigenvalues means it is also capturing lots of smaller-scale patterns. It turns out that as networks train, they develop these heavy-tailed spectra spontaneously. This is thought to reflect an implicit form of self-regularization: instead of relying on explicit forms of regularization (like dropout or weight decay), the network’s weight matrices organize themselves so that their spectra have no sharp edge, reminiscent of disordered physical systems. Mahoney & Martin’s ICML19 paper describes how even without any extra regularization, trained layers exhibit either a “bulk-plus-spike” (traditional) spectrum or a fully heavy-tailed one. For our purposes, the key is: the tail exponent α encodes information about the layer’s training. If α is large, the tail is thinner, suggesting the layer’s weights haven’t become very structured yet. If α is small, the layer is doing something complex and is likely more “mature.” LLR uses this idea: it periodically computes an estimate of α for each layer and then uses those values to set that layer’s learning rate. Layers with higher α (less heavy tail) get scaled up in LR; those with lower α get scaled down. Practically, how does one estimate α? A common method is to fit a power-law to the largest eigenvalues. One approach (as used by LLR) is the Hill estimator. Sort the eigenvalues (\lambda_i) in ascending order. Let (k) be the number of largest ones we trust to follow a power law (for example, half of them). Then the Hill estimate of the tail exponent is [ \alpha_{\mathrm{Hill}} = 1 + \left(\frac{1}{k}\sum_{i=1}^k \ln\frac{\lambda_{n-i+1}}{\lambda_{n-k}}\right). ] This basically fits a theoretical tail to ({\lambda_{n-k+1},\ldots,\lambda_n}). LLR uses a similar calculation on each layer’s weight correlation matrix to get a per-layer α value. Once we have α for each layer, we map it to a learning rate. If (\alpha_l) is the exponent for layer (l), and (\alpha_{\min}, \alpha_{\max}) are the min/max across layers, LLR uses a simple linear mapping. In code, one sets [ f(l) = \eta \times \Bigl( \frac{\alpha_l - \alpha_{\min}}{\alpha_{\max}-\alpha_{\min}},(s-1) + 1 \Bigr), ] where (\eta) is a base LR and (s>1) is a scaling factor (the paper often uses (s=5)). In this scheme, the layer with (\alpha=\alpha_{\max}) (the largest exponent, least heavy-tailed) gets the highest LR (s\eta), and the layer with (\alpha=\alpha_{\min}) gets the base LR (\eta). In effect, layers deemed “under-trained” (weak heavy tails) are accelerated with up to 5× the base LR, while “mature” layers get the smaller base LR. There are a few important tweaks to make this work well in practice:. Embedding and final head layers: It turns out that the word embedding layer (and output head) often shows very high α values persistently. If we applied the rule blindly, those layers would always get extremely low LRs (near (\eta)), potentially starving them of learning. So LLR fixes this by capping the embedding/output layers at the maximum LR (i.e. treating them as though they have (\alpha=\alpha_{\max})). The intuition is that embeddings tend to start “undertrained” and need a strong push. In short, we override the scaling for those layers so they always use the top of the range. Soft switching: One-shot jumps in learning rates can destabilize training. Thus LLR doesn’t instantaneously teleport a layer’s LR from old to new; it transitions linearly over a small number of steps (a parameter (t_{\text{switch}})). In practice, once we compute new target LRs for each layer, we ramp each layer’s LR gradually over (t_{\text{switch}}) steps. This smoothing avoids any large gradient spikes. Efficient scheduling: Computing eigenvalues intermittently for every layer and fitting power-laws is costly. The authors observe that after an initial period (<∼20% of total training tokens), the α values stabilize for each layer. So they only perform these heavy computations early on. Specifically, LLR does periodic spectrum measurements during an “active phase” (the first 20% of training), updating each layer’s LR scheme as it goes. After that, the per-layer LRs remain fixed and follow the ordinary decay schedule. This makes the overhead slight compared to the full training time. Overall, LLR can be implemented on top of any standard optimizer (the paper shows results with AdamW and a new optimizer “Muon”) and any base LR schedule. Conceptually, you take your tuned global LR (\eta) and then allow LLR to stretch/shrink each layer’s LR around that anchor. One nice property is that if (\eta) was already well-tuned globally, LLR’s layerwise scaling keeps you in a good regime, so there’s little extra hyperparameter tuning needed. How LLR Fits into Training. Let’s walk through how LLR works during training:. Initialize with uniform LR: In the beginning we use a warmup or the usual schedule with a single LR for all layers (perhaps after a warmup period). Begin active phase: After some initial steps (once the model is in a dynamic regime), pause and compute the weight spectrum. For each layer (l), take its weight matrix (W_l) (or weight correlation (W_l^T W_l)), perform eigen-decomposition, and estimate the power-law exponent (\alpha_l) (e.g. via the Hill method using the largest half of eigenvalues). Compute layerwise LRs: Collect (\alpha_l) across layers to find (\alpha_{\min}) and (\alpha_{\max}). Then for each layer compute its target LR: larger LRs for larger (\alpha_l), smaller for smaller (\alpha_l). Clip or fix the embedding/output layers to the maximum LR. Smooth transition: Instead of changing immediately, slowly ramp each layer’s LR to its new target over (t_{\text{switch}}) steps (say a thousand steps). This is simply by incrementally adjusting each layer’s LR in the optimizer. Continue training: After reaching the new layerwise LRs, continue with the global LR schedule (often a cosine or linear decay of the base LR (\eta)). No more heavy spectral computations are done (due to the efficient-active-phase trick). Training proceeds until convergence or the token budget runs out. The effect is that in the first phase of training, layers that have underdeveloped weights get an extra speed boost, while layers that are already quite “learned” are slightly slowed. The goal is to keep all layers learning at a more balanced pace. Empirically, this helps avoid the situation where some layers finish learning quickly and others lag, which could otherwise waste training time or hamper final accuracy. It’s worth noting that LLR adds very little extra tuning. You still choose a base learning rate (\eta) and schedule as usual; LLR simply warps those into a band ([\eta, s\eta]) per layer. The main new hyperparameters are the scaling factor (s) (defaults to 5 in the paper) and how often to compute spectra. But the authors find that taking (s\approx5) works well across the board, and they only needed to compute spectra every few thousand steps for 20% of training. In practice, one can often “lift” the already-tuned (\eta) of the uniform-LR run and let LLR inherit those settings, with minimal extra search. Experimental Gains: Speed and Accuracy. Di He et al. tested LLR on a range of Transformer setups, spanning multiple model sizes and optimizers. Highlights include pre-training and fine-tuning LLaMA-style and GPT-nano models (from tens of millions to a few billion parameters) using both AdamW (the standard optimizer) and a new optimizer called Muon. The datasets and tasks were typical LLM benchmarks (language modeling and downstream zero-shot tasks). Across the board, LLR delivered strong benefits in both speed and performance. Most strikingly, LLR achieved up to ~1.5× faster convergence than the uniform-LR baseline. In practical terms, this means that LLR could reach the same loss or accuracy with only two-thirds the number of training steps (or tokens) as the standard recipe. In wall-clock terms, that’s a 50% time savings, which is huge for large models. This speedup was robust across scales (from 60M to 1B parameters in their tests) and even held up when moving to even larger runs. Notably, the speed gain came without any drop in final model quality – in fact, the models with LLR often ended up a bit better. For example, on a 1B-parameter LLaMA model, the standard uniform-LR training yielded an average zero-shot accuracy of around 47.1% on the chosen tasks. Using LLR, the model’s accuracy rose to roughly 49.0%, while using 33% fewer tokens. In other words, not only did it train faster, it also generalized a bit better. These gains of about 1–2 percentage points in accuracy (on top of strong baselines) were consistent across other model sizes. LLR also improved the model’s perplexity (a standard language modeling score) and balanced the heavy-tail exponents among layers. After training with LLR, the inter-layer variance of (\alpha) was noticeably smaller, suggesting a more uniform state of training across the network. The authors compared LLR against a host of alternative strategies as well. This included widely-used adaptive optimizers: LARS/LAMB (popular in large-batch training), Sharpness-Aware Minimization (SAM), and various recently proposed layerwise schedulers. LLR consistently outperformed these; for instance, no other method achieved as large a speedup on the LLaMA-1B runs. The authors even list custom schemes like “AlphaDecay” and “TempBalance” (concurrent heavy-tail based methods), and LLR still led. The takeaway is that a surprisingly simple heuristic – measure heavy tails and scale LRs – outshines many more intricate optimizers and schedules, at least in this setting. Importantly, the overhead of LLR was small. Because spectral measurements are only done early on, and only involve reasonably-sized matrices (layer widths, not the enormous attention tensors), the extra compute was negligible compared to full training. And once a good layerwise schedule is found, it can be reused: the authors note that LLR incurs “low tuning overhead” because it essentially inherits the global hyperparameters from the uniform baseline. So adopting LLR doesn’t mean hours of extra hyperparameter sweeps; typically one can just take the best LR from the usual grid search and let LLR do the rest. Finally, beyond from-scratch pre-training, the method proved applicable for fine-tuning tasks as well. The paper reports that even when fine-tuning on downstream tasks, allowing layerwise LR adjustments improved perplexity and sometimes accuracy over naive fine-tuning. This suggests that LLR’s insight — that different layers benefit from different training rates — holds whenever a transformer is being tuned. Its overhead during fine-tuning is even smaller (since we fine-tune far fewer steps), so it can be an easy win. Relation to Other Techniques. LLR sits at the intersection of a few trends in modern deep learning. On one hand it’s a learning-rate scheduling method, akin to things like warmup/cosine schedules or cyclical LRs. On the other hand it’s a form of differential learning rate across parameters. In transfer learning and fine-tuning, it’s common to slow the LR of early layers (often via an exponential decay across depth, called layer-wise LR decay) while using a higher LR for task-specific heads. LLR generalizes this idea: instead of a heuristic decay, it dynamically decides which layers are “low” or “high” LR based on their spectra. Some large-model training methods also adjust learning rates in a per-parameter way. For example, adaptive optimizers like Adam adjust effective step sizes individually based on gradient statistics. Others like LAMB rescale updates to account for layerwise norm differences. However, those methods do online adaptation based on gradients, whereas LLR operates on weights’ spectral structure. In practice, LLR can wrap around AdamW (as the authors showed) – it simply replaces the single learning rate with a per-layer rate schedule and leaves the underlying optimizer intact. Another class of related ideas involves “learning rate warmup” and sophisticated scheduling (like Adam’s bias correction) to stabilize deep net training. LLR is compatible with all of this: it still uses warmup and overall decay. Its soft-switch mechanism for LR transitions is reminiscent of smooth warmup itself. In fact, one can think of LLR as adding an extra warmup/cooldown per layer, guided by the network’s own state. One could also view LLR through the lens of balancing “layer capacity.” Some recent work argues that deeper layers have more sparsity or “idle capacity” early in training. LLR in effect pushes on those idle weights to activate them sooner. It also shares spirit with layer-specific regularization: a layer that’s undertrained relative to others is implicitly under-regularized, so LLR increases its training “force”. An immediate question is whether simply picking some fixed layerwise schedule without spectral feedback might do similarly. The answer seems to be no: the heavy-tail signals adapt to the data and model size, whereas any fixed assignment would have to be tuned. The paper’s experiments suggest that naive schemes (like a static decay schedule or one-hot schemes on the embedding only) give smaller or inconsistent gains. What LLR brings is a data-driven measure: it automatically discovers which layers need help. Practical Takeaways. For practitioners, the message is tantalizingly simple: building a quick spectral monitor into your training loop and scaling per-layer LRs can pay dividends with minimal fuss. The authors have provided code (see their GitHub) that hooks into standard training loops. In practice, you would take your existing config (say a 100M-token LLaMA run on AdamW with LR 1e-3), and add an LLR wrapper. During the first 10–20% of tokens, the code occasionally computes eigenvalues of each transformer weight matrix (this is only a few hundred dimensions at most, far cheaper than other model ops). It then rescales each layer’s LR within the range [base_LR, 5×base_LR], perhaps on a cosine schedule afterward. This costs maybe a few percent extra GPU time overall but returns a 30–50% reduction in steps needed to converge. The approach seems to be model-agnostic, relying only on the existence of linear layers to analyze. So it should carry over to vision transformers, speech models, any large network with repeated blocks. (Heavy-tail behavior has indeed been observed broadly.) The only caveat is that very small models or early in training the spectra might not be clear power laws; in those cases LLR just defaults to near-uniform LR. But once you’re in the scale of hundreds of millions of parameters, the effect reportedly kicks in. It’s also worth noting that LLR is orthogonal to many other training improvements. You can still use dropout, weight decay, Adafactor, quantization, etc. LLR mostly sits on top of the LR schedule. In particular, the authors show it working with a new optimizer (Muon) developed concurrently, which suggests even further gains if combined with the latest optim techniques. In summary, “one learning rate for all layers” is an oversimplification. Deep Transformer models are layer-heterogeneous beasts, and exploiting that fact can yield practical gains. The heavy-tail theory gives us a principled tool to quantify each layer’s state, and turning it into an adaptive LR scheme is both elegant and effective. As large-scale model training becomes even more expensive, tricks like LLR that squeeze extra efficiency out of the same compute budget are likely to become standard. Whether you’re training a multi-billion parameter LLM or fine-tuning a smaller transformer for robotics perception, consider checking if your layers truly learn at the same pace. If not, a little spectral analysis and some per-layer rate tuning might make your training lightning-fast, bridging the gap between “one learning rate” and “the right learning rate” for each layer. References: The ideas discussed here are drawn from [Di He et al., 2026] on heavy-tail guided layerwise learning rates for transformers, and the broader theory of heavy-tailed spectra in neural networks. The concept of layerwise learning rate decay in transfer learning is discussed in practical guides, highlighting that different layers often benefit from different LRs.