Banded Square Root Matrix Factorization for Differentially Private Model Training

Thank you for the encouraging feedback. In the following we hope to address all your concerns.

> Using CVXPY with SCS to compute AOF is not efficient.

There are, of course, many ways to solve AOF (4). We used SCS, because it is a dedicated SPD-solver which allows for a simple and reproducible implementation without additional hyperparameters (Algorithm 2). General-purpose gradient-based solvers, such as L-BFGS, were meant to be included in our comment about approximate solvers (lines 122-125). These can be faster for solving (4) up to a certain precision, but they require additional design choices, such as the number of update steps, learning rates, a barrier function to enforce the constraints, and potentially projections to ensure positive definiteness. Importantly, however, note that our claim in the work is not that “BSR is x-times faster than AOF”. Actual runtimes always depend on many factors, including hardware, programming language, software packages, etc. We mention runtimes rather as a qualitative guidance to the reader, which is why we do not report a formal table or plot of runtimes. Our core argument is that (4) is a challenging optimization problem to solve. Any numeric solver will be slower and less convenient to use than our $O(n \log n)$ closed-form expressions. In real-world problems solving (4) creates a bottleneck, especially because it has to be solved anew for any change of hyperparameters. BSR overcomes this bottleneck.

> References to larger-scale matrix factorizations.

Thank you for the references; we will include them into our manuscript. Ref [1] indeed aims to solve the optimization problem (4). Could you please provide a reference for the timing results you mentioned? We did not spot them in the repository or associated report arXiv:2404.06430, but would be happy to adapt our presentation accordingly. [2] solves a simpler problem (from [McKenna et al., VLDB 2018]). Although adapting the formulation might be possible, the additional constraints could complicate the numerical solution. Nevertheless, we will be happy to include this in our discussion.

> The current experiment only demonstrates results for matrix sizes up to 2000, which is insufficient

We report experiments to a size of 2000 due to the computational limits of calculating AOF. For BSR, much larger workload sizes are no problem, and we will be happy to add them to the manuscript. The speed advantage of BSR over AOF will increase with $n$, regardless of the solver. We illustrate this in the response PDF: a plain python implementation of Theorem 1 (Alg.3 in PDF) computes the BSR-coefficients for $n=10,000$ in 24ms , for $n=100,000$ in 2.4s and for $n=500,000$ in ~67s (Table 1 in PDF, single-core CPU, double precision). In practice, one would likely not use vanilla MF-SGD for $n \ge 500,000$, because the resulting matrices (e.g. $C^{-1}$) would be inconveniently large (1 TB of memory in dense fp32 form).

> Enforcing positive semi-definiteness (PSD).

We did not include this step in Algorithm 2, because we considered it a post-processing operation only needed when trying to extract $C$ from the solution matrix $S$ (which despite the use of SCS, is not always perfectly PSD). We will be happy to add it, though. The actual operation is a standard procedure (Alg.4 in the response PDF). Our experiments use $tol=1/\sqrt{n}$ as a heuristic, because we observed it to give good results across a wide range of problem sizes, in particular $C^{-1}$ remaining stable. Constant values for tol or values proportional to $1/n$ did not work as well. We did not find this issue discussed in the literature.

> Equation (6).

Thanks for spotting this typo. We will fix it.

> Different learning rates.

Indeed, different learning rates would result in a non-Toeplitz workload matrix, so our analytical expressions no longer apply. However, the banded-square-root method remains applicable, only that the matrix square root has to be computed numerically. This step remains quite efficient because the workload matrix is still lower triangular, for which specialized routines exist, e.g. [Björck& Hammarling. “A Schur method for the square root of a matrix”, 1983]. We will include a discussion of the impact of varying learning rates in the revised version.

> General introduction to differential privacy.

Indeed, due to space constraints, we were not able to give an introduction in the manuscript, which primarily targets an audience already familiar with DP. We will add a section in the appendix.

Thank you for the valuable feedback. Below we clarify the raised points:

> Implementation of AOF seems suboptimal.

(note that Reviewer VxsQ had a similar question, and for easier reading we provide our answer in both our replies)

There are, of course, many ways to solve AOF (4). We used SCS, because it is a dedicated SPD-solver that allows a simple and reproducible implementation without additional hyperparameters (Algorithm 2). General-purpose gradient-based solvers, such as L-BFGS, were meant by our comment about approximate solvers (lines 122-125). These can be faster for solving (4) up to a certain precision, but they require additional design choices, such as the learning rates, number of update steps, a barrier function to enforce the constraints, and potentially projections to ensure positive definiteness.

> References on larger-scale matrix factorizations.

Ref arXiv:2306.08153 uses an in-house implementation by Google that is not freely available. The manuscript does not state runtimes or hardware. However, the fact that they also only report experiments up to $n=2,052$, we take as confirmation that solving (4) for large $n$ is indeed challenging. Regarding the KDD2016 reference, we do not believe it is comparable. It solves a different optimization problem (in fact, AOF’s formulation (4) was just proposed in 2023), which is also an SDP in the context of DP but with a simpler constraint set. Nevertheless, we’ll be happy to include it in our discussion.

Importantly, however, note that our claim in the work is not that “BSR is x-times faster than AOF”. Actual runtimes always depend on many factors, including hardware, programming language, software packages, etc. We mention runtimes rather as a qualitative guidance to the reader, which is why we do not report a formal table or plot of runtimes. Our core argument is that solving (4) is a challenging optimization problem. Any numeric solver will be slower and less convenient to use than our $O(n \log n)$ closed-form expressions. In real-world problems solving (4) creates a bottleneck, especially because it has to be solved anew for any change of hyperparameters. BSR overcomes this bottleneck.

> Scalability of BSR.

We reported experiments to a size of 2000 only due to the computational limits of AOF. For BSR, much larger workload sizes are no problem. We will be happy to add them to the manuscript. The speed advantage of BSR over AOF will increase with $n$, regardless of the solver. We illustrate this in the response PDF: a plain python implementation of Theorem 1 (Alg.3 in PDF) computes the BSR-coefficients for $n=10,000$ in 24ms , for $n=100,000$ in 2.4s and for $n=500,000$ in ~66s (CPU, single-threaded). If a dense matrix is wanted for C, times are approximately 3x higher due to the memory access overhead (Table 1 in PDF).

> Scale of $n$ for modern DP + ML applications?

In this context, $n$ is simply the total number of model updates, so large values of $n$ are quite common for standard network training. E.g., training MNIST (60,000 examples) with batchsize 128 for 10 epochs requires $n=(60000/128)*10= 4680$. A standard ResNet-Training on ImageNet1K (1.2M examples, batchsize 512, 90 epochs) yields $n=225,180$. Llama2-like LLM training (2T tokens, batchsize 4M tokens, 1 epoch) would result in $n=500,000$. However, in practice one would not use vanilla MF-SGD for such sizes, because the resulting $C^{-1}$ matrix would require 1 TB of memory to store (in dense form). It remains an open research question how to scale MF-SGD to such sizes. Our contribution only covers a part of that, namely avoiding the need to solve the preparatory (4).

> Theorem 2

We use this theorem to compute b-min-sep sensitivity for all cases, including Theorems 6 and 7. The idea of using Theorem 2 to find a better (potentially non-banded) factorization is indeed interesting. However, implementing this is not straightforward, because Theorem 2 requires the matrix $C$ to have decreasing coefficients. It is unclear how one could impose this property during the optimization, where one only has access to $S = C^TC$.

> Is the factorization time the dominating cost?

For AOF: yes. The factorization (4) is by far the most costly step, regardless of the solver. The Cholesky-type decomposition (line 118) is also costly but faster (seconds for $n=10000$, minutes for $n=50000$). BSR avoids both steps. Instead, one computes $C$ explicitly, which is efficient due to Theorem 1 (see above), and computing C’s sensitivity is trivial due to Theorem 2. The speed of AOF and BSR are both functions of $n$ and the bandwidth. The runtime of numeric solvers for AOF will also depend on the desired precision and the entries of the workload matrix (which determine the hardness of the optimization problem).

> Notation and typos.

Thank you for the suggestions, we will include them.

> It is remarkable that the approximation quality of BSR and AOF is so close

We found BSR to be consistently close to AOF across all settings we tested, not just the one in the manuscript. We’ll be happy to extend the manuscript with more results. Regarding bandwidth: Fig.3 in the response PDF shows that for the recommended p=b, BSR is typically close to AOF even for extreme values of $b$ ($b=2$ or $b=n$). For $p\ll b$ or $p\gg b$, indeed the distance to (b-banded) AOF grows. Some numeric results for $p=n$ can already be found in Appendix C in column “sqrt” (Tables 10-17 for $b=n$, Tables 2-9 for $b=100$). Unfortunately, their captions were mislabeled, which we will fix.

> Figures: the baselines were meant as a reference to help judge the significance of the differences between the methods. However, we could put the plots without baselines in the appendix, or if preferred, we could put the non-banded “sqrt” instead. Please let us know what you would prefer. We hope that also the numeric tables in the appendix allow judging the differences.

Thank you for the insightful review. Here, we address the individual concerns:

> Relation to arXiv:2404.16706.

Indeed, this interesting preprint (which we cite as [Dvijotham et al., 2024]) is concurrent work. It studies only the case of MF-SGD without momentum or weight decay (“prefix sum”) and is limited to streaming data ($k=1$, “single epoch”). This setting enables a deeper analytic treatment (in particular the sensitivity computation, which e.g. does not benefit from bandedness), but it comes at the cost of covering fewer real-world scenarios.

> While the previous work did require solving an SDP, to my understanding this is for the case where can be is only required to be a banded lower-triangular matrix. If is required to be Toeplitz as well, as the banded square root factorization is, then one only needs to optimize over variables $b$ instead $nb$ [...]

This is an interesting observation that could lead to a hybrid between [Choquette-Choo et al., 2023a] and our work. As a compromise between speed and optimality, one could keep the objective of the problem (4) but restrict it to Toeplitz matrices, thereby parametrizing the problem with $b$ variables. Unfortunately, the objective function would be quite more involved. It would no longer have a standard structure, such as the “matrix fractional function” form of the SDP (4), and it would not be convex in its variables (AOF’s (4) is only convex in the product matrix $S=C^T C$, which enters the objective via $S^{-1}$). Gradient-based optimizers could converge quickly but one would lose the guarantee of being close to a global optimum within a specifiable tolerance.

> Did the authors consider this alternate approach / if so, do you believe banded square root still offers speedups in this setting?

We had not considered this approach, but it sounds like promising future work that we would be happy to follow up on (with the permission of the reviewer). That said, we believe it would be unfair to consider it a shortcoming of our work that an alternative path exists which, to our knowledge, has not appeared in any prior work and will come with its own challenges.

Regarding the speed: even if the above idea might accelerate the numerical optimization, it will not surpass the simplicity and efficiency of BSR’s closed-form expressions. Furthermore, even if solved optimally, the quality of such a solution should lie between BSR and AOF, which are close together already.

> Privacy amplification.

Another interesting point, thank you for bringing it up. Since the BSR matrix has the same band structure as the AOF matrix, the privacy properties amplification arguments from Choquette-Choo et al. will still be applicable to our results. We will add a note to this in the manuscript. We did not include this aspect in the original submission because the material was already quite dense and we considered it somewhat orthogonal to our main analysis. But indeed, it would be interesting future work.

Thank you for the encouraging feedback and detailed comments.

> Experimental Evaluation.

We believe there may be a slight misunderstanding. Our main contribution lies in the algorithm and properties of BSR. Specifically, we demonstrate that in the context of MF-SGD, using BSR requires adding less noise compared to the baselines (for identical levels of privacy) and it is essentially on par with the computationally much more costly AOF (e.g., Figure 1). This fact is independent of any subsequent model training steps and requires no data, not even synthetic. We do not introduce a new technique for differential private model training, which would indeed require more extensive experimental validation and comparison to other private training techniques.

The MF-SGD experiments on CIFAR that we report (Figure 2) are only meant as a proof-of-concept to illustrate that, by keeping all other components fixed, the lower added noise tends to translate into higher accuracy. The exact experimental setup we will be happy to clarify in the manuscript: like prior work, we split the dataset uniformly at random into batches and train in a centralized way.

> Limitations of BSR.

We will be happy to extend our discussion in Section 5 to provide a broader view. In particular, we observe that BSR achieves results comparable to AOF, although we currently cannot prove this, as the theoretical properties of AOF are not well understood enough. Regarding real-world usability, the analytic expression for BSR does not apply to variable learning rates, which would be desirable.

> Post-processing step (line 275)

Note that this step is unrelated to BSR; it was only needed for AOF. We believe the problem is both numerical and principled. By this, we mean that, on the one hand, the problem is indeed numerical: if we could solve (4) with infinite mathematical precision, post-processing would not be needed. On the other hand, the numerical issues arise because of a principled drawback, namely that the AOF procedure (solving (4), computing a Cholesky decomposition, inverting the resulting matrix) is quite sensitive to small numeric deviations. Consequently, errors from solving (4) only to a certain precision, as all numerical solvers do, can lead to disproportionately large undesirable effects.

> Comparison to non-private approaches.

We will be happy to include such information in the manuscript. For example, non-private training in the setting of Figure 2 achieves an accuracy of approximately 65%. However, we would like to emphasize that our contribution is orthogonal to the contrast between private and non-private training. Our submission shows that by using a different processing step (BSR), one benefits in terms of speed (over AOF) and accuracy (over the baselines). Outside the context of MF-SGD, such a comparison cannot be made.

> Notation ($\Pi$), clarity of proof and sketches.

Thank you for bringing this to our attention. We will clarify the proof and sketches and we will remind the reader of the definition of $\Pi$.

> Minor issues.

Thank you, we will fix these items.