FedDG: Federated Domain Generalization on Medical Image Segmentation via Episodic Learning

FedDG: Federated Domain Generalization on Medical Image Segmentation via Episodic Learning in Continuous Frequency Space

论文标题

FedDG: Federated Domain Generalization on Medical Image Segmentation via Episodic Learning in Continuous Frequency Space

论文来源

CVPR 2021

Paper: https://arxiv.org/abs/2103.06030

Code: https://github.com/liuquande/FedDG-ELCFS

1 背景介绍

在医学影像分割任务中，越来越多的工作倾向于收集多个医疗机构的数据并由此构建出更加精确且鲁棒的深度网络。但是由于医疗影像的特殊性(存在大量不同的隐私保护规定)，导致这种依靠大规模数据集提升模型性能和鲁棒性的方法十分困难。最近，很多研究者注意到可以使用联邦学习(Federated Learning, FL)解决隐私保护的问题。

联邦学习可以在多分布的数据集中进行模型训练，每个单独的数据集只需在本地进行处理。网络参数更新时只需要将各个本地数据训练的模型定期推送到中心服务器，中心服务器随后将各个本地模型进行聚合得到新的网络参数并更新到本地终端。在联邦学习的过程中只有网络参数的交互，避免出现数据交互，因此不违反医疗数据的隐私保护政策。

现有的联邦学习方法往往只关注已有的数据集中模型的表现，忽视了模型在未知的新数据分布下的泛化性和鲁棒性。实际应用场景中，我们往往会遇到将现有的网络模型部署到新的医疗机构和地区的需求，不同机构使用的设备差异、地区间人群分布差异等都会导致医疗影像的数据分布差异(domain shift)。这种数据分布差异是阻碍联邦学习在实际中广泛应用的一个关键问题。

2 论文贡献

1. 这篇文章基于医疗影像分割任务的实际应用场景，提出了一个新的问题设定：联邦学习的领域泛化问题(Federated Domain Generalization, FedDG)。(如Figure 1 (a)所示)

2. 这篇文章提出一个基于连续频域空间差值(Continuous Frequency Space Interpolation, CFSI)的泛化联邦学习方法，可以在隐私保护的前提下提高联邦学习模型的泛化性。

3. 这篇文章为提升本地模型训练的领域泛化能力，提出一个新的基于边界方向的交替学习策略(Boundary-oriented Episodic Learning scheme, BEL)。

4. 这篇文章在两个医疗影像分割的DG基准测试上取得了FedDG设定下的最好成绩，并通过实验验证了方法的有效性。

3 方法

3.1 Federated Domain Generalization

首先介绍一下联邦学习的领域泛化问题设定。我们将数据和标注的联合分布定义为$S=(X,Y)$，有$K$个不同分布的训练数据集 S = { S 1 , S 2 , … , S K } S=\{ S^1, S^2, \dots, S^K\} S={S1,S2,…,SK}。其中每个数据集 S k = { ( x i k , y i k ) } i = 1 N k S^k=\{(x_i^k, y_i^k) \}_{i=1}^{N^k} Sk={(xik​,yik​)}i=1Nk​。算法目标是使用$K$个不同分布的数据集在对应的$K$个本地模型的训练过程中学习一个模型$f_{\theta }:X\to Y$，使之可以直接泛化到一个新分布的数据集 T = { X t , Y } T=\{X^t,Y\} T={Xt,Y}上。

FedDG设定的主要挑战有：

• 训练数据为多分布的数据集，并且每个本地模型只能接触到一个单一分布。如何保证模型跨分布的泛化性。
• 由于不同数据集之间的巨大差异，如何将这些本地模型进行整合从而得到domain-invariant的统一模型。
• 医疗影像分割任务在分割边界存在较大的不确定性，如何在FedDG设定下提高边界区域特征的领域不变性(domain invariance)。

关于将本地模型整合为统一模型的方式，这篇文章直接采用联邦学习中广泛使用的参数平均的方式，即将各个本地模型的网络参数根据数据集大小加权后取平均。

3.2 Continuous Frequency Space Interpolation

为了解决本地模型训练时无法使用其他数据集的问题，作者提出使用连续频域空间插值的方法提取不同数据集的分布信息的方法。具体而言，作者指出在视觉心理学领域，通常认为频域图像$F(x_i^k)$的相位信息$p_i^k$代表了high-level的语义信息，而幅值信息$A_i^k$代表了low-level的分布信息(style/texture)。因此作者采集了每个图像的幅值信息，构成分布集合(distribution bank)$A=[A^1,A^2,\dots ,A^K]$，其中 A k = { A i k } i = 1 N k A^k=\{ A^k_i\}_{i=1}^{N^k} Ak={Aik​}i=1Nk​。

接下来作者在训练过程中对每副图像进行幅值上的插值：随机从分布集合中选取一个不同分布的幅值信息，再对其中的低频部分(超参数$\alpha$控制)进行插值，插值比例$\lambda$为$[0.0,1.0]$的随机采样。

将插值后的幅值信息和原始的相位信息结合，再进行傅里叶反变换得到新的图像。

整体过程如Figure 2的左边部分所示。

3.3 Boundary-oriented Episodic Learning

针对本地模型训练过程中增强边界区域特征的领域不变性问题，这篇文章提出了一个基于元学习的学习策略。由于上一步中使用不同分布的幅值信息插值可以得到不同分布的图像，因此可以在本地模型的训练过程中模拟domain shift的情况。

Episodic learning at local client: 使用Dice loss $L_{seg}$对网络参数${\theta }^k$进行更新，学习率为$\beta$。

Boundary-oriented meta optimization: 为了保持边界预测在不同分布间的不变性，作者首先使用形态学滤波对分割标注进行处理，分别提取出边界部分$y_{i\_bd}^k$和背景部分$y_{i\_bg}^k$的mask，然后将$l$层的网络特征$Z_i^k$按照mask提取出边界特征和背景特征：

对边界特征和背景特征使用InfoNCE[1] loss约束，针对相同标注的不同分布数据进行特征相似度最大化约束，而针对不同标注的数据进行特征相似度最小化约束，其中特征相似度使用余弦相似度。

最后得到边界特征的领域不变性约束。

Overall local learning objective:

meta test过程的损失函数$L_{meta}$为

整体优化过程

4 实验

4.1 Datasets

这篇文章分别在视网膜眼底影像的视盘和视杯分割任务与MRI影像的卵巢分割任务上进行实验。其中视网膜数据由4个不同的公开数据集构成，MRI数据来自6个不同的公开数据源。具体的数据预处理方法和超参数设置参考原文[2]。

4.2 Comparison with DG methods

视网膜视盘与视杯分割结果如Table 1 所示，评价指标为Dice值。

卵巢MRI分割结果如Table 2所示。

Figure 4 给出了不同DG方法的泛化性对比的可视化结果。

4.3 Ablation Study

作者还分别验证了频域空间插值CFSI和交替学习策略BEL单独对分割任务的影响，如Figure 5所示。其中频域空间插值的影响相对更大一些，但是对频域插值图像进行InfoNCE loss[1]的约束也很重要，可以看到在同时使用CFSI和BEL相比单独的方法有很大的提升。

Figure 6 (a)展示了频域空间插值的降维可视化结果，Figure 6 (b)展示了不同插值系数对分割结果的影响，可以分割效果随着插值系数的增大而有一定程度的提高。

Figure 7 (a)展示了使用InfoNCE loss[1]对特征进行领域不变的相似性约束在训练数据以及测试数据上的影响。可以看到使用此约束可以显著提高边界特征与背景特征的差异，提高预测边界的一致性。Figure 7 (b)展示了在元学习过程中使用对边界特征的约束对网络泛化能力的影响。

5 总结与讨论

这篇文章提出了一个新的问题设定FedDG，考虑了医疗图像处理领域的隐私保护问题，同时更符合实际应用场景，未来可以在此方向进行进一步的研究。并且提出了有效的多分布数据增强方法和特征领域不变性的约束方法，在多个DG的实验中取得了现阶段最好的结果。

但是，这篇文章提出的频域插值方法虽然不直接使用原始图像数据，但是仍然收集了全部数据的频域信息，和联邦学习的初衷有些许差异。并且，本地模型的训练过程中使用元学习框架，训练开销大，在很多场景下不利于方法的部署。

6 参考文献

[1]. Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. arXiv preprint arXiv:2002.05709, 2020. 5

[2] Quande Liu, Cheng Chen, Jing Qin, Qi Dou and Pheng-Ann Heng. FedDG: Federated Domain Generalization on Medical Image Segmentation via Episodic Learning in Continuous Frequency Space. CVPR 2021