[论文解读] Re-Weighted Softmax Cross-Entropy to Control Forgetting in Federated Learning

Gwen Legate, Lucas Caccia+1 more|arXiv (Cornell University)|Apr 11, 2023

Privacy-Preserving Technologies in Data被引用 4

一句话总结

该论文提出加权Softmax（WSM），一种在联邦学习中对客户端的Softmax logits进行重加权的方法，以减轻由数据异质性引起的灾难性遗忘。通过在本地训练期间抑制分布外类别的梯度，WSM缓解了客户端漂移问题，并提升了全局模型性能，尤其在高数据异质性和低客户端参与度的情况下，CIFAR-100上准确率提升达7.6%，CIFAR-10上提升1.02%。

ABSTRACT

In Federated Learning, a global model is learned by aggregating model updates computed at a set of independent client nodes, to reduce communication costs multiple gradient steps are performed at each node prior to aggregation. A key challenge in this setting is data heterogeneity across clients resulting in differing local objectives which can lead clients to overly minimize their own local objective, diverging from the global solution. We demonstrate that individual client models experience a catastrophic forgetting with respect to data from other clients and propose an efficient approach that modifies the cross-entropy objective on a per-client basis by re-weighting the softmax logits prior to computing the loss. This approach shields classes outside a client's label set from abrupt representation change and we empirically demonstrate it can alleviate client forgetting and provide consistent improvements to standard federated learning algorithms. Our method is particularly beneficial under the most challenging federated learning settings where data heterogeneity is high and client participation in each round is low.

研究动机与目标

• 解决联邦学习中本地客户端遗忘的问题，即客户端在其本地数据上过拟合，从而遗忘其他客户端数据的表征。
• 识别出数据异质性和客户端参与度有限会加剧客户端漂移，导致全局模型性能下降。
• 提出一种轻量级、通信高效的缓解方法，无需数据共享或额外模型参数。
• 证明通过客户端特定的Softmax重加权修改交叉熵目标函数，可提升客户端间的泛化能力并增强收敛性。
• 在多种联邦学习算法（FedAvg、SCAFFOLD、FedProx）和数据集（CIFAR-10、CIFAR-100）下评估该方法，覆盖不同异质性和客户端参与度水平。

提出的方法

• 在计算交叉熵损失前，对每个客户端的Softmax logits实施基于类别的特定重加权方案。
• 采用修改后的交叉熵目标函数，对客户端标签集中以外类别的logits进行下采样，以降低其对本地更新的影响。
• 使用温度缩放的Softmax并结合类别特定权重，以稳定训练过程，并减少对分布外类别的梯度更新。
• 将加权Softmax（WSM）目标函数无缝集成到标准联邦优化算法（如FedAvg、SCAFFOLD、FedProx）中，无需修改聚合或通信协议。
• 设计权重函数以优先考虑分布内类别，同时抑制分布外类别的更新，从而在本地优化过程中最小化干扰。
• 确保该方法通信高效，并与隐私保护联邦学习兼容，避免数据共享或超出损失函数的模型正则化。

Figure 1: Illustration of catastrophic forgetting within client rounds. A global model with knowledge of all classes is sent to all clients participating in a given FL round. Local training increases the client model performance on the client’s local distribution but tends to simultaneously decrease

实验结果

研究问题

• RQ1 在非独立同分布的联邦学习设置下，本地客户端遗忘在多大程度上会损害全局模型性能？
• RQ2 在不依赖数据共享或模型正则化的情况下，对交叉熵损失中的Softmax logits进行客户端特定重加权，能否有效减少灾难性遗忘？
• RQ3 在高数据异质性和低客户端参与度——现实联邦学习中的关键挑战——下，所提出的WSM方法表现如何？
• RQ4 WSM是否不仅提升FedAvg，还能改善其他联邦优化算法（如SCAFFOLD和FedProx）的泛化能力并增强收敛性？
• RQ5 本地训练步数和客户端参与率对WSM在缓解遗忘方面有效性的影响如何？

主要发现

• WSM通过抑制分布外类别梯度更新，减少本地客户端遗忘，使客户端模型更加稳定和可泛化。
• 在高数据异质性（α=0.1）的CIFAR-100上，FedAvg+WSM达到27.4%的准确率，相比FedAvg的19.8%提升了7.6%。
• 在CIFAR-10上，FedAvg+WSM达到62.4%的准确率，相比FedAvg的60.8%提升了1.02%，且在最优学习率下表现最佳。
• 当客户端参与度较低时（如1%），FedAvg与FedAvg+WSM之间的性能差距最为显著，随着参与客户端增多，增益逐渐减小。
• 将本地迭代次数从7增加到21时，FedAvg的准确率下降更剧烈，而FedAvg+WSM表现出显著更稳定的性能，表明其具有更好的遗忘控制能力。
• WSM不仅提升了FedAvg的性能，也显著改善了SCAFFOLD和FedProx的性能，证明其在多种联邦优化框架中具有广泛适用性。

Figure 2: CIFAR-10 - Distributions of ten randomly selected clients with data partitioned according to a Dirichlet distribution parameterized by $\alpha=0.1$ .