Stratify: Rethinking Federated Learning for Non-IID Data through Balanced Sampling
作者: Hui Yeok Wong, Chee Kau Lim, Chee Seng Chan
分类: cs.LG
发布日期: 2025-04-18
💡 一句话要点
Stratify：通过平衡采样重思考非独立同分布数据下的联邦学习
🎯 匹配领域: 支柱五：交互与反应 (Interaction & Reaction)
关键词: 联邦学习 非独立同分布数据 分层抽样 数据异构性 标签调度 客户端选择
📋 核心要点
- 现有联邦学习方法在非独立同分布数据上表现不佳，主要因为它们侧重于缓解数据异构性的表象，而忽略了FedAvg本身的设计缺陷。
- Stratify借鉴分层抽样思想，通过分层标签调度（SLS）和标签感知客户端选择，确保训练过程中各类标签的平衡暴露，从而解决数据异构性问题。
- 实验结果表明，Stratify在多个数据集上取得了与IID基线相当的性能，同时加速了收敛并降低了客户端计算负担。
📝 摘要（中文）
在非独立同分布（non-IID）数据上的联邦学习（FL）仍然是一个关键挑战，因为现有方法难以应对严重的数据异构性。目前的方法主要通过对联邦平均（FedAvg）进行增量调整来解决non-IID的症状，而不是直接解决其固有的设计局限性。因此，在高度异构的条件下，性能会显著下降，因为对不同类别和特征分布的不平衡暴露的根本问题仍然没有解决。本文介绍了一种新的FL框架Stratify，旨在系统地管理整个训练过程中的类别和特征分布，从而有效地解决non-IID挑战的根本原因。受到经典分层抽样的启发，我们的方法采用分层标签调度（SLS）来确保跨标签的平衡暴露，从而显著降低聚合梯度中的偏差和方差。作为SLS的补充，我们提出了一种标签感知的客户端选择策略，将参与限制在仅拥有与计划标签相关数据的客户端。此外，Stratify还结合了一种细粒度、高频率的更新方案，加速收敛并进一步减轻数据异构性。为了维护隐私，我们实施了一种利用同态加密的安全客户端选择协议，从而能够在不泄露敏感客户端信息的情况下获得精确的全局标签统计信息。在MNIST、CIFAR-10、CIFAR-100、Tiny-ImageNet、COVTYPE、PACS和Digits-DG上的广泛评估表明，Stratify获得了与IID基线相当的性能，加速了收敛，并减少了与最先进方法相比的客户端计算，突显了其在实际联邦学习场景中的实际有效性。
🔬 方法详解
问题定义：论文旨在解决联邦学习中非独立同分布（non-IID）数据带来的挑战。现有方法，如FedAvg及其变体，在数据高度异构的情况下性能显著下降，因为它们未能有效解决客户端数据分布不平衡的问题，导致模型训练过程中对不同类别和特征的暴露不均。
核心思路：Stratify的核心思路是通过模拟经典的分层抽样方法，在联邦学习过程中对不同类别的数据进行平衡采样。具体来说，它通过分层标签调度（SLS）来控制每个客户端参与训练的频率，确保模型在训练过程中能够充分学习到各个类别的信息，从而降低偏差和方差。
技术框架：Stratify框架主要包含以下几个模块：1) 分层标签调度（SLS）：根据全局标签分布，动态调整每个训练轮次中需要关注的标签。2) 标签感知客户端选择：只选择拥有与当前调度标签相关数据的客户端参与训练。3) 高频更新：采用更频繁的本地模型更新，加速收敛。4) 安全客户端选择协议：使用同态加密技术，在不泄露客户端隐私的前提下，实现标签感知的客户端选择。
关键创新：Stratify的关键创新在于其将经典的分层抽样思想引入联邦学习，通过显式地控制训练过程中各类数据的暴露比例，从根本上解决了非独立同分布数据带来的问题。与现有方法通过调整FedAvg的聚合方式不同，Stratify直接干预了客户端数据的选择和使用，从而更有效地缓解了数据异构性。
关键设计：Stratify的关键设计包括：1) SLS的调度策略：具体如何根据全局标签分布确定每个训练轮次的标签优先级。2) 标签感知客户端选择的阈值：如何确定客户端数据与当前调度标签的相关性。3) 安全客户端选择协议的同态加密方案：选择哪种同态加密算法，以及如何设计加密和解密过程以保证效率和安全性。这些细节在论文中应该有更详细的描述，但根据摘要无法得知具体实现。
🖼️ 关键图片
📊 实验亮点
Stratify在多个数据集（MNIST、CIFAR-10、CIFAR-100、Tiny-ImageNet、COVTYPE、PACS和Digits-DG）上进行了广泛的评估，结果表明其性能可与在独立同分布数据上训练的基线模型相媲美。此外，Stratify还表现出更快的收敛速度和更低的客户端计算负担，优于现有的联邦学习方法。
🎯 应用场景
Stratify适用于各种存在数据异构性的联邦学习场景，例如：医疗健康领域，不同医院的数据分布可能存在差异；金融领域，不同地区用户的交易数据可能存在差异；物联网领域，不同设备的传感器数据可能存在差异。该方法能够提升模型在这些场景下的泛化能力和鲁棒性，具有重要的实际应用价值。
📄 摘要（原文）
Federated Learning (FL) on non-independently and identically distributed (non-IID) data remains a critical challenge, as existing approaches struggle with severe data heterogeneity. Current methods primarily address symptoms of non-IID by applying incremental adjustments to Federated Averaging (FedAvg), rather than directly resolving its inherent design limitations. Consequently, performance significantly deteriorates under highly heterogeneous conditions, as the fundamental issue of imbalanced exposure to diverse class and feature distributions remains unresolved. This paper introduces Stratify, a novel FL framework designed to systematically manage class and feature distributions throughout training, effectively tackling the root cause of non-IID challenges. Inspired by classical stratified sampling, our approach employs a Stratified Label Schedule (SLS) to ensure balanced exposure across labels, significantly reducing bias and variance in aggregated gradients. Complementing SLS, we propose a label-aware client selection strategy, restricting participation exclusively to clients possessing data relevant to scheduled labels. Additionally, Stratify incorporates a fine-grained, high-frequency update scheme, accelerating convergence and further mitigating data heterogeneity. To uphold privacy, we implement a secure client selection protocol leveraging homomorphic encryption, enabling precise global label statistics without disclosing sensitive client information. Extensive evaluations on MNIST, CIFAR-10, CIFAR-100, Tiny-ImageNet, COVTYPE, PACS, and Digits-DG demonstrate that Stratify attains performance comparable to IID baselines, accelerates convergence, and reduces client-side computation compared to state-of-the-art methods, underscoring its practical effectiveness in realistic federated learning scenarios.