[论文解读] Parallel Multi-Dimensional LSTM, With Application to Fast Biomedical Volumetric Image Segmentation

Marijn Stollenga, Wonmin Byeon+2 more|arXiv (Cornell University)|Jun 24, 2015

Advanced Neural Network Applications参考文献 31被引用 153

一句话总结

本文提出了一种新型可并行化的多维长短期记忆网络 PyraMiD-LSTM，通过将 LSTM 计算重新组织为金字塔形、网格对齐的拓扑结构，实现了高效的 GPU 并行化。与传统 MD-LSTM 相比，后者因存在顺序依赖性且 GPU 利用率低下而受限，PyraMiD-LSTM 在 3D 生物医学图像分割任务中实现了最先进性能，在 MRBrainS13 数据集上取得了新的 SOTA 结果，同时在 EM-ISBI12 上也表现出色。

ABSTRACT

Convolutional Neural Networks (CNNs) can be shifted across 2D images or 3D videos to segment them. They have a fixed input size and typically perceive only small local contexts of the pixels to be classified as foreground or background. In contrast, Multi-Dimensional Recurrent NNs (MD-RNNs) can perceive the entire spatio-temporal context of each pixel in a few sweeps through all pixels, especially when the RNN is a Long Short-Term Memory (LSTM). Despite these theoretical advantages, however, unlike CNNs, previous MD-LSTM variants were hard to parallelize on GPUs. Here we re-arrange the traditional cuboid order of computations in MD-LSTM in pyramidal fashion. The resulting PyraMiD-LSTM is easy to parallelize, especially for 3D data such as stacks of brain slice images. PyraMiD-LSTM achieved best known pixel-wise brain image segmentation results on MRBrainS13 (and competitive results on EM-ISBI12).

研究动机与目标

• 为克服传统多维长短期记忆网络（MD-LSTM）在 GPU 并行化方面的不足，从而限制其在大规模 3D 生物医学图像分割中的应用。
• 通过重新思考 MD-LSTM 的计算拓扑结构，实现在 MRI 和电子显微镜图像堆栈等 3D 体数据上的高效、可扩展推理。
• 通过采用能够捕捉完整时空上下文的循环架构，在 MRBrainS13 和 EM-ISBI12 等具有挑战性的基准上实现最先进分割性能。
• 证明当针对现代 GPU 硬件进行适当并行化时，MD-LSTM 可以在分割任务中超越卷积网络。

提出的方法

• 将标准 MD-LSTM 的计算顺序重构为金字塔形拓扑，使每一层以允许行内所有像素完全并行处理的方式处理 3D 体数据的 2D 切片。
• 采用改进的连接方式，增加对角连接，以消除由标准 45° 扫描拓扑旋转引入的上下文间隙。
• 通过将金字塔扫描视为扩张或重叠卷积的一种形式，利用高效的 CUDA 卷积操作在 GPU 上实现高吞吐量计算。
• 通过在所有 3D 体素位置上扩展金字塔扫描，将相同架构应用于 3D 数据，使用 8 个方向的扫描以捕获完整的体素上下文。
• 使用标准 LSTM 方程，包含输入门、遗忘门、输出门和单元状态，但重新排列计算顺序，以实现对整行或整列的独立处理。
• 由于金字塔结构，获得更大的有效感受野（大于 3×3），从而产生重叠且冗余的上下文窗口，提升了鲁棒性和性能。

实验结果

研究问题

• RQ1 通过重构 MD-LSTM 拓扑结构，是否能够在保持循环网络完整上下文建模能力的同时实现高效的 GPU 并行化？
• RQ2 所提出的 PyraMiD-LSTM 是否在 3D 生物医学图像数据集上的分割精度上优于现有的 CNN 和 MD-LSTM 基线方法？
• RQ3 金字塔形计算顺序是否能减少计算瓶颈并提升 3D 体积分割任务的推理速度？
• RQ4 PyraMiD-LSTM 中的重叠、多方向上下文在多大程度上提升了分割的鲁棒性和泛化能力？
• RQ5 该架构在真实世界高分辨率 3D 医学影像应用中是否具备可扩展性和有效性？

主要发现

• PyraMiD-LSTM 在 MRBrainS13 基准上取得了最佳已知结果，在所有指标（包括 DICE、修改后的豪斯多夫距离和体积差异）上均达到新的 SOTA 水平。
• 在 EM-ISBI12 数据集上，PyraMiD-LSTM 表现具有竞争力，非后处理方法中 Rand 错误率最低，优于其他领先方法。
• 在 MRBrainS13 数据集中，模型在灰质上达到 84.82% 的 DICE 分数，白质为 88.33%，脑脊液为 83.72%，在最终排名中超越所有其他团队。
• 模型性能未因使用 dropout 正则化而提升，事实上，dropout 反而增加了训练时间且未带来性能增益，表明该架构本身已具备高度鲁棒性。
• 金字塔形拓扑实现了高效的 GPU 并行化，减少了计算瓶颈，使 3D 体数据的推理速度相比标准 MD-LSTM 显著提升。
• 该方法证明，当针对现代硬件进行适当并行化时，具备完整上下文感知能力的循环网络可在分割任务中超越卷积网络。