论文笔记：Confidence-Aware Personalized Federated Learning via Variational Expectation Maximization

论文 Confidence-Aware Personalized Federated Learning via Variational Expectation Maximization（2023 CVPR）的阅读笔记，可在 arXiv 在线阅读

1. Abstract

本文针对 FL 中各个客户的数据集非独立同分布以及数据集大小不同的场景。pFedVEM 是一个基于层次贝叶斯模型和变分推断的 PFL 框架。其中每个客户有一个本地模型，所有客户共享一个全局模型，本地模型和全局模型采用同样的模型架构。全局模型被建模为一个潜变量，作用是增强客户本地模型参数的联合分布，以及捕获不同客户的共同趋势。

算法给出了一个置信度值的估计，用于提供每个客户在聚合阶段的权重参数。

2. Symbol Table

符号	含义
${\mathcal{D}}_j$	第 $j$ 个客户的数据集
$n_j$	第 $j$ 个客户数据中实例数量
${\mathbf{x}}_i^{\left(j\right)}$	第 $j$ 个客户数据集中的第 $i$ 条实例
$y_i^{\left(j\right)}$	第 $j$ 个客户数据集中的第 $i$ 条实例的标签（单标签）
${\mathbf{w}}_j$	第 $j$ 个客户的模型权重
$J$	客户数量

3. Problem Formulation

在本文的设置中，一共有 $J$ 个客户，$J$ 个本地模型和 $1$ 个全局模型，所有的模型采用相同的模型架构（区别只是模型参数不同）。全局模型不用于预测，它只作为辅助的潜变量优化局部模型。每个局部模型各自在自己的数据集上进行预测。

在 PFL 中，要解决的是如下优化问题，也就是要获得使得总的平均损失最小的本地模型集合
$$\underset{{\mathbf{w}}_{1:J}\in {\mathbb{R}}^d}{\min }f\left({\mathbf{w}}_{1:J};\mathcal{D}\right):=\frac{1}{J}\sum _{j=1}^J{f}_j\left({\mathbf{w}}_j;{\mathcal{D}}_j\right)$$
$f_j\left({\mathbf{w}}_j;{\mathcal{D}}_j\right)$ 是第 $j$ 个客户的损失
$$f_j\left({\mathbf{w}}_j;{\mathcal{D}}_j\right)=\frac{1}{n_j}\sum _{i=1}^{n_j}l\left({\mathbf{x}}_i^{\left(j\right)},y_i^{\left(j\right)};\mathbf{w}\right)$$
$l\left(\cdot ,\cdot ;\mathbf{w}\right)$ 是当 ${\mathbf{w}}_j=\mathbf{w}$ 时，模型在数据点 $\left(\cdot ,\cdot \right)$ 上的损失。

4. Confidence-aware PFL

本文为 PFL 提出了一个通用的贝叶斯框架，然后基于变分期望最大化来进行模型优化。达到的效果是每个客户可以为自己的本地训练结果估计一个置信度分数，该分数被用在模型聚合过程中自适应的调教权重（全局模型由本地模型加权聚合得到，做的事情就是如何自适应的得到这个权重）

4.1 Hierarchical Bayesian modeling

从贝叶斯的角度来说，PFL 要做的事情就是为任意客户 $j$ 获取后验概率密度函数 $p\left({\mathbf{w}}_j\mid {\mathcal{D}}_j\right)$。由于在本文的设置中，所有的客户都运行相似的任务，因此对一个客户而言，其余客户的数据集也应该有助于构建该客户的后验。但是，每个客户不能访问其余客户的隐私数据集，因此，通过构建一个潜变量全局模型来捕获不同客户之间的相关性，每个客户可借助该全局模型从其他客户那里获得帮助，这就是本文模型的核心思想。

用 $\mathbf{w}$ 代表全局模型的模型参数，本文假设在不同的客户之间有条件独立性如下
$$p\left({\mathbf{w}}_i\mid \mathbf{w}\right)p\left({\mathbf{w}}_j\mid \mathbf{w}\right)=p\left({\mathbf{w}}_i,{\mathbf{w}}_j\mid \mathbf{w}\right)$$
$p\left({\mathbf{w}}_j\mid \mathbf{w}\right)$ 的含义就是在基于所有客户数据集的共有趋势的前提下，对第 $j$ 个客户的模型进行定制。由于添加了全局模型 $\mathbf{w}$，因此，要求解的最佳模型参数集扩展为 { w , w 1 : J } \left\{\boldsymbol{w}, \boldsymbol{w}_{1:J}\right\} {w,w1:J​}，因此，要做的就是获得后验
$$p\left(\mathbf{w},{\mathbf{w}}_{1:J}\mid \mathcal{D}\right)$$
根据贝叶斯规则有
$$p\left(\mathbf{w},{\mathbf{w}}_{1:J}\mid \mathcal{D}\right)=\frac{p\left(\mathcal{D}\mid \mathbf{w},{\mathbf{w}}_{1:J}\right)p\left(\mathbf{w},{\mathbf{w}}_{1:J}\right)}{p\left(\mathcal{D}\right)}$$

其正比于等式右方的分子部分，又有
$$\begin{array}{rl}p\left(\mathcal{D}\mid {\mathbf{w}}_j\right)& =\exp \left(-n_j{f}_j\left({\mathbf{w}}_j;{\mathcal{D}}_j\right)\right)\\ p\left(\mathbf{w},{\mathbf{w}}_{1:J}\right)& =p\left(\mathbf{w}\right)\prod _{j=1}^{J}p\left({\mathbf{w}}_j\mid \mathbf{w}\right)\end{array}$$
其中第一个等式为似然函数定义的展开（有一点小变换，相当于先对似然函数取对数将连乘变成连加，然后又进行指数操作，还原为似然函数），第二个等式因为前面的条件独立性假设而成立

因此，模型转变为
$$p\left(\mathbf{w}\right)\prod _{j=1}^{J}p\left({\mathbf{w}}_j\mid \mathbf{w}\right)\exp \left(-n_j{f}_j\left({\mathbf{w}}_j;{\mathcal{D}}_j\right)\right)$$
这表明边际联合分布 $p\left({\mathbf{w}}_{1:J}\mid \mathcal{D}\right)=\int p\left(\mathbf{w},{\mathbf{w}}_{1:J}\mid \mathcal{D}\right)d\mathbf{w}$ 在引入了全局模型作为潜变量后变得更加灵活，能够允许客户之间进行复杂的通信。

4.2 Variational expectation maximization

本文选择先验（推理问题不关注先验，选一个即可）$p\left({\mathbf{w}}_j\mid \mathbf{w}\right)$ 为一个各项同性高斯条件分布，即
$$p\left({\mathbf{w}}_j\mid \mathbf{w}\right)=\mathcal{N}\left({\mathbf{w}}_j\mid \mathbf{w},{\rho }_j^2\mathbf{I}\right)$$
其中 ${\rho }_j^2$ 代表分布的方差

4.3 Maximizing the marginal likelihood

优化贝叶斯模型的一种方法是使用 MAP，但是选择的先验给出了一组超参 ${\rho }_{1:J}$，这很难选择，因此，本文采用的方法是用分布 $q\left({\mathbf{w}}_{1:J}\right)$ 来近似后验分布 $p\left({\mathbf{w}}_{1:J}\mid \mathcal{D}\right)$，假设各客户独立，分布 $q\left({\mathbf{w}}_{1:J}\right)$ 被定义为一组分布的乘积
$$q\left({\mathbf{w}}_{1:J}\right):=\prod _{j=1}^J{q}_j\left({\mathbf{w}}_j\right)$$
分布 $q_j\left({\mathbf{w}}_j\right)$ 被选择为多元高斯分布，即
$$q_j\left({\mathbf{w}}_j\right)=\mathcal{N}\left({\mathbf{w}}_{\mathbf{j}}\mid {\mathbf{\mu }}_j,{\mathbf{\Sigma }}_j\right)$$
KL 散度用于衡量两个分布的相似程度，因此，需要的最佳 $q\left({\mathbf{w}}_{1:J}\right)$ 由下式得到
$$q^{\ast }\left({\mathbf{w}}_{1:J}\right)=\underset{q\left({\mathbf{w}}_{1:J}\right)}{\mathrm{argmin}}\mathrm{KL}\left(q\left({\mathbf{w}}_{1:J}\right)\Vert p\left({\mathbf{w}}_{1:J}\mid \mathcal{D}\right)\right)$$
又有 KL 散度展开有
$$\mathrm{KL}\left(q\left({\mathbf{w}}_{1:J}\right)\Vert p\left({\mathbf{w}}_{1:J}\mid \mathcal{D}\right)\right)=\log p\left(\mathcal{D};{\mathbf{w}}_{1:J}\right)-\mathrm{ELBO}\left(q\left({\mathbf{w}}_{1:J}\right);{\rho }_{1:J}^2,\mathbf{w}\right)$$
又由于 KL 散度的非负性，因此有
$$\log p\left(\mathcal{D};{\mathbf{w}}_{1:J}\right)>\mathrm{ELBO}\left(q\left({\mathbf{w}}_{1:J}\right);{\rho }_{1:J}^2,\mathbf{w}\right)$$
因此，最优化问题转为极大化证据下界 ELBO，将其展开如下
$$\begin{array}{rl}\mathrm{ELBO}\left(q\left({\mathbf{w}}_{1:J}\right);{\rho }_{1:J}^2,\mathbf{w}\right)& =\sum _{j=1}^J\left({\mathbb{E}}_{q\left({\mathbf{w}}_j\right)}\left[\log p\left({\mathcal{D}}_j\mid {\mathbf{w}}_j\right)\right]-\mathrm{KL}\left[q\left({\mathbf{w}}_j\Vert p\left({\mathbf{w}}_j\mid \mathbf{w},{\rho }_j^2\right)\right)\right]\right)\\ & =\sum _{j=1}^J\left({\mathbb{E}}_{q\left({\mathbf{w}}_j\right)}\left[\log p\left({\mathcal{D}}_j\mid {\mathbf{w}}_j\right)\right]-{\mathbb{E}}_{q\left({\mathbf{w}}_j\right)}\left[\log q\left({\mathbf{w}}_j\right)\right]+{\mathbb{E}}_{q\left({\mathbf{w}}_j\right)}\left[\log p\left({\mathbf{w}}_j,\mathbf{w},{\rho }_j^2\right)\right]\right)-\sum _{j=1}^J\log p\left(\mathbf{w},{\rho }_j^2\right)\\ & =\sum _{j=1}^J\left({\mathbb{E}}_{q\left({\mathbf{w}}_j\right)}\left[\log p\left({\mathcal{D}}_j,{\mathbf{w}}_j\right)-\log p\left({\mathbf{w}}_j\right)-\log q\left({\mathbf{w}}_j\right)+\log p\left({\mathbf{w}}_j,\mathbf{w},{\rho }_j^2\right)-\log p\left(\mathbf{w},{\rho }_j^2\right)\right]\right)\\ & =\sum _{j=1}^J\left({\mathbb{E}}_{q\left({\mathbf{w}}_j\right)}\left[\log p\left({\mathcal{D}}_j,{\mathbf{w}}_j\right)-\log q\left({\mathbf{w}}_j\right)\right]\right)\\ & \propto \sum _{j=1}^J\left({\mathbb{E}}_{q\left({\mathbf{w}}_j\right)}\left[\log p\left({\mathcal{D}}_j\mid {\mathbf{w}}_j,\mathbf{w}\right)+\log p\left({\mathbf{w}}_j\mid \mathbf{w},{\rho }_j^2\right)\right]\right)\\ & \propto \sum _{j=1}^J{\mathbb{E}}_{q\left({\mathbf{w}}_j\right)}\left[\log p\left({\mathbf{w}}_j\mid \mathbf{w},{\rho }_j^2\right)\right]\end{array}$$
可以得到闭式解，同时定义每个客户的置信度分数为 ${\tau }_j:=1/{\rho }_j^2$，有
$${\tau }_j=d/\left(\mathrm{Tr}\left({\mathbf{\Sigma }}_{\mathbf{j}}\right)+{\Vert {\mathbf{\mu }}_j-\mathbf{w}\Vert }^2\right)$$
其中 $d$ 是 $\mathbf{w}$ 的维度，同时定义服务器聚合权重的公式如下
$${\mathbf{w}}^{\ast }=\frac{\sum _{j=1}^J{\tau }_j{\mathbf{\mu }}_j}{\sum _{j=1}^J{\tau }_j}$$
直观来说，客户的不确定性越小，模型偏差越小，置信度分数就越高，置信度分数越高的客户，在模型聚合时的权重越大。

5. Algorithm Implementation

为了能够用梯度下降进行训练，需要对分布进行重参数化并采样，具体来说，用 $\left({\mathbf{\mu }}_{1:J},{\mathbf{\pi }}_{1:J}\right)$ 参数化高斯分布，并以如下形式从中采样
$$q\left({\mathbf{w}}_j\right)={\mathbf{\mu }}_j+\mathrm{diag}\left(\log \left(1+\exp \left({\mathbf{\pi }}_j\right)\right)\right)\cdot \mathcal{N}\left(0,{\mathbf{I}}_d\right)$$
此外，使用蒙特卡洛算法近似 ELBO 中第一项，即 ${\mathbb{E}}_{q\left({\mathbf{w}}_j\right)}\left[\log p\left({\mathcal{D}}_j\mid {\mathbf{w}}_j\right)\right]$ 的值，因此，当进行 $K$ 次采样时，目标函数变为
$$\frac{n_j}{K}\sum _{k=1}^K{f}_j\left({\mathcal{D}}_j;{\mathbf{w}}_{j,k}\right)-\mathrm{KL}\left[q\left({\mathbf{w}}_j\Vert p\left({\mathbf{w}}_j\mid \mathbf{w},{\rho }_j^2\right)\right)\right]$$
算法伪代码如下，一共进行 $T$ 轮迭代，每轮迭代中，所有客户先用全局模型参数本地更新自己的模型参数，其中 $\mathbf{\mu }$ 代表 head model 参数（用于学习特征表达），$\mathbf{\theta }$ 代表 base model 参数（用于执行分类）。然后服务器随机选择一批客户，被选中的客户将自己的置信度分数与模型参数一并上传服务器。服务器根据权重计算公式加权聚合得到新一轮的全局模型参数并将其返还给客户。