An Optimization Framework for Differentially Private Sparse Fine-Tuning

Abstract: Differentially private stochastic gradient descent (DP-SGD) is broadly considered to be the gold standard for training and fine-tuning neural networks under differential privacy (DP). With the increasing availability of high-quality pre-trained model checkpoints (e.g., vision and language models), fine-tuning has become a popular strategy. However, despite recent progress in understanding and applying DP-SGD for private transfer learning tasks, significant challenges remain -- most notably, the performance gap between models fine-tuned with DP-SGD and their non-private counterparts. Sparse fine-tuning on private data has emerged as an alternative to full-model fine-tuning; recent work has shown that privately fine-tuning only a small subset of model weights and keeping the rest of the weights fixed can lead to better performance. In this work, we propose a new approach for sparse fine-tuning of neural networks under DP. Existing work on private sparse finetuning often used fixed choice of trainable weights (e.g., updating only the last layer), or relied on public model's weights to choose the subset of weights to modify. Such choice of weights remains suboptimal. In contrast, we explore an optimization-based approach, where our selection method makes use of the private gradient information, while using off the shelf privacy accounting techniques. Our numerical experiments on several computer vision models and datasets show that our selection method leads to better prediction accuracy, compared to full-model private fine-tuning or existing private sparse fine-tuning approaches.

Introduction

Training deep learning models with Differentially Private Stochastic Gradient Descent (DP-SGD) is essential for protecting user data but often incurs a significant "utility tax." DP-SGD operates by clipping gradients and adding random noise to them. As the number of parameters ($d$) increases, the total amount of noise injected grows proportionally to $\sqrt{d}$, frequently resulting in a substantial accuracy drop compared to non-private training.

To address this "curse of dimensionality," researchers have turned to sparse fine-tuning, where only a small fraction of the model's weights are updated. However, selecting which weights to tune is non-trivial. Previous approaches relied on static rules (e.g., tuning only the final layer), public data, or pre-trained weight magnitudes. These methods are often suboptimal as they fail to account for the specific characteristics of the private downstream dataset.

Methodology: The SPARTA Framework

We introduce SPARTA (Sparse & PrivAte Row–gradient Thresholding Algorithm), an optimization-based approach to subnetwork selection. The core innovation lies in using private gradient information to intelligently select a "high-impact" subset of weights to update.

Expand

Source, p. 1

The framework consists of several key components:

1. Optimization Formulation: Mask selection is treated as a mathematical optimization problem. Using a Taylor expansion, SPARTA approximates which weights would most effectively reduce the loss if updated.
2. Row-Grouping: To mitigate high noise levels in DP, SPARTA aggregates weights into groups (specifically rows of the weight matrix). Averaging the "importance" within a group increases the signal-to-noise ratio, enabling the identification of an effective mask even under tight privacy budgets.
3. Privacy-Preserving Selection: The selection process itself is differentially private. The computational overhead for mask selection is minimal, roughly equivalent to one additional epoch of standard DP-SGD.
4. Two-Phase Training: Training begins with a brief "warm-up" phase. Subsequently, SPARTA selects the optimal sparse mask, and the model is fine-tuned exclusively on these selected weights for the remaining epochs.

Experimental Setup

The effectiveness of SPARTA was evaluated on standard computer vision benchmarks:

• Datasets: CIFAR-10 and CIFAR-100 served as the private downstream tasks.
• Pre-training: Models were pre-trained on larger public datasets like ImageNet and ImageNet-32.
• Architectures: The framework was tested across diverse models, including ResNet18, Wide-ResNet (WRN-28-10), and various Vision Transformers (DeiT-Tiny, Small, and Base).

Discussion and Key Insights

A primary finding of this work is that individual noisy gradients are unreliable for mask selection in high-privacy regimes. However, by aggregating gradients into row groups, the noise cancels out, revealing the parameters most critical for the task. This confirms the intuition that for DP-SGD, tuning a small number of well-chosen parameters is generally superior to full-model fine-tuning.

Furthermore, the use of row-grouping results in structured sparsity. This structure facilitates hardware-level optimizations, leading to a 15% wall-clock time speedup compared to full-model fine-tuning.

Results

SPARTA consistently outperformed state-of-the-art private fine-tuning methods, such as DP-BitFit and magnitude-based pruning, across various privacy budgets ($\epsilon = 2, 4, 8$):

• Accuracy Boost: On CIFAR-100 using a DeiT-Tiny model at $\epsilon=2$, SPARTA achieved 69.58% accuracy, significantly surpassing full-model fine-tuning (61.22%) and last-layer-only approaches (58.10%).
• Robustness: The method's effectiveness across both convolutional and transformer-based architectures demonstrates its broad applicability.
• Efficiency: By tuning only approximately 20% of the parameters, SPARTA achieved better results than tuning 100%, while also reducing training time.

Conclusion

While SPARTA offers significant improvements, it introduces hyperparameters such as the mask-finding epoch and the sparsity level, which currently require computational effort to tune. Future research could extend this optimization framework to other Parameter-Efficient Fine-Tuning (PEFT) methods, such as LoRA, by privately selecting optimal ranks or subspaces for adapters.