[Paper Review] On the effectiveness of partial variance reduction in federated learning with heterogeneous data

Bo Li, Mikkel N. Schmidt+2 more|arXiv (Cornell University)|Dec 5, 2022

Privacy-Preserving Technologies in Data4 citations

TL;DR

This paper proposes FedPVR, a partial variance reduction method in federated learning that applies variance reduction only to the final classification layers of deep neural networks. By reducing gradient variance specifically in the last layers, FedPVR achieves faster convergence and higher accuracy than FedAvg and other baselines on heterogeneous data, while maintaining low communication cost and proving convergence in both convex and non-convex settings.

ABSTRACT

Data heterogeneity across clients is a key challenge in federated learning. Prior works address this by either aligning client and server models or using control variates to correct client model drift. Although these methods achieve fast convergence in convex or simple non-convex problems, the performance in over-parameterized models such as deep neural networks is lacking. In this paper, we first revisit the widely used FedAvg algorithm in a deep neural network to understand how data heterogeneity influences the gradient updates across the neural network layers. We observe that while the feature extraction layers are learned efficiently by FedAvg, the substantial diversity of the final classification layers across clients impedes the performance. Motivated by this, we propose to correct model drift by variance reduction only on the final layers. We demonstrate that this significantly outperforms existing benchmarks at a similar or lower communication cost. We furthermore provide proof for the convergence rate of our algorithm.

Motivation & Objective

• To investigate why FedAvg underperforms in non-IID, heterogeneous data settings with deep neural networks.
• To identify which layers in a neural network are most affected by data heterogeneity and model drift.
• To design a communication-efficient federated learning method that targets only the most problematic layers for drift correction.
• To prove convergence of the proposed method under both convex and non-convex settings.
• To demonstrate that partial variance reduction outperforms full-model variance reduction and baseline methods in accuracy and communication efficiency.

Proposed method

• The method introduces a control variate based on the difference between client and server classifier gradients to reduce variance in the final fully connected layers.
• Variance reduction is applied selectively to the last few layers (e.g., layers 27–28 in VGG-11), while other layers use standard FedAvg updates.
• A drift diversity metric is introduced to quantify gradient direction and scale variation across clients per communication round.
• The algorithm uses momentum on non-variance-reduced layers to stabilize training.
• Conformal prediction is applied to improve reliability in high-stakes applications, using a small portion of the validation set for calibration.
• The convergence rate is theoretically analyzed, showing weak dependence on data heterogeneity and fast convergence comparable to centralized SGD.

Experimental results

Research questions

• RQ1 Which layers in a deep neural network are most affected by data heterogeneity in federated learning?
• RQ2 Can partial variance reduction—applied only to final layers—improve convergence and accuracy in non-IID settings?
• RQ3 How does the proposed method compare to FedAvg, FedProx, FedDyn, and SCAFFOLD in terms of communication efficiency and final model performance?
• RQ4 Does the proposed method provably converge under non-convex settings, and how does its convergence rate scale with data heterogeneity?
• RQ5 Can conformal prediction further improve reliability without increasing communication cost?

Key findings

• FedPVR achieves faster convergence than FedAvg, FedProx, FedDyn, and SCAFFOLD on CIFAR10 and CIFAR100 with varying levels of data heterogeneity (α=0.1 to α=1.0).
• With a fixed number of communication rounds, FedPVR achieves top-1 accuracy on par with or better than centralized learning in several settings.
• Applying variance reduction only to the final layers results in higher gradient diversity in feature extraction layers, which enhances model expressiveness.
• The method requires fewer communication rounds than baselines to reach the same accuracy, while transmitting a similar or slightly larger number of parameters than FedAvg.
• Theoretical analysis shows that FedPVR converges at a rate comparable to centralized SGD, with weak dependence on data heterogeneity measures.
• Conformal prediction improves predictive reliability, with adaptive prediction sets achieving high empirical coverage while maintaining small average predictive set sizes.