PaperlayerView exact revision on arXiv

Source-linked AI summary

FedMef: Towards Memory-efficient Federated Dynamic Pruning

Hong Huang, Weiming Zhuang, Chen Chen, Lingjuan Lyu

arXiv:2403.14737v1cs.LGcs.DC

TL;DR

Resource-constrained devices face high computation and memory demands when training federated models, and existing pruning approaches leave post-pruning accuracy and activation-memory challenges. FedMef combines budget-aware extrusion with scaled activation pruning. It reports a 28.5% reduction in MobileNetV2 memory footprint while improving accuracy by more than 2%.

  • Problem

    Training deep models for cross-device federated learning is challenging on resource-constrained devices, and existing pruning approaches face post-pruning accuracy degradation and high activation memory usage.

    Source: §1 · pp. 1–2

  • Method

    FedMef combines budget-aware extrusion, which transfers information from parameters marked for pruning within a preset budget, with scaled activation pruning to reduce activation-cache memory.

    Source: §1 · p. 2

  • Results

    28.5% lower MobileNetV2 memory footprint and more than 2% higher accuracy; FedMef outperforms state-of-the-art methods across three datasets and two models.

    Source: §1 · p. 2

  • Takeaways & Limitations

    FedMef is designed to generate specialized models for resource-constrained cross-device devices while addressing both accuracy and memory constraints.

    Source: §5 · p. 8

    The reported training-memory estimate excludes hyper-parameters and momentum.

    Source: §7.2 · p. 14

Abstract

from arXiv · show

Federated learning (FL) promotes decentralized training while prioritizing data confidentiality. However, its application on resource-constrained devices is challenging due to the high demand for computation and memory resources to train deep learning models. Neural network pruning techniques, such as dynamic pruning, could enhance model efficiency, but directly adopting them in FL still poses substantial challenges, including post-pruning performance degradation, high activation memory usage, etc. To address these challenges, we propose FedMef, a novel and memory-efficient federated dynamic pruning framework. FedMef comprises two key components. First, we introduce the budget-aware extrusion that maintains pruning efficiency while preserving post-pruning performance by salvaging crucial information from parameters marked for pruning within a given budget. Second, we propose scaled activation pruning to effectively reduce activation memory footprints, which is particularly beneficial for deploying FL to memory-limited devices. Extensive experiments demonstrate the effectiveness of our proposed FedMef. In particular, it achieves a significant reduction of 28.5% in memory footprint compared to state-of-the-art methods while obtaining superior accuracy.

1. Introduction

Cross-device federated learning faces high computation and memory costs, while existing federated dynamic pruning still struggles with post-pruning accuracy and activation memory. FedMef combines budget-aware extrusion and scaled activation pruning to address these challenges, and outperforms state-of-the-art methods across the reported datasets and models.

1. Introduction

  • Training deep models' high computation and memory costs challenge deployment of cross-device federated learning on resource-constrained devices.Prior federated pruning methods that first train dense models retain these resource demands.Source: §1 · p. 1
  • Existing federated dynamic pruning can degrade post-pruning accuracy by discarding informative low-magnitude parameters and does not reduce activation memory.The methods may not recover pre-pruning accuracy before the next pruning iteration.Source: §1 · p. 1
  • Scaled activation pruning uses Normalized Sparse Convolution to produce near-zero activations, then prunes activation caches by magnitude for gradient computation.The figure caption reports more than 3 times activation-memory savings on CIFAR-10 with MobileNetV2.Source: Figure 2 · p. 2, §1 · p. 2
  • Budget-aware extrusion transfers essential information from pruning candidates to other parameters through a surrogate loss within a preset budget.It is designed to preserve post-pruning accuracy rather than simply discard low-magnitude parameters.Source: §1 · p. 2
  • 28.5% lower MobileNetV2 memory footprint and more than 2% higher accuracy: FedMef outperforms state-of-the-art methods across three datasets and two models.The experiments use CIFAR-10, CINIC-10, and TinyImageNet with ResNet18 and MobileNetV2.Source: §1 · p. 2

2. Related Work

Related work describes pruning approaches that reduce model size but leave challenges in training cost, activation memory, and accuracy. In cross-device federated learning, these gaps remain across federated pruning and activation-cache compression methods.

2.1. Neural Network Pruning

  • Traditional pruning scores and discards parameters in a pretrained dense model, while dynamic sparse training adjusts pruned architectures during training but can neglect activation-cache memory and accuracy.The cited dynamic sparse methods maintain desired sparsity while adjusting model architecture.Source: §2.1 · pp. 2–3

2.2. Federated Neural Network Pruning

  • Earlier federated pruning often requires dense on-device training; dynamic on-device methods avoid that step but still leave activation-memory and post-pruning accuracy issues.FedMef is presented as the first work to address accuracy and memory constraints simultaneously.Source: §2.2 · p. 3

2.3. Activation Cache Compression

  • Activation-cache compression alternatives include costly recomputation and quantization overhead, while activation pruning relies on batch normalization and can be unsuitable for small batches and non-i.i.d. data.FedMef's scaled activation pruning compresses caches without relying on batch normalization.Source: §2.3 · p. 3

3. Methodology

FedMef targets post-pruning accuracy loss and activation-memory use in federated dynamic pruning with budget-aware extrusion and scaled activation pruning. It trains sparse models on resource-constrained devices while transferring information from parameters marked for pruning and pruning activation caches.

3.1. Problem Setup

  • The problem setup minimizes federated prediction loss over device-specific datasets while requiring model and activation-cache sparsity to meet memory-derived targets.Source: §3.1 · p. 3

3.2. Design Principles

  • FedMef iterates sparse on-device training with activation-cache pruning, budget-aware extrusion before server magnitude pruning, and gradient-based growth, then resumes training toward convergence.Source: §3.2 · pp. 3–4

3.4. Scaled Activation Pruning

  • Scaled activation pruning removes activation-cache entries after the forward pass and uses sparse caches for backward computation; small batches weaken BN and challenge DropIT’s zero-centering requirement.Source: §3.2 · p. 4, §3.4 · p. 5

3.3. Budget-aware Extrusion

  • Budget-aware extrusion uses a surrogate loss on low-magnitude parameters to transfer their information to surviving parameters before pruning, limiting information and accuracy loss.Source: §3.3 · p. 4
  • Its budget-aware learning rate scales the REX schedule by extrusion status and takes the maximum of the original and budgeted rates, addressing late-stage decay within a limited training budget.Source: §3.3 · p. 4

3.4. Scaled Activation Pruning

  • NSConv standardizes sparse filters without pruned weights; its γ parameter controls output variance, concentrating activations near zero to reduce the disparity between pruned and original activations.NSConv supports small-batch training without interdependencies between batch elements and keeps training and testing uniform.Source: §3.4 · pp. 5–6

4. Evaluation

FedMef is evaluated against federated pruning baselines across datasets, architectures, and constrained training settings. The section reports benchmark comparisons, cost measurements, and component ablations.

4.1. Experimental Setup

  • Experiments use CIFAR-10, CINIC-10, and TinyImageNet with ResNet18 and MobileNetV2 across 100 devices, with α = 0.5 partitions and 500 rounds.The default batch size is 64; comparisons include FL-PQSU, FedDST, and FedTiny, with structure adjustments every 10 rounds until round 300.Source: §4.1 · p. 6

4.2. Performance Evaluation

  • On CIFAR-10 with MobileNetV2, FedMef improves accuracy by 2.13% while saving 28.5% memory usage versus FedTiny, the best baseline.The comparison spans three datasets and two architectures; Figure 4 compares accuracy and memory footprint, with memory normalized to full-size FedAvg training.Source: §4.2 · p. 6, Figure 4 · p. 7

4.3. Ablation Study

  • FedMef's accuracy gains are attributed to budget-aware extrusion, while scaled activation pruning primarily augments memory conservation.Removing either component boosts performance less; removing activation scaling causes information loss during backpropagation and performance degradation.Source: §4.2 · p. 6, §4.3 · p. 8, Figure 5 · p. 8

4.2. Performance Evaluation

  • ResNet18 has higher accuracy than MobileNetV2 across all datasets, while FedTiny generally outperforms other baselines at comparable memory footprints.The authors therefore use ResNet18 as the default model and FedTiny as the primary reference in subsequent experiments.Source: §4.2 · pp. 6–7
  • At target density stm = 90%, FedMef adds 0.003× computational overhead and 0.005× communication overhead while improving accuracy by 2% versus other baselines.Table 1 reports maximum training FLOPs and data exchange per device in one adjustment round.Source: §4.2 · p. 7, Table 1 · p. 7
  • At batch size 1 on CIFAR-10 with ResNet18, FedMef outperforms all baseline methodologies.Figure 5's left panel compares frameworks at batch size 1 against the FedAvg accuracy reference.Source: §4.2 · p. 7, Figure 5 · p. 8
  • With one local epoch, FedMef maintains performance when FedTiny declines under resource-constrained adjustment periods; at ΔR=5, its convergence is more stable and final accuracy is higher.Table 2 reports mean accuracy and standard deviation for FedMef and FedTiny across adjustment periods; Figure 6 shows validation accuracy during training.Source: §4.2 · p. 8, Table 2 · p. 8, Figure 6 · p. 8
  • FedMef retains superior performance to FedTiny across tested degrees of data heterogeneity on CIFAR-10.Figure 5's middle panel compares FedMef with FedAvg and FedTiny as the Dirichlet factor α varies; lower α indicates greater heterogeneity.Source: §4.2 · p. 8, Figure 5 · p. 8

5. Conclusion

FedMef combines budget-aware extrusion and scaled activation pruning to address accuracy loss and activation memory in cross-device federated learning. Experiments show better accuracy and memory footprint than existing approaches.

5. Conclusion

  • FedMef uses budget-aware extrusion to transfer information from parameters marked for pruning to other parameters within a limited budget.Scaled activation pruning prunes activation caches after each layer’s forward pass and uses the sparse caches during backpropagation to save memory without compromising accuracy.Source: §5 · p. 8
  • FedMef is designed to generate specialized models on resource-constrained devices in cross-device federated learning.The framework addresses post-pruning accuracy degradation and high activation memory usage.Source: §5 · p. 8
  • Experiments show that FedMef outperforms existing approaches in both accuracy and memory footprint.Source: §5 · p. 8

6. Proof of Theorem 1

The proof analyzes convolution outputs through a linearized form and derives how normalized sparse convolution controls their mean and variance. Experiments report that NSConv behaves similarly to batch normalization, while unnormalized activation values contract with depth.

6.1. Internal Covariate Shift

  • The proof recasts convolution as a matrix multiplication by flattening filters into weights and stacking input patches into a linear input.The product is reshaped to recover the original convolution output.Source: §6.1 · p. 11, Figure 7 · p. 13
  • Without normalization, convolution output mean and variance scale with 1/c_in and 1/c_in^2, respectively.The analysis assumes ReLU inputs with positive mean and notes that filter means are difficult to keep at zero.Source: §6.1 · p. 11

6.2. Proof

  • NSConv centers convolution outputs around zero and sets their variance using γ, allowing many activation values to be zero or close to zero.For the stated ReLU-Conv setting, the theorem gives zero output mean and variance γ^2(σ_f^2 + μ_f^2).Source: §6.2 · p. 11

6.3. Experiment Result

  • NSConv behaves similarly to batch normalization in ResNet18, while unnormalized activation values decrease and become more centralized in deeper layers.The reported observations support the derived scaling of the output mean and variance.Source: §6.3 · p. 12

7. Memory, FLOPs and Communication Costs

The section estimates training memory, computation, and communication costs for FedMef and baselines, using density-dependent sparse storage schemes and explicit accounting assumptions.

7.1. Compression Schemes

  • Sparse storage switches among dense, bitmap, coordinate-offset, and CSR/CSC schemes according to parameter density.The section reshapes tensor data before compression to estimate parameter-training memory.Source: §7.1 · p. 13

7.2. The Memory Footprint of Training Models

  • Training-memory estimates include parameters, activations, activation gradients, and parameter gradients, with activation-gradient memory set equal to activation memory.The estimate omits hyper-parameter and momentum memory; FedMef’s scaled activation pruning reduces activation-cache memory, while activation gradients remain unpruned.Source: §7.2 · pp. 14–15

7.3. Training FLOPs

  • For FedMef, assuming dense and static-sparse inference FLOPs Fd and Fs and local iteration count E, maximum training FLOPs per round are 3(Fs + Fo)(E − 1) + (Fs + Fo) + 2Fd.Fo is the computing overhead of BaE and SAP, estimated as 4(1 − sm)nθ + na log na; 4(1 − sm)nθ represents regularization and WSConv FLOPs, while na log na denotes activation-pruning FLOPs.Source: §7.3 · p. 15

7.4. The Communication Cost

  • FedMef’s maximum data exchange per round is 2O_s + O_ξ, matching FedTiny’s pattern of exchanging sparse parameters and periodically uploading TopK gradients.The model mask requires no extra download overhead because the sparse storage formats encode element positions.Source: §7.4 · p. 15

8. More Experiments

The additional experiments examine FedMef across local-epoch and client-participation settings, regularization strengths, and comparisons with dense models.

8.1. Impact of Local Epochs Number

  • FedMef consistently outperforms other baselines as the number of local epochs varies, while very low local epochs substantially reduce FedDST performance.The evaluation uses CIFAR-10 with ResNet34; the authors note that FedDST needs sufficient local training before adjusting model architecture.Source: §8.1 · p. 15, Table 3 · p. 16

8.3. Impact of Regularization Coefficient λ

  • FedMef accuracy initially rises and then drops sharply as λ increases.The authors attribute the later decline to large λ values rapidly zeroing parameters and making the model excessively sparse.Source: Figure 9 · p. 16, §8.3 · p. 16

8.2. Impact of Selected Clients Number

  • FedMef outperforms other baselines in most selected-client settings, while FedDST performs substantially worse than FedMef and FedTiny.The evaluation varies clients selected per round using ResNet50.Source: §8.2 · p. 16, Table 4 · p. 16

8.4. Empirical comparison with dense model

  • On CIFAR-10, scaled activation pruning maintains dense-model performance while reducing activation memory, and FedMef with budget-aware extrusion outperforms FedAVG+SAP.Source: §8.4 · p. 16

9. More Detail and Information

The section reports the deployment hardware and explains why NSConv and unstructured pruning suit the paper’s low-memory setting.

9.1. Hardware information

  • Training is simulated on GPUs, while deployment is on a Raspberry Pi 4B with 1GB RAM.Source: §9.1 · p. 16

9.2. NSConv v.s. BN

  • NSConv outperforms BN at small batch sizes and matches BN at large batch sizes, with small batches relevant to low-memory devices.NSConv is described as more suitable for CNNs and sparse weights, but it may add computational overhead on dense weights.Source: §9.2 · p. 16

9.3. Structured pruning v.s. Unstructured pruning

  • Unstructured pruning can reach 80%+ sparsity while maintaining accuracy, whereas structured pruning can sharply reduce performance above 10% sparsity and has limited memory impact [41].Source: §9.3 · p. 16

Related papers

Loading 2403.14737v1…