SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jingfeng Wu*
Wennan Zhu
wennanzhu@google.com
Peter Kairouz
kairouz@google.com
In federated frequency estimation (FFE), multiple clients work together to estimate the frequencies of their collective data by communicating with a server that respects the privacy constraints of Secure Summation (SecSum), a cryptographic multi-party computation protocol that ensures that the server can only access the sum of client-held vectors. For single-round FFE, it is known that count sketching is nearly information-theoretically optimal for achieving the fundamental accuracy-communication trade-offs [Chen et al., 2022]. However, we show that under the more practical multi-round FEE setting, simple adaptations of count sketching are strictly sub-optimal, and we propose a novel hybrid sketching algorithm that is provably more accurate. We also address the following fundamental question: how should a practitioner set the sketch size in a way that adapts to the hardness of the underlying problem? We propose a two-phase approach that allows for the use of a smaller sketch size for simpler problems (e.g. near-sparse or light-tailed distributions). We conclude our work by showing how differential privacy can be added to our algorithm and verifying its superior performance through extensive experiments conducted on large-scale datasets.
In many distributed learning applications, a server seeks to compute population information about data that is distributed across multiple clients (users). For example, consider a distributed frequency estimation problem where there are  $n$  clients, each holding a local data from a domain of size  $d$ , and a server that aims to estimate the frequency of the items from the  $n$  clients with the minimum communication cost. This task can be done efficiently by letting each client binary encode their data and send the encoding to the server, at a local communication bandwidth cost of  $\log(d)$  bits. With the binary encoding, the server can faithfully decode each local data and compute the global frequency vector (i.e., the normalized histogram vector).
However, the local data could be sensitive or private, and the clients may wish to keep it hidden from the server. The above binary encoding communication method, unfortunately, allows the server to observe each individual local data, and therefore may not satisfy the users' privacy concerns. Federated Analytics (FA) [Ramage and Mazzocchi, 2020, Zhu et al., 2020] addresses this issue by developing new methods that enable the server to learn population information about the clients while preventing the server from prying on any individual local data. In particular, a cryptographic
*Work done during an internship at Google Research
multi-party computation protocol, Secure Summation (SecSum) [Bonawitz et al., 2016], has become a widely adopted solution to provide data minimization guarantees for FA [Bonawitz et al., 2021]. Specifically, SecSum sets up a communication protocol between clients and the server, which injects carefully designed additive noise to each data that cancels out when all of the local data are summed together, but blurs out (information theoretically) each individual local data. Under SecSum, the server is able to faithfully obtain the correct summation of the data from all clients but is unable to read a single local data. The federated frequency estimation (FFE) problem refers to the distributed frequency estimation problem under the constraint of SecSum. Clearly, the binary encoding method is not compatible with SecSum, because when the binary encoding is passed to the server through SecSum, the server only gets the summation of the binary encodings of the users' data, which does not provide sufficient information for computing the global frequency vector.
A naive approach to FFE can be accomplished by employing one-hot encoding: each client encodes its local data into a  $d$ -dimensional one-hot vector that represents the local frequency vector and sends it to the server through  $\mathrm{SecSum}$ . Then the server observes the summation of the local frequency vectors using  $\mathrm{SecSum}$  and scales it by the number of clients to obtain the true frequency vector. However, this one-hot encoding approach costs  $\Theta(d \log(n))$  bits of communication per client. This is because  $\mathrm{SecSum}$  adds noise from a finite group of size  $\Theta(\log n)$  to each component of the  $d$ -dimensional local frequency vector [Bonawitz et al., 2016] to avoid overflows. With a linear dependence on domain size  $d$ , the one-hot encoding approach is inefficient for large domain problems, which is a common setting in practice. In what follows, we focus on the regime where  $d > n$ .
Recently, linear compression methods were applied to mitigate the high communication cost issue for FFE with large domains [Chen et al., 2021, 2022]. The idea is to first linearly compress the local frequency vector into a lower dimensional vector before sending it to the server through SecSum; as linear compression operators commute with the summation operator, the server equivalently observes a linearly compressed global frequency vector though SecSum (after rescaling by the number of clients). The server then applies standard decoding methods to approximately recover the global frequency vector from the linearly compressed one. In particular, Chen et al. [2022] show that CountSketch [Charikar et al., 2002] (among other sparse recovery methods) can be used as a linear compressor for the above purpose, which leads to a communication bandwidth cost of  $\mathcal{O}(n\log(d)\log(n))$  bits. Therefore when  $d > n$ , CountSketch achieves a saving in local communication bandwidth compared to the one-hot encoding method that requires  $\Theta(d\log(n))$  bits. Moreover, Chen et al. [2022] show that for FFE with a single communication round, an  $\Omega(n\log(d))$  local communication cost is information-theoretically unavoidable for worst-case data distributions, i.e., we cannot do better without making additional assumptions on the global frequency vector.
Contributions. In this work, we make three notable extensions to CountSketch for FFE problems.
1. We show that the way Chen et al. [2022] set up the sketch size (linear in the number of clients  $n$ ) is often pessimistic (see Corollary 2.4). In fact, in the streaming literature, the estimation error induced by CountSketch is known to adapt to the tail norm of the global frequency vector [Minton and Price, 2014], which is often sub-linear in  $n$ . Motivated by this, we provide an easy-to-use, two-phase approach that allows practitioners to determine the necessary sketch size by automatically adapting to the hardness of the FFE problem instance.
2. We consider FFE with multiple communication rounds, which better models practical deployments of FA where aggregating over (hundreds of) millions of clients in a single round is not possible due to device availability and limited server bandwidth. We propose a new multi-round sketch algorithm called HybridSketch that provably performs better than simple adaptations of CountSketch in the multi-round setting, leading to further improvements in the communication cost. Surprisingly, we show that HybridSketch adapts to the tail norm of a heterogeneity vector (see Theorem 3.2). Moreover, the tail of the heterogeneity vector is always no heavier, and could be much lighter, than that of the global frequency vector, explaining the advantage of HybridSketch. For instance, on the C4 dataset [Bowman et al., 2020] with a domain size of  $d = 150,868$  and 150,000 users, we show that our method can reduce the sketch size by  $83\%$  relative to simple sketching methods when the number of sketch rows is not very large.
3. We extend the Gaussian mechanism for CountSketch proposed by Pagh and Thorup [2022], Zhao et al. [2022] to the multi-round FFE setting to show how our sketching methods can be made differentially private [Dwork et al., 2006]. We also characterize the trade-offs between accuracy and privacy for our proposed method.
We conclude by verifying the performance of our methods through experiments conducted on several large-scale datasets. All proofs and additional experimental results are differed to the appendices.
2 Adapting CountSketch to the Hardness of the Instance
In this part, we focus on single-round FFE and show how CountSketch can achieve better results when the underlying problem is simpler. Motivated by this, we also provide a two-phase method for auto-tuning the hyperparameters of CountSketch, allowing it to automatically adapt to the hardness of the instance.
Single-Round FFE. Consider  $n$  clients, each holding an item from a discrete domain of size  $d$ . The items are denoted by  $x_{t} \in [d]$  for  $t = 1, \dots, n$ . Then the frequency of item  $j$  is denoted by
$$
f _ {j} := \frac {1}{n} \sum_ {t = 1} ^ {n} \mathbb {1} \left[ x _ {t} = j \right].
$$
We use  $\mathbf{x}_t$  to denote the one-hot representation of  $x_{t}$ , i.e.,  $\mathbf{x}_t = \mathbf{e}_{x_t}$  where  $(\mathbf{e}_t)_{t=1}^d$  refers to the canonical basis. Then the frequency vector can be denoted by
$$
\mathbf {f} := \left(f _ {1}, \ldots , f _ {d}\right) ^ {\top} = \frac {1}{n} \sum_ {t = 1} ^ {n} \mathbf {x} _ {t} \in [ 0, 1 ] ^ {d}.
$$
In single-round FFE, the  $n$  clients communicate with a server once under the constraint of SecSum, and aim to estimate the frequency vector  $\mathbf{f}$ . Note that SecSum ensures that the server can only observe the sum of the local data.
Algorithm 1 COUNT SKETCH FOR FEDERATED FREQUENCY ESTIMATION
Require:  $n$  clients with local data  $x_{t} \in [d]$  for  $t = 1, \dots, n$ . Sketch length  $L$  and width  $W$ .
1: The server prepares independent hash functions and broadcasts them to each client:
$$
h _ {\ell}: [ d ] \rightarrow [ W ], \sigma_ {\ell}: [ d ] \rightarrow \{\pm 1 \}, \text {f o r} \ell \in [ L ].
$$
2: for Client  $t = 1, \dots, n$  in parallel do
3: Client  $t$  encodes the local data  $x_{t}\in [d]$  to  $\mathsf{enc}(x_t)\in \mathbb{R}^{L\times W}$  where
$$
\left(\operatorname {e n c} \left(x _ {t}\right)\right) _ {\ell , k} = \mathbb {1} \left[ h _ {\ell} \left(x _ {t}\right) = k \right] \cdot \sigma_ {\ell} \left(x _ {t}\right) \text {f o r} \ell \in [ L ], k \in [ W ].
$$
4: Client  $t$  sends  $\mathsf{enc}(x_t)\in \mathbb{R}^{L\times W}$  to SecSum.
5: end for
6: SecSum receives  $\left(\operatorname{enc}(x_t)\right)_{t=1}^n$  and only reveals the summation  $\sum_{t=1}^n \operatorname{enc}(x_t)$  to the server.
7: for Item  $j = 1,\dots ,d$  in parallel do
8: Server produces  $L$  estimators for  $f_{j}$ :
$$
\mathbf {d e c} (j; \ell) := \sigma_ {\ell} (j) \cdot \left(\frac {1}{n} \sum_ {t = 1} ^ {n} \mathbf {e n c} (x _ {t})\right) _ {\ell , h _ {\ell} (j)} \text {f o r} \ell \in [ L ].
$$
9: Server computes the median of the  $L$  estimators:
$$
\operatorname {d e c} (j) := \operatorname {m e d i a n} \left\{\operatorname {d e c} (j; \ell): \ell \in [ L ] \right\}.
$$
10: end for
11: return  $(\mathsf{dec}(j))_{j=1}^{d}$  as estimate to  $(f_j)_{j=1}^d$ .
Count Sketch. CountSketch is a classic streaming algorithm that dates back to [Charikar et al., 2002]. In the literature of streaming algorithms, CountSketch has been extensively studied and is known to be able to adapt to the hardness of the problem instance. Specifically, CountSketch of a fixed size induces an estimation error adapting to the tail norm of the global frequency vector [Minton and Price, 2014].
A recent work by Chen et al. [2022] apply CountSketch to single-round FFE. See Algorithm 1 for details. They show that CountSketch approximately solves single-round FFE with a communication cost of  $\mathcal{O}(n\log (d)\log (n))$  bits per client. Moreover, they show  $\Omega (n\log (d))$  bits of communication per client is unavoidable for worst-case data distributions (unless additional assumptions are made),
confirming its near optimality. However, the results by Chen et al. [2022] are pessimistic as they ignore the ability of CountSketch to adapt to the hardness of the problem instance. In what follows, we show how the performance of CountSketch can be improved when the underlying problem becomes simpler.
We first present a problem-dependent accuracy guarantee for CountSketch of a fixed size,  $L \times W$ , that gives the sharpest bound to our knowledge. The bound is due to Minton and Price [2014] and is restated for our purpose.
Proposition 2.1 (Restated Theorem 4.1 in Minton and Price [2014]). Let  $(\hat{f}_j)_{j=1}^d$  be estimates produced by CountSketch (see Algorithm 1). Then for each  $p \in (0,1)$ ,  $W \geq 2$  and  $L \geq \log(1/p)$ , it holds that: for each  $j \in [d]$ , with probability at least  $1 - p$ ,
$$
| \hat {f} _ {j} - f _ {j} | <   C \cdot \sqrt {\frac {\log (1 / p)}{L} \cdot \frac {1}{W} \cdot \sum_ {i > W} \left(f _ {i} ^ {*}\right) ^ {2}},
$$
where  $(f_i^*)_{i\geq 1}$  refers to  $(f_i)_{i\geq 1}$  sorted in non-increasing order, and  $C > 0$  is an absolute constant.
For the concreteness of discussion, we will focus on  $\ell_{\infty}$  as a measure of estimation error in the remainder of the paper. Our discussions can be easily extended to  $\ell_2$  or other types of error measures. Proposition 2.1 directly implies the following  $\ell_{\infty}$ -error bounds for CountSketch (by an application of union bound).
Corollary 2.2 ( $\ell_{\infty}$ -error bounds for CountSketch). Consider Algorithm 1. Then for each  $p \in (0,1)$ ,  $L = \log(d/p)$  and  $W \geq 2$ , it holds that: with probability at least  $1 - p$ ,
$$
\left\| \operatorname {d e c} (\cdot) - \mathbf {f} \right\| _ {\infty} <   C \cdot \sqrt {\frac {1}{W} \cdot \sum_ {i > W} \left(f _ {i} ^ {*}\right) ^ {2}}, \tag {1}
$$
where  $C > 0$  is an absolute constant. In particular, (1) implies that
$$
\left\| \operatorname {d e c} (\cdot) - \mathbf {f} \right\| _ {\infty} <   C / W.
$$
According to Corollary 2.2, the estimation error is smaller when the underlying frequency vector  $(f_{i}^{*})_{i\geq 1}$  has a lighter tail. In other words, CountSketch requires a smaller communication bandwidth when the global frequency vector has a lighter tail. Our next Corollary 2.3 precisely characterizes this adaptive property in terms of the required communication bandwidth. To show this, we will need the following definition on the probable approximate correctness of an estimate.
Definition 1  $((\tau, p)$ -correctness). An estimate  $\hat{\mathbf{f}} := (\hat{f}_i)_{i=1}^d$  of the global frequency vector  $\mathbf{f} := (f_i)_{i=1}^d$  is  $(\tau, p)$ -correct if
$$
\mathbb {P} \left\{\| \hat {\mathbf {f}} - \mathbf {f} \| _ {\infty} := \max  _ {i} | \hat {f} _ {i} - f _ {i} | > \tau \right\} <   p.
$$
Corollary 2.3 (Oracle sketch size). Fix parameters  $\tau, p \in (0,1)$ . Then for CountSketch (see Algorithm 1) to produce an  $(\tau, p)$ -correct estimate, it suffices to set the sketch size to  $L = \log(d / p)$  and
$$
W = C \cdot \min  \left\{\left(\# \left\{f _ {i}: f _ {i} \geq \tau \right\} + \frac {1}{\tau^ {2}} \cdot \sum_ {f _ {i} <   \tau} f _ {i} ^ {2}\right), n \right\}, \tag {2}
$$
where  $C > 0$  is an absolute constant. In particular, the width  $W$  in (2) satisfies
$$
W \leq W _ {\text {w o r s t}} := C \cdot \min  \left\{2 / \tau , n \right\}. \tag {3}
$$
Corollary 2.3 suggests that the sketch size can be set smaller if the underlying frequency vector has a lighter tail. When translated to the communication bits per client (that is  $\mathcal{O}(L\cdot W\cdot \log (n))$ , where  $\log (n)$  accounts for the cost of SecSum), Corollary 2.3 implies that CountSketch requires
$$
\mathcal {O} \left(\min  \left\{\# \left\{f _ {i} \geq \tau \right\} + \frac {1}{\tau^ {2}} \sum_ {f _ {i} <   \tau} f _ {i} ^ {2}, n \right\} \log (d) \log (n)\right) \leq \mathcal {O} \left(\min  \left\{1 / \tau , n \right\} \log (d) \log (n)\right) \tag {4}
$$
bits of communication per client to be  $(\tau, p)$ -correct. In the worst case where  $(f_i)_{i=1}^d$  is  $\Theta(n)$ -sparse and  $\tau = \mathcal{O}(1/n)$ , (4) nearly matches the  $\Omega(n \log(d))$  information-theoretic worst-case communication cost shown in Chen et al. [2022], ignoring the  $\log(n)$  factor from SecSum. However, in practice,  $(f_i)_{i=1}^d$  has a fast-decaying tail, and (4) suggests that CountSketch can use less communication to solve the problem. We provide the following examples for a better illustration of the sharp contrast between the worst and typical cases.
Corollary 2.4 (Examples). Fix parameters  $\tau, p \in (0,1)$ . Consider Algorithm 1 with sketch length  $L = \log(d/p)$ . Then in each case for Algorithm 1 to produce an  $(\tau, p)$ -correct estimate for  $\tau > 1/n$ :
1. When  $f_{i} \propto 2^{-i}$ , it suffices to set  $W = \Theta (\log (1 / \tau))$ .
2. When  $f_{i} \propto i^{-a}$  for  $a > 1$ , it suffices to set  $W = \Theta (\tau^{-1 / a})$ .
3. When  $f_{i} \propto i^{-1} \log^{-b}(i)$  for  $b > 1$ , it suffices to set  $W = \Theta (\tau^{-1} \log^{-b}(1 / \tau))$ .
4. When  $f_{i} = 10 / n$  for  $i = 1, \dots, n / 10$ , it suffices to set  $W = \Theta(1 / \tau)$ .
A Two-Phase Method for Hyperparameter Setup. Corollary 2.3 allows to use CountSketch with a smaller width for an easier single-round FFE problem, saving communication bandwidth. However, the sketch size formula given by (2) in Corollary 2.3 relies on crucial information of the frequency  $(f_{i})_{i\geq 1}$ , i.e.,  $\# \{f_i:f_i\geq \tau \}$  and  $\sum_{f_i < \tau}f_i^2$ , which are unknown to who sets the sketch size. Thus, it is unclear if and how these gains can be realized in practical deployments.
We resolve this quandary by observing that in practice, the frequency vector often follows Zipf's law [Cevher, 2009, Powers, 1998]. This motives us to conservatively model the global frequency vector by a polynomial with parameters. By doing so, we can first run a small CountSketch to collect data from a (randomly sampled) fraction of the clients for estimating the parameters. Then based on the estimated parameter, we can set up an appropriate sketch size for a CountSketch to solve the FFE problem. This two-phase method is formally stated as follows.
We approximate the (sorted) global frequency vector  $(f_{i}^{*})_{i = 1}^{d}$  by a polynomial [Cevher, 2009] with two parameters  $\alpha >0$  and  $\beta >0$ , such that
$$
f _ {i} ^ {*} \approx \operatorname {p o l y} \left(i; \alpha , \beta\right), \operatorname {p o l y} \left(i; \alpha , \beta\right) := \left\{ \begin{array}{l l} \beta \cdot i ^ {- \alpha}, & i \leq i ^ {*}; \\ 0, & i > i ^ {*}, \end{array} \right.
$$
where  $i^{*} \coloneqq \max \{i : \sum_{j=1}^{i} \beta \cdot j^{-\alpha} \leq 1\}$  is set such that poly  $(i; \alpha, \beta)$  is a valid frequency vector. Here's an executive summary of the proposed approach for setting the sketch size.
1. Randomly select a subset of clients (e.g., 5,000 out of  $10^{6}$ .)
2. Fix a small sketch (e.g.,  $16 \times 100$ ) and run Algorithm 1 with the subset of clients to obtain an estimate  $(\tilde{f}_i)$ .
3. Use the top- $k$  values (e.g., top 20) from  $\tilde{f}_i$  to fit a polynomial with parameter  $\alpha$  and  $\beta$  (under squared error).
4. Solve Equation (4) under the approximation that  $f_{i}^{*} \approx \beta \cdot i^{\alpha}$  and output  $W$  according to the result.
Experiments. We conduct three sets of experiments to verify our methods. In the first set of experiments, we simulate a single-round FFE problem with the Gowalla dataset [Cho et al., 2011]. The dataset contains 6,442,892 lists of location information. We first construct a domain of size  $d = 175,000$ , which corresponds to a grid over the US map. Then we sample  $n = d / 10 = 17,500$  lists of the location information (that all belong to the domain created) to represent the data of  $n$  clients, uniformly at random. This way, we set up a single-round FFE problem with  $n = 17,500$  clients in a domain of size  $d = 175,000$ . In the experiments, we fix the confidence parameter to be  $p = 0.1$  and the sketch length to be  $L = \ln(2d / p) \approx 16$ . The targeted  $\ell_{\infty}$ -error  $\tau$  is chosen evenly from  $(10^{-3}, 10^{-1})$ . We only test  $\tau > 20 / n$  because it is less important to estimate frequencies over items with small counts (say, 20). For CountSketch, we compute sketch width with three strategies, using (2) (called "instance optimal"), using (3) (called "minimax optimal"), and using the two-phase method. We emphasize that the "instance optimal" method is not a practical algorithm as it requires access to unknown information about the frequency; we use it only for demonstrating the correctness of our theory. We set all constant factors to be 2. The results are presented in Figures 1(a) and (b). We observe that the "minimax optimal" way of hyperparameter choice is in fact suboptimal in practice, and is improved by the "instance optimal" and the two-phase strategies.
In the second set of experiments, we run simulations on the "Colossal Clean Crawled Corpus" (C4) dataset [Bowman et al., 2020], which consists of clean English text scraped from the web. We treat each domain in the dataset as a user and calculate the number of examples each user has. The domain size  $d = 150,868$ , which is the maximum example count per user. We randomly sample  $n = 150,000$  users from the dataset. We fix the sketch length to be  $L = 5$ . Other parameters are the same as the Gowalla dataset. The results are presented in Figures 1(d) and (e), and are consistent with what we have observed in the Gowalla simulations.
(a) Gowalla, single round
(c) Gowalla, multi-round
(d) C4, single round
(b) Gowalla, single round
(f) C4, multi-round
(g) Setiment-140, single round
(e) C4, single round
(h) Setiment-140, single round
(i) Setiment-140, multi-round
Figure 1: Single-round and multi-round FFE simulations. Subfigures (a) and (b) compare different hyperparameter strategies for CountSketch in a single-round FFE problem on the Gowalla dataset [Cho et al., 2011]. Subfigure (c) compares three sketch methods in a multi-round FFE problem on the Gowalla dataset. Subfigures (d), (e), and (f) are counterparts of subfigures (a), (b), and (c), respectively, but on the C4 [Bowman et al., 2020] dataset. Similarly, subfigures (g), (h), and (i) are counterparts of subfigures (a), (b), and (c), respectively, but on the Sentiment-140 [Go et al., 2009] dataset.
In the third set of experiments, we run simulations on a Twitter dataset Sentiment-140 [Go et al., 2009]. The dataset contains  $d = 739,972$  unique words from  $N = 659,497$  users. We randomly sample one word from each user to construct our experiment dataset. The number of rounds  $M = 10$ , and in each round,  $n = N / 10 = 65,949$  clients participate. The algorithm setup is the same as in the Gowalla experiments. Results are provided in Figures 1(g) and (h), and are consistent with our prior understandings.
3 Sketch Methods for Multi-Round Federated Frequency Estimation
In practice, having all clients participate in a single communication round is usually infeasible due to the large number of devices, their unpredictable availability, and limited server bandwidth [Bonawitz et al., 2019]. This motivates us to consider a multi-round FFE setting.
Multi-Round FFE. Consider a FFE problem with  $M$  rounds of communication. In each round,  $n$  clients participate, each holding an item from a universe of size  $d$ . The items are denoted by  $x_{t}^{(m)} \in [d]$ , where  $t \in [n]$  denotes the client index and  $m \in [M]$  denotes the round index. For simplicity, we assume in each round a new set of clients participate. So in total there are  $N = Mn$  clients. Then the frequency of item  $j$  is now denoted by
$$
f _ {j} := \frac {1}{M n} \sum_ {m = 1} ^ {M} \sum_ {t = 1} ^ {n} \mathbb {1} \left[ x _ {t} ^ {(m)} = j \right].
$$
For the  $m$ -th round, the local frequency is denoted by  $f_{j}^{(m)} \coloneqq \frac{1}{n}\sum_{t=1}^{n}\mathbb{1}\left[x_{t}^{(m)} = j\right]$ . Clearly, we have  $f_{j} = \frac{1}{M}\sum_{m=1}^{M}f_{j}^{(m)}$ . Similarly, we use  $\mathbf{x}_{t}^{(m)}$  to denote the one-hot representation of  $x_{t}^{(m)}$ , i.e.,
(a)  $f_{i}\propto i^{-1.1}$