FedDG: Federated Domain Generalization on Medical Image Segmentation

论文：FedDG: Federated Domain Generalization on Medical Image Segmentation via Episodic Learning in Continuous Frequency Space
收录情况：CVPR2021
pdf链接：https://arxiv.org/abs/2103.06030
代码：https://github.com/liuquande/FedDG-ELCFS
主要方法：使用多个域作为源域训练，但值得注意的是多个源域数据不是直接拿到，而是只拿到不同源域数据的幅度谱，另外增加了一个优化边界的损失函数，然模型泛化的未见过的目标域。
关键词：multi-source domain generalization

1. Abstract

联邦学习可以让分散的医疗中心获得共享的模型来分别训练自己的数据。但是在临床部署中，联邦学习训练的模型在完全没见过的数据集上性能仍然会下降。在这篇文章中，我们提出并解决了一个新的问题设定联邦预适应(FedDG)，其目标是通过来自多个源域的数据学习一个联邦的模型从而能泛化到未见过的目标域数据。我们提出了一个新颖的方法叫做连续频域空间中的情景学习(ELCFS)，使得每个分散的客户端能够利用多源域的数据分布。此方法通过连续频域空间插值机制传输数据信息，能有效保护客户端之间的数据隐私。另外，我们设计了一个针对优化边界的情景学习范式，来满足医学图像分割中模型泛化的问题。

2. Introduction

主要贡献点：
（1）提出了联邦预适应问题（Federated Domain Generalization, FedDG）；
（2）使用了一种保护隐私的数据共享解决方法，即跨客户端的连续频率空间插值机制；
（3）针对客户端提出了一种面向边界的情景学习策略，能够改善因为域漂移导致的边界分割不清晰问题；
（4）在前列腺和眼底数据集上进行充分实验验证了方法的有效性。

3. Method

Related Work主要介绍Federated Learning和Domain Generalization，这里略过。

3.1 Federated Domain Generalization

• 问题设定：在FedDG中，$(\mathcal{X},\mathcal{Y})$表示图像空间和标签空间的联合分布，$\mathcal{S}=\left\{{\mathcal{S}}^1,{\mathcal{S}}^2,\dots ,{\mathcal{S}}^K\right\}$表示$K$个源域数据。${\mathcal{S}}^k={\left\{\left(x_i^k,y_i^k\right)\right\}}_{i=1}^{N^k}$表示其中一个源域$({\mathcal{X}}^k,\mathcal{Y})$的数据分布。FedDG的任务是利用$K$个源域数据训练模型$f_{\theta }:\mathcal{X}\to \mathcal{Y}$，使其在完全未见过的域$\mathcal{T}$上有较好的表现。
• 挑战：（1）多源域数据是分布式存储的，每个客户端只能访问自己的数据集，这就限制了模型学习不同的数据分布；（2）能见到的多源域数据可能存在很大的异质性（就是data distribution差异较大），因此泛化到目标域时不能保证模型有良好的通用性；（3）医学图像通常分割边界较模糊，这对DG方法有较大挑战。

针对上面提出的挑战，文章后续提出了相应的解决方法。

3.2 Continuous Frequency Space Interpolation

• 对于某一张图像$x_i^k\in {\mathbb{R}}^{H\times W\times C}$，对其进行FFT变换到频域：$$\mathcal{F}\left(x_i^k\right)(u,v,c)=\sum _{h=0}^{H-1}\sum _{w=0}^{W-1}{x}_i^k(h,w,c)e^{-j2\pi \left(\frac{h}{H}u+\frac{w}{W}v\right)}$$上面的频谱可以分解成幅度谱$A_i^k\in {\mathbb{R}}^{H\times W\times C}$和相位谱$P_i^k\in {\mathbb{R}}^{H\times W\times C}$，其中幅度谱反映了low-level的特征信息（比如图像风格），相位谱反映了high-level的图像语义信息。
  为了构建不同客户端数据分布字典$\mathcal{A}=\left[{\mathcal{A}}^1,\dots ,{\mathcal{A}}^K\right]$，其中${\mathcal{A}}^k={\left\{{\mathcal{A}}_i^k\right\}}_{i=1}^{N^k}$表示第$k$个客户端的图像幅度谱。这个字典在每一个客户端时可以共享的。
• 然后通过如下公式获得两个域$k\to n$插值的幅度谱：$${\mathcal{A}}_{i,\lambda }^{k\to n}=(1-\lambda ){\mathcal{A}}_i^k\ast (1-\mathcal{M})+\lambda {\mathcal{A}}_j^n\ast \mathcal{M}$$其中，${\mathcal{A}}_j^n$是从distribution bank中随机采样的幅度谱，${\mathcal{A}}_i^k$是训练过程中输入的源域中图像的幅度谱，$\mathcal{M}=1_{(h,w)\in [-\alpha H:\alpha H,-\alpha W:\alpha W]}$是对幅度谱低频部分取mask，$\lambda$是$[0.0,1.0]$随机采样。
  对插值后的幅度谱进行反傅立叶变换：$$x_{i,\lambda }^{k\to n}={\mathcal{F}}^{-1}\left({\mathcal{A}}_{i,\lambda }^{k\to n},{\mathcal{P}}_i^k\right)$$另外，$t_i^k={\left\{x_{i,\lambda }^{k\to n}\right\}}_{n\ne k}$
• 注意：distribution bank是对$K$个源域的图像进行傅立叶变化取幅度谱得到的，这一步计算容易，可以快速获得；另外对于每一个客户端训练模型时，随机采样共享的distribution bank中的幅度谱进行插值操作。
  

3.3 Boundary-oriented Episodic Learning

Episodic learning at local client

在不同客户端训练自己的数据看作是元学习，原始数据$x_i^k$是训练集，变换后的数据$t_i^k$为测试集。
首先，客户端模型根据dice loss更新参数：$${\hat{\theta }}^k={\theta }^k-\beta {\nabla }_{{\theta }^k}{\mathcal{L}}_{seg}\left(x_i^k;{\theta }^k\right)$$其次，挑选最优模型时用的是$t_i^k$作为验证集。

Boundary-oriented meta optimization

提出了一种优化边界的损失函数

• 原始训练图提取目标边缘和背景边缘一维特征向量： 根据$x_i^k$的label $y_i^k$，使用形态学算子可以提取到boundary mask $y_{i\_bd}^k$ 和background mask $y_{i\_bg}^k$，上图右边所示。
  分割网络最后一层的特征图$Z_i^k$与$y_{i\_bd}^k$，$y_{i\_bg}^k$通过如下操作得到每张训练图片的一维特征向量，这个特征向量代表了边界的区域特征$$h_{i_{-}bd}^k=\frac{{\sum }_{h,w}{Z}_i^k\ast y_{i-bd}^k}{{\sum }_{h,w}{y}_{i_{-bd}}^k};h_{i_{-}bg}^k=\frac{{\sum }_{h,w}{Z}_i^k\ast y_{i_{-}bg}^k}{{\sum }_{h,w}{y}_{i_{-}bg}^k}$$
• 变换图提取目标边缘和背景边缘一维特征向量： 对于每一张训练图$x_i^k$，都对应着$K-1$张变换图$t_i^k$，每一张变换图也可以根据label $y_i^k$提取$h_{i_{-}bd}^k$和$h_{i_{-}bg}^k$。因此对于每一张训练图，最后都能得到$K$个boundary-related和$K$个background-related特征向量。
• 定义边界损失函数： 在$2K$个向量中，定义特征对$(h_m,h_p)$，正类对为都属于boundary-related或background-related的向量，负类对为一个是boundary-related另一个是background-related向量。对于每一个正类对，使用InfoNCE loss：$$\ell \left(h_m,h_p\right)=-\log \frac{\exp \left(h_m\odot h_p/\tau \right)}{{\sum }_{q=1,q\ne m}^{2K}\mathbb{F}\left(h_m,h_q\right)\cdot \exp \left(h_m\odot h_q/\tau \right)}$$ $\odot$为余弦相似度，正类对$\mathbb{F}\left(h_m,h_q\right)$为0，负类对为1，$\tau$为温度系数。最后的${\mathcal{L}}_{\text{boundary}}$为$\ell$ over所有正类对的平均：$${\mathcal{L}}_{\text{boundary }}=\sum _{m=1}^{2K}\sum _{p=m+1}^{2K}\frac{\left(1-\mathbb{F}\left(h_m,h_p\right)\right)\cdot \ell \left(h_m,h_p\right)}{B(K,2)\times 2}$$其中$B(K,2)$为$K$取2
• 总的损失函数： $${\mathcal{L}}_{\text{meta }}={\mathcal{L}}_{\text{seg }}\left(t_i^k;{\hat{\theta }}^k\right)+\gamma {\mathcal{L}}_{\text{boundary }}\left(x_i^k,t_i^k;{\hat{\theta }}^k\right)$$ $$\underset{{\theta }^k}{\mathrm{argmin}}{\mathcal{L}}_{seg}\left(x_i^k;{\theta }^k\right)+{\mathcal{L}}_{meta}\left(x_i^k,t_i^k;{\hat{\theta }}^k\right)$$

Experiments

采用leave-one-domain-out策略，即一个数据域作为unseen domain，其余数据域为source domain