Banded Square Root Matrix Factorization for Differentially Private Model Training

Summary

The paper proposes the Banded Square Root (BSR) matrix factorization to speed up banded matrix factorization. The authors demonstrate that the workload matrix for SGD with momentum and weight decay can be expressed as a lower triangular Toeplitz matrix. They utilize explicit recursive expressions to compute the square root of this workload matrix, and Theorem 1 provides a closed-form expression for the square root matrix. Theoretical analysis shows that the calculation of the square root matrix is efficient, and the sensitivity of the mechanism is discussed. Detailed analysis of the approximation error is also provided, demonstrating that the square root factorization has asymptotically optimal approximation quality. Experiments show that BSR can achieve performance comparable to previous state-of-the-art methods.

Strengths

1. The discussion regarding SGD with momentum and weight decay is interesting. The workload matrix in this context is a lower triangular Toeplitz matrix, which allows for efficient factorization algorithms. Prior work has not explored this setting; they typically treat the workload matrix with momentum as a general lower triangular matrix. This finding is significant for future research, suggesting potential improvements over current solvers for banded matrix factorization. 2. The theoretical analysis demonstrates that Banded Square Root (BSR) matrix factorization can be computed efficiently even for large problem sizes. The expected approximation error for BSR indicates that it achieves asymptotically optimal approximation quality. These analyses ensure that BSR is both fast and accurate.

Weaknesses

1. Using CVXPY with SCS to compute AOF is not efficient; L-BFGS can solve this problem more efficiently. For instance, as shown in [1], an implementation using GPUs can solve a problem size of 10,000 in less than an hour. Even with CPUs, the implementation in [2] can solve a problem size of 2,000 in less than an hour, with slight modifications to add the b-min-step constraint. 2. The current experiment only demonstrates results for matrix sizes up to 2000, which is insufficient. A comparison involving matrix sizes of 10,000 or more would be more convincing. How long would it take to obtain a solution for problem sizes of 10,000 and 100,000 using BSR? [1] https://github.com/apple/pfl-research/blob/2a9d65bd66dc89ef80a24e1be0d28446a0422133/pfl/privacy/ftrl_mechanism.py#L140 [2] https://github.com/dpcomp-org/hdmm/blob/7a5079a7d4f1a06b0be78019adadf83c538d0514/src/hdmm/templates.py#L391

Questions

1. Enforcing positive semi-definiteness (PSD) for the matrix is crucial for convergence in AOF. In line 274, you mentioned a post-processing step for S to ensure that all its eigenvalues are at least $\sqrt{1/n}$. Could you provide a more detailed description of this step? I couldn't find it in the source code provided in Appendix A. 2. In equation (6), b should be p, right?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

1. If the learning rate is not constant, then the workload matrix will not be a Toeplitz matrix, right? This would significantly impact the efficiency of BSR. During the training process, a proper learning rate schedule can enhance convergence. Therefore, a discussion about varying learning rates should be included. 2. The paper does not provide a general introduction to differential privacy, and Figure 2 presents (epsilon, delta) without sufficient background information. This lack of context makes the paper not self-contained.

Summary

This paper proposes the Banded Square Root Matrix Factorization, an efficient approximation of the optimal banded matrix factorization for differentially private machine learning applications. The authors give a closed form expression for the BSR C matrix for any SGD + Momentum + Weight Decay workload by exploiting the Toepltiz structure of the workload and factors. These factorizations can be used in place of the prior work and enjoy some efficiency advantages.

Strengths

* The scope, contributions, and key results are clearly written. * Near-optimal and efficiency computable banded matrix factorizations could make it feasible to use this mechanism in some new settings where it was previously not possible. * Authors demonstrated deep understanding of the area, and made a compelling story with a strong mix of theoretical derivations and empirical observations.

Weaknesses

* The implementation of AOF seems suboptimal and hence the comparison appears to be biased. Multiple-day runtimes for n ~ 700 seems high when prior work conducts experiments for n >= 2000 (https://arxiv.org/pdf/2306.08153). Doing a little digging, I found https://www.kdd.org/kdd2016/papers/files/rpp0224-yuanA.pdf, which ran experiments up to n ~ 8192. The limits of the BSR approach are not really demonstrated besides the statement "Even sizes of n = 10,000 or more take at most a few seconds." How much beyond this can you scale? Would be good to add a scalability experiment.

Questions

* Assuming your method scales well beyond the prior work of n ~ 8192, can you make a convincing case that in modern DP + ML applications we need to scale beyond that point? I'm not super familiar with how many iterations DP-SGD is usually run for. * Thm 2 is an interesting result that allows you to compute sensitivity efficiently under this special Toepltiz structure without requiring bandedness, but it's not clear where you are using this Thm. Is it better to find a banded factorization with bands = minsep or use a non-banded factorization with Thm 2? >* Also minor point: might be better to write Eq 10 in terms of m_0, ... m_{n-1} instead of M (or maybe include both expressions). >* There is an incorrect reference to Eq 10 in the experiments >* Typo: n_{n-1} * Regarding the statement: "Apart from the factorization itself, the computational cost of BSR and AOF are nearly identical." Is the factorization time usually the dominating term in the cost for AOF? Is that still the case for BSR, even when scaling up well beyond n ~ 8192? What factor(s) other than n (if any) impact the complexity of these mechanisms? * It is remarkable that the approximation quality of BSR and AOF is so close. It looks like you evaluated a few settings, one of them being b=100, k=n/100 in Fig 1. Does this hold more generally across other settings, and can you demonstrate that with an experiment? There are two extremes where it I might expect *some* gap: (a) 2 bands (b) n bands. >* In Fig 1, the most interesting comparison to me is BSR vs. AOF, the baselines have already been established to be far from optimal in prior work and their presence makes it difficult to perceive difference between the two most interesting lines, would be good to update the plot accordingly.

Rating

7

Confidence

5

Soundness

4

Presentation

4

Contribution

3

Limitations

The authors have discussed some limitations, but perhaps adding some of the things I mentioned in this review (or addressing them directly) would be good.

Summary

This paper addresses optimal matrix factorization mechanism for differential privacy, focussing on stochastic gradient descent (SGD) optimization. It introduces the banded squared root (BSR) factorization and provide formal lower and upper bounds to the approximation error. The BSR achieves competitive formal approximation error compared to state-of-the-art methods and demonstrates practical utility in real-world training tasks. The authors provide theoretical bounds for single and repeated participation training, showing that BSR maintains high training accuracy with provable privacy guarantees while being computationally efficient.

Strengths

- The authors make a significant contribution to the field of differentially private machine learning under matrix factorization mechanisms by leveraging efficient BSR for the SGD workload matrix. - The incorporation of the SDG workload matrix in this line of research is elegant. - The banded squared root (BSR) factorization is computationally efficient and scalable to large workload matrices, making it a valuable tool. - The proof sketches are intuitive and convincing, although I haven't thoroughly reviewed the technical proofs in the appendix. - A notable advantage of BSR is its independence from specific training objectives. - BSR achieves competitive approximation error, on par with state-of-the-art methods in differentially private training. - The approach demonstrates practical utility by maintaining high training accuracy in real-world training tasks. - The authors provide theoretical guarantees for both single and repeated participation scenarios, adding to the method's credibility. - The theoretical explanations are (for the most part) well-written and easy to follow. - The paper includes extensive technical supplementary material. - The inclusion of useful, copy-pasteable Python code listings is interesting. - The authors show respect for existing work by referencing relevant StackOverflow answers and software packages.

Weaknesses

- The experimental evaluation is limited to a small set of real-world and synthetic experiments. - The range of synthetic datasets used in the experiments is restricted, which may not be representative of diverse scenarios. - The data generation process is unclear, and only homogeneous partitions are considered, with no variation in data distributions. - The real-world experiments are limited to a single dataset, CIFAR-10, which is insufficient to draw general conclusions. - The post-processing step (line 275) lacks clear explanation, leaving it unclear whether it addresses a principled problem or a numerical issue. - The experiments do not provide a comparison with non-private approaches, which makes it difficult to understand the practical limitations of the method. - The evaluation is limited to a single real-world dataset, CIFAR-10, and a simple ConvNet model. - The experiments do not discuss potential practical limitations of BSR, which may impact its applicability. - The comparative analysis is limited and does not include a wide range of existing differential privacy methods. - The formal proof of Theorem 2 (Appendix D.2) employs unclear and undefined notation (Π), which makes it difficult to follow. - Neither the proof sketch nor the formal proof are straightforwardly understandable, which may hinder comprehension. Minor: - The tone of "supposedly optimal" (line 254) could be perceived as inappropriate and may benefit from rephrasing. - The proof sketch of Theorem 2 refers to the wrong appendix (Appendix D.3 instead of Appendix D.2, presumably).

Questions

./.

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

./.

Summary

The paper considers the problem of adding correlated noise $C^{-1}z$ instead of independent noise $z$ across iterations in continual counting (equivalently, DP-SGD). Past work gave an algorithm to choose $C^{-1}$ that optimizes some objective on the noise under b-min-separated participation, but this algorithm requires solving an SDP which is infeasible when using a large number of iterations, i.e. optimizing over $C$ that has a large number of rows / columns $n$. The authors propose a choice of $C$, which they call banded square root, which has a succinct representation that is efficiently computable in time independent of $n$. Furthermore, their choice of $C$ is banded, i.e. is non-zero below the $b$-th diagonal for $b \geq 1$, which gives it some nice properties such as being able to compute $C^{-1}z$ efficiently in a streaming manner. To define the banded square root $C$, they authors take $A$ which represents the workload, effectively a representation of the updates in SGD which can include momentum and weight decay, and take its matrix square root. They then truncate the matrix square root to its first $b$ diagonals. The authors give an explicitly and efficiently calculable formula for the computing the first column of this matrix, which specifies the whole matrix as it is Toeplitz. The authors analyze the asymptotic error guarantees of different choices of $C$, including the banded square root, in the single- and multiple-participation settings. Specifically, their error guarantees show that banded square root improves on $C = I$ or $C = A$, and nearly matches a lower bound on any factorization if using enough bands. The authors conduct experiments training models for classification on CIFAR10 and show that banded square root is competitive with the choice of $C$ given by solving a more expensive SDP.

Strengths

* The work makes theoretical progress on a practical problem. Correlated noise/DP-MF is now being used in practice to train models with DP, and especially as model sizes and training runs get larger, more efficient ways to implement DP-MF will lead to better models in practice. In particular, itmakes the results of the DP-MF literature more accessible to those who do not have large amounts of compute to use the techniques they introduced. * The theoretical analysis of the error of DP-MF as a function of the number of bands is novel and adds theoretical understanding of banded $C$ that was not present in the previous paper of Choquette-Choo et al. Furthermore, the set of theoretical results is rather extensive and gives a pretty complete theoretical understanding of the authors' methods. * While DP-MF is a relatively niche topic, the paper does a good job slowly introducing the problem, their approach, and their theoretical results.

Weaknesses

* While the previous work did require solving an SDP, to my understanding this is for the case where $C$ can be is only required to be a banded lower-triangular matrix. If $C$ is required to be Toeplitz as well, as the banded square root factorization is, then one only needs to optimize over $b$ variables instead of $\approx nb$ and it is not clear that the computations in the previous work are expensive, which mitigates the improvements in this paper. * It is worth pointing out there is a work of https://arxiv.org/abs/2404.16706, that also gives an efficient-to-compute way to choose $C$, and leads to faster streaming noise generation for than banded matrices. This work I would consider concurrent, so I did not account for any possible overlap between the two papers' results and impact when assigning my score.

Questions

* wrt the first weakness, did the authors consider this alternate approach / if so, do you believe banded square root still offers speedups in this setting? * The previous work of Choquette-Choo et al. shows that banded $C$ are compatible with privacy amplification with sampling. Have you considered combining your error analysis with their privacy analysis (e.g., maybe for RDP to simplify) in the setting where batches are sampled in DP-SGD? It might be an interesting question to see how the number of bands affects error in the presence of amplification, since they observed more bands reduce the benefits of amplification.

Rating

6

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

Yes