Fine-tuning Learning Rates › Layer-wise Learning Rate Decay
Language AI HandbookAll chapters Fine-tuning Learning Rates › Layer-wise Learning Rate DecayOne section · Sections in this chapter4 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Layer-wise learning rate decay (LLRD) is the most common implementation of discriminative fine-tuning for transformer models. It was popularized by the ULMFiT paper and has become standard practice for fine-tuning BERT and similar models. While the underlying mathematics are identical to discriminative fine-tuning, LLRD specifically stresses the layer-by-layer structure of modern transformer architectures and has been refined through extensive empirical study across dozens of NLP benchmarks.
Think of LLRD as a zoning ordinance for your network. Historical districts near the foundation of the city (the lower layers) face strict rules about what can be changed: minor renovations only, no demolition. Commercial zones at the top of the network have much more permissive rules, letting rapid development and transformation. The zoning shows not arbitrary preference but the actual character and function of each zone in the overall system.
What makes LLRD particularly effective for transformers is that the layer-by-layer structure maps directly onto the parameter group structure required by modern optimizers. You define one parameter group per layer and assign each group its own learning rate. The AdamW optimizer, which is the standard choice for transformer fine-tuning, handles per-group learning rates natively, so the implementation adds very little complexity over a single-learning-rate setup. The engineering cost is low, but the empirical benefits, particularly on smaller datasets, can be substantial.
Why Layer Position Matters
LLRD works because transformer representations are hierarchical. Varying learning rates by layer is based on research into what each layer learns during pre-training. Research on probing tasks, where simple classifiers are trained on frozen layer representations, has systematically mapped out this hierarchy with remarkable precision:
- Embedding layers capture lexical semantics and basic word relationships. At this level, the model encodes what words mean in isolation and how they relate to other words with similar meanings. These representations are highly general and transfer well to virtually any downstream task. Damaging them would undermine the entire network.
- Lower transformer layers (1-4 in BERT) capture syntactic structure and local dependencies. These layers develop sensitivity to part-of-speech patterns, phrase boundaries, and local word order. They learn to distinguish nouns from verbs, identify subject-verb relationships, and recognize basic grammatical patterns. These structural properties are shared across all syntactically regular language tasks.
- Middle layers (5-8) develop richer semantic representations that integrate local syntactic information with broader contextual meaning. These layers begin to capture coreference, semantic roles, and the flow of meaning across sentences. They are more specialized than lower layers but still broadly useful.
- Upper layers (9-12) capture task-relevant features that are more specific to the pre-training objective. For a model trained with masked language modeling, these layers specialize in predicting missing words given context. This specialization is useful but also means these representations are most tightly coupled to the pre-training task and farthest from what any given fine-tuning task might need.
When fine-tuning for a downstream task, the upper layers typically need the most adaptation because they were optimized for the pre-training task rather than your specific objective. Lower layers, which capture more universal linguistic properties that any NLP task would benefit from, need less adjustment. Aggressive updates to lower layers risk disrupting the carefully learned syntactic and semantic foundations that the entire model relies upon.
Implementing LLRD for Transformers
For transformer models, we typically define parameter groups corresponding to:
- Embedding layer: Token embeddings and position embeddings
- Encoder/decoder layers: Each transformer block as a separate group
- Task head: The newly added classification or regression layer
!uv pip install transformers torch
import torch
import torch.nn as nn
import math
from transformers import AutoModel, BertConfig, BertModel
import matplotlib.pyplot as plt
import numpy as np
def create_llrd_parameter_groups(model, base_lr=2e-5, decay_factor=0.95, weight_decay=0.01):
"""
Create parameter groups with layer-wise learning rate decay.
Args:
model: HuggingFace transformer model
base_lr: Learning rate for the top layer
decay_factor: Multiplicative decay per layer
weight_decay: L2 regularization strength
Returns:
List of parameter group dictionaries for optimizer
"""
parameter_groups = []
# Get all encoder layers
encoder_layers = model.encoder.layer if hasattr(model, 'encoder') else model.transformer.h
num_layers = len(encoder_layers)
# Embedding parameters - lowest learning rate
embedding_params = []
for name, param in model.named_parameters():
if ('embedding' in name.lower() or 'embed' in name.lower()) and param.requires_grad:
embedding_params.append(param)
if embedding_params:
embedding_lr = base_lr * (decay_factor ** num_layers)
parameter_groups.append({
'params': embedding_params,
'lr': embedding_lr,
'weight_decay': weight_decay,
'group_name': 'embeddings'
})
# Layer-wise parameters
for layer_idx, layer in enumerate(encoder_layers):
layer_lr = base_lr * (decay_factor ** (num_layers - 1 - layer_idx))
layer_params = [p for p in layer.parameters() if p.requires_grad]
if layer_params:
parameter_groups.append({
'params': layer_params,
'lr': layer_lr,
'weight_decay': weight_decay,
'group_name': f'layer_{layer_idx}'
})
# Pooler and any remaining top-level parameters
remaining_params = []
assigned_params = set()
for group in parameter_groups:
for p in group['params']:
assigned_params.add(id(p))
for name, param in model.named_parameters():
if id(param) not in assigned_params and param.requires_grad:
remaining_params.append(param)
if remaining_params:
parameter_groups.append({
'params': remaining_params,
'lr': base_lr,
'weight_decay': weight_decay,
'group_name': 'top'
})
return parameter_groupsLet's examine the learning rate distribution across layers:
# Create a sample configuration to visualize
base_lr = 2e-5
decay_factor = 0.95
num_layers = 12
lr_distribution = {}
lr_distribution["embeddings"] = base_lr * (decay_factor**num_layers)
for layer_idx in range(num_layers):
lr = base_lr * (decay_factor ** (num_layers - 1 - layer_idx))
lr_distribution[f"layer_{layer_idx}"] = lr
lr_distribution["task_head"] = base_lr / decay_factor
# Prepare lists for visualization
layer_names = list(lr_distribution.keys())
learning_rates = list(lr_distribution.values())Learning rate distribution with LLRD: Layer Learning Rate ------------------------------ embeddings 1.08e-05 layer_0 1.14e-05 layer_1 1.20e-05 layer_2 1.26e-05 layer_3 1.33e-05 layer_4 1.40e-05 layer_5 1.47e-05 layer_6 1.55e-05 layer_7 1.63e-05 layer_8 1.71e-05 layer_9 1.81e-05 layer_10 1.90e-05 layer_11 2.00e-05 task_head 2.11e-05
Learning rates decay geometrically through the layers. The embeddings receive the lowest rate (). This keeps the model's basic vocabulary representations remain stable. This stability is important because the embedding layer defines the basic semantic space that all other layers build upon. If embeddings shift substantially during fine-tuning, the higher layers must relearn how to interpret these shifted representations, potentially causing cascading disruptions throughout the network.
import matplotlib.pyplot as plt
import numpy as np
use_book_style()
plt.rcParams["figure.figsize"] = (6.0, 4.0)
fig, ax = plt.subplots()
x_positions = np.arange(len(layer_names))
colors = plt.cm.viridis(np.linspace(0.2, 0.8, len(layer_names)))
bars = ax.bar(x_positions, [lr * 1e5 for lr in learning_rates], color=colors)
ax.set_xticks(x_positions)
ax.set_xticklabels(layer_names, rotation=45, ha="right")
ax.set_ylabel(r"Learning Rate ($\times 10^{-5}$)")
ax.set_xlabel("Model Component")
ax.set_title(
f"Layer-wise Learning Rate Decay (LLRD)\nBase LR = {base_lr:.0e}, Decay = {decay_factor}"
)
polish_axes(ax)
# Add value labels on bars
for bar, lr in zip(bars, learning_rates):
ax.text(
bar.get_x() + bar.get_width() / 2,
bar.get_height() + 0.02,
f"{lr:.1e}",
ha="center",
va="bottom",
fontsize=8,
rotation=45,
)This visualization shows the learning rate hierarchy. The embeddings and lower layers remain relatively stable to preserve basic language understanding, while the learning rate increases exponentially toward the task head to facilitate rapid adaptation to the new problem. The smooth progression ensures that no single layer boundary experiences a large discontinuity in adaptation rates, which could create optimization instabilities where one layer is changing rapidly relative to its neighbors.
Choosing the Decay Factor
The decay factor controls how aggressively learning rates differ between layers, and choosing the right value requires balancing two competing concerns. A decay factor too close to 1.0 gives insufficient differentiation between layers, losing the benefits of discriminative fine-tuning. A decay factor too small may prevent lower layers from adapting at all, which can be problematic if your task requires the model to learn new low-level patterns not well-represented in the pre-training data. The optimal value depends on your task and how much the pre-trained representations need to change:
- : No decay; all layers receive the same learning rate (standard fine-tuning). This is equivalent to ignoring the hierarchical structure of the network and may be appropriate when you have abundant data and the task differs substantially from pre-training.
- : Gentle decay; good default for most NLP tasks. This creates roughly a 2:1 ratio between the top and bottom layer learning rates. This gives differentiation while still allowing all layers to adapt.
- : Moderate decay; useful when lower layers should change minimally. This creates approximately a 3:1 ratio and is appropriate when you have limited data or when the pre-training domain closely matches your target task.
- : Aggressive decay; for tasks very similar to pre-training. This can create ratios of 5:1 or greater and in effect "freezes" the lower layers by giving them very small learning rates, putting all adaptation in the upper layers.
Empirically, values between 0.9 and 0.95 work well for most transformer fine-tuning scenarios. When in doubt, 0.95 gives a safe starting point that rarely hurts performance and often helps. For domain adaptation tasks where the target text differs substantially from the pre-training corpus, a slightly larger factor (closer to 1.0) may be warranted because even the lower layers need to adjust to the new vocabulary distribution and stylistic patterns.
