(Amplified) Banded Matrix Factorization: A unified approach to private training

Introduction

In datacenter applications, precise control of the sampling/shuffling of training data is possible, and so dp-sgd with privacy amplification is one of the most popular ways to train machine learning models with formal privacy guarantees. However, Choquette-Choo et al. recently demonstrated that a multi-epoch extension of the mf-dp-ftrl algorithm can outperform amplified dp-sgd in some settings, depending on the privacy and computational budget (typically larger budgets above $\epsilon\approx 2$ and a small number of training epochs). This leaves the state-of-the-art for centralized DP training in the unsatisfactory state where one must try both algorithms to be assured of the best performance.

In cross-device federated learning, devices choose when they are available to participate in training, and so precise sampling and shuffling is generally not possible (see Sec. 3 for more details). Motivated by these limitations which make amplified dp-sgd infeasible, Kairouz et al. developed the (tree-aggregation-based) dp-ftrl algorithm. Their dp-ftrl does not rely on (or benefit from) privacy amplification and instead adds carefully correlated noise to the gradients to boost utility. Denisov et al. proposed mf-dp-ftrl, replacing the tree-aggregation scheme of Kairouz et al. with a general matrix-factorization mechanism. By optimizing over this space to find mechanisms with optimal error, substantial performance improvements were possible. However, the work of Denisov et al. applies only to the single participation (single epoch) setting. Hence, for cross-device FL the state-of-the-art also requires considering multiple algorithms: tree-aggregation-based dp-ftrl when devices may participate more than one time, or mf-dp-ftrl when devices participate only once. Importantly, the extension of mf-dp-ftrl to multiple epochs of Choquette-Choo et al. only applies in the centralized setting, as it again requires precise control of the participation pattern.

In this work, we address the limitations of mf-dp-ftrl (mf for short) noted above, and show that it can, in fact, achieve across-the-board state-of-the-art performance in both settings across all $\epsilon$. To accomplish this, we define a family of banded MF mechanisms, shown in Fig. 4 (c.f. Fig. 8 of App. D for visualizations of other factorization structures). We summarize our main contributions below.

Here, the $(k,b)$-participation schema of Choquette-Choo et al. cannot be easily enforced. We propose a strict generalization, $b$-min-sep-participation, which can be practically enforced by FL infrastructure. We show how to efficiently and exactly bound the sensitivity for banded matrices in Thm. 2, allowing formal DP guarantees and the numerical optimization of optimal mechanisms (Sec. 4). These innovations lead to significant privacy-utility benefits in a production deployment (Fig. 6 of Sec. 6). Our work also generalizes the sensitivity calculations of Choquette-Choo et al. to provide a general upper-bound on $b$-min-sep-participation sensitivity (Thm. 3), which allows the matrices of Choquette-Choo et al. to be used in the FL setting, as well as removing the need to exactly bound $b$ before training (see Sec. 6 and App. K).

Contributions for centralized training

The existing privacy amplification analysis of dp-sgd does not allow for the correlated noise that is applied in mf-dp-ftrl. Our paper introduces a novel partitioning of the BandMF iterates into independent queries. This allows us to prove in Thm. 4 of Sec. 5 that banded matrices enjoy the benefits of privacy amplification, and show that dp-sgd is a special case, giving us the best of both algorithms. This enables us to always pareto-dominate DP-SGD, unlike Choquette-Choo et al. which only does so for large enough $\epsilon$ as observed in Fig. 1. Further, this allows us to improve on both baselines, between $1-4\%$-points. Informally:

Suppose we partition the dataset into $b$ equal-size subsets, and in step $i$ each example in the $i\pmod{b}$-th subset participates with probability $\frac{Bb}{m}$ where there are $m$ examples and the batch size is $B$. Then, a $\hat{b}$-banded $n$-iteration matrix mechanism with $\hat{b}\leq b$ satisfies the same privacy guarantees as answering $n/b$ queries on a dataset, where each element of the dataset is independently included in each query with probability $\frac{Bb}{m}$.

As an example of Thm. 1, consider doing $n=2,000$ iterations of dp-sgd on CIFAR-10, which has $m=50,000$ examples, using a minibatch of $B=500$ examples in each round. This has the same DP guarantees as answering $2,000$ queries using a subsampled Gaussian mechanism with sampling probability $p=500/50,000=0.01$. If we instead use, e.g., BandMF with $\hat{b}=b=10$, our suggested sampling scheme is the following: Partition CIFAR-10 into 10 subsets of $5,000$ examples each, $D_{1},D_{2},…D_{10}$. In rounds 1, 11, 21… we sample $500$ examples from $D_{1}$, in rounds 2, 12, 22… we sample $500$ examples from $D_{2}$, and so on. We sample from each $D_{i}$ a total of $2000/10=200$ times, and each time our sampling probability is $500/5000=0.1$ within the subset. So Theorem 1 shows $\hat{b}=10$ BandMF satisfies the same DP guarantees as answering $200$ queries with $p=0.1$. As a special case of Theorem 1, dp-sgd is simply mf with a suitable diagonal matrix with $\hat{b}=1$, and thus Thm. 1 recovers the privacy guarantees of dp-sgd with amplification by sampling. Empirically, we show that mf with amplification has privacy-utility tradeoffs that are no worse than dp-sgd for all $\epsilon$, and often significantly better as can be seen in Fig. 1.

Finally, we explore the computational tradeoffs of our approach. We find that banded matrices with $b$-min-sep-participation are equally efficient to optimize as those under $(k,b)$-participation but significantly reduce the memory and time complexity of the per-iteration noise generation from $\mathcal{O}(n)$ to a constant $\mathcal{O}(\hat{b})$ (where often total steps $n\gg d$). We will release all code with the final manuscript.

Related work

The matrix mechanism (mf mechanism or mf) has a rich history in offline, statistical queries , with many applications including to online PCA , estimating marginals , and top-k selection . Recently, this has been studied under the adaptive streaming setting, where privacy analysis must account for an adversary adaptively defining the inputs at each step . Denisov et al. showed a connection with the DP-FTRL algorithm of Kairouz et al. and with DP ML broadly; they showed that computing optimal MF significantly improves the privacy-utility-computation tradeoffs when making only a single pass (epoch) over the training data. Choquette-Choo et al. showed that MF achieves state-of-the-art results in DP ML by showing how to optimize MF under arbitrary passes over the training data. Henzinger and Upadhyay study the problem of DP-continual observation and the explicit factorization of the workload matrix $\mathbf{A}$ that minimizes the completely bounded norm, which is only off from the optimal by an additive constant. The connection between DP empirical risk minimization and DP online regret minimization has been studied for a long time. Asi et al. demonstrated that DP-FTRL style algorithms achieve the best known regret in certain classes of online learning problems (a.k.a. the realizable regime). An important question that still remains open is whether DP-FTRL style algorithms can obtain optimal population risk guarantees under DP .

Matrix Factorization, Sensitivity, and Efficient Implementations

where $\mathbf{z}$ is suitably calibrated noise to the sensitivity of the so-called ‘query matrix’ $\mathbf{C}$.

Multiple participations

Intuition for (anti-)correlated noise in mf-dp-ftrl

Fig. 2 compares dp-sgd and mf-dp-ftrl. To gain an intuition for why mf-dp-ftrl can perform better than dp-sgd, observe that vanilla SGD has iterates $\theta_{t}=\theta_{0}-\eta\sum_{i=1}^{t}\hat{\mathbf{x}}_{i}$, and hence when the noisy gradients $\hat{\mathbf{x}}_{i}$ are added, the $\mathbf{C}^{-1}_{[i,j]}\mathbf{z}_{[i,:]}$ terms in mf-dp-ftrl serve to cancel out some of the noise introduced on previous rounds. This reduces the total error in the final model (i.e., the prefix sums). However, this worsens the sensitivity of the mechanism to $>\!\!1$ (as it is for dp-sgd). This is because an adversary trying to learn $\mathbf{x}_{1}$ via $\hat{\mathbf{x}}_{1}$ can partially learn the value of $\mathbf{z}_{1}$ from $\hat{\mathbf{x}}_{2}$, whereas in dp-sgd $\hat{\mathbf{x}}_{1}$ and $\hat{\mathbf{x}}_{2}$ are uncorrelated. This tradeoff is what mf-dp-ftrl aims to minimize. More details on this intuition are in Sec. A.1.

Adjacency and participation schemas

Optimizing factorizations

Different factorizations $\mathbf{A}=\mathbf{B}\mathbf{C}$ can have very different performance in practice. Thus, in mf applications it is common to optimize over the space of factorizations, where the objective function is the expected total squared error on $\mathbf{A}$, given as $\mathcal{L}(\mathbf{B},\mathbf{C})=\operatorname{sens}_{\Pi}^{2}(\mathbf{C})\left\|\mathbf{B}\right\|_{F}^{2}$. We define the expected root-mean-squared error (RMSE) as $\sigma\sqrt{\mathcal{L}(\mathbf{B},\mathbf{C})/n}$, where $\sigma$ is the standard deviation of the Gaussian noise. We take $\sigma=1$ when simply comparing mechanisms, or (in Sec. 6), calibrate $\sigma$ to achieve specific $(\epsilon,\delta)$-DP guarantees.

To facilitate optimization, utilizing the fact that the optimal-for-squared-error decoder $\mathbf{B}$ is $\mathbf{A}\mathbf{C}^{\dagger}$ , we note $\mathcal{L}(\mathbf{B},\mathbf{C})=\mathcal{L}(\mathbf{A}\mathbf{C}^{\dagger},\mathbf{C})$. The expected total squared error is invariant to scaling $\mathbf{C}$ by a constant, and hence it is sufficient to optimize under a sensitivity $1$ constraint. Further expressing the sensitivity and error in terms of $\mathbf{X}=\mathbf{C}^{\top}\mathbf{C}$ (note $\mathbf{X}$ is unrelated to the data $\mathbf{x}$), we have

assuming $\operatorname{sens}_{\Pi}(\mathbf{C})=1$ and $\mathbf{A}$ is in the rowspace of $\mathbf{C}$. Thus, we arrive at:

The matrix factorization optimization problem is to solve the convex optimization

and then find $\mathbf{C}$ so that $\mathbf{C}^{\top}\mathbf{C}=\mathbf{X}$, e.g., via Cholesky decomposition.

A Participation Schema for FL and the Sensitivity of Banded Matrices

In cross-device FL, devices locally evaluate eligibility criteria to deterimine when they might participate in training , for example only checking-in to the coordinating server when they are plugged in, on unmetered wifi, and idle. This makes it practically difficult to enforce the $(k,b)$-participation of Choquette-Choo et al. , where devices are assumed to participate at the same relative position in each epoch: devices are unlikely to meet the eligibility criteria during the narrow windows of both step $i$ and $i+b$. Further, precise sampling cannot provide the same level of privacy amplification as in the centralized setting. Consider if 6500 devices are needed to complete a round . An extreme (but realistic depending on time of day) setting may have only 6500 devices meeting eligibility criteria. Thus, either the protocol proceed without any sampling/amplification or wait until more devices are available; neither are desirable. We avoid amplification in the cross-device setting and instead proceed by addressing the question: Can mf-dp-ftrl be extended to the cross-device federated learning setting with multiple client participations?

With $(k,b)$-participation difficult to enforce in practice, our first challenge is to define a new participation schema with several properties: (a) the sensitivity of any matrix mechanism under this query can be bounded; (b) this bound is tight over an expressive class of matrices; (c) this bound can be efficiently represented as a constraint in a mathematical program so as to be able to find a near-optimal factorization $\mathbf{A}=\mathbf{B}\mathbf{C}$. In Defn. 1, we propose $b$-min-sep-participation, a generalization of $(k,b)$-participation which can be practically enforced by cross-device FL systems, thus enabling us to leverage BandMF in this setting (see Sec. 6).While $b$-min-sep-participation is required for the cross-device FL setting, $(k,b)$-participation is preferred in centralized settings where it can be enforced and will also satisfy our BandMF amplification guarantees leading to the empirical results we show in Fig. 1. In $b$-min-sep-participation, the distance between any two participations is at least $b$, rather than exactly $b$ as in $(k,b)$-participation:

The $b$-min-sep-participation schema is given by

Observe this participation schema is easy for devices to enforce: each device remembers the last step $i$ in which it participated, and when it again becomes eligible, it checks in to the server, and participates in training as long as the current step is at least $i+b$; it does not need to check in during a narrow (and unknown to the device) time window for a specific step.

We now turn to computing sensitivity under $b$-min-sep-participation. For $(k,b)$-participation, $|\Pi_{(k,b)}|=b$, a fact Choquette-Choo et al. [15, Eq. 3] critically exploited when computing sensitivity via brute force computation of a maximum over the elements in $\Pi$. With $b$-min-sep-participation, we have $|\Pi_{b}|=\mathcal{O}(\exp(n))$, and hence any brute force approach which requires checking some value for all $\pi\in\Pi_{b}$ will be impractical. Following the formalism of [15, Section 2], a participation schema $\Pi$ (plus a specification of model dimension $d$) yields an expression for the sensitivity of the function $\mathbf{x}\mapsto\mathbf{C}\mathbf{x}$ assuming that the contributions of any given user to the data structure $\mathbf{x}$ are restricted to the rows in $\mathbf{x}$ indexed by some $\pi\in\Pi$. By Prop. E.1 of Sec. E.2, independent of model dimension $d$, we show sensitivity for any schema $\Pi$ may be bounded by

Eq. 5 highlights several subtleties in computing sensitivity. First, is the challenge presented by the exponentially large number of patterns in $\Pi_{b}$. Second is the question of tightness of the inequality in Eq. 5: how much are we losing by effectively ignoring any cancellation in the matrix $\mathbf{X}$?

Fortunately, banded matrices render Eq. 5 both exactly computable and tight (independent of dimension $d$), showing that $\Pi_{b}$ satisfies the requirements of (b) and (c) above. We say a (general) matrix $\mathbf{X}$ is $\hat{b}$-banded if for all $i,j\in[n]$, $|i-j|\geq\hat{b}$ implies $\mathbf{X}_{[i,j]}=0$. While this is off-by-one from the bandwidth ($\mathbf{X}$ has bandwidth $b-1$), our definition will be useful as it will be natural to match $\hat{b}$-banded matrices with $b$-min-separation. Further, for $\hat{b}$-banded lower-triangular matrices (which will play a central role), $\hat{b}$ intuitively gives the number of bands in the matrix.

For non-banded matrices, the right-hand side of Eq. 5 remains efficiently computable (but not easily expressible in a mathematical program), enabling us to provide nontrivial privacy guarantees for matrices which are not $b$ banded under $b$-min-sep, showing that $\Pi_{b}$ satisfies (a) as well. The key subroutine is Alg. 3, which gives an efficient dynamic program for solving linear optimization over $\Pi_{b}$. Define $\mathbf{u}(\pi)\in\{0,1\}^{n}$ by $\mathbf{u}(\pi)_{i}=1$ if $i\in\pi$ and 0 otherwise. Then, Alg. 3 solves

This is the key subroutine in Alg. 4 and Alg. 5. Proofs for Thm. 2 and Thm. 3 are deferred to Sec. E.2.

For an arbitrary (non-banded) $\mathbf{C}$, let $\mathbf{X}=\mathbf{C}^{\top}\mathbf{C}$ and $b,k^{\prime}$ as in Thm. 2. Then Alg. 4 of App. E upper-bounds $\operatorname{sens}(\mathbf{C})$ under schema $\Pi_{b}$ in polynomial time for any dimension $d$.

Optimizing Banded Matrices

To enjoy the benefits of banded matrices within the framework of mf-dp-ftrl, we need to design an algorithm that can efficiently optimize over the space of $\hat{b}$-banded $\mathbf{C}$ matrices. To solve this problem, we will work in the domain of $\mathbf{X}=\mathbf{C}^{\top}\mathbf{C}$, and utilize the following fact:

Utilizing Prop. 4.1, we can modify 1 by introducing the constraint $\mathbf{X}_{[i,j]}=0\>\text{ if }|i-j|\geq\hat{b}$. This additional linear constraint preserves convexity of the optimization problem, and makes the sensitivity calculation tractable as well. However, it is still not immediately obvious how to solve the optimization problem, since we need to run the dynamic program defined in Alg. 5 of App. E to compute sensitivity. For this reason, we impose the additional constraint that $\operatorname{diag}{(\mathbf{X})}=\bf 1$. This constraint, together with bandedness, ensures that the squared sensitivity is equal to $k$ for all $\mathbf{X}$ by Thm. 2. The final optimization problem we seek to solve is stated below:

The matrix factorization optimization problem for banded matrices is to solve

and then find $\mathbf{C}$ so that $\mathbf{C}^{\top}\mathbf{C}=\mathbf{X}$ via Prop. 4.1.

We would like to remark on the similarity between 2 and the single-participation version of 1. The two problems are identical modulo the bandedness constraint, which is an equality constraint on individual entries of $\mathbf{X}$. Therefore, existing primal-optimization based solvers for the single-participation matrix mechanism can be extended to optimize over this new space of matrices with little modification. Specifically, the only modification necessary is to initialize to an appropriately banded feasible $\mathbf{X}$ matrix, like $\mathbf{X}=\mathbf{I}$, and to post-process the gradient w.r.t $\mathbf{X}$ by setting $\frac{\partial L}{\partial\mathbf{X}_{[i,j]}}=0$ if $|i-j|\geq\hat{b}$ in each step. Since the equality constraints exactly specify individual entries of $\mathbf{X}$, 2 can be solved as an unconstrained optimization problem (over the remaining entries in $\mathbf{X}$), using any number of off-the-shelf unconstrained optimization algorithms.Note that the constraint that $\mathbf{X}$ is positive definite is typically handled implicitly by taking sufficiently small step sizes to avoid violating that constraint . As recommended by McKenna et al. , we use the LBFGS algorithm to solve this problem.

The constraint on $\operatorname{diag}{(\mathbf{X})}$ serves multiple purposes. First, $\operatorname{diag}{(\mathbf{X})}=\bf 1$ implies that $\|\mathbf{C}_{[:,i]}\|_{2}=1$ for all $i$, i.e., that $\mathbf{C}$ has equal column norms. This ensures that BandMF reduces to dp-sgd when $\hat{b}=1$, which is desirable. Second $\operatorname{diag}{(\mathbf{X})}=\bf 1$ simplifies the optimization problem greatly, as the sensitivity computation for both $(k,b)$-participation and $b$-min-sep are trivial and tight under this constraint (Thm. 2 Claim (1)). Third, imposing this constraint does not drastically change the search landscape, or cost much in terms of RMSE; see Table 2 for a comparison of matrices with and without this constraint, and Fig. 8 for a visualization. Fourth, this constraint allows us to solve a single optimization problem that is simultaneously tailored for $(k,b)$-participation and $b$-min-sep-participation. In Appendices B and C, we formulate an optimization problem without the $\operatorname{diag}{(\mathbf{X})}=\mathbf{1}$ constraint, discuss how we solve it, and compare matrices generated with and without this constraint empirically.

Amplification for Banded Matrix Mechanisms

In the centralized setting where we can control the participation patterns of individual examples, the privacy guarantees of BandMF can be amplified. We focus on amplification by sampling with fixed batch size in this section, but give a more general statement in App. F.

Existing privacy analysis of mf-dp-ftrl is based on the reduction to the batch release of the entire $\mathbf{C}\mathbf{x}+\mathbf{z}$ as a single Gaussian mechanism event so standard amplification techniques don’t directly apply. Instead, for each participation by an example, we consider the set of rows in $\mathbf{C}\mathbf{x}$ affected by this participation as a Gaussian mechanism (see the groups of rows in Fig. 4). Then as long as the sets of rows corresponding to different participations do not interact, which is ensured by the bandedness of $\mathbf{C}$, we can apply amplification to them separately.

Observe from Fig. 4 that the structure of BandMF guarantees that the set of rows of $\mathbf{C}\mathbf{x}$ which depend on each of $\mathbf{x}_{j},\mathbf{x}_{\hat{b}+j},\mathbf{x}_{2\hat{b}+j},\ldots$ are disjoint sets. Thus, we use the following sampling scheme for determining which examples participate in each step which is made formal as Alg. 2. Let $D_{1},D_{2},\ldots,D_{\hat{b}}$ be an arbitrary partition of $D$ into $\hat{b}$ indexed subsets of size $\check{m}:=\lfloor m/\hat{b}\rfloor$ (for simplicity, we discard extra examples so all $D_{j}$ have size exactly $\check{m}$). In steps $j,\hat{b}+j,2\hat{b}+j,\ldots$, we will only use examples in $D_{j}$. Hence, participation follows $(k,b)$-participation for $b=\hat{b}$;and in our case, $k=1$ always. because it is optimal for $(k,b)$-participation to have the number of bands $\hat{b}$ equal the min-seperation $b$, in the remainder of this section and the associated appendices we simply write $b$ instead of $\hat{b}$. Within each of these steps, we sample a size $B$ subset of $D_{j}$ uniformly at random to use in computing $\mathbf{x}$.

Roughly speaking, Thm. 4 below shows that if we use Alg. 2, BandMF satisfies any standard privacy guarantees satisfied by dp-sgd run for $k$ rounds, where in each round we sample $B$ examples from a dataset of size $\check{m}$. In other words, it is equivalent to running dp-sgd for $1/b$ times as many rounds, but with the sampling probability multiplied by $b$.

Suppose $\mathbf{C}$ is $b$-banded and lower triangular, and the examples participating in each step are chosen according to Alg. 2. Then BandMF satisfies any standard DP guaranteeStandard DP includes $\epsilon$-DP, $(\epsilon,\delta)$-DP, Rényi DP, zCDP, and Gaussian DP. satisfied by performing $k$ sensitivity-$\kappa$ queries on a dataset of size $\check{m}$ using the Gaussian mechanism, where each query is run on a random subset of examples of size $B$. $\kappa$ is the maximum column norm of $\mathbf{C}$.

The key idea behind Thm. 4 is the following: assume we have two datasets $D,D^{\prime}$ that differ in an example in $D_{1}$ (WLOG), i.e., the differing example can only participate in steps $1,b+1,\ldots,(k-1)b+1$. Then by the banded structure of $\mathbf{C}$ and the standard technique of reducing adaptive queries to non-adaptive queries (see e.g. Claim D.1 in ), the first $b$ rows of $\mathbf{C}\mathbf{x}$ (i.e., all the rows where examples in step 1 influence the output) can be viewed as a query on the examples in $D_{1}$ that were included in step 1, the next $b$ rows can be viewed as an adaptively chosen query on the examples included in steps $b+1$, and so on. See Fig. 4 for a visualization.

Thm. 5 in Sec. F.3 provides a strong a generalization of Thm. 4. It shows that shuffling, another common amplification technique often used for dp-sgd, can also be applied to our BandMF. We also provide an explicit algorithm for accounting for amplification via sampling in terms of the dp_accounting library , along with examples of concrete privacy parameters derived from these corollaries.

Optimizing the number of bands

Thm. 4 shows that different numbers of bands (with a corresponding sampling scheme) give different privacy amplification guarantees. This implies given a particular privacy (or RMSE) target, one should optimize the number of bands to get the best RMSE (or privacy) possible. This can be done efficiently. Generally, for large values of $\epsilon$ larger numbers of bands perform better, and as $\epsilon\rightarrow 0$, eventually amplification dominates so $\hat{b}=1$ becomes optimal. Details are in Sec. F.4.

Experiments

Our experiments on example-level DP for image classification (of CIFAR10) and user-level DP for next word prediction (NWP) of Stack Overflow (SONWP) focus on comparing our BandMF with the existing state-of-the-art Multi-epoch MF and dp-sgd . We finish by showing that BandMF improves over state-of-the-art for a production mobile keyboard next word prediction model. In all cases, noise $\sigma$ is calibrated using privacy loss distributions to achieve the stated privacy guarantees for zero-out adjacency , as implemented in . Following the common convention, amplified results are based on privacy accounting for Poisson sampling, though shuffling was used in non-production training. We find BandMF can outperform both mechanisms across a wide range of privacy budgets to as low as $\epsilon\approx 0.5$. Past this, it is no worse than either.

We begin by comparing dp-sgd with Multi-epoch MF and BandMF in terms of their RMSE (Sec. 2) on the prefix workload. This is one measure for how much noise these mechanisms add during training, and is a reasonable proxy for learning performance.Observe in Fig. 5(a) that there are regimes where dp-sgd outperforms Multi-epoch MF and vice versa in terms of RMSE. When then number of epochs equals $n$, dp-sgd reduces to full gradient descent (GD) (and there is no amplification benefit), and the optimal mf mechanism is close to the identity matrix (that is, GD), and so the algorithms become almost identical (mf has a very small advantage, as the identity matrix is not quite optimal for RMSE). However, as shown in Fig. 5(b), we see that BandMF is always at least as good as dp-sgd in terms of RMSE. The improvement is most potent in the lower number of epochs and higher $\epsilon$ regime, which is standard for large model training. In Fig. 5(c), we see that for fixed $n$ as the number of epochs increases, all mechanisms enjoy improved RMSE, and BandMF in fact reduces to dp-sgd in that regime (though BandMF may be outperformed in RMSE by Multi-epoch MF due to our imposed optimization constraint of constant diagonals in $\mathbf{X}$).

Centralized training with amplification

Our full experimental setup is described in App. H, and closely follows prior work . We train for $20$ epochs on CIFAR-10, and tune all mechanisms to achieve their best performance for each $\epsilon$, using $12$ repeated runs. Fig. 1(a) shows that BandMF with amplification and an optimized number of bands $\hat{b}$ can obtain utility benefits over both prior mechanisms. We find that for $\epsilon\in$, BandMF achieves a consistent $\approx 1$ percentage point boost in performance over Multi-epoch MF. Below $\epsilon\approx 2$, where dp-sgd previously dominated, we find that BandMF obtains a benefit around $3$ percentage points. These two findings show that BandMF is able to balance and leverage the benefits of both amplification and correlated noise effectively. As the budget $\epsilon$ gets smaller, we find that BandMF is equivalent to dp-sgd.

We next consider the now-standard StackOverflow next-word-prediction (NWP) task with user-level differential privacy, again following (full details in App. I), in particular 2052 steps and 6 epochs, with $B=1000$. The previous state-of-the-art for centralized training at $\epsilon\geq 2$ corresponds to their Multi-epoch MF.All StackOverflow experiments use the $\hat{b}{=}2052$ matrices optimized with equal-column-norms via Eq. 6 for Multi-epoch MF, which we observed to be minimally different in terms of RMSE and accuracy. We again tune $\hat{b}$ under amplification for optimal RMSE, selecting $\hat{b}=9,18,32,64$ for $\epsilon=1,2,4,8$ respectively. At $\epsilon=16$ we find Multi-epoch MF is optimal. Fig. 1(b) shows substantial improvements for combining amplification with mf for $\epsilon\in$; Table 5 of App. I gives the hyperparameters and accuracy values for this figure.

Cross-device federated learning

We consider SO NWP again, but now assuming each user’s data corresponds to the data on one device. We assume 2052 rounds and 6 epochs as above. Amplification is generally not possible in the cross-device setting, and so the prior state-of-the-art was (1) Single-epoch MF of Denisov et al. for single-epoch training (using $B=167$ rather than $B=1000$), and (2) Optimal TreeAgg, which essentially takes the binary-tree matrix $\mathbf{C}_{\mathcal{T}}$, and instead of using the less-efficient “online” estimator of Honaker , uses the pseudo-inverse $\mathbf{C}_{\mathcal{T}}^{\dagger}$ for noise generation (see Sec. 2). The $b$-min-sep-participation sensitivity of $\mathbf{C}_{\mathcal{T}}$ can still be calculated using the dynamic program of Kairouz et al. , while the use of $\mathbf{C}_{\mathcal{T}}^{\dagger}$ requires the machinery of Denisov et al. ; see e.g. the OptDecoderHonaker results of Choquette-Choo et al. [15, Fig. 6]. Our Thm. 3 further enables an upper-bound on the $b$-min-sep-participation sensitivity of the Multi-epoch MF matrices of Choquette-Choo et al. ; this incurs a penalty of about 15% (see Table 2 in App. C) compared to $(k,b)$-sensitivity. Fig. 6(a) shows that our BandMF and Multi-epoch MF again outperform prior baselines (though the previously untested multi-epoch Optimal TreeAgg performs quite well); Table 4 gives the hyperparameters and accuracy values for this figure.

Application in production cross-device FL

We fine-tune a Spanish next word prediction model, pretrained on the multilingual C4 dataset , with on-device user data using FL. Our setup follows , and is described in full in App. K. We compared to an existing implementation of the Online TreeAgg algorithm of Kairouz et al. (not the optimal version using $\mathbf{C}_{\mathcal{T}}^{\dagger}$ in simulation). Both algorithms ran for $n=2000$ training rounds. The BandMF matrix was optimized for $\hat{b}=400$ bands; however, the production system only allows approximate control of the separation between participations, and post-hoc we could only bound $b$ by 390 rounds for Online TreeAgg and 385 for BandMF, necessitating the use of Thm. 3 for the analysis of BandMF as $b<\hat{b}$.

We used the same clients/round goal of 6500 for both, and tuned noise multipliers to achieve comparable RMSE, hence tuning for a stronger privacy guarantee rather than improved accuracy. Fig. 6(b) shows our results, and we see BandMF actually achieves a slight improvement in accuracy, possibly due to learning-rate cooldown (which was only implemented for BandMF). Our primary result is then that we are able to improve the privacy guarantee from $\rho{=}0.52$-zCDP for Online TreeAgg to $\rho{=}0.24$-zCDP for BandMF, or $(\epsilon{=}6.69,\delta{=}10^{-10})$-DP to $(\epsilon{=}4.35,\delta{=}10^{-10})$-DP. Details of the privacy guarantee following the best practices of Ponomareva et al. are in Sec. K.1.

Discussion and Limitations

In this paper, we proposed the BandMF mechanism, which extends mf-dp-ftrl and enjoys the benefits of privacy amplification. This allows it to solely operate above the previous Pareto frontier defined by both amplified dp-sgd and mf-dp-ftrl in centralized training scenarios. Moreover, BandMF is well-suited to federated training scenarios, and improves state-of-the-art there as well. Additionally, the computational overhead of BandMF is less than mf-dp-ftrl by a factor of $\nicefrac{{b}}{{n}}$. It still has a $b\times$ time and space overhead compared to dp-sgd, which can be prohibitive for very large models with billions of parameters. This is an interesting and important future research direction.

The authors thank the early feedback and discussion from Natalia Ponomareva, and the support of FL production training from Yanxiang Zhang and Yuanbo Zhang.

References

Appendix A Notation summary

Fig. 7 compares dp-sgd and mf-dp-ftrl. To gain an intuition for why mf-dp-ftrl can perform better than dp-sgd, observe that vanilla SGD has iterates $\theta_{t}=\theta_{0}-\eta\sum_{i=1}^{t}\hat{\mathbf{x}}_{i}$, and hence when the noisy gradients $\hat{\mathbf{x}}_{i}$ are added, the $\mathbf{C}^{-1}_{[i,j]}\mathbf{z}_{[i,:]}$ terms in mf-dp-ftrl serve to cancel out some of the noise introduced on previous rounds. The noise cancellation reduces the total error in all the prefix sums of gradients $\sum_{i=1}^{t}\hat{\mathbf{x}}_{i}$ for $t\in[n]$, but also worsens the privacy guarantee of the mechanism, i.e. increases its sensitivity. The privacy worsens as e.g., an adversary trying to learn $\mathbf{x}_{1}$ via $\hat{\mathbf{x}}_{1}$ can partially learn the value of $\mathbf{z}_{1}$ from $\hat{\mathbf{x}}_{2}$, whereas in dp-sgd $\hat{\mathbf{x}}_{1}$ and $\hat{\mathbf{x}}_{2}$ are uncorrelated. Hence there is a tradeoff between the total error and sensitivity (see next paragraph): dp-sgd sets $\mathbf{C}^{-1}_{[i,j]}=0$ below the main diagonal, effectively minimizing sensitivity (assuming a fixed normalization of the main diagonal), but with a large total error due to no noise cancellation. On the other hand, mf-dp-ftrl can arrive at a better compromise between mechanism sensitivity and the total error. This is formalized in the optimization problem of Sec. 4.

Without a sampling assumption, the implied DP adversary knows which examples participated in batch $S_{i}$, and for DP-SGD with uncorrelated noise, knows they only need to “attack” $\hat{\mathbf{x}}_{i}$. However, with mf-dp-ftrl, the information from $S_{i}$ can potentially be masked with a larger amount of initial noise in $\hat{\mathbf{x}}_{i}$, which is then canceled out over subsequent rounds. “Spreading out” the release of information about batch $S_{i}$ over a larger number of iterations in this way can intuitively provide better privacy, while still allowing for accurate partial sums of gradients (and hence SGD iterates). This is, in a loose sense, similar to the way SGD with sampling “hides” information about a particular example at randomly chosen iterations.

Appendix B Dropping the diag⁡(𝐗)=𝟏diag𝐗1\operatorname{diag}{(\mathbf{X})}=\mathbf{1} constraint

As discussed in Sec. 4, BandMF by default imposes an equal column norm constraint on the generated factorization. In the optimization problem, this is accomplished by imposing the constraint $\operatorname{diag}{(\mathbf{X})}=\mathbf{1}$. In this section we show how we can solve the optimization problem without this constraint for $(k,b)$-participation. This optimization problem is formulated to minimize total squared error with respect to $(k,b)$-participation, although in principle the optimized matrices could be used in the $b$-min-sep-participation setting with some degradation in solution quality. Prop. B.1 provides an expression for efficiently computing the sensitivity of a $b$-banded matrix.

Let $\mathbf{X}\in\mathbf{S}_{+}^{n}$ be a $b$-banded matrix, and let $\Pi$ denote the $(k,b)$-participation schema. Then

To integrate this expression into the optimization problem, we can replace the $\operatorname{diag}{(\mathbf{X})}=\mathbf{1}$ constraint with $b$ linear constraints on $\operatorname{diag}{(\mathbf{X})}$. This modification does not affect the convexity of the problem, although it does slightly complicate the algorithms needed to solve it. To handle this new problem, our approach is to replace the gradient with respect to $\mathbf{X}$ at each iteration, with a projected gradient, which is obtained by setting $v_{i}=\sum_{j=0}^{k-1}\operatorname{diag}(\Delta\mathbf{X})_{i+jb}$ for all $i=1,\dots,b$, and setting $\operatorname{diag}(\Delta\mathbf{X})_{i+jb}=\operatorname{diag}(\Delta\mathbf{X})_{i+jb}-\nicefrac{{v_{i}}}{{k}}$. This ensures that sensitivity does not change between iterations of the numerical optimization procedure.

For the reasons mentioned in Sec. 4, by default we impose the simpler constraint $\operatorname{diag}(\mathbf{X})=\mathbf{1}$. In App. C, we provide some numerical comparisons between these two approaches. Specifically, the rows of Table 2 with Equal column norms? (F)alse correspond to the approach described here; observe this results in slightly improved RMSE under $(k,b)$-participation compared to the corresponding rows with Equal column norms? (T)rue.

Appendix C Empirical evaluation of banded matrices

Table 2 compares the matrix mechanisms studied under different participation patterns but normalized to have sensitivity $\operatorname{sens}(\mathbf{C})=1$ under $(k{=}6,b{=}342)$-participation. The sensitivity under single participation $k=1$ is lowest as expected. With column normalization, sensitivity is also $1$ under $b{\geq}{342}$-min-sep-participation. We make the following observations:

For the MF mechanisms, column normalization hurts RMSE for $(k,b)$-participation compared to the approach of App. B (as it is an additional constraint), but actually improves RMSE under $b$-min-sep-participation.

We conjecture that the $(k,b)$-participation optimized matrices (MF without column normalization, App. B) are optimal for the prefix-sum workloadThis conjecture is not trivially true, as we still enforce a non-negativity or orthogonality constraint; see Choquette-Choo et al. [15, Appendix I.3]. Hence the conjecture is that these constraints are already satisfied by the optimal matrix for this workload.; With this in mind, we see there is at most a small increase in RMSE for switching to the more challenging $b$-min-sep-participation schema ($1.00\rightarrow 1.05$) . If (as we further conjecture) the optimal matrices for prefix-sum in fact are $k$-banded, the gap is even smaller (at most $1.04\rightarrow 1.05)$. Hence, at least for the prefix-sum workload $\mathbf{A}$, there is limited room for improvement in developing optimization procedures that directly optimize over the larger feasible set offered by $b$-min-sep-participation.

Using fewer than $b$ bands does degrade performance on the RMSE metric, with DP-SGD being the extreme case, yielding prefix sum estimates almost $10\times$ worse than the mf mechanisms.

The results of Denisov et al. imply that the binary-tree $\mathbf{C}$ matrix can in fact be used in the online setting, with the Moore-Penrose pseudo-inverse giving the optimal decoder for RMSE , corresponding to the ‘full’ estimator of Honaker . We include this in the table as a baseline, and see that it is in general outperformed by our MF mechanisms by about $1.5\times$ in RMSE.

Appendix D Example structures of MF

Figs. 8 and 9 show the structure of some of the key matrix factorization approaches considered in this work. One can immediately see the impact of the $(k,b)$-participation schema in the optimal matrices, in particular for the non-banded Multi-epoch MF matrices (the two top-right matrices), where $\mathbf{C}$ contains diagonals of negative entries separated by $b$ steps. In the bottom two rows, we see that requiring equal column norms (“EN-” for equal norms) has a relatively minor impact on the structure of the matrices.

Appendix E Algorithms and Analysis for Sec. 3

E.2 Analysis

The sensitivity of $\mathbf{C}$ for a given participation schema $\Pi$ may be expressed as:

where $\mathbf{X}=\mathbf{C}^{\top}\mathbf{C}$. This upper bound is tight when $\mathbf{P}_{\pi}\mathbf{C}^{\top}\mathbf{C}\mathbf{P}_{\pi}\geq 0\forall\pi\in\Pi$, and is independent of the dimension $d$ of the rows of $\mathbf{u}$.

This implies the statement Eq. 7 by the definition of sensitivity and neighboring in our setting.

Now, let $\mathbf{X}_{\pi}:=\mathbf{P}_{\pi}\mathbf{C}^{\top}\mathbf{C}\mathbf{P}_{\pi}$ be the matrix formed by zeroing out the rows and columns not indexed by $\pi$ from $\mathbf{X}$. Assume that every $\mathbf{u}\in\mathfrak{D}$ has row norms bounded by 1. Expanding the trace in Eq. 7, writing $x_{ij}$ for the elements of $\mathbf{X}_{\pi}$ and $\mathbf{u}_{[j,:]}$ for the $j^{th}$ row of $\mathbf{u}$, we have

which yields the claimed bound. When $\mathbf{X}_{\pi}$ is elementwise nonnegative, taking $\mathbf{u}_{[i,:]}=\mathbf{u}_{[j,:]}$ for any unit vector shows the claimed tightness in this case.

This statement can be viewed as a partial extension of [15, Theorem G.1]. It does not imply every case handled there, but also implies results which cannot be derived from that Theorem.

Conclusion (1) is implied by (2), noting that the conditions on $\mathbf{C}$ imply that Alg. 5 will return a value at most $\kappa\sqrt{k^{\prime}}$ in this setting.

where $\mathbf{u}(\pi)\in\{0,1\}^{n}$ is given by $\mathbf{u}(\pi)_{i}=1$ if $i\in\pi$ and otherwise. The last equality follows from the orthogonality condition on sufficiently separated columns of $\mathbf{C}$ trivially implied by bandedness. It is straightforward to verify the dynamic program of Alg. 3 constructs a feasible $\pi$ which attains the maximum. ∎

The proof will be constructive. Let $\mathbf{J}$ be the $n\times n$ exchange matrix defined as

Let $\mathbf{Y}=\mathbf{J}\mathbf{X}\mathbf{J}$ and note that $\mathbf{Y}$ is symmetric and positive definite. Let $\mathbf{H}=\text{Cholesky}(\mathbf{Y})^{\top}$ and note that (1) $\mathbf{H}^{\top}\mathbf{H}=\mathbf{Y}$ by definition of Cholesky decomposition, (2) $\mathbf{H}$ is upper triangular, and (3) $\mathbf{H}$ is $\hat{b}$-banded by Du Croz et al. .

We will show that for $\mathbf{C}=\mathbf{J}\mathbf{H}\mathbf{J}$, we have (1) $\mathbf{X}=\mathbf{C}^{\top}\mathbf{C}$, (2) $\mathbf{C}$ is lower triangular, and (3) $\mathbf{C}$ is $\hat{b}$-banded.

For Claim (2) and (3), note that left-multiplying by $\mathbf{J}$ reverses the rows and right-multiplying by $\mathbf{J}$ reverses the columns, and therefore $\mathbf{C}_{[i,j]}=\mathbf{H}_{n-i+1,n-j+1}$.

For Claim (2), we need to show $\mathbf{C}_{[i,j]}=0$ if $i<j$. If $i<j$ then $n-i+1>n-j+1$, and since $\mathbf{H}$ is upper triangular, we know $\mathbf{H}_{[n-i+1,n-j+1]}=0$, as desired.

For Claim (3), we need to show that $\mathbf{C}_{[i,j]}=0if|i-j|\geq\hat{b}$. Observe that if $|(n-i+1)-(n-j+1)|=|i-j|$ and therefore since $\mathbf{H}$ is $\hat{b}$-banded, so is $\mathbf{C}$.

Appendix F Additional Analysis for Sec. 5

In this section we prove our general amplification statement Theorem 5, of which Theorem 4 is a corollary. Recall that we use $b$ instead of $\hat{b}$ in this appendix since our sampling scheme enforces $(k,b)$-participation. Throughout this section, we slightly abuse notation by letting $i\pmod{b}=b$ instead of if $i/b$ is integer.

We first give the general sampling scheme (Alg. 6) as well as sequence of queries (Alg. 7) that provides an upper bound on the privacy guarantees of DP-MF using this sampling scheme.

F.2 General Amplification Statement and Proof

Given the sampling scheme and query sequence, we can now state our general amplification statement:

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 guaranteeStandard DP includes $\epsilon$-DP, $(\epsilon,\delta)$-DP, Rényi DP, zCDP, and Gaussian DP. satisfied by Alg. 7 in Sec. F.1 with $\kappa=\max_{i\in[n]}\left\|\mathbf{C}\mathbf{e}_{i}\right\|_{2}=\max_{i\in[n]}\sqrt{\mathbf{X}_{i,i}}$ and the same choice of $\mathcal{S}$.

Consider two datasets $D,D^{\prime}$ that differ by an example contained in the partition subset $D_{j}$. We argue about the privacy of $\mathbf{C}\mathbf{x}+\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

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^{\prime}$, each row of $(\mathbf{C}\mathbf{x}+\mathbf{z})_{R_{\emptyset}}=\mathbf{C}_{R_{\emptyset}}\mathbf{x}+\mathbf{z}_{R_{\emptyset}}$ has a distribution independent of whether $D$ or $D^{\prime}$ was used. So it suffices to give privacy guarantees for outputting only $(\mathbf{C}\mathbf{x}+\mathbf{z})_{R_{j}},(\mathbf{C}\mathbf{x}+\mathbf{z})_{R_{b+j}},\ldots,(\mathbf{C}\mathbf{x}+\mathbf{z})_{R_{(k-1)b+j}}$.

We can decompose rows $R_{j}$ of $\mathbf{C}\mathbf{x}+\mathbf{z}$ as follows:

The same logic applies to each of $(\mathbf{C}\mathbf{x}+\mathbf{z})_{R_{b+j}},\ldots,(\mathbf{C}\mathbf{x}+\mathbf{z})_{R_{(k-1)b+j}}$. Putting it all together and taking a max over the sensitivity of the individual queries, releasing $\mathbf{C}\mathbf{x}+\mathbf{z}$ satisfies any standard privacy guarantee satisfied by answering $k$ adaptively chosen queries, with sensitivity $\max_{i\in[n]}\left\|\mathbf{C}\mathbf{e}_{i}\right\|_{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,\mathcal{S}$. ∎

F.3 Corollaries of Thm. 5

We give here several corollaries of Thm. 5 that are of interest.

Note that when $b=1$, the partition contains a single subset, i.e., is the 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^{[\check{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 . 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:

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{C}\mathbf{x}+\mathbf{z}$ satisfies $(\alpha,\frac{\alpha n}{2\sigma^{2}b})$-RDP.

Using Theorem 11 of 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 plugging this in gives $(\alpha,\frac{2\alpha B^{2}bn}{\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\pmod{m}},d_{(i-1)B+2\pmod{m}},\ldots,d_{iB\pmod{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 $\mathcal{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:

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]}\left\|\mathbf{C}\mathbf{e}_{i}\right\|_{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$.

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{C}\mathbf{x}+\mathbf{z}$ satisfies $(\epsilon,\delta)$-DP. With e.g., the amplification for shuffled $(\epsilon,\delta)$-DP mechanisms given by Theorem 5.1 of and Cor. F.1, if $\epsilon$ is a constant, we instead get that $\mathbf{C}\mathbf{x}+\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

where $\mathbf{C}_{b}$ is a $b$-banded lower triangular matrix optimized via 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 searchSee for example the calibrate_dp_mechanism function of DP Team . 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 ), 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. 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 $\lvert D_{j}\rvert=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 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).

Appendix G Additional RMSE Experiment Details

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}$ = Identity).

Appendix H Additional CIFAR-10 Experiment Details

We tune all jobs on a learning rate grid of coefficients in {1, 2, 5} on powers in . 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

Appendix I Additional StackOverflow Next-Word-Prediction Experiment Details

We follow the experimental setup for StackOverflow NWP from Denisov et al. and Choquette-Choo et al. . 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 , 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.

Learning rate warmup and cooldown

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 uses the effective prefix-sum RMSE (the objective we use for optimizing (banded) matrices $\mathbf{C}$),

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{I}\mathbf{C}^{-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):

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.

Appendix J Efficient Multiplication and Inverse of Banded Matrices

Algorithms 8 and 9 (Fig. 15) give algorithms for lower triangular banded matrix-vector multiplication and inverse banded matrix-vector multiplication. Note that both algorithms are compatible with the streaming nature of gradients. As soon as the next input $\mathbf{x}_{i}$ is received, the algorithm can immediately output $\mathbf{y}_{i}$. Both algorithms require storing a state of size $\hat{b}$, and run in $O(n\cdot\hat{b})$ time. While the algorithms are described with respect to computing matrix-vector products, they can also be used to compute matrix-matrix products where the right-hand-side is a $n\times d$ matrix by multiplying by each column independently. In this setting, these algorithms require $O(\hat{b}\cdot d)$ space and $O(n\cdot\hat{b}\cdot d)$ time. Both algorithms have appeared previously in the literature on Monte Carlo methods, which have a similar problem at their core to that of noise generation for mf; see e.g. [56, Section 2].

Appendix K Application to a Real-World Cross-Device FL System

We train a one-layer LSTM language model of $\sim$2.4 million parameters in a practical cross-device FL system following . The model is used for predicting the next word of Spanish in a mobile virtual keyboard. We pretrain the model on public multilingual C4 dataset , and then fine-tune with on-device user data in FL. In a common practical FL system, clients have to satisfy criteria like being charged, idle and connected to unmetered network to participate in a round , hence only a subset of clients can be reached and there is a strong diurnal pattern of client participation . It is very challenging to hold a fixed set of clients for evaluation, or develop random sampling for privacy amplification. Though the current implementation of client participation control is feasible through the client timer, the tuning of separation $b$ in practice can be challenging. Therefore, the training of Fig. 6 (b) can achieve smaller separation $b$ than in .

We follow the guidelines outlined in [48, Sec. 5.3] to report privacy guarantees.

DP setting. This a central DP guarantee where the service provider is trusted to correctly implement the mechanism.

Data accesses covered: The DP guarantee applies to all well-behaved clients Clients that faithfully follow the algorithm including participation limits. Due to the design of the algorithm, a mis-behaved client does not adversely affect the DP guarantee of any well-behaved clients. in a single training run. We do not account for hyperparameter tuning, or the selection of the final model checkpoint using evaluation metrics or A/B testing in our guarantees. Public multilingual C4 data is used for pre-training.

Final mechanism output: Only the final model checkpoint is released for use in production, however the mechanism’s output is technically the full sequence of privatized gradients, and so the guarantee also applies at this level, and hence all intermediate models are protected (including those sent to devices participating in federated learning).

Unit of privacy. Device-level DP is considered, i.e., the notion of adjacency is with respect to arbitrary training datasets on each client device, and the device might have an arbitrarily large local dataset containing arbitrary training examples. For user’s with a single device, this corresponds directly to user-level DP; for devices shared with multiple users, this provides a stronger notion of DP than user-level; for a user with multiple devices that happen to both participate in training the model, the notion is weaker, but group privacy can be used to obtain a user-level guarantee.

Adjacency definition for “neigbouring” datasets: We use the zero-out definition . This is a a special form of the add-or-remove definition, where neighboring data sets differ by addition/removal of a single client. In the absence of a client at any training step, we assume that the client’s model update gets replaced with the all zeros vector. This assumption enforces a subtle modification to the traditional definition of the add/remove notion of DP which allows neighboring data sets to have the same number of records.

Type of accounting used: Both $\rho-$zCDP accounting, and PLD accounting for $(\epsilon,\delta)-$DP are used.

Accounting assumptions : Each client only participates limited times during the training, and there are at least a min-separation of $b$ rounds between two consecutive participation of a client. This is enforced by a timer on clients in the cross-device FL system.

The formal DP statement: The privacy guarantees are $\rho{=}0.52$-zCDP and $(\epsilon{=}6.69,\delta{=}10^{-10})$-DP for Online TreeAgg, while BandMF achieves $\rho{=}0.24$-zCDP and $(\epsilon{=}4.35,\delta{=}10^{-10})$-DP.

Transparency and verifiability: We are going to open source our code based on TensorFlow Federated and Tensorflow Privacy. Key portions of the cross-device FL system will also open sourced.