PaperlayerView exact revision on arXiv

Source-linked AI summary

Local Learning Matters: Rethinking Data Heterogeneity in Federated Learning

Matias Mendieta, Taojiannan Yang, Pu Wang, Minwoo Lee, Zhengming Ding, Chen Chen

arXiv:2111.14213v3cs.LGcs.CVcs.DC

TL;DR

Federated learning must optimize across clients with non-IID data, while common proximal approaches can add resource costs and restrict local learning. This paper studies local learning generality using regularization and second-order indicators, then proposes FedAlign. FedAlign delivers competitive or state-of-the-art accuracy across varied settings with lower computation and memory overhead, though the experiments focus only on image tasks and models.

  • Problem

    Non-IID client data makes federated optimization difficult, while proximal solutions can incur computation or memory overhead and restrict local updates.

  • Method

    The paper analyzes regularization using second-order indicators of learning generality and proposes FedAlign, which regularizes final-block Lipschitz constants through distillation.

  • Results

    FedAlign achieves state-of-the-art accuracy across multiple datasets and varied FL settings while requiring significantly less computation and memory overhead than other state-of-the-art methods.

  • Takeaways & Limitations

    The findings support addressing federated data heterogeneity through local learning generality rather than solely restraining local models toward the global model.

  • Takeaways & Limitations

    The experiments focus only on image tasks and models, leaving natural language processing applications for future work.

Abstract

from arXiv · show

Federated learning (FL) is a promising strategy for performing privacy-preserving, distributed learning with a network of clients (i.e., edge devices). However, the data distribution among clients is often non-IID in nature, making efficient optimization difficult. To alleviate this issue, many FL algorithms focus on mitigating the effects of data heterogeneity across clients by introducing a variety of proximal terms, some incurring considerable compute and/or memory overheads, to restrain local updates with respect to the global model. Instead, we consider rethinking solutions to data heterogeneity in FL with a focus on local learning generality rather than proximal restriction. To this end, we first present a systematic study informed by second-order indicators to better understand algorithm effectiveness in FL. Interestingly, we find that standard regularization methods are surprisingly strong performers in mitigating data heterogeneity effects. Based on our findings, we further propose a simple and effective method, FedAlign, to overcome data heterogeneity and the pitfalls of previous methods. FedAlign achieves competitive accuracy with state-of-the-art FL methods across a variety of settings while minimizing computation and memory overhead. Code is available at https://github.com/mmendiet/FedAlign

1. Introduction

Federated learning is challenged by heterogeneous client data, which causes local models to drift toward misaligned objectives and can make convergence costly. The paper reframes the problem around local learning generality, studies regularization with second-order indicators, and proposes resource-efficient FedAlign.

  • Problem: Non-IID client data makes federated optimization difficult because local models converge toward objectives misaligned with the global model.This client drift causes overfitting to local objectives and hinders the performance of the aggregated model.
  • Motivation: Limiting local epochs reduces drift but slows convergence and increases communication rounds, creating impractical time and communication costs.The paper identifies effective treatment of data heterogeneity as a central concern for real-world federated systems.
  • Motivation: Existing proximal approaches restrain local updates but also limit local convergence potential, while some incur substantial computation or memory overhead.The paper links restricted local updates to less new information gathered per communication round and notes unstable gains across non-IID settings.
  • Approach: The paper rethinks data heterogeneity through local learning generality, analyzing data and structural regularization methods with second-order indicators and varied FL evaluations.Hessian eigenvalue, trace, and cross-client Hessian matching measurements are identified as meaningful indicators for successful optimization.
  • Approach: FedAlign is a distillation-based regularization method that targets final-block Lipschitz constants to promote generality while minimizing additional resource needs.By regularizing only the final block, FedAlign focuses on the portion described as most prone to overfitting.
  • Outcome: FedAlign achieves state-of-the-art accuracy across multiple datasets and FL settings while requiring significantly less computation and memory overhead than other state-of-the-art methods.The paper presents this as the method’s principal accuracy-efficiency outcome.

2. Related Work

Prior federated learning work improves local training or server aggregation, with local methods commonly using proximal or contrastive terms to control client drift. The paper highlights their trade-offs and studies regularization as an alternative local-training strategy.

  • Federated Learning: Federated learning algorithms seek a collective model minimizing training loss across clients, with FedAvg optimizing local loss and aggregating weighted client updates.The weighting uses each client’s number of training samples relative to the total.
  • Federated Learning: Recent methods target either client-side local training or server-side aggregation; this work focuses on local training and client drift.Server-side aggregation methods are treated as orthogonal to the paper’s focus.
  • Proximal and Contrastive Methods: FedProx adds a proximal term to limit extreme local updates, while MOON uses a contrastive loss to align current local representations with the global model.MOON also pushes the current local model away from a previous-round local copy in representation space.
  • Limitations: Proximal and contrastive approaches can reduce local convergence potential and gather less new information per communication round.These methods may also require substantial memory and computation; MOON stores and processes three full-size models simultaneously.
  • Alternative Directions: Server-side neuron matching and normalized averaging address aggregation, while federated Mixup methods can compromise privacy and increase communication burden by sharing additional data.The paper separates these aggregation-oriented approaches from its local-training focus.
  • Regularization: The paper evaluates FedAvg, FedProx, and MOON against local regularization methods including Mixup, Stochastic Depth, and GradAug.These regularization methods are applied within a standard FedAvg setup.

3. Empirical Study

The empirical study compares standard FL algorithms with regularization methods across heterogeneous federated settings, using accuracy and second-order indicators to assess performance and learning generality. Regularization methods, particularly GradAug, often match or outperform proximal approaches while requiring less resource overhead.

  • Experimental setup: Standard regularization methods are evaluated alongside FedAvg, FedProx, and MOON within a common CIFAR-100 federated setting.The study uses ResNet56 and Dirichlet-partitioned client data, with α = 0.5 by default.
  • Results comparison: GradAug achieves approximately 2% higher accuracy than MOON and approximately 4% higher accuracy than FedAvg and FedProx.StochDepth also achieves accuracy similar to MOON, while regularization methods use less memory and/or compute than MOON.
  • Second-order analysis: Hessian eigenvalue and trace measurements, together with cross-client Hessian matching, serve as indicators of loss-space smoothness, generality, and consistency.The analysis estimates the top Hessian eigenvalue and trace for final averaged global models and relates these measures to method effectiveness.
  • Second-order analysis: GradAug has the lowest top Hessian eigenvalue and Hessian trace, and its loss landscape is considerably smoother than FedAvg’s.Regularization methods reduce λmax and HT more effectively than FedProx and MOON.
  • Data heterogeneity: As heterogeneity decreases, GradAug remains best, while FedProx and MOON improve more slowly and retain low accuracy in the homogeneous setting.At α = 0.1, StochDepth improves over MOON by approximately 1.7% and also improves in more homogeneous settings.
  • Local training epochs: From 20 to 30 local epochs, regularization methods continue improving while FedProx and MOON remain relatively stagnant.FedAvg improves by approximately 0.3% over the same interval, whereas structural regularization maintains productive training across a wider epoch range.
  • Number of clients: Across client-count and partial-sampling settings, regularization methods maintain higher accuracy than FedAvg, with GradAug performing best in all cases.These methods remain competitive despite not being designed specifically for federated learning or partial client sampling.

4. Proposed Method – FedAlign

FedAlign promotes local learning generality by aligning Lipschitz-constant approximations between full-width and slimmed final blocks. It targets strong accuracy with substantially lower compute and memory overhead than competing regularization methods.

  • Method: FedAlign is a distillation-based regularizer that aligns Lipschitz constants between full-width and reduced-width versions of the final network block.The reduced-width sub-block uses temporary uniform pruning and reuses the preceding block’s features.
  • Method: The local objective combines cross-entropy with a Lipschitz-alignment loss computed from spectral-norm approximations of full-width and reduced-width transmitting matrices.The spectral norms are approximated with Power Iteration, avoiding singular value decomposition.
  • Method: FedAlign regularizes local representations toward generality rather than forcing local models to remain close to the global model.This distinguishes its correction term from traditional proximal approaches that directly restrain local updates.
  • Results: 1.02x difference in FLOPs accompanies a ∼4.0% accuracy improvement over FedAvg on the reported CIFAR-100 comparison.FedAlign also improves accuracy over MOON by ∼1.9% while reducing local compute by over 65% and memory requirements by over 70%.
  • Results: FedAlign reduces compute needs by ∼47% versus GradAug (n = 2) and ∼33% versus GradAug (n = 1) without sacrificing accuracy.Across additional settings, it maintains competitive accuracy, including against GradAug in ImageNet-200 sampling experiments.
  • Experiments: The method was evaluated across ablations and additional CIFAR-10 and ImageNet-200 settings, using fixed FedAlign hyperparameters ωS = 0.25 and µ = 0.45.The additional experiments used ResNet56 on CIFAR-10 and ResNet18 on ImageNet-200.

5. Conclusion and Discussion

The study reframes federated optimization around local learning generality and proposes FedAlign as a resource-efficient approach. The authors report strong improvements over classic baselines and state-of-the-art methods, while identifying limits in experimental scope.

  • The paper studies federated data heterogeneity through the perspective of local learning generality.
  • The experiments focus only on image tasks and models, leaving natural language processing applications for future work.The authors state that FedAlign can be applied to language applications because it operates in feature space and does not fundamentally rely on input type.
  • FedAlign provides strong improvement over classic baselines and state-of-the-art federated learning methods while addressing local computational constraints.

Supplementary Material

The supplementary material documents efficiency analysis, second-order analysis, hyperparameter studies, data partitioning, and implementation details. These sections cover both evaluation and practical training considerations.

  • Supplementary Material: Section A analyzes communication and compute efficiency for all explored methods.
  • Supplementary Material: Section B presents second-order analysis of FedAlign.
  • Supplementary Material: Section C reports hyperparameter ablations for FedAlign.
  • Supplementary Material: The supplementary material details and visualizes the non-IID data partitioning scheme.
  • Supplementary Material: Section E gives implementation details for transmitting matrices and training FedAlign.

A. Communication and Compute Efficiency

The efficiency analysis evaluates how much local computation and communication each method requires to reach a target accuracy on CIFAR-100. It measures total computation and parameter-transfer communication across completed rounds.

  • A. Communication and Compute Efficiency: The evaluation measures rounds, local compute, and communication cost required to achieve 60% accuracy on CIFAR-100.The setting allows methods to run without limiting the number of rounds.
  • A. Communication and Compute Efficiency: Local computation is calculated as the total MFLOPs across all nodes, samples, and completed rounds.
  • A. Communication and Compute Efficiency: Communication cost is calculated from the model parameters transferred as 32 bit weights over all completed rounds.

B. Second-order Analysis

The second-order analysis finds that FedAlign produces the lowest maximum Lipschitz constant among the evaluated methods. However, it does not directly produce a strong reduction in HN.

  • B. Second-order Analysis: FedAlign achieves the lowest λmax across all evaluated methods.Its Lipschitz-focused distillation loss reduces λmax as intended.
  • B. Second-order Analysis: The reduced λmax is associated with stronger generalization and performance.
  • B. Second-order Analysis: FedAlign does not directly translate to a strong reduction in HN.The authors identify extending FedAlign to address HN as a promising direction for future work.

C. Hyperparameter Ablations of FedAlign

FedAlign uses default hyperparameters ωS = 0.25 and µ = 0.45, with ablations indicating greater sensitivity to ωS than µ. Table 9 reports accuracy and second-order metrics for comparing methods on CIFAR-100.

  • FedAlign’s default hyperparameters are ωS = 0.25 and µ = 0.45.
  • FedAlign is more sensitive to ωS than µ in the CIFAR-100 basic-setting ablations.The ablations vary µ with ωS fixed at 0.25, and vary ωS with µ fixed at 0.45.
  • Table 9 reports CIFAR-100 accuracy alongside λmax and HT for loss-space smoothness, and HN and HD for cross-client consistency.

D. Data Partitioning

The datasets are partitioned into K unbalanced client subsets using Dir(α), producing substantial variation in per-class sample counts across clients. The visualization covers α = 0.5 across three datasets and α = 0.1 for CIFAR-100.

  • The employed datasets are partitioned into K unbalanced subsets using a Dirichlet distribution Dir(α).
  • At α = 0.5, the distribution is visualized for all three datasets, while CIFAR-100 is also shown at α = 0.1.The α = 0.1 CIFAR-100 setting is studied in Section 3.5 of the main paper.
  • Class sample counts vary considerably across clients, and clients often have no samples from some classes.
  • Figure 3 columns show the number of samples per class allocated to each client.

E. Additional Implementation Details

The implementation aligns feature-map spatial sizes before calculating XF and XS, rescales LLip relative to LCE before backpropagation, and applies gradient clipping.

  • Adaptive average pooling reduces the larger feature map to the smaller feature map’s spatial size when calculating XF and XS.The implementation uses PyTorch’s adaptive average pool operation for this alignment.
  • Before backpropagation, LLip is multiplied by the µ scaling parameter and relatively rescaled against LCE.The rescaling uses the ratio loss ce.item()/loss lip.item() multiplied by loss lip.
  • A gradient clip is applied during the loss-processing procedure.
Loading 2111.14213v3…