2025-11-11T23:28:21.956833

Bridging Memory Gaps: Scaling Federated Learning for Heterogeneous Clients

Wu, Li, Tian et al.
Federated Learning (FL) enables multiple clients to collaboratively train a shared model while preserving data privacy. However, the high memory demand during model training severely limits the deployment of FL on resource-constrained clients. To this end, we propose \our, a scalable and inclusive FL framework designed to overcome memory limitations through sequential block-wise training. The core idea of \our is to partition the global model into blocks and train them sequentially, thereby reducing training memory requirements. To mitigate information loss during block-wise training, \our introduces a Curriculum Mentor that crafts curriculum-aware training objectives for each block to steer their learning process. Moreover, \our incorporates a Training Harmonizer that designs a parameter co-adaptation training scheme to coordinate block updates, effectively breaking inter-block information isolation. Extensive experiments on both simulation and hardware testbeds demonstrate that \our significantly improves model performance by up to 84.2\%, reduces peak memory usage by up to 50.4\%, and accelerates training by up to 1.9$\times$.
academic

Bridging Memory Gaps: Scaling Federated Learning for Heterogeneous Clients

基本信息

  • 论文ID: 2408.10826
  • 标题: Bridging Memory Gaps: Scaling Federated Learning for Heterogeneous Clients
  • 作者: Yebo Wu, Jingguang Li, Chunlin Tian, KaHou Tam, Li Li, Chengzhong Xu (澳门大学)
  • 分类: cs.DC (分布式计算)
  • 发表时间: 2024年8月 (arXiv v2: 2025年10月)
  • 论文链接: https://arxiv.org/abs/2408.10826v2

摘要

联邦学习(FL)使多个客户端能够在保护数据隐私的同时协作训练共享模型。然而，模型训练过程中的高内存需求严重限制了FL在资源受限客户端上的部署。为此，本文提出SCALEFL，一个可扩展且包容的FL框架，通过顺序块式训练克服内存限制。SCALEFL的核心思想是将全局模型分割为块并顺序训练，从而减少训练内存需求。为缓解块式训练中的信息损失，SCALEFL引入课程导师(Curriculum Mentor)为每个块制定课程感知的训练目标。此外，SCALEFL集成训练协调器(Training Harmonizer)设计参数协同适应训练方案，有效打破块间信息隔离。

研究背景与动机

核心问题

  1. 内存墙问题：联邦学习训练过程中需要在内存中保留所有中间激活、模型权重和优化器状态，导致高内存占用。例如，在ImageNet上训练ResNet34消耗超过12GB内存，而现成移动设备的RAM通常仅为4-12GB。
  2. 设备异构性：资源受限的边缘设备无法参与本地训练，阻止了其宝贵数据对全局模型的贡献。
  3. 现有方法局限性：
    • 模型异构训练：需要高质量公共数据集进行知识蒸馏，在FL中难以获得
    • 部分训练：宽度缩放破坏模型架构，深度缩放受最大内存容量客户端限制

研究动机

随着模型架构变得更深更宽以获得更高分析能力，内存问题进一步恶化。本文旨在设计一个既能显著降低内存需求又能保持模型性能的FL框架。

核心贡献

  1. 提出SCALEFL框架：通过顺序块式训练显著降低训练内存需求，使资源受限设备能够有效参与
  2. 设计两个核心组件：课程导师和训练协调器协同塑造每个块的学习行为，促进连贯的结构化特征学习
  3. 全面实验验证：在多个基准数据集上证明SCALEFL的有效性和鲁棒性
  4. 理论分析：提供收敛性分析，证明方法的理论可靠性

方法详解

任务定义

在包含N个客户端的FL系统中，每个客户端n拥有本地数据集Dn。目标是训练全局模型Θ，同时满足各客户端的内存约束。

顺序块式训练范式

基本流程：

  1. 模型构建：服务器为当前训练阶段t构建子模型 Θg,t = θ1,F, θ2,F, ..., θt, θOp
  2. 本地训练：仅更新块θt和输出模块θOp
  3. 模型聚合：使用加权平均聚合参数更新
  4. 进度评估：监控块θt的训练进度并判断收敛
  5. 模型增长：冻结收敛块并引入新块

核心技术组件

1. 课程导师 (Curriculum Mentor)

问题分析：基于信息瓶颈理论，发现顺序块式训练导致严重信息损失。通过nHSIC平面动态分析显示，SBT在第一个块训练后丢失大量输入信息，导致后续块无法提取关键特征。

解决方案：设计课程感知训练目标

L_θt = L_CE - λt · nHSIC(X;Zt) - γt · nHSIC(Y;Zt)

其中：

  • L_CE为交叉熵损失
  • nHSIC(X;Zt)衡量输入信息保留
  • nHSIC(Y;Zt)衡量任务相关性
  • λt和γt根据训练阶段动态调整

策略：初期较高λt和较低γt强调输入信息保留，后期逐渐降低λt增加γt转向任务特定特征提取。

2. 训练协调器 (Training Harmonizer)

问题识别：

  • 前向信息流受限：下游块仅在前序块收敛后开始训练
  • 反向信息流受限：梯度局限于块内，导致梯度隔离

参数协同适应方案：

  1. 动态模型增长：每轮动态编排各块学习过程，使下游块能实时适应前序块更新
  2. 并发训练策略：当前块与前序块的后几层同时训练，促进梯度流动

更新公式：

θ^(k+1)_(n,t) + L^(k+1)_(n,t-1) ← (θ^k_(n,t) + L^k_(n,t-1)) - η · ∂L^k_(n,t)/∂(θ^k_(n,t) + L^k_(n,t-1))

完整训练目标

结合L2正则化处理数据异构性：

L^r_t = L_θt + (μ/2)||θ^r_t - θ^(r-1)_t||^2_2

实验设置

数据集

  • CIFAR10/CIFAR100：经典图像分类数据集
  • CINIC10：CIFAR10扩展版本
  • Mini-ImageNet：小规模ImageNet
  • FEMNIST：大规模FL数据集(805,263张图像)

模型架构

  • ResNet18/ResNet34：深度残差网络
  • VGG11 BN：经典卷积网络
  • SqueezeNet：轻量级网络
  • Vision Transformer (ViT)：Transformer架构

实验环境

  • 混合设置：仿真和真实设备测试床
  • 设备配置：100个异构移动设备，每轮随机选择10%参与
  • 内存预算：100-1000MB随机分配
  • 优化器：SGD，权重衰减5e-4，本地epoch=5

对比方法

  1. AllSmall：基于最弱设备内存缩小全局模型
  2. ExclusiveFL：仅允许内存充足设备参与
  3. DepthFL：深度缩放适应异构设备
  4. HeteroFL：静态通道缩放
  5. FedRolex：动态宽度缩放
  6. SmartFreeze：简单顺序块式训练
  7. ProFL：分解式顺序训练

实验结果

主要结果

非IID场景下的性能表现：

方法CIFAR10 (ResNet18/VGG11/SqueezeNet)参与率
AllSmall69.5%/75.1%/49.6%100%/100%/100%
ExclusiveFL76.8%/79.3%/40.6%18%/22%/11%
SCALEFL80.4%/87.6%/58.0%100%/100%/100%

关键发现：

  1. 显著性能提升：相比AllSmall提升10.9%、12.5%、8.4%
  2. 全设备参与：实现100%设备参与率，而ExclusiveFL仅18-22%
  3. 内存效率：峰值内存使用减少高达50.4%
  4. 训练加速：收敛速度提升1.9倍

扩展性分析

不同内存约束下的鲁棒性：

  • ResNet34场景下，ExclusiveFL完全不可行(0%参与率)
  • SCALEFL相比其他方法提升高达27.4%

大规模数据集：

  • FEMNIST数据集上比FedAvg提升3%准确率
  • 支持120-500个设备规模

Transformer兼容性：

  • ViT模型上仅比理论基线低2%，但理论基线在实际中不可行

硬件评估

内存效率：

  • Jetson TX2上测试显示峰值内存使用减少50.4%
  • 单轮训练时间减少1.84-2.31倍

训练效率：

  • 相比端到端训练显著减少单轮训练时间
  • ViT上实现1.9倍加速

消融实验

组件贡献分析：

  • 移除课程导师：CIFAR100 IID场景下准确率下降1.2%
  • 移除训练协调器：准确率显著下降9.0%
  • 两个组件协同作用对性能至关重要

相关工作

资源受限FL

  1. 模型异构训练：FedMD等方法需要公共数据集进行知识蒸馏
  2. 部分训练：HeteroFL、FedRolex通过宽度缩放，DepthFL、InclusiveFL通过深度缩放

块式训练

  1. ProgFed：渐进式引入新块但仍需端到端训练
  2. SmartFreeze：顺序训练每个块但忽略信息损失问题
  3. ProFL：分解为收缩和增长阶段但未解决核心挑战

理论分析

收敛性证明

论文提供了SCALEFL的收敛性分析，在标准假设下(平滑性、有界梯度)证明：

(1/R) Σ E[||∇L^r_t(Θ^r_(g,t))||^2] ≤ Ψ/√R

即平均梯度范数收敛到0，模型收敛到稳定点。

结论与讨论

主要结论

  1. SCALEFL成功解决了FL中的内存墙问题，使资源受限设备能够参与训练
  2. 课程导师和训练协调器有效缓解了顺序块式训练的核心挑战
  3. 在多个数据集和模型上实现了显著的性能提升和内存节省

局限性

  1. 块划分策略：论文未深入讨论最优块划分方法
  2. 通信开销：虽然减少了内存使用，但可能增加通信轮数
  3. 超参数敏感性：λt和γt的设置需要仔细调优

未来方向

  1. 自适应块划分策略
  2. 与其他FL优化技术的结合
  3. 在更大规模实际部署中的验证

深度评价

优点

  1. 问题重要性：解决了FL实际部署的关键瓶颈
  2. 方法创新性：课程感知训练目标和参数协同适应方案具有原创性
  3. 理论基础：基于信息瓶颈理论的分析提供了坚实的理论支撑
  4. 实验全面性：涵盖多种模型、数据集和真实硬件测试
  5. 实用价值：显著的内存节省和性能提升具有实际应用价值

不足

  1. 复杂性：引入的两个组件增加了系统复杂性
  2. 超参数调优：λt、γt等参数需要针对不同场景调优
  3. 通信分析：缺乏对通信开销的详细分析
  4. 收敛速度：虽然单轮更快，但总体收敛轮数可能增加

影响力

  1. 学术贡献：为资源受限FL提供了新的解决思路
  2. 实用价值：能够实际部署在移动设备等资源受限环境
  3. 可复现性：提供了详细的实验设置和参数配置

适用场景

  1. 移动设备FL：智能手机、IoT设备等内存受限场景
  2. 边缘计算：边缘服务器资源有限的环境
  3. 大模型训练：需要训练大型模型但设备资源不足的场景

参考文献

论文引用了FL领域的重要工作，包括FedAvg、HeteroFL、FedRolex等经典方法，以及信息瓶颈理论、HSIC等理论基础，文献引用较为全面和权威。

总体评价：这是一篇高质量的联邦学习论文，针对实际部署中的关键问题提出了创新解决方案。方法设计合理，实验验证充分，理论分析完整，具有重要的学术价值和实用价值。