Understanding Linear Probing then Fine-tuning Language Models from NTK Perspective

TL;DR

This paper analyzes the training dynamics of linear probing followed by fine-tuning for language models using NTK theory, revealing how the linear head norm impacts model performance and calibration, and extending insights to LoRA.

Contribution

It provides a theoretical NTK-based analysis of LP-FT, decomposing the NTK to understand the role of the linear head norm and its effects on fine-tuning dynamics.

Findings

LP-FT outperforms LP and FT alone on ID and OOD data.

Increased linear head norm during LP reduces feature changes and affects calibration.

NTK analysis extends to LoRA, confirming its effectiveness.

Abstract

The two-stage fine-tuning (FT) method, linear probing (LP) then fine-tuning (LP-FT), outperforms linear probing and FT alone. This holds true for both in-distribution (ID) and out-of-distribution (OOD) data. One key reason for its success is the preservation of pre-trained features, achieved by obtaining a near-optimal linear head during LP. However, despite the widespread use of large language models, there has been limited exploration of more complex architectures such as Transformers. In this paper, we analyze the training dynamics of LP-FT for classification tasks on the basis of the neural tangent kernel (NTK) theory. Our analysis decomposes the NTK matrix into two components. This decomposition highlights the importance of the linear head norm alongside the prediction accuracy at the start of the FT stage. We also observe a significant increase in the linear head norm during LP, which stems from training with the cross-entropy (CE) loss. This increase in the linear head norm effectively reduces changes in learned features. Furthermore, we find that this increased norm can adversely affect model calibration, which can be corrected using temperature scaling. Additionally, we extend our analysis with the NTK to the low-rank adaptation (LoRA) method and validate its effectiveness. Our experiments using a Transformer-based model on multiple natural language processing datasets confirm our theoretical analysis. Our study demonstrates the effectiveness of LP-FT for fine-tuning language models. Code is available at https://github.com/tom4649/lp-ft_ntk.