DAdaQuant: Doubly-adaptive quantization for communication-efficient Federated Learning

Federated Learning assumes a client-server topology with a set ${\mathbb{C}}=\{c_{i}|i\in\{1,2...N\}\}$ of $N$ clients that are connected to a single server. Each client $c_{k}$ has a local dataset $D_{k}$ from the local data distribution $\mathcal{D}_{k}$. Given a model $M$ with parameters ${\bm{p}}$, a loss function $f_{{\bm{p}}}(d\in D_{k})$ and the local loss $F_{k}({\bm{p}})=\frac{1}{|D_{k}|}\sum_{d\in D_{k}}f_{\bm{p}}(d)$, FL seeks to minimize the global loss $G({\bm{p}})=\sum_{k=1}^{N}\frac{|D_{k}|}{\sum_{l}|D_{l}|}F_{k}({\bm{p}})$.

DAdaQuant makes only minimal assumptions about the FL algorithm. Crucially, DAdaquant can complement FedAvg (McMahan et al., 2017), which is representative of a large class of FL algorithms.

FedAvg trains the model $M$ over several rounds. In each round $t$, FedAvg sends the model parameters ${\bm{p}}_{t}$ to a random subset ${\mathbb{S}}_{t}$ of $K$ clients who then optimize their local objectives $F_{k}({\bm{p}}_{t})$ and send the updated model parameters ${\bm{p}}_{t+1}^{k}$ back to the server. The server accumulates all parameters into the new global model ${\bm{p}}_{t+1}=\sum_{k\in{\mathbb{S}}_{t}}\frac{|D_{k}|}{\sum_{j}|D_{j}|}{\bm{p}}_{t+1}^{k}$  and starts the next round. Algorithm 1 lists FedAvg in detail. For our experiments, we use the FedProx (Li et al., 2018) adaptation of FedAvg. FedProx improves the convergence of FedAvg by adding the proximal term $\frac{\mu}{2}\|{\bm{p}}_{t+1}^{k}-{\bm{p}}_{t}\|^{2}$ to the local objective $F_{k}({\bm{p}}_{t+1}^{k})$ in Algorithm 1 of Algorithm 1.

While DAdaQuant can be applied to any quantizer with a configurable quantization level, it is optimized for fixed-point quantization. We introduce Federated QSGD as a competitive fixed-point quantizer on top of which DAdaQuant is applied.

In general, fixed-point quantization uses a quantizer $\textrm{{Q}}_{q}$ with quantization level $q$ that splits $\mathbb{R}_{\geq 0}$ and $\mathbb{R}_{\leq 0}$ into $q$ intervals each. $\textrm{{Q}}_{q}(p)$ then returns the sign of $p$ and $|p|$ rounded to one of the endpoints of its encompassing interval. $\textrm{{Q}}_{q}({\bm{p}})$ quantizes the vector ${\bm{p}}$ elementwise.

We design DAdaQuant’s quantization stage based on QSGD, an efficient fixed-point quantizer for state-of-the-art gradient compression. QSGD quantizes a vector ${\bm{p}}$ in three steps:

1. 1.

   Quantize ${\bm{p}}$ as $\textrm{{Q}}_{q}(\frac{{\bm{p}}}{||{\bm{p}}||_{2}})$ into $q$ bins in $[0,1]$, storing signs and $||{\bm{p}}||_{2}$ separately. (lossy)

2. 2.

   Encode the resulting integers with 0 run-length encoding. (lossless)

3. 3.

   Encode the resulting integers with Elias $\omega$ coding. (lossless)

QSGD has been designed specifically for quantizing gradients. This makes it not directly applicable to parameter compression. To overcome this limitation, we apply difference coding to uplink compression, first introduced to FL by FedPAQ. Each client $c_{k}$ applies $\textrm{{Q}}_{q}$ to the parameter updates ${\bm{p}}^{k}_{t+1}-{\bm{p}}_{t}$ (cf. Algorithm 1 of Algorithm 1) and sends them to the server. The server keeps track of the previous parameters ${\bm{p}}_{t}$ and accumulates the quantized parameter updates into the new parameters as ${\bm{p}}_{t+1}={\bm{p}}_{t}+\sum_{k\in{\mathbb{S}}_{t}}\frac{|D_{k}|}{\sum_{l}|D_{l}|}\textrm{{Q}}_{q}({\bm{p}}^{k}_{t+1}-{\bm{p}}_{t})$ (cf. Algorithm 1 of Algorithm 1). We find that QSGD works well with parameter updates, which can be regarded as an accumulation of gradients over several training steps. We call this adaptation of QSGD Federated QSGD.

Time-adaptive quantization uses a different quantization level $q_{t}$ for each round $t$ of FL training. DAdaQuant chooses $q_{t}$ to minimize communication costs without sacrificing accuracy. To this end, we find that lower quantization levels suffice to initially reduce the loss, while partly trained models require higher quantization levels to further improve (as illustrated in Figure 2). FracTrain is built on similar observations for non-distributed training. Therefore, we design DAdaQuant to mimic FracTrain in monotonically increasing $q_{t}$ as a function of $t$ and using the training loss to inform increases in $q_{t}$.

When $q$ is too low, FL converges prematurely. Like FracTrain, DAdaQuant monitors the FL loss and increases $q$ when it converges. Unlike FracTrain, there is no single centralized loss function to evaluate and unlike AdaQuantFL, we do not assume availability of global training loss $G({\bm{p}}_{t})$. Instead, we estimate $G({\bm{p}}_{t})$ as the average local loss $\hat{G}_{t}=\sum_{k\in{\mathbb{S}}_{t}}\frac{|D_{k}|}{\sum_{l}|D_{l}|}F_{k}({\bm{p}}_{t})$ where ${\mathbb{S}}_{t}$ is the set of clients sampled at round $t$. Since ${\mathbb{S}}_{t}$ typically consists of only a small fraction of all clients, $\hat{G}_{t}$ is a very noisy estimate of $G({\bm{p}}_{t})$. This makes it unsuitable for convergence detection. Instead, DAdaQuant tracks a running average loss $\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{{G}}}}_{t}=\psi\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{{G}}}}_{t-1}+(1-\psi)\hat{G}_{t}$.

We initialize $q_{1}=q_{\text{min}}$ for some $q_{\text{min}}\in{\mathbb{N}}$. DAdaQuant determines training to converge whenever $\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{{G}}}}_{t}\geq\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{{G}}}}_{t+1-\phi}$ for some $\phi\in{\mathbb{N}}$ that specifies the number of rounds across which we compare $\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{{G}}}}$. On convergence, DAdaQuant sets $q_{t}=2q_{t-1}$ and keeps the quantization level fixed for at least $\phi$ rounds to enable reductions in $G$ to manifest in $\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{{G}}}}$. Eventually, the training loss converges regardless of the quantization level. To avoid unconstrained quantization increases on convergence, we limit the quantization level to $q_{\text{max}}$.

The following equation summarizes DAdaQuant’s time-adaptive quantization:

FL algorithms typically accumulate each parameter $p_{i}$ over all clients into a weighted average $p=\sum_{i=1}^{K}{w_{i}p_{i}}$ (see Algorithm 1). Quantized FL communicates and accumulates quantized parameters $\textrm{{Q}}_{q}(p)=\sum_{i=1}^{K}{w_{i}\textrm{{Q}}_{q}(p_{i})}$ where $q$ is the quantization level. We define the quantization error $e^{q}_{p}$ as $e^{q}_{p}=|p-\textrm{{Q}}_{q}(p)|$. We observe that $\mathbb{E}_{p_{1}\ldots p_{K}}[\mathrm{Var}(\textrm{{Q}}_{q}(p))]$ is a useful statistic of the quantization error because it strongly correlates with the loss added by quantization. For a stochastic, unbiased fixed-point compressor like Federated QSGD, $\mathbb{E}_{p_{1}\ldots p_{K}}[\mathrm{Var}(\textrm{{Q}}_{q}(p))]$ equals $\mathbb{E}_{p_{1}\ldots p_{K}}[\mathrm{Var}(e^{q}_{p})]$ and can be evaluated analytically.

We observe in our experiments that communication cost per client is roughly a linear function of Federated QSGD’s quantization level $q$. This means that the communication cost per round is proportional to $Q=Kq$. We call $Q$ the communication budget and use it as a proxy measure of communication cost.

Client-adaptive quantization dynamically adjusts the quantization level of each client. This means that even within a single round, each client $c_{k}$ can be assigned a different quantization level $q_{k}$. The communication budget of client-adaptive quantization is then $Q=\sum_{k=1}^{K}{q_{k}}$ and $\textrm{{Q}}_{q}(p)$ generalizes to $\textrm{{Q}}_{q_{1}\ldots q_{K}}(p)=\sum_{i=1}^{K}{w_{i}\textrm{{Q}}_{q_{i}}(p_{i})}$. We devise an algorithm that chooses $q_{k}$ to minimize $Q$ subject to $\mathbb{E}_{p_{1}\ldots p_{K}}[\mathrm{Var}(e^{q_{1}\ldots q_{K}}_{p})]=\mathbb{E}_{p_{1}\ldots p_{K}}[\mathrm{Var}(e^{q}_{p})]$ for a given $q$. Thus, our algorithm effectively minimizes communication costs while maintaining a quantization error similar to static quantization. Theorem 1 provides us with an analytical formula for quantization levels $q_{1}\ldots q_{K}$.

DAdaQuant applies Theorem 1 to lower communication costs while maintaining the same loss as static quantization does with a fixed $q$. To ensure that quantization levels are natural numbers, DAdaQuant approximates the optimal real-valued solution as $q_{i}=\max(1,\text{round}(\sqrt{\frac{a}{b}}\times w_{i}^{2/3}))$. Appendix B gives a detailed proof of Theorem 1. To the best of our knowledge, DAdaQuant is the first algorithm to use client-adaptive quantization.

DAdaQuant combines the time-adaptive and client-adaptive quantization algorithms described in the previous sections. At each round $t$, time-adaptive quantization determines a preliminary quantization level $q_{t}$. Client-adaptive quantization then finds the client quantization levels $q_{t}^{k},k\in\{1,\ldots,K\}$ that minimize $\sum_{i=1}^{K}{q_{i}}$ subject to $\mathbb{E}_{p_{1}\ldots p_{K}}[\mathrm{Var}(e^{q_{1}\ldots q_{K}}_{p})]=\mathbb{E}_{p_{1}\ldots p_{K}}[\mathrm{Var}(e^{q}_{p})]$. Algorithm 1 lists DAdaQuant in detail. Figure 3 gives an example of how our time-adaptive, client-adaptive and doubly-adaptive quantization algorithms set quantization levels.

Reisizadeh et al. (2019) prove the convergence of FL with quantization for convex and non-convex cases as long as the quantizer Q is (1) unbiased and (2) has a bounded variance. These convergence results extend to DAdaQuant when combined with any quantizer that satisfies (1) and (2) for DAdaQuant’s minimum quantization level $q=1$. Crucially, this includes Federated QSGD.

We highlight DAdaQuant’s low overhead and general applicability. The computational overhead is dominated by an additional evaluation epoch per round per client to compute $\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{{G}}}}_{t}$, which is negligible when training for many epochs per round. DAdaQuant can compliment any FL algorithm that trains models over several rounds and accumulates a weighted average of client parameters. Most FL algorithms, including FedAvg, follow this design.

Evaluation We use DAdaQuant with Federated QSGD to train different models with FedProx on different datasets for a fixed number of rounds. We monitor the test loss and accuracy at fixed intervals and measure uplink communication at every round across all devices.

Models & datasets We select a broad and diverse set of five models and datasets to demonstrate the general applicability of DAdaQuant. To this end, we use DAdaQuant to train a linear model, CNNs and LSTMs of varying complexity on a federated synthetic dataset (Synthetic), as well as two federated image datasets (FEMNIST and CelebA) and two federated natural language datasets (Sent140 and Shakespeare) from the LEAF (Caldas et al., 2018) project for standardized FL research. We refer to Section A.1 for more information on the models, datasets, training objectives and implementation.

System heterogeneity In practice, FL has to cope with clients that have different compute capabilities. We follow Li et al. (2018) and simulate this system heterogeneity by randomly reducing the number of epochs to $E^{\prime}$ for a random subset ${\mathbb{S}}_{t}^{\prime}\subset{\mathbb{S}}_{t}$ of clients at each round $t$, where $E^{\prime}$ is sampled from $[1,\ldots,E]$ and $|{\mathbb{S}}_{t}^{\prime}|=0.9K$.

Baselines We compare DAdaQuant against competing quantization-based algorithms for FL parameter compression, namely Federated QSGD, FedPAQ (Reisizadeh et al., 2019), GZip with fixed-point quantization (FxPQ + GZip), UVeQFed (Shlezinger et al., 2020) and FP8. Federated QSGD (see section 3.3) is our most important baseline because it outperforms the other algorithms. FedPAQ only applies fixed-point quantization, which is equivalent to Federated QSGD without lossless compression. Similarly, FxPQ + GZip is equivalent to Federated QSGD with Gzip for its lossless compression stages. UVeQFed generalizes scalar quantization to vector quantization, followed by arithmetic coding. We apply UVeQFed with the optimal hyperparameters reported by its authors. FP8 (Wang et al., 2018a) is a floating-point quantizer that uses an 8-bit floating-point format designed for storing neural network gradients. We also evaluate all experiments without compression to establish an accuracy benchmark.

Hyperparameters With the exception of CelebA, all our datasets and models are also used by Li et al.. We therefore adopt most of the hyperparameters from Li et al. and use LEAF’s hyperparameters for CelebA Caldas et al. (2018). For all experiments, we sample 10 clients each round. We train Synthetic, FEMNIST and CelebA for 500 rounds each. We train Sent140 for 1000 rounds due to slow convergence and Shakespeare for 50 rounds due to rapid convergence. We use batch size 10, learning rates 0.01, 0.003, 0.3, 0.8, 0.1 and $\mu$s (FedProx’s proximal term coefficient) 1, 1, 1, 0.001, 0 for Synthetic, FEMNIST, Sent140, Shakespeare, CelebA respectively. We randomly split the local datasets into 80% training set and 20% test set.

To select the quantization level $q$ for static quantization with Federated QSGD, FedPAQ and FxPQ + GZip, we run a gridsearch over $q=1,2,4,8,\ldots$ and choose for each dataset the lowest $q$ for which Federated QSGD exceeds uncompressed training in accuracy. We set UVeQFed’s “coding rate” hyperparameter $R=4$, which is the lowest value for which UVeQFed achieves negligible accuracy differences compared to uncompressed training. We set the remaining hyperparameters of UVeQFed to the optimal values reported by its authors. Section A.3 shows further experiments that compare against UVeQFed with $R$ chosen to maximize its compression factor.

For DAdaQuant’s time-adaptive quantization, we set $\psi$ to 0.9, $\phi$ to ${1/10}^{th}$ of the number of rounds and $q_{\textrm{max}}$ to the quantization level $q$ for each experiment. For Synthetic and FEMNIST, we set $q_{\textrm{min}}$ to 1. We find that Sent140, Shakespeare and CelebA require a high quantization level to achieve top accuracies and/or converge in few rounds. This prevents time-adaptive quantization from increasing the quantization level quickly enough, resulting in prolonged low-precision training that hurts model performance. To counter this effect, we set $q_{\textrm{min}}$ to $q_{\textrm{max}}/2$. This effectively results in binary time-adaptive quantization with an initial low-precision phase with $q=q_{\textrm{max}}/2$, followed by a high-precision phase with $q=q_{\textrm{max}}$.

We repeat the main experiments three times and report average results and their standard deviation (where applicable). Section 4.2 shows the highest accuracy and total communication for each experiment. Figure 5 plots the maximum accuracy achieved for any given amount of communication.