Title: A Unified Fast Gradient Clipping Framework for DP-SGD
Venue: NeurIPS 2023
Accepted: yes
Keywords: differential privacy

dp-sgd

gradient clipping

computational complexity
Abstract: A well-known numerical bottleneck in the differentially-private stochastic gradient descent (DP-SGD) algorithm is the computation of the gradient norm for each example in a large input batch. When the loss function in DP-SGD is consists of an intermediate linear operation, existing methods in the literature have proposed decompositions of gradients that are amenable to fast norm computations. In this paper, we present a framework that generalizes the above approach to arbitrary (possibly nonlinear) intermediate operations. Moreover, we show that for certain operations, such as fully-connected and embedding layer computations, further improvements to the runtime and storage costs of existing decompositions can be deduced using certain components of our framework. Finally, preliminary numerical experiments are given to demonstrate the substantial effects of the aforementioned improvements.

A Unified Fast Gradient Clipping Framework for DP-SGD

Abstract

A well-known numerical bottleneck in the differentially-private stochastic gradient descent (DP-SGD) algorithm is the computation of the gradient norm for each example in a large input batch. When the loss function in DP-SGD consists of an intermediate linear operation, existing methods in the literature have proposed decompositions of gradients that are amenable to fast norm computations. In this paper, we present a framework that generalizes the above approach to arbitrary (possibly nonlinear) intermediate operations. Moreover, we show that for certain operations, such as fully-connected and embedding layer computations, further improvements to the runtime and storage costs of existing decompositions can be deduced using certain components of our framework. Finally, preliminary numerical experiments are given to demonstrate the substantial effects of the aforementioned improvements.

1 Introduction

Machine learning models - more specifically, neural network-based models - are becoming more popular in industrial applications, in user-facing products, and different scientific fields. The popularity of these models lies on their flexibility and ability to be trained on ever bigger datasets, which may contain personal information about individuals. As these models become bigger and more descriptive, ML practitioners need to ensure that the models, and their black-box interactions, do not reveal information about the data used to train the model. In fact, it has been shown repeatedly [18] that large neural-network models can be used to know if a particular example was used in the training data. Another line of attacks has demonstrated that one can actually reconstruct some training instances with simple interactions with a trained model [3].

The only known robust way of protecting against these attacks is to train models using differential privacy [10, 11]. Using this approach, ML practitioners provide an information theoretic guarantee that ensures that the final model does not depend on any individual example 1 . To date, the most popular method for training models with differential privacy is the differentially private stochastic gradient descent (DP-SGD) [1] method. In theory, the DP-SGD algorithm requires only minimal changes with respect to its non-private counterpart; one only requires to clip gradients observed in the training process and add some noise proportional to the clipping value. However, in practice, the (na¨ ıve) gradient clipping step has been shown to increase memory and computational costs in all popular learning platforms (JAX, Tensorflow, and Pytorch). More precisely, for a batch of size n a na¨ ıve implementation of DP-SGD requires calculating n gradients (one for each example in the batch) so they can be clipped. This is in stark contrast with most back-propagation-based training algorithms which calculate a single gradient. This implies a dependency on the runtime and memory in proportion to the batch size used to train the model, and for large models, this cost makes the

prospect of training with differential privacy simply not viable. Several papers [8, 17, 21] have also observed that large models need to be trained with large batch sizes to obtain competitive levels of utility, under reasonable amounts of privacy.

In [12], it was recognized that one could clip gradients without actually materializing every example gradient. This technique is referred to as ghost clipping . The initial proposal from [12] was specialized to fully connected, feed-forward neural networks that consist of only dense layers. The ghost clipping algorithm was later extended to handle convolution layers [13, 15] and attention layers [14].

However, each of the above solutions required an ad-hoc analysis on the correctness of their implementation. In this work we present a general analysis of the ghost clipping algorithm. To produce our results, we study this algorithm from the lens of linear operator theory. Our paper has three main contributions:

We unify years of ad-hoc analysis and interpretations of the ghost clipping algorithm under a single framework.

We provide a future-proof way of expanding the ghost-clipping technique to new layers with (possibly) nonlinear dependencies on their weights.

We demonstrate that by framing the ghost clipping problem in the language of linear operators, we can obtain a better performance of DP-SGD on models with embedding layers (crucial for personalization models consisting of embeddings of tens of millions of parameters) and fully connected layers with linear bias broadcasting operators.

Additionally, we show in our Appendix how to apply our framework to more complex tranforms/layers such as layer normalization and multi-head attention.

To complement the results of this paper, we open-sourced the general interface of the code using the TensorFlow Keras API 2 . By introducing an abstract interface, we also expect practitioners to easily extend the ghost clipping algorithm to any type of layer.

2 Notation and preliminaries

Throughout the paper, ( W , 〈· , ·〉 ) denotes a Hilbert space and ‖ · ‖ denotes its induced norm. Examples of W are R d with the standard dot product and the space of matrices W = R p × q with the Euclidean (Frobenius) inner product between A,B ∈ W given by 〈 A,B 〉 = tr ( A glyph[latticetop] B ) . For two matrics A,B we let ‖ A ‖ denote the Frobenius norm of A .

Given two Hilbert spaces ( W , 〈· , ·〉 W ) , ( Y , 〈· , ·〉 Y ) , we denote linear operators between them by italicized letters ( A : W → Y ) and A ∗ : Y → W to be the adjoint of A . That is, A ∗ is the unique linear operator that satisfies

Let ( W , 〈· , ·〉 W ) and ( Y , 〈· , ·〉 Y ) denote two Hilbert spaces with respective induced norms ‖ · ‖ W and ‖ · ‖ Y . Moreover, let ψ : W → Y be an arbitrary function. The Fr´ echet derivative of ψ at w 0 ∈ W is given by the unique bounded linear operator Dψ ( w 0 ): W → Y satisfying

We say ψ is differentiable if its Fr´ echet derivative exists for all w 0 ∈ W . Throughout this paper we will use two special properties of the Fr´ echet derivative: the chain rule and the existence of gradients. Let ( Z , 〈· , ·〉 Z ) be another Hilbert space and φ : Y → Z be given. The chain rule provides us with a simple way to calculate the derivative of the function φ ◦ ψ : W → Z , namely,

The Fr´ echet derivative of ψ at w 0 with respect to a subset of variables u is denoted by D u φ ( w 0 ) . Finally, ∇ ψ ( w 0 ) ∈ W denotes the (unique) gradient of a function ψ at w 0 , which satisfies

The existence of the gradient is guaranteed by the well-known Riesz-Fr´ echet Representation Theorem [16]. The gradient of ψ at w 0 with respect to a set of variables u is denoted by ∇ u ψ ( w 0 ) .

Algorithm 1 DP-SGD algorithm

Input: Data sample S , parameter initialization w 0 , number of rounds T , learning rate η , noise multiplier σ , clipping norm C .

Private Stochastic Gradient Descent We now turn our attention to the DP-SGD algorithm and specifically, to its instantiation for neural network-based models. Let X be an arbitrary space and S = { x 1 , . . . , x n } ⊂ X be a sample of examples. Let W denote a parameter space and h : X × W → R denote an arbitrary function. The stochastic gradient descent (SGD) algorithm solves the optimization problem

by iteratively updating the model paramters using gradients of the loss over a batch B ⊂ S of data.

The celebrated DP-SGD algorithm was introduced by [1] as a simple modification on the SGD algorithm to make it private 3 . Specifically, the DP-SGD algorithm, given in Algorithm 1, is identical to the SGD algorithm except in two steps:

The DP-SGD algorithm needs to calculate | B | (so-called) per-example gradients in order to clip them to have a bounded norm.

Adding Gaussian noise N (0 , C 2 σ 2 / | B | 2 ) for some noise multiplier σ > 0 .

Notice that the first step has a prohibitively large cost for large networks. Specifically, in the per example gradient calculation, notice that that the computational and memory usage of a na¨ ıve implementation 4 of DP-SGD algorithm increases as Θ( n w | B | ) , where n w is the number of parameters in the network. This increase in resources effectively negates the advantages of the backpropagation algorithm. In this paper we show how to run the DP-SGD algorithm with only a small constant increment in both the memory footprint and runtime of traditional SGD.

We will focus on the scenario where the learner is trying to minimize the loss across the space of neural networks. That is, the parameter vector w is a concatenation of k parameters ( w 1 , . . . , w k ) and there exists k functions φ 1 , . . . , φ k , and a loss function glyph[lscript] such that

Note that { φ j } correspond to the layers of the network and { w j } are the vectors parameterizing these layers.

3 Previous work

The DP-SGD algorithm was first introduced by [1]. Due to its simplicity to adapt to standard machine learning frameworks it is now probably the most popular method for training private

machine learning models. In the past few years, a large number of papers have been devoted to improving the privacy-utility trade-offs of DP-SGD [4] as well as demonstrating that DP-SGD can be used on multiple tasks [7, 9, 2]. While the main focus on this line of research has been in understanding the privacy-utility trade-offs, some authors did touch on the performance issues of calculating per-example gradients. A simple solution to reduce the memory and computational blow-up of the DP-SGD algorithm is to calculate the gradient with respect to a micro-batch of examples instead. While this reduced the computational cost of running DP-SGD, it came at the cost of decreased quality of the model learned. Some authors [2] got around this performance issues by using some JIT compilation features of JAX. However these authors still observed an increase in the memory use of their private implementations.

In order to speed up the gradient clipping step without sacrificing quality, [12] proposed the ghost clipping technique (see Section 4). ITs general idea is that one can obtain a private gradient estimate ˜ g without materializing each per example gradient as long as one knows the norm of each per-example gradient. Moreover, the authors show that one can easily calculate these norms by using information already materialized in the forward and backward passes of the back propagation algorithm for training neural networks. The results however were limited to neural networks consisting of only dense layers. The ghost clipping trick was later on applied to train transformers [14] and image classification models which consist of networks with convolution layers [15].

The above works provide ad-hoc proofs that their implementation of the ghost clipping algorithm is correct for their particular layer implementation but do not attempt to extend their results to arbitrary layers. In this work we present a way to extend the ghost clipping techinque to arbitrary layers and pinpoint what properties of each layer are crucial for enabling an efficient implementation of the DP-SGD algorithm.

Finally, it is also worth mentioning that paper [6] presents an efficient implementation of the ghost clipping technique in PyTorch for models with fully-connected, embedding, convolution, and normalization layers using variable caching to avoid recomputing certain intermediate parameter gradients. Our implementation also takes advantage of variable caching, but does this implicitly through TensorFlow's GradientTape API 5 .

4 The Ghost Clipping Algorithm

Algorithm 1 suggests that in order to implement DP-SGD one must calculate every per-example gradient, clip it to achieve a bounded norm and then aggregate it to obtain the private gradient estimate ˜ g . The main observation from [12] was that, with knowledge of the norms ‖ g ‖ x one could estimate ˜ g without explicitly materializing each per example gradient. To do this, the authors define weights r x = min { C/ ‖ g x ‖ , 1 } and a new weighted sum

Treating the weights r x as fixed (with respect to the parameters w ), one can use the linearity of the gradient operator to see that ∇ S ( w ) = ˜ g . Crucially, for neural networks, the gradient of S can be efficiently calculated using a single back-propagation step. The second contribution of [12] was showing that, for neural networks consisting of dense layers, one can also efficiently calculate ‖ g x ‖ for every x ∈ B using only a single forward and backward pass of the back propagation algorithm. That is, the cost of calculating a private gradient is simply twice that of calculating a non-private gradient step.

In the subsections below, we present a general formulation for calculating ‖ g x ‖ in a single forward and backward pass for arbitrary feed-forward neural networks, under the assumption that:

glyph[triangleright] each function φ i ( · , x ) can be decomposed as the composition of at least two Fr´ echet differentiable subfunctions.

Algorithm 2 General Gradient Norm Framework

Input: Data x ∈ X , parameter w = ( w 1 , . . . , w k ) , and functions { ψ x,i } k i =1 , { Z x,i } k i =1 satisfying

k

(

for any

˜

w

= ( ˜

w

1

, . . . ,

Compute

ζ

x,i

:=

Z

˜

w

x,i

)

w

i

)

Store a 'nice' representation Ω x,i ( · ) of the squared-norm function g ↦→ ‖ [ DZ x,i ( w i )] ∗ ( g ) ‖ 2 for any gradient g of the function glyph[lscript] x ◦ ψ x,i ( · )

end for

4.1 The General Norm Calculation Algorithm

Let w = ( w 1 , . . . , w k ) parameterize a neural network. We first note that ‖∇ w h ( x, w ) ‖ 2 = ∑ k j =1 ‖∇ w k h ( x, w ) ‖ 2 . Hence, we can focus on efficiently calculating the norm of the gradient corresponding to the parameters of each layer. Let us then fix a layer and denote ¯ w ∈ ¯ W to be its parameter vector for some restricted space ¯ W . Moreover, let the input x be fixed. We see that, as a function of the layer parameters ¯ w only, the loss h has the form

where φ x corresponds to operations performed by the fixed layer and all subsequent layers in the network, and glyph[lscript] x ( · ) ≡ glyph[lscript] ( x, · ) corresponds to the loss function, e.g., mean-squared error, for the input example x . The following result, whose proof can be found in the Appendix, gives an expression for the gradient of h under any decomposition of φ x .

Proposition 4.1. Let x ∈ X be fixed and let ( glyph[lscript] x , φ x ) be a pair of Fr´ echet differentiable functions satisfying (3) . Moreover, let ( ψ x , Z x ) be a pair of Fr´ echet differentiable functions satisfying φ x = ψ x ◦ Z x , and denote

The above result gives us an alternative way of computing ‖∇ ¯ w h ( x, ¯ w ) ‖ 2 , namely, (i) pick a decomposition ( ψ x , Z x ) of φ x and (ii) compute ‖A ∗ g ‖ 2 in some efficient manner. In later subsections, we provide examples of layers and decompositions where this two-step approach is drastically more efficient than the naive approach of materializing ∇ ¯ w h ( x, ¯ w ) and computing its Euclidean norm.

In view of the decomposition in Proposition 4.1, we present a unified method for computing ‖∇ h ( x, w ) ‖ in Algorithm 2. Specifically, it consists of two loops over the parameters w i . The initial loop is performed to obtain the intermediate outputs ζ i,x and some sort of 'nice' or efficient representation of the squared-norm functions Ω i,x . The follow-up loop, also over the parameters w i , computes some necessary intermediate gradients and combines them with the functions Ω i,x to obtain the gradient decomposition (and subsequent norm computation) given by Proposition 4.1.

4.2 Efficient Implementation of Algorithm 2

We now describe how Algorithm 2 can be implemented efficiently when h ( · , · ) is formed by a neural network N . First, it is shown in Appendix B.1 that the intermediate quantities { ζ i,x } k i =1 and { g i,x } k i =1 can be computed in a single forward and backward pass of N (as opposed to na¨ ıvely traversing the network Θ( k ) times). Second, it is shown in Appendix B.2 that for any i , the batch gradients {∇ ζ x,i ( glyph[lscript] x ◦ ψ x,i )( ζ x,i ) } x ∈ B can be obtained a single (batched) backward pass of N if (i) condition (5) holds for every x ∈ B and (ii) the gradient g x,i is replaced by ∇ ζ x,i { ∑ x ∈ B glyph[lscript] x ◦ ψ x,i ( ζ x,i ) } . Consequently, if we ignore the costs of forming and evaluating Ω x,i ( · ) for each x ∈ B

, where glyph[lscript]

x

(

·

)

≡

glyph[lscript]

(

x,

·

)

.

and i = 1 , . . . , k , then the above facts imply that we can obtain the norms {‖∇ h ( x, w ) ‖} x ∈ B using only one additional forward and (batched) backward pass of N .

5 Efficient Squared-Norm Functions

This section presents efficient representations of the functions Ω i ( · ) in Algorithm 2 for some basic layer functions φ x ( w ) ≡ φ ( w,x ) . Examples involving more complicated layer functions (e.g., layer normalization and multi-head attention) can be found in the Appendix.

Each subsection begins with a precise description of φ x ( · ) , presents a decomposition φ ( · ) = ψ x ◦ Z x ( · ) , and gives an analysis of the runtime and storage costs of an efficient representation of the function Ω x ( · ) ≡ ‖ [ DZ x ( w )] ∗ ( · ) ‖ 2 . For the sake of conciseness, the proofs of these decompositions are given in the Appendix.

In addition, we make comparisons of our approach with the naive approach (of computing gradients for each example) and other ghost clipping-like approaches. The results of these comparisons are summarized in Tables 1-2 (see the first paragraphs in Subsections 5.1-5.3 for descriptions of the variables/dimensions).

5.1 Fully-Connected Layers

Given variables V ∈ R p × q and b ∈ R m , a layer input U x ∈ R r × p , an activation function α : R r × q ↦→ R r × q , and a linear broadcasting operator Q : R m ↦→ R r × q satisfying m | rq , the standard fully-connected layer function φ x ( · ) is given by

Typically, p , q , r , and m are called the input dimension , output dimension , channel dimension , and bias dimension , respectively. Usually, it is the case that r glyph[lessmuch] min { p, q } and a common case is r = 1 .

Wenowconsider the squared-norm function Ω x : R r × q ↦→ R generated by the choice of Z x ( V, b ) := U x V + Q b . Denoting A as in (4), for some ¯ w = ( V, b ) , we have that

for any g ∈ R r × q . Hence, Ω x ( · ) can be efficiently represented by U x U ∗ x and an efficient representation of the function g ↦→‖Q ∗ g ‖ 2 . For an example of the latter, suppose Q is the operator that repeats b in an r -byq matrix row-wise ρ := rq/m times. Then, defining the row/column maps

it follows from the definition of the adjoint that

For this example, we incur a compute (resp. storage) cost of Θ( r 2 p ) (resp. Θ( r 2 ) ) for materializing U x U ∗ x and a compute and storage cost of O (1) for the function g ↦→‖Q ∗ g ‖ 2 , which does not depend on the data x . On the other hand, evaluation of Ω x ( · ) incurs: (i) a compute (resp. storage) cost of Θ( r 2 q | B | ) (resp. Θ( r 2 | B | ) ) for materializing each gg ∗ in the batch B and then computing the inner product term 〈 U x U ∗ x , gg ∗ 〉 and (ii) a compute cost of Θ( rq | B | ) for evaluating ‖Q ∗ g ‖ , which only requires a single pass through the entries of g for each example. Consequently, the total runtime (resp. storage) cost for a batch of examples B is

Let us now make a comparison of this decomposition with the naive and ghost clipping approaches. For our discussion, suppose Q is as in the previous example.

In the naive approach of computing ∇ ( V,b ) h ( x, ( V, b )) for each x ∈ B , it is straightforward 8 to show that the compute (resp. storage) costs are Θ( | B |{ pq + m } ) (resp. Θ( | B |{ p + r } q ).

For the classic ghost clipping approach considered in [12, 13], our approach is identical in the case of Q ≡ 0 . When Q glyph[negationslash]≡ 0 , this approach instead groups Q together with U x and suggests the decomposition in (7) but with U x = [ U x ; Q ] and Q = 0 . Let us now analyze the costs in the worst-case scenario when QQ ∗ is dense, e.g., when m = 1 . In order to evaluate 〈QQ ∗ , gg ∗ 〉 using a consistent representation of QQ ∗ , we must materialize Q (resp. g ) as a rq -bym matrix ( rq -by1 column vector) and compute QQ ∗ ∈ R rq × rq (resp. gg ∗ ∈ R rq × rq ) using matrix multiplication. Clearly, the compute (resp. storage) costs for 〈QQ ∗ , gg ∗ 〉 are then Θ( m { rq } 2 + | B |{ rq } 2 ) (resp. Θ( mrq + | B |{ rq } 2 ) ) given Q which, together with the costs for ‖ U ∗ x g ‖ 2 , yield the complexity in Table 2 (resp. Table 1).

5.2 Embedding Layers

Given variables W ∈ R r × d , input indices π ( x ) = [ π 1 ( x ) , . . . , π q ( x )] ∈ { 1 , . . . , r } q , the standard embedding layer function φ x ( · ) is given by

where Y π ( x ) ∈ R q × r is the one-hot matrix whose glyph[lscript] -th row is the π glyph[lscript] ( x ) -th basis vector in R r . Typically, r , d , and q are called the vocabulary size , embedding dimension , and number of queries , respectively. Usually, it is the case that max { d, q } glyph[lessmuch] r , and a common case is q = 1 .

We now consider the squared-norm function Ω x : R q × d ↦→ R generated by the choice of Z x ( W ) = φ x = Y π ( x ) W for any gradient g = ∇ ( glyph[lscript] x ◦ ψ x )( Z x ( W )) . First, define the quantities

for k = 1 , . . . , q , where n k ( x ) denotes the number times index k appears in π ( x ) . Next, notice that g contains at most ˜ q unique rows, due to the definition of the gradient and the fact that Z x ( W ) has at most ˜ q unique rows. Denoting A as in (4), for some ¯ w = W , it can be shown that

where ˜ g i denotes any row j of g where π j ( x ) = i (note that all rows of g that satisfy this condition are the same). Hence, Ω x ( · ) can be efficiently represented by the nonzero values of { n k ( x ) } r k =1 . An efficient algorithm using O (˜ q log q ) runtime and O (˜ q ) extra storage, for computing { n x ( x ) } r k =1 , for a fixed x ∈ B , is given in Algorithm 3 of the Appendix. On the other hand, the evaluation of Ω x ( · ) only requires evaluating the ˜ q unique rows of an input g ∈ R q × d (given by π ( x ) ) and multiplying

them with the ˜ q nonzero cached values { n k ( x ) } r k =1 , which can be done with a Θ( | B | ˜ qd ) runtime cost for a batch B . Consequently, the total runtime (resp. storage) cost for a batch of examples B is Θ( | B |{ q log ˜ q + ˜ qd } ) (resp. Θ( | B | ˜ q ) ) using Algorithm 3.

Let us now make a comparison of this decomposition with the naive approach and the ghost clipping approach in [13].

In the naive approach of computing ∇ W h ( x, W ) for each x ∈ B , it is straightforward 9 to show that the compute (resp. storage) costs are Θ( | B | qd ) (resp. Θ( | B | qd ) ).

In the classic ghost clipping approach, we treat φ x ( · ) as a linear operator, compute Y π ( x ) Y ∗ π ( x ) , and choose the representation Ω x ( g ) = 〈 Y π ( x ) Y ∗ π ( x ) , gg ∗ 〉 . In the worst-case scenario, it is straightforward to see that Y π ( x ) Y ∗ π ( x ) can be fully dense, e.g., ˜ q = 1 . In this setting, even if Y π ( x ) is sparsely represented by a small set of ˜ q unique indices, computing Y π ( x ) Y ∗ π ( x ) still incurs a compute and storage cost of O ( q 2 ) . Now, since each example gradient g consists of q embedding vectors in R d , for a batch B , the compute (resp. storage) cost of materializing gg ∗ is Θ( | B | q 2 d ) (resp. Θ( | B | q 2 ) ). Combining the above complexities with the Θ( | B | q 2 ) runtime cost of computing the desired inner products, i.e. (7) with U x = Y π ( x ) for x ∈ B , we obtain the complexities in Tables 1-2.

5.3 Low Rank Approximation Layer

Given input matrix U x ∈ R n × n , one way [19, 20] to encourage a rankk (or lower) approximation of U x is to add the intermediate layer transform

for some V ∈ R n × k , where ρ ( · ) is a sparsity promoting regularizer (e.g., glyph[lscript] 1 norm, SCAD, MCP) and σ ( · ) is the function that maps matrices to their singular values. Expanding the norm term, note that the above function can be equivalently (ignoring terms depending solely on U x ) expressed as

for some function R : R n × k ↦→ R . Note that R ( V ) does not depend on x and, hence, its computation does not depend the batch B .

In view of (8), we now consider the squared-norm function Ω x : R ↦→ R generated by the choice of Z x ( V ) = 〈 U x , V V ∗ 〉 / 2 , where, clearly, one has φ x ( V ) = -4 Z x ( V ) + R ( V ) . Denoting A as in (4), for some ¯ w = V , it can be shown that

for g ∈ R . Hence, Ω x ( · ) can be efficiently represented by the scalar ‖ ( U x + U ∗ x ) V ‖ 2 . It is straightforward to see that computing ‖ ( U x + U ∗ x ) V ‖ requires only a Θ( n 2 k ) runtime cost and a Θ(1) storage cost. Moreover, evaluation of Ω x ( · ) , given ‖ ( U x + U ∗ x ) V ‖ , requires only a O (1) runtime cost.

In the naive approach of computing ∇ V h ( x, V ) for each x ∈ B , it is straightforward to see that the compute (resp. storage) costs is Θ( | B | n 2 k ) (resp. Θ( | B | nk ) ) due to the excessive computation (resp. storage) of g ( U x + U ∗ x ) V for x ∈ B . The authors are not aware of any ghost clipping-like techniques in the nonlinear setting.

6 Numerical Experiments

This section presents numerical experiments that compare our proposed adjoint-based framework ( Adjoint ) against the na¨ ıve implementation of DP-SGD ( Naive ), which computes gradients for each example in a batch, and the classic ghost clipping frameworks ( GhostClip ) that are described in Subsections 5.1 and 5.2. Specifically, it presents runtimes and memory costs for the gradient norm computation of fully-connected and embedding layers.

Each problem instances was run on a cloud computing platform consisting of (i) 112 Intel(R) Xeon(R) Platinum processors with 28 cores each, (ii) 64 GB of RAM, (iii) Python 3.10.11, and (iv) Tensorflow 2.14. For simplicity, memory usage was measured as the peak amount of heap memory utilized in the run of single gradient norm computation. We also simplify our computations by utilizing batches of size | B | = 1 for the first two subsections. The loss function used in our experiment, glyph[lscript] x ( · ) , is the mean-squared-error. To reduce the variance of the results in the first two subsections, we repeat each problem instance 20 times and report only the median runtime and memory cost over the repetitions.

6.1 Fully-Connected Layer

Figure 1 presents numerical results for the setting considered in Subsection 5.1, where Q is the (linear) broadcasting operator that duplicates the bias rq/m times to match the dimension of the layer outputs. It specifically plots the effect of the bias dimension m for various values of the output dimension q . For simplicity, all problem instances fix an input and channel dimension of 2 and 4096 , respectively. Additional experiments, involving the effect of batch size, are given in Appendix F.

The results in Figure 1 demonstrate that the runtime and memory costs of Adjoint are marginal compared to those of GhostClip for fully-connected layers. These results also support the analysis of Subsection 5.1 in that: (i) GhostClip 's runtime has a stronger positive dependence on the bias dimension m than Adjoint 's runtime and (iii) GhostClip 's runtime has a strong positive dependence on the output dimension q .

6.2 Embedding Layer

Figure 2 presents numerical results for the setting considered in Subsection 5.2. It specifically plots the effect of the number of queries q for various values of the vocabulary size r . For simplicity, all problem instances fix the embedding dimension to be 10 and the embedding indices { π i ( x ) } q i =1 are chosen uniformly at randomly from the set { 1 , . . . , r } .

The results in Figure 2 demonstrate that the runtime and memory costs of Adjoint are marginal compared to those of GhostClip for embedding layers. Moreover, the memory cost graph supports

the analysis in Subsection 5.2 in that GhostClip 's memory cost has a significantly stronger positive dependence on the query size q compared to Adjoint 's memory cost.

6.3 Small BERT model

Figure 3 presents numerical results for an end-to-end training run of a small BERT model that consists of dense, embedding, normalization, and multi-head attention sublayers. It specifically plots the effect of the batch size | B | on the runtime and peak memory usage of training the model. To be more precise, the BERT model is an instance of the TensorFlow BertEncoder model with a vocabulary size of 100, one intermediate transformer layer, and all other parameters set to their default values. Each experiment consists of a single training loop of 50 iterations with uniformly sampled random input data of a query size of 5. The efficient squared norm functions were taken from the descriptions in Section 5 and Appendices E.1 and E.2.

The results in Figure 3 demonstrate that the runtimes and peak memory usages of Adjoint scale better and are comparatively smaller than the corresponding ones for Naive .

7 Concluding Remarks

The analysis in Subsection 5.1 may also be applied to more complex layers whose parameter transformations φ x ( · ) primarily involve linear transforms. For example, it is shown in [5, Subsection 2.3] that 2D-convolution layers are equivalent to fully-connected layers when the example image is appropriately transformed to form the matrix U x in (6) and, hence, our storage and runtime savings for fully-connected layers easily apply to 2D- (and generally n D-) convolution layers. Another example is the multihead attention layer, whose parameter transforms are simple matrix multiplications (see Appendix E.1 for a derivation).

7.1 Limitations

The proposed framework only applies to layers with at least one differentiable intermediate transformation and models with differentiable losses. Moreover, the framework does not support shared trainable layers.

References

Martin Abadi, Andy Chu, Ian Goodfellow, H Brendan McMahan, Ilya Mironov, Kunal Talwar, and Li Zhang. Deep learning with differential privacy. In Proceedings of the 2016 ACM SIGSAC conference on computer and communications security , pages 308-318, 2016.

Rohan Anil, Badih Ghazi, Vineet Gupta, Ravi Kumar, and Pasin Manurangsi. Large-scale differentially private bert. arXiv preprint arXiv:2108.01624 , 2021.

Borja Balle, Giovanni Cherubin, and Jamie Hayes. Reconstructing training data with informed adversaries. In 2022 IEEE Symposium on Security and Privacy (SP) , pages 1138-1156. IEEE, 2022.

Raef Bassily, Vitaly Feldman, Kunal Talwar, and Abhradeep Guha Thakurta. Private stochastic convex optimization with optimal rates. Advances in neural information processing systems , 32, 2019.

Zhiqi Bu, Jialin Mao, and Shiyun Xu. Scalable and efficient training of large convolutional neural networks with differential privacy. Advances in Neural Information Processing Systems , 35:38305-38318, 2022.

Zhiqi Bu, Yu-Xiang Wang, Sheng Zha, and George Karypis. Differentially private optimization on large model at small cost. In International Conference on Machine Learning , pages 31923218. PMLR, 2023.

Aldo Gael Carranza, Rezsa Farahani, Natalia Ponomareva, Alex Kurakin, Matthew Jagielski, and Milad Nasr. Privacy-preserving recommender systems with synthetic query generation using differentially private large language models. arXiv preprint arXiv:2305.05973 , 2023.

Soham De, Leonard Berrada, Jamie Hayes, Samuel L Smith, and Borja Balle. Unlocking high-accuracy differentially private image classification through scale. arXiv preprint arXiv:2204.13650 , 2022.

Carson Denison, Badih Ghazi, Pritish Kamath, Ravi Kumar, Pasin Manurangsi, Krishna Giri Narra, Amer Sinha, Avinash Varadarajan, and Chiyuan Zhang. Private ad modeling with dpsgd. arXiv preprint arXiv:2211.11896 , 2022.

Cynthia Dwork, Frank McSherry, Kobbi Nissim, and Adam Smith. Calibrating noise to sensitivity in private data analysis. In Theory of Cryptography: Third Theory of Cryptography Conference, TCC 2006, New York, NY, USA, March 4-7, 2006. Proceedings 3 , pages 265-284. Springer, 2006.

Cynthia Dwork, Aaron Roth, et al. The algorithmic foundations of differential privacy. Foundations and Trends in Theoretical Computer Science , 9(3-4):211-407, 2014.

Ian Goodfellow. Efficient per-example gradient computations. arXiv preprint arXiv:1510.01799 , 2015.

Jaewoo Lee and Daniel Kifer. Scaling up differentially private deep learning with fast perexample gradient clipping. Proceedings on Privacy Enhancing Technologies , 2021(1), 2021.

Xuechen Li, Florian Tramer, Percy Liang, and Tatsunori Hashimoto. Large language models can be strong differentially private learners. arXiv preprint arXiv:2110.05679 , 2021.

Gaspar Rochette, Andre Manoel, and Eric W Tramel. Efficient per-example gradient computations in convolutional neural networks. arXiv preprint arXiv:1912.06015 , 2019.

Walter Rudin et al. Principles of mathematical analysis , volume 3. McGraw-hill New York, 1976.

Tom Sander, Pierre Stock, and Alexandre Sablayrolles. Tan without a burn: Scaling laws of dp-sgd. In International Conference on Machine Learning , pages 29937-29949. PMLR, 2023.

Reza Shokri, Marco Stronati, Congzheng Song, and Vitaly Shmatikov. Membership inference attacks against machine learning models. In 2017 IEEE symposium on security and privacy (SP) , pages 3-18. IEEE, 2017.

Fei Wen, Lei Chu, Peilin Liu, and Robert C Qiu. A survey on nonconvex regularization-based sparse and low-rank recovery in signal processing, statistics, and machine learning. IEEE Access , 6:69883-69906, 2018.

Quanming Yao and James Kwok. Efficient learning with a family of nonconvex regularizers by redistributing nonconvexity. In International Conference on Machine Learning , pages 26452654. PMLR, 2016.

Yaodong Yu, Maziar Sanjabi, Yi Ma, Kamalika Chaudhuri, and Chuan Guo. Vip: A differentially private foundation model for computer vision. arXiv preprint arXiv:2306.08842 , 2023.

Peer reviews:
comments: SUMMARY:
The paper provides a unification of ad-hoc analysis and interpretations of the ghost clipping algorithm under a single framework. It also shows that for certain operations, such as fully-connected and embedding layer computations, further improvements to the runtime and storage costs of existing decompositions can be deduced using certain components of the proposed framework.

STRENGTHS:
- **Originality**: The paper proposes a novel framework for fast gradient clipping used in DP-SGD. This is a significant theoretical contribution to the field of privacy of machine learning.
- **Quality**: The paper is well-organized. A thorough explanation of the proposed framework, including mathematical derivations and experimental results is provided.
- **Clarity**: The paper is easy to follow and understand, even for readers who are not experts in the field. The authors provide clear explanations of the background, motivations, concepts, and techniques used, and the experimental results are presented in a concise and understandable manner.
- **Significance**: The paper has significant implications for the specific field of privacy of machine learning. It provides a unified perspective on the ghost clipping algorithm in DP-SGD.

WEAKNESSES:
The paper focuses only on a specific technique of gradient clipping, i.e., ghost clipping. Whether this ghost clipping lies at the heart of the big field of privacy of deep learning or not is unclear to me.

QUESTIONS:
How can the proposed framework be generalized beyond the ghost-clipping algorithm?

LIMITATIONS:
Limitations are not discussed by the authors. Perhaps the possible relation or extension to other clipping algorithms can be mentioned in the final as a separate concluding section which is now missing.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-30

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provides a unified framework for efficiently computing the gradient norms of individual examples for a wide range of neural network architectures, which significantly decreases runtime and storage costs for implementing DP-SGD.\n\nSTRENGTHS:\n1. The considered problem is important: Computing per-example gradients is a computation bottleneck for previous implementations of DP-SGD.\n2. Theoretical and empirical improvements over baselines are clearly presented and significant. Tables 1 and 2 summarize the time and storage complexity for previous work and the proposed algorithm, and Figures 1 and 2 summarize the observed improvements in experiments. The theoretical and the empirical improvements are both convincing.\n3. The presentation is rigorous and clear.\n\nWEAKNESSES:\n1. The experimental evaluation does not provide a big-picture idea of the end-to-end savings provided by the proposed framework. Figures 1 and 2 exhibit the runtime/storage costs only for the gradient computation operation, and it does not show how the overall costs are affected. If gradient computation accounts for a small portion of the overall runtime, then the overall time savings will not be significant. As it stands, the experimental results do not give a concrete idea for the overall savings in runtime, which is the metric of interest for practitioners.\n2. It would be nice to see an empirical evaluation for the more popular layer types mentioned in the text, e.g. convolution and attention.\n\nQUESTIONS:\n1. What is the overall reduction in runtime/storage cost provided by the proposed framework?\n2. How does the empirical performance of the proposed framework compare to that of ghost clipping for convolution layers and attention layers?\n\nLIMITATIONS:\nThere is no discussion of societal impact beyond the underlying (implied) importance of user privacy. However, I don't think any further discussion of societal impact is necessary for this work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-03"}

comments: SUMMARY:
This paper unifies a framework that enables us to apply ghost clipping algorithms to new layers of neural networks, apart from fully connected, feed-forward layers. In particular, the paper discusses efficient implementations of computing norms for running DP-SGD to train neural networks.

STRENGTHS:
The contributions of this paper are as follows: 

1. It shows how to utilize the proposed framework to apply ghost clipping techniques to new layers for DP-SGD on training neural networks. Three main examples are included to discuss how to implement efficient algorithms to compute the gradient norm and to run DP-SGD. 

2. Experimental results evaluate the effectiveness of their proposed framework over the classical ghost clipping framework in run time and memory requirements.

WEAKNESSES:
Literature on the importance of DP-SGD and clipping operators is lacking in this paper. It is better to include them in Section 3 (Previous work) or in the introduction section to motivate the study of this paper in a broader impact.

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the memory/time complexity of the per-sample gradients norm computation step in the DP-SGD algorithm. The authors present a general theoretical framework that generalizes the idea of ghost norm computation (computing the norm of the persample gradients without having to store the gradients), and can be applied to arbitrary layers of a DNN (if an ad hoc decomposition can be found). The authors show early experiments on the benefits of their methods compared to previous ghost clipping techniques.\n\nSTRENGTHS:\noriginality: - Computing persamble gradient norms is a real bottleneck when training DNNs with DP-SGD. Typically, the Opacus library (Pytorch for privacy) stores the gradients per sample before computing the norms, which can be memory intensive, especially for large architectures. Other methods have been proposed to compute the norm without storing the gradients (ghost norm), but no general framework has been proposed yet. This paper gives a clear problem formulation and a general theoretical framework to approach this problem. I was very happy to see a theoretical work that tries to tackles this issue.\nquality: The math are correct\nclarity:  ok\nsignificance: using the theory of linear operator to improve time and space complexity of DP-SGD is good novelty that could benefit the community\n\nWEAKNESSES:\n- Typos: l108 “authors”\n- Overall, the experiments are very insufficient. They are performed with batches of size 1, for which the vanilla per-sample gradient storage method has a better running time, but with which it is never compared. The authors should investigate in detail how their method would work with GPUs to see how it would be beneficial in practice: the complexity bounds given are linear in the batch size, but in practice we benefit greatly from parallel computation/optimized matrix multiplication. \n- Experiments should be performed on whole (even small) networks used in practice, with batches of size greater than one. In Li et al (2022), they compare ghost norm to opacus when training a GPT2-like architecture and find similar runtimes; a similar comparison could be done here (even with much smaller versions of GPT2 when computational power is limited). \n- It is not the same abstract in OpenReview and in the article. The first states that savings can be as high as 150x: this could be the case in very specific situations chosen for the experiments, but the authors do not elaborate on why these choices are realistic in practice. \n- The section on related work is sparse, with very few papers cited (no mention of why gradients need to be clipped, no mention of the whole deep learning literature where this is a real bottleneck (Opacus))\n- No code for reproduction\n\nQUESTIONS:\n- Typically Opacus (storing per sample gradients) or ghost norm will be 2x slower than non DP training. How can the proposed method be 150x faster? That would mean that the proposed method is faster than non DP training?\n- Why does the green curve start decreasing at some point in the memory figure for the embedding layer?\n- For the embedding layer, the authors state that q=1 in practice but show memory/runtime for q = {1000, 2000, …, 10000}. Can you please clarify on these choices for comparison?\n- The paper states that it can be applied to arbitrary layers, but the examples are linear layers (which were already tackled in ghost norm previous work). What are the conditions for the existence of the decomposition in proposition 4.1? What do the authors have in mind by “efficient manner” (line 156)? When is that possible?\n\nLIMITATIONS:\nThe paper does not include a limitation section. \nSuggestions for improvement would be to perform a rigorous empirical study of the proposed method, with entire architectures and using micro batches, and see if can easily used with an existing DL library (pytorch, JAX...).", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work unifies the analysis of the ghost clipping algorithm, which was previously analyzed for specific architectures only. The proposed framework extends previous results to a wide class of network architectures including new applications and provides asymptotically faster computation times than the previous case-by-case approaches. Some numerical simulations confirm the effectiveness of the proposed algorithm demonstrating faster computation time in practice.\n\nSTRENGTHS:\nThe paper is nicely structured and well motivated. The considered family of functions (NN architectures) is quite general and the proposed approach introduced by Proposition 4.1 seems to be very principled. Despite of the generality, the strategy proposed in Proposition 4.1 can be implemented for a number of important special classes of functions such as fully-connected networks, embedding layers and low rank approximation layers. The authors demonstrate improved runtime and storage costs for this algorithm in these three special cases.\n\nWEAKNESSES:\n-- The main motivation of the work is that when the batch sizes of DP-SGD is large, the runtime and storage can be large. Why this problem cannot be simply solved by using parallel computing (assigning different batches to different workers)? If this is one possible approach, I believe this should be mentioned and discussed in the related work. \n\n-- It seems that the assumptions on the loss function and the layers $\\phi_x$ are not explicitly stated. It is unclear to me if it is sufficient to have all $\\phi$ and the loss function to be Frechet differentiable or any additional technical assumptions are needed for Proposition 4.1. Why the sub gradient of $Z_x$ might not be unique while $\\nabla h$ and $g$ in this proposition are unique?\n\nQUESTIONS:\n- Why in Tables 1 and 2 some entries come with $\\Theta(\\cdot)$ and some $O(\\cdot)$? Are there lower bounds for storage and runtime costs for these settings? \n\n- Does equation (3) hold for the last layer only? \n\n- Typo in line 54. \n\n\n---\nUpdate: The authors addressed my concerns and I would like to increase my score from 5 to 6 to further support the acceptance.\n\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}