Characterizing the Training Dynamics of Private Fine-tuning with Langevin Diffusion

## Abstract

Abstract Differentially private (DP) machine learning pipelines typically involve a two-phase process: non-private pre-training on a public dataset, followed by fine-tuning on private data using DP optimization techniques. In the DP setting, it has been observed that full fine-tuning may not always yield the best test accuracy, even for in-distribution data. This paper (1) analyzes the training dynamics of DP linear probing (LP) and full fine-tuning (FT), and (2) explores the phenomenon of sequential fine-tuning, starting with linear probing and transitioning to full fine-tuning (LP-FT), and its impact on test loss. We provide theoretical insights into the convergence of DP fine-tuning within an overparameterized neural network and establish a utility curve that determines the allocation of privacy budget between linear probing and full fine-tuning. The theoretical results are supported by empirical evaluations on various benchmarks and models. The findings reveal the complex nature of DP fine-tuning methods. These results contribute to a deeper understanding of DP machine learning and highlight the importance of considering the allocation of privacy budget in the fine-tuning process.

## 1 Introduction

Today, many differentially-private (DP) machine learning pipelines proceed in two phases: (1) A model is pre-trained (non-privately) on a public dataset.
(2) The model is then fine-tuned on private data, typically using DP optimization techniques such as DP stochastic gradient descent (DP-SGD) and its variants Hoory et al. (); Bie et al. (); De et al. (); Tang et al. (); Ganesh et al. ().
The non-private pre-training phase is believed to guide model parameters into a favorable region within the loss landscape, thereby facilitating easier local optimization under DP constraints. For example,
pre-training a GPT-2 model on public data before privately fine-tuning it has been shown to significantly improve the BLEU score from 24.25 to 63.19, while also reducing the privacy cost (i.e., the DP $\epsilon$ parameter) from 41.48 to 6.27 during the private fine-tuning phase Yu et al. (); Ganesh et al. ().

Figure: Figure 1: Linear probing freezes the lower layers and optimizes the last linear layer while full fine-tuning optimizes the whole network.
Refer to caption: x1.png

Two common strategies for fine-tuning models (with or without DP) are shown in Figure 1: (1) full fine-tuning (FT) updates all parameters of the model, and (2) linear probing (LP) freezes the lower layers, and only updates the final linear layer.
In the non-DP setting, full fine-tuning is known to exhibit better in-distribution accuracy than linear probing on datasets like ImageNet, CIFAR-10, Ent-30 and Liv-17 with few exceptions(^1^11Non-DP linear probing achieves better in-distribution test accuracies on DomainNet Kumar et al. ().) Kornblith et al. (); Zhai et al. (); He et al. (); Kumar et al. (). On the other hand, for out-of-distribution tasks, linear probing can sometimes outperform full-finetuning Kumar et al. ().

In the DP setting, Tang et al. () have observed empirically that full fine-tuning does not always yield the best test accuracy, even for in-distribution(^2^22In this work, we focus on in-distribution settings. We will just use “test accuracy” rather than “in-distribution test accuracy” in the rest of this paper.) data. Rather, a sequential fine-tuning strategy that starts with linear probing for some epochs, then proceeds with full fine-tuning for the remaining epochs (we refer to this strategy as LP-FT), empirically achieves better test accuracy than either full fine-tuning or linear probing alone Tang et al. ().

Figure: (a) DP Utility Curves, $\sigma=0.3$
Refer to caption: x2.png

A natural question is whether this phenomenon is generally true for DP machine learning, and if so, *how much* linear probing is enough?
In other words, given a total privacy budget, how much should we allocate to linear probing vs. full fine-tuning to achieve the best test accuracy?
These questions are currently difficult to answer because we lack a theoretical understanding of the impact of DP fine-tuning on model optimization. Analyzing fine-tuning is challenging because we need to understand the training dynamics, rather than just the loss function and its overall minimizers Neyshabur et al. (); Kumar et al. (); Wang et al. (). Moreover, existing techniques and DP convergence results only consider the dimension of model parameters and ignore the model architecture Ganesh et al. (). We require new theoretical methods that capture the key structural difference between linear probing and full fine-tuning (Figure 1).

In this work,
we provide both theoretical and empirical justification for the conjecture that in the DP setting, linear probing followed by full fine-tuning achieves better test loss than linear probing or full fine-tuning alone.
We present a theoretical explanation of DP fine-tuning within an overparameterized two-layer linear network.
We establish that under certain initial conditions, DP fine-tuning can be less effective than DP linear probing with high-quality pre-trained features. Furthermore, we characterize a non-monotonic and concave utility curve: as one increases the fraction of rounds used for linear probing, test accuracy increases up to a point, then decreases.
Figure 2 illustrates this trend on a ResNet-50 architecture; experimental settings are provided in Section 4.2. We observe that for non-private optimization (bottom), full fine-tuning achieves the highest test accuracy. However, for DP-SGD (top), linear probing outperforms full fine-tuning, and LP-FT outperforms both LP and FT.

Our theoretical results analyze the convergence of a Langevin diffusion-based model of DP gradient descent (DP-GD).
To prove our main results, we introduce to the Lanvegin diffusion setting the *imbalance matrix*, an invariant previously studied in the context of gradient flows Arora et al. (); Du et al. (); Min et al. (, ); Marcotte et al. ().
We prove that in DP fine-tuning, (1) the diagonal of the expected imbalance matrix linearly increases and other entries remain invariant, which is different from the case for (noise-free) gradient flows.
(2) However, we find that the expectation of the imbalance matrix is invariant if we scale the DP noise of the linear head with respect to its size. This invariant property of multi-layer networks helps us disentangle the nonconvexity of DP-GD and leads to new utility bounds.

Technically, these tools allow us to bound the convergence rates of both DP linear probing and nonconvex DP full fine-tuning.
Unlike prior work in this field, our results do not require
convexity conditions, other “model-free assumptions” (e.g. Lipschitz conditions, dissipativity, smoothness, the Poincaré inequality), or approximations (e.g. neural tangent kernel) in DP gradient descent analysis Bassily et al. (); Raginsky et al. (); Bubeck et al. (); Asi et al. (); Tu (); Bu et al. (); Ganesh et al. (); Ye et al. ().
We believe our analysis techniques may thus be of broader interest.

To summarize, our contributions are:

- •
DP Fine-Tuning (FT) vs. Linear Probing (LP):
We provide the first theoretical evidence that DP linear probing can outperform DP full fine-tuning, even on in-distribution data.
In particular, our analysis on a two-layer linear network regression task indicates that DP full fine-tuning may underperform relative to linear probing when privacy budgets are low.
- •
Two-Phase Sequential Approach (LP-FT): We analyze a two-phase fine-tuning method, beginning with linear probing and then transitioning to full fine-tuning.
Our results theoretically demonstrate
a trade-off in allocating the privacy budget between the two phases. Contrary to initial expectations, this sequential approach can outperform mere linear probing, even when pre-trained features are already highly effective.
- •
We empirically evaluate the trends predicted by our theoretical model over nonlinear neural networks. We experiment on four fine-tuning benchmarks (ImageNet-1K$\rightarrow$CIFAR-10,ImageNet-1K$\rightarrow$STL-10, CIFAR-10$\rightarrow$STL-10, RandP$\rightarrow$CIFAR-10) and five model architectures (MobileNet-v3, Mini-DeiT, ResNet-18, ResNet-50, WideResNet-16-4). Our results confirm our theoretical predictions, showing that concave utility curves generally exist for deep neural networks on real datasets.

## 2 Setup and Background

Task and notation. For our theoretical analyses, we consider a regression task with $n$ training data points: $\{(x_{i},y_{i}):x_{i}\in\mathbb{R}^{d},y_{i}\in\mathbb{R}\}_{i=1}^{n}$. We collect the data inputs in a matrix $X\in\mathbb{R}^{n\times d}$ and put the labels in a vector $Y\in\mathbb{R}^{n}$. The objective is to train a predictor $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ that minimizes the empirical risk function $\mathcal{L}=\sum_{i=1}^{n}({1}/{2})(f(x_{i})-y_{i})^{2}$.

### 2.1 Differentially-Private Optimization

Differential privacy (DP) is a widely-used method to evaluate privacy leakage in a database accessed through queries Dwork and Roth (). In the context of machine learning, it ensures that an adversary cannot confidently ascertain whether a specific training sample (or set of training samples) was in the training dataset.

###### Definition 2.1 ( $(\epsilon,\delta)$ -Differential Privacy Dwork and Roth ( 2014 ) ) .

We say two datasets $\mathcal{D}$, $\mathcal{D}^{\prime}$ with $n$ records each are neighboring if they only differ in one record. Given privacy parameters $\epsilon>0,\delta\in(0,1)$, we say a random algorithm $\mathcal{A}$ satisfies $(\epsilon,\delta)$-differential privacy if for any measurable $B\subseteq\mathrm{range}(A)$ and for any pair of neighboring datasets $\mathcal{D},\mathcal{D}^{\prime}$, we have

$$ $\mathbb{P}(\mathcal{A}(\mathcal{D})\in B)\leq e^{\epsilon}\mathbb{P}(\mathcal{ A}(\mathcal{D}^{\prime})\in B)+\delta$ (1) $$

where $\mathcal{A}(\mathcal{D})$ denotes the output distribution of algorithm $\mathcal{A}$ on dataset $\mathcal{D}$.

###### Definition 2.2 ( $\ell_{2}$ sensitivity (Dwork and Roth, 2014 ) ) .

Let $g:\mathbb{R}^{n\times d}\to\mathbb{R}^{p}$ be a function operating on datasets. Let $D=(X,Y),D^{\prime}=(X^{\prime},Y^{\prime})\in\mathbb{R}^{n\times(d+1)}$ be *neighboring* datasets. The $\ell_{2}$-sensitivity of $g$
is defined as
$\Delta(g):=\sup_{D\sim D^{\prime}}\|g(D)-g(D^{\prime})\|_{2}$.

The privacy guarantee of Langevin diffusion depends on the $\ell_{2}$ sensitivity of the loss gradients Ganesh et al. (); Ye et al. ().

DP Gradient Descent (DP-GD). DP-GD is a widely-adopted DP variant of Gradient Descent. Each iteration of DP-GD consists of two parts: (1) clipping the gradients of each data point to a prefixed norm of at most $C$, and (2) adding zero-mean spherical Gaussian noise of a prefixed scale $\sigma$ to the mean of the clipped gradients Abadi et al. (); Sander et al. (). DP-GD is $(\epsilon,\delta)$-differentially private if we choose $\sigma=\Omega((C/n)\sqrt{T\log(1/\delta)}/\epsilon)$, where $T$ is the number of iterations Dwork and Roth ().

Differentially Private Gradient Descent
{algorithmic}[1]
\Requireinitial parameter $\theta^{(0)}$, learning rate $\alpha$, noise scale $\sigma$, gradient norm bound $C$, train data ${\cal D}$, $T\in{\mathbb{Z}}$.
\Ensure$\theta^{(T)}$ and compute the overall privacy cost $(\epsilon,\delta)$ using the Rényi Differential Privacy accountant.
\For$t=0,\dots,T-1$
\State${\boldsymbol{g}}^{(t)}\leftarrow\frac{1}{n}\left(\sum_{(x_{i},y_{i})\in{\cal D
}}\mathrm{clip}_{C}(\nabla_{w}\ell(w^{(t)},x_{i},y_{i})))\right)$
\State${\boldsymbol{g}}^{(t)}\leftarrow{\boldsymbol{g}}^{(t)}+\mathcal{N}(0,\sigma^{2
}\mathrm{I})$
\State$\theta^{(t+1)}\leftarrow\theta^{(t)}-\alpha{\boldsymbol{g}}^{(t)}$
\EndFor

Langevin Dynamics. A Langevin diffusion is a stochastic differential equation that models the dynamics of a system under the influence of both deterministic and fluctuating forces, also known as random forces Lemons and Gythiel (). For example, we can define an $m$-dimensional Langevin diffusion to model DP-GD as follows:

$$ $\partial\theta(t)=-\nabla\mathcal{L}(\theta(t))\partial t+\sqrt{2\sigma^{2}} \partial W(t)$ (2) $$

where $\theta\in\mathbb{R}^{m}$ is the model parameter evolving in continuous time, $\mathcal{L}$ is the empirical risk, $\sigma>0$ is the noise scale, and $\{W(t)\}_{t\geq 0}$ is the standard Brownian motion in $\mathbb{R}^{m}$. This Langevin diffusion has been used for the analysis of gradient-based stochastic learning algorithms like DP-(S)GD Xu et al. (); Li et al. (); Chourasia et al. (); Ye et al. (); Ganesh et al. (). Notice that (2) captures the additive noise of DP-GD, but it does not capture the effect of gradient clipping. Like prior work Xu et al. (); Chourasia et al. (); Ganesh et al. (); Ye et al. (), we do not explicitly model clipping. Instead, we have assumed the noise scale is upper bounded in 3.4; other works typically handle this issue by assuming Lipschitz conditions or making analogous assumptions upper bounding the gradient norm Raginsky et al. (); Ganesh et al. ().

###### Definition 2.1 ( ( ϵ , δ ) italic-ϵ 𝛿 (\epsilon,\delta) ( italic_ϵ , italic_δ ) -Differential Privacy Dwork and Roth ( 2014 ) ) .

###### Definition 2.2 ( ℓ 2 subscript ℓ 2 \ell_{2} roman_ℓ start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT sensitivity (Dwork and Roth, 2014 ) ) .

## 3 Theoretical Analysis

### 3.1 Linear overparameterized setting

Model. We focus on a setting adapted from prior works that utilizes pre-trained representations Arora et al. (); Tarmoun et al. (); Kumar et al. (); Min et al. ().
We define a two-layer linear network as function $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$, $f(x)=\langle v,Bx\rangle$ where $v\in\mathbb{R}^{k},B\in\mathbb{R}^{k\times d}$. Then we can write the empirical risk function as

$$ $\mathcal{L}(v,B)=\frac{1}{2}\|XB^{T}v-Y\|_{2}^{2}.$ (3) $$

For simplicity, we assume that $n\geq d$ and $X^{T}X=I_{d\times d}$. Note that the loss function is nonconvex in the parameters being fine-tuned, so the gradient-descent training becomes a nonlinear dynamical system.
This significantly complicates theoretical analysis Min et al. (). Prior works have dealt with this challenge by using heavy approximations based on neural tangent kernel and first-order Taylor expansions Bu et al. (); Ye et al. (). We overcome the theoretical difficulties by using conservation laws and geometric properties of Langevin dynamics. We introduce the *imbalance matrix*, an invariant previously studied in the context of gradient flows, to Langevin diffusion Arora et al. (); Du et al. (); Min et al. (, ); Marcotte et al. ().

Pretrained features. We evaluate a feature extractor $B$ by the least square error $\gamma(B):=\min_{u\in\mathbb{R}^{k}}\mathcal{L}(u,B)=Y^{T}(I_{n}-XB^{T}(XB^{T}
)^{{\dagger}})Y$, where $(\cdot)^{{\dagger}}$ denotes the pseudo inverse of a matrix. This metric measures the optimal loss we can obtain by linear probing when fixing the current features. We denote the pretrained feature extractor as $B_{0}$. We use $\gamma$ as shorthand for the initial least square error of our pretrained feature extractor, i.e., $\gamma(B_{0})$. We follow the common assumption in prior works Tripuraneni et al. (); Kumar et al. () and suppose $B_{0}$ has orthonormal rows.

Fine-tuning methods. To analyze different fine-tuning approaches, we use Langevin diffusion. Langevin diffusion is a continuous generalization of DP-GD Kumar et al. (). It has been extensively studied as a way to understand SGD and its variants Bubeck et al. (); Cheng et al. (); Xu et al. (); Chourasia et al. (); Ganesh et al. (); Ye et al. ().
Compared to discrete time-step analysis of SGD, continuous schemes like Langevin diffusion (1) are amenable to theoretical analysis; and (2) provide a more systematic way to analyze and understand a variety of algorithms Su et al. (); Li et al. (); Elkabetz and Cohen (); Kumar et al. (); Miyagawa (). Based on previous works Erdogdu and Hosseinzadeh (); Chewi et al. (); Mou et al. (); Erdogdu et al. (); Balasubramanian et al. (); Zhang et al. (); Johnston and Sabanis (); Nguyen et al. (); Mousavi-Hosseini et al. (), we define the Langevin diffusion for linear probing and full fine-tuning below.

###### Definition 3.1 (Langevin diffusion for linear probing) .

Let $W_{t}$ be the standard $k$-dimensional Brownian motion. Then the Langevin diffusion for linear probing is defined by the following stochastic differential equation:

$$ $\displaystyle\partial v=$ $\displaystyle-\nabla_{v}\mathcal{L}(v,B_{0})\partial t+\sqrt{2\sigma^{2}} \partial W_{t}$ $\displaystyle=$ $\displaystyle-B_{0}X^{T}(XB_{0}^{T}v-Y)\partial t+\sqrt{2\sigma^{2}}\partial W _{t}.$ (4) $$

Here we use “$\partial$” as the differential notation.

###### Definition 3.2 (Langevin diffusion for fine-tuning) .

Let $W_{t}$ be the standard $k$-dimensional brownian motion and $W_{t}^{\prime}$ be a matrix whose entries are standard and independent brownian motions. Then we define the Langevin diffusion for fine-tuning a two-layer linear network as

$$ $\displaystyle\partial v$ $\displaystyle=-\nabla_{v}\mathcal{L}(v,B)\partial t+\sqrt{2\sigma^{2}}\partial W _{t}$ (5) $\displaystyle=-BX^{T}(XB^{T}v-Y)\partial t+\sqrt{2\sigma^{2}}\partial W_{t}$ $\displaystyle\partial B$ $\displaystyle=-\nabla_{B}\mathcal{L}(v,B)\partial t+\sqrt{2\sigma^{2}}\partial W _{t}^{\prime}$ $\displaystyle=-v(XB^{T}v-Y)^{T}X\partial t+\sqrt{2\sigma^{2}}\partial W_{t}^{ \prime}.$ $$

Assumptions. We introduce two assumptions for our analysis. The first assumption models the widely-used random initialization of neural network parameters.

###### Assumption 3.3 (Random initialization of linear head He et al. ( 2015 ) ) .

Initialize the linear head according to a normal distribution $v_{0}\sim\mathcal{N}(0,\beta I_{k\times k})$ where $\beta>0$.

The second assumption is based on the first one. It describes a common phenomenon in differential privacy deployment: the loss might not converge if the privacy mechanism perturbs the gradients too much. To ensure that DP-GD works for full fine-tuning, we assume that the noise scale (or variance) in the privacy mechanism is upper bounded by a constant.

###### Assumption 3.4 (Upper bounded noise scale) .

Let $\beta>\frac{-\|X^{T}Y\|+\sqrt{\|X^{T}Y\|^{2}+4(1+d)\|X^{T}Y\|+4d}}{2k}$. Then we assume that the noise scale $\sigma>0$ we add for privacy in the fine-tuning process is upper-bounded by

$$ $\sigma^{2}<\min\left\{\frac{k\beta+\|B_{0}X^{T}Y\|^{2}}{2k},\frac{k\beta-1}{ \sqrt{2}(1+d)},\frac{1}{1+\sqrt{2}(1+d)}\left[\frac{k\beta(k\beta+\|X^{T}Y\|^{ 2})}{(1+d)\|X^{T}Y\|+d}-1\right]\right\}$ (6) $$

By lower bounding $\beta$ and upper bounding $\sigma$, we ensure that the signal of the gradient is significant enough so that the loss is more likely to decrease in full fine-tuning.

###### Definition 3.1 (Langevin diffusion for linear probing) .

###### Definition 3.2 (Langevin diffusion for fine-tuning) .

###### Assumption 3.3 (Random initialization of linear head He et al. ( 2015 ) ) .

###### Assumption 3.4 (Upper bounded noise scale) .

### 3.2 Structural invariance of linear networks

To understand the properties of a dynamics analysis problem, it can be useful to identify *invariants*, or functions whose output is conserved during optimization.
Such conservation laws can be seen as a "weaker" form of implicit bias, helping to elucidate which properties (e.g., sparsity, low-rank) are preferred by the optimization dynamics among a potentially infinite set of minimizers Marcotte et al. (). To prove the convergence of our optimization, we study the *imbalance matrix*, an invariant for multi-layer linear networks.
The imbalance matrix has previously been studied in the context of gradient flows Arora et al. (); Du et al. (); Min et al. (, ); Marcotte et al. (). However, it has not, to the best of our knowledge, been used in the context of Langevin dynamics.

###### Definition 3.5 (Imbalance matrix) .

For a two-layer linear network, we define the imbalance matrix as

$$ $D:=vv^{T}-BB^{T}.$ (7) $$

Prior work on gradient flows has found that the imbalance matrix remains invariant over the evolution of gradient flows modeling gradient descent Arora et al. (); Du et al. (); Marcotte et al. (). This property can be used to derive tight convergence bounds Min et al. (, ).
However, a similar analysis has not materialized for Langevin diffusion models of DP-GD.

We observe that prior work on Langevin diffusion to analyze private optimization has implicitly assumed that the sensitivity of each layer in a neural network is the same Ganesh et al. (); Ye et al. ().
Hence, they fix a uniform noise scale for every parameter of the network.
Under these conditions, we show that, when we ignore the sensitivity of each layer and use a uniform noise scale $\sigma$, the imbalance matrix is *not* invariant in expectation, unlike in (noise-free) gradient flow Arora et al. (); Du et al. (); Marcotte et al. (). This complicates the use of the imbalance matrix for theoretical analysis Ye and Du ().

###### Lemma 3.6 (Imbalance matrix in fine-tuning)

During fine-tuning (Equation (5)), the derivative of the imbalance matrix $D$ in Definition 3.5 is

$$ $\frac{\partial}{\partial t}\mathbb{E}[D]=(1-d)\sigma^{2}I_{k\times k},$ (8) $$

where $d$ is the dimension of data inputs ($B\in\mathbb{R}^{k\times d}$).

Our main observation is that by modeling differences in sensitivity of different layers, we can recover the invariance property of the imbalance matrix.
The following proposition characterizes the sensitivity of the linear head and the feature extractor, and illustrates why they have differing sensitivities at initialization.

###### Proposition 3.7

We assume that the training dataset $\mathcal{D}=(X,Y)$ is normalized such that $X^{T}X=I_{d\times d},\|Y\|_{2}=1$.
We initialize the linear head by $v_{0}\sim\mathcal{N}(0,\beta I_{k\times k})$ and $\beta=k/\sqrt{d}$.
At the initialization of full fine-tuning, the linear head $v$ has a greater sensitivity than the feature extractor $B$:

$$ $\Delta(\nabla_{v}\mathcal{L}(v_{0},B_{0}))=\Theta\left(\sqrt{d}\cdot\Delta( \nabla_{B}\mathcal{L}(v_{0},B_{0}))\right)$ (9) $$

Based on this observation, we propose a modified version of Langevin diffusion for full fine-tuning, which accounts for layer-wise sensitivity. With this modified definition, the imbalance matrix is again invariant in expectation.

###### Definition 3.8 (Modified Langevin diffusion for fine-tuning) .

Let $W_{t}$ be the standard $k$-dimensional brownian motion. Let $W_{t}^{\prime}$ be a $k\times d$ matrix whose entries are standard and independent brownian motions. Then we define the modified Langevin diffusion for fine-tuning a two-layer linear network as

$$ $\displaystyle\partial v=$ $\displaystyle-\nabla_{v}\mathcal{L}(v,B)\partial t+\sqrt{2\sigma^{2}{\color[ rgb]{1,0,0}d}}\partial W_{t}$ (10) $\displaystyle=$ $\displaystyle-BX^{T}X(B^{T}v-X^{T}Y)\partial t+\sqrt{2\sigma^{2}{\color[rgb]{ 1,0,0}d}}\partial W_{t}$ $\displaystyle\partial B=$ $\displaystyle-\nabla_{B}\mathcal{L}(v,B)\partial t+\sqrt{2\sigma^{2}}\partial W _{t}^{\prime}$ $\displaystyle=$ $\displaystyle-v(XB^{T}v-Y)^{T}X\partial t+\sqrt{2\sigma^{2}}\partial W_{t}^{ \prime}.$ $$

The only difference between this diffusion and Equation 5 is the additional factor of $d$, shown in red, reflecting the fact that the linear head has greater function sensitivity than the feature extractor.

###### Lemma 3.9 (Imbalance matrix in the modified fine-tuning dynamics)

In the modified fine-tuning dynamics (Equation (10)), the derivative of the imbalance matrix $D$ in Definition 3.5 is zero in expectation

$$ $\mathbb{E}\left[\frac{dD}{dt}\right]=0.$ (11) $$

As a result, the expected value of the imbalance matrix is a constant in the fine-tuning process.

Note that by modeling the different sensitivities of different layers, we are also modeling a slightly different variant of DP-GD, which applies different noise scales to different layers. This approach has been used successfully in practice, though it is not the most common variant of DP-GD Cummings et al. ().

###### Definition 3.5 (Imbalance matrix) .

###### Lemma 3.6 (Imbalance matrix in fine-tuning)

###### Proposition 3.7

###### Definition 3.8 (Modified Langevin diffusion for fine-tuning) .

###### Lemma 3.9 (Imbalance matrix in the modified fine-tuning dynamics)

### 3.3 Convergence of fine-tuning methods

###### Theorem 3.10 (Convergence of linear probing)

For simplicity, we assume that $X^{T}X=I_{d\times d}$. If Assumption 3.3 holds and we run linear probing for time $t$, then the loss at time $t$ satisfies:

$$ $\mathbb{E}[\mathcal{L}_{\mathrm{lp}}(t)]\leq\frac{1}{2}(k\beta+\|Y\|^{2})e^{-t }+(\gamma+k\sigma^{2})(1-e^{-t})$ (12) $$

where $k$ is the dimension of the linear head and $\gamma$ denotes the least square error of the pretrained feature extractor $B_{0}$. Equation (12) upper monotonically decreases in time if Assumption 3.4 also holds.

The first coefficient $\frac{1}{2}(k\beta+\|Y\|^{2})$ in Equation (12) is the initial loss that depends on the random initialization and the training data.
In particular, notice that the initial loss does not depend on our choice of $B_{0}$; this is due to looseness in our analysis; which considers a worst-case initialization.
The second coefficient $\gamma+k\sigma^{2}$ depends on the noise scale $\sigma$ and the minimum loss $\gamma$ linear probing could achieve when we fix the feature extractor. The first term $\frac{1}{2}(k\beta+\|Y\|^{2})e^{-t}$ describes that the loss tends to exponentially decrease, while the second term describes the limiting behavior induced by linear probing and the added noise.

###### Theorem 3.11 (Convergence of fine-tuning)

For simplicity, we assume that $X^{T}X=I_{d\times d}$. If Assumption 3.3 and 3.4 hold, and we run fine-tuning (Equation (10)) for time $t$, then the loss converges:

$$ $\mathbb{E}[\mathcal{L}(t)]\leq\frac{1}{2}(k\beta+\|Y\|^{2})e^{-ct}+L^{\square} (1-e^{-ct})$ (13) $$

where $\begin{cases}c=k\beta-1-\sqrt{2}\sigma^{2}(1+d)>0\\
L^{\square}=\sigma^{2}\frac{(1+d)\|X^{T}Y\|+d}{c}\end{cases}$.

This upper bound has a similar form to Equation (12) while the factor $c$ of the exponential terms depends on the initialization and the noise. If we initialize the parameters far from the origin, i.e. $\beta>>0$, the nonlinearity of fine-tuning dynamics would force the loss to converge faster. On the other hand, if the random mechanism perturbs the gradient more, i.e. $\sigma$ increases, then the noise slows down the convergence. When we take limit $\sigma\rightarrow 0$ in Theorem 3.10 and 3.11, the Langevin diffusion degenerates to a gradient flow and the loss converges exponentially to zero as $t\rightarrow\infty$. This recovers known results from the non-private optimization literature Min et al. ().

There is still room for improvement in Theorem 3.11. The upper bound is loose in the sense that it does not incorporate $\gamma$, the pretraining quality. Nonetheless, we can develop new theoretical results based on this upper bound.

###### Theorem 3.10 (Convergence of linear probing)

###### Theorem 3.11 (Convergence of fine-tuning)

### 3.4 Linear probing vs. fine-tuning

The following result identifies a condition under which linear probing may work better than fine-tuning. Note that we cannot conclude that linear probing is definitively better than fine-tuning for our model, as both Theorem 3.10 and Theorem 3.11 are upper bounds. However, we observe empirically that the predictions that come from this condition appear to qualitatively hold.

###### Corollary 3.12

For simplicity, we assume that $X^{T}X=I_{d\times d}$. Suppose Assumption 3.3 and 3.4 hold. Then linear probing achieves a smaller loss upper bound in expectation if:

$$ $t\rightarrow\infty\text{ and }\gamma<\frac{\sigma^{2}}{c}((1+d)\|X^{T}Y\|+d)-k \sigma^{2}.$ (14) $$

The intuition of Corollary 3.12 is that a linear head works better given better pretrained features (i.e. smaller $\gamma$), a smaller size of the linear head, a larger noise scale, and enough running time. Otherwise, fine-tuning will eventually achieve a smaller loss than linear probing as $t\rightarrow\infty$.

###### Corollary 3.12

### 3.5 First-linear-probing-then-finetuning

We next study the two-phase fine-tuning method of linear probing followed by full fine-tuning. We refer to this strategy as LP-FT, and find that it may work better in some settings than linear probing or full fine-tuning alone. Linear probing first can accelerate fine-tuning by aligning the linear head.
The following result provides a convergence bound for LP-FT when we linear-probe for time $t_{\mathrm{lp}}$, and then fully fine-tune for time $t$.

###### Proposition 3.13 (Linear probing accelerates fine-tuning)

For simplicity, we assume that $X^{T}X=I_{d\times d}$. If Assumption 3.3 and 3.4 hold, and we run linear probing for time $t_{\mathrm{lp}}$ and then fine-tuning (Equation (10)) for time $t$, then the loss is upper bounded by:

$$ $\mathbb{E}[\mathcal{L}(t)]\leq\mathbb{E}[\mathcal{L}_{\mathrm{lp}}]e^{-ct}+L^{ \square}(1-e^{-ct})$ (15) $$

where $\mathcal{L}_{\mathrm{lp}}$ is the expected loss after linear probing. The coefficient $c=\mathbb{E}[\lambda_{\max}(D)]>0$ increases as $t_{\mathrm{lp}}$ increases
when we run linear probing in a finite time interval $t_{\mathrm{lp}}\in\left(0,\ln\left[3+\frac{k(\sigma^{2}-\beta)}{\|B_{0}X^{T}Y
\|^{2}}\right]\right)$.

###### Corollary 3.14

For simplicity, we assume that $X^{T}X=I_{d\times d}$. Suppose Assumption 3.3 and 3.4 hold. Then the two-phase method, first-linear-probing-then-finetuning (LP-FT), could achieve a tighter loss upper bound than linear probing or fine-tuning in expectation
if we first run linear probing for $t_{\mathrm{lp}}\in\left(0,\ln\left[3+\frac{k(\sigma^{2}-\beta)}{\|B_{0}X^{T}Y
\|^{2}}\right]\right)$.

By Corollary 3.14, we predict that: when we fix other hyperparameters (e.g. the total training time $T$), the performance of LP-FT depends on the noise scale $\sigma$.
If $\sigma$ is large enough such that $T<\ln\left[3+\frac{k(\sigma^{2}-\beta)}{\|B_{0}X^{T}Y\|^{2}}\right]$, then LP is the best; if $\sigma$ is small enough such that $\ln\left[3+\frac{k(\sigma^{2}-\beta)}{\|B_{0}X^{T}Y\|^{2}}\right]\leq 0$, then FT is the best; LP-FT could achieve the best performance when the noise scale is in a proper interval $\sigma^{2}\in\left(\beta-2\frac{\|B_{0}X^{T}Y\|^{2}}{k},\beta+(e^{T}-3)\frac{
\|B_{0}X^{T}Y\|^{2}}{k}\right)$. We qualitatively verify this prediction in our experiments in Section 4.3.

###### Proposition 3.13 (Linear probing accelerates fine-tuning)

###### Corollary 3.14

## 4 Experiments

### 4.1 Experimental settings

**Table 1: Selected results on four DP fine-tuning benchmarks and five model architectures. As suggested by Corollary 3.14, as the noise scale $\sigma$ increases, the best fine-tuning strategy changes from full fine-tuning (FT) (small $\sigma$, low privacy regime) to linear-probing-then-full-finetuning (LP-FT), to linear probing (LP) (large $\sigma$, high privacy regime). We sort the rows by the number of parameters of each model and the noise scale in an ascending order. For each experiment setting, we report average test accuracies with standard errors.**
| Model | # of param. | benchmark | $\sigma$ | LP | LP-FT | FT |
| --- | --- | --- | --- | --- | --- | --- |
| MobileNet-v3 | $1.52\times 10^{6}$ | CIFAR-10$\rightarrow$STL-10 | 0.0 | 71.12 (0.31) | 71.30 (0.11) | 77.02 (0.31) |
| MobileNet-v3 | $1.52\times 10^{6}$ | CIFAR-10$\rightarrow$STL-10 | 0.5 | 69.31 (0.03) | 71.96 (0.11) | 61.97 (0.14) |
| MobileNet-v3 | $1.52\times 10^{6}$ | CIFAR-10$\rightarrow$STL-10 | 1.0 | 69.25 (0.05) | 70.98 (0.14) | 60.02 (0.21) |
| MobileNet-v3 | $1.52\times 10^{6}$ | CIFAR-10$\rightarrow$STL-10 | 5.0 | 67.09 (0.18) | 65.59 (0.18) | 43.64 (0.17) |
| Mini-DeiT-Ti | $2.89\times 10^{6}$ | ImageNet-1K$\rightarrow$STL-10 | 0.0 | 95.74 (0.04) | 96.82 (0.08) | 96.17 (0.08) |
| Mini-DeiT-Ti | $2.89\times 10^{6}$ | ImageNet-1K$\rightarrow$STL-10 | 0.1 | 95.29 (0.05) | 96.09 (0.01) | 94.52 (0.16) |
| Mini-DeiT-Ti | $2.89\times 10^{6}$ | ImageNet-1K$\rightarrow$STL-10 | 5.0 | 93.17 (0.09) | 91.52 (0.33) | 53.94 (1.65) |
| ResNet-18 | $1.11\times 10^{7}$ | CIFAR-10$\rightarrow$STL-10 | 0.0 | 67.92 (0.22) | 67.71 (0.26) | 71.45 (0.36) |
| ResNet-18 | $1.11\times 10^{7}$ | CIFAR-10$\rightarrow$STL-10 | 0.5 | 63.92 (0.04) | 64.85 (0.11) | 52.17 (0.11) |
| ResNet-18 | $1.11\times 10^{7}$ | CIFAR-10$\rightarrow$STL-10 | 1.0 | 63.91 (0.02) | 64.68 (0.26) | 51.73 (0.15) |
| ResNet-18 | $1.11\times 10^{7}$ | CIFAR-10$\rightarrow$STL-10 | 3.0 | 63.52 (0.07) | 62.11 (0.69) | 47.87 (0.82) |
| ResNet-50 | $2.56\times 10^{7}$ | ImageNet-1K$\rightarrow$CIFAR-10 | 0.0 | 82.35 (0.07) | 95.54 (0.06) | 96.82 (0.01) |
| ResNet-50 | $2.56\times 10^{7}$ | ImageNet-1K$\rightarrow$CIFAR-10 | 0.1 | 61.31 (0.55) | 69.55 (0.44) | 69.93 (0.24) |
| ResNet-50 | $2.56\times 10^{7}$ | ImageNet-1K$\rightarrow$CIFAR-10 | 0.3 | 61.06 (0.27) | 66.13 (1.01) | 53.30 (0.36) |
| ResNet-50 | $2.56\times 10^{7}$ | ImageNet-1K$\rightarrow$CIFAR-10 | 0.5 | 61.12 (0.54) | 60.81 (0.72) | 40.70 (0.31) |
| ResNet-50 | $2.56\times 10^{7}$ | ImageNet-1K$\rightarrow$CIFAR-10 | 1.0 | 60.36 (0.49) | 59.83 (0.74) | 30.92 (1.12) |
| WRN-16-4 | $4.37\times 10^{7}$ | RandP$\rightarrow$CIFAR-10 | 9.3 | 61.2 (-) | 72.32 (-) | 69.1 (-) |

We run experiments on five deep learning models and four transfer learning benchmarks to verify if our theoretical prediction, the existence of concave utility curves, generalizes to deep neural networks and real datasets.
Each experimental setting comprises: (1) a model architecture, (2) a (larger) dataset for public pretraining, and (3) a (smaller) dataset as the private data for fine-tuning.
The benchmarks we use are:

- •
ImageNet-1K$\rightarrow$CIFAR-10. ImageNet-1K is a large-scale dataset. We initialize pretrained features of ResNet-50 from MoCo-v2 Chen et al. () and MoCo-v3 Chen* et al. (), trained on ImageNet-1K Russakovsky et al. () without privacy. We then privately fine-tune the ResNet-50 on CIFAR-10.
- •
ImageNet-1K$\rightarrow$STL-10. We pretrain a DeiT model on ImageNet then pretrain a Mini-DeiT-Ti model with weight distillation from the DeiT model Touvron et al. (); Zhang et al. (). After that, we privately fine-tune the Mini-DeiT-Ti model on STL-10 Coates et al. () for 20 epochs.
- •
CIFAR-10$\rightarrow$STL-10. We pretrain the feature extractor on CIFAR-10 Krizhevsky () using stochastic gradient descent without privacy mechanisms. Then we finetune the pretrained features and a randomly initialized linear head on STL-10. This benchmark has been studied in the context of domain adaptation French et al. (); Kumar et al. (). The training subset of STL-10 only contains 500 images. To align with the small scale fine-tuning data, we run the experiments with smaller and data-efficient models: MobileNet-v3 and ResNet-18.
- •
RandP$\rightarrow$CIFAR-10. To reproduce the results of Tang et al. () and verify the general existence of concave utility curves, we also consider a slightly non-standard pretraining protocol. We pretrain a wide residual network (WRN) Zagoruyko and Komodakis () on synthetic images generated by random diffusion processes. We follow the settings in Tang et al. ().

For each fine-tuning method, {linear probing (LP), full fine-tuning (FT), first-linear-probing-then-full-finetuning (LP-FT)}, we use a batch size of $64$ and sweep over 3 learning rates $\{0.1,0.05,0.025\}$. We fine-tune for $e=10$ epochs in total, employ early stopping, and select the optimal learning rate based on the accuracy of the in-distribution validation.
To demonstrate utility curves on LP-FT (e.g., Figure 2), we vary the number of epochs of linear probing from $e_{LP}=0$ to $e_{LP}=10$; all remaining epochs are allocated to full fine-tuning, i.e., $e_{FT}=10-e_{LP}$.
Note that full fine-tuning corresponds to $e_{LP}=0$ (the leftmost point of our plots), and linear probing corresponds to $e_{LP}=10$.
For experiments that simply compare LP, FT, and LP-FT (Table 1), we implement the LP-FT baseline by running $e_{LP}=5$ epochs of linear probing and $e_{FT}=5$ epochs of full fine-tuning.
For all private optimization steps, we apply DP-SGD (a variant of DP-GD that operates over minibatches in each iteration) as specified by Abadi et al. () (i.e., we do not vary the noise by layer). This was done to observe whether our predictions hold under the most common usage of private optimization with DP-SGD.
The test accuracy of each setting is an average over 3 random seeds.

### 4.2 Existence of non-monotonic utility curves

In this part, we empirically characterize the utility curve in more details.

DP settings. For the ImageNet-1K$\rightarrow$CIFAR-10 benchmark, we change the epochs allocated to linear probing (see Figure 2). As one increases the fraction of rounds used for linear probing, test accuracy increases up to a point, then decreases. This kind of utility curves also exist in the RandP$\rightarrow$CIFAR-10 benchmark with $\sigma=9.3$ Tang et al. ().

Non-DP settings. We fine-tune the pretrained ResNet-50 on CIFAR-10 without privacy mechanisms. The utility curves are monotonically decreasing as we increase the number of epochs allocated to linear probing
(Figure 2).

### 4.3 Examining the theoretical predictions

One prediction from our theory (Corollary 3.14) is that there exist two thresholds $0<c_{1}<c_{2}$ for the noise scale $\sigma$ such that we can qualitatively rank the expected empirical losses $\mathbb{E}[\mathcal{L}_{\mathrm{lp}}],\mathbb{E}[\mathcal{L}_{\mathrm{ft}}],
\mathbb{E}[\mathcal{L}_{\mathrm{lp-ft}}]$ of LP, FT, and LP-FT.
In this comparison, our LP-FT baseline consists of half linear probing epochs, and half full fine-tuning epochs.

- •
if $\sigma<c_{1}$, then $\mathbb{E}[\mathcal{L}_{\mathrm{ft}}]<\mathbb{E}[\mathcal{L}_{\mathrm{lp-ft}}]
<\mathbb{E}[\mathcal{L}_{\mathrm{lp}}]$, i.e. full fine-tuning is the best choice;
- •
if $\sigma\in(c_{1},c_{2})$, then $\mathbb{E}[\mathcal{L}_{\mathrm{lp-ft}}]<\min\{\mathbb{E}[\mathcal{L}_{\mathrm
{lp}}],\mathbb{E}[\mathcal{L}_{\mathrm{ft}}]\}$, i.e. first-linear-probing-then-full-finetuning is the best choice;
- •
if $\sigma>c_{2}$, then $\mathbb{E}[\mathcal{L}_{\mathrm{lp}}]<\mathbb{E}[\mathcal{L}_{\mathrm{lp-ft}}]
<\mathbb{E}[\mathcal{L}_{\mathrm{ft}}]$, i.e. linear probing is the best choice.

To qualitatively test this prediction, we sweep over different noise scales $\sigma$ and fix other hyperparameters in each benchmark and model architecture (see Table 1). As expected, among the three fine-tuning methods, FT almost always does the best under small noise scales (including the non-private setting where $\sigma=0$), LP-FT does the best under moderate noise scales, and LP does the best under large noise scales.
This trend holds across all four benchmarks we evaluated. The final row for WRN-16-4 was conducted only to reproduce the results of Tang et al. ().

## 5 Conclusion

Our work explores the impact of differentially-private (DP) fine-tuning strategies on model optimization in machine learning pipelines. We observe that the traditional approach of full fine-tuning, where all model parameters are updated, does not always yield the best test accuracy, even for in-distribution data. Instead, we explore a sequential fine-tuning strategy that combines linear probing and full fine-tuning (LP-FT), and empirically demonstrate its superiority over individual approaches.

Through theoretical analysis on a two-layer linear network regression task, we provide evidence that DP linear probing can outperform DP full fine-tuning, particularly when privacy budgets are limited. For the two-phase sequential approach (LP-FT), our theoretical results reveal a utility curve, suggesting LP-FT can outperform either of linear probing or full fine-tuning alone, even when pre-trained features are highly effective.
Our experiments across various fine-tuning benchmarks and model architectures consistently demonstrate the existence of utility curves, reinforcing the validity of our theoretical model.

A limitation of our analysis is that we do not have lower bounds on the convergence rates for linear probing or full fine-tuning. This prevents us from being able to conclusively state that one fine-tuning method is better than the other in any regime. Developing a lower bound for these convergence rates is an interesting question for future work.

## Appendix A Proofs

### A.1 Itô’s formula and its consequences

###### Theorem A.1 (Itô’s formula)

Let $X_{t}$ be a $\mathbb{R}^{n}$-valued Itô process satisfying the stochastic differential equation $\partial X_{t}=A_{1}(t,X_{t})\partial t+A_{2}(t,X_{t})\partial W_{t}$ with $A_{1}(t,X_{t})$ being $\mathbb{R}^{n}$-valued, $A_{2}(t,X_{t})$ being $M_{m,n}(\mathbb{R})$-valued, and $W_{t}$ being a standard $n$-dimensional brownian motion. Let $f:[0,\infty)\times\mathbb{R}^{n}\rightarrow\mathbb{R}$ be a function with continuous partial derivatives. Then $Y_{t}:=f(t,X_{t})$ is also an Itô process, and its stochastic differential equation is

$$ $\partial Y_{t}=\frac{\partial f(t,X_{t})}{\partial t}\partial t+\langle\nabla f (t,X_{t}),A_{1}(t,X_{t})\partial t+A_{2}(t,X_{t})\partial W_{t}\rangle+\frac{1 }{2}\langle A_{2}(t,X_{t})\partial W_{t},H_{f}A_{2}(t,X_{t})\partial W_{t}\rangle$ (16) $$

where $H_{f}$ is the Hessian matrix of $f$ over $X_{t}$ defined as $(H_{f})_{ij}=\frac{\partial^{2}f}{\partial(X_{t})_{i}\partial(X_{t})_{j}}$ and $(X_{t})_{i}$ denotes the i-th entry of random vector $X_{t}$.

###### Corollary A.2 (Loss dynamics during linear probing)

During linear probing (Equation (4)), the stochastic differential equation describing the loss dynamics is

$$ $\partial\mathcal{L}_{\mathrm{lp}}=-(B_{0}^{T}v-X^{T}Y)^{T}B_{0}^{T}B_{0}(B_{0} ^{T}v-X^{T}Y)\partial t+\sqrt{2\sigma^{2}}(B_{0}^{T}v-X^{T}Y)^{T}B_{0}^{T} \partial W_{t}+k\sigma^{2}\partial t.$ (17) $$

###### Proof of Corollary A.2 .

By Itô’s formula (Equation (A.1)), the loss dynamics is

$$ $\displaystyle\partial\mathcal{L}_{\mathrm{lp}}=$ $\displaystyle\partial\frac{1}{2}\|XB_{0}^{T}v-Y\|^{2}$ (18) $\displaystyle=$ $\displaystyle(XB_{0}^{T}v-Y)^{T}XB_{0}^{T}\partial v+\frac{1}{2}(\partial v)^{ T}B_{0}X^{T}XB_{0}^{T}(\partial v)$ (19) $\displaystyle=$ $\displaystyle(XB_{0}^{T}v-Y)^{T}XB_{0}^{T}\partial v+\frac{1}{2}(\partial v)^{ T}(\partial v)$ (20) $\displaystyle=$ $\displaystyle(XB_{0}^{T}v-Y)^{T}XB_{0}^{T}[-B_{0}X^{T}(XB_{0}^{T}v-Y)\partial t +\sqrt{2\sigma^{2}}\partial W_{t}]+k\sigma^{2}\partial t$ (21) $\displaystyle=$ $\displaystyle(B_{0}^{T}v-X^{T}Y)^{T}B_{0}^{T}[-B_{0}(B_{0}^{T}v-X^{T}Y) \partial t+\sqrt{2\sigma^{2}}\partial W_{t}]+k\sigma^{2}\partial t$ (22) $\displaystyle=$ $\displaystyle-(B_{0}^{T}v-X^{T}Y)^{T}B_{0}^{T}B_{0}(B_{0}^{T}v-X^{T}Y)\partial t +\sqrt{2\sigma^{2}}(B_{0}^{T}v-X^{T}Y)^{T}B_{0}^{T}\partial W_{t}+k\sigma^{2}\partial t$ (23) $$

∎

###### Corollary A.3 (Loss dynamics during fine-tuning)

During fine-tuning (Equation (5)), the stochastic differential equation describing the loss dynamics is

$$ $\displaystyle\partial\mathcal{L}_{\mathrm{ft}}=$ $\displaystyle-(B^{T}v-X^{T}Y)^{T}B^{T}B(B^{T}v-X^{T}Y)\partial t+(B^{T}v-X^{T} Y)^{T}B^{T}\sqrt{2\sigma^{2}}\partial W_{t}$ (24) $\displaystyle-(B^{T}v-X^{T}Y)^{T}(B^{T}v-X^{T}Y)v^{T}v\partial t+(B^{T}v-X^{T} Y)^{T}(\sqrt{2\sigma^{2}}\partial W_{t}^{\prime})v$ (25) $\displaystyle+\sigma^{2}\|B\|_{F}^{2}\partial t+\sigma^{2}d\|v\|_{2}^{2} \partial t.$ (26) $$

where we use $\partial$ as the differential sign and use $d$ as the data input dimension.

###### Proof of Corollary A.3 .

$$ $\displaystyle\partial\mathcal{L}_{\mathrm{ft}}=$ $\displaystyle\partial\frac{1}{2}\|XB^{T}v-Y\|^{2}$ (27) $\displaystyle=$ $\displaystyle\frac{1}{2}\left\langle\nabla_{v}\|(XB^{T}v-Y)\|^{2},\partial v \right\rangle+\frac{1}{2}\left\langle\nabla_{B}\|(XB^{T}v-Y)\|^{2},\mathrm{vec }(\partial B)\right\rangle$ (28) $\displaystyle+\frac{1}{4}(\partial v)^{T}H_{\|(XB^{T}v-Y)\|^{2}}(\partial v)+ \frac{1}{4}[\mathrm{vec}(\partial B)]^{T}H_{\|(XB^{T}v-Y)\|^{2}}\mathrm{vec}( \partial B)$ (29) $\displaystyle=$ $\displaystyle(XB^{T}v-Y)^{T}XB^{T}\partial v+(XB^{T}v-Y)^{T}X(\partial B)^{T}v$ (30) $\displaystyle+\frac{1}{2}(\partial v)^{T}BX^{T}XB^{T}(\partial v)+\frac{1}{2}[ \mathrm{vec}(\partial B)]^{T}\begin{bmatrix}v_{1}\\ 0\\ \vdots\\ v_{k}\end{bmatrix}\underbrace{\begin{bmatrix}v_{1}&0&\cdots&v_{k}\end{bmatrix} }_{d\times k}\mathrm{vec}(\partial B)$ (36) $\displaystyle=$ $\displaystyle-(B^{T}v-X^{T}Y)^{T}B^{T}B(B^{T}v-X^{T}Y)\partial t+(B^{T}v-X^{T} Y)^{T}B^{T}\sqrt{2\sigma^{2}}\partial W_{t}$ (37) $\displaystyle-(B^{T}v-X^{T}Y)^{T}(B^{T}v-X^{T}Y)v^{T}v\partial t+(B^{T}v-X^{T} Y)^{T}(\sqrt{2\sigma^{2}}\partial W_{t}^{\prime})v$ (38) $\displaystyle+\sigma^{2}\mathrm{trace}(BB^{T})\partial t+\sigma^{2}d\|v\|^{2}\partial t$ (39) $\displaystyle=$ $\displaystyle-(B^{T}v-X^{T}Y)^{T}B^{T}B(B^{T}v-X^{T}Y)\partial t+(B^{T}v-X^{T} Y)^{T}B^{T}\sqrt{2\sigma^{2}}\partial W_{t}$ (40) $\displaystyle-(B^{T}v-X^{T}Y)^{T}(B^{T}v-X^{T}Y)v^{T}v\partial t+(B^{T}v-X^{T} Y)^{T}(\sqrt{2\sigma^{2}}\partial W_{t}^{\prime})v$ (41) $\displaystyle+\sigma^{2}\|B\|_{F}^{2}\partial t+\sigma^{2}d\|v\|_{2}^{2}\partial t$ (42) $$

∎

###### Remark A.4 (Noise effects on linear networks) .

In the loss dynamics of fine-tuning (Corollary A.3), the noise induced deterministic terms

$$ $\sigma^{2}(\|B\|_{F}^{2}+d\|v\|_{2}^{2})\partial t$ $$

does not explicitly depend on the linear head size $k$. We do a sanity check for this result in a discretized setting (so that we skip Itô’s lemma and stochastic calculus). Say we inject noise $\Delta B$ to $B$, where $\Delta B$ is a $k\times d$-matrix, and its entries are independent and follow Gaussian distribution $\mathcal{N}(0,\sigma)$. Then the expectation of the perturbed loss is:

$$ $\displaystyle\mathbb{E}[\mathcal{L}]=$ $\displaystyle\frac{1}{2}\mathbb{E}[\|X(B+\Delta B)^{T}v-Y\|^{2}]$ (43) $\displaystyle=$ $\displaystyle\frac{1}{2}\|XB^{T}v-Y\|^{2}+\mathbb{E}[(XB^{T}v-Y)^{T}X(\Delta B )^{T}v]+\frac{1}{2}\mathbb{E}[v^{T}\Delta B(\Delta B)^{T}v]$ (44) $\displaystyle=$ $\displaystyle\frac{1}{2}\|XB^{T}v-Y\|^{2}+\frac{1}{2}\mathbb{E}[v^{T}\Delta B( \Delta B)^{T}v]$ (45) $\displaystyle=$ $\displaystyle\frac{1}{2}\|XB^{T}v-Y\|^{2}+\frac{1}{2}\sigma^{2}\cdot d\cdot\|v \|^{2}$ (46) $$

###### Theorem A.1 (Itô’s formula)

###### Corollary A.2 (Loss dynamics during linear probing)

###### Proof of Corollary A.2 .

###### Corollary A.3 (Loss dynamics during fine-tuning)

###### Proof of Corollary A.3 .

###### Remark A.4 (Noise effects on linear networks) .

### A.2 Linear probing loss upper bound

###### Lemma A.5 ((Strong) convexity of linear probing phase)

The empirical risk $\mathcal{L}=\frac{1}{2}\sum_{i=1}^{n}\ell(f(x_{i}),y_{i})$ is $1$-strongly convex.

###### Lemma A.6 (Initial loss before linear probing)

If Assumption 3.3 holds, then the expected empirical risk before linear probing is

$$ $\mathbb{E}[\mathcal{L}_{0}]=\frac{1}{2}(k\beta+\|Y\|^{2})$ (47) $$

###### Proof of Lemma A.6 .

We initialize the linear head with a Gaussian distribution $\mathcal{N}(0,\beta I_{k\times k})$. So the expected initial loss is:

$$ $\displaystyle\mathbb{E}[\mathcal{L}_{0}]=$ $\displaystyle\frac{1}{2}\mathbb{E}[\|XB_{0}^{T}v_{0}-Y\|^{2}]$ (48) $\displaystyle=$ $\displaystyle\frac{1}{2}\mathbb{E}[v_{0}^{T}B_{0}X^{T}XB_{0}^{T}v_{0}+Y^{T}Y-2 Y^{T}XB_{0}^{T}v_{0}]$ (49) $\displaystyle=$ $\displaystyle\frac{1}{2}\mathbb{E}[v_{0}^{T}B_{0}B_{0}^{T}v_{0}+Y^{T}Y]$ (50) $\displaystyle=$ $\displaystyle\frac{1}{2}\mathbb{E}[v_{0}^{T}v_{0}+Y^{T}Y]=\frac{1}{2}(k\beta+ \|Y\|^{2})$ (51) $$

∎

###### Theorem A.7 (Expected loss upper bound of linear probing)

The expected empirical risk in linear probing is upper bounded by

$$ $\mathbb{E}[\mathcal{L}_{\mathrm{lp}}(t)]\leq e^{-t}\mathbb{E}[\mathcal{L}_{0}] +(1-e^{-t})(\gamma+k\sigma^{2})$ (52) $$

###### Proof of Theorem 3.10 .

By Lemma A.5, $\mathcal{L}$ is $1$-strongly convex, we have the Lojasiewicz inequality

$$ $\displaystyle\mathcal{L}(v)-\{\min_{v}\mathcal{L}\}\leq\frac{1}{2}\|\nabla_{v} \mathcal{L}(v)\|_{2}^{2}$ (53) $$

For simplicity, we denote $\mathbb{E}[\mathcal{L}]:=\hat{\mathcal{L}}$. Consider the Langevin diffusion in Equation (4) when $\mathcal{L}(v)-\{\min_{v}\mathcal{L}\}-k\sigma^{2}>0$, by Corollary A.2:

$$ $\displaystyle\partial\mathcal{L}(v)=$ $\displaystyle\langle\nabla_{v}\mathcal{L}(v),-\nabla_{v}\mathcal{L}(v)\partial t +\sqrt{2\sigma^{2}}\partial W_{t}\rangle+k\sigma^{2}\partial t$ (54) $\displaystyle\partial\mathcal{L}(v)\leq$ $\displaystyle-\|\nabla_{v}\mathcal{L}(v)\|_{2}^{2}\partial t+\langle\nabla_{v} \mathcal{L}(v),\sqrt{2\sigma^{2}}\partial W_{t}\rangle+k\sigma^{2}\partial t$ (55) //By Lojasiewicz inequality (56) $\displaystyle\partial\mathcal{L}(v)\leq$ $\displaystyle(-\mathcal{L}(v)+\{\min_{v}\mathcal{L}\})\partial t+\langle\nabla _{v}\mathcal{L}(v),\sqrt{2\sigma^{2}}\partial W_{t}\rangle+k\sigma^{2}\partial t$ (57) $\displaystyle\partial(\mathbb{E}[\mathcal{L}(v)]-\{\min_{v}\mathcal{L}\}-k \sigma^{2})\leq$ $\displaystyle-(\mathbb{E}[\mathcal{L}(v)]-\{\min_{v}\mathcal{L}\})\partial t+k \sigma^{2}\partial t$ (58) $\displaystyle\partial(\hat{\mathcal{L}}-\{\min_{v}\mathcal{L}\}-k\sigma^{2})\leq$ $\displaystyle-(\hat{\mathcal{L}}-\{\min_{v}\mathcal{L}\}-k\sigma^{2})\partial t$ (59) $\displaystyle\text{//When }\hat{\mathcal{L}}-\{\min_{v}\mathcal{L}\}-k\sigma^{ 2}>0$ (60) $\displaystyle\partial\ln|\hat{\mathcal{L}}-\{\min_{v}\mathcal{L}\}-k\sigma^{2}|\leq$ $\displaystyle-1\partial t$ (61) $\displaystyle\ln|\hat{\mathcal{L}}-\{\min_{v}\mathcal{L}\}-k\sigma^{2}|\leq$ $\displaystyle\ln|\widehat{\mathcal{L}(v_{0})}-\{\min_{v}\mathcal{L}\}-k\sigma^ {2}|-t$ (62) $\displaystyle\hat{\mathcal{L}}-\{\min_{v}\mathcal{L}\}-k\sigma^{2}\leq$ $\displaystyle e^{-t}(\widehat{\mathcal{L}(v_{0})}-\{\min_{v}\mathcal{L}\}-k \sigma^{2})$ (63) $\displaystyle\hat{\mathcal{L}}\leq$ $\displaystyle e^{-t}(\widehat{\mathcal{L}(v_{0})}-\{\min_{v}\mathcal{L}\}-k \sigma^{2})+\{\min_{v}\mathcal{L}\}+k\sigma^{2}$ (64) $\displaystyle\hat{\mathcal{L}}\leq$ $\displaystyle e^{-t}\widehat{\mathcal{L}(v_{0})}+(1-e^{-t})(\{\min_{v}\mathcal {L}\}+k\sigma^{2})$ (65) $\displaystyle\hat{\mathcal{L}}\leq$ $\displaystyle e^{-t}\widehat{\mathcal{L}(v_{0})}+(1-e^{-t})(\gamma+k\sigma^{2})$ (66) $$

∎

###### Corollary A.8 (Expected loss upper bound of linear probing from Kaiming initialization)

If Assumption 3.3 holds, then the expected loss is upper bounded by

$$ $\mathbb{E}[\mathcal{L}_{\mathrm{lp}}(t)]\leq\frac{1}{2}(k\beta+\|Y\|^{2})e^{-t }+(1-e^{-t})(\gamma+k\sigma^{2})$ (67) $$

###### Proof of Corollary A.8 .

The result is immediate when we combine Lemma A.6 and Theorem 3.10.
∎

###### Lemma A.5 ((Strong) convexity of linear probing phase)

###### Lemma A.6 (Initial loss before linear probing)

###### Proof of Lemma A.6 .

###### Theorem A.7 (Expected loss upper bound of linear probing)

###### Proof of Theorem 3.10 .

###### Corollary A.8 (Expected loss upper bound of linear probing from Kaiming initialization)

###### Proof of Corollary A.8 .

### A.3 Imbalance matrix from linear probing

###### Lemma A.9 (Eigenvalues of imbalance matrix at the beginning of fine-tuning)

During the linear probing phase (Equation (4)), for the imbalance matrix defined in Definition 3.5,

- 1.
the minimum eigenvalue of the imbalance matrix is always $-1$;
- 2.
other eigenvalues evolve in this way:
$\mathbb{E}[\lambda]=\mathbb{E}\left[\|v\|_{2}^{2}\right]-1\geq-1$
(68)

###### Proof of Lemma A.9 .

Consider any eigenpair $(\lambda,u)$ of matrix $D$, we have

$$ $\displaystyle Du=$ $\displaystyle\lambda u$ (69) $\displaystyle(vv^{T}-B_{0}B_{0}^{T})u=$ $\displaystyle\lambda u$ (70) $\displaystyle(vv^{T}-I_{k\times k})u=$ $\displaystyle\lambda u$ (71) $\displaystyle(v^{T}u)v=$ $\displaystyle(\lambda+1)u$ (72) $$

We can take any $u\perp v$ and $(u,-1)$ is an eigenpair of $D$. So $-1$ is always an eigenvalue of $D$. We need to discuss two different cases here:

- 1.
If $\lambda=-1$, we only know that $u\perp v$.
- 2.
If $\lambda\not=-1$, then $v$ and $u$ are parallel. Say $u=\alpha v$, then
$\displaystyle u=$
$\displaystyle\frac{v^{T}u}{\lambda+1}v$
(74)
$\displaystyle\alpha v=$
$\displaystyle\frac{\alpha\|v\|_{2}^{2}}{\lambda+1}v$
(75)
$\displaystyle\Longrightarrow\lambda=$
$\displaystyle\|v\|_{2}^{2}-1\geq-1$
(76)

∎

###### Proposition A.10 (Expected eigenvalue of imbalance matrix at the beginning of fine-tuning)

Say we run linear probing for time $t$. If Assumption 3.3 holds, then for the imbalance matrix defined in Definition 3.5, we have

$$ $\mathbb{E}[\|v\|^{2}]=k\beta e^{-2t}+2\|B_{0}X^{T}Y\|^{2}(e^{-t}-e^{-2t})+(\|B _{0}X^{T}Y\|^{2}+k\sigma^{2})(1-e^{-2t})$ (77) $$

throughout the linear probing process. Then by Lemma A.9, for those eigenvalues not equal to $-1$, we have

$$ $\mathbb{E}[\lambda]=\mathbb{E}\left[\|v\|_{2}^{2}\right]-1=k\beta e^{-2t}+2\|B _{0}X^{T}Y\|^{2}(e^{-t}-e^{-2t})+(\|B_{0}X^{T}Y\|^{2}+k\sigma^{2})(1-e^{-2t})-1$ (78) $$

at the beginning of fine-tuning after linear probing.

###### Proof of Proposition A.10 .

By Equation (4), the Langevin diffusion of linear probing is:

$$ $\displaystyle\partial v=-B_{0}X^{T}(XB_{0}^{T}v-Y)\partial t+\sqrt{2\sigma^{2} }\partial W_{t}=-v\partial t+B_{0}X^{T}Y\partial t+\sqrt{2\sigma^{2}}\partial W _{t}$ (79) $$

We consider the evolution of $v^{T}v$: by Itô’s formula (Equation (A.1))

$$ $\displaystyle\partial v^{T}v=$ $\displaystyle 2v^{T}\partial v+(\partial v)^{T}I_{k}(\partial v)$ (80) $\displaystyle\partial v^{T}v=$ $\displaystyle-2v^{T}(v-B_{0}X^{T}Y)\partial t+2v^{T}\sqrt{2\sigma^{2}}\partial W _{t}+2k\sigma^{2}\partial t$ (81) $\displaystyle\partial v^{T}v=$ $\displaystyle(-2v^{T}v+2v^{T}B_{0}X^{T}Y)\partial t+2v^{T}\sqrt{2\sigma^{2}} \partial W_{t}+2k\sigma^{2}\partial t$ (82) $$

To solve the above equation, we need to solve the dynamics of $v^{T}B_{0}X^{T}Y$:

$$ $\displaystyle\partial Y^{T}XB_{0}^{T}v=$ $\displaystyle-Y^{T}XB_{0}^{T}(v-B_{0}X^{T}Y)\partial t+\sqrt{2\sigma^{2}} \partial W_{t}$ (83) $\displaystyle\partial\mathbb{E}[Y^{T}XB_{0}^{T}v]=$ $\displaystyle-\mathbb{E}[Y^{T}XB_{0}^{T}v]dt+\|B_{0}X^{T}Y\|^{2}\partial t$ (84) $\displaystyle\frac{\partial}{\partial t}\mathbb{E}[Y^{T}XB_{0}^{T}v-\|B_{0}X^{ T}Y\|^{2}]=$ $\displaystyle-\mathbb{E}[Y^{T}XB_{0}^{T}v-\|B_{0}X^{T}Y\|^{2}]$ (85) $\displaystyle\frac{\partial}{\partial t}\ln|\mathbb{E}[Y^{T}XB_{0}^{T}v-\|B_{0 }X^{T}Y\|^{2}]|=$ $\displaystyle-1$ (86) $\displaystyle|\mathbb{E}[Y^{T}XB_{0}^{T}v_{t}-\|B_{0}X^{T}Y\|^{2}]|=$ $\displaystyle|\mathbb{E}[Y^{T}XB_{0}^{T}v_{0}-\|B_{0}X^{T}Y\|^{2}]|\cdot\exp(-t)$ (87) $$

When Assumption 3.3 holds, $\mathbb{E}[Y^{T}XB_{0}^{T}v_{0}]=0$. Then

$$ $\displaystyle|\mathbb{E}[Y^{T}XB_{0}^{T}v_{t}-\|B_{0}X^{T}Y\|^{2}]|=$ $\displaystyle|\mathbb{E}[Y^{T}XB_{0}^{T}v_{0}-\|B_{0}X^{T}Y\|^{2}]|\cdot\exp(-t)$ (88) $\displaystyle\mathbb{E}[\|B_{0}X^{T}Y\|^{2}-Y^{T}XB_{0}^{T}v_{t}]=$ $\displaystyle\mathbb{E}[\|B_{0}X^{T}Y\|^{2}-Y^{T}XB_{0}^{T}v_{0}]\cdot\exp(-t)$ (89) $$

So we can rewrite Equation (82) as:

$$ $\displaystyle\partial\mathbb{E}[\|v\|^{2}]=$ $\displaystyle(-2\mathbb{E}[\|v\|^{2}]+2\mathbb{E}[v^{T}B_{0}X^{T}Y])\partial t +2k\sigma^{2}\partial t$ (90) $\displaystyle\partial\mathbb{E}[\|v\|^{2}]=$ $\displaystyle(-2\mathbb{E}[\|v\|^{2}]+2(\mathbb{E}[\|B_{0}X^{T}Y\|^{2}-Y^{T}XB _{0}^{T}v_{0}]\cdot\exp(-t)+\|B_{0}X^{T}Y\|^{2}))\partial t+2k\sigma^{2}\partial t$ (91) $\displaystyle\frac{1}{2}\frac{\partial}{\partial t}\mathbb{E}[\|v\|^{2}]=$ $\displaystyle-\mathbb{E}[\|v\|^{2}]+\mathbb{E}[\|B_{0}X^{T}Y\|^{2}-Y^{T}XB_{0} ^{T}v_{0}]\cdot\exp(-t)+(\|B_{0}X^{T}Y\|^{2}+k\sigma^{2})$ (92) $$

Let $a_{1}=\mathbb{E}[\|B_{0}X^{T}Y\|^{2}-Y^{T}XB_{0}^{T}v_{0}],a_{2}=\|B_{0}X^{T}Y
\|^{2}+k\sigma^{2},f(t)=\mathbb{E}[\|v\|^{2}]$ and rewrite the above equation:

$$ $\displaystyle\frac{1}{2}f^{\prime}(t)+f(t)=$ $\displaystyle a_{1}e^{-t}+a_{2}$ (93) $\displaystyle f^{\prime}(t)+2f(t)=$ $\displaystyle 2a_{1}e^{-t}+2a_{2}$ (94) $\displaystyle e^{2t}f^{\prime}(t)+2e^{2t}f(t)=$ $\displaystyle 2a_{1}e^{t}+2a_{2}e^{2t}$ (95) $\displaystyle e^{2t}f(t)\bigg{|}_{0}^{t}=$ $\displaystyle(2a_{1}e^{t}+a_{2}e^{2t})\bigg{|}_{0}^{t}$ (96) $\displaystyle e^{2t}f(t)=$ $\displaystyle f(0)+2a_{1}(e^{t}-1)+a_{2}(e^{2t}-1)$ (97) $\displaystyle f(t)=$ $\displaystyle f(0)e^{-2t}+2a_{1}(e^{-t}-e^{-2t})+a_{2}(1-e^{-2t})$ (98) $$

by Assumption 3.3 we have $f(0)=k\beta$ and $a_{1}=\|B_{0}X^{T}Y\|^{2}$.
∎

###### Lemma A.11 (Imbalance matrix in fine-tuning)

During fine-tuning (Equation (5)), the imbalance matrix $D$ in Definition 3.5 evolves as

$$ $\frac{\partial}{\partial t}\mathbb{E}[D]=(1-d)\sigma^{2}I_{k\times k}$ (99) $$

where $d$ is the dimension of data inputs ($B\in\mathbb{R}^{k\times d}$).

###### Proof of Lemma 3.6 .

We prove this lemma by analyzing the infinitesimal generator $A$ of imbalance matrix $D$ at any time:

$$ $\displaystyle A(D)_{ij}:=$ $\displaystyle\lim\limits_{t\downarrow 0}\frac{\mathbb{E}^{D}[(D(t))_{ij}]-{(D) }_{ij}}{t}$ (100) $\displaystyle=$ $\displaystyle 0+\sigma^{2}\sum_{i^{\prime}\in[k]}\sum_{j^{\prime}\in[k]} \mathbf{1}[i^{\prime}=j^{\prime}=i=j]$ (101) $\displaystyle-\sigma^{2}\sum_{i^{\prime}\in[k],j^{\prime}\in[d]}\sum_{i^{ \prime\prime}\in[k],j^{\prime\prime}\in[d]}\mathbf{1}[i^{\prime}=i^{\prime \prime}=i=j\text{ and }j^{\prime}=j^{\prime\prime}]$ (102) $$

the generator is zero for $i\not=j$. So we can just consider the case where $i=j$.

$$ $\displaystyle A(D)_{ii}=$ $\displaystyle\sigma^{2}\sum_{i^{\prime}\in[k]}\sum_{j^{\prime}\in[k]}\mathbf{1 }[i^{\prime}=j^{\prime}=i]$ (103) $\displaystyle-\sigma^{2}\sum_{i^{\prime}\in[k],j^{\prime}\in[d]}\sum_{i^{ \prime\prime}\in[k],j^{\prime\prime}\in[d]}\mathbf{1}[i^{\prime}=i^{\prime \prime}=i\text{ and }j^{\prime}=j^{\prime\prime}]$ (104) $\displaystyle=$ $\displaystyle(1-d)\sigma^{2}$ (105) $$

∎

###### Lemma A.12 (Monotonic eigenvalue of imbalance matrix in fine-tuning)

Denote $D_{\mathrm{lp}}$ as the imbalance matrix right after linear probing phase. All eigenvalues of the imbalance matrix are decreasing in expectation during fine-tuning. Specifically,

$$ $\mathbb{E}[\lambda(D)]=\mathbb{E}[\lambda(D_{\mathrm{lp}})]+(1-d)\sigma^{2}t$ (106) $$

where $t$ is the time-span of fine-tuning process.

###### Proof of Lemma A.12 .

Pick any eigenpair $(\lambda,u)$ of imbalance matrix $D$ (Definition 3.5) such that $\|u\|_{2}=1$. By Itô’s lemma (Equation (A.1)):

$$ $\displaystyle\partial\lambda=$ $\displaystyle u^{T}(\partial D)u+u^{T}(\partial D)(\lambda I-D)^{{\dagger}}( \partial D)u^{T}$ (107) $\displaystyle=$ $\displaystyle(1-d)\sigma^{2}\|u\|_{2}^{2}\partial t+\partial M_{t}+(1-d)^{2} \sigma^{4}u^{T}(\lambda I-D)^{{\dagger}}u^{T}$ (108) $\displaystyle=$ $\displaystyle(1-d)\sigma^{2}\partial t+\partial M_{t}+(1-d)^{2}\sigma^{4}u^{T} (\lambda I-D)^{{\dagger}}u^{T}$ (109) $$

where $M_{t}$ is the martingale induced by the Brownian noise and $(\cdot)^{{\dagger}}$ denotes the pseudo inverse of a certain matrix. Say the the singular value decomposition (SVD) of $D$ is

$$ $D=U\Sigma U^{T}=U\begin{bmatrix}\lambda_{1}&&\boldsymbol{0}\\ &\lambda_{2}\\ \boldsymbol{0}&&\ddots\end{bmatrix}U^{T}$ (110) $$

where we have $\lambda\in\mathrm{diag}\Sigma$ and $u$ being a column vector in $U$. So we can write the SVD of $(\lambda I-D)$ as:

$$ $\lambda I-D=V\Sigma^{\prime}V^{T}=V\begin{bmatrix}\lambda-\lambda_{1}&& \boldsymbol{0}\\ &\lambda-\lambda_{2}\\ \boldsymbol{0}&&\ddots\end{bmatrix}V^{T}$ (111) $$

where we obtain $V$ by removing $u$ in the columns of $U$ and we obtain $\Sigma^{\prime}$ by removing $\lambda$ in $\Sigma$. Then the pseudo inverse of $(\lambda I-D)$ is

$$ $(\lambda I-D)^{{\dagger}}=V\Sigma^{\prime}V^{T}=V\begin{bmatrix}\frac{1}{ \lambda-\lambda_{1}}&&\boldsymbol{0}\\ &\frac{1}{\lambda-\lambda_{2}}\\ \boldsymbol{0}&&\ddots\end{bmatrix}V^{T}$ (112) $$

Since $U$ is orthogonal, we shall have $V^{T}u=\boldsymbol{0}$. Then we can rewrite the stochastic dynamics of $D$ as:

$$ $\frac{\partial}{\partial t}\mathbb{E}[\lambda]=(1-d)\sigma^{2}$ (113) $$

∎

###### Lemma A.9 (Eigenvalues of imbalance matrix at the beginning of fine-tuning)

###### Proof of Lemma A.9 .

###### Proposition A.10 (Expected eigenvalue of imbalance matrix at the beginning of fine-tuning)

###### Proof of Proposition A.10 .

###### Lemma A.11 (Imbalance matrix in fine-tuning)

###### Proof of Lemma 3.6 .

###### Lemma A.12 (Monotonic eigenvalue of imbalance matrix in fine-tuning)

###### Proof of Lemma A.12 .

### A.4 Fine-tuning loss

###### Lemma A.13 (Bounding the norm of linear head $\|v\|_{2}^{2}$ )

During fine-tuning (Equation (5)), we can bound the norm of $\|v\|_{2}^{2}$ with the imbalance matrix $D$ in Definition 3.5 as

$$ $\frac{\underline{\lambda}+\sqrt{\underline{\lambda}^{2}+4\|w\|^{2}}}{2}\leq\|v \|_{2}^{2}\leq\frac{\bar{\lambda}+\sqrt{\bar{\lambda}^{2}+4\|w\|^{2}}}{2}$ (114) $$

where we denote $\underline{\lambda}=\lambda_{\min}(\hat{D}),\bar{\lambda}=\lambda_{\max}(\hat{
D})$.

###### Proof of Lemma A.13 .

Given the information of imbalance matrix, we can bound the linear head norm. Denote $\underline{\lambda}=\lambda_{\min}(D),\bar{\lambda}=\lambda_{\max}(D)$. Denote $w=B^{T}v$ and multiply $D$ with $v$ on both sides:

$$ $\displaystyle v^{T}Dv=$ $\displaystyle(v^{T}v)^{2}-(v^{T}B)(B^{T}v)$ (115) $\displaystyle v^{T}Dv=$ $\displaystyle\|v\|_{2}^{4}-\|w\|_{2}^{2}$ (116) $$

We have a range for the Rayleigh quotient: $\frac{x^{T}Dx}{x^{T}x}\in[\underline{\lambda},\bar{\lambda}]$. So we obtain two inequalities:

$$ $\displaystyle\begin{cases}\|v\|_{2}^{4}-\|w\|_{2}^{2}\geq\underline{\lambda}\| v\|_{2}^{2}\\ \|v\|_{2}^{4}-\|w\|_{2}^{2}\leq\bar{\lambda}\|v\|_{2}^{2}\end{cases}$ (117) $\displaystyle=$ $\displaystyle\begin{cases}\|v\|^{4}-\underline{\lambda}\|v\|^{2}-\|w\|^{2}\geq 0 \\ \|v\|^{4}-\bar{\lambda}\|v\|^{2}-\|w\|^{2}\leq 0\end{cases}$ (118) $$

To get a lower bound of $v$, we can solve two quadratic inequalities.
For the first quadratic equation, since the smaller root is non-positive, $\underline{\lambda}-\sqrt{\underline{\lambda}^{2}+4\|w\|^{2}}\leq 0$, we just bound $\|v\|^{2}$ with the larger root:

$$ $\|v\|^{2}\geq\frac{\underline{\lambda}+\sqrt{\underline{\lambda}^{2}+4\|w\|^{2 }}}{2}$ (119) $$

similarly, for the second quadratic equation, we obtain an upper bound for $\|v\|^{2}$ with the right-side zero point:

$$ $\|v\|^{2}\leq\frac{\bar{\lambda}+\sqrt{\bar{\lambda}^{2}+4\|w\|^{2}}}{2}$ (120) $$

∎

###### Lemma A.14 (Bounding eigenvalues of $B^{T}B$ (re-stated from Min et al. [ 2023b ] ))

During fine-tuning (Equation (5)), we can bound any nonzero eigenvalue $\lambda_{i}$ of $B^{T}B$ as

$$ $\lambda_{i}\in\left[\frac{-\bar{\lambda}+\sqrt{\bar{\lambda}^{2}+4(z_{i}^{T}w) ^{2}}}{2},\frac{-\underline{\lambda}+\sqrt{\underline{\lambda}^{2}+4(z_{i}^{T} w)^{2}}}{2}\right]$ (121) $$

where we use the imbalance matrix $D$ in Definition 3.5 and denote

$$ $\begin{cases}\bar{\lambda}=\lambda_{\max}(D)\\ \underline{\lambda}=\lambda_{\min}(D)\end{cases}$ (122) $$

###### Proof of Lemma A.14 .

The proof of this lemma follows the proof of Lemma 3 in Min et al. []. $B^{T}B$ is symmetric and positive semidefinite ($x^{T}B^{T}Bx=\|Bx\|_{2}^{2}\geq 0$). So every eigenvalue of $B^{T}B$ is non-negative.

$D$ has at most one positive eigenvalue: if $D$ has more than one eigenvalues, then the subspace of $\mathbb{R}^{k}$ spanned by the all positive eigenvectors has dimension at least $2$, which must have non-trivial intersection with $\ker(v^{T})$ as $\dim(\ker(v^{T}))=k-1$. Then there exists a nonzero vector $z\in\ker(v^{T})$ such that $z^{T}Dz>0$, which would imply $-z^{T}BB^{T}z=z^{T}Dz>0$, a contradiction.

For any eigenvalue-eigenvector pair $(\lambda_{i},z_{i})$ of $B^{T}B$ where $\lambda_{i}\not=0$ and $z_{i}\in\mathbb{S}^{d-1}$,

$$ $\displaystyle\lambda_{i}^{2}=$ $\displaystyle z_{i}^{T}(B^{T}B)^{2}z_{i}$ (123) //replace something with imbalance matrix (124) $\displaystyle\lambda_{i}^{2}=$ $\displaystyle(z_{i}^{T}w)^{2}-z_{i}^{T}B^{T}DBz_{i}$ (125) $\displaystyle\lambda_{i}^{2}-(z_{i}^{T}w)^{2}=$ $\displaystyle-z_{i}^{T}B^{T}DBz_{i}$ (126) $\displaystyle\lambda_{i}^{2}-(z_{i}^{T}w)^{2}\in$ $\displaystyle(z_{i}^{T}(B^{T}B)z_{i})\cdot[-\lambda_{\max},-\lambda_{\min}]$ (127) $\displaystyle\lambda_{i}^{2}-(z_{i}^{T}w)^{2}\in$ $\displaystyle\lambda_{i}\cdot[-\lambda_{\max},-\lambda_{\min}]$ (128) $$

again, we can rewrite this as two quadratic inequalities

$$ $\begin{cases}\lambda_{i}^{2}+\lambda_{\max}\lambda_{i}-(z_{i}^{T}w)^{2}\geq 0 \\ \lambda_{i}^{2}+\lambda_{\min}\lambda_{i}-(z_{i}^{T}w)^{2}\leq 0\end{cases}$ (129) $$

from them we know that there are two possible intervals:

$$ $\begin{cases}\lambda_{i}\in\left[-\infty,\frac{-\lambda_{\max}-\sqrt{\lambda_{ \max}^{2}+4(z_{i}^{T}w)^{2}}}{2}\right]\cup\left[\frac{-\lambda_{\max}+\sqrt{ \lambda_{\max}^{2}+4(z_{i}^{T}w)^{2}}}{2},+\infty\right]\\ \lambda_{i}\in\left[\frac{-\lambda_{\min}-\sqrt{\lambda_{\min}^{2}+4(z_{i}^{T} w)^{2}}}{2},\frac{-\lambda_{\min}+\sqrt{\lambda_{\min}^{2}+4(z_{i}^{T}w)^{2}}} {2}\right]\end{cases}$ (130) $$

Note that we must have $\lambda_{i}\geq 0$ since $B^{T}B$ is positive semidefinite. So we can rewrite the bounds:

$$ $\lambda_{i}\in\left[\frac{-\lambda_{\max}+\sqrt{\lambda_{\max}^{2}+4(z_{i}^{T} w)^{2}}}{2},\frac{-\lambda_{\min}+\sqrt{\lambda_{\min}^{2}+4(z_{i}^{T}w)^{2}}} {2}\right]$ (131) $$

since the function $f(x)=-x+\sqrt{x+c^{2}}$ is monotonically decreasing, we have $f(\lambda_{\max})\leq f(\lambda_{\min})$, i.e. the lower bound is no greater than the upper bound, i.e. the above interval is always non-empty.
∎

###### Lemma A.13 (Bounding the norm of linear head ‖ v ‖ 2 2 superscript subscript norm 𝑣 2 2 \|v\|_{2}^{2} ∥ italic_v ∥ start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT )

###### Proof of Lemma A.13 .

###### Lemma A.14 (Bounding eigenvalues of B T ⁢ B superscript 𝐵 𝑇 𝐵 B^{T}B italic_B start_POSTSUPERSCRIPT italic_T end_POSTSUPERSCRIPT italic_B (re-stated from Min et al. [ 2023b ] ))

###### Proof of Lemma A.14 .

### A.5 Layerwise sensitivity

###### Lemma A.15 (Sensitivity upper bounds for two-layer linear networks)

Consider two neighboring and normalized datasets $\mathcal{D}=(X,Y)$ and $\mathcal{D}^{\prime}=(X^{\prime},Y^{\prime})$ where $X,X^{\prime}\in\{A^{T}A=I_{d\times d}:A\in\mathbb{R}^{n\times d}\}$ and $\|Y\|_{2}=\|Y^{\prime}\|_{2}=1$. Then the $L_{2}$-distance of each layer is upper bounded by

$$ $\begin{cases}\|\nabla_{B}\mathcal{L}(v,B;\mathcal{D})-\nabla_{B}\mathcal{L}(v, B;\mathcal{D}^{\prime})\|_{2}\leq\|v\|_{2}\|(X^{\prime})^{T}Y^{\prime}-X^{T}Y \|_{2}\\ \|\nabla_{v}\mathcal{L}(v,B;\mathcal{D})-\nabla_{v}\mathcal{L}(v,B;\mathcal{D} ^{\prime})\|_{2}\leq\|B\|_{2}\|(X^{\prime})^{T}Y^{\prime}-X^{T}Y\|_{2}\end{cases}$ (132) $$

###### Proof of Lemma A.15 .

The $L_{2}$-distance of $B$-layer gradients are

$$ $\displaystyle\|\nabla_{B}\mathcal{L}(v,B;X)-\nabla_{B}\mathcal{L}(v,B;X^{ \prime})\|_{2}$ $\displaystyle=$ $\displaystyle\|v(XB^{T}v-Y)^{T}X-v(X^{\prime}B^{T}v-Y^{\prime})^{T}X^{\prime} \|_{2}$ $\displaystyle=$ $\displaystyle\|v(X^{T}XB^{T}v-X^{T}Y)^{T}-v((X^{\prime})^{T}X^{\prime}B^{T}v-( X^{\prime})^{T}Y^{\prime})^{T}\|_{2}$ $\displaystyle\leq$ $\displaystyle\|v\|_{2}\|X^{T}XB^{T}v-X^{T}Y-(X^{\prime})^{T}X^{\prime}B^{T}v+( X^{\prime})^{T}Y^{\prime}\|_{2}$ $\displaystyle=$ $\displaystyle\|v\|_{2}\|(X^{\prime})^{T}Y^{\prime}-X^{T}Y\|_{2}$ $$

Similarly, the $L_{2}$-distance of $v$-layer gradients are

$$ $\displaystyle\|\nabla_{v}\mathcal{L}(v,B;X)-\nabla_{v}\mathcal{L}(v,B;X^{ \prime})\|_{F}\leq$ $\displaystyle\|B\|_{F}\|X^{T}XB^{T}v-X^{T}Y-(X^{\prime})^{T}X^{\prime}B^{T}v+( X^{\prime})^{T}Y^{\prime}\|_{2}$ (133) $\displaystyle=$ $\displaystyle\|B\|_{F}\|(X^{\prime})^{T}Y^{\prime}-X^{T}Y\|_{2}$ (134) $$

∎

###### Proof of Proposition 3.7 .

At the initialization of full fine-tuning, we have $\|B_{0}\|_{F}=k$ and $\|v_{0}\|_{2}=k/\sqrt{d}$. When we pick $\mathcal{D},\mathcal{D}^{\prime}$ such that $X=X^{\prime}$, $Y\parallel v_{0}$ and $Y^{\prime}=-Y$, we have

$$ $\displaystyle\|\nabla_{B}\mathcal{L}(v_{0},B_{0};\mathcal{D})-\nabla_{B} \mathcal{L}(v_{0},B_{0};\mathcal{D}^{\prime})\|_{2}=\sqrt{2}\|v_{0}\|_{2}$ (135) $$

So we can bound the local sensitivity of $\nabla_{B}\mathcal{L}$ at $(v_{0},B_{0})$ by:

$$ $\Delta(\nabla_{B}\mathcal{L}(v_{0},B_{0})):=\sup_{\mathcal{D},\mathcal{D}^{ \prime}}\|\nabla_{B}\mathcal{L}(v_{0},B_{0};\mathcal{D})-\nabla_{B}\mathcal{L} (v_{0},B_{0};\mathcal{D}^{\prime})\|_{2}\in\left[\sqrt{2}\|v_{0}\|_{2},2\|v_{0 }\|_{2}\right]$ (136) $$

By taking expectation on $v_{0}$, we get $\mathbb{E}\left[\Delta(\nabla_{B}\mathcal{L}(v_{0},B_{0}))\right]\in\left[k
\sqrt{2/d},2k/\sqrt{d}\right]$.
When we pick $\mathcal{D},\mathcal{D}^{\prime}$ such that $X=X^{\prime}$, $(1,Y)$ being an eigenvalue-eigenvector pair of $B_{0}$, and $Y=-Y^{\prime}$, we have

$$ $\displaystyle\|\nabla_{v}\mathcal{L}(v_{0},B_{0};\mathcal{D})-\nabla_{v} \mathcal{L}(v_{0},B_{0};\mathcal{D}^{\prime})\|_{2}=\sqrt{2}k$ (137) $$

So we can bound the local sensitivity of $\nabla_{v}\mathcal{L}$ at $(v_{0},B_{0})$ by:

$$ $\Delta(\nabla_{v}\mathcal{L}(v_{0},B_{0})):=\sup_{\mathcal{D},\mathcal{D}^{ \prime}}\|\nabla_{v}\mathcal{L}(v_{0},B_{0};\mathcal{D})-\nabla_{v}\mathcal{L} (v_{0},B_{0};\mathcal{D}^{\prime})\|_{2}\in\left[\sqrt{2}k,2k\right]$ (138) $$

Combining the bounds of local sensitivities of $\nabla_{v}\mathcal{L}$ and $\nabla_{B}\mathcal{L}$, we have

$$ $\Delta(\nabla_{v}\mathcal{L}(v_{0},B_{0}))=\Theta(\sqrt{d}\Delta(\nabla_{B} \mathcal{L}(v_{0},B_{0})))$ (139) $$

∎

###### Lemma A.15 (Sensitivity upper bounds for two-layer linear networks)

###### Proof of Lemma A.15 .

###### Proof of Proposition 3.7 .

### A.6 Fine-tuning loss upper bound

###### Lemma A.16 (Imbalance matrix in fine-tuning under layerwise noise)

During fine-tuning (Equation (10)), the imbalance matrix $D$ in Definition 3.5 evolves as

$$ $\mathbb{E}\left[\frac{dD}{dt}\right]=0$ (140) $$

###### Proof of Lemma 3.9 .

We prove this lemma by analyzing the infinitesimal generator $A$ of imbalance matrix $D$:

$$ $\displaystyle A(D_{0}(v,B))_{ij}:=$ $\displaystyle\lim\limits_{t\downarrow 0}\frac{\mathbb{E}^{D_{0}}[D_{ij}]-{(D_{ 0})}_{ij}}{t}$ (141) $\displaystyle=$ $\displaystyle 0+\sigma^{2}\sum_{i^{\prime}\in[k]}\sum_{j^{\prime}\in[k]} \mathbf{1}[i^{\prime}=j^{\prime}=i=j]$ (142) $\displaystyle-\sigma^{2}\sum_{i^{\prime}\in[k],j^{\prime}\in[d]}\sum_{i^{ \prime\prime}\in[k],j^{\prime\prime}\in[d]}\mathbf{1}[i^{\prime}=i^{\prime \prime}=i=j\text{ and }j^{\prime}=j^{\prime\prime}]$ (143) $$

the generator is zero for $i\not=j$. So we can just consider the case where $i=j$.

$$ $\displaystyle A(D_{0}(v,B))_{ii}=$ $\displaystyle\sigma^{2}\sum_{i^{\prime}\in[k]}\sum_{j^{\prime}\in[k]}\mathbf{1 }[i^{\prime}=j^{\prime}=i]$ (144) $\displaystyle-\sigma^{2}\sum_{i^{\prime}\in[k],j^{\prime}\in[d]}\sum_{i^{ \prime\prime}\in[k],j^{\prime\prime}\in[d]}\mathbf{1}[i^{\prime}=i^{\prime \prime}=i\text{ and }j^{\prime}=j^{\prime\prime}]$ (145) $\displaystyle=$ $\displaystyle(d-d)\sigma^{2}$ (146) $\displaystyle=$ $\displaystyle 0$ (147) $$

∎

###### Theorem A.17 (Loss upper bound of fine-tuning)

In fine-tuning under layerwise noise (Equation (10)), we have

$$ $\mathbb{E}[\mathcal{L}]\lessapprox\mathbb{E}[\mathcal{L}]e^{(-\bar{\lambda}+ \sqrt{2}\sigma^{2}(1+d))t}+L^{\square}(1-e^{(-\bar{\lambda}+\sqrt{2}\sigma^{2} (1+d))t})$ (148) $$

where $L^{\square}=\sigma^{2}\frac{(1+d)\|X^{T}Y\|-d\underline{\lambda}}{\bar{\lambda
}-\sqrt{2}\sigma^{2}(1+d)}$.

###### Proof of Theorem 3.11 .

We first simplify the loss dynamics:

$$ $\displaystyle\partial\mathcal{L}=$ $\displaystyle\partial\frac{1}{2}\|XB^{T}v-Y\|^{2}$ (149) $\displaystyle=$ $\displaystyle\frac{1}{2}\left\langle\nabla_{v}\|XB^{T}v-Y\|^{2},\partial v \right\rangle+\frac{1}{2}\left\langle\nabla_{B}\|XB^{T}v-Y\|^{2},\mathrm{vec}( \partial B)\right\rangle$ (150) $\displaystyle+\frac{1}{4}(\partial v)^{T}H_{\|XB^{T}v-Y\|^{2}}(\partial v)+ \frac{1}{4}[\mathrm{vec}(\partial B)]^{T}H_{\|XB^{T}v-Y\|^{2}}\mathrm{vec}( \partial B)$ (151) $\displaystyle=$ $\displaystyle(XB^{T}v-Y)^{T}XB^{T}\partial v+(XB^{T}v-Y)^{T}X(\partial B)^{T}v$ (152) $\displaystyle+\frac{1}{2}(\partial v)^{T}BB^{T}(\partial v)+\frac{1}{2}[ \mathrm{vec}(\partial B)]^{T}H_{\|XB^{T}v-Y\|^{2}}\mathrm{vec}(\partial B)$ (153) $\displaystyle=$ $\displaystyle-(XB^{T}v-Y)^{T}XB^{T}BX^{T}(XB^{T}v-Y)\partial t+(XB^{T}v-Y)^{T} XB^{T}\sqrt{2\sigma^{2}d}\partial W_{t}$ (154) $\displaystyle-(XB^{T}v-Y)^{T}XX^{T}(XB^{T}v-Y)v^{T}v\partial t+(XB^{T}v-Y)^{T} X(\sqrt{2\sigma^{2}}\partial W_{t}^{\prime})v$ (155) $\displaystyle+\sigma^{2}\mathrm{trace}(BB^{T})\partial t+\sigma^{2}d\|v\|^{2}\partial t$ (156) $\displaystyle=$ $\displaystyle-(B^{T}v-X^{T}Y)^{T}B^{T}B(B^{T}v-X^{T}Y)\partial t+(B^{T}v-X^{T} Y)^{T}B^{T}\sqrt{2\sigma^{2}}\partial W_{t}$ (157) $\displaystyle-(B^{T}v-X^{T}Y)^{T}(B^{T}v-X^{T}Y)v^{T}v\partial t+(B^{T}v-X^{T} Y)^{T}(\sqrt{2\sigma^{2}}\partial W_{t}^{\prime})v$ (158) $\displaystyle+\sigma^{2}\mathrm{trace}(B^{T}B)\partial t+\sigma^{2}d\|v\|^{2}\partial t$ (159) $$

By Lemma A.13 and Lemma A.14, we have

$$ $\displaystyle\partial\mathbb{E}\mathcal{L}=$ $\displaystyle-\mathbb{E}[(w-X^{T}Y)^{T}(B^{T}B+v^{T}vI_{d\times d})(w-X^{T}Y)] \partial t+\sigma^{2}\mathbb{E}[\|B\|_{F}^{2}+d\|v\|_{2}^{2}]\partial t$ (160) $\displaystyle\leq$ $\displaystyle\mathbb{E}\left\{-\|w-X^{T}Y\|_{2}^{2}\frac{\underline{\lambda}+ \sqrt{\underline{\lambda}^{2}+4\|w\|^{2}}}{2}\partial t-\|w-X^{T}Y\|_{2}^{2} \frac{-\bar{\lambda}+\sqrt{\bar{\lambda}^{2}+4(z_{\min}^{T}w)^{2}}}{2}\partial t\right\}$ (161) $\displaystyle+\mathbb{E}\left\{\sigma^{2}d\frac{-\underline{\lambda}+\sqrt{ \underline{\lambda}^{2}+4(z_{\min}^{T}w)^{2}}}{2}\partial t+\sigma^{2}d\frac{ \bar{\lambda}+\sqrt{\bar{\lambda}^{2}+4\|w\|^{2}}}{2}\partial t\right\}$ (162) $\displaystyle\leq$ $\displaystyle-\frac{1}{2}\mathbb{E}[\|w-X^{T}Y\|_{2}^{2}(\Lambda_{\min}+ \Lambda_{\max})]\partial t+\frac{1}{2}\sigma^{2}\mathbb{E}[d\Gamma_{\min}+ \Gamma_{\max}]\partial t$ (163) $$

where we define

$$ $\begin{cases}\Lambda_{\min}=\underline{\lambda}+\sqrt{\underline{\lambda}^{2}+ 4\|w\|^{2}}\geq\max\left(0,2\underline{\lambda}\right)\\ \Lambda_{\max}=-\bar{\lambda}+\sqrt{\bar{\lambda}^{2}+4(z_{\min}^{T}w)^{2}} \geq\max\left(0,-2\bar{\lambda}\right)\\ \Gamma_{\min}=-\underline{\lambda}+\sqrt{\underline{\lambda}^{2}+4(z_{\min}^{T }w)^{2}}\leq\max\left(2\|w\|,2\|w\|-2\underline{\lambda}\right)=2\|w\|+2\max(0 ,-\underline{\lambda})\\ \Gamma_{\max}=\bar{\lambda}+\sqrt{\bar{\lambda}^{2}+4\|w\|^{2}}\leq\max(2\|w\| ,2\|w\|+2\bar{\lambda})=2\|w\|+2\max(0,\bar{\lambda})\end{cases}$ (164) $$

Then we rewrite the upper bound:

$$ $\displaystyle\partial\mathbb{E}[\mathcal{L}]\leq$ $\displaystyle-\frac{1}{2}\mathbb{E}[\|w-X^{T}Y\|_{2}^{2}(\Lambda_{\min}+ \Lambda_{\max})]\partial t+\frac{1}{2}\sigma^{2}\mathbb{E}[d\Gamma_{\min}+ \Gamma_{\max}]\partial t$ (165) $\displaystyle\partial\mathbb{E}[\mathcal{L}]\lessapprox$ $\displaystyle-\bar{\lambda}\mathbb{E}[\mathcal{L}]\partial t+\sigma^{2}(\sqrt{ 2}(1+d)\mathbb{E}[\mathcal{L}]^{1/2}+(1+d)\|X^{T}Y\|-d\underline{\lambda})\partial t$ (166) $\displaystyle\partial\mathbb{E}[\mathcal{L}]\lessapprox$ $\displaystyle(-\bar{\lambda}+\sqrt{2}\sigma^{2}(1+d))\mathbb{E}[\mathcal{L}] \partial t+\sigma^{2}((1+d)\|X^{T}Y\|-d\underline{\lambda})\partial t$ (167) $\displaystyle\mathbb{E}[\mathcal{L}]\lessapprox$ $\displaystyle\mathbb{E}[\mathcal{L}]e^{(-\bar{\lambda}+\sqrt{2}\sigma^{2}(1+d) )t}+L^{\square}(1-e^{(-\bar{\lambda}+\sqrt{2}\sigma^{2}(1+d))t})$ (168) $$

where $L^{\square}=\sigma^{2}\frac{(1+d)\|X^{T}Y\|-d\underline{\lambda}}{\bar{\lambda
}-\sqrt{2}\sigma^{2}(1+d)}$.
∎

###### Lemma A.16 (Imbalance matrix in fine-tuning under layerwise noise)

###### Proof of Lemma 3.9 .

###### Theorem A.17 (Loss upper bound of fine-tuning)

###### Proof of Theorem 3.11 .

## Appendix B More information on experiments

This Appendix provides additional information regarding the datasets, pretraining methods, and adaptation methods.

### B.1 Dataset and method details

- •
ImageNet-1K$\rightarrow$CIFAR-10/STL-10. We pretrain the feature extractor without privacy mechanisms on ImageNet-1K and privately fine-tune the feature extractor and a linear head on CIFAR-10.
–
ResNet-50: We consider unsupervised pretraining methods and use a larger model architecture for this benchmark. We initialize pretrained features of ResNet-50 from MoCo-v2 Chen et al. [] & MoCo-v3 Chen* et al. [] on ImageNet-1K Russakovsky et al. [] without privacy mechanisms and then privately fine-tune the ResNet-50 model on CIFAR-10 for 10 epochs.
–
Mini-DeiT-Ti: We run Vision Transformer (ViT) models like mini Data-efficient image Transformers (Mini-DeiT) on ImageNet-1K$\rightarrow$STL-10 benchmark Touvron et al. [], Zhang et al. []. We first pretrain a DeiT model on ImageNet then pretrain a Mini-DeiT-Ti model with weight distillation from the DeiT model Zhang et al. []. The weight of the Mini-DeiT-Ti model is publicly available in the official Github repository of Zhang et al. []. After that, we privately fine-tune the Mini-DeiT-Ti model on STL-10 for 20 epochs.
- •
CIFAR-10$\rightarrow$STL-10. This benchmark has been extensively studied in the context of domain adaptation French et al. [], Kumar et al. []. We pretrain the feature extractor on CIFAR-10 Krizhevsky [] using stochastic gradient descent without privacy mechanisms. Then we finetune the pretrained features and a randomly initialized linear head on STL-10 for 10 epochs Coates et al. []. The training subset of STL-10 is small and only contains 500 images. To align with the small scale fine-tuning data, we run the experiments with smaller and data-efficient models: MobileNet-v3 and ResNet-18.
- •
RandP$\rightarrow$CIFAR-10. To verify the general existence of concave utility curves, we also consider a slightly non-standard pretraining protocol: pretrain a wide residual network (WRN) Zagoruyko and Komodakis [] on synthetic images generated by random diffusion processes. We basically follow the experiment settings in Tang et al. [] and reproduce their results. We first sample synthetic images from StyleGAN-based datasets and employ representation learning to train a feature extractor using the loss function proposed by Wang and Isola []. Then we privately fine-tune the model on CIFAR-10.

### B.2 Hyperparameter Sweeps

For vision transformer models, we sweep over 5 learning rates {0.1, 0.05, 0.025, 0.0125, 0.00625} for both linear probing (LP), fine-tuning (FT), and first-linear-probing-then-finetuning (LP-FT). Other models are empirically not sensitive to learning rates and we only sweep over 3 learning rates {0.1, 0.05, 0.025}.

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
questions: Most of my concerns are mentioned in the Weaknesses part.

**Q1.** Is there a reason why Langevin diffusion is defined differently from Ganesh et al. (2023b)?

**Q2.** How realistic are Assumptions 3.1 and 3.2? Were they experimentally validated?

**Q3.** Public pre-training and private fine-tuning approach has been seriously questioned recently. I would appreciate the author's thoughts on the recent position paper by Tramèr et al. (2024).

___

Tramèr, F., Kamath, G., & Carlini, N. (2024). Position: Considerations for Differentially Private Learning with Large-Scale Public Pretraining. Proceedings of the 41st International Conference on Machine Learning.
soundness: 3
strengths: **S1.** The importance and relevance of the problem tackled — differentially private deep learning — is highlighted well, given the challenge posed by the high dimensionality of typical models.

**S2.** The concept of splitting the privacy budget between full parameter tuning (or Full Fine-Tuning, FFT) and Linear Probing (LP: tuning only the last layer) appears intriguing and potentially valuable for practical applications.

**S3.** The attempt to provide rigorous theoretical analysis addresses a gap in the current literature.

**S4.** The paper includes multiple experiments across diverse datasets and models, supporting the benefits of the proposed hybrid approach.
summary: This paper addresses the challenge of differentially private fine-tuning of deep learning models. The authors highlight that naïve full-parameter fine-tuning leads to misalignment between the pre-trained features and the last layer. They propose a hybrid strategy that combines linear probing with fine-tuning, demonstrating theoretically and empirically that this approach mitigates feature distortion. The theoretical framework is based on a simplified Langevin diffusion and a two-layer ReLU neural network. The paper’s theoretical insights are further supported by experimental evaluations on various vision tasks and models.
weaknesses: **W1.** The idea of combining FFT and LP is not entirely novel, as similar approaches were introduced and experimentally tested by Tang et al. (2023). A more detailed discussion of how this paper builds on or extends prior work is needed.

**W2.** The theoretical model appears oversimplified. Using a 2-layer neural network with ReLU activations while enabling a decoupling of feature learning and classification seems far removed from the practical settings under consideration. It is unclear how this model aligns with realistic applications, particularly since results in Section 4.1.1 suggest exponential convergence, implying that the underlying problem is no harder than strongly convex optimization (or that PL conditions hold for the loss).

In addition, the analysis is done for a non-standard simplified Langevin diffusion instead of DP-SGD. The benefits of such an approach seem questionable. The authors *"Apply a zeroth-order asymptotic expansion"* which sounds very complicated but looks like a simple removal of Brownian motion representing Gaussian noise. This noise is the core component of DP training, and its removal makes method (1) equivalent to simple gradient flow. In addition, it ignores the crucial per-sample clipping operation, which makes DP deep learning feasible. Thus, the proposed approximation loses the main features of differentially private training. The authors say on line 164
> our modeling preserves the noisy behavior characteristic of DP-SGD

which need to be justified.

**W3.** Assumptions 3.1 and 3.2 are unconventional and strongly restrictive. Only one previous work referenced these assumptions, making it difficult to accept them as standard or practical in the context of simple binary classification.

**W4.** The experimental setup lacks details on the clipping threshold, a key parameter influencing DP-SGD performance and necessary for reproducibility.

**W5.** The paper’s theoretical results and mathematical proofs in the Appendix are difficult to follow, partly due to unclear notation and insufficient explanation. Specific issues and questions include:

- The notation $\sim$ in line 214 is ambiguous.

- The choice of zero-mean Gaussian initialization for the linear head is restrictive.

- The meaning of *“optimality"* of $w_j$ in line 232 is unclear.

- After a quick look at the work of Ganesh et al. (2023b), I did not find the exact statement of Theorem 4.1. Could the authors please point me to the exact place in the original paper?

- Theorems 4.2 and 4.3 are hard to comprehend and presented without almost any commentary or explanation.

- It is unclear to me why the authors say that
> According to Theorem 5.1, a greater proportion of the privacy budget should be allocated to DP-LP
when the total privacy budget is smaller.

- Moreover, I would like to ask what is parameter $\rho$. It seems like it can make $r$ arbitrarily large.

- Section 5.2 relies on assumptions from the Appendix, which are also very strong (like E.7). Such an approach makes it hard to access the theoretical contributions of the paper adequately.

- What does variable $D$ denote on line 442? How big can $t_{lp}$ be in practice for realistic parameter values? It is unclear how Corollary 5.3 is obtained for the main paper text.

- The steps of Theorem 3.3 proof are unclear. For instance, how is formula (28) obtained? Why the probability on line 958 equals $1-2^{-h}$?

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: In line 199 , why are subspaces separated by $\mathbb{I}\_{\boldsymbol{x}\_i^{\top} z>0}=\mathbb{I}\_{y_i=-1}$ or $\mathbb{I}\_{\boldsymbol{x}\_i^{\top} z>0}=\mathbb{I}\_{y_i=1}$ ?
soundness: 3
strengths: - The paper investigates an interesting phenomenon in fine-tuning methods, where non-private fine-tuning distorts pretrained features and leads to degraded OOD performance. A similar empirical effect is observed in private settings, as shown in Figure 3.

- The proposed DP-LP running before DP-FFT can theoretically reduce feature distortion. It would be beneficial to provide experimental results to compare with Figure 3.

- The proposed LP mechanism in private fine-tuning is effective, as demonstrated in Tables 1 and 2.
summary: The paper examines differentially private (DP) fine-tuning (FT) methods, showing through theoretical and empirical analysis that DP full FT can distort pretrained backbone features due to misalignment between the pretrained backbone and the randomly initialized linear head. To address this, the authors propose DP-LP_FFT, which first performs linear probing and then fine-tunes. Additionally, the paper provides convergence rates for the proposed methods using two-layer neural networks.
weaknesses: There is a lack of comparison of feature distortion with DP-LP.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. theorem 3.4 says after $\Delta t$, DP-FFT does not distort the pre-trained features. But the Eq (10) is stated for $\forall t\in(0, \Delta t)$. Is there a typo in the range of $t$? I guess it should be $∀t\in(\Delta t, \infty)$?
2. typos around line 1148
soundness: 3
strengths: 1. Use Langevin diffusion as an approximation of DP-SGD to study the dynamics of DP tuning. While obtained from an approximation model, findings are very interesting and align with empirical observatons. 
2. The proposed hybrid tuning method is supported by theoretical guaranteens and evidence from extensive experiments. 
3. Insights into privacy budget allocation are of practical interest.
summary: This work studies the training dynamics of DP fine-tuning using a Langevin diffusion approximation of DP-SGD, and provides theoretical understandings of how fune-tuning techniques (such as DP-LP, DP-FFT, or a hybrid) affect the out-of-distribution performance of a 2-layer ReLU network.
Using a zeroth-order approximation, the authors provide a theoretical explanation of an empirical phenomenon that randomly initialized linear heads distort pre-trained backbone features in the early stages of DP-FFT.
To mitigate or even avoid the feature distortion, they propose a hybrid method that combines DP-LP and DP-FFT, and further examine the privacy budget allocation across DP-LP and DP-FFT. 
Extensive numerical experiments support theoretical findings and demonstrate the effectiveness of the proposed hybrid method.
Overall, this work deepens our understanding of DP fine-tuning by providing a solid theoretical explanation. 
I found this work interesting and thus recommend an Accept.
weaknesses: In practice, trainings are usually done in a discrete manner (say t=1,2,...), but Lagenvin diffusion is a continuous approximation for t>0. This gap may hamper the generality of this work's findings. For example, if $\Delta t$ in Theorem 3.3 belongs to [0,1], then feature distortion might not be an issue, because we start from $t=1$. Therefore, it would be helpful if authors could provide a short discussion of the value of $\Delta t$, or name some driving factors that may significantly affect the value of $\Delta t$.

Meta-review: This paper is in borderline! Although the authors have addressed some comments and questions, the reviewer is still not convinced about the theoretical contributions of the paper as they cannot see the additional value to experimental results. Moreover, the reviewer also thinks that the result of the paper may not be correct due to the sum over N, the upper bound is expected to scale linearly with N (assuming dataset size due to undefined notation), which makes the result poor from a Differential Privacy perspective. Although the authors have updated the theorem statement and the proof, the reviewer is still skeptical and did not want to support the publication of the paper. Due to some technical issue, I cannot recommend the acceptance of this paper. However, I strongly encourage the authors to revise the paper carefully and resubmit it to the next publication venues.

Decision: Reject