SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Algorithm 7 Sequence of queries that bounds privacy of Alg. 6
Inputs: Dataset  $\tilde{D} = \{d_1,d_2,\dots ,d_{\tilde{m}}\}$  sampling distribution  $\mathcal{S}$  over  $(2^{[\tilde{m}]})^k$  Sample  $(S_{1},S_{2},\ldots ,S_{k})\sim \mathcal{S}$    
for  $i = 1,2,\ldots ,k$  do  $\tilde{D}_i\gets \{d_j:j\in S_i\}$  Perform (adaptively chosen) sensitivity  $\Delta$  query on  $\tilde{D}_i$  with noise  $\mathcal{N}(0,\sigma^2)$
Figure 10: A visualization of how we can decompose a banded matrix mechanism into independent queries on  $D_{j}$  (as in Alg. 7) under our sampling scheme.
F.2 General Amplification Statement and Proof
Given the sampling scheme and query sequence, we can now state our general amplification statement: Theorem 5. Suppose  $\mathbf{C}$  is  $b$ -banded and lower triangular, and the examples participating in each step are chosen according to Alg. 6 with a given choice of  $\mathcal{S}$ . Then BANDMF satisfies any standard DP guarantee satisfied by Alg. 7 in App. F.1 with  $\kappa = \max_{i\in [n]}\| \mathbf{C}\mathbf{e}_i\| _2 = \max_{i\in [n]}\sqrt{\mathbf{X}_{i,i}}$  and the same choice of  $\mathcal{S}$ .
Proof. Consider two datasets  $D, D'$  that differ by an example contained in the partition subset  $D_j$ . We argue about the privacy of  $\mathbf{Cx} + \mathbf{z}$ . For simplicity we assume  $j$  is such that  $(k - 1)b + j \leq n$ ; elements in  $D_j$  such that  $j$  does not satisfy this condition can potentially participate  $k - 1$  times instead of  $k$ , and in turn the privacy guarantee we can prove for these elements can only be stronger.
Since  $\mathbf{C}$  is  $b$ -banded, we can partition the rows of  $\mathbf{C}$  into  $k + 1$  subsets
$$
R _ {j}, R _ {b + j}, R _ {2 b + j} \dots R _ {(k - 1) b + j}, R _ {\emptyset},
$$
where  $R_{j}$  (resp.  $R_{b + j}, R_{2b + j} \ldots R_{(k - 1)b + j}$ ) denotes the set of rows in  $\mathbf{C}$  for which the  $j$ th entry is non-zero, and  $R_{\emptyset} = [n] \setminus (R_{j} \cup R_{b + j} \cup \ldots)$ , i.e.,  $R_{\emptyset}$  are the rows not included in any of these sets, i.e., rows of  $\mathbf{C}$  where entries  $j, b + j, \ldots$  are all zero. The fact that  $\mathbf{C}$  is lower-triangular and  $b$ -banded ensures that these subsets do not overlap, i.e., this is a valid partition as can be observed in Fig. 4.
Let  $\mathbf{C}_R$  denote  $\mathbf{C}$  restricted to the set of rows in  $R$ . From the perspective of an adversary distinguishing  $D$  from  $D'$ , each row of  $(\mathbf{Cx} + \mathbf{z})_{R_\emptyset} = \mathbf{C}_{R_\emptyset} \mathbf{x} + \mathbf{z}_{R_\emptyset}$  has a distribution independent of whether  $D$  or  $D'$  was used. So it suffices to give privacy guarantees for outputting only  $(\mathbf{Cx} + \mathbf{z})_{R_j}, (\mathbf{Cx} + \mathbf{z})_{R_{b+j}}, \ldots, (\mathbf{Cx} + \mathbf{z})_{R_{(k-1)b+j}}$ .
We can decompose rows  $R_{j}$  of  $\mathbf{Cx} + \mathbf{z}$  as follows:
$$
\left(\mathbf {C} \mathbf {x} + \mathbf {z}\right) _ {R _ {j}} = \mathbf {C} _ {R _ {j}} \mathbf {x} + \mathbf {z} _ {R _ {j}} = \mathbf {C} _ {R _ {j}} \mathbf {x} _ {j} + \mathbf {C} _ {R _ {j}} \mathbf {x} _ {- j} + \mathbf {z} _ {R _ {j}}. \tag {10}
$$
Where  $\mathbf{x}_j$  denotes  $\mathbf{x}$  with all rows except  $j$  zeroed out, and  $\mathbf{x}_{-j}$  denotes  $\mathbf{x} - \mathbf{x}_j$ , i.e.,  $\mathbf{x}$  with row  $j$  zeroed out. By the  $b$ -bonded property of  $\mathbf{C}$ ,  $\mathbf{C}_{R_j}\mathbf{x}_{-j}$  has 0 sensitivity to the examples in  $D\setminus D_{j}$ . Then,
${}^{9}$  Standard DP includes  $\epsilon$  -DP,  $\left( {\epsilon ,\delta }\right)$  -DP,Renyi DP,zCDP,and Gaussian DP.
by Eq. (10), for  $i \in R_j$ , we observe that the  $i$ th row of  $(\mathbf{Cx} + \mathbf{z})_{R_j}$  corresponds to an (adaptive) query made with  $\ell_2$ -sensitivity  $\mathbf{e}_i^\top \mathbf{C}\mathbf{e}_j$  to the examples used in step  $j$ , i.e., those given by  $D_j$  and  $S_j$ , and noise  $N(0, \sigma^2)^d$ . So  $(\mathbf{Cx} + \mathbf{z})_{R_j}$  corresponds to a sequence of adaptive queries on the examples used in step  $j$ , and answering this sequence of queries satisfies any standard privacy guarantee satisfied by answering a single (scalar, adaptively chosen) query with sensitivity  $\| \mathbf{C}\mathbf{e}_j\|_2$  to the example chosen in step  $j$  and noise  $N(0, \sigma^2)$  by Claim D.1 in [17].
The same logic applies to each of  $(\mathbf{Cx} + \mathbf{z})_{R_{b + j}},\ldots ,(\mathbf{Cx} + \mathbf{z})_{R_{(k - 1)b + j}}$ . Putting it all together and taking a max over the sensitivity of the individual queries, releasing  $\mathbf{Cx} + \mathbf{z}$  satisfies any standard privacy guarantee satisfied by answering  $k$  adaptively chosen queries, with sensitivity  $\max_{i\in [n]}\| \mathbf{Ce}_i\| _2$  to the examples used in steps  $j,b + j,\ldots ,(k - 1)b + j$  respectively. This is exactly Alg. 7 with the specified choice of  $\Delta ,S$
F.3 Corollaries of Thm. 5
We give here several corollaries of Thm. 5 that are of interest.
Equivalence to DP-SGD: Note that when  $b = 1$ , the partition contains a single subset, i.e., is the entire dataset. In particular, in this setting Thm. 5 recovers the privacy guarantees of amplified DP-SGD under any amplification scheme, e.g. including the ones discussed below.
Amplification via sampling: To recover Thm. 4 from Thm. 5, we take the distribution over  $2^{[\tilde{m} ]}$  corresponding to the uniform distribution over subsets of size  $B$ , and let  $\mathcal{S}$  be the product of this distribution with itself  $k$  times. This is equivalent to the following: in step  $i$ , we include each element of  $D_{i\pmod{b}}$  independently with probability  $q$ . For this choice of  $\mathcal{S}$ , Alg. 6 reduces to Alg. 2 and Thm. 4. We next make the amplified privacy guarantee explicit in terms of the dp_accounting Python library [18]. Given  $n,m,b$  and a target per-step batch size  $B$ , we could write a dp_accounting.DpEvent capturing the privacy guarantees of the matrix factorization mechanism as follows:
Example F.1.
gaussian_event  $=$  dp_accounting.GaussianDpEvent(noise-multiplier)   
 $\mathrm{q} = \mathrm{B}$  /math.floor(n/b)   
sampled_event  $=$  dp_accounting.PoissonSampledDpEvent( q, gaussian_event   
）   
composed_event  $=$  dp_accounting.SelfComposedDpEvent( sampled_event，math.ceil(m/b)   
）
Example F.2. To give an example of the amplification guarantee, for simplicity assume  $n / b, m / b$  are integer. If all column norms in  $\mathbf{C}$  are 1, each row of  $\mathbf{x}$  has sensitivity 1, and each entry of  $\mathbf{z}$  has standard deviation  $\sigma$ , then outputting  $\mathbf{Cx} + \mathbf{z}$  satisfies  $(\alpha, \frac{\alpha n}{2\sigma^2b})$ -RDP.
Using Theorem 11 of [46] and Thm. 4, for appropriate choice of  $\alpha$  and  $q$ , this improves to  $(\alpha, q^2 \cdot \frac{2\alpha n}{\sigma^2 b})$ -RDP with amplification by sampling. In particular, if we have a target per-step batch size  $B$ , then we should choose  $q = \frac{Bb}{m}$ , and if this choice of  $q$  satisfies the conditions in [46] plugging this in gives  $(\alpha, \frac{2\alpha B^2 b n}{\sigma^2 m^2})$ -RDP. Notice that  $b = 1$  recovers the privacy guarantees of DP-SGD with sampling probability  $B / m$ , and this privacy guarantee weakens as  $b$  increases.
Amplification via shuffling: Fix a per-step batch size  $B$ . Then, suppose we shuffle the list of examples, and cyclically iterate over batches of size  $B$  in this list as the sets of examples to use in each step of matrix factorization. That is, we shuffle  $D$  into an ordered list  $d_{1}, d_{2}, \ldots$ , and in step  $i$  use examples  $d_{(i-1)B+1} (\bmod m), d_{(i-1)B+2} (\bmod m), \ldots, d_{iB} (\bmod m)$ .
For simplicity let's consider the case where  $m / (Bb)$  is integer. In particular, this means in this shuffling scheme, each example appears once every  $m / B$  steps, and for each of these steps  $i$ ,  $i \pmod{b}$  is the same. Then this shuffling scheme is equivalent to the following: First, rather than choose an arbitrary partition to apply Thm. 5, we choose a uniformly random partition into  $b$  subsets
of size  $m / b$ . Then, we choose  $S$  to be the distribution giving by shuffling  $[m / b]$  and then cyclically iterating over the shuffled list in batches of size  $B$ . Given this equivalence, we get the following:
Corollary F.1. Suppose the examples in matrix factorization are chosen by shuffling  $D$  and then iterating over batches of size  $B$ . If  $n / (Bb)$  is integer, then the matrix factorization mechanism satisfies any standard privacy guarantee satisfied by  $k$  adaptive scalar queries with sensitivity  $\max_{i\in [n]}\| \mathbf{C}\mathbf{e}_i\| _2$  and noise  $N(0,\sigma^2)$ , with the examples in each query given by shuffling a dataset of size  $m / b$  and cyclically iterating over this list in batches of size  $B$ .
Example F.3. Consider the simplified case where  $m = n$ , we choose a random permutation  $\pi$ , and in step  $i$  query example  $d_{\pi(i)}$ . In this case, if all the column norms of  $\mathbf{C}$  are 1,  $\mathbf{x}$ 's rows have sensitivity 1, and  $\mathbf{z}$ 's entries have standard deviation  $\sigma = \mathcal{O}\left(\frac{\sqrt{\ln(1 / \delta)}}{\epsilon}\right)$ , we get that  $\mathbf{Cx} + \mathbf{z}$  satisfies  $(\epsilon, \delta)$ -DP. With e.g., the amplification for shuffled  $(\epsilon, \delta)$ -DP mechanisms given by Theorem 5.1 of [6] and Cor. F.1, if  $\epsilon$  is a constant, we instead get that  $\mathbf{Cx} + \mathbf{z}$  satisfies  $\left(\epsilon \cdot \mathcal{O}\left(\sqrt{\frac{b\log(1 / \delta)}{n}}\right), \delta \cdot \mathcal{O}\left(\frac{n\ln(1 / \delta)}{b}\right)\right)$ -DP.
F.4 Optimizing the number of bands
Let  $\sigma_{\epsilon, \delta}(b)$  be the required Gaussian noise magnitude for a  $b$ -banded MF run for  $n$  iterations using e.g. Alg. 2 to achieve  $(\epsilon, \delta)$ -DP with per-step batch size  $B$ . Then, the expected total squared error introduced while achieving  $(\epsilon, \delta)$ -DP with amplification can be calculated as
$$
\sigma_ {\epsilon , \delta} (b) ^ {2} \mathcal {L} \left(\mathbf {A C} _ {b} ^ {- 1}, \mathbf {C} _ {b}\right)
$$
where  $\mathbf{C}_b$  is a  $b$ -banded lower triangular matrix optimized via Problem 2. Generally, smaller values of  $b$  will allow for more amplification, and hence a smaller  $\sigma$ ; however, this introduces a stronger set of constraints on the optimization problem, likely increasing the  $\mathcal{L}$  term. Hence, the choice of  $b$  should be optimized. Fortunately,  $\sigma_{\epsilon,\delta}(\cdot)$  can be computed efficiently: Thm. 4 implies a procedure to compute  $\epsilon$  given  $\sigma, \delta, b$ , and then one can use binary search<sup>10</sup> and this procedure to find the  $\sigma$  giving a desired  $\epsilon$ . In addition, one can pre-compute the optimal matrices  $\mathbf{C}_b$  for different numbers of bands. The search can be restricted to  $b \in \{1,2,\dots,\frac{m}{B},n\}$  since for  $b = \frac{m}{B}$  we have  $|D_j| = B$ , i.e. Alg. 2 and Thm. 4 provide no privacy amplification.
Unlike the un-amplified version of MF, now the best factorization depends on the privacy parameters  $(\epsilon, \delta)$ . The benefits of amplification are generally stronger for small  $\epsilon$ : For example, amplification by sampling with probability  $p$  roughly improves  $\epsilon$  to  $\log(1 + p(e^{\epsilon} - 1))$  (see e.g. Section 6 of [55]), which is approximately  $p\epsilon$  for  $\epsilon \leq 1$ , and approximately  $\epsilon - \log(1/p)$  if  $e^{\epsilon} \gg 1/p$ . Hence, with smaller values of  $\epsilon$ , we expect the benefits of amplification to outweigh the benefits of correlated noise, in which case  $b = 1$  will be optimal. With larger values of  $\epsilon$ , we expect the benefits of correlated noise to outweigh the benefits of amplification, and in this regime  $b = n$  will be optimal. For moderate values of  $\epsilon$ , we expect the optimal  $b$  to be somewhere in the middle.
F.5 Applying BANDMF with Privacy Amplification
Consider the setup of Sec. 5 with our CIFAR10 setting described fully in App. H. As mentioned in that section, we use the convention that our privacy analysis will assume Poisson sampling, even though we are using passes over a shuffled dataset. We have  $m = 50,000$  and train for  $k = 20$  epochs with a batch size  $B = 500$ . Choquette-Choo et al. [15] lets us bound the sensitivity and optimize matrices for this setting, however, without privacy amplification. Suppose we choose  $\hat{b} = 100$ . Because  $m / \hat{b} = 500 = B$ , we get that the sampling probability is  $100\%$  (see Example F.1) and thus get no benefits from amplification. For all  $\hat{b} \in [1,100)$  we get amplification benefits which can be seen intuitively as follows.
If  $\hat{b} = 2$ , we get that there are two partitions  $D_{1}, D_{2}$ . Then or first event will be the simultaneous release of  $\mathbf{x}_1, \mathbf{x}_2$ , our second of  $\mathbf{x}_3, \mathbf{x}_4$ , and so on. Because each partition is of size  $|D_j| = m / \hat{b} = 25,000$  and  $B = 500$ , we have a sampling probability  $q = 2\%$ . Given our parameters, we also have  $d = k \cdot m / B = 2,000$ , and so we must compose  $d / \hat{b} = 1,000$  events (as seen in Example F.1). Because in this setting each event is the batch release of  $\hat{b}$  steps of  $\mathbf{x}$ , where each example participates at most
10See for example the calibrate_dp_mechanism function of DP Team [18].
once on each release, observe that we need only normalize the sensitivity of this mechanism under  $(k = 1, b = \hat{b} = 2)$ -participation. Generally, as  $\hat{b}$  increases we have a higher sampling probability, but fewer events to compose, and vice versa. As can be seen, when  $\hat{b} = 1$ , this reduces to the standard accounting for DP-SGD. It can also be seen that we desire each  $D_{j}$  to be a non-overlapping partition of  $D$ , as otherwise, a single example may participate multiple times in the same event (and thus have a higher sampling probability).
G Additional RMSE Experiment Details
G.1 Optimal Number of Bands
In this section, we provide supplementary data surrounding the RMSE experiments in Fig. 5. Table 3 shows the optimal number of bands for each  $(\epsilon, k)$  pair considered in the RMSE experiments. It shows the general trend that as  $\epsilon$  decreases, or  $k$  increases, the optimal number of bands decreases.
G.2 Explaining many-epoch setting
Observe in Fig. 5 that as BANDMF and MULTI-EPOCH MF incur similar RMSE as DP-SGD when the number of epochs increases.
This phenomenon occurs because as the number of epochs changes, so does the sensitivity of the  $\mathbf{X}$  matrix, and subsequently the geometry of the optimization problem. By Eq. (5), we see that sensitivity can be calculated as a sum of absolute values of entries of  $\mathbf{X}$  corresponding to iterations where a single user might participate. As the number of epochs increases, the number of entries of  $\mathbf{X}$  that we have to sum up also increases. It turns out that under the constraint of constant sensitivity, it is better to put more weight on the diagonal of  $\mathbf{X}$  than the off-diagonal. This is an observation we have made by solving this optimization problem numerically in a number of settings. Hence, as the number of epochs increases,  $\mathbf{X}$  becomes more and more diagonally dominant, making the mechanism closer to DP-SGD ( $\mathbf{X} = \mathbf{Identity}$ ).
ε/k	1	2	4	8	16	32	64	128	256	512	1024
0.03125	2	2	1	1	1	1	1	1	1	1	1
0.0625	4	2	1	1	1	1	1	1	1	1	1
0.125	8	4	2	1	1	1	1	1	1	1	1
0.25	8	4	4	2	1	1	1	1	1	1	1
0.5	16	8	4	4	2	1	1	1	1	1	1
1.0	32	16	8	4	2	2	1	1	1	1	1
2.0	64	32	16	8	4	2	2	1	1	1	1
4.0	128	64	32	16	8	4	2	2	1	1	1
8.0	1024	512	256	32	16	8	4	2	2	1	1
16.0	1024	512	256	128	64	32	8	4	4	2	1
Table 3: Optimal number of bands for each  $\left( {\epsilon ,k}\right)$  pair,when  $n = {1024}$  and  $\delta  = {10}^{-6}$  .
H Additional CIFAR-10 Experiment Details
H.1 Setup and Tuning
We tune all jobs on a learning rate grid of coefficients in  $\{1,2,5\}$  on powers in [-2, 3]. We find that no momentum works best for DP-SGD and momentum=0.95 works best for MF-DP-FTRL mechanisms on average in tuning; though initial tuning found that tuning momentum as well could lead to slightly better results at some  $\epsilon$  budgets, we found that a more refined grid of learning rates nearly always led to a fixed momentum being optimal, and so we fix this parameter. We also found that a learning rate cooldown to  $0.05\times$  the initial learning rate over the last 500 steps of training improved all runs and so we fix this parameter. All models trained for 20 epochs on CIFAR10 with a batch size of 500. We repeat each setting 12 times and show  $95\%$  bootstrapped confidence intervals.
H.2 Additional Figures
Figure 11: On CIFAR-10, BANDMF is at least as good as DP-SGD across all  $\epsilon$ , and often significantly better. BANDMF is better than the prior MF-DP-FTRL from Choquette-Choo et al. [15] up to  $\epsilon \approx 5$ . We compare the ratio of the total error (see Sec. 4) of BANDMF with either mechanism. Lower values indicate that BANDMF is better. The yellow markers indicate the best BANDMF mechanism that was better for that  $\epsilon$  budget if one existed. Unlike in Fig. 5(b), We only optimize the Band MF over  $\hat{b} \in [0, n/k]$  which leads to a regime around  $\epsilon > 5$  where the it performs worse than the Multi-epoch MF of Choquette-Choo et al. [15].
Figure 12: Our banded matrices consistently perform at least as well as the best prior method in each range of  $\epsilon$ . Around  $\epsilon \approx 1$ , we observe significant utility benefits from the banded mechanism around  $2 - 3$  percentage points over DP-SGD. We only optimize the Band MF over  $\hat{b} \in [0, n / k]$  which leads to a regime around  $\epsilon > 5$  where the it performs worse than the Multi-epoch MF of Choquette-Choo et al. [15];  $\hat{b} = n$  is equivalent to this approach modulo the sensitivity definition which we exclude to emphasize the regime we improve on. Empirical setup is in App. H.
I Additional StackOverflow Next-Word-Prediction Experiment Details
We follow the experimental setup for StackOverflow NWP from Denisov et al. [17] and Choquette-Choo et al. [15]. Except for SINGLE-EPOCH MF (which uses  $B = 167$  clients/round for 1 epoch), all privacy guarantees and accuracy results are for 6 epochs of training using  $B = 1000$  clients/round for 2052 rounds (also 1 epoch). The matrices used in these experiments are included in Table 2.
For computational efficiency in estimating model accuracy at a given privacy guarantee, we actually compute in simulation updates from only 100 clients/round, and scale the noise multiplier by a corresponding factor  $(\frac{100}{1000}$  for 6 epoch experiments,  $\frac{100}{167}$  for SINGLE-EPOCH MF). This approach has been used previously [35, 44], and we independently verified it has a negligible impact on the estimates of accuracy figures we report. Tables 4 and 5 include the unscaled noise multipliers  $\sigma$  for our experiments.
Figure 13: Correlation between optimal server learning rates  $\eta_{s}$  and the effective RMSE during training, see Eq. (11).
Optimizer and learning-rate tuning For all SO NWP experiments we use the FedSGDM optimizer [50]. This optimization approach first takes multiple local SGD steps (with learning rate 1.0 in our experiments) on the training data of each user in the batch (cohort) before clipping to  $\zeta = 1$ , summing, and passing the result  $\mathbf{x}_i$  into the DP mechanism which adds noise  $[\mathbf{C}^{\dagger}\mathbf{z}]_{[i,:]} \in \mathbb{R}^{d}$  on each iteration  $i$ . The resulting privatized sum is then divided by the batch size  $B$  and passed to the "server" (post-aggregation) optimizer, in our case SGDM with momentum parameter  $\beta = 0.95$  and learning rate  $\eta_s$ . We find tuning  $\eta_s$  depending on the noise level is critical. By using the computationally efficient approach mentioned above, we were able to conduct rigorous tuning over a learning rate grid of  $1.7^i$  for powers  $i$  in  $\{-9,\dots,4\}$ , estimating good initial guesses based on prior work. Table 6 gives the full set of results, and Fig. 14 shows convergence as a function of the number of rounds (iters).
Learning rate warmup and cooldown Denisov et al. [17] found learning rate cooldown was effective, and Choquette-Choo et al. [15] found that zeroing-out client updates with large  $\ell_{\infty}$  norms was critical to stability in early training. We find that additionally introducing a learning-rate warmup schedule reduces the need for this zeroing-out (though we still enable it), and generally decreases the variance in training results. Hence, all of our experiments (for all algorithms) using a linear learning rate warmup from  $0.05\eta_{s}$  to  $1.0\eta_{s}$  over the first  $15\%$  of rounds (309), and a linear decay from  $1.0\eta_{s}$  to  $0.05\eta_{s}$  over the last  $25\%$  of rounds (513).
Using RMSE to tune optimal server learning rates Fig. 13 plots the server learning rates  $\eta_{s}$  from Table 6 on the  $y$ -axis (with the optimal rates shown as larger symbols, and sub-optimal rates as small symbols, versus two different measures of the error for the DP mechanism on the  $x$ -axis: The left plot gives the effective prefix-sum RMSE (the objective we use for optimizing (banded) matrices C),
$$
(\text {M e c h a n i s m} \times \text {n o i s e - m u l t i p l i e r} / (\text {c l i e n t s - p e r - r o u n d}) = \sqrt {\mathcal {L} \left(\mathbf {S C} ^ {- 1} , \mathbf {C}\right) / n} \times \sigma / B, \tag {11}
$$
where  $\mathbf{S}$  is the prefix-sum workload (lower-triangular matrix of ones) and  $\sigma$  and  $B$  are as given in Table 4. The right plot uses the RMSE in error of individual gradients, computed by replacing the  $\mathcal{L}$  term in the above with  $\mathcal{L}(\mathbf{IC}^{-1},\mathbf{C})$  where we take the workload  $\mathbf{A}$  to be the identity matrix  $\mathbf{I}$  rather than the prefix sum matrix  $\mathbf{S}$ .
We see a strong linear correlation between the prefix-sum RMSE and optimal learning rate in the left plot; this does not hold for individual gradient errors (right plot). Based on this, we use the following linear regression to choose learning rates for the non-federated (amplified) SO NWP experiments (still rounding to the nearest  $1.7^{i}$  for consistency):
$$
\log (\eta_ {s}) = - 0. 9 5 \cdot \log (L _ {e}) - 4. 6 4
$$
This allowed us to estimate learning rates for the amplified experiments with a high degree of accuracy; Table 7 gives the final selected learning rates.
Figure 14: Convergence plots for all cross-device federated learning simulation experiments.