Fine-tuning Architecture

BERT: Pre-training of Deep Bidirectional Transformers

Hard

Problem

Theory

Fine-tuning is the second stage of BERT's two-stage paradigm (Devlin et al., 2019): first pre-train a deep bidirectional Transformer on unlabeled corpora, then fine-tune the entire model on a smaller labeled dataset for a specific downstream task. As the paper states: "Fine-tuning is straightforward since the self-attention mechanism in the Transformer allows BERT to model many downstream tasks. For each task, we simply plug in the task-specific inputs and outputs and fine-tune all the parameters end-to-end."

What It Is / What It Does

Fine-tuning takes the pre-trained BERT encoder and adds a small, randomly initialized task-specific layer (the "head") on top. The entire system is trained end-to-end on labeled data. This contrasts with feature-based approaches (like ELMo) where pre-trained representations are frozen.

• Stage 1 (Pre-training): BERT learns general language understanding from massive unlabeled text via masked language modeling and next sentence prediction, building deep contextual representations across 12 or 24 Transformer layers.
• Stage 2 (Fine-tuning): All parameters are updated on a small labeled dataset. The pre-trained weights provide an excellent initialization, so the model converges quickly with minimal data.

The task head is typically a single linear layer. For sequence classification it operates on the [CLS] hidden state; for token classification it operates on every token's hidden state independently. The head is the only component initialized from scratch.

Key Equations

Let H \in \mathbb{R}^{L \times d} be the final encoder output, where L is sequence length and d is hidden dimension (768 for BERT-Base, 1024 for BERT-Large). $ \in \mathbb{R}^d$ is the [CLS] hidden state, and h_i \in \mathbb{R}^d is the hidden state at position i.

Sequence classification uses only the [CLS] token's representation:

\text{logits} = h_{\text{CLS}} \cdot W + b

where W \in \mathbb{R}^{d \times K} is the classifier weight matrix, b \in \mathbb{R}^K is the bias vector, and K is the number of classes.

Token classification applies the same linear transformation to every token independently:

\text{logits}_i = h_i \cdot W + b \quad \text{for } i = 1, 2, \ldots, L

where W \in \mathbb{R}^{d \times K} and b \in \mathbb{R}^K are shared across all positions.

Probability distribution via softmax:

P(y = k \mid h) = \frac{\exp(\text{logits}_k)}{\sum_{j=1}^{K} \exp(\text{logits}_j)}

Training loss is standard cross-entropy. For sequence classification with ground-truth class c:

\mathcal{L} = -\log P(y = c \mid h_{\text{CLS}})

For token classification, the loss sums over all token positions (ignoring padding and special tokens):

\mathcal{L} = -\sum_{i=1}^{L} \log P(y_i = c_i \mid h_i)

Gradients from this loss flow back through the entire encoder, adjusting all pre-trained parameters. This end-to-end gradient flow is what distinguishes fine-tuning from feature extraction.

Sequence vs Token Classification

The same BERT encoder serves both tasks. The difference lies entirely in which hidden states the classifier head reads.

Sequence classification produces one prediction per input:

• Input representation: The [CLS] token aggregates information from the entire sequence into h_{\text{CLS}} via self-attention across all layers.
• Classifier: A single linear layer maps h_{\text{CLS}} \in \mathbb{R}^d to \mathbb{R}^K.
• Tasks: Sentiment analysis, natural language inference, topic classification, paraphrase detection.

Token classification produces one prediction per token:

• Input representation: Every token position's hidden state h_i is used, not just [CLS].
• Classifier: The same linear layer is applied independently to each token's hidden state. Weights W and b are shared across positions.
• Tasks: Named entity recognition (NER), part-of-speech tagging, slot filling in dialog systems.

For sentence-pair tasks, BERT packs both sentences into one input: [CLS] sentence_A [SEP] sentence_B [SEP]. Segment embeddings distinguish which sentence each token belongs to, and [CLS] captures the cross-sentence relationship.

Layer Freezing

Layer freezing disables gradient updates for selected parameters by setting requires_grad = False, so they retain pre-trained values throughout training.

Why freeze layers:

• Prevent catastrophic forgetting: On small datasets, updating all parameters risks overwriting the general linguistic knowledge learned during pre-training.
• Reduce overfitting: Fewer trainable parameters means lower model capacity relative to dataset size, acting as implicit regularization.
• Speed up training: Frozen layers skip gradient computation and weight updates, reducing memory and training time.

Which layers to freeze:

• Bottom layers (0-3 in BERT-Base): Capture general linguistic features like syntax and morphology. Transfer well across tasks and are safest to freeze.
• Middle layers (4-7): Capture intermediate representations. Whether to freeze depends on domain similarity to pre-training data.
• Top layers (8-11): Capture task-relevant, high-level semantic features. Benefit most from fine-tuning and should generally remain trainable.
• Embedding layer: Often frozen since subword embeddings are well-learned during pre-training.

Guidelines by dataset size:

• Very small (< 1K examples): Freeze all encoder layers, train only the classifier head.
• Small (1K-10K examples): Freeze bottom 6-8 layers, fine-tune top layers and head.
• Medium (10K-100K examples): Freeze bottom 2-4 layers or none. Use a small learning rate.
• Large (100K+ examples): Fine-tune all layers. Sufficient data prevents catastrophic forgetting.

Gradual unfreezing is a related technique: start with only the top layer unfrozen, train briefly, then unfreeze the next layer down and repeat.

Paper Context

Devlin et al. (2019) introduced a fine-tuning recipe that became the standard. The recommended hyperparameters were deliberately conservative.

Recommended hyperparameters:

• Learning rate: 2e-5, 3e-5, 4e-5, or 5e-5 -- roughly 100x smaller than typical training-from-scratch rates.
• Batch size: 16 or 32.
• Epochs: 2, 3, or 4. Very few epochs needed because the model starts from strong representations.
• Warmup: Linear warmup over the first 10% of steps, then linear decay.
• Dropout: 0.1 on the classifier head.

Key results from the paper:

• GLUE benchmark: 80.5 accuracy, a 7.7 point improvement over prior state of the art.
• SQuAD 1.1: F1 of 93.2, surpassing human performance (91.2 F1).
• SQuAD 2.0: F1 of 83.1, handling unanswerable questions via [CLS].
• CoNLL-2003 NER: F1 of 92.8, showing the same model excels at both sequence-level and token-level tasks.

Fine-tuning outperformed feature-based approaches on most tasks, validating the end-to-end strategy.

Catastrophic Forgetting

Catastrophic forgetting occurs when fine-tuning overwrites the general knowledge learned during pre-training. The model becomes specialized for the new task at the cost of losing broad linguistic understanding.

Why small learning rates matter:

• Small updates preserve knowledge: A rate of 2e-5 to 5e-5 nudges pre-trained weights rather than replacing them.
• Large rates destroy representations: At 1e-3 or higher, gradients rapidly overwrite attention patterns learned from billions of tokens.

Why few epochs matter:

• Drift compounds: Every step pushes parameters further from pre-trained values. Many epochs on small data cause substantial cumulative drift.

Connection to layer freezing:

• Freezing as direct remedy: If parameters do not update, they cannot forget.
• Complementary strategies: Layer freezing, small learning rates, few epochs, and warmup address the same problem from different angles and are often combined.

Discriminative learning rates offer a middle ground: assign smaller rates to lower layers and larger rates to upper layers. For example, layer 0 uses 1e-6, layer 6 uses 1e-5, layer 11 uses 5e-5.

Numerical Example

Consider 3-class sentiment classification (negative, neutral, positive) with hidden dimension d = 4.

Step 1: Encoder output. After passing "[CLS] This movie is great [SEP]" through BERT, the [CLS] hidden state is:

h_{\text{CLS}} = [0.5, -0.3, 0.8, 0.1]

Step 2: Classifier weights. W \in \mathbb{R}^{4 \times 3} and b \in \mathbb{R}^3:

W = \begin{bmatrix} 0.2 & -0.1 & 0.4 \\ 0.3 & 0.5 & -0.2 \\ -0.1 & 0.2 & 0.6 \\ 0.4 & -0.3 & 0.1 \end{bmatrix}, \quad b = [0.1, -0.1, 0.2]

Step 3: Compute logits. \text{logits} = h_{\text{CLS}} \cdot W + b:

\text{logits}_0 = (0.5)(0.2) + (-0.3)(0.3) + (0.8)(-0.1) + (0.1)(0.4) + 0.1 = 0.07

\text{logits}_1 = (0.5)(-0.1) + (-0.3)(0.5) + (0.8)(0.2) + (0.1)(-0.3) - 0.1 = -0.17

\text{logits}_2 = (0.5)(0.4) + (-0.3)(-0.2) + (0.8)(0.6) + (0.1)(0.1) + 0.2 = 0.95

\text{logits} = [0.07, -0.17, 0.95]

Step 4: Softmax.

\exp(\text{logits}) = [e^{0.07}, e^{-0.17}, e^{0.95}] = [1.073, 0.844, 2.586]

\text{sum} = 1.073 + 0.844 + 2.586 = 4.503

P = [0.238, 0.187, 0.574]

The model predicts class 2 (positive) with 57.4% probability, matching "This movie is great."

Step 5: Loss. With ground-truth label class 2:

\mathcal{L} = -\log(0.574) = 0.555

This loss backpropagates through the classifier head and the entire BERT encoder.

Token Classification Example

NER on "John works at Google" with labels {O, PER, ORG} (K = 3). The same W and b are applied to every token's hidden state independently.

After encoding:

h_{\text{John}} = [0.9, -0.2, 0.4, 0.3], \quad h_{\text{works}} = [0.1, 0.5, -0.1, 0.2], \quad h_{\text{Google}} = [0.7, 0.1, 0.6, -0.4]

"John": \text{logits} = h_{\text{John}} \cdot W + b = [0.30, -0.30, 0.87]. Softmax: [0.243, 0.133, 0.624]. Prediction: ORG. With random weights this is incorrect -- a trained model would predict PER.

"works": \text{logits} = [0.36, 0.06, 0.10]. Softmax: [0.402, 0.298, 0.310]. Prediction: O. Correct -- "works" is not an entity.

"Google": \text{logits} = [0.05, 0.12, 0.78]. Softmax: [0.225, 0.242, 0.533]. Prediction: ORG. Correct -- "Google" is an organization.

The total loss sums cross-entropy across all token positions.

Modern Context

As models grew to hundreds of billions of parameters, full fine-tuning became impractical, driving parameter-efficient alternatives.

• LoRA (Low-Rank Adaptation): Freezes pre-trained weights and injects trainable low-rank matrices. The update is \Delta W = AB with A \in \mathbb{R}^{d \times r}, B \in \mathbb{R}^{r \times d}, r \ll d. Reduces trainable parameters by 1000x while matching full fine-tuning on many tasks.
• Adapter layers: Small bottleneck modules inserted between Transformer layers. Only adapters train; the base model is frozen.
• Prompt tuning: Learnable continuous vectors prepended to input embeddings. The model is completely frozen; only prompt vectors are optimized.
• In-context learning (ICL): Introduced by GPT-3. Task demonstrations given in the input prompt; the model performs zero-shot or few-shot without parameter updates.
• QLoRA: Combines 4-bit quantization with LoRA, enabling fine-tuning of 65B+ parameter models on a single GPU.

Despite these advances, full fine-tuning remains the gold standard when maximum performance is needed. BERT's core insight -- that pre-trained representations can be adapted with minimal architecture changes -- underpins all of these methods.

Pitfalls

• Learning rate too high: Using 1e-3 or higher rapidly destroys pre-trained representations. Always use 2e-5 to 5e-5.
• Too many epochs on small data: Training 10+ epochs on a few thousand examples causes severe overfitting and catastrophic forgetting. Stick to 2-4 epochs.
• Forgetting to use [CLS] for sequence tasks: Averaging all token states or using the last token misaligns with pre-training, which used [CLS] as the aggregate representation.
• Applying sequence head to token tasks: Using only [CLS] for NER or POS tagging discards per-token information. Token classification requires every token's hidden state.
• Skipping warmup: The randomly initialized head produces large, noisy gradients that can permanently damage pre-trained representations. Use warmup for 5-10% of steps.
• Freezing too many layers: Training only the head limits adaptation. Underperforms on domain-specific tasks.
• Freezing too few layers on small data: Fine-tuning all 12 layers on 500 examples gives too many degrees of freedom. Freeze the bottom 8-10 layers.
• Ignoring input formatting: Omitting [CLS]/[SEP], wrong segment IDs, or exceeding 512 tokens produces degraded results without an obvious error.

Examples

Hints

1. Subtract each row maximum before exponentiating logits.
2. Set grad_logits[np.arange(batch_size), labels] -= 1.
3. Use cls_states.T @ grad_logits for the weight gradient.

Requirements

• Use NumPy.
• Select the position-zero state from each sequence.
• Compute numerically stable softmax probabilities and mean cross-entropy.
• Compute analytical gradients for the classifier weights and bias.
• Apply exactly one gradient-descent update.
• Return exactly the documented arrays and pre-update loss.

Constraints

• hidden_states has shape (B, S, D) and dtype float64.
• labels has shape (B), dtype int64, and values from zero through C minus one.
• classifier_W has shape (D, C) and dtype float64.
• classifier_b has shape (C) and dtype float64.
• learning_rate is a positive float.
• B, S, D, and C are positive.

Starter Code

import numpy as np

def bert_fine_tuning_step(hidden_states: np.ndarray, labels: np.ndarray,
                          classifier_W: np.ndarray, classifier_b: np.ndarray,
                          learning_rate: float) -> dict:
    """
    Returns updated classifier parameters and the pre-update loss.
    """
    pass

Test Cases