Ternary Compression for Communication-Efficient Federated Learning

Jinjin Xu; Wenli Du; Ran Cheng; Wangli He; Yaochu Jin

Source-linked AI summary

Ternary Compression for Communication-Efficient Federated Learning

Jinjin Xu, Wenli Du, Ran Cheng, Wangli He, Yaochu Jin

arXiv:2003.03564v2Mar 7, 2020cs.LGcs.DCstat.ML

TL;DR

Federated learning must coordinate models across private, distributed data while controlling communication costs and handling non-IID distributions. The paper trains ternary-quantized client models with FTTQ and builds T-FedAvg to compress communication in both directions. Experiments and analyses report reduced communication costs alongside competitive federated-learning performance.

Problem
Non-IID data and the large parameter size of full-precision models make federated learning costly and difficult for resource-constrained devices.
Method
FTTQ optimizes ternary-quantized client networks during training, and T-FedAvg uses them to compress upstream and downstream federated communication.
Results
T-FedAvg reduced communication costs substantially, including 88% in upload and download for MLP, while achieving competitive test accuracy and convergence behavior.
Takeaways & Limitations
Ternary quantization can reduce federated-learning communication while retaining effective model training and inference on edge devices.

Abstract

from arXiv · show

Learning over massive data stored in different locations is essential in many real-world applications. However, sharing data is full of challenges due to the increasing demands of privacy and security with the growing use of smart mobile devices and IoT devices. Federated learning provides a potential solution to privacy-preserving and secure machine learning, by means of jointly training a global model without uploading data distributed on multiple devices to a central server. However, most existing work on federated learning adopts machine learning models with full-precision weights, and almost all these models contain a large number of redundant parameters that do not need to be transmitted to the server, consuming an excessive amount of communication costs. To address this issue, we propose a federated trained ternary quantization (FTTQ) algorithm, which optimizes the quantized networks on the clients through a self-learning quantization factor. Theoretical proofs of the convergence of quantization factors, unbiasedness of FTTQ, as well as a reduced weight divergence are given. On the basis of FTTQ, we propose a ternary federated averaging protocol (T-FedAvg) to reduce the upstream and downstream communication of federated learning systems. Empirical experiments are conducted to train widely used deep learning models on publicly available datasets, and our results demonstrate that the proposed T-FedAvg is effective in reducing communication costs and can even achieve slightly better performance on non-IID data in contrast to the canonical federated learning algorithms.

I. INTRODUCTION

Federated learning keeps data on distributed clients while coordinating a global model, but non-IID data and substantial bidirectional communication costs remain major challenges. This paper addresses these challenges by training ternary-quantized local models and introducing a communication-efficient federated protocol.

Non-IID and unbalanced client data can degrade federated learning performance, while increasing model depth and computation raise communication demands.
Existing approaches often use full-precision models or compress them only after client training, making deployment on resource-constrained IoT systems difficult.
Federated learning trains local models on private client data, aggregates them at a server, and redistributes the global model for continued local updates.
The paper introduces ternary quantization during local-model training and inference, producing models suited to network-edge devices.
T-FedAvg compresses both upstream and downstream communication, while weight quantization may make reverse engineering of model parameters more difficult.
The proposed methods are theoretically analyzed for quantization convergence, unbiasedness, and reduced weight divergence, then evaluated with MLP and ResNet models on MNIST and CIFAR10.

B. Quantization

The paper quantizes normalized network weights into ternary values using thresholds and learned scaling factors, then integrates this process into federated training to reduce computation and communication.

Quantization approach: Ternary quantization maps full-precision weights into {-1, 0, +1} using layer-wise thresholds and scaling factors.TTQ uses positive and negative factors, while the proposed approach adopts a trainable layer-specific quantization factor.
FTTQ: FTTQ normalizes weights layer by layer before quantization to avoid precision loss from magnitude imbalance.The normalization function maps weights into [-1, 1].
FTTQ: FTTQ trains the quantization factor together with local model weights and uses adapted thresholds influenced by layer sparsity.Controlled randomness is introduced into the client threshold parameter to make reverse engineering more difficult while limiting performance loss.
FTTQ: A single quantization factor replaces TTQ’s two factors because theoretical analysis shows they converge to the same absolute value.The change is intended to reduce communication and computation costs and may alleviate weight divergence.
Communication efficiency: Quantized local updates reduce upstream communication, while additional global-model quantization is required to reduce downstream communication.The paper therefore introduces a ternary federated averaging protocol after FTTQ.

B. Ternary Federated Averaging

T-FedAvg performs quantization during local training and communicates quantized models and quantization factors to support ternary federated averaging.

Client-server workflow: Clients normalize and quantize local models during training before uploading quantized parameters and quantization factors to the server.The server aggregates the uploaded local models to obtain a global model.
Client-server workflow: The server aggregates the local models and sends a quantized global model back to participating clients.At inference, only the quantized model is used for prediction.

1) Upstream:

T-FedAvg reduces communication in both directions by transmitting quantized local models upstream and quantized global models downstream, with a fallback for model collapse.

Upstream: Quantized local models and factors are uploaded to the server, where uploaded models are converted into continuous form and aggregated.The server then applies a strategy for returning the global model to clients.
Downstream: Strategy I re-quantizes the aggregated global model before broadcasting it, making downstream communication more efficient.The server uses a threshold of ∆S = 0.05 × max(|θr|) in the described strategy.
Downstream: Strategy II sends the full-precision global model if the quantized global model collapses, defined as a performance drop greater than 3%.The fallback is used when quantization causes that specified performance decline.
Downstream: T-FedAvg reduces communication costs in both upload and download phases, with Strategy I identified as more communication-efficient.The paper reports that Strategy I was used when MLP was the learning model.

IV. THEORETICAL ANALYSIS

The theoretical analysis studies convergence of TTQ quantization factors, unbiasedness of FTTQ and T-FedAvg, and reduced weight divergence under non-IID data.

Quantization-factor convergence: The analysis begins by proving convergence of the two TTQ quantization factors to the same absolute value.This result motivates replacing them with a single quantization factor.
Assumptions: The convergence proposition assumes normalized full-precision parameters are initialized from a symmetric distribution centered at 0.The stated example is θi ∽ U(−1, 1), with a fixed threshold and iteratively adapted factors.
Proof strategy: The proof treats the converged negative factor as an optimal quantization solution that reduces Euclidean distance between full-precision and quantized weights.The weights are partitioned into positive, zero, and negative index sets according to the threshold.

B. The Unbiasedness of FTTQ

This section establishes that FTTQ can produce an unbiased quantizer output under the paper’s stated uniform-distribution assumption. The proof analyzes the quantizer’s random masking and parameter distributions.

The proof analyzes quantized elements indexed by Ip, whose values follow a uniform distribution between ∆ and 1.The threshold ∆ is fixed once parameters are generated under Assumption 4.1.
The quantizer’s expectation is derived from random masking, the quantization variable u, and the number of indexed elements |Ip|.These quantities appear in the proof’s expectation calculation.
Under Assumption 4.1, FTTQ’s quantizer output is an unbiased estimator of the input.The paper first states the assumption and then concludes unbiasedness when weights are uniformly distributed.

C. The Properties of T-FedAvg

This section analyzes T-FedAvg’s behavior under shared initialization and shows that ternary quantization can reduce expected weight divergence in non-IID federated learning.

Under Assumption 4.2, all clients are initialized with the same global model.The assumption concerns a federated system with K clients and one server.
The divergence analysis compares the quantized federated model with the centralized model through the weight-divergence criterion.The criterion is attributed to Zhao et al. and uses θCen as the centralized-learning model.
The proof evaluates expectations using uniform variables and the triangular distribution of the sum of two independent uniform variables.The derivation includes the resulting density on [−2, 2].
Quantization can reduce the expected weight divergence when learning from non-IID data in federated learning.This conclusion follows the section’s theoretical proof.

D. Convergence Analysis

This section establishes a convergence guarantee for T-FedAvg under smooth, strongly convex objectives, bounded gradient variance, full participation, and IID data. With suitable local epochs, T-FedAvg matches the stated FedAvg convergence rate.

Under the stated smoothness, strong-convexity, gradient-noise, and gradient-bound assumptions, T-FedAvg has convergence rate O(1/(NR)).The result applies to N devices with full participation and IID data, where R is each client’s total SGD iterations.
The convergence proof relies heavily on prior proofs for federated optimization.The paper cites Qu et al. and other works for the detailed arguments.
The analysis sets γ = max{32κ, E} and uses step size ηr = 1/[4µ(γ+r)].These choices are made according to the cited convergence analysis.
Choosing local epochs E = O(√(R/N)) makes O(E^2/R^2) equal O(1/(NR)).Under this choice, the paper states that FedAvg and T-FedAvg share the convergence rate O(1/(NR)).

V. EXPERIMENTAL RESULTS

This section evaluates the proposed method on widely used benchmark datasets through controlled experiments comparing it with the standard federated learning algorithm. The evaluation focuses on test accuracy and communication costs.

The experiments use widely used benchmark datasets to evaluate the proposed method.
Controlled experiments compare the proposed method with the standard federated learning algorithm.
The comparisons measure test accuracy and communication costs.

A. Settings

The experiments evaluate centralized, federated, communication-mitigated, and ternary federated learning across benchmark datasets, model architectures, and data distributions. Settings include both simulated and physical client systems with controlled training comparisons.

System configuration: The physical system uses one CPU laptop server and five GPU workstation clients connected through a wireless LAN, with client-server-only communication.Simulations typically use 100 clients; the physical CIFAR10 setup uses N = 5 clients, while MNIST simulation uses N = 100.
Compared algorithms: The compared methods are centralized SGD, FedAvg, CMFL, and the proposed T-FedAvg, whose first and last feature layers remain full-precision.T-FedAvg is the quantized federated approach evaluated against the centralized and federated baselines.
Datasets: Experiments use MNIST and CIFAR10, representative classification benchmarks for small and larger models respectively.MNIST contains 60,000 training and 10,000 testing grayscale images; CIFAR10 contains 50,000 training and 10,000 testing color images across 10 classes.
Models: The evaluated architectures include an MLP, a shallow CNN, and a simplified ResNet18 with convolutional channels reduced to 64.The MLP has hidden layers of 30 and 20 neurons, while the CNN has five convolutional and three fully connected layers.
Data distribution: Data conditions include IID subsets, non-IID class assignments, and unequal client dataset sizes.Non-IID settings assign each client samples from Nc classes, while unbalancedness varies dataset sizes across clients.
Training configuration: Federated training uses λ = 0.1 for MNIST and λ = 1 for CIFAR10, with Nc = 10, batch size 64, and E = 5 or 10 local epochs.MNIST uses E = 5 and CIFAR10 uses E = 10; CIFAR10 learning rates decay by 0.95 every five communication rounds.
Evaluation protocol: Algorithms are compared after 100 training iterations using the same training-sample budget, with results averaged over five independent runs.For MNIST, federated learning runs for 100 rounds with 100 clients and 600 samples per client, while centralized methods train for 100 epochs.

B. Results on IID Data

On IID data, T-FedAvg substantially reduces communication while retaining competitive test accuracy across MLP, CNN, and ResNet. It also converges fastest when communication is measured in uploaded bits and performs comparably or better under the included non-IID evaluation.

Communication costs: 88% lower upload and download communication costs are achieved by T-FedAvg than FedAvg for MLP over 100 rounds.Strategy II is not used, and the first MLP layer is not transferred to the server.
Communication costs: 13% of FedAvg’s upload cost and 82% of FedAvg and CMFL’s download cost are used by T-FedAvg for CNN.The server uses Strategy II several times to maintain performance.
Communication costs: Nearly 94% of upstream and 25% of downstream communication are compressed by T-FedAvg for ResNet, while outperforming CMFL.These comparisons use fixed-round communication-cost evaluations on IID data.
Test accuracy: 91.95% test accuracy is achieved by T-FedAvg on MLP, exceeding FedAvg but remaining below centralized learning’s 92.25%.CMFL reaches 90.91%, while the centralized baseline reaches 92.25%.
Test accuracy: 84.46% test accuracy is achieved by T-FedAvg on CNN versus 85.47% for FedAvg, while ResNet reaches 87.87% versus 88.34%.These are the reported CIFAR10 test accuracies for the two models.
Convergence and non-IID performance: T-FedAvg converges fastest over uploaded bits in all three IID instances and achieves 5.22% and 6.43% enhancements over FedAvg and CMFL on CIFAR10 when Nc = 5.Over rounds, T-FedAvg converges faster than CMFL, while FedAvg is slightly faster for CNN on CIFAR10.

E. Influence of Unbalancedness in Data Size

The study evaluates federated learning under uneven client data sizes and examines computational runtime alongside accuracy. Unbalancedness has little effect on performance in the tested IID setting, while deeper models increase T-FedAvg’s computational burden.

Data unbalancedness: β ranges from 0.1, where samples concentrate on a few clients, to 1, where clients hold nearly equal sample counts.Experiments use an average of 30 participating clients out of 100.
Experimental setup: FedAvg and T-FedAvg are compared across varying β values to assess accuracy under unbalanced data sizes.The evaluation uses the test accuracies illustrated in Fig. 8.
Accuracy results: Unbalancedness does not significantly affect either federated learning algorithm’s performance in the tested IID setting.The authors attribute this to local models learning properly from IID data despite uneven client sample counts.
Computational complexity: T-FedAvg has per-iteration client complexity O((B+3/2)D), compared with FedAvg’s O(BD), while server averaging costs are O(λND) and O((λN−1)D), respectively.Here B is batch size, D the number of model elements, λ the participation ratio, and N the client count.
Runtime results: On the tested workstation, T-FedAvg’s MLP runtime is almost the same as FedAvg’s and much lower than CMFL’s.The runtime comparison counts five local epochs over 20 independent runs.
Limitations: As global-model depth increases, T-FedAvg’s time complexity increases significantly, creating a computational limitation for deep models.The authors propose quantizing only during the last local epoch as a possible future solution.