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

Authors: Shuqi Ke, Charlie Hou, Sewoong Oh, Giulia Fanti
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: differential privacy, convergence, fine-tuning theory, transfer learning theory, langevin diffusion, gradient flow

Abstract:
We show that differentially private full fine-tuning (DP-FFT) can distort pre-trained backbone features based on both theoretical and empirical results. We identify the cause of the distortion as the misalignment between the pre-trained backbone and the randomly initialized linear head. We prove that a sequential fine-tuning strategy can mitigate the feature distortion: first-linear-probing-then-fine-tuning (DP-LP-FFT). A new approximation scheme allows us to derive approximate upper and lower bounds on the training loss of DP-LP and DP-FFT, in a simple but canonical setting of 2-layer neural networks with ReLU activation. Experiments on real-world datasets and architectures are consistent with our theoretical insights.   We also derive new upper bounds for 2-layer linear networks without the approximation. Moreover, our theory suggests a trade-off of privacy budget allocation in multi-phase fine-tuning methods like DP-LP-FFT.

Decision: reject
Reviewer scores: 5, 6, 8

Reviews and discussion:
<thread>
creation date:
1730358254394

thread_type:
official_review

<thread_object>
creator: reviewer_kixw

<content>
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.

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.

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$.

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
</content>
</thread_object>

</thread>
<thread>
creation date:
1730599753489

thread_type:
official_review

<thread_object>
creator: reviewer_hkiw

<content>
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.

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.

weaknesses:
There is a lack of comparison of feature distortion with DP-LP.

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}$ ?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730649199483

thread_type:
official_review

<thread_object>
creator: reviewer_avt2

<content>
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.

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.

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}$?

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731644800182

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response

comment:
We would like to extend our sincere gratitude to all the reviewers for their time and effort in evaluating our work. We are especially grateful to reviewer avt2 for the detailed and insightful feedback, which has been invaluable in refining our contributions. We also appreciate the positive feedback from reviewers HKiw and kixW, whose encouraging comments have reinforced our confidence in the value of this research.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733273064750

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of the discussion

comment:
Dear AC, 

Thank you for overseeing our paper; we appreciate your time and efforts. We wanted to provide a brief summary of the discussion from our perspective. Overall, our work introduces a novel analysis and understanding of differentially-private finetuning, demonstrating a phenomenon we call representation distortion for DP full fine-tuning. We confirm our theory with extensive empirical evidence on several architectures and settings. 

**Strengths**:

The reviewers highlighted a number of strengths of the work: 

- **From reviewer avt2 and HKiw**: 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. The paper includes multiple experiments across diverse datasets and models, supporting the benefits of the proposed hybrid approach.
- **From reviewer avt2 and kixW**: The attempt to provide rigorous theoretical analysis addresses a gap in the current literature. While obtained from an approximation model, findings are very interesting and align with empirical observations.
- **From our rebuttal**: We theoretically analyze feature distortion in DP-LP-FFT universally across private and non-private settings.

**Key Weakness and response**:

- **From reviewer avt2**: The theoretical model (2-layer ReLU network) appears oversimplified.
    - Response: Two-layer neural networks are widely accepted in the theoretical community as valuable tools for deriving intuition. The theoretical insights derived from these models align with empirical results on complex architectures.
- **From reviewer avt2**: Assumptions 3.1 and 3.2 are unconventional and restrictive.
    - Response: These assumptions are rooted in a well-established line of research within neural alignment and are consistent with empirical observations.
- **From reviewer avt2**: The analysis ignores the crucial per-sample clipping operation.
    - Response: We have updated our paper in openreview and we provide analysis of Langevin diffusion with clipping in Appendix F:
        - We note that this is the first analysis of Langevin diffusion with clipping, and DP-GD with clipping based on stochastic differential equations in the literature, to our knowledge. Our results include:
            - The existence of a unique strong solution of the Langevin diffusion with clipping.
            - The zeroth order approximation error of Langevin diffusion with clipping.
            - A clipping version of Theorem 3.3 (feature distortion).
- **From reviewer avt2**: The paper lacks relevance to real-world applications and practical recommendations based on the empirical results.
    - Response: We provide a systematic study of DP-LP-FFT methods and feature distortion in fine-tuning. Our paper provides diagnostic tools for feature distortion and our experiments provide architecture-specific recommendations.
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.