[Paper Review] Large Scale Transfer Learning for Differentially Private Image Classification

Harsh Mehta, Abhradeep Thakurta+2 more|arXiv (Cornell University)|May 6, 2022

Privacy-Preserving Technologies in Data5 citations

TL;DR

This paper demonstrates that large-scale transfer learning significantly improves differentially private image classification on ImageNet. By pretraining vision transformers on massive public datasets (JFT-300M/4B), fine-tuning just the final layer for a single full-batch step with LAMB optimizer and near-zero initialization achieves state-of-the-art accuracy of 81.7% across a wide privacy budget range (ε ∈ [4,10], δ = 10⁻⁶), drastically reducing computational cost while maintaining high utility under strong privacy guarantees.

ABSTRACT

Differential Privacy (DP) provides a formal framework for training machine learning models with individual example level privacy. In the field of deep learning, Differentially Private Stochastic Gradient Descent (DP-SGD) has emerged as a popular private training algorithm. Unfortunately, the computational cost of training large-scale models with DP-SGD is substantially higher than non-private training. This is further exacerbated by the fact that increasing the number of parameters leads to larger degradation in utility with DP. In this work, we zoom in on the ImageNet dataset and demonstrate that, similar to the non-private case, pre-training over-parameterized models on a large public dataset can lead to substantial gains when the model is finetuned privately. Moreover, by systematically comparing private and non-private models across a range of large batch sizes, we find that similar to non-private setting, choice of optimizer can further improve performance substantially with DP. By using LAMB optimizer with DP-SGD we saw improvement of up to 20$\%$ points (absolute). Finally, we show that finetuning just the last layer for a \emph{single step} in the full batch setting, combined with extremely small-scale (near-zero) initialization leads to both SOTA results of 81.7 $\%$ under a wide privacy budget range of $ε\in [4, 10]$ and $δ$ = $10^{-6}$ while minimizing the computational overhead substantially.

Motivation & Objective

• Improve the privacy-utility trade-off in large-scale private image classification using transfer learning.
• Address the high computational cost and performance degradation of differentially private training on large models.
• Investigate the impact of model size, pretraining data scale, optimizer choice, and hyperparameters like learning rate and weight initialization on DP performance.
• Minimize computational overhead while maximizing model accuracy under strong privacy constraints (ε ∈ [4,10], δ = 10⁻⁶).

Proposed method

• Pre-train vision transformers (ViT-H/14-4b) on large public datasets (JFT-300M and JFT-4B) without privacy constraints.
• Fine-tune the pre-trained model on ImageNet using Differentially Private Stochastic Gradient Descent (DP-SGD) with full-batch, single-step updates.
• Apply the LAMB optimizer during private fine-tuning to improve convergence and performance with large batch sizes.
• Use near-zero or extremely small initialization for the final classification layer to enhance privacy-utility trade-off.
• Optimize hyperparameters including input resolution (256×256), cropping strategy (central crop), and learning rate scaling.
• Control privacy budget via noise multiplier σ and ensure (ε, δ)-differential privacy using standard DP-SGD mechanisms.

Experimental results

Research questions

• RQ1 Can large-scale pretraining on public datasets improve the accuracy of differentially private image classifiers on ImageNet?
• RQ2 How does the choice of optimizer (e.g., LAMB vs. SGD with momentum) affect performance under DP-SGD with large batch sizes?
• RQ3 What is the impact of last-layer weight initialization magnitude on final accuracy in private fine-tuning?
• RQ4 Can a single-step full-batch fine-tuning procedure achieve state-of-the-art results while minimizing computational cost?
• RQ5 How does input resolution influence model performance when training with large batches under differential privacy?

Key findings

• Using ViT-H/14-4b with pretraining on JFT-4B and fine-tuning the last layer for a single full-batch step achieved 81.7% top-1 accuracy on ImageNet under ε ∈ [4,10], δ = 10⁻⁶.
• The LAMB optimizer improved accuracy by up to 20 percentage points (absolute) compared to SGD with momentum under the same privacy budget (ε=10, δ=10⁻⁶).
• Initializing the last layer with near-zero weights improved performance significantly; larger initializations degraded accuracy, especially under DP.
• A 256×256 input resolution yielded the best performance in large-batch training, indicating a non-trivial relationship between resolution and DP utility.
• Training for a single epoch (equivalent to one full-batch step) outperformed longer training schedules, likely due to reduced noise accumulation and better privacy accounting.
• The combination of large-scale pretraining, LAMB optimization, and minimal fine-tuning achieved SOTA results that surpassed prior work, including Kurakin et al. (2022) and De et al. (2022), across all ε values tested.