Title: Banded Square Root Matrix Factorization for Differentially Private Model Training
Venue: NeurIPS 2024
Accepted: yes
Keywords: Differential Privacy

Machine Learning

Federated Learning
Abstract: Current state-of-the-art methods for differentially private model training are based on matrix factorization techniques. However, these methods suffer from high computational overhead because they require numerically solving a demanding optimization problem to determine an approximately optimal factorization prior to the actual model training. In this work, we present a new matrix factorization approach, BSR, which overcomes this computational bottleneck. By exploiting properties of the standard matrix square root, BSR allows to efficiently handle also large-scale problems. For the key scenario of stochastic gradient descent with momentum and weight decay, we even derive analytical expressions for BSR that render the computational overhead negligible. We prove bounds on the approximation quality that hold both in the centralized and in the federated learning setting. Our numerical experiments demonstrate that models trained using BSR perform on par with the best existing methods, while completely avoiding their computational overhead.

Banded Square Root Matrix Factorization for Differentially Private Model Training

Abstract

Current state-of-the-art methods for differentially private model training are based on matrix factorization techniques. However, these methods suffer from high computational overhead because they require numerically solving a demanding optimization problem to determine an approximately optimal factorization prior to the actual model training. In this work, we present a new matrix factorization approach, BSR, which overcomes this computational bottleneck. By exploiting properties of the standard matrix square root, BSR allows to efficiently handle also large-scale problems. For the key scenario of stochastic gradient descent with momentum and weight decay, we even derive analytical expressions for BSR that render the computational overhead negligible. We prove bounds on the approximation quality that hold both in the centralized and in the federated learning setting. Our numerical experiments demonstrate that models trained using BSR perform on par with the best existing methods, while completely avoiding their computational overhead.

1 Introduction

We study the problem of differentially private (DP) model training with stochastic gradient descent (SGD) in the setting of either federated or centralized learning. This task has recently emerged as one of the most promising ways to train powerful machine learning models but nevertheless guarantee the privacy of the used data, which led to a number of studies, both theoretical as well as application-driven [Abadi et al., 2016, Yu et al., 2020, Zhang et al., 2021, Kairouz et al., 2021, Denisov et al., 2022]. The state of the art in the field are approaches based on the matrix factorization (MF) mechanism [Li et al., 2015, Henzinger et al., 2024], which combines theoretical guarantees with practical applicability [Choquette-Choo et al., 2023a,c,b, 2024] 1 It is based on the observation that all iterates of SGD are simply linear combinations of model gradients, which are computed at intermediate time steps. Consequently, the iterates can be written formally as the result of multiplying the matrix of coefficients, called workload matrix , with the row-stacked gradient vectors. To preserve the privacy of the training data in this process one adds suitably scaled Gaussian noise at intermediate steps of the computation. The MF mechanism provides a way to select the noise covariance structure based on a factorization of the workload matrix into two matrices.

Identifying the minimal amount of noise necessary to achieve a desired privacy level requires solving an optimization problem over all possible factorizations, subject to data participation constraints . For some specific settings, the optimal solutions have been characterized: for streaming learning , when each data batch contributes at most once to the gradients, Li et al. [2015] presented a formulation of this problem as a semi-definite program. Henzinger et al. [2024] proved that a square root factorization

of the workload matrix is asymptotically optimal for different linear workloads, including continual summation and decaying sums.

In this work, our focus lies on the settings that are most relevant to machine learning tasks: workload matrices that reflect SGD-like optimization, and participation schemes in which each data batch can potentially contribute to more than one gradient vector, as it is the case for standard multi-epoch training. Assuming that this happens at most once every b steps, for some value b ≥ 1 , leads to the problem of optimal matrix factorization in the context of b -min-separated participation sensitivity. Unfortunately, as shown in Choquette-Choo et al. [2023a], finding the optimal matrix factorization in this setting is computationally intractable. Instead, the authors proposed an approximate solution by posing additional constraints on the solution set. The result is a semi-definite program that is tractable, but still has high computational cost, making it practical only for small to medium-sized problem settings.

Subsequent work concentrated on improving or better understanding the factorizations for specific algorithms, such as plain SGD without gradient clipping, momentum, or weight decay [Koloskova et al., 2023] or specific, e.g. convex, objective functions [Choquette-Choo et al., 2024]. Often, streaming data was asssumed, i.e. each data item can contribute at most to one model update, which is easier to analyze theoretical, but further removed from real-world applications [Dvijotham et al., 2024]. A concurrent line of works has also focused on scaling matrix factorizations for large-scale training. McMahan et al. [2024] extended the Buffered Linear Toeplitz (BLT) mechanism to support multi-participation in federated learning, improving privacy-utility tradeoffs, while ensuring memory efficiency. McKenna [2024] improved DP Banded Matrix Factorization's scalability, enabling it to handle millions of iterations and large models efficiently. These advancements enhance the practicality of differentially private training in real-world applications.

Our work aims at a general-purpose solution that covers as many realistic scenarios as possible. Our ultimate goal is to make general-purpose differentially private model training as simple and efficient to use as currently dominating non-private technique. Our main contributions are:

We introduce a new factorization, the banded squared root (BSR) , which is efficiently computable even for large workload matrices and agnostic to the underlying training objective. For matrices stemming from SGD optimization potentially with momentum and/or weight decay, we even provide closed form expressions.

We provide general lower bounds on the approximate error for any factorization, along with specific upper and lower bounds for the BSR , Square Root, and baseline factorizations, in the contexts of both single participation (streaming) and repeated participation (e.g., multi-epoch) training.

We demonstrate experimentally that BSR's approximation error is comparable to the state-of-the-art method, and that both methods also perform comparably in real-world training tasks .

Overall, the proposed BSR factorization achieves training high-accuracy models with provable privacy guarantees while staying computationally efficient even for large-scale training tasks.

2 Background

Our work falls into the areas of differentially private (stochastic) optimization , of which we remind the reader here, following mostly the description of Denisov et al. [2022] and Choquette-Choo et al. [2023a]. The goal is to estimate a sequence of parameter vectors, Θ = ( θ 1 , . . . , θ n ) ∈ R d , where each θ i is a linear combination of update vectors , x 1 , . . . , x i ∈ R d , that were computed in previous steps, typically as gradients of a model with respect to training data that is meant to stay private. We assume that all x i have a bounded norm, ‖ x i ‖ ≤ ζ . Compactly, we write Θ = AX , where the lower triangular workload matrix A contains the coefficients and X ∈ R n × d is formed by stacking the update vectors as rows. With different choices of A , the setting then reflects many popular first-order optimization algorithms, in particular stochastic gradient descent (SGD), potentially with momentum and/or weight decay. Depending on how exactly x 1 , . . . , x n are obtained, the setting can express different centralized as well as federated training paradigms. To formalize this aspect, we adopt the concept of b -min-separated participation [Choquette-Choo et al., 2023a]. For some integer b ≥ 1 , it states that if a data item (e.g. a single training example in central training, or a client batch in

Algorithm 1 Differentially Private SGD with Matrix Factorization

Input: Initial model θ 0 ∈ R d , dataset D , batchsize b , matrix C ∈ R n × n , model loss glyph[lscript] ( θ, d ) , clipnorm ζ , noise matrix Z ∈ R n × d with i.i.d. entries ∼ N (0 , s 2 ) , where s = σ sens k,b ( C ) .

for i = 1 , 2 , . . . , n do

S i ←{ d 1 , . . . , d m } ⊆ D select a data batch, respecting the data participation constraints g i ←∇ θ glyph[lscript] ( θ i -1 , d j )) for j = 1 , . . . , m

x

ˆ

x

m

i

i

∑

j

=1

x

i

+

ζ

clip

[

C

ζ

-

(

1

g

Z

j

]

)

[

i,

θ

i

update

(

θ

i

-

1

,

ˆ

x

i

Output: Θ = ( θ 1 , . . . , θ n )

federated learning) contributed to an update x i , the earliest it can contribute again is the update x i + b . Additionally, let 1 ≤ k ≤ n b be the maximal number any data point can contribute. In particular, this notion also allows us to treat in a unified way streaming data ( b = n or k = 1 ), as well as unrestricted access patterns ( k = n with b = 1 ), but also intermediate settings, such as multi-epoch training on a fixed-size dataset.

The matrix factorization approach [Li et al., 2015] adopts a factorization A = BC of the workload matrix and computes Θ MF = B ( CX + Z ) , where Z is Gaussian noise that is chosen appropriately to make the intermediate result CX + Z private to the desired level. Algorithm 1 shows the resulting algorithm in pseudocode. It exploits the fact that instead of explicit multiplication by C and B , standard optimization toolboxes can be employed with suitably modified update vectors, because also Θ MF = A ( X + C -1 Z ) , and multiplication by A corresponds to performing the optimization.

Different factorizations recover different algorithms from the literature. For example, B = A , C = Id recovers DP-SGD [Abadi et al., 2016], where noise is added directly to the gradients. Conversely, B = Id , C = A simply adds noise to each iterate of the optimization [Dwork et al., 2006]. However, better choices than these baselines are possible, in the sense that they can guarantee the same levels of privacy with less added noise, and therefore potentially with higher retained accuracy. The reason lies in the fact that B and C play different roles: B acts as a post-processing operation of already private data. Hence, it has no further effect on privacy, but it influences to what amount the added noise affects the expected error in the approximation of Θ . Specifically, for Z ∼ N (0; s Id) ,

In contrast, CX is the quantity that is meant to be made private. Doing so requires noise of a strength proportional to C 's sensitivity , sens( C ) := sup X ∼ X ′ ‖ CX -CX ′ ‖ F , where the neighborhood relation , X ∼ X ′ , indicates that the two sequences of update vectors differ only in those entries that correspond to a single data item 2 As shown in Choquette-Choo et al. [2023a], in the setting of b -min-separated repeated participation, it holds that

where Π k,b = { π ⊂ { 1 , . . . , n } : | π | ≤ k ∧ ( { i, j } ⊂ π ⇒ i = j ∨ | i -j | ≥ b ) } , is the set of possible b -min-separated index sets with at most k participation. Furthermore, (2) holds even with equality if all entries of C glyph[latticetop] C are non-negative.

Combining (1) with s = sens k,b ( C ) yields a quantitative measure for the quality of a factorization.

Definition 1. For any factorization A = BC , its expected approximation error is

where the 1 / √ n factor is meant to make the quantity comparable across different problem sizes.

·

)

,

where clip

ζ

(

d

) = min(1

, ζ/

// SGD model updates

||

||

d

)

d

]

The optimal factorization by this reasoning would be the one of smallest expected approximation error. Unfortunately, minimizing (3) across all factorizations it is generally computationally intractable. Instead, Choquette-Choo et al. [2023a] propose an approximately optimal factorization .

Definition 2. For a workload matrix A , let S be the solution to the optimization problem

where S n + is the cone of positive definite n × n matrices. Then, A = BC is called the approximately optimal factorization (AOF) , if C is lower triangular and fulfills C glyph[latticetop] C = S .

The optimization problem (4) is a semi-definite program (SDP), and can therefore be solved numerically using standard packages. However, this is computationally costly, and for large problems (e.g. n > 5000 ) computing the AOF solution is impractical. This poses a problem for real-world training tasks, where the number of update steps are commonly thousands or tens of thousands.

Solving (4) itself only approximately can mitigate this problem to some extent, but as we will discuss in Section 4, this can lead to robustness problems, especially because the recovery of C from S in Definition 2, e.g. by a Cholesky decomposition, tends to be sensitive to numerical errors.

3 Banded Square Root Factorization

In the following section, we introduce our main contribution: a general-purpose factorization for the task of differentially private stochastic optimization that can be computed efficiently even for large problem sizes.

Definition 3 (Banded Square Root Factorization) . Let A ∈ R n × n be a lower triangular workload matrix with strictly positive diagonal entries. Then, we call A = C 2 the square root factorization (SR) , when C denotes the unique matrix square root that also has strictly positive diagonal entries. Furthermore, for any bandwidth p ∈ { 1 , . . . , n } , we define the banded square root factorization of bandwidth p ( p -BSR) of A , as

where C | p | is created from C by setting all entries below the p -th diagonal to 0 ,

Note that determining the SR, and therefore any p -BSR, is generally efficient even for large workload matrices, because explicit recursive expressions exist for computing the square root of a lower triangular matrix [Björck and Hammarling, 1983, Deadman et al., 2012].

In the rest of this work, we focus on the case where the workload matrix stems from SGD with momentum and/or weight decay , and we show that then even closed form expressions for the entries of C | p | exist that renders the computational cost negligible.

3.1 Banded Square Root Factorization for SGD with Momentum and Weight Decay

We recall the update steps of SGD with momentum and weight decay:

where x 1 , . . . , x n are the update vectors, η > 0 is the learning rate and 0 ≤ β < 1 is the momentum strength and 0 < α ≤ 1 is the weight decay parameter. Note that our results also hold for β = 0 , i.e. without momentum, and for α = 1 , i.e., without weight decay. In line with real algorithms and to avoid degenerate cases, we assume β < α throughout this work. The update vectors are typically gradients of the model with respect to its parameters, but additional operations such as normalization or clipping might have been applied.

Unrolling the recursion, we obtain an expression for θ i as a linear combination of update vectors as

Consequently, the workload matrix has the explicit form A = ηA α,β for

As one can see, A α,β is a lower triangular Toeplitz-matrix, so it is completely determined by the entries of its first column. In the following, we use the notation LDToep( m 1 , . . . , m n ) to denote a lower triangular Toeplitz matrix with first column m 1 , . . . , m n , i.e. A α,β = LDToep( a 0 , . . . , a n -1 ) .

Our first result is an explicit expression for the positive square root of A α,β (and thereby its p -BSR). Theorem 1 (Square-Root of SGD Workload Matrix) . Let A α,β be the workload matrix (9) . Then A α,β = C 2 α,β for C α,β = LDToep( c 0 , . . . , c n -1 ) , with c 0 = 1 and c j = ∑ j i =0 α j -i r j -i r i β i for j = 1 , . . . , n -1 with coefficients r i = ∣ ∣ ( -1 / 2 i )∣ ∣ . For any p ∈ { 1 , . . . , n } , the p -banded BSR matrix C | p | α,β is obtained from this by setting all coefficients c j = 0 for j ≥ p .

Proof sketch. The proof be found in Appendix F.1. Its main idea is to factorize A α,β into a product of two simpler lower triangular matrices, each of which has a closed-form square root. We show that the two roots commute and that the matrix C α,β is their product, which implies the theorem.

3.2 Efficiency

We first establish that the p -BSR for SGD can be computed efficiently even for large problem sizes. Lemma 1 (Efficiency of BSR) . The entries of C | p | α,β can be determined in runtime O ( p log p ) , i.e., in particular independent of n .

Proof sketch. As a lower triangular Toeplitz matrix, C | p | α,β is fully determined by the values of its first column. By construction c p +1 , . . . , c n = 0 , so only the complexity of computing c 1 , . . . , c p -1 matters. These can be computed efficiently by writing them as the convolution of vectors ( α i r i ) i =0 ,...,p -1 and ( β i r i ) i =0 ,...,p -1 and, e.g., employing the fast Fourier transform.

Note that for running Algorithm 1, the matrix B of the factorization is not actually required. However, one needs to know the sensitivity of C | p | α,β , as this determines the necessary amount of noise. The following theorem establishes that for a large class of matrices, including the BSR in the SGD setting, this is possible exactly and in closed form.

Theorem 2 (Sensitivity for decreasing non-negative Toeplitz matrices) . Let M = LDToep( m 0 , . . . , m n -1 ) be a lower triangular Toeplitz matrix with decreasing non-negative entries, i.e.

m 0 ≥ m 1 ≥ m 2 ≥ . . . m n -1 ≥ 0 . Then its sensitivity (2) in the setting of b -min-separation is

where M [ · , 1+ jb ] denotes the (1 + jb ) -th column of M .

Proof sketch. The proof can be found in Appendix F.2. It builds on the identity (2), which holds with equality because of the non-negative entries of M . Using the fact that the entries of M are non-increasing one establishes that an optimal b -separated index set is { 1 , 1 + b, · · · , 1 + ( k -1) b } . From this, the identity (10) follows.

Corollary 1. The sensitivity of the p -BSR for SGD can be computed using formula (10) .

Proof sketch. It suffices to show that the coefficients c 0 , . . . , c n -1 of Theorem 1 are monotonically decreasing. We do so by an explicit computation, see Appendix F.3.

3.3 Approximation Quality - Single Participation

Having established the efficiency of BSR, we now demonstrate its suitability for high-quality model training. To avoid corner cases, we assume that n b is an integer, which does not affect the asymptotic behavior. We also discuss only the case in which the update vectors have bounded norm ζ = 1 . Results for general ζ can readily be derived using the linearity of the sensitivity with respect to ζ .

We first discuss the case of model training with single participation ( k = 1 ), where more precise results are possible than the general case. Our main result are bounds on the expected approximation error of the square root factorization that, in particular, prove its asymptotic optimality.

Theorem 3 (Expected approximation error with single participation) . Let A α,β ∈ R n × n be the workload matrix (9) of SGD with momentum 0 ≤ β < 1 and weight decay parameter 0 < α ≤ 1 , where α > β . Assume that each data item can contribute at most once to an update vector (e.g. single participation, k = 1 ). Then, the expected approximation error of the square root factorization , A α,β = C 2 α,β , fulfills

for α < 1 , and

Proof sketch. For the proof, we establish a relations between sens 1 ,n ( C ) and ‖ C α,β ‖ F , and then we bound the resulting expressions by an explicit analysis of the norm. For details, see Appendix F.5.

The following two results provide context for the interpretation of Theorem 3.

Theorem 4. Assume the setting of Theorem 3. Then, for any factorization A α,β = BC with C glyph[latticetop] C ≥ 0 , the expected approximation error fulfills

Proof sketch. The theorem is the special case k = 1 of Theorem 8, which we state in the next section and prove in Section F.9.

Theorem 5. Assume the setting of Theorem 3. Then, the baseline factorizations A α,β = A α,β · Id and A α,β = Id · A α,β fulfill, for α < 1 ,

Proof sketch. The result follows from an explicit analysis of the coefficients, see Appendix F.6.

Discussion. Theorems 3 to 5 provide a full characterization of the approximation quality of the square root factorization as well as its alternatives: 1) the square root factorization has asymptotically optimal approximation quality, because the upper bounds in Equation (12) match the lower bounds in Equation (13); 2) the AOF from Definition 2 also fulfills the conditions of Theorem 4. Therefore, it must also adhere to the lower bound (13) and cannot be asymptotically better than the square root factorization; 3) the approximation qualities of the baseline factorizations in Equation (14) and (15) are asymptotically worse than optimal in the α = 1 setting, and worse by a constant factor for α < 1 . The BSR factorization can be applied even in more general scenarios, such as with varying learning rates. However, in this case, the workload matrix will no longer be Toeplitz. This makes it difficult to provide analytical guarantees for the matrix, but it can still be applied numerically.

3.4 Approximation Quality - Repeated Participation.

We now provide mostly asymptotic statements about the approximation quality of BSR and baselines in the setting where data items can contribute more than once to the update vectors.

Theorem 6 (Approximation error of BSR) . Let A α,β ∈ R n × n be the workload matrix (9) of SGD with momentum 0 ≤ β < 1 and weight decay 0 < α ≤ 1 , with α > β . Let A α,β = B | p | α,β C | p | α,β , be its banded square root factorization as in Definition 3. Then, for any b ∈ { 1 , . . . , n } , p ≤ b , and k ∈ { 1 , . . . , n b } it holds:

Proof sketch. For the proof, we separately bound the sensitivity of C | p | α,β and the Frobenius norm of B | p | α,β . The former is straightforward because of the matrix's band structure. The latter requires an in-depth analysis of the inverse matrix' coefficient. Both steps are detailed in Appendix F.7.

The following results provide context for the interpretation of Theorem 6.

Theorem 7 (Approximation error of Square Root Factorization) . Let A α,β ∈ R n × n be the workload matrix (9) of SGD with momentum 0 ≤ β < 1 and weight decay 0 < α ≤ 1 , with α > β . Let A α,β = C 2 α,β be its square root factorization . Then, for any b ∈ { 1 , . . . , n } and k = n b it holds:

Proof sketch. We bound sens k,b ( C α,β ) and ‖ C α,β ‖ F using the explicit entries for C α,β from Theorem 1. Details are provided in Appendix F.8.

Theorem 8. Assume the setting of Theorem 6. Then, for any factorization A α,β = BC with C glyph[latticetop] C ≥ 0 , the approximation error fulfills

Proof sketch. The proof is based on the observation that ‖ X ‖ F ‖ Y ‖ F ≥ ‖ XY ‖ ∗ for any matrices X,Y , where ‖ · ‖ ∗ denotes the nuclear norm. To derive (18), we show that sens k,b ( C ) is lower bounded by √ k n ‖ C ‖ F , and derive explicit bounds on the singular values of A α,β .

Theorem 9. Assume the setting of Theorem 6. Then, the baseline factorizations A α,β = A α,β · Id and A α,β = Id · A α,β fulfill

Proof sketch. The proof relies on the fact that the workload matrices can be lower bounded componentwise by simpler matrices: A α,β ≥ A α, 0 and A α, 0 ≥ Id . For the simpler matrices, the bounds (19) can then be derived analytically, and the general case follows by monotonicity.

Discussion. Analogously to the case of single participation, Theorems 6 to 9 again establish that the proposed BSR is asymptotically superior to the baseline factorizations if α = 1 . A comparison of Theorems 6 and 7 suggests that, at least for maximal participation, k = n b and p = b , the bandedness of the p -BSR improves the approximation quality, specifically in the practically relevant regime where b glyph[lessmuch] n . While none of the methods match the lower bound of Theorem 6, we conjecture that this is not because any asymptotically better methods would exist, but rather a sign of Equation (18) is not tight. Both theoretical consideration and experiments suggest that a term linear in k should appear there. For α < 1 , all studied methods are asymptotically identical and, in fact, optimal.

4 Experiments

To demonstrate that BSR can achieve high accuracy not only in theory but also in practice, we compare it to AOF and baselines in numerical experiments. Our results show that BSR achieves quality comparable to the AOF, but without the computational overhead, and it clearly outperforms the baseline factorizations. The privacy guarantees are identical for all methods, so we do not discuss them explicitly.

Implementation and computational cost. We implement BSR by the closed-form expressions of Theorem (1). For single data participation, we use the square root decomposition directly. For repeated data participation we use p -BSR with p = b . Using standard python/numpy code, computing the BSR as dense matrices are memory-bound rather than compute-bound. Even sizes of n = 10 , 000 or more take at most a few seconds. Computing only the Toeplitz coefficients is even faster, of course.

To compute AOF, we solve the optimization problem (4) using the cvxpy package with SCS backend, see Algorithm B for the source code 3 . With the default numerical tolerance, 10 -4 , each factorization took a few minutes ( n ≤ 100 ) to hours ( n ≤ 500 ) to several days ( n ≥ 700 ) of CPU time. Note that this overhead reappears for any change in the number of update steps, n , weight decay, α , or momentum, β , as these induce different workload matrices. In our experiments, when the optimization for AOF did not terminate within 10 days, we reran the optimization problem with the tolerance increased by a factor of 10. The runtime depends not only on the matrix size but also on the entries. In particular, we observe matrices with momentum to be harder to factorize than without. For large matrix sizes we frequently encountered numerical problems: the intermediate matrices, S , in (4), often did not fulfill the positive definiteness condition required to solve the subsequent Cholesky decomposition for C . Unfortunately, simply projecting the intermediates back to the cone of positive semi-definite matrices is not enough, because the resulting C matrices also have to be invertible and not too badly conditioned. Ultimately, we adopted a postprocessing step for S that ensures that all its eigenvalues were at least of value √ 1 /n , which we find to be a reasonable modification to ensure the stability of the convergence. Enforcing this empirically found value leads to generally good results, as our experiments below show, but it does add an undesirable extra level of complexity to the process. In contrast, due to its analytic expressions, BSR does not suffer from numerical problems. It also does not possess additional hyperparameters, such as a numeric tolerance or the number of optimization steps.

Apart from the factorization itself, the computational cost of BSR and AOF are nearly identical. Both methods produce (banded) lower triangular matrices, so computing the inverse matrices or solving linear systems can be done within milliseconds to seconds using forward substitution. Note that, in principle, one could even exploit the Toeplitz structure of p -BSR, but we found this not to yield any practical benefit in our experiments. Computing the sensitivity is trivial for p -BSR using Corollary 1, and it is still efficient for AOF by the dynamic program proposed in Choquette-Choo et al. [2023a].

Expected Approximation Error. As a first numeric experiment, we evaluate the expected approximation error for workload matrices that reflect different SGD settings. Specifically, we use workload matrices (9) for n ∈ { 100 , 200 , . . . , 1000 , 1500 , 2000 } , with α = { 0 . 99 , 0 . 999 , 0 . 9999 , 1 } , and β ∈ { 0 , 0 . 9 } , either with single participation, k = 1 , or repeated participation, b = 100 , k = n/ 100 . Figure 1 shows the expected approximate error, E ( B,C ) , of the proposed BSR, AOF, as well as the baseline factorizations, A = A · Id and A = Id · A in two exemplary cases. Additional results for other privacy levels can be found in Appendix C.

200

400

600

800 1000

1500

2000

workload matrix size,

n

0

20

40

60

80

approximation quality,

=0.999

=0.0

b

=100

BSR

AOF

id

A

200 400 600 8001000

1500

2000

workload matrix size,

n

0

500

1000

1500

approximation quality,

=1.0

=0.9

b

=100

BSR

AOF

id

A

25 200

500

1000

2000

number of update steps

10

20

30

40

accuracy

CIFAR-10 (single participation)

BSR

AOF

id

A

5

10

15

20

maximum number of epochs

20

30

40

50

accuracy

CIFAR-10 (repeated participation)

BSR

AOF

id

A

The results confirm our expectations from the theoretical analysis: in particular, BSR's expected approximation error is quite close to AOF's, typically within a few percent (left plot). Both methods are clearly superior to the naive factorizations. For large matrix sizes, BSR sometimes even yields slightly better values than AOF (right plot). However, we believe this to be a numeric artifact of us having to solve AOF with less-than-perfect precision.

Private Model Training on CIFAR-10. To demonstrate the usefulness of BSR in practical settings, we follow the setup of Kairouz et al. [2021] and report results for training a simple ConvNet on the CIFAR-10 dataset (see Table 1 in Appendix C for the architecture). Specifically, we adapt Google's reference implementation of DP-SGD in jax Bradbury et al. [2018] to work with the different matrix factorizations: BSR, AOF, and the two baselines. To reflect the setting of single-participation training, we split the 50,000 training examples into batches of size m ∈ { 1000 , 500 , 250 , 200 , 100 , 50 , 25 } , resulting in n ∈ { 100 , 200 , 400 , 500 , 1000 , 2000 } update steps. For repeated participation, we fix the batch size to 500 and run k ∈ { 1 , 2 , . . . , 10 , 15 , 20 } epoch of training, i.e. n = 100 k and b = 100 . In both cases, 20% of the training examples are used as validation sets to determine the learning rate η ∈ { 0 . 01 , 0 . 05 , 0 . 1 , 0 . 5 , 1 } , weight decay parameters α ∈ { 0 . 99 , 0 . 999 , 0 . 9999 , 1 } , and momentum β ∈ { 0 , 0 . 9 } . Figure 2 shows the test set accuracy of the model trained with hyperparameters that achieved the highest validation accuracy. 4 One can see the expected effect that in DP model training, more update steps/epochs do not necessary lead to higher accuracy due to the need to add more noise. The quality of models trained with BSR is mostly identical to AOF. When training for a large number of epochs it achieves even better slightly results, but this could also be an artifact of us having to solve AOF with reduced precision in this regime. Both methods are clearly superior to the baselines.

5 Conclusion and Discussion

We introduce an efficient and effective approach to the matrix factorization mechanism for SGDbased model training with differential privacy. The proposed banded square root factorization (BSR) factorization achieves results on par with the previous state-of-the-art, and clearly superior to baseline methods. At the same time, it does not suffer from the previous method's computational overhead, thereby making differentially private model training practical even for large scale problems.

Despite the promising results, some open questions remain. On the theoretical side, the asymptotic optimality of BSR without weight decay is still unresolved because the current upper bounds on the expected approximation error do not match the provided lower bounds. Based on the experimental results, we believe this discrepancy lies with the lower bounds, which we suspect should be linear in the number of participations. We observe that BSR achieves results comparable to AOF, although we cannot currently prove this due to the insufficient understanding of AOF's theoretical properties; nonetheless, we consider it a promising research direction. On the practical side, it would be interesting to extend the guarantees to even more learning scenarios, such as variable learning rates.

Acknowledgments and Disclosure of Funding

This research was supported by the Scientific Service Units (SSU) of ISTA through resources provided by Scientific Computing (SciComp). We thank Monika Henzinger and Jalaj Upadhyay for their valuable comments on the earlier versions of this manuscript.

References

Abadi, A. Chu, I. Goodfellow, H. B. McMahan, I. Mironov, K. Talwar, and L. Zhang. Deep learning with differential privacy. In ACM Special Interest Group on Security, Audit and Control (SIGSAC) , 2016.

Batir, H. Kücük, and S. Sorgun. Convolution identities involving the central binomial coefficients and Catalan numbers. Transactions on Combinatorics , 2021.

Å. Björck and S. Hammarling. A Schur method for the square root of a matrix. Linear Algebra and its Applications , 1983.

Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax .

A. Choquette-Choo, A. Ganesh, R. McKenna, H. B. McMahan, J. K. Rush, A. G. Thakurta, and X. Zheng. (Amplified) banded matrix factorization: A unified approach to private training. In Conference on Neural Information Processing Systems (NeurIPS) , 2023a.

A. Choquette-Choo, A. Ganesh, T. Steinke, and A. Thakurta. Privacy amplification for matrix mechanisms. In International Conference on Learning Representations (ICLR) , 2023b.

A. Choquette-Choo, H. B. McMahan, K. Rush, and A. Thakurta. Multi-epoch matrix factorization mechanisms for private machine learning. International Conference on Machine Learing (ICML) , 2023c.

A. Choquette-Choo, K. Dvijotham, K. Pillutla, A. Ganesh, T. Steinke, and A. Thakurta. Correlated noise provably beats independent noise for differentially private learning. In International Conference on Learning Representations (ICLR) , 2024.

Deadman, N. J. Higham, and R. Ralha. Blocked schur algorithms for computing the matrix square root. In International Workshop on Applied Parallel Computing (PARA) , 2012.

Denisov, H. B. McMahan, J. Rush, A. Smith, and G. A. Thakurta. Improved Differential Privacy for SGD via optimal private linear operators on adaptive streams. In Conference on Neural Information Processing Systems (NeurIPS) , 2022.

Dvijotham, H. B. McMahan, K. Pillutla, T. Steinke, and A. Thakurta. Efficient and near-optimal noise generation for streaming differential privacy, 2024. arXiv:2404.16706 [cs.DS].

Dwork. Differential privacy. In International colloquium on automata, languages, and programming (ICALP) , 2006.

Dwork and A. Roth. The algorithmic foundations of differential privacy . Foundations and Trends in Theoretical Computer Science, 2014.

Dwork, F. McSherry, K. Nissim, and A. Smith. Calibrating noise to sensitivity in private data analysis. In Theory of Cryptography Conference (TCC) , 2006.

F. Elliott. The characteristic roots of certain real symmetric matrices, 1953. M. S. thesis, University of Tennessee,.

A. Gershgorin. Über die Abgrenzung der Eigenwerte einer Matrix. Proceedings of the USSR Academy of Sciences. Mathematics Series , (6):749-754, 1931.

Granqvist, C. Song, Áine Cahill, R. van Dalen, M. Pelikan, Y. S. Chan, X. Feng, N. Krishnaswami, V. Jina, and M. Chitnis. pfl-research: simulation framework for accelerating research in private federated learning, 2024. URL https://arxiv.org/abs/2404.06430 .

H. Greene and D. E. Knuth. Mathematics for the Analysis of Algorithms . Springer, 1990.

Henzinger, J. Upadhyay, and S. Upadhyay. A unifying framework for differentially private sums under continual observation. In Symposium on Discrete Algorithms (SODA) , 2024.

Kairouz, B. McMahan, S. Song, O. Thakkar, A. Thakurta, and Z. Xu. Practical and private (deep) learning without sampling or shuffling. In International Conference on Machine Learing (ICML) , 2021.

Koloskova, R. McKenna, Z. Charles, J. Rush, and H. B. McMahan. Gradient descent with linearly correlated noise: Theory and applications to differential privacy. In Conference on Neural Information Processing Systems (NeurIPS) , 2023.

Kurakin, S. Song, S. Chien, R. Geambasu, A. Terzis, and A. Thakurta. Toward training at imagenet scale with differential privacy, 2022. arXiv:2201.12328.

Li, G. Miklau, M. Hay, A. McGregor, and V. Rastogi. The matrix mechanism: Optimizing linear counting queries under Differential Privacy. International Conference on Very Large Data Bases (VLDB) , 2015.

Li, B. Ding, C. Zhang, N. Li, and J. Zhou. Federated matrix factorization with privacy guarantee. International Conference on Very Large Data Bases (VLDB) , 2021.

J. MacWilliams and N. J. A. Sloane. The theory of error-correcting codes . Elsevier, 1977.

McKenna. Scaling up the banded matrix factorization mechanism for differentially private ml, 2024. arXiv:2405.15913 [cs.DS].

B. McMahan, Z. Xu, and Y. Zhang. A hassle-free algorithm for private learning in practice: Don't use tree aggregation, use blts, 2024. arXiv:2408.08868 [cs.DS].

V. Nguyen, Y. Mori, and T. Mori. Relaxed monotonic conditions for Schur stability of real polynomials. IEICE Transactions on Fundamentals of Electronics, Communications and Computer Sciences , 90(10):2326-2328, 2007.

Papernot and T. Steinke. Hyperparameter tuning with renyi differential privacy. In International Conference on Learning Representations (ICLR) , 2021.

Ponomareva, H. Hazimeh, A. Kurakin, Z. Xu, C. Denison, H. B. McMahan, S. Vassilvitskii, S. Chien, and A. G. Thakurta. How to DP-fy ML: A practical guide to machine learning with differential privacy. Journal of Artificial Intelligence Research (JAIR) , 77:1113-1201, 2023.

SebastienB. Trace norm of a triangular matrix with only ones above the diagonal, 2017. URL https://math.stackexchange.com/questions/1857078/ trace-norm-of-a-triangular-matrix-with-only-ones-above-the-diagonal .

Shin, S. Kim, J. Shin, and X. Xiao. Privacy enhanced matrix factorization for recommendation with local differential privacy. IEEE Transactions on Knowledge and Data Engineering , 30(9), 2018.

Vadhan. The complexity of differential privacy . Tutorials on the Foundations of Cryptography, 2017.

Yu, H. Zhang, W. Chen, J. Yin, and T.-Y. Liu. Gradient perturbation is underrated for differentially private convex optimization. In International Joint Conference on Artificial Intelligence (IJCAI) , 2020.

Zhang, J. Ding, M. Wu, S. T. Wong, H. Van Nguyen, and M. Pan. Adaptive privacy preserving deep learning algorithms for medical data. In Winter Conference on Applications of Computer Vision (WACV) , 2021.

Appendix

A General introduction to differential privacy.

Differential privacy [Dwork, 2006] is a robust framework designed to provide strong privacy guarantees for statistical analyses and data sharing. It aims to protect individual data points in a dataset while still allowing meaningful aggregate information to be extracted. Unlike traditional data anonymization techniques, which might involve removing identifiers or aggregating data, differential privacy offers a mathematical definition of privacy that quantifies the amount of privacy loss and ensures that the risk of identifying any individual's data remains low, even when combined with other data sources.To formalize this concept, a randomized mechanism M is said to provide ( ε, δ ) -differential privacy if, for all data sets D and D ′ that differ in one element, and for all subsets of the mechanism's output space S :

For a detailed introduction to differential privacy, we recommend the books "The Algorithmic Foundations of Differential Privacy" by Dwork and Roth [2014] and "The Complexity of Differential Privacy" by Vadhan [2017].

B Source code for computing AOF

Algorithm 2 Source code for computing AOF using cvxpy .

C Network architecture for CIFAR-10 experiments

D Additional Experimental Results

In this section we provide additional experiments comparing BSR, AOF and baselines: Figures 3 and 4 and following tables show their expected approximation error (lower is better) for workload matrices stemming from SGD with different hyperparameter settings. Figure 5) and following tables show the accuracy of resulting classifiers on CIFAR-10 (higher is better) for different privacy levels.

The results show the same trends as the one in Section 4. BSR achieves almost identical expected approximation error as AOF, and results in equally good classifiers. In some cases, results for BSR even improve over AOF's. Presumably this is because of numerical issues in solving the optimization problem for AOF.

200 400 600 800 1000

1500

2000

workload matrix size,

n

0

50

100

150

approximation quality,

=1.0

=0.0

b

=100

BSR

AOF

id

A

200 400 600 8001000

1500

2000

workload matrix size,

n

0

500

1000

1500

approximation quality,

=1.0

=0.9

b

=100

BSR

AOF

id

A

200 400 600 800 1000

1500

2000

workload matrix size,

n

0

50

100

approximation quality,

=0.9999

=0.0

b

=100

BSR

AOF

id

A

200 400 600 8001000

1500

2000

workload matrix size,

n

0

500

1000

approximation quality,

=0.9999

=0.9

b

=100

BSR

AOF

id

A

200

400

600

800 1000

1500

2000

workload matrix size,

n

0

20

40

60

80

approximation quality,

=0.999

=0.0

b

=100

BSR

AOF

id

A

200 400 600 800 1000

1500

2000

workload matrix size,

n

0

200

400

600

800

approximation quality,

=0.999

=0.9

b

=100

BSR

AOF

id

A

200

400

600

800 1000

1500

2000

workload matrix size,

n

0

10

20

30

approximation quality,

=0.99

=0.0

b

=100

BSR

AOF

id

A

200 400 600 800 1000

1500

2000

workload matrix size,

n

0

100

200

300

approximation quality,

=0.99

=0.9

b

=100

BSR

AOF

id

A

200

500

1000

2000

workload matrix size,

n

0

10

20

30

approximation quality,

=1.0

=0.0

k

=1

BSR

AOF

id

A

200

500

1000

2000

workload matrix size,

n

0

100

200

300

approximation quality,

=1.0

=0.9

k

=1

BSR

AOF

id

A

200

500

1000

2000

workload matrix size,

n

0

10

20

30

approximation quality,

=0.9999

=0.0

k

=1

BSR

AOF

id

A

200

500

1000

2000

workload matrix size,

n

0

100

200

300

approximation quality,

=0.9999

=0.9

k

=1

BSR

AOF

id

A

200

500

1000

2000

workload matrix size,

n

0

5

10

15

20

approximation quality,

=0.999

=0.0

k

=1

BSR

AOF

id

A

200

500

1000

2000

workload matrix size,

n

0

50

100

150

200

approximation quality,

=0.999

=0.9

k

=1

BSR

AOF

id

A

200

500

1000

2000

workload matrix size,

n

0

2

4

6

approximation quality,

=0.99

=0.0

k

=1

BSR

AOF

id

A

200

500

1000

2000

workload matrix size,

n

0

20

40

60

approximation quality,

=0.99

=0.9

k

=1

BSR

AOF

id

A

25 200

500

1000

2000

number of update steps

10

20

30

accuracy

CIFAR-10 (single participation)

BSR

AOF

id

A

5

10

15

20

maximum number of epochs

20

30

40

accuracy

CIFAR-10 (repeated participation)

BSR

AOF

id

A

25 200

500

1000

2000

number of update steps

20

30

40

accuracy

CIFAR-10 (single participation)

BSR

AOF

id

A

5

10

15

20

maximum number of epochs

20

30

40

50

accuracy

CIFAR-10 (repeated participation)

BSR

AOF

id

A

E Experimental Results for Different Optimizers

In this section we report on experimental results when different optimizers are used to (approximately) solve the AOF optimization problem (4). Besides cvxpy (CVX) these are standard gradient descent (GD) and the Limited-Memory Broyden-Fletcher-Goldfarb-Shanno algorithm (LBFGS) . The latter two we implement in jax using the optax toolbox. Similar to [Granqvist et al., 2024, ftrl_mechanism.py ], we use an adaptive line-search for the step size of the gradient-based methods, which at the same time ensures the positive definiteness constraints of the optimization problem. Our implementation differs from theirs, however, in that our learning rate is not restricted to shrink monotonically, thereby avoiding premature termination.

E.1 Runtime

We report the runtimes for the different methods in Tables 24 to 39. For comparison, we also include results for BSR and the CVX optimizers with three tolerance levels in the same settings, where practically feasible. Note that while the experiments for BSR and CVX used a single-core CPU-only environment, the experiments for GD and LBFGS were run on an NVIDIA H100 GPU with 16 available CPU cores. As a consequence, the absolute runtimes are not directly comparable between the methods, but they should rather be seen as illustrations of the scaling behavior of the method for different workload types and problem sizes.

Indeed, the results show a clear trend: BSR is the fastest, with almost no overhead. Even for the largest problem sizes of n = 10000 , BSR never took more than 2.5s to despite running in the single-core CPU-only setup. GD and LBFGS benefit strongly from the GPU hardware. In the multiple participation setting ( p = 100 , k = n/p ), they solve most workload sizes within a few

minutes, except the largest ones, which for GD can take a few hours. In the single participation setting ( k = 1 ), LBFGS also occasionally need several hours to converge. In general, stronger weight decay (smaller α ) tends to lead to lower runtimes, while the use of momentum ( β = 0 . 9 ) to higher times until convergence. CVX (on weak hardware) is orders of magnitude slower than the other methods. Furthermore, its runtime grow approximately cubic with the problem size, whereas for GD and LBFGS the relation is not too far from linear. Note that despite the stable patterns described above, all runtime results should be taken with caution, because internal parameters of the optimization, such as the convergence criterion and the specific implementation of the line search can substantially influence the overall runtime as well.

E.2 Expected Approximation Error

Figures 6 and 7 report the expected approximation errors achieved by the different optimizers of AOF (4) and by BSR. For CVX, we report the smallest error value across all tolerance levels for which the optimization converged.

The curves show several clear trends. GD and LBFGS generally perform similarly, and achieve expected approximation errors slightly (at most a few percent) lower than BSR. An exception are the problems with momentum ( β = 0 . 9 ) in the single participation setting, where it appears that SGD occasionally fails to find the optimum for large problem sizes ( n ≥ 1500 ). CVXperforms comparably to the other methods for problems without momentum ( β = 0 ). With momentum, however, the solutions it find are often worse than the other methods, especially in the single-participation setting and for medium to large problem sizes ( n ≥ 500 ). Presumably, even smaller tolerance values would be require here, which, however, would result in even longer runtimes.

100

1000

10000

workload matrix size,

n

10 2

approximation quality,

=1.0,

=0.9,

p

=

b

=100,

k

=

n

/

b

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

10 1

10 2

approximation quality,

=1.0,

=0.0,

p

=

b

=100,

k

=

n

/

b

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

10 2

approximation quality,

=0.9999,

=0.9,

p

=

b

=100,

k

=

n

/

b

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

10 1

approximation quality,

=0.9999,

=0.0,

p

=

b

=100,

k

=

n

/

b

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

10 2

approximation quality,

=0.999,

=0.9,

p

=

b

=100,

k

=

n

/

b

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

10 1

approximation quality,

=0.999,

=0.0,

p

=

b

=100,

k

=

n

/

b

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

10 1

10 2

approximation quality,

=0.99,

=0.9,

p

=

b

=100,

k

=

n

/

b

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

10 1

approximation quality,

=0.99,

=0.0,

p

=

b

=100,

k

=

n

/

b

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

2×10 1

3×10 1

4×10 1

approximation quality,

=1.0,

=0.9,

k

=1

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

3×10 0

4×10 0

approximation quality,

=1.0,

=0.0,

k

=1

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

2×10 1

3×10 1

approximation quality,

=0.9999,

=0.9,

k

=1

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

2.25 × 10 0

2.5 × 10 0

2.75 × 10 0

3×10 0

3.25 × 10 0

3.5 × 10 0

3.75 × 10 0

approximation quality,

=0.9999,

=0.0,

k

=1

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

1.2 × 10 1

1.4 × 10 1

1.6 × 10 1

1.8 × 10 1

2×10 1

2.2 × 10 1

2.4 × 10 1

approximation quality,

=0.999,

=0.9,

k

=1

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

2×10 0

3×10 0

approximation quality,

=0.999,

=0.0,

k

=1

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

10 1

1.1 × 10 1

1.2 × 10 1

1.3 × 10 1

1.4 × 10 1

1.5 × 10 1

approximation quality,

=0.99,

=0.9,

k

=1

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

100

1000

10000

workload matrix size,

n

1.8 × 10 0

1.9 × 10 0

2×10 0

2.1 × 10 0

2.2 × 10 0

2.3 × 10 0

approximation quality,

=0.99,

=0.0,

k

=1

BSR

AOF-CVX

AOF-GD

AOF-LBFGS

F Complete Proofs

In this section, we provide the complete proofs for our results from the main manuscript. For the convenience of the reader, we also restate the statements themselves.

F.1 Proof of Theorem 1

Theorem 1 (Square-Root of SGD Workload Matrix) . Let A α,β be the workload matrix (9) . Then A α,β = C 2 α,β for C α,β = LDToep( c 0 , . . . , c n -1 ) , with c 0 = 1 and c j = ∑ j i =0 α j -i r j -i r i β i for j = 1 , . . . , n -1 with coefficients r i = ∣ ∣ ( -1 / 2 i )∣ ∣ . For any p ∈ { 1 , . . . , n } , the p -banded BSR matrix C | p | α,β is obtained from this by setting all coefficients c j = 0 for j ≥ p .

Proof. We observe that A α,β can be written as

Relying on the result from Henzinger et al. [2024], that E 1 / 2 1 =     1 0 . . . 0 r 1 1 . . . 0 . . . . . . . . . . . . r n -1 r n -2 . . . 1     with r k = ∣ ∣ ∣ ( -1 / 2 k ) ∣ ∣ ∣ , one can check that the square roots of the matrices E α , E β are:

An explicit check yields that these matrices commute, i.e. E 1 / 2 α E 1 / 2 β = E 1 / 2 β E 1 / 2 α . Therefore

F.2 Proof of Theorem 2

Theorem 2 (Sensitivity for decreasing non-negative Toeplitz matrices) . Let M = LDToep( m 0 , . . . , m n -1 ) be a lower triangular Toeplitz matrix with decreasing non-negative entries, i.e.

m 0 ≥ m 1 ≥ m 2 ≥ . . . m n -1 ≥ 0 . Then its sensitivity (2) in the setting of b -min-separation is

where M [ · , 1+ jb ] denotes the (1 + jb ) -th column of M .

Proof. Because all entries of M are positive, so are the entries of M glyph[latticetop] M . Therefore, the condition is fulfilled such that (2) holds with equality, and

where Π k,b = { π ⊂ { 1 , . . . , n } : | π | ≤ k ∧ ( { i, j } ⊂ π ⇒ i = j ∨ | i -j | ≥ b ) } . We now establish that { 1 , 1 + b, . . . , 1 + ( k -1) b } is an optimal index set, which implies the statement of the theorem.

To see this, let π ∗ be any optimal solution and let i ∗ ∈ π ∗ be a column index that is not as far left as possible, i.e., π ∗ \ { i ∗ } ∪ { i ∗ -1 } would be a valid index set in Π k,b . If such an index i ∗ exists, we split π ∗ into the left indices, which are smaller than i ∗ , and the remaining, right , ones: π ∗ = π ∗ L ˙ ∪ π ∗ R with π ∗ L = { i | i ∈ π ∧ i < i ∗ } , π ∗ R = { i | i ∈ π ∧ i ≥ i ∗ } . Then, we construct a new index set in which the left indices are kept but all right ones are shifted by one position to the left: π ′ = π ∗ L ∪ ← -π ∗ R with ← -π ∗ R = { i -1 | i ∈ π ∗ R } . By the condition on i ∗ , we know π ′ ∈ Π k,b .

Wenow prove that f ( π ′ , π ′ ) ≥ f ( π, π ) , so π ′ must also be optimal. First, we observe two inequalities: for any i, j > 1 :

and for i ≥ 1 , j > 1 :

where the last inequality holds because by assumption m l -i -1 ≥ m l -i for l ≥ i +1 .

Now, we split f ( π ′ , π ′ ) and f ( π ′ , π ′ ) into three terms: the inner products of indices below i ∗ , the ones of terms above i ∗ and the ones between both,

The first term appears identically in both expressions. The second term fulfills

by Equation (24). The third term fulfills

by Equation (28). In combination, this establishes f ( π ′ , π ′ ) ≥ f ( π ∗ , π ∗ ) , and since π ∗ was already optimal, the same must hold for π ′ .

Using the above construction, we can create a new optimal index sets, π ∗ , until reaching one that does not contain any index i ∗ as described anymore. Then π ∗ = { 1 , 1+ b, . . . , 1+( l -1) b } for some l ∈ N must hold. If l = k , the statement of Theorem 2 is confirmed. Otherwise, π ′ = { 1 , . . . , 1+( k -1) b } is superset of π ∗ , so because of the positivity of entries, f ( π ′ , π ′ ) ≥ f ( π ∗ , π ∗ ) must hold. Once again, because π ∗ was optimal, the same must hold for π ′ , which concludes the proof.

F.3 Proof of Corollary 1

Corollary 1. The sensitivity of the p -BSR for SGD can be computed using formula (10) .

Proof. From (1) we know that C α,β is a Toeplitz matrix with coefficients (1 , c 1 , . . . , c n -1 ) , where c j = ∑ j i =0 α i r i r j -i β j -i for 0 ≤ β < α ≤ 1 , with r i = | ( -1 / 2 i ) | = B i 4 i , where B i = ( 2 i i ) is the i -central binomial coefficient . It suffices to show that c j ≥ c j +1 for any j ∈ { 1 , . . . , n -1 } .

First, we show for the r i coefficients:

where C j = 1 j +1 B j = 1 j +1 ( 2 j j ) is the j -th Catalan number .

Now, we study the case α = 1 . If β = 0 , then c 1 = c 2 = · · · = c n = 1 , so monotonicity is fulfilled. Otherwise, i.e. 0 < β < 1 , we write

Using the classic identity between Catalan numbers, 2 ∑ k i =0 B k -i C i = B k +1 , e.g. [Batir et al., 2021, Identity 4.2] we obtain

where the last inequality follow from the fact that β -i -β > 0 for each i = 0 , . . . , k and any β < 1 . This proves the monotonicity of c k .

For α < 1 , we observe that c j = α j ∑ j i =1 r i r k -i γ j -i for γ = α β < 1 . Clearly, the sequence α j is decreasing, and by the above argument, the sum is decreasing, too. Consequently, c j is the product of two decreasing sequences, so it is also decreasing, which concludes the proof.

F.4 Useful Lemmas

Before the remaining proofs, we establish a number of useful lemmas.

Lemma 2. For any C ∈ R n × n with C glyph[latticetop] C ≥ 0 it holds for any b ∈ { 1 , . . . , n } that

where ‖ C ‖ 2 2 , ∞ = max i =1 ,...,n ‖ C [ · ,i ] ‖ 2 .

Proof. This follows directly from Theorem 2:

Lemma 3. For any C ∈ R n × n with C glyph[latticetop] C ≥ 0 it holds for any b ∈ { 1 , . . . , n } and k ∈ { 1 , . . . , b n } that

Proof. We first show the upper bound. Observe that for any π ⊂ [ n ] :

Therefore, using Theorem 2:

For the lower bound, we introduce some additional notation. Let ˜ Π k be the set of b -separated index sets with exactly k elements. Then, from Theorem 2, we obtain

with the notation S ( I ) = ∑ i ∈ I ‖ C [ · ,i ] ‖ 2 for any index set I ⊂ { 1 , . . . , n } .

Now, we prove by backwards induction over k = 1 , . . . , n b :

As base case, let k = n b . Denote by π i := { i, i + b, i +2 b, . . . , i +( n -b ) } for i = 1 , . . . , b the uniformly spaced index sets. By construction they all fulfill π i ∈ ˜ Π n/b and ⋃ n i =1 π i = [ n ] , where the union is disjoint. Therefore

This proves the statement (49), because 1 b = k n in this case.

As an induction step, we prove that if (49) holds for some value k ≤ n b , then it also holds for k -1 ≥ 1 .

Let π ∗ ∈ argmax π ∈ ˜ Π k S ( π ) and j ∗ = argmin j ∈ π ∗ S ( { j } ) , such that we know that S ( { j } ) ≤ 1 k S ( π ∗ ) . Now, set π ′ = π ∗ \ { j ∗ } . Because π ′ ∈ ˜ Π k -1 , it follows that

where in the last step we used the induction hypothesis. This concludes the proof.

Lemma 4. For C α,β as in (1) , and k = 1 , it holds that

Proof. From Lemmas 3 and 2 we obtain

where the last identify follows from the explicit form of C α,β . The lower bound follows from

and again the explicit form of ‖ C ‖ 2

F .

Lemma 5. For r j = | ( -1 / 2 j ) | = 1 4 j ( 2 j j ) it holds that:

Proof. The double inequality is a particular case of a more general pair of binomial inequalities when k = j and m = 2 j :

where H ( k/m ) is the binary entropy function, with H (1 / 2) = 1 . The proof of the general result (56), can be found in MacWilliams and Sloane [1977, Chapter 10, Lemma 7, p309].

Lemma 6. Let c k = ∑ k j =0 α j r j r k -j β k -j as in (1) . Then c 0 = 1 , and for j ≥ 1 :

Proof. We exploit the upper and lower bounds from Lemma 5. First, we write c k = α k ∑ k j =0 r j r k -j γ j with γ := β α . Then we check immediately that c 0 = 1 and c 1 = 1 2 ( α + β ) = α 2 (1 + γ ) ≤ α 2 1 1 -γ .

For j ≥ 2 we derive the upper bound by

which proves the upper bound on a j . The lower bound for j ≥ 1 follows trivially from

Lemma 7. For j ∈ { 1 , . . . , n } it holds

Proof. We first prove the result for α = 1 . Combining Lemmas 4 and 6 we obtain

For α < 1 , if follows analogously:

where the last term emerges from ∑ ∞ i = j +1 α 2 i i ≥ α 2( j +1) j +1 ∑ ∞ i =0 α 2 i = α 2( j +1) j +1 1 1 -α 2 .

Lemma 8. Let 0 ≤ β < α ≤ 1 . Let σ 1 ≥ · · · ≥ σ n be the sorted list of singular values of A α,β . If α < 1 , then for j = 1 , . . . , n :

If α = 1 , then for j = 1 , . . . , n

and consequently

Proof. The statements on the singular values follow from the following Lemma 9, because A α,β = E α E β . Because E α and E β are diagonalizable and they commute, we have σ n ( E β ) σ j ( E α ) ≤ σ j ( E α E β ) ≤ σ 1 ( E β ) σ j ( E α ) . For α < 1 the lower bound follows from ‖ A α,β ‖ ∗ ≥ trace A α,β , and the upper bound follows from the identity ‖ A α,β ‖ ∗ = ∑ n j =1 σ j .

For α = 1 , the bounds follow from the same identity together with the fact that

for α = 1 , and otherwise and

Lemma 9 (Singular values of E t ) . For 0 ≤ t ≤ 1 , let E t = LDToep(1 , t, . . . , t n -1 ) ∈ R n × n . Then the singular values σ 1 ( E t ) ≥ · · · ≥ σ n ( E t ) fulfill for i = 1 , . . . , n :

Proof. Wefollow the steps of SebastienB [2017], and use that the singular values of E t are the reciprocals of the singular values of E -1 t , which themselves are the eigenvalues of ( E t ) -1 (( E t ) -1 ) glyph[latticetop] =: T , i.e., for i = 1 , . . . , n :

The E -1 t and T can be computed explicitly as

Lemma 10. All eigenvalues, µ , of T fulfill

Proof. By Gershgorin's circle theorem [Gershgorin, 1931], we know that µ fulfills i) | 1 -µ | ≤ t , i.e. 1 -t ≤ t ≤ µ ≤ 1 + t or ii) | 1 + t 2 -µ | ≤ 2 t , i.e. 1 -2 t + t 2 ≤ µ ≤ 1 + 2 t + t 2 . For t ∈ [0 , 1] the first condition implies the second, so (79) must hold.

Case I: For t < 1 , the statement (76) follows from Lemma 10 in combination with (77).

exactly Toeplitz, because of the top left entry, so closed-form expressions for the eigenvalues of tridiagonal Toeplitz matrices do not apply to it. Instead, we can compute its eigenvalues explicitly. Matrices of this form have been studied by Elliott [1953]; for completeness, we provide a full proof here.

Let µ be an eigenvalue of T with eigenvector Ψ = (Ψ 0 , . . . , Ψ n -1 ) . From the eigenvector equation T Ψ = µ Ψ we obtain

which yields a linear recurrence relation

with two boundary conditions

We solve the recurrence relation using the polynomial method [Greene and Knuth, 1990]. The characteristic polynomial of (83) is P ( z ) = z 2 +( µ -2) z +1 . Its roots are

for some value θ ∈ [0 , 2 π ) . Note that the expression under the second square root is positive, because of Lemma 10. The last equation is a consequence of, | r ± | 2 = 1 4 ( (2 -µ ) 2 +(4 -(2 -µ ) 2 ) ) = 1 . Consequently,

From standard results on linear recurrence, it follows that any solution to (83) has the form Ψ j = c 1 ( r + ) j + c 2 ( r -) j for some constants c 1 , c 2 ∈ C . The fact that Ψ j must be real-valued implies that c 1 = c 2 =: αe iφ for some values α ∈ R , φ ∈ [0 , 2 π ) . Dropping the normalization constant (which we could recover later if needed), we obtain

Next, we use the boundary conditions to establish values for φ and θ .

Equation (85) can be rewritten as

which, using 2 cos( α + β ) = cos( a + b ) + cos( a -b ) , simplifies to

Consequently, φ + nθ = 1 2 π + kπ must hold for some k ∈ N .

Equation (84) can be rewritten as

One solution to this would be θ = 0 , but that would implies µ = 0 , which is inconsistent with T being an invertible matrix. So instead, it must hold that φ = θ 2 + kπ for some k ∈ N .

Combining both conditions and solving for θ we obtain

Each such value θ k for k ∈ { 0 , . . . , n -1 } yields an eigenvector with associated eigenvalue µ = 2 -2 cos θ k = 4sin 2 ( θ k / 2) . Now, (76) follows from this in combination with (77).

F.5 Proof of Theorem 3

Theorem 3 (Expected approximation error with single participation) . Let A α,β ∈ R n × n be the workload matrix (9) of SGD with momentum 0 ≤ β < 1 and weight decay parameter 0 < α ≤ 1 , where α > β . Assume that each data item can contribute at most once to an update vector (e.g. single participation, k = 1 ). Then, the expected approximation error of the square root factorization , A α,β = C 2 α,β , fulfills

which simplifies to for α < 1 , and

Proof. The proof consists of a combination of Lemmas 4 and 7. Because in the single participation k = 1 , so we need just the first column of matrix C α,β :

which proves the upper bounds. For the lower bounds, for any α ≤ 1 :

Also, for α = 1 :

because log( ( n +1)!) ≥ ( n +1)log( n +1) -n .

F.6 Proof of Theorem 5

Theorem 5. Assume the setting of Theorem 3. Then, the baseline factorizations A α,β = A α,β · Id and A α,β = Id · A α,β fulfill, for α < 1 ,

Proof. For α = 1 , by Lemma 2 we have:

For α < 1 :

Together with

as β < 1 and the sum ∑ n -1 j =0 jβ j +1 is uniformly bounded by β ∑ ∞ j =0 jβ j = β 2 (1 -β ) 2 For α < 1 :

where the second equality due to the fact that we average over the sequence jx j +1 which converges to 0 for | x | < 1 .

F.7 Proof of Theorem 6

Theorem 6 (Approximation error of BSR) . Let A α,β ∈ R n × n be the workload matrix (9) of SGD with momentum 0 ≤ β < 1 and weight decay 0 < α ≤ 1 , with α > β . Let A α,β = B | p | α,β C | p | α,β , be its banded square root factorization as in Definition 3. Then, for any b ∈ { 1 , . . . , n } , p ≤ b , and k ∈ { 1 , . . . , n b } it holds:

Proof. Consider a Lower Triangular Toeplitz (LTT) matrix multiplication:

where c j = ∑ j i =1 a i b n +1 -i so the LTT structure is preserved with multiplication that allows us to work with sequences and their convolutions rather than matrix multiplication. For instance, we would write the previous product in the form:

The inverse of the Lower Triangular Toeplitz matrix remains a Lower Triangular Toeplitz (LTT) matrix because we can find a unique sequence ( c 1 , . . . , c n ) such that:

with the restriction that a 1 = 0 ; otherwise, the original matrix was not invertible. We consider the banded square root factorization A α,β = B | p | α,β C | p | α,β which is characterized by the following identity:

glyph[negationslash]

We will bound the Frobenius norm of the LTT matrix ( b 0 , . . . , b n -1 ) . By the uniqueness of the solution, we obtain that for the first p values we have b i = c i . For the next p values we have the following formula:

By induction argument, we can see that b p + j = 2 c p + j for 0 ≤ j ≤ p -1 . For the remaining n -2 p values we will prove convergence to a constant.

We make an ansatz for the solution of the linear recurrence in the form:

where y j represents the error terms, which will be proven to converge to 0 . The sequence y j satisfies the following recurrence formula:

We denote w j = c j α -j which is a decreasing sequence because the values correspond to the C 1 ,β/α matrix. We rewrite the recurrence in matrix notation:

To show that the error terms y j goes to 0 as j goes to infinity, we first study the characteristic polynomial of the associate homogeneous relations:

Because 1 > w 1 > w 2 > · · · > w b -1 > 0 , it follows from Schur's (relaxed) stability condition [Nguyen et al., 2007, Theorem 1] that all its (complex) roots lie inside of the open unit circle. Therefore, all solutions to the homogeneous relation converge to zero at a rate exponential in j and y j = o (1) and ∑ ∞ j =0 y 2 j = O α,β,p (1) . Then we can bound the Frobenious norm of the matrix B | p | α,β as:

We use the following lower bound for the sum of w j :

Combining these bounds we can upper bound the Frobenious norm of the matrix B | p | α,β the following way:

Simplifying for the leading terms in asymptotics, we have:

Sensitivity of C | p | β . For the b -min-separation participation sensitivity we have the following bound for any p ≤ b :

Combining sensitivity with the upper bound for the Frobenious norm we obtain:

F.8 Proof of Theorem 7 for Square Root Factorization

Theorem 7 (Approximation error of Square Root Factorization) . Let A α,β ∈ R n × n be the workload matrix (9) of SGD with momentum 0 ≤ β < 1 and weight decay 0 < α ≤ 1 , with α > β . Let A α,β = C 2 α,β be its square root factorization . Then, for any b ∈ { 1 , . . . , n } and k = n b it holds:

Proof. We prove the case without weight decay ( α = 1 ) and with weight decay ( α < 1 ) separately. Case 1) no weight decay ( α = 1 ).

We start by bounding the b -min-separation sensitivity:

Consider a scalar product for a general pair of indices, j > i :

Using the bounds on c k (6) for α = 1 we can lower and upper bound this sum by:

We can compute the indefinite integral explicitly:

for F ( a ) = 2 log ( √ 1 a +1+ √ 1 a ) . In combination, we obtain the upper and lower bound for (126):

Now we are ready to bound the sensitivity of the matrix C 1 ,β :

and, analogously

Firstly, using ( k e ) k ≤ k ! ≤ k k , we bound the sum of the logarithms:

To upper bound the last term in sensitivity lower bound (134), we use the auxiliary inequality f ( a ) = 2 log (√ 1 a +1+ √ 1 a ) ≤ 4 √ a to derive:

To bound the final term we establish the following inequalities for f ( a ) :

Then we can bound the first double sum in (134) as

To bound the term ∑ 0 ≤ i<j ≤ k -1 log ( k -i j -i ) we use the following identities:

Now, using that x log x is a monotonically increasing function,

As a lower bound, we obtain

Therefore, combining the upper bound (146) and the lower bound (149) yields

Combining all three terms together we obtain the following bounds for the squared sensitivity (134):

Now, we recall the bounds for the Frobenius norm of the matrix C 1 ,β 7 and (96):

With the auxiliary inequality √ a 2 + √ b 2 ≤ √ a + b ≤ √ a + √ b and combining (155) and the bounds on Frobenius norm (7) and (96) we get that:

Making the lower bound well-defined requires n ≥ 45 , otherwise one can simply take E ( C 1 ,β , C 1 ,β ) ≥ 1 . As a final step, we combine both inequalities in the following asymptotic statement:

which concludes the proof of the case without weight decay.

Case 2) with weight decay ( α < 1 ).

As above, we first express the b -min-separation sensitivity of the matrix C α,β in terms of inner products,

and then consider a scalar product for a general pair of indexes j > i :

Now, we use the bounds on c t from Lemma 6 for α < 1 , to upper and lower bound this sum with the following expression, where γ = β α :

We substitute these bounds into Equation (161) to obtain the following upper bound for sensitivity of matrix C α,β :

E

(

C

α,β

, C

α,β

)

k

where the second inequality is due to Equation (95).

A lower bound for the sensitivity follows directly from Lemma 3:

The Frobenius norm of the matrix C α,β is the same as that for one round participation; thus, we could reuse Inequalities (95):

By merging the bounds for sensitivity and Frobenius norm, we derive the following bounds for error:

√

[

1

-

2

log

1

-

2

+

2

α

2

-

(

α

β

)

1

α

(1

-

γ

)

(1

α

)(1

b

2

b

√

(170)

The combination of these results yields the following asymptotic statement:

which concludes the proof.

F.9 Proof of Theorem 8

Theorem 8. Assume the setting of Theorem 6. Then, for any factorization A α,β = BC with C glyph[latticetop] C ≥ 0 , the approximation error fulfills

Proof. Let A α,β = BC be any factorization with CC glyph[latticetop] ≥ 0 . From Lemma 3 it follows that

where ‖ · ‖ ∗ denotes the nuclear norm , and the last inequality follows from its variational form, ‖ M ‖ ∗ = min { X,Y : XY T = M } ‖ X ‖ F ‖ Y ‖ F . The statement of the Theorem follows by inserting the corresponding bounds on ‖ A α,β ‖ ∗ from Lemma 8.

F.10 Proof of Theorem 9

Theorem 9. Assume the setting of Theorem 6. Then, the baseline factorizations A α,β = A α,β · Id and A α,β = Id · A α,β fulfill

]

-

α

)

b

Proof. Case 1) A α,β = BC with B = A α,β and C = Id . It is easy to check that sens k,b ( C ) = √ k , so E ( B,C ) = √ k n ‖ A α,β ‖ F . Because A α, 0 ≤ A α,β ≤ 1 α -β A α, 0 componentwise, we have for α = 1 ,

which implies the corresponding statement of the theorem. For 0 < α < 1 , we use that A 1 , 0 > Id componentwise, so ‖ A α, 0 ‖ 2 F > ‖ Id ‖ 2 F = n , which conclude the proof of this case.

Again, we use the fact that A α, 0 ≤ A α,β ≤ 1 α -β A α, 0 . Now A α, 0 fulfills the conditions of Theorem 2, so from Equation (10) we know

We first study (176) for α = 1 . Then, from the explicit structure of A 1 , 0 = LDToep(1 , 1 , . . . , 1) one sees that the vectors inside the norm have a block structure

for i = 1 , . . . , k , and v ′ =    k . . . k    ∈ R n -bk , appears only if k < n b . Now we check

which concludes the proof of this case. For α < 1 , A α,β ≥ Id componentwise readily implies

which implies the statement of the theorem.

because bk ≥ n . Consequently

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Our abstract and introduction clearly states our claims, the contributions made in the paper, the main assumptions and limitations.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitation of the approximately optimal matrix factorization methods. We emphasize that the bounds for the banded square root factorization are asymptotic and should be perceived together with the numerical experiments.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All lemmas and theorem state the necessary assumptions. We provide proof sketches in the main text and proofs for all results in the appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: For the factorization error comparison, which we consider to be the main experimental results, we have explicit descriptions of the procedures together with the code included in the appendix. For the model training, we adapted the public code from prior work [Choquette-Choo et al., 2023c], and we describe the changes in the appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide the code for the approximately optimal matrix factorization. The code for model training is adapted from the publicly accessible code, see above.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so No is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We use the same setting as Choquette-Choo et al. [2023c] and report to have the same hyperparameters for the training.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: we report error bars and indicate that experiments were performed multiple times to ensure reliability and statistical significance.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer:[Yes]

Justification: Our experiments run on commodity hardware.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have acquainted with the guidelines and preserve the anonymity of the authors in the paper.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: Our work is about model training while preserving data privacy. We mention this in the manuscript.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have cited all the papers which have results and code used in our paper.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets

has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\n2. 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.\n\nWEAKNESSES:\n1. 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.\n\n2. 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?\n\n[1] https://github.com/apple/pfl-research/blob/2a9d65bd66dc89ef80a24e1be0d28446a0422133/pfl/privacy/ftrl_mechanism.py#L140 \n\n[2] https://github.com/dpcomp-org/hdmm/blob/7a5079a7d4f1a06b0be78019adadf83c538d0514/src/hdmm/templates.py#L391\n\nQUESTIONS:\n1. 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.\n\n2. In equation (6), b should be p, right?\n\nLIMITATIONS:\n1. 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.\n\n2. 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* The scope, contributions, and key results are clearly written.\n* 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.\n* Authors demonstrated deep understanding of the area, and made a compelling story with a strong mix of theoretical derivations and empirical observations.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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.\n\n* 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?\n>* 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).\n>* There is an incorrect reference to Eq 10 in the experiments\n>* Typo: n_{n-1}\n\n* 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?\n\n* 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.  \n>* 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.\n\nLIMITATIONS:\nThe authors have discussed some limitations, but perhaps adding some of the things I mentioned in this review (or addressing them directly) would be good.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses optimal matrix factorization mechanism for differential privacy, focussing on stochastic gradient descent (SGD) optimization.\nIt introduces the banded squared root (BSR) factorization and provide formal lower and upper bounds to the approximation error.\nThe BSR achieves competitive formal approximation error compared to state-of-the-art methods and demonstrates practical utility in real-world training tasks. \nThe 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.\n\nSTRENGTHS:\n- 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.\n- The incorporation of the SDG workload matrix in this line of research is elegant.\n- The banded squared root (BSR) factorization is computationally efficient and scalable to large workload matrices, making it a valuable tool.\n- The proof sketches are intuitive and convincing, although I haven't thoroughly reviewed the technical proofs in the appendix.\n- A notable advantage of BSR is its independence from specific training objectives.\n- BSR achieves competitive approximation error, on par with state-of-the-art methods in differentially private training.\n- The approach demonstrates practical utility by maintaining high training accuracy in real-world training tasks.\n- The authors provide theoretical guarantees for both single and repeated participation scenarios, adding to the method's credibility.\n- The theoretical explanations are (for the most part) well-written and easy to follow.\n- The paper includes extensive technical supplementary material.\n- The inclusion of useful, copy-pasteable Python code listings is interesting.\n- The authors show respect for existing work by referencing relevant StackOverflow answers and software packages.\n\nWEAKNESSES:\n- The experimental evaluation is limited to a small set of real-world and synthetic experiments.\n- The range of synthetic datasets used in the experiments is restricted, which may not be representative of diverse scenarios.\n- The data generation process is unclear, and only homogeneous partitions are considered, with no variation in data distributions.\n- The real-world experiments are limited to a single dataset, CIFAR-10, which is insufficient to draw general conclusions.\n- The post-processing step (line 275) lacks clear explanation, leaving it unclear whether it addresses a principled problem or a numerical issue.\n- The experiments do not provide a comparison with non-private approaches, which makes it difficult to understand the practical limitations of the method.\n- The evaluation is limited to a single real-world dataset, CIFAR-10, and a simple ConvNet model.\n- The experiments do not discuss potential practical limitations of BSR, which may impact its applicability.\n- The comparative analysis is limited and does not include a wide range of existing differential privacy methods.\n- The formal proof of Theorem 2 (Appendix D.2) employs unclear and undefined notation (Π), which makes it difficult to follow.\n- Neither the proof sketch nor the formal proof are straightforwardly understandable, which may hinder comprehension.\n\nMinor:\n\n- The tone of \"supposedly optimal\" (line 254) could be perceived as inappropriate and may benefit from rephrasing.\n- The proof sketch of Theorem 2 refers to the wrong appendix (Appendix D.3 instead of Appendix D.2, presumably).\n\nQUESTIONS:\n./.\n\nLIMITATIONS:\n./.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n* 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.\n* While DP-MF is a relatively niche topic, the paper does a good job slowly introducing the problem, their approach, and their theoretical results.\n\nWEAKNESSES:\n* 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.\n* 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.\n\nQUESTIONS:\n* 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?\n* 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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}