Learning Rate Scheduling with Matrix Factorization for Private Training

Privacy has become a major concern as machine learning systems are trained on sensitive data such as personal communications, financial transactions, and medical records. Beyond the risk of direct data exposure, models themselves may memorize and unintentionally reveal private information, creating serious ethical and security challenges. These concerns are especially pressing for production-level large language models trained on vast and heterogeneous datasets.

A widely studied approach to mitigating these risks is differential privacy (DP), which provides formal mathematical guarantees that the output of a learning algorithm does not reveal sensitive information about any individual training example [12]. In practice, DP is often achieved by injecting carefully calibrated noise into either the gradients, ensuring that an adversary cannot infer the presence or absence of a single data point with high confidence. More recently, large-scale efforts such as VaultGemma [30] have demonstrated that it is possible to train billion-parameter models with rigorous privacy guarantees, showing that DP can be integrated into state-of-the-art architectures without prohibitive utility loss.

To make model training differentially private, algorithms typically inject noise into the gradients to mask the contribution of any individual data point. The most common approach, DP-SGD, adds independent Gaussian noise at each update, which provides strong privacy guarantees but can significantly reduce accuracy [1]. Matrix factorization has emerged as a more general alternative that introduces correlations in the injected noise, enabling improved accuracy while preserving privacy [8, 9]. The method is also applied beyond centralized model training to applications such as federated learning [6, 32] as well as decentralized learning [5]. The approach has also seen practical adoption, with Google reporting its use for training production on-device language models in their 2024 blog post “Advances in private training for production on-device language models” [31].

Recent work has focused on making matrix factorization memory efficient [25, 2, 20, 26], and it has also been analyzed theoretically, mostly in the setting of Toeplitz workloads [13, 16, 18, 11]. However, existing utility analyses assume a constant learning rate. While Denisov, McMahan, Rush, Smith and Thakurta [10] introduced a non-Toeplitz workload with varying learning rates, its theoretical properties remain largely unexplored. In this work, we address this gap by studying matrix factorization under learning rate schedules.

Learning rate scheduling plays a critical role in the optimization of machine learning models, and a variety of strategies have been proposed in the literature. Popular approaches include cosine annealing [23] and cyclical learning rates [29], which adapt the step size during training to improve convergence. Another common technique is warm-starting [15], where models begin with a small learning rate that is gradually increased, as used in large-scale training setups [14]. In this work, we focus on learning rate decays available in PyTorch [27] such as exponential, linear, polynomial and cosine, which are widely used and can be easily applied in practice.

Learning rate scheduling can be particularly useful in private training when the number of iterations is limited. By accelerating convergence, it enables higher accuracy in settings such as warm-up training [22], private fine-tuning [24], and training under computational constraints. It has also been combined with matrix factorization as a form of learning rate cool-down [9, 8, 7], and was shown to provide improvements over fixed learning rates in [10], where the workload of our interest was originally introduced.

The most common way to train differentially private models is by using DP-SGD [1]. At each step we receive a gradient $g_i\in {ℝ}^d$, clip it to a fixed ${\mathrm{\ell }}_2$ norm $\zeta >0$, and add appropriately scaled independent Gaussian noise $z_i\sim 𝒩(0,I_d{\sigma }^2)$, where $\sigma$ depends on the target privacy level $(ϵ,\delta )$. The model is then updated as

where ${\eta }_i$ is the learning rate at step $i$.

This procedure can be improved by correlating the noise across iterations. To formalize this, we define a matrix $G\in {ℝ}^{n\times d}$ of stacked gradients, a matrix $\mathrm{\Theta }\in {ℝ}^{n\times d}$ of intermediate models, and a workload matrix $A\in {ℝ}^{n\times n}$ that encodes the training process such that $\mathrm{\Theta }=AG$. If we use a constant learning rate, this matrix, denoted $A_1$, is a lower triangular matrix of ones. For varying learning rates we instead use a matrix $A_{\chi }$, described later.

To ensure that the intermediate models are differentially private, we can apply a matrix factorization mechanism. Specifically, we factorize $A$ into two matrices $B$ and $C$, then compute $CG$, add Gaussian noise $Z\in 𝒩{(0,{\sigma }^2)}^{n\times d}$ to ensure differential privacy, and finally multiply the result by $B$ as post-processing:

which is equivalent to adding correlated Gaussian noise with covariance structure induced by $C^{-1}$ to the gradients.

The remaining question is: how much noise must be added, and does this procedure remain differentially private when the gradients are not known in advance but depend on the current model? The foundational work of [10] shows that the procedure is indeed differentially private, even when the gradients adaptively depend on the model, provided we add noise of scale

where $\zeta$ denotes the clipping norm, and ${\sigma }_{ϵ,\delta }$ is the noise multiplier of the standard Gaussian mechanism, which can be computed numerically [4]. The term $\text{sens}(C)$ represents the global sensitivity of the Gaussian mechanism for the product $CG$ when the row or rows corresponding to a single datapoint in $G$ change; it can be computed explicitly as ${\Vert C\Vert }_{1\to 2}$, the maximum column norm of $C$. The case of multi-participation (multi-epoch) is discussed in Section 4.2. Now we have all the steps to train the model with differential privacy as presented in Algorithm 1.

The choice of factorization $A=BC$ significantly impacts the quality of the private estimation. Following the work of [10, 16] we quantify the approximation quality by either the mean squared error (MeanSE) or the maximum expected squared error (MaxSE), which can be computed as

where $\parallel \cdot {\parallel }_{\mathrm{F}}$ denotes the Frobenius norm and $\parallel \cdot {\parallel }_{2\to \mathrm{\infty }}$ the maximum row ${\mathrm{\ell }}_2$-norm. These approximation errors are independent of $G$, and the term ${\sigma }_{ϵ,\delta }\zeta$ is independent of the matrix factorization. To isolate the contribution of the factorization $(B,C)$, we will use the notation $\text{MeanSE}(B,C)$, $\text{MaxSE}(B,C)$ assuming $\zeta ={\sigma }_{ϵ,\delta }=1$ in the theoretical analysis.

We now turn to the workload of stochastic gradient descent (SGD) with learning rate scheduling. Let ${\chi }_1,{\chi }_2,\mathrm{\dots },{\chi }_n$ be a sequence with $\min {\chi }_t=\beta >0$ and $\max {\chi }_t=1$, representing a learning rate scheduler such that the actual learning rate at time $t$ is ${\eta }_t=\eta {\chi }_t$. We assume that $\beta$ is reasonably separated from $1$, as the regime $\beta \to 1$ is not of interest since it nullifies the benefits of scheduling. The workload matrix of interest is then:

where $A_1$ is a prefix-sum matrix (lower triangular matrix of all ones) and $D$ is a diagonal matrix of learning rates, i.e., $D=\mathrm{diag}({\chi }_1,\mathrm{\dots },{\chi }_n)$. We will study the problem of optimal matrix factorization in MaxSE and MeanSE metrics for the matrix $A_{\chi }$ with the learning rate decays given in Table 1. We emphasize that ${\chi }_n=\beta$ for all the listed decays. For the experiments, we will also include the constant learning rate $({\chi }_k=1)$.

In this work, we prove general lower and upper bounds for the MaxSE and MeanSE errors. For the upper bound, we use a prefix-sum-based factorization given by $B=A_{\chi }{(A_1)}^{-1/2}$ and $C=A_1^{1/2}$, which has been shown to be nearly optimal up to the next asymptotic term for the prefix sum problem (${\chi }_k=1$) [17]. To further improve the bounds, we propose a learning-rate-aware factorization. To define it, let $A_{\chi }^{\text{Toep}}$ denote the Toeplitz matrix with ${\chi }_1,\mathrm{\dots },{\chi }_n$ on its subdiagonals:

We propose $C_{\chi }={(A_{\chi }^{\text{Toep}})}^{1/2}$ as a learning-rate-aware correlation matrix. To analyze its properties, we consider the exponentially decaying learning rate ${\chi }_t={\beta }^{\frac{t-1}{n-1}}={\alpha }^{t-1}$ with $\alpha ={\beta }^{\frac{1}{n-1}}$. In this setting, the correlation matrix can be computed explicitly as

where the coefficients are $r_j=\left|\left(\genfrac{}{}{0pt}{}{-1/2}{j}\right)\right|=\frac{1}{4^j}\left(\genfrac{}{}{0pt}{}{2j}{j}\right)$.

In this work, we derive upper and lower bounds on the MaxSE and MeanSE errors of the learning rate scheduling workload $A_{\chi }$ for a large class of learning rate schedulers ${\chi }_1,\mathrm{\dots },{\chi }_n$. In the following theorem, we prove an upper bound based on the prefix-sum factorization $A_{\chi }{A}_1^{-1/2}\times A_1^{1/2}$.

The conditions assumed in Theorem 1 are satisfied for all learning rate decays presented in Table 1, more formally:

Moreover, in this work we also prove general lower bounds for any learning rate schedules:

In particular, plugging in the exponential learning rate decay ${\chi }_k={\beta }^{\frac{k-1}{n-1}}$ yields the following upper and lower bounds.

Note that our results do not report the leading constant in the error bounds. The reason is simply that the workload matrices we study are harder to analyze than the much-studied lower-triangular matrix of all-ones, for which more fine-grained analysis has been performed. For comparison, even for the square-root factorization of the Toeplitz workload with exponential weight decay, the exact leading constant has not been determined [16, 19].

We further improve the upper bound by considering a learning-rate–aware factorization $C={(A_{\chi }^{\text{Toep}})}^{1/2}$, which can be computed explicitly for the exponential learning rate decay ${\chi }_k={\beta }^{\frac{k-1}{n-1}}={\alpha }^{k-1}$. This yields the factorization $A_{\chi }=B_{\alpha }\times C_{\alpha }$, where $C_{\alpha }$ is defined in equation (8), and $B_{\alpha }$ is obtained as $A_{\chi }{(C_{\alpha })}^{-1}$.

In Lemma 7 of [19], the sensitivity of the matrix $C_{\alpha }$ has been computed as:

We then bound both the maximum row norm and the Frobenius norm of $B_{\alpha }$, which leads to the following lemma.

This factorization achieves the optimal rate for the MaxSE error and, asymptotically, performs better than alternative factorizations for the MeanSE error.

We summarize the errors for the exponential learning rate decay in Table 2. In addition, we consider four alternative factorizations: the trivial factorizations $A_{\chi }\times I$ and $I\times A_{\chi }$, the prefix-sum–inspired factorization $A_1^{1/2}\times A_1^{1/2}D$, and the square-root factorization $A_{\chi }^{1/2}\times A_{\chi }^{1/2}$. The square-root factorization is highly nontrivial to obtain since the matrix is not Toeplitz; we defer its treatment to Section 4.1.

We then numerically compare the proposed factorizations in the single-epoch (single-participation) setting using the MaxSE and MeanSE metrics, as functions of the learning rate decay $\beta$ and the matrix size $n$ (see Figure 1 for exponential decay and Figure 5 in the Appendix for other learning rate decays). As an approximation of the actual optimal value for MeanSE, we use a dense factorization [10] implemented in the jax-privacy library [3]. On the plots, we refer to this approximation as “dense”. For MaxSE, it is computationally infeasible to compute the exact optimal value for large matrix sizes. Therefore, we rely on the lower bound derived in Theorem 3, which we denote on the plots as “lower bound”. We observe that our learning-rate-aware factorization outperforms the others in terms of MaxSE. However, for the proposed values of $n$ and $\beta$, it performs worse than the prefix sum based factorization in terms of MeanSE. To further investigate this, we plot the colormap of the gain over the prefix sum based approach (see Figure 1). In the blue regions, our method performs worse, while in the red regions it performs better. As can be seen, for any fixed $n$, sufficiently small values of $\beta$ lead to the learning-rate-aware factorization outperforming the prefix sum based approach, thereby numerically validating our theoretical findings.

We demonstrate the practical benefits of learning rate scheduling in Figure 3 on CIFAR-10 dataset. All experiments satisfy $(9,{10}^{-5})$-DP and use a 3-block CNN trained for 10 epochs with batch size 128 and clipping norm 1. For privacy accounting, we use Poisson subsampling with PLD accounting [21] for DP-SGD and amplification by Ball-in-Bins subsampling with Monte Carlo accounting [7] for all matrix mechanisms. Subfigure (a) shows validation accuracy across different initial learning rates, where exponential learning rate scheduling improves performance compared to DP-SGD with a fixed learning rate ($\beta =1$). Subfigure (b) reports test accuracy using the best learning rate chosen on the validation set. All factorizations benefit substantially from scheduling, and the learning-rate–aware factorization (denoted as BISR w/ LRS) achieves even further improvements. However, optimizing the factorization with respect to learning rate workload does not necessarily lead to additional gains: while RMSE can serve as a proxy for performance, it does not perfectly predict it. In practice, workload optimization increases the added noise per iteration, and this effect is not fully compensated during training due to the non-linearity introduced by large noise. This was also stated as an open problem in a recent survey on matrix factorization [28].

In Figure 4, we compare different learning rate schedulers with a constant one. We observe that learning rate scheduling improves accuracy for DP-SGD in all cases except for the polynomial decay with $\gamma =2$, which deteriorates performance. The other schedulers substantially improve the accuracy of BISR. Moreover, our proposed learning-rate-aware factorization (BISR w/ LRS) further improves the quality, with the largest improvement for linear LRS, making it a suitable factorization for high-performance private training.

Learning rate scheduling has been shown to improve convergence in both private and non-private machine learning. In this work, we combine learning rate scheduling with matrix factorization and propose a learning-rate-aware factorization, which in the case of exponential learning rate decay is theoretically shown to improve the error. Through numerical experiments using the MaxSE and MeanSE metrics, as well as CIFAR-10 model training, we demonstrate its benefits.

We have primarily studied learning rate decay, but similar techniques can be applied to warm-starting, where the learning rate is initially small and then gradually increased. Optimization-based approaches for matrix factorization are generally agnostic to the choice of learning rate scheduling, but adapting our learning-rate-aware factorization to this setting may pose extra challenges.