[论文解读] Invert to Learn to Invert

Patrick Putzky, Max Welling|arXiv (Cornell University)|Nov 25, 2019

Generative Adversarial Networks and Image Synthesis被引用 38

一句话总结

本论文提出一种内存高效、全可逆的迭代反向模型（i-RIM），由可逆神经网络构成，使得3D MRI 重建可以非常深，并在训练时保持恒定内存，达到最先进的结果。

ABSTRACT

Iterative learning to infer approaches have become popular solvers for inverse problems. However, their memory requirements during training grow linearly with model depth, limiting in practice model expressiveness. In this work, we propose an iterative inverse model with constant memory that relies on invertible networks to avoid storing intermediate activations. As a result, the proposed approach allows us to train models with 400 layers on 3D volumes in an MRI image reconstruction task. In experiments on a public data set, we demonstrate that these deeper, and thus more expressive, networks perform state-of-the-art image reconstruction.

研究动机与目标

• 使用可逆网络来解决训练深度迭代反问题模型时的内存瓶颈。
• 在不存储激活值的情况下，实现对逆问题进行非常深的架构训练（多达 400 层）。
• 将迭代重建方法扩展到 MRI 的大规模 3D 观测。
• 提出一种稳定的可逆层以支持深层 i-RIM 并提升训练稳定性。

提出的方法

• 将迭代反问题模型与可逆学习结合，在训练期间实现恒定内存。
• 通过重构更新方程使其可逆，开发可逆的 Recurrent Inference Machines（i-RIM）。
• 引入具有正交 1x1 卷积的可逆层，以确保稳定的逆操作与简易的逆推计算。
• 在可逆框架内使用带空间下采样的残差块和门控线性单元，以实现稳定训练。
• 在来自 fastMRI 挑战的加速 MRI 数据上进行训练和评估，并与 U-Net 和 RIM 基线进行比较。

实验结果

研究问题

• RQ1 能否使用可逆神经网络在训练中实现恒定内存来训练迭代反问题模型？
• RQ2 为使 RIMs 完全可逆且在深层 3D 模型中保持稳定，需要哪些结构调整？
• RQ3 i-RIM 在 4x 和 8x 加速下的 MRI 重建是否达到最先进水平？
• RQ4 i-RIM 的内存高效性如何使得对大型 3D 体积训练更深的网络成为可能？

主要发现

• 提出的 i-RIM 在训练中实现了 O(1) 的内存复杂度，使得非常深的模型（多达 400 层）成为可能。
• 在 fastMRI 数据上的 MRI 重建中，i-RIM 在 4x 和 8x 加速设置下在 NMSE、PSNR 和 SSIM 指标的验证和测试中，优于 U-Net 基线和标准 RIM。
• 3D i-RIMs 接近 2D i-RIMs 的性能，表明对体积数据具备很强的可扩展性。
• Table 1 展示了内存节省，支持原本在非可逆训练下不可行的 3D 模型。
• 该方法在发表时使 i-RIM 位列 single-coil fastMRI 榜首。