The label distribution skew induced data heterogeniety has been shown to be a significant obstacle that limits the model performance in federated learning, which is particularly developed for collaborative model training over decentralized data sources while preserving user privacy. This challenge could be more serious when the participating clients are in unstable circumstances and dropout frequently. Previous work and our empirical observations demonstrate that the classifier head for classification task is more sensitive to label skew and the unstable performance of FedAvg mainly lies in the imbalanced training samples across different classes. The biased classifier head will also impact the learning of feature representations. Therefore, maintaining a balanced classifier head is of significant importance for building a better global model. To this end, we propose a simple yet effective framework by introducing a prior-calibrated softmax function for computing the cross-entropy loss and a prototype-based feature augmentation scheme to re-balance the local training, which are lightweight for edge devices and can facilitate the global model aggregation. The improved model performance over existing baselines in the presence of non-IID data and client dropout is demonstrated by conducting extensive experiments on benchmark classification tasks.

翻译：标签分布偏移导致的数据异质性已被证明是限制联邦学习模型性能的重要障碍，该技术专为在保护用户隐私的前提下对分散数据源进行协作模型训练而设计。当参与客户端处于不稳定环境并频繁退出时，这一挑战可能更为严峻。先前工作及我们的实证观察表明，分类任务的分类器头部对标签偏移更为敏感，而FedAvg不稳定的性能主要源于不同类别间训练样本的不平衡。有偏的分类器头部还会影响特征表示的学习。因此，维护平衡的分类器头部对于构建更优的全局模型至关重要。为此，我们提出一个简洁而有效的框架，通过引入先验校准的softmax函数计算交叉熵损失，以及基于原型的特征增强方案来重新平衡局部训练——这些方法对边缘设备轻量化，并能促进全局模型聚合。通过在基准分类任务上进行大量实验，我们证明了该方法在非独立同分布数据和客户端丢失场景下相较于现有基线的性能提升。