论文分享：「FED BN」使用 LOCAL BATCH NORMALIZATION 方法解决 Non-iid 问题

‍

‍

图片

更多干货内容，请关注公众号：隐语的小剧场

本次分享内容基于 ICLR 2021 收录的一篇文章：《FED BN: FEDERATED LEARNING ON NON-IID FEATURES VIA LOCAL BATCH NORMALIZATION》，这篇论文主要探讨了使用 LOCAL BATCH NORMALIZATION 方法解决 Non-iid 问题。围绕这篇论文的分享将分为 4 个部分：

1、BATCH NORMALIZATION 及其解决 Non-iid 问题的方法；

2、Non-iid 问题及常用的 sota 解决方法；

3、FED BN 及其解决 Non-iid 问题的方法；

4、围绕论文实验分析 FED BN 效果与价值。

话题一、BATCH NORMALIZATION 及其解决 Non-iid 问题的方法

1. BATCH NORMALIZATION

BATCH NORMALIZATION 是在明文机器学习中比较常用且效果较好的一种方法，它可以很好的解决 feature scaling 的问题，包括层与层之间 internal Covariate Shift 的问题。

那么什么是 feature shift 呢，我们来举个例子

图片

如图任务是识别图片内容是不是玫瑰花，右侧绿色小圈是玫瑰花，红色则不是，蓝色则表分类边界。虽然上下 2 组图的任务相同，但是肉眼可见其特征分布明显不同。这种分布状态在训练时就会导致收敛困难，需要更多的步骤来抵消 scale 不同所带来的影响才能最终收敛。这种 feature scale 差异在和权重进行矩阵相乘时，会产生一些偏离较大的差异值，这个差异值在网络中会影响后面的层，并且偏离越大影响越明显。BN 的作用则是将输入标准化，缩小 scale 的范围。其结果就是可以提高梯度的收敛程度，并提升训练速度。

那么 BN 是怎么来的呢？

在早些时期，网络相对扁平的时候有一种方法 Whiten（白化）：对输入数据做单位方差的正态分布转换，可以加速网络收敛。在深度网络中随着网络深度的不断增加，每层特征值分布会逐渐的向激活函数的输出区间的上下两端（激活函数饱和区间）靠近，如此继续就会导致梯度消失，影响训练收敛。一些工作包括 Relu, Resnet 和 BN 等就是尝试解决这个问题。既然 whiten 可以加速收敛，那么是否深度网络中对每一层做 whiten 就可以加速收敛？由这种设想便产生了 BN，其作用就是保证机器学习的 iid 假设，保证每个隐层节点的激活输入分布固定，降低 feature shift 对模型训练造成的影响。

那 BN 怎么做？

训练阶段

图片

1、取 batch 数据计算均值

2、取 batch 数据计算方差

3、做 normalize

4、通过训练得到 scale 和 shift 参数

这里 normalize 在后的过程中还涉及了两个参数：γ及β，一个对应 scale，一个对应 shift，是对其进行一个平移操作的参数。为什么需要这两个参数？我们以 sigmoid 激活函数来进行举例；论文作者在设计时默认将 BN 添加在激活函数之前，如果在得出 normalize 后不添加线性变化，则会导致大部分输入结果停留在线性空间。但线性空间的多层叠加却是无效的，相较之一层线性网络没有任何区别，即削弱了整个网络的表达能力。故而需要添加γ及β参数进行转换，使其找到线性和非线性之间的平衡点，如此既能享受到非线性表达的好处，又可以避免它落在两端影响收敛速度。

图片

在预测流程中，则没有 batch 的概念，我们需要使用全部的待预测数据求均值，然后求抽样方差，最后使用这两个参数，对其进行转化并在预测中使用。

优点小结：

1、可以极大地提升训练速度及收敛速度，可以使用较大的学习率

2、可以增强学习的效果

3、调参过程会变简单，因为 BN 的使用对初始值的敏感度降低

缺点小结：

1、训练和预测的结果存在些许的不一致

2、单个 batch 的计算需要增加计算量

话题二、iid

iid，即独立同分布，假设训练数据和测试数据满足相同分布，是通过训练数据获得的模型能够在测试集上获得较好效果的基本前提保证。但是在部分联邦场景下，数据源不同使得这一假设很难被满足。关于 iid 的分类参考相关论文可划分为五种情形：

1、Feature distribution skew，即特征分布偏移（以前文玫瑰花识别为例，x 的特征不一致，但 label 分布一致）；

2、Label distribution skew，即特征分布呈现一致，但每一方的 label 都是不相同（如不同的人所认得的字的合集不同）；

3、Quantity skew，即数据量不同；

4、Same label but different feature，即我们所熟知的垂直场景，每一方都有一些数据，这些数据共享一个 label；

5、Same feature but different label，特征相同但 label 不同，即多任务学习。

图片

如上图（图引自于）所示，将数据及常用方法与对应可解决的问题进行划分，FED BN 解决的就是特征分布偏移的问题。

解决 iid 问题的经典联邦算法

1. Fed-AVG

Fed-AVG 是联邦学习中的经典算法之一，主要解决两个问题：通信问题和 Client 数量问题。

通信问题，即相较于 SGD 方法需要在每一代迭代之后将梯度（gradient）或权重（weight）发至服务端进行聚合，需要增加很多通信量和计算量。Fed-AVG 则允许在客户端在本地完成一些 step 之后再到服务端做聚合，通过增加客户端的计算减少整体的通信，并且 Fed-AVG 并不需要每一次所有 Client 都参加，进一步可以降低通信和计算量。该工作从理论上证明了 Fed-avg 的收敛性。

Fedavg 的收敛路径如下图所示：

图片

图片

Fed-AVG 的流程为：

1、server 把全局模型下发至每个 client

2、各个 client 使用 local 数据通过等其他优化方法进行迭代，得到 gradient 或 weight，并加密发送给 server

3、Server 对收到的 gradient 和 weight 进行聚合得到 global g/w，发送给各个 clients

4、Client 收到 global g/w 对其本地的模型进行更新

2. FedProx

FedProx 是在 Fed-AVG 基础上进行了一些演化，作者引入了 proximal term 这一概念做约束，将学习目标由 F（x） 变为了 H（x），目的是使得本地更新不要太过远离初始 global model，在容忍系统异构性的前提下减少 Non-IID 的影响。用γ作为本地迭代的 proxy，值越小更新精度越高。

图片

FedProx 整体流程基本上和 fed-avg 是一样的，差异在于 loss 函数做了修改，增加了约束项。

图片

所以每轮下降的梯度变成了

图片

3. Scaffold

Scaffold 也是从 fed-avg 的基础上演化过来的 non-iid 策略，通过在随机过程 control variates（guess）来克服 client drift 来进行联邦学习。scaffold 想既能看到自己的数据，又能看到 server 的数据啊，这样的话自己学习的方向就可以尽量往 global 这边靠，避免发散不收敛。但是 server 的数据是看不到的，所以就去猜，为了猜 server 的梯度方向，作者给了 C 和 Ci 的概念，Ci 是梯度， C 是合相梯度，这里假设合相梯度是 server 想要优化的方向，拿合相梯度和本次计算出来的梯度求了一个差，就可以得到优化应该去的方向。所以对于这个算法，不仅要更新参数，还要更新猜测项。

算法流程如下：

图片

可以看到和 fed-avg 的区别在于多了 ci 和 c，每一个 client 计算的 ci 是对本方多轮迭代的 gradient 进行平均得到的，server 的 C 是对各个方面得到的 ci 进行平均得到。

图片

话题三：FedBN

FedBN 这篇工作最大的贡献在于他提出了如何在联邦学习场景下使用 batch normalization 策略，来解决 feature shift 问题，同时可以利用 BN 来加速训练。

•文章通过在局部模型中加入批量归一化层（BN）解决联邦学习数据异构性中 feature shift 这种情况

•Feature shift：

•y 为标签，x 为特征，文章将 feature shift 定义为以下情况：

•例如：医学成像中不同的扫描仪/传感器，自动驾驶中不同的环境分布(公路和城市)，使得本地客户端的样本分布不同于其他客户端。

•与 FedAvg 类似，FedBN 也进行局部更新和模型聚合，不同的是，FedBN 假设局部模型有批量归一化层（BN），且 BN 的参数不参与聚合。

图片

收敛性分析：

作者用 NTK 做了一系列形式化证明，得出的结论是：在 feature shift 场景下，FedBN 的收敛比 FedAvg 更快更光滑

Experiment：

实验选了 5 个来自不同域且带有 feature shift 性质的数据集（不同域的数据具有异构性和相同的标签分布），具体包括 SVHN、USPS、 MNIST-M、MNIST、SynthDigits。文章预先对这五个数据集进行预处理，控制无关因素，例如客户之间样本数量不平衡问题，使 BN 的作用在实验中更加易于观察。

图片

各个数据集之间分布的差异图

图片

模型提供了一个简单的 CNN 网络，每一层后边都加了一个 BN ReLU

实验一

图片

从上面 5 个数据集和模型 setting 可以看出，不管从收敛质量还是从收敛速度来说，都是要比 FedBN 以及 FedAVG 要好。

图片

上图是 3 个实验：

a 图表明了不同的 local update step 对收敛性的影响。横轴是 Epochs 其实是 local update step 的概念。纵轴是 test accuracy。

1、可以看出 local update step 和 testing\_acc 负相关

2、fedBn 全程比 fed-avg 表现优异

3、尤其是我们看下最终收敛后的结果，fedbn 表现的更加稳定

图 b 是表示不同的本地数据集大小对于收敛的影响，横轴 percentage，即本地拿出百分之多少的数据进行训练，纵轴是 testing accuracy。我们可以看出，测试的准确度在 local client 只贡献 20%数据时候明显下降，但是对应于 singleset 的提升是反向的，数据量越少对于性能的提升越大，所以表明 fedbn 适合在单方非持有大量数据的联合训练中使用。

图 c 是表明不同异构性对于收敛的影响，作者这里试图回答一个问题异构性在什么程度时 FedBN 会比 FedAVG 位置好。这里的实验把每 1 方的数据都切分成了 10 份，一共有 5 个数据，总共 50 份数据，所以横坐标总数为 50。

1、开始时每个 dataset 选一个数据集来参与训练

2、然后逐渐的增加 clients，要保持 n 倍 clients 数量

3、More clients -> less heterogeneity

假设数据集更多，客户端更多，他的异构性就会更低。得出结论是，在任何设置下，FedBN 都比 FedAVG 效果要更好一点。

接下来是和 SOTA 的比较，对比的方法为 Fed Prox，Fed Avg，还有 Single set（把数据都拉到一方去进行训练，仍有 non-iid 特性）。因为 FedBN 的模型里边加了 BN，所以结果相较 Singleset 要更好一点，因为 BN 把一些“data shift”做了处理，降低了 non-iid 数据的影响。FedBN 和其他的相比的话，指标都要好一些，尤其是在 svhn 这个数据上，提升的幅度是最明显，将近有 10%。这个数据的特点就是特征很高，差异很大。所以这个方法对特征差异大的数据是比较有效的。同时 FedBN 在多次跑的时候，error 的方差最小，训练出的模型稳定性更好。

最后我们 FedBN 看一下在真实数据集下的表现：这里文章选择了自然图像数据集 Office-Caltech10，它有 4 个数据源 Office-31（3 个数据源）、Caltech256，每个 clients 分配了 4 个数据源中的一个，所以数据在 clients 之间是 Non-iid 的。

第二个数据源是 DomainNet,包含了从 6 个数据来源的自然图像数据集分别是：(Clipart,Infograph,Painting,Quickdraw,Real,Sketch)，划分方法和上面的 Cal10 是一样的。

第三个数据集从 ABIDE I 数据集选取，包含了 4 个医学数据集（NYU，USM，UM，UCLA）每一方被看做一个 client。

图片

结果和分析：

从结果可以看出，FedBN 在各方面全面超越了当前的 SOTA 工作。而且，我们可以看到 FedBN 除了 QuickDraw 数据集外，都获得了相对优势的结果，甚至优于 SingleSet 的结果，BN 在解决 feature shift 提升模型效果上有其显著的优势。

上述的结果也给我们在医学健康领域使用 FedBN 提供了信心，因为医疗场景下的数据的特征往往是有限数量，离散分布，而且是有 feature shift 的。

总结：

这篇工作提出了 FedBN，给出了在联邦场景下使用 BN 来缓解 feature shift 的方法，并证明了其收敛性和有效性。FedBN 的方法和通信、优化器方法是正交的，在实际使用中可以组合使用。FedBN 在算法是对 FedAvg 只做了很小的改动，在实际中还可以与其他联邦 strategy 进行结合使用（Pysyft、TFF 等已集成）。FedBN 已经在医学健康，自动驾驶等领域被证明是很有效的策略，这些场景的数据的特点是，不同方的本地数据由于存在 feature shift 导致的 non-iid。关于数据安全性，BN 层的数据对于整个交互和聚合是不可见的（Invasible），所以一定程度上增加了对本地数据攻击的困难。

更多干货内容，请关注公众号：隐语的小剧场