
这篇论文是新加坡国立大学的博士生施宇钧发表在ICLR2023上的,主要研究了数据异构问题影响联邦学习模型的原因在于表征维度的坍缩,并提出了改进方案。这篇论文写的深入浅出,是真正的Understanding类型的文章,值得学习。
链接:https://arxiv.org/abs/2210.00226
1 介绍
1.1 论文主要工作
指出了non-iid对联邦学习的影响机理:non-iid导致了表征的维度坍缩,由此影响了全局模型的性能。
1.2 全文逻辑
- 数据异构的影响首先体现在全局模型表征的维度崩塌
- 每个客户端上的本地模型也观察到了维度崩塌的现象
- 对本地模型进行了梯度流动动力学分析,从理论上阐述了non-iid导致本地模型维度崩塌的过程。
- 针对以上的观察,对正则项进行改进。
1.3 相关论文
杨立昆组首先用类似的方法研究了对比学习的维度崩塌情况
Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian. Understanding dimensional collapse in contrastive self-supervised learning. arXiv preprint arXiv:2110.09348, 2021.
2 观察过程
2.1 全局模型实验
设置了不同non-iid的情况:$\alpha={0.01, 0.05, 0.25, \infty}$
计算模型输出表征的协方差矩阵:$\sum = \frac{1}{N} \sum^N_{i=1}(z_i - \bar{z})(z_i - \bar{z})^T$, $\bar{z}=\frac{1}{N}\sum^N_{i=1} z_i$
然后对每个协方差矩阵进行奇异值分解,排列奇异值的对数:

x轴是奇异值的索引,y轴是奇异值的对数
曲线只有前面少数的奇异值较大,后面大部分奇异值较小,即表征分布被压缩在少数特征方向上,即维度坍缩现象。
实验表示:non-iid程度越高,维度探索的坍缩就越严重。
2.2 本地模型实验
推断:全局模型的维度坍缩是从本地模型继承过来的。
在本地模型上使用了跟实验1类似的方法,针对不同程度non-iid程度下得到的局部模型进行了可视化:

Local和Global趋势相同:数据异质性越强,维度崩溃越明显
因此得出结论,全局模型的维度坍缩是从本地客户端继承过来的
3 维度坍缩的理论分析
通过分析在本地训练过程中模型权值的梯度流动力学,分析显示异构本地训练数据会驱动模型的权值向低秩发展,从而导致表征的维度崩塌。
t:优化时间步长
$z_i(t)$输出的表征向量
将前L层的权重矩阵的乘积表示为:
$$ \prod(t) = W_{L}(t)W_{L-1}(t)...W_{1}(t) $$
设$\sigma_k(t)$$(k \in[d])$是$\prod(t)$的第$k$大的奇异值,则有:

其中$G(t)$为:

$\mu_c$是$c$类数据所占数据总量的比例。当non-iid程度高的时候,某些类的$\mu_c$将趋于0。
当越来越多的$\mu_c$趋向于0的时候,$G(t)$会偏向于低秩矩阵。
当$G(t)$低秩时,式(3)中的$\sigma_k(t)$只有在$k$值较小时数值才会很大。因此$\prod(t)$会演变成低秩矩阵
而协方差矩阵$\sum(t)$可以用$\prod(t)$表示为:

因此若$\prod(t)$演变为低秩矩阵,协方差矩阵$\sum(t)$也会倾向于低秩。
4 解决方案
4.1 用正则项约束奇异值
由于Global上的维崩塌源于Local,因此从Local的训练过程来缓解问题。
在本地训练时,在表征上添加以下正则化项:

$\lambda_i$是表征的协方差矩阵的第$i$个奇异值,该式对奇异值之间的方差进行了正则,保证了尾部的奇异值不会折叠到0。
该式的缺点:需要算出来所有的奇异值,开销大
4.2 降低成本,引入F范数
首先对所有表征向量$z_i$应用z值归一化:$\hat{z}_i=(z_i-\bar{z})/ \sqrt{Var(z)}$,使$z_i$的协方差矩阵等于其相关系数矩阵。下面的建议提出了一个更方便的成本函数来规范化。
命题1:对于一个具有奇异值的d行d列矩阵,有:

$\|K\|_F$是F范数。作者证明了正则化相关系数矩阵的$\|K\|_F$具有与最小化奇异数相同的效果。与奇异值相比,$\|K\|_F$更易计算。
4.3 FedDecorr
利用上面的命题,在训练过程中对表征向量相关矩阵的F范数进行正则化:

本地客户端的总目标变为:

5 实验验证
数据集:CIFAR10、CIFAR100、TinyImageNe
Baseline:FedAvg、FedAvgM、FedProx、MOON
Non-iid设置:$\alpha \in \{{0.05, 0.1, 0.5, \infty}\}$
5.1 是否缓解了维度坍缩

缓解了Local的维度崩溃,且其改善的结果也继承给了Global
5.2 对Baseline的改进效果


推测:除了标签的异质性之外,还有其他因素会导致FL中的维度崩溃
5.3 对通信效率的提高

5.4 消融实验
客户端数量的影响

客户端数量越多,提升效果越显著。
正则项系数$\beta$的影响

5.5 附录实验
计算效率
一轮训练中花费的总时间:

其他non-iid设置的实验

小结
这篇论文应该是ICLR最喜欢的文章形式,先针对一个现象进行实验和观察,探究出现这个现象的原因,再根据这个现象进行改进,文中大部分篇幅都在于探究的过程,而其方法的改进反而非常精简,比如这篇论文他的改进方案只是添加了一个正则项这么简单,但是对实验结果的改进却很显著。
如果非要吹毛求疵的话,那就是他在附录中进行的non-iid分片实验中,并没有进行$C=1$的实验,他的这个方法应该也对$C=1$的情况束手无策,毕竟本地模型只有一个类的表征,再怎么正则也拦不住他奇异值变0。
总的来说,这篇论文有很多地方都值得我学习,也是自己努力的目标。
Comments | NOTHING