论文解读《BiX-NAS: Searching Efficient Bi-directional Architecture for Medical Image Segmentation》
论文解读《BiX-NAS:搜索医学图像分割的高效双向架构》
文章发表出处：MICCAI2021(国际医学图像计算和计算机辅助干预协会(Medical Image Computing and Computer Assisted Intervention Society) 举办，跨医学影像计算(MIC)和计算机辅助介入 (CAI) 两个领域的综合性学术会议，是该领域的顶级会议)
项目网页（含代码以及论文）：https://bionets.github.io/
论文：论文地址
代码：代码地址
摘要：最近，可循环的U-Net网络变体已被引入到各种医学图像分割任务中。现有的研究侧重于通过重用构建块来实现网络循环，在不显著增加参数的同时提升网络的性能。但这种可循环的模块仍然不可避免地增加了计算成本。
本文研究了双向跳跃连接 (bi-directional skip connection) 网络的多尺度（multi-scale) 变形，然后通过一种新颖的两阶段式神经结构搜索 (NAS) 算法BiX-NAS，自动搜索出一种超轻量且有效的可循环结构。本文提出的方法通过过滤掉无效的多尺度跳跃连接来降低网络的计算成本。在三个不同的医学图像数据集上的实验表明，BiX-NAS搜索出的结构（BiX-Net）以显着降低的计算成本实现了最先进的性能。
一、引言：
基于深度学习的方法在医学图像分析中占有优势。U-Net是一种广泛使用的分段网络，它通过构建前向跳过连接(skip connections)来将编码器中的编码特征与解码特征聚合在一起。利用可重用构建块，对此类体系结构进行迭代推理，取得了最新进展。提出了一个递归U-Net，它在每次迭代时递归成对编码和解码块的子集。BiO-Net引入了反向跳过，将解码器中的语义传递到同一级别的编码器。虽然这种反复的设计可以大大缩小网络的规模，按照预先设定的运行时间，计算成本仍不可避免地增加。同时，多尺度方法的成功表明，可以使用融合细粒度特征和粗粒度语义的多尺度跳过。为此，引入多个语义尺度的前向/后向跳转，以较低的计算成本搜索多尺度特征的有效聚合是一个非常重要的研究方向。
神经结构搜索(NAS)方法自动感知结构，以获得有效和经济的性能增益。经典的NAS方法是通过在搜索和分别验证每个样本模型时随机抽取样本来进化的。可微NAS算法将离散搜索空间放松为连续，并委托反向传播来搜索最佳候选。Auto-DeepLab将可微的NAS算法应用到图像分割任务中，以确定每个构建块的最佳操作符和拓扑。同时，NAS-Unet应用了基于梯度的细胞结构自动搜索来构建类似u网的架构。尽管在相同级别的特征融合上取得了成功，引入了一个多尺度搜索空间，使他们提出的MS-NAS具有安排多级特征聚合的能力。然而，通过上述NAS算法搜索的体系结构仅比手工设计的体系结构带来了边际的改进，而且它们的搜索过程效率很低。
在本文中，我们提出了一种有效的多尺度(抽象为“X”)NAS方法，即BiX-NAS，它通过不断跳过多尺度特征，同时丢弃不重要的特征来搜索最优的双向结构(BiX-Net)。
本文的研究成果主要体现在以下三个方面:(1)以双向循环的方式研究了多尺度融合方案，提出了一种有效的两阶段神经网络体系结构搜索(NAS)策略，即BiX-NAS，该策略可自动搜索最优的双向体系结构。(2)分析了经典进化搜索算法搜索缺陷的瓶颈，提出了一种新的渐进式进化算法，进一步发现候选跳过子集，加快搜索过程。(3)以三个医学图像分割数据集为基准，与其他先进的图像分割方法相比有较大的差距。
二、方法简介：
我们首先讨论了将多尺度跳跃连接引入生物网作为直观升级的有效性，然后我们演示了BiX-NAS每个阶段的细节，最后，我们提出了跳跃公平性原则，以确保搜索公平性和效率。
图1所示。BiO-Net、BiO-Net++和BiX-Net的概述，具有4级和2次迭代。同一级别的编码器和解码器块被重复利用。
**2.1 **：
**1.BiO-Net++**轻量的BiO-Net多尺度升级版
如何轻量化模型是一个非常重要的问题。在一个前序工作中，BiO-Net会触发多个编码（encoding) 和解码 (decoding) 阶段，将同一尺度的特征用双向跳跃连接融合起来（图1左）。尽管这种循环设计可以大大缩小网络规模，但根据其预设的迭代次数，计算成本仍然不可避免地增加。
本文首先将多尺度的跳跃连接引入BiO-Net，提出了BiO-Net++（图1中）。注意，特征融合策略由（连接）concatnation优化为element-wise（元素）平均，以降低网络总体复杂度。
虽然上述设计促进了多尺度特征融合并缩小了网络规模，但从实验中发现，这种密集的连接会增加过量的计算量却仅仅在整体性能方面带来微小改善（表1）。本文目标寻找一种更稀疏连接的BiO-Net++子结构，让模型不仅可以从多尺度融合中受益，还可以最大程度地减轻计算负担。
2.BiX-NAS:超高效搜索BiO-Net++子结构
为此，本文提出了一种两阶段式搜索算法BiX-NAS，以找到稀疏连接的子结构。在阶段一中，其采用了一种可训练的选择矩阵（selection matrix）来缩小整体的搜索空间，并在阶段二中引入进化 NAS来逐步进化出最佳的子结构。
2.1. Phase1（第一阶段）：通过（选择矩阵）selection matrix缩小搜索空间
为了降低整体搜索难度，本文采用可微分NAS算法来快速筛掉那些易于发现的无效跳跃连接。
假设待搜索的模块有N个输入跳跃连接，本文预设在稀疏连接的结构中，只有 个候选的跳跃连接可以被保留。因此，对于有L层和T次迭代的BiO-Net++ 来说, 搜索空间大概是：
本文通过构建一个可学习的（选择矩阵）selection matrix M来对N个传入连接和k个候选连接之间的映射 进行建模，并将 公式化为以下完全可微的方程：
2.2. Phase2（第二阶段）: 渐进式进化搜索确定最佳子结构
为了进一步减少阶段一中搜到结构的潜在冗余，本文引入了一种新颖的进化搜索算法以找到更好的网络。
具体来说，此算法将同时搜索一对extraction stage（提取阶段）之间所有的跳跃连接 （见上图黄色标注），然后在当前搜索结束后逐步移动到下一对。注意，由于网络依赖从前往后的拓扑顺序， 因此在搜索时需要从最后一对extraction stage（提取阶段）开始，并逐步移动到第一个。
在传统的进化NAS方法中，最直接的策略是从SuperNet采样不同的子结构（由某些跳跃连接组成），然后分别进行独立的训练。这种传统策略有两个主要缺陷：第一，分开训练子结构可能会导致不公平的结果；其次，根据实验和经验，搜索过程非常缓慢且低效。
2.3. 分析搜索的公平性：
为了克服上述第一个缺陷，本文提出了skip fairness（跳跃公平）的概念，并声称所有连接搜索算法都需要满足这一原则。请注意，每个采样的子结构都是从不断更新的SuperNet中随机抽取的，它构成了每次迭代的总体P。
上述概念要求，在不同的采样子结构中（e.g. ），任何对应的跳跃连接 （e.g. ）需要携带相同特征。否则，不一致的特征会影响搜索决策，从而导致不公平结果。由于单独训练采样的子结构产生的传入特征不一致，传统的策略违反了这一原则，因此不直接适用于搜索跳跃链接。
本文提出的阶段二NAS算法通过共享部分的传入特征来满足（跳跃公平）skip fairness。具体来说，假设在第t 和第t+1个extraction stage（提取阶段）之间搜索时：从第1个到第t-1个extraction stage（提取阶段）的结构是固定的，且可以共享这些stages的生成特征。这些stages（阶段）组成的子结构被称为head network （头网络）（图2a）。相反，第t到最后一个extraction stage（提取阶段）的网络结构随着不同的采样而变化，这种不固定的stages（阶段）组成的子结构集被称为tail networks（尾网络）（图2a)。这些tail networks（尾网络）共享相同的超网络权重，但具有不同的连接方式。如图2b，head network（头网络）生成的特征被传入到所有采样的tail networks（尾网络）中。然后对所有tail networks（尾网络）的loss进行平均，并且仅计算gradients一次。因此，阶段二的实际搜索过程非常的高效，克服了第二个缺陷。
在每对提取阶段之间的搜索完成后，如图2c，本文遵循多目标选择标准，该标准基于验证准确性 (IoU) 和计算复杂性 (MACs) 保留在帕累托前沿的子结构。算法细节和搜索到的BiX-Net详见论文。
三、实验结果
BiX-Net在nuclei segmentation（核分割）和organ segmentation（器官分割）上都达到了SOTA的准确率，但是仅需0.28M参数量和28G MACs。U-Net的计算量超过BiX-Net将近23倍，BiO-Net的计算量则是超过将近40倍。
细胞核分割：我们的方法与普通U-Net、最先进的循环U-Net变体和同类状态的技术NAS搜索网络进行了比较。所有模型都是从头开始训练的，最终结果报告为三次独立运行的平均值。表1显示，我们的普通BiX-Net的性能明显优于先进的NAS对等物，并且在显著降低网络复杂性的情况下取得了与BiO-Net相当的结果。值得注意的是，我们在TNBC上的结果远高于其他所有人，说明我们的泛化能力较强。所有数据集的定性比较如图3所示。
多器官分割：CHAOS挑战的目标是在CT或MRI上分别精确分割四个腹部器官:肝脏、左肾、右肾和脾脏。我们没有将每个类上的网络训练成几个独立的二值分割任务，而是复制所有模型，直接输出所有类的日志。与核分割相似，从表2可以看出，在最先进的NAS搜索网络中，BiX-Net在所有类中都取得了最好的性能，并且计算复杂度低得多。虽然手工制作的生物网比所有的比较方法都要好，但其计算复杂度是4.1倍，可训练参数数是39.4倍。此外，当所有器官在单个切片中显示时，BiX-Net产生了更好的分割掩模(图3)。
此外，我们还通过在核数据集上直接训练SuperNet BiO-Net++ (BiO-Net的多尺度升级)和Phase1（阶段一）搜索体系结构(Phase2没有搜索过)进行了两项消融研究。如表1所示，与BiO-Net++相比，我们的Phase1搜索算法减少了8.6%的MACs, BiX-Net最终实现了18.5%的MACs减少，这验证了BiX-NAS的（阶段一）Phase1和（阶段二）Phase2搜索的必要性。与之前的NAS工作不同，我们最终搜索的BiXNet遵循循环的双向范式，在不同的迭代中重复使用相同的构建块。注意，有一个构建块(级别4的编码器)在所有迭代中被跳过(补充材料图1)，导致总网络参数进一步减少。
对于所有指标，BiX-Net在所有数据集上获得的分数高于竞争的NAS同行，并使用我们提出的BiO-Net++以更少的计算实现了相同的结果。此外，我们进行了两个尾部配对t检验，以分析我们的方法与其他竞争NAS方法之间的统计显著性。BiX-Net实现p值<原子核数据集上的0.05和<在CHAOS数据集上，验证了我们方法的重要性。
四、结论：
在本文中，我们提出了一种高效的搜索编码器和解码器之间双向多尺度跳跃连接的两阶段NAS算法，即BiX-NAS。我们首先用一个新的选择矩阵跟踪可微分NAS以缩小搜索空间。在此基础上，提出了一种有效的渐进式进化搜索算法，以进一步减少跳跃冗余。在各种分割任务上的实验结果表明，搜索的BiX-Net以更少的参数和计算成本超过了最先进的NAS同行。