PaperlayerView exact revision on arXiv

Source-linked AI summary

Layer-wise Adaptive Model Aggregation for Scalable Federated Learning

Sunwoo Lee, Tuo Zhang, Chaoyang He, Salman Avestimehr

arXiv:2110.10302v3cs.LG

TL;DR

Periodic full aggregation overlooks differences in model discrepancy across layers and can spend bandwidth synchronizing similar parameters. FedLAMA adapts aggregation intervals layer by layer, and experiments report a smaller accuracy drop than periodic full aggregation as intervals increase, with comparable communication efficiency.

  • Problem

    Periodic full aggregation synchronizes all model parameters together despite layer-wise differences in model discrepancy, so synchronizing similar parameters can increase communication cost without meaningful training progress.

    Source: §1 · p. 1

  • Method

    FedLAMA estimates layer-wise model discrepancy at runtime and assigns longer aggregation intervals to layers with smaller discrepancy.

    Source: §1 · p. 1, §4 · p. 3

  • Results

    Across three image-classification datasets, FedLAMA shows a smaller accuracy drop than periodic full aggregation as intervals increase, while achieving comparable communication efficiency.

    Source: §1 · p. 2

  • Takeaways & Limitations

    FedLAMA adjusts aggregation frequency to balance communication efficiency and accuracy, increasing intervals only for layers with the weakest contribution to total model discrepancy.

    Source: §1 · p. 2

    The paper does not directly compare FedLAMA with compression methods and leaves combining layer-wise aggregation with them for future work.

    Source: §2 · p. 2

Abstract

from arXiv · show

In Federated Learning, a common approach for aggregating local models across clients is periodic averaging of the full model parameters. It is, however, known that different layers of neural networks can have a different degree of model discrepancy across the clients. The conventional full aggregation scheme does not consider such a difference and synchronizes the whole model parameters at once, resulting in inefficient network bandwidth consumption. Aggregating the parameters that are similar across the clients does not make meaningful training progress while increasing the communication cost. We propose FedLAMA, a layer-wise model aggregation scheme for scalable Federated Learning. FedLAMA adaptively adjusts the aggregation interval in a layer-wise manner, jointly considering the model discrepancy and the communication cost. The layer-wise aggregation method enables to finely control the aggregation interval to relax the aggregation frequency without a significant impact on the model accuracy. Our empirical study shows that FedLAMA reduces the communication cost by up to 60% for IID data and 70% for non-IID data while achieving a comparable accuracy to FedAvg.

1 Introduction

Periodic full aggregation ignores differences in model discrepancy across layers, potentially spending bandwidth synchronizing parameters that contribute little to training progress. FedLAMA adapts aggregation intervals layer by layer; experiments report a smaller accuracy drop as intervals increase, with comparable communication efficiency.

1 Introduction

  • Periodic full aggregation synchronizes all model parameters at once, despite layers having different degrees of discrepancy, which can waste bandwidth on similar parameters.The passage identifies limited network bandwidth as a concern for federated-learning scalability.Source: §1 · p. 1
  • FedLAMA assigns longer aggregation intervals to layers with smaller model discrepancy, using a runtime metric to estimate layer-wise discrepancy.It prioritizes layers by their contribution to total model discrepancy and adjusts intervals accordingly.Source: §1 · pp. 1–2
  • FedLAMA targets a practical trade-off between communication cost and accuracy by increasing intervals only for layers with the weakest contribution to total discrepancy.The paper also reports a theoretical convergence analysis for smooth, non-convex problems under non-IID data settings.Source: §1 · p. 2
  • Across three image-classification datasets, FedLAMA has a smaller accuracy drop than periodic full aggregation as intervals increase, while achieving comparable communication efficiency.The evaluation covers CIFAR-10, CIFAR-100, and FEMNIST.Source: §1 · p. 2

2 Related Works

Related work reduces communication through update compression, adaptive aggregation, and heterogeneous architectures, while commonly retaining periodic full-model aggregation. FedLAMA’s related work also discusses similarity measures and layer freezing, and leaves combining its approach with compression for future work.

2 Related Works

  • Compression replaces local updates with smaller representations through structured updates, such as low-rank approximation and random masks, or sketched updates, such as compression, sparsification, and quantization.The cited passage describes both categories as studied approaches to communication-efficient updates.Source: §2 · p. 2
  • FedLAMA does not directly compare its aggregation-frequency approach with compression methods, although compression can be applied to each layer’s update before aggregation.Combining layer-wise aggregation with compression is left as future work.Source: §2 · p. 2
  • CCA and its extension CKA estimate representational similarity, but these methods can be computationally expensive and do not indicate gradient differences relevant to convergence.The paper proposes a runtime-cheap metric for layer-wise model discrepancy.Source: §2 · p. 2
  • Layer-freezing methods stop updating layers in a bottom-up direction and can reduce training time without strongly affecting accuracy.The passage says these empirical techniques are supported by analysis of sequential layer convergence from the input side.Source: §2 · p. 2

3 Background

The background frames federated optimization in terms of local objectives and model synchronization. It describes FedAvg’s response to data heterogeneity, defines model discrepancy, and outlines the FedLAMA algorithm’s interval-adjustment steps.

3 Background

  • FedLAMA’s algorithm takes a base aggregation interval and an interval-increase factor, updates local layers, then adjusts intervals across layers and outputs the synchronized model.Source: §3 · p. 3
  • The optimization setup associates each client with a local objective defined over its local data distribution.The passage defines m as the number of local models and F_i(x) as the expected local objective for distribution D_i.Source: §3 · pp. 2–3
  • FedAvg uses periodic full aggregation, and its convergence slows as data heterogeneity increases.Source: §3 · p. 3
  • Local updates can move models in different parameter-space directions because of stochastic-gradient variance and heterogeneous data, producing model discrepancy.The passage notes that some discrepancy remains even with IID data because of stochastic-gradient variance.Source: §3 · p. 3

4 Layer-wise Adaptive Model Aggregation

FedLAMA prioritizes layers by their contribution to model discrepancy, then adjusts their aggregation intervals to balance discrepancy against communication cost. It lengthens intervals for layers with smaller discrepancy, using the balance point between the two factors to guide the adjustment.

4 Layer-wise Adaptive Model Aggregation

  • The layer-wise unit model discrepancy prioritizes layers by how much they contribute to total model discrepancy.Under a smoothness assumption, the difference between local and global gradients is bounded by a scaled difference between their parameters.Source: §4 · p. 3
  • FedLAMA estimates the contribution of the l layers with the smallest discrepancy scores and compares it with their parameter share to choose an interval setting.As l increases, the discrepancy contribution δ_l increases while 1 − λ_l decreases; the algorithm seeks a point where they are similar.Source: §4 · p. 4
  • For CIFAR-10 ResNet20, increasing intervals for 9 layers meets a 20% model-discrepancy increase with an 80% communication-cost decrease, both by factor φ.The crossing point is near 0.2; the figure also compares the two factors for CIFAR-100 Wide-ResNet28-10.Source: §4 · p. 4, Figure 1 · p. 5
  • FedLAMA can be extended to improve convergence rate with minor extra communications, but this work omits that setting because it can increase latency cost.The extension is intended for environments where latency cost can be ignored, such as high-performance computing platforms.Source: §4 · p. 4
  • Increasing an interval by factor φ reduces network-bandwidth consumption by φ without increasing communication latency cost.FedAvg is the special case φ = 1; for φ > 1, a subset of layers is synchronized less often, while the whole model is synchronized every φτ′ iterations.Source: §4 · pp. 4–5

5 Convergence Analysis

The convergence analysis uses smoothness, unbiased-gradient, bounded-variance, and bounded-dissimilarity assumptions to analyze FedLAMA. It establishes bounds for the algorithm and describes how interval size affects model discrepancy and communication frequency.

5.1 Preliminaries

  • The analysis assumes smooth local objectives, unbiased stochastic gradients, bounded gradient variance, and bounded dissimilarity across clients.When local objective functions are identical, the bounded-dissimilarity constants are β² = 1 and κ² = 0.Source: §5.1 · p. 5

5.2 Analysis

  • Under Assumptions 1–4 and a learning-rate condition, Theorem 5.3 bounds FedLAMA’s average-squared gradient norm, with τ defined by the largest layer aggregation interval.The theorem assumes all local models start from the same point.Source: §5.2 · p. 6
  • With sufficiently small diminishing learning rates and many iterations, FedLAMA achieves linear speedup.The analysis also states that when T > m³, the first term in its bound becomes dominant.Source: §5.2 · p. 6
  • A larger interval-increase factor φ increases model-discrepancy terms while proportionally reducing communication frequency at selected layers.The analysis recommends tuning φ to reduce communication cost without much increasing model discrepancy.Source: §5.2 · p. 6

6 Experiments

Across IID and non-IID benchmarks, FedLAMA maintains comparable classification performance while reducing communication relative to periodic full aggregation. The experiments also examine where communication savings arise and how they vary with data heterogeneity.

6 Experiments

  • The experiments use three image-classification benchmarks and simulate federated training by sequentially training multiple models per process before aggregating across processes.The evaluation uses 128 clients and averages accuracy across at least three runs; serialized training means communication cost, rather than wall-clock time, is considered.Source: §6 · p. 6

6.1 Classification Performance Analysis

  • FedLAMA achieves comparable classification performance with reduced communication on IID CIFAR-10/100 and non-IID FEMNIST and CIFAR-10/100.On non-IID CIFAR-100, it has a minor accuracy drop versus periodic full averaging at τ′ = 6, but remains above averaging at τ′ = 12.Source: §6.1 · p. 7

6.2 Communication Efficiency Analysis

  • FedLAMA increases aggregation intervals mostly at large output-side layers, whose share of model parameters makes them consequential targets for reducing communication cost.The reported extra computation is almost negligible because the layer discrepancy metric is calculated only after each communication round.Source: §6.2 · p. 7
  • The communication-cost reduction is similar across tested Dirichlet coefficients and device participation ratios, indicating effectiveness across the evaluated heterogeneity settings.The passage refers to Tables 4 and 5.Source: §6.2 · p. 8

7 Conclusion

FedLAMA adaptively adjusts aggregation intervals by layer and offers a convergence guarantee under stated assumptions. The conclusion reports empirical efficacy and points to combining FedLAMA with other methods as future work.

7 Conclusion

  • FedLAMA adaptively adjusts model aggregation intervals layer by layer at run-time.The paper also reports a convergence guarantee under assumptions on non-convex objectives and non-IID data.Source: §7 · p. 8
  • The conclusion reports a convergence guarantee under assumptions on non-convex objectives and non-IID data.The stated assumptions qualify the guarantee.Source: §7 · pp. 8–9
  • Figure 2 counts communications at individual layers over the full training run on non-IID data.The caption identifies the measure and setting, but does not state the plotted values.Source: Figure 2 · p. 9
  • Figure 3 compares total communication data size corresponding to Figure 2.Its caption says the comparison shows where FedLAMA's performance gain comes from, without specifying the plotted values.Source: Figure 3 · p. 9
  • Harmonizing FedLAMA with advanced optimizers, gradient compression, and low-rank approximation is identified as future work.Source: §7 · p. 9

A Appendix

The appendix introduces its proof material and summarizes the assumptions used in the analysis to make the appendix self-contained.

A Appendix

Assumptions

The analysis lists smoothness, unbiased stochastic gradients, bounded variance, and bounded dissimilarity among its assumptions.

Assumptions

  • Each local objective function is assumed to be L-smooth.Source: Assumptions · p. 12
  • Each client's stochastic gradient is assumed to be an unbiased estimator of its local full-batch gradient.Source: Assumptions · p. 12
  • Each client's stochastic gradient is assumed to have bounded variance.Source: Assumptions · p. 12
  • Bounded dissimilarity is parameterized by constants β2 ≥ 1 and κ2 ≥ 0.When local objectives are identical, β2 = 1 and κ2 = 0.Source: Assumptions · p. 12

A.1 Proofs

The proofs establish results under stated initialization, assumption, and learning-rate conditions, including bounds for the framework and model discrepancy. A CIFAR-10 figure contrasts FedAvg and FedLAMA learning curves under IID and non-IID data.

A.1 Proofs

  • Theorem 5.3 assumes all local models share initialization and relies on Assumptions 1–4.Source: §A.1 · p. 12
  • Theorem 5.3's learning-rate constraints are subsequently rearranged into a single constraint.The supplied passages do not state the resulting constraint's full expression.Source: §A.1 · p. 13
  • Lemma 5.1 gives a framework result under Assumptions 1–3 when η ≤ 1/(Lτ).The proof text also states the learning-rate constraint η ≤ 1/(τL).Source: §A.1 · pp. 13, 15
  • Lemma 5.2 addresses model discrepancy under Assumptions 1–4, with A := 2η^2L^2τ(τ−1) < 1.Source: §A.1 · p. 15
  • Figure 4 compares CIFAR-10 ResNet20 learning curves for IID and non-IID data across FedAvg and FedLAMA interval settings.Its caption reports that FedAvg loses convergence speed as its interval increases, while FedLAMA increases intervals while maintaining convergence speed.Source: Figure 4 · p. 17

A.2 Additional Experimental Results

Additional experiments show that FedLAMA preserves convergence and accuracy better than periodic full aggregation as aggregation intervals increase, across CIFAR-10, CIFAR-100, and FEMNIST. Its convergence speed is comparable to FedAvg at the base interval while communication cost is similar to FedAvg at an increased interval.

A.2 Additional Experimental Results

  • Across the extra experiments, FedLAMA matched convergence speed at base interval τ′ while using communication comparable to FedAvg at increased interval φτ′.The experiments used 128 clients, local batch size 32, and a 10-epoch learning-rate warm-up.Source: §A.2 · p. 17

Experimental settings

  • For CIFAR-10 and CIFAR-100, non-IID data used Dirichlet coefficients α = 0.1, 0.5, and 1.0; smaller α represented greater heterogeneity in client samples and label balance.The experiments used 128 clients, local batch size 32, and a 10-epoch learning-rate warm-up.Source: §A.2 · p. 17

CIFAR-10

  • On CIFAR-10, FedLAMA curves changed little as intervals rose from 6 to 24; at φ = 2, its accuracy matched or slightly exceeded FedAvg’s.FedAvg converged more slowly and reached lower validation accuracy as its interval increased; the trend held for IID and non-IID data.Source: §A.2 · p. 17, Table 6 · p. 18, Table 7 · p. 19
  • For CIFAR-10, accuracy dropped significantly as FedAvg’s base interval τ′ increased, while φ could be tuned to network bandwidth.With limited bandwidth, the authors suggest increasing φ and slightly increasing the epoch budget to achieve good accuracy.Source: §A.2 · p. 17, Table 8 · p. 19

CIFAR-100

  • On CIFAR-100, FedLAMA achieved accuracy comparable to short-interval FedAvg even with 25% device sampling and Dirichlet coefficient 0.1.FedAvg at τ ≥ 12 did not achieve reasonable accuracy under strongly heterogeneous data; its learning curves were also more affected by longer intervals.Source: §A.2 · p. 17, Figure 5 · p. 18, Table 9 · p. 19

FEMNIST

  • On FEMNIST, FedLAMA matched FedAvg with τ′ = 10 at interval increase factors φ ≥ 4.With increased intervals, FedLAMA’s curves were less affected and reached higher validation accuracy after the same number of iterations than periodic full aggregation.Source: §A.2 · p. 18, Figure 6 · p. 18, Table 12 · p. 20
Loading 2110.10302v3…