Linearization Explains Fine-Tuning in Large Language Models

📝 Paper Summary

Parameter-Efficient Fine-Tuning (PEFT) Neural Tangent Kernel (NTK) Theory Optimization Dynamics

Regularizing fine-tuning to stay close to pretrained weights explicitly forces the process into a linearized regime, enabling performance prediction and optimal layer selection using Neural Tangent Kernel properties.

Core Problem

Parameter-Efficient Fine-Tuning (PEFT) methods lack a rigorous theoretical understanding of their training dynamics and layer selection choices, making informed algorithmic exploration difficult.

Why it matters:

Full fine-tuning of Large Language Models (LLMs) is computationally prohibitive, necessitating efficient alternatives
Current PEFT methods heuristically select layers or apply rank-limited updates without theoretical guarantees of success
Understanding fine-tuning dynamics allows for better-informed, scalable, and more efficient model adaptation

Concrete Example: When selecting layers for Low Rank Adaptation (LoRA), practitioners currently rely on trial and error. This theoretical framework uses spectral perturbation bounds to predict which layers will minimize error before actual training begins.

Key Novelty

Linearized Fine-Tuning via Explicit Inductive Bias

Applies an explicit weight decay towards the original pretrained model to constrain fine-tuning to the lazy training regime
Formulates the fine-tuning process as kernel regression using the fixed Neural Tangent Kernel (NTK)
Uses the eigenvalue spectrum of the NTK to theoretically bound empirical risk and evaluate the impact of selecting specific layers for adaptation

Architecture

The conceptual approximation of regularized fine-tuning in parameter space by its dual in the tangent space

Breakthrough Assessment

7/10

Provides strong theoretical grounding for PEFT using the NTK, linking layer selection to spectral properties, though empirical results are missing from the provided text to fully verify practical impact.

⚙️ Technical Details

Problem Definition

Setting: Fine-tuning a pretrained model via empirical risk minimization with an explicit L2-distance inductive bias towards pretrained parameters

Inputs: Target task dataset D_T = {(x_i, y_i)} and pretrained model parameters theta_0

Outputs: Fine-tuned model parameters theta_star

Pipeline Flow

Compute NTK induced by currently selected trainable parameters
Evaluate candidate layer by calculating spectral perturbation on the NTK
Predict empirical risk bounds for the candidate layer
Select optimal layers and execute regularized fine-tuning

System Modules

NTK Spectrum Analyzer (Kernel Evaluation)

Computes the Neural Tangent Kernel (NTK) induced by the current set of fine-tuning parameters

Model or implementation: Linearized NTK model

Candidate Layer Evaluator (Kernel Evaluation)

Calculates the spectral perturbation on the NTK when adding a candidate layer's parameters

Model or implementation: Linearized NTK model

Risk Predictor

Estimates the empirical risk bounds using Corollary 1 to determine if the candidate layer should be included

Model or implementation: Mathematical bound estimator

Novel Architectural Elements

Theoretical algorithm for layer selection based on NTK spectral perturbation bounds (Corollary 1) prior to actual training execution

Modeling

Base Model: Not reported in the paper (mentions LLMs generally)

Training Method: Selectively Regularized Gradient Descent

Objective Functions:

Purpose: Minimize the downstream task error while staying close to the pretrained model.

Formally: min_theta R_tilde(theta) + (lambda/2) * ||theta - theta_0||^2

Adaptation: LoRA (Low Rank Adaptation)

Trainable Parameters: Not reported in the paper

Key Hyperparameters:

regularization_strength_lambda: Determines proximity to pretrained model and controls the time scale at which the linear approximation holds

Compute: Not reported in the paper

Comparison to Prior Work

vs. Malladi et al. (2023): Introduces an explicit L2 inductive bias to theoretically guarantee and constrain fine-tuning to the lazy training regime
vs. Deshpande et al. (2021) [not cited in paper]: Provides exact spectral perturbation bounds for layer selection rather than just using the linearized model for inference

Limitations

The text provided does not contain the experimental section, making it impossible to verify the empirical claims numerically
The theoretical bounds rely on the assumption of squared loss, which may not fully capture cross-entropy loss dynamics common in LLM fine-tuning
The regularization parameter lambda requires careful tuning; too large a value restricts learning, while too small a value breaks the linearization approximation

Reproducibility

Code: https://github.com/zahrahimi/linearization

Available: Code repository (https://github.com/zahrahimi/linearization). Missing: The provided text snippet cuts off before the experiments section, so specific datasets, hyperparameters, and model sizes evaluated are not available.

📊 Experiments & Results

Evaluation Setup

Empirical validation of theoretical bounds on Low Rank Adaptation (LoRA) for Large Language Models (LLMs) - specific details missing due to truncated text

Metrics:

Statistical methodology: Not explicitly reported in the paper

Main Takeaways

Regularized fine-tuning with an explicit inductive bias towards pretrained weights forces the model into a linearized (lazy training) regime
The distance between the fine-tuned model and its linearized approximation can be strictly bounded based on the regularization strength lambda
The empirical risk of the fine-tuned model is bounded by the minimum and maximum eigenvalues of the Neural Tangent Kernel (NTK)
Adding trainable parameters (e.g., selecting specific layers for LoRA) perturbs the NTK spectrum, and this perturbation can be bounded to predict which layers will yield the lowest empirical risk before actual training

📚 Prerequisite Knowledge

Prerequisites

Neural Tangent Kernel (NTK) theory
Parameter-Efficient Fine-Tuning (PEFT)
Kernel regression and Reproducing Kernel Hilbert Spaces (RKHS)

Key Terms

PEFT: Parameter-Efficient Fine-Tuning—methods that adapt large models by training only a small subset of parameters to save computational resources

NTK: Neural Tangent Kernel—a mathematical tool describing how neural networks evolve during training, treating them like kernel machines in the infinite-width limit

lazy training: A training regime where the model's weights barely change from their initialization, allowing the network function to be approximated as a strictly linear model

LoRA: Low Rank Adaptation—a PEFT technique that injects trainable low-rank matrices into transformer layers while freezing the original weights

RKHS: Reproducing Kernel Hilbert Space—a mathematical space of functions associated with a kernel, used here to formulate fine-tuning as regularized function estimation

inductive bias: A set of assumptions the learning algorithm uses to predict outputs, implemented here as an explicit penalty for straying far from the pretrained weights

empirical risk: The measured average loss or error of the model on the training dataset

LLM: Large Language Model—a foundation model with an enormous number of parameters trained on vast amounts of text