[论文解读] Towards Understanding and Mitigating Dimensional Collapse in Heterogeneous Federated Learning

Yujun Shi, Jian Liang+3 more|arXiv (Cornell University)|Oct 1, 2022

Privacy-Preserving Technologies in Data被引用 12

一句话总结

本文识别并解决了异构联邦学习中的维数坍缩问题，其中数据异质性导致表征退化为低维流形。作者提出FedDecorr，一种轻量级、即插即用的正则化方法，在本地训练过程中最小化表征维度之间的相关性，显著提升了在强数据异质性条件下CIFAR10、CIFAR100和TinyImageNet上的模型性能。

ABSTRACT

Federated learning aims to train models collaboratively across different clients without the sharing of data for privacy considerations. However, one major challenge for this learning paradigm is the {\em data heterogeneity} problem, which refers to the discrepancies between the local data distributions among various clients. To tackle this problem, we first study how data heterogeneity affects the representations of the globally aggregated models. Interestingly, we find that heterogeneous data results in the global model suffering from severe {\em dimensional collapse}, in which representations tend to reside in a lower-dimensional space instead of the ambient space. Moreover, we observe a similar phenomenon on models locally trained on each client and deduce that the dimensional collapse on the global model is inherited from local models. In addition, we theoretically analyze the gradient flow dynamics to shed light on how data heterogeneity result in dimensional collapse for local models. To remedy this problem caused by the data heterogeneity, we propose {\sc FedDecorr}, a novel method that can effectively mitigate dimensional collapse in federated learning. Specifically, {\sc FedDecorr} applies a regularization term during local training that encourages different dimensions of representations to be uncorrelated. {\sc FedDecorr}, which is implementation-friendly and computationally-efficient, yields consistent improvements over baselines on standard benchmark datasets. Code: https://github.com/bytedance/FedDecorr.

研究动机与目标

• 研究数据异质性如何影响联邦学习中的表征学习。
• 识别在数据异质性下全局和本地模型中的维数坍缩作为关键故障模式。
• 从理论上分析将数据异质性与低秩权重矩阵及表征坍缩联系起来的梯度流动力学机制。
• 提出一种实用且计算高效的缓解联邦学习中维数坍缩的方法。

提出的方法

• 提出FedDecorr，一种在本地训练期间最小化表征相关性矩阵的Frobenius范数的正则化项。
• 在计算相关性矩阵之前对表征进行z分数标准化，以稳定训练过程。
• 将FedDecorr作为即插即用的损失组件集成，可对任何现有联邦学习基线模型进行最小修改。
• 使用归一化相关性矩阵的平方Frobenius范数作为正则化损失，以鼓励表征之间不相关。
• 使用主任务损失与FedDecorr正则化损失的组合训练本地模型，然后通过标准FedAvg进行聚合。
• 证明FedDecorr计算效率高，且与现有联邦学习框架完全兼容。

实验结果

研究问题

• RQ1 数据异质性如何影响联邦学习模型中表征的内在维度？
• RQ2 全局模型中的维数坍缩在多大程度上源自在异质数据上训练的本地模型？
• RQ3 在数据异质性下，导致低秩权重矩阵和表征坍缩的底层梯度动力学机制是什么？
• RQ4 一个简单且轻量级的正则化项能否在不增加计算或通信成本的前提下有效缓解联邦学习中的维数坍缩？

主要发现

• 更强的数据异质性导致更严重的维数坍缩，表现为表征协方差矩阵的奇异值更趋近于零。
• 维数坍缩不仅出现在全局模型中，也出现在本地训练的模型中，表明其根源在于本地数据分布的系统性问题。
• 理论分析表明，数据异质性会使本地模型权重偏向低秩结构，这直接导致表征坍缩。
• FedDecorr能有效降低表征相关性并缓解维数坍缩，从而在多个基准数据集上实现一致的性能提升。
• 在高数据异质性或客户端数量较多等挑战性设置下，该方法带来的性能提升更为显著。
• FedDecorr引入的计算开销可忽略不计，可无缝集成到任何联邦学习流程中作为即插即用的组件。