Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

BERT Fine-tuning › Fine-tuning Hyperparameters

Michael BrenndoerferUpdated July 23, 20259 min read

Language AI HandbookAll chapters

Part 24 of 67BERT and Variants

Chapter 171BERT Fine-tuning › Fine-tuning Hyperparameters

Reading section 6 of 12 · Fine-tuning Hyperparameters · 9 minSection 6 of 12 · 9 min

Section 6 of 12 · 9 min read

BERT Fine-tuning › Fine-tuning HyperparametersOne section · Sections in this chapter6 / 12 ↓

Sections in this chapter12

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

Fine-tuning is sensitive to hyperparameter choices. Unlike pre-training, where you have billions of tokens, fine-tuning datasets are often small (thousands of examples), making overfitting a constant concern.

Learning RateLink Copied

The learning rate is the most critical hyperparameter. BERT's original paper recommends values in the range 2e-5 to 5e-5, much smaller than typical neural network training.

Why Small Learning Rates?

Pre-trained weights encode valuable knowledge. Large learning rates would rapidly overwrite this knowledge with task-specific but potentially less general patterns. Small learning rates allow gradual adaptation while preserving useful representations.

Out[34]:

Visualization

Line plot showing validation accuracy over epochs for three different learning rates, with 2e-5 performing best. — Effect of learning rate on fine-tuning. Very small rates (1e-6) learn too slowly. Very large rates (1e-3) cause instability. The sweet spot (2e-5 to 5e-5) balances learning speed with stability.

Batch Size and Training StepsLink Copied

BERT fine-tuning typically uses batch sizes of 16 or 32. Larger batches can work with learning rate scaling, but smaller datasets may not benefit.

Training duration is measured in epochs. Most tasks converge within 2-4 epochs. More epochs risk overfitting, especially on small datasets.

In[35]:

Code

def compute_training_steps(dataset_size, batch_size, num_epochs):
    """Calculate total training steps and recommended warmup."""
    steps_per_epoch = dataset_size // batch_size
    total_steps = steps_per_epoch * num_epochs
    warmup_steps = int(total_steps * 0.1)  # 10% warmup is common

    return {
        "steps_per_epoch": steps_per_epoch,
        "total_steps": total_steps,
        "warmup_steps": warmup_steps,
    }

In[36]:

Code

# Common dataset sizes
datasets = [
    ("Small (1K examples)", 1000),
    ("Medium (10K examples)", 10000),
    ("Large (100K examples)", 100000),
]

configs = []
for name, size in datasets:
    batch_size = 16 if size <= 10000 else 32
    epochs = 4 if size <= 10000 else 3
    stats = compute_training_steps(size, batch_size, epochs)
    configs.append((name, batch_size, epochs, stats["total_steps"]))

Out[37]:

Console

Training Configuration Recommendations
============================================================
Dataset                   Batch    Epochs   Total Steps    
------------------------------------------------------------
Small (1K examples)       16       4        248            
Medium (10K examples)     16       4        2500           
Large (100K examples)     32       3        9375

Smaller datasets benefit from smaller batch sizes (16) to see more gradient updates per epoch, and more epochs (4) to learn from limited data. Larger datasets can use bigger batches (32) for stability and fewer epochs (3) since each epoch provides more learning signal. The total steps column shows how training duration scales with dataset size.

Warmup and Learning Rate SchedulingLink Copied

Why not just use a constant learning rate? Two problems arise. First, in early training, the classifier head's weights are random, producing random gradients that might destabilize BERT's carefully tuned representations. Second, in late training, the model is close to convergence, and large updates can overshoot the optimum.

Learning rate scheduling addresses both problems by varying the learning rate throughout training. The standard BERT schedule has two phases:

Warmup: The learning rate starts at zero and increases linearly to its peak value. This gives the classifier time to produce meaningful gradients before larger updates hit BERT's layers.
Linear decay: The learning rate decreases linearly from its peak to zero. This allows fine-grained adjustments as the model converges.

The complete schedule can be expressed mathematically as:

\eta_t = \begin{cases} \eta_{\text{peak}} \cdot \frac{t}{t_{\text{warmup}}} & \text{if } t < t_{\text{warmup}} \\ \eta_{\text{peak}} \cdot \frac{T - t}{T - t_{\text{warmup}}} & \text{otherwise} \end{cases}

where:

$\eta_t$ : the learning rate at step $t$
$\eta_{\text{peak}}$ : the maximum learning rate (typically 2e-5 to 5e-5 for BERT)
$t_{\text{warmup}}$ : the number of warmup steps (typically 10% of total steps)
$T$ : the total number of training steps

Let's trace through the schedule. At step 0, $\eta_0 = 0$ : the model doesn't update at all. By step $t_{\text{warmup}}$ , the learning rate reaches its peak: $\eta_{t_{\text{warmup}}} = \eta_{\text{peak}}$ . Then it begins declining, reaching zero exactly at step $T$ . The denominator $T - t_{\text{warmup}}$ in the decay phase ensures the linear decrease is calibrated to hit zero at the final step.

Out[38]:

Visualization

Line plot showing learning rate starting at 0, rising to peak during warmup, then decaying linearly to 0. — Linear warmup followed by linear decay. The learning rate ramps up during the first 10% of training steps, then decreases linearly to zero.

Layer-wise Learning Rate DecayLink Copied

So far, we've treated all of BERT's parameters equally: every layer receives the same learning rate. But should they? Research on transfer learning suggests that different layers encode different kinds of information:

Lower layers (near the input) encode general linguistic features: part-of-speech patterns, syntactic structures, word relationships. These are useful across many tasks.
Upper layers (near the output) encode more abstract, task-specific features. In pre-training, these became tuned to masked language modeling and next sentence prediction.

For fine-tuning, this layered structure suggests a strategy: preserve the general features in lower layers by updating them slowly, while allowing upper layers to adapt more aggressively to the new task. We can implement this by giving each layer its own learning rate, with lower layers receiving smaller rates.

In[39]:

Code

def get_layerwise_lr_params(model, base_lr, decay_rate=0.9):
    """
    Create parameter groups with layer-wise learning rate decay.
    Lower layers get smaller learning rates.
    """
    param_groups = []
    num_layers = len(model.bert.encoder.layer)

    # Embeddings: lowest learning rate
    param_groups.append(
        {
            "params": model.bert.embeddings.parameters(),
            "lr": base_lr * (decay_rate**num_layers),
            "name": "embeddings",
        }
    )

    # Encoder layers: decreasing LR from bottom to top
    for i, layer in enumerate(model.bert.encoder.layer):
        layer_lr = base_lr * (decay_rate ** (num_layers - i - 1))
        param_groups.append(
            {
                "params": layer.parameters(),
                "lr": layer_lr,
                "name": f"layer_{i}",
            }
        )

    # Classifier: full learning rate
    param_groups.append(
        {
            "params": model.classifier.parameters(),
            "lr": base_lr,
            "name": "classifier",
        }
    )

    return param_groups

In[40]:

Code

# Compute learning rates for each layer group
base_lr = 2e-5
decay_rate = 0.9
num_layers = 12

layer_lrs = []
layer_lrs.append(("Embeddings", base_lr * (decay_rate**num_layers)))
for i in range(num_layers):
    lr = base_lr * (decay_rate ** (num_layers - i - 1))
    layer_lrs.append((f"Layer {i}", lr))
layer_lrs.append(("Classifier", base_lr))

Out[41]:

Visualization

Horizontal bar chart showing learning rates for each BERT layer, with embeddings having the smallest rate and classifier having the largest. — Layer-wise learning rate decay with decay factor 0.9. Lower layers receive smaller learning rates to preserve general linguistic features, while upper layers and the classifier adapt more aggressively.

The pattern is clear: the embeddings, which encode the most fundamental token representations, receive the smallest learning rate (about 28% of the base). Each successive layer receives a slightly higher rate, culminating in the classifier head, which gets the full base learning rate since it must learn entirely from scratch.

The layer-wise learning rate for layer $i$ (counting from 0 at the bottom) is computed as:

\eta_i = \eta_{\text{base}} \cdot \xi^{(L - i - 1)}

where:

$\eta_i$ : the learning rate for layer $i$
$\eta_{\text{base}}$ : the base learning rate (used for the classifier)
$\xi$ : the decay rate (typically 0.9 or 0.95)
$L$ : the total number of transformer layers (12 for BERT-Base)

Let's unpack this formula. The exponent $(L - i - 1)$ counts how many layers are above layer $i$ . For the top layer (layer 11 in BERT-Base), $L - i - 1 = 0$ , so $\eta_{11} = \eta_{\text{base}} \cdot \xi^0 = \eta_{\text{base}}$ : it gets the full learning rate. For layer 10, $\eta_{10} = \eta_{\text{base}} \cdot \xi^1$ : it gets 90% of the full rate. For layer 0 (the bottom transformer layer), $\eta_0 = \eta_{\text{base}} \cdot \xi^{11}$ : it gets about 31% of the full rate.

The embeddings sit below all transformer layers. By convention, they receive $\eta_{\text{base}} \cdot \xi^L = \eta_{\text{base}} \cdot 0.9^{12} \approx 0.28 \cdot \eta_{\text{base}}$ . This 28% rate reflects how important it is to preserve the token embeddings that encode fundamental word meanings.