[Paper Review] Fed-GraB: Federated Long-tailed Learning with Self-Adjusting Gradient Balancer

Zikai Xiao, Zihan Chen+8 more|arXiv (Cornell University)|Oct 11, 2023

Privacy-Preserving Technologies in Data4 citations

TL;DR

Fed-GraB proposes a self-adjusting gradient balancer (SGB) for federated long-tailed learning that dynamically reweights clients' gradients using global long-tailed priors derived via a Direct Prior Analyzer (DPA), enabling improved performance on minority classes without compromising majority class accuracy. It achieves state-of-the-art results on CIFAR-10-LT, CIFAR-100-LT, ImageNet-LT, and iNaturalist under privacy-preserving federated settings.

ABSTRACT

Data privacy and long-tailed distribution are the norms rather than the exception in many real-world tasks. This paper investigates a federated long-tailed learning (Fed-LT) task in which each client holds a locally heterogeneous dataset; if the datasets can be globally aggregated, they jointly exhibit a long-tailed distribution. Under such a setting, existing federated optimization and/or centralized long-tailed learning methods hardly apply due to challenges in (a) characterizing the global long-tailed distribution under privacy constraints and (b) adjusting the local learning strategy to cope with the head-tail imbalance. In response, we propose a method termed $ exttt{Fed-GraB}$, comprised of a Self-adjusting Gradient Balancer (SGB) module that re-weights clients' gradients in a closed-loop manner, based on the feedback of global long-tailed distribution evaluated by a Direct Prior Analyzer (DPA) module. Using $ exttt{Fed-GraB}$, clients can effectively alleviate the distribution drift caused by data heterogeneity during the model training process and obtain a global model with better performance on the minority classes while maintaining the performance of the majority classes. Extensive experiments demonstrate that $ exttt{Fed-GraB}$ achieves state-of-the-art performance on representative datasets such as CIFAR-10-LT, CIFAR-100-LT, ImageNet-LT, and iNaturalist.

Motivation & Objective

• To address the challenge of training robust global models in federated learning when clients hold non-IID, long-tailed data distributions without exposing private data.
• To overcome the limitation of existing methods that fail to balance tail-class performance due to lack of global prior knowledge under privacy constraints.
• To design a local training strategy that synergistically improves performance on both head and tail classes in federated long-tailed learning.
• To enable dynamic, closed-loop gradient reweighting that adapts to global class imbalance without requiring manual hyperparameter tuning.

Proposed method

• Introduces a Direct Prior Analyzer (DPA) module that estimates the global long-tailed class distribution from aggregated client gradients, enabling privacy-preserving global prior inference.
• Proposes a Self-adjusting Gradient Balancer (SGB) that uses a closed-loop control mechanism to reweight local gradients based on feedback from the global prior.
• Employs a PID-like control strategy to compute re-weighting coefficients β using the error signal Δ(t) − z(t), where Δ(t) tracks the imbalance and z(t) is a target reference.
• Uses a feedback mechanism where Δ(t) is updated based on the difference between current and target imbalance, enabling adaptive gradient rebalancing during training.
• Integrates SGB into each client’s local training loop, allowing real-time adjustment of gradient contributions to reduce distribution drift.
• Leverages federated averaging with gradient reweighting to ensure global model convergence while preserving class-specific performance.

Figure 1: (a): An illustration of the Fed-GraB framework, in which SGB is mounted on all classes of each client based on the prior vector $\mathcal{P}_{c}$ derived by the DPA module, ensuring federated balanced local training; (b): Comparison for the difference of positive and negative gradients.

Experimental results

Research questions

• RQ1 How can global long-tailed class statistics be estimated in a privacy-preserving manner without direct data sharing in federated learning?
• RQ2 How can local gradient updates be dynamically reweighted to improve minority class performance without degrading majority class accuracy?
• RQ3 Can a closed-loop control mechanism effectively balance gradients across clients with heterogeneous and potentially non-long-tailed local data distributions?
• RQ4 How robust is the proposed method to varying degrees of global data imbalance and client-level data heterogeneity?
• RQ5 What is the impact of different control parameters (KP, KI, KD) on the stability and performance of the gradient balancing mechanism?

Key findings

• Fed-GraB achieves state-of-the-art performance on CIFAR-10-LT, CIFAR-100-LT, ImageNet-LT, and iNaturalist, significantly improving tail-class accuracy while maintaining head-class performance.
• The Direct Prior Analyzer (DPA) effectively infers global long-tailed priors from client gradients, enabling accurate global imbalance estimation without data leakage.
• The Self-adjusting Gradient Balancer (SGB) reduces performance variance across clients and mitigates distribution drift caused by data heterogeneity.
• Ablation studies confirm that SGB’s closed-loop control mechanism is robust to different target values z(t), with minimal impact from initialization bias.
• SGB demonstrates strong generalization across imbalance factors from 5 to 50, with larger gains observed in highly imbalanced settings.
• The computational and storage overhead of SGB is minimal, comparable to standard optimization methods like Adam, making it practical for real-world deployment.

Figure 2: An example of the global and local distributions of CIFAR-10-LT (Left) and CIFAR-100-LT (Right) under non-IID setting. The bar on the top demonstrates the sorted global long-tailed distribution. Each row represents the diverse statistics of class imbalance of each client.