论文阅读笔记--Data-Free Knowledge Distillation for Heterogeneous Federated Learning（FEDGEN）

链接：https://arxiv.org/abs/2105.10056v2
这篇文章发表在ICML 2021，讲的是如何解决联邦学习中的数据异构问题。作者认为现有的知识蒸馏（Knowledge Distillation，K-D）的方法大多都需要server有一批代理数据（proxy data），协助K-D完成知识传授的过程，这有时候是做不到的。而且只修改全局模型并不能解决用户之间的异质性，反过来这还会影响全局聚合的质量。FEDGEN相比之前的算法有3个优越性：
1）不需要server拥有额外的数据；
2）K-D得到的知识被用于限制client的local update而不是限制全局模型，使client在non-iid的数据下表现更好（FL中的个性化）；
3）不需要client贡献自己的整个模型，而只需要贡献自己的预测层（prediction layer），隐私得到更好的保护，通信也降低了；

Knowledge Distillation

知识蒸馏就是把一个模型的知识传授给另外一个模型的过程，也称为Teacher-Student模型，通常会用于模型压缩等任务。典型的知识蒸馏方法需要有一个代理数据集，一个教师模型${\theta }_T$和一个学生模型${\theta }_S$，同样一条数据样本经过${\theta }_T$和${\theta }_S$，在prediction layer会分别得到一个输出向量（logits output），我们通过迭代代理数据集训练学生模型，最小化${\theta }_T$和${\theta }_S$logits output的分布之间的距离，这个距离通常使用KL散度（Kullback-Leibler divergence）来度量，最终使得学生模型能够在同样的数据样本下，得到和教师模型相似的输出，达到知识传授的目的。K-D在实际的使用中可能还涉及到一些技巧，比如温度参数什么的，这里就不详细讨论了，总之它可以形式化地描述为：

其中，神经网络模型被划分为两个部分，${\theta }^f$和${\theta }^p$，分别表示除prediction layer之外的前面的层和prediction layer，函数f表示前面的层的输出，g表示预测层的输出，$\sigma$表示预测层的激活函数（有些工作没有把激活函数加入知识蒸馏的过程，应该也是可以的）。
以往的server持有一个代理数据集的知识蒸馏方法可以描述为：

Knowledge Extraction

作者是如何解决K-D需要代理数据集的问题，来实现data-free的呢。首先想到的是使用生成的数据来做蒸馏，但是生成的数据需要符合全局的数据分布，但是由于数据隐私问题，server是无法获取每个client的数据分布构建全局分布的。为解决这个问题，该联邦学习系统要求client能够提供给server自己样本标签的分布，也就是client共享y而不共享x。server根据全局的标签的分布，学习一个映射$Q^{\ast }:\mathcal{Y}\to \mathcal{X}$，这样生成的样本是符合全局的数据分布的

$Q^{\ast }$这里是一个MLE，y的分布p(y)可以通过client共享的标签分布来估计，
p(y|x)可以由client的模型集成来估计，

但是这里还是有问题，如果直接优化$Q^{\ast }$，$\mathcal{X}$的维度过高会带来很大的计算负载，并且生成$\mathcal{X}$也可能造成数据泄露。于是作者想到可以在一个latent space里学习一个$G^{\ast }:\mathcal{Y}\to \mathcal{Z}$来代替$Q^{\ast }$做K-D，也就是

如果我们用$w$来描述$G^{\ast }$的参数的话，优化目标就变成了：

为实现多样化的输出，作者还给生成器的输入加入了高斯噪声。
server生成数据是带有全局数据分布的信息的，server将这些数据分发给client，这些数据对client而言相当于进行了数据增强，client用增强过后的样本进行训练，损失函数就变成了两个部分：
式子的第二项就是local model在生成样本上的损失，

FEDGEN

FEDGEN的算法框架如上图所示。server选择client，广播生成模型$w$，全局模型$\theta$(实际上只是模型的prediction layer），和全局的标签分布给client，client用收到的全局模型替代自己的模型，构建增强的数据样本（第7行），统计本地的标签分布（第8行）。然后本地依据Equation 5进行更新，将更新的局部模型${\theta }_k$发送到server。server收集局部模型进行聚合，收集局部的标签信息构建全局的数据分布，然后依据Equation 4优化数据生成器G。为什么只需要共享prediction layer就可以了呢，可以看Equation 4，y输入到生成器$w$ ，得到输出z，将z再次输入到联邦聚合得到的prediction layer种中，我们希望能再次输出y，也就是在训练$w$学习prediction layer的反函数。在Equation 5的第二项中，我们也只是将生成的z直接输入到prediction layer。下图展示了FEDGEN的overview。
右半部分，client将依据本地训练得到的${\theta }^p$上传到server，server利用这些prediction layer加入高斯噪声训练生成器$G_w$，$G_w$可以依据y在latent space生成样本。在左半部分的第k个client的局部训练当中，增强的样本通过直接加入prediction layer，参与局部训练。

（转载请注明出处）