SlowGuess/NeurIPS_2023_OCR huggingface dataset:
for any  $g \in \mathbb{R}^{r \times q}$ . Hence,  $\Omega_x(\cdot)$  can be efficiently represented by  $U_x U_x^*$  and an efficient representation of the function  $g \mapsto \| \mathcal{Q}^* g \|^2$ . For an example of the latter, suppose  $\mathcal{Q}$  is the operator that repeats  $b$  in an  $r$ -by-  $q$  matrix row-wise  $\rho \coloneqq rq / m$  times. Then, defining the row/column maps
$$
\pi (\ell , k) := 1 + \left\lfloor \frac {(\ell - 1) m + (k - 1)}{q} \right\rfloor , \quad \xi (\ell , k) := 1 + \left\{\left[ (\ell - 1) m + (k - 1) \right] \bmod r \right\},
$$
Based on the example in Section 5.
Using Algorithm 3 for ours.
it follows from the definition of the adjoint that
$$
\left[ \mathcal {Q} ^ {*} g \right] _ {k} = \sum_ {\ell = 1} ^ {\rho} g _ {\pi (\ell , k), \xi (\ell , k)} \quad \forall g \in \mathbb {R} ^ {r \times q}, \quad k = 1, \dots , m.
$$
For this example, we incur a compute (resp. storage) cost of  $\Theta(r^2p)$  (resp.  $\Theta(r^2)$ ) for materializing  $U_x U_x^*$  and a compute and storage cost of  $O(1)$  for the function  $g \mapsto \|\mathcal{Q}^* g\|^2$ , which does not depend on the data  $x$ . On the other hand, evaluation of  $\Omega_x(\cdot)$  incurs: (i) a compute (resp. storage) cost of  $\Theta(r^2q|B|)$  (resp.  $\Theta(r^2|B|)$ ) for materializing each  $gg^*$  in the batch  $B$  and then computing the inner product term  $\langle U_x U_x^*, gg^* \rangle$  and (ii) a compute cost of  $\Theta(rq|B|)$  for evaluating  $\|\mathcal{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
$$
\Theta \left(r ^ {2} p + r ^ {2} q | B | + r q | B |\right) = \Theta \left(r ^ {2} \{p + | B | q \}\right)
$$
$$
[ \operatorname {r e s p.} \Theta (1 + r ^ {2} + r ^ {2} | B |) = \Theta (| B | r ^ {2}) ].
$$
Let us now make a comparison of this decomposition with the naive and ghost clipping approaches. For our discussion, suppose  $\mathcal{Q}$  is as in the previous example.
In the naive approach of computing  $\nabla_{(V,b)}h(x,(V,b))$  for each  $x\in B$ , it is straightforward to show that the compute (resp. storage) costs are  $\Theta (|B|\{pq + m\})$  (resp.  $\Theta (|B|\{p + r\} q)$ ).
For the classic ghost clipping approach considered in [12, 13], our approach is identical in the case of  $\mathcal{Q} \equiv 0$ . When  $\mathcal{Q} \neq 0$ , this approach instead groups  $\mathcal{Q}$  together with  $U_x$  and suggests the decomposition in (7) but with  $U_x = [U_x; \mathcal{Q}]$  and  $\mathcal{Q} = 0$ . Let us now analyze the costs in the worst-case scenario when  $\mathcal{QQ}^*$  is dense, e.g., when  $m = 1$ . In order to evaluate  $\langle \mathcal{QQ}^*, gg^* \rangle$  using a consistent representation of  $\mathcal{QQ}^*$ , we must materialize  $\mathcal{Q}$  (resp.  $g$ ) as a  $rq$ -by- $m$  matrix ( $rq$ -by-1 column vector) and compute  $\mathcal{QQ}^* \in \mathbb{R}^{rq \times rq}$  (resp.  $gg^* \in \mathbb{R}^{rq \times rq}$ ) using matrix multiplication. Clearly, the compute (resp. storage) costs for  $\langle \mathcal{QQ}^*, gg^* \rangle$  are then  $\Theta(m\{rq\}^2 + |B|\{rq\}^2)$  (resp.  $\Theta(mrq + |B|\{rq\}^2)$ ) given  $\mathcal{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 \in \mathbb{R}^{r \times d}$ , input indices  $\pi(x) = [\pi_1(x), \dots, \pi_q(x)] \in \{1, \dots, r\}^q$ , the standard embedding layer function  $\phi_x(\cdot)$  is given by
$$
\phi_ {x} (W) = Y _ {\pi (x)} W,
$$
where  $Y_{\pi (x)} \in \mathbb{R}^{q\times r}$  is the one-hot matrix whose  $\ell$ -th row is the  $\pi_{\ell}(x)$ -th basis vector in  $\mathbb{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\} \ll r$ , and a common case is  $q = 1$ .
We now consider the squared-norm function  $\Omega_x: \mathbb{R}^{q \times d} \mapsto \mathbb{R}$  generated by the choice of  $Z_x(W) = \phi_x = Y_{\pi(x)}W$  for any gradient  $g = \nabla (\ell_x \circ \psi_x)(Z_x(W))$ . First, define the quantities
$$
n _ {k} (x) := | \{i: \pi_ {i} (x) = k, i = 1, \dots , q \} |, \tilde {q} := | \{i: n _ {i} (x) > 0 \} |,
$$
for  $k = 1, \ldots, q$ , where  $n_k(x)$  denotes the number times index  $k$  appears in  $\pi(x)$ . Next, notice that  $g$  contains at most  $\tilde{q}$  unique rows, due to the definition of the gradient and the fact that  $Z_x(W)$  has at most  $\tilde{q}$  unique rows. Denoting  $\mathcal{A}$  as in (4), for some  $\bar{w} = W$ , it can be shown that
$$
\Omega_ {x} (g) = \| \mathcal {A} ^ {*} g \| ^ {2} = \| Y _ {\pi (x)} ^ {*} g \| ^ {2} = \sum_ {i = 1} ^ {r} n _ {i} ^ {2} (x) \cdot \| \tilde {g} _ {i} \| ^ {2},
$$
where  $\tilde{g}_i$  denotes any row  $j$  of  $g$  where  $\pi_j(x) = i$  (note that all rows of  $g$  that satisfy this condition are the same). Hence,  $\Omega_x(\cdot)$  can be efficiently represented by the nonzero values of  $\{n_k(x)\}_{k=1}^r$ . An efficient algorithm using  $O(\tilde{q}\log q)$  runtime and  $O(\tilde{q})$  extra storage, for computing  $\{n_x(x)\}_{k=1}^r$ , for a fixed  $x \in B$ , is given in Algorithm 3 of the Appendix. On the other hand, the evaluation of  $\Omega_x(\cdot)$  only requires evaluating the  $\tilde{q}$  unique rows of an input  $g \in \mathbb{R}^{q \times d}$  (given by  $\pi(x)$ ) and multiplying
Hint: the cost of storing the gradients for a single batch is equal to the cost of storing the weights.
them with the  $\tilde{q}$  nonzero cached values  $\{n_k(x)\}_{k=1}^r$ , which can be done with a  $\Theta(|B| \tilde{q}d)$  runtime cost for a batch  $B$ . Consequently, the total runtime (resp. storage) cost for a batch of examples  $B$  is  $\Theta(|B|\{q \log \tilde{q} + \tilde{q}d\})$  (resp.  $\Theta(|B| \tilde{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  $\nabla_W h(x, W)$  for each  $x \in B$ , it is straightforward to show that the compute (resp. storage) costs are  $\Theta(|B|qd)$  (resp.  $\Theta(|B|qd)$ ).
In the classic ghost clipping approach, we treat  $\phi_x(\cdot)$  as a linear operator, compute  $Y_{\pi (x)}Y_{\pi (x)}^*$ , and choose the representation  $\Omega_{x}(g) = \langle Y_{\pi (x)}Y_{\pi (x)}^{*},gg^{*}\rangle$ . In the worst-case scenario, it is straightforward to see that  $Y_{\pi (x)}Y_{\pi (x)}^{*}$  can be fully dense, e.g.,  $\tilde{q} = 1$ . In this setting, even if  $Y_{\pi (x)}$  is sparsely represented by a small set of  $\tilde{q}$  unique indices, computing  $Y_{\pi (x)}Y_{\pi (x)}^{*}$  still incurs a compute and storage cost of  $O(q^{2})$ . Now, since each example gradient  $g$  consists of  $q$  embedding vectors in  $\mathbb{R}^d$ , for a batch  $B$ , the compute (resp. storage) cost of materializing  $gg^{*}$  is  $\Theta(|B|q^{2}d)$  (resp.  $\Theta(|B|q^{2})$ ). Combining the above complexities with the  $\Theta(|B|q^{2})$  runtime cost of computing the desired inner products, i.e. (7) with  $U_{x} = Y_{\pi (x)}$  for  $x\in B$ , we obtain the complexities in Tables 1-2.
5.3 Low Rank Approximation Layer
Given input matrix  $U_x \in \mathbb{R}^{n \times n}$ , one way [19, 20] to encourage a rank- $k$  (or lower) approximation of  $U_x$  is to add the intermediate layer transform
$$
\phi_ {x} (V) = \left\| U _ {x} - V V ^ {*} \right\| ^ {2} + \rho \circ \sigma \left(V V ^ {*}\right)
$$
for some  $V \in \mathbb{R}^{n \times k}$ , where  $\rho(\cdot)$  is a sparsity promoting regularizer (e.g.,  $\ell_1$  norm, SCAD, MCP) and  $\sigma(\cdot)$  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
$$
\phi_ {x} (V) = - 2 \left\langle U _ {x}, V V ^ {*} \right\rangle + \mathcal {R} (V), \tag {8}
$$
for some function  $\mathcal{R}:\mathbb{R}^{n\times k}\mapsto \mathbb{R}$ . Note that  $\mathcal{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  $\Omega_x: \mathbb{R} \mapsto \mathbb{R}$  generated by the choice of  $Z_x(V) = \langle U_x, VV^* \rangle / 2$ , where, clearly, one has  $\phi_x(V) = -4Z_x(V) + \mathcal{R}(V)$ . Denoting  $\mathcal{A}$  as in (4), for some  $\bar{w} = V$ , it can be shown that
$$
\Omega_ {x} (g) = \| \mathcal {A} ^ {*} g \| ^ {2} = \frac {g ^ {2}}{4} \| (U _ {x} + U _ {x} ^ {*}) V \| ^ {2},
$$
for  $g \in \mathbb{R}$ . Hence,  $\Omega_x(\cdot)$  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  $\Theta(n^2k)$  runtime cost and a  $\Theta(1)$  storage cost. Moreover, evaluation of  $\Omega_x(\cdot)$ , given  $\| (U_x + U_x^*)V\|$ , requires only a  $O(1)$  runtime cost.
In the naive approach of computing  $\nabla_V h(x, V)$  for each  $x \in B$ , it is straightforward to see that the compute (resp. storage) costs is  $\Theta(|B| n^2 k)$  (resp.  $\Theta(|B| nk)$ ) due to the excessive computation (resp. storage) of  $g(U_x + U_x^*) V$  for  $x \in B$ . The authors are not aware of any ghost clipping-like techniques in the nonlinear setting.
This section presents numerical experiments that compare our proposed adjoint-based framework (Adjoint) against the naive 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.
Hint: we need to store (and take the norm of)  $q$  embedding vectors in  $\mathbb{R}^d$  for each example in the batch  $B$ .
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,  $\ell_x(\cdot)$ , 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  $\mathcal{Q}$  is the (linear) broadcasting operator that duplicates the bias  $r q / 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.
Figure 1: Runtime and memory cost graphs for fully-connected layer computations with bias dimensions  $m = \{2^1, 2^2, \ldots, 2^{11}\}$  and output dimensions  $q = 3, 4, 5$ .
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 (ii) 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  $\{\pi_i(x)\}_{i=1}^q$  are chosen uniformly at randomly from the set  $\{1, \ldots, r\}$ .
Figure 2: Runtime and memory cost graphs for embedding computations with query sizes  $q = \{1000, 2000, \dots, 10000\}$  and vocabulary sizes  $r = 5000, 7500, 10000$ .
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.
Figure 3: Runtime and memory cost graphs for training a small BERT model with batch sizes  $|B| = \{100,200,400,800,1600\}$ .
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.
The analysis in Subsection 5.1 may also be applied to more complex layers whose parameter transformations  $\phi_x(\cdot)$  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 nD-) 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.
The authors acknowledge the generous support from Walid Krichene (Google) and Li Zhang (Microsoft), who have meticulously reviewed the derivation and implementation of the proposed framework.
[1] 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.
[2] Rohan Anil, Badih Ghazi, Vineet Gupta, Ravi Kumar, and Pasin Manurangsi. Large-scale differentially private bert. arXiv preprint arXiv:2108.01624, 2021.
[3] 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.
[4] 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.
[5] 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.
[6] 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 3192-3218. PMLR, 2023.
[7] 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.
[8] 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.
[9] 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.
[10] 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.
[11] 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.
[12] Ian Goodfellow. Efficient per-example gradient computations. arXiv preprint arXiv:1510.01799, 2015.
[13] Jaewoo Lee and Daniel Kifer. Scaling up differentially private deep learning with fast perexample gradient clipping. Proceedings on Privacy Enhancing Technologies, 2021(1), 2021.
[14] Xuechen Li, Florian Tramer, Percy Liang, and Tatsunori Hashimoto. Large language models can be strong differentially private learners. arXiv preprint arXiv:2110.05679, 2021.
[15] Gaspar Rochette, Andre Manoel, and Eric W Tramel. Efficient per-example gradient computations in convolutional neural networks. arXiv preprint arXiv:1912.06015, 2019.
[16] Walter Rudin et al. Principles of mathematical analysis, volume 3. McGraw-hill New York, 1976.
[17] 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.
[18] 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.
[19] 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.
[20] Quanming Yao and James Kwok. Efficient learning with a family of nonconvex regularizers by redistributing nonconvexity. In International Conference on Machine Learning, pages 2645-2654. PMLR, 2016.
[21] 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.