NanoBatch Privacy: Enabling fast Differentially Private learning on the IPU

TL;DR

NanoBatch Privacy enables fast, scalable differentially private deep learning on IPUs by using gradient accumulation and batch size of 1, significantly improving speed and efficiency over existing GPU-based methods.

Contribution

The paper introduces NanoBatch Privacy, a lightweight add-on for TensorFlow Privacy that leverages IPUs with gradient accumulation to improve DP training speed and accuracy.

Findings

Achieves over 15x speedup on ImageNet compared to TFDP on GPUs.

Maintains high accuracy with larger batch sizes in DP training.

Demonstrates effective DP training for Covid-19 chest CT prediction.

Abstract

Differentially private SGD (DPSGD) has recently shown promise in deep learning. However, compared to non-private SGD, the DPSGD algorithm places computational overheads that can undo the benefit of batching in GPUs. Micro-batching is a common method to alleviate this and is fully supported in the TensorFlow Privacy library (TFDP). However, it degrades accuracy. We propose NanoBatch Privacy, a lightweight add-on to TFDP to be used on Graphcore IPUs by leveraging batch size of 1 (without microbatching) and gradient accumulation. This allows us to achieve large total batch sizes with minimal impacts to throughput. Second, we illustrate using Cifar-10 how larger batch sizes are not necessarily optimal from a privacy versus utility perspective. On ImageNet, we achieve more than 15x speedup over TFDP versus 8x A100s and significant speedups even across libraries such as Opacus. We also provide two extensions: 1) DPSGD for pipelined models and 2) per-layer clipping that is 15x faster than the Opacus implementation on 8x A100s. Finally as an application case study, we apply NanoBatch training for use on private Covid-19 chest CT prediction.