[论文解读] Searching Collaborative Agents for Multi-plane Localization in 3D Ultrasound

Yuhao Huang, Xin Yang+11 more|arXiv (Cornell University)|Jul 30, 2020

Face recognition and analysis参考文献 26被引用 4

一句话总结

该论文提出了一种多智能体强化学习（MARL）框架，结合一次性的神经架构搜索（NAS）与基于RNN的协作机制，用于在3D超声图像中同时定位多个标准平面（正中矢状面、横断面、冠状面）。通过使用基于梯度的可微分神经架构搜索（GDAS）优化各智能体网络，并利用RNN建模平面之间的空间关系，该方法在降低模型规模和计算成本的同时，实现了冠状面定位的最先进精度（5.93°/0.89 mm）。

ABSTRACT

3D ultrasound (US) is widely used due to its rich diagnostic information, portability and low cost. Automated standard plane (SP) localization in US volume not only improves efficiency and reduces user-dependence, but also boosts 3D US interpretation. In this study, we propose a novel Multi-Agent Reinforcement Learning (MARL) framework to localize multiple uterine SPs in 3D US simultaneously. Our contribution is two-fold. First, we equip the MARL with a one-shot neural architecture search (NAS) module to obtain the optimal agent for each plane. Specifically, Gradient-based search using Differentiable Architecture Sampler (GDAS) is employed to accelerate and stabilize the training process. Second, we propose a novel collaborative strategy to strengthen agents' communication. Our strategy uses recurrent neural network (RNN) to learn the spatial relationship among SPs effectively. Extensively validated on a large dataset, our approach achieves the accuracy of 7.05 degree/2.21mm, 8.62 degree/2.36mm and 5.93 degree/0.89mm for the mid-sagittal, transverse and coronal plane localization, respectively. The proposed MARL framework can significantly increase the plane localization accuracy and reduce the computational cost and model size.

研究动机与目标

• 解决在3D超声容积中自动、同步定位多个标准平面的挑战。
• 通过实现智能体间通信以建模平面间空间关系，克服单智能体强化学习的局限性。
• 通过为每个平面特定智能体进行自动化的神经架构搜索（NAS），降低模型复杂度与计算成本。
• 提升在3D超声中常见噪声多、解剖结构多变环境下的鲁棒性，尤其针对子宫影像。
• 开发一种可泛化的框架，适用于3D医学影像中的多平面定位。

提出的方法

• 设计了一种多智能体强化学习（MARL）框架，包含三个独立智能体，分别负责定位正中矢状面、横断面和冠状面。
• 采用一次性的基于梯度的可微分神经架构搜索（GDAS），自动发现每个平面特定智能体的最优神经网络架构，提升性能并减少模型规模。
• 引入基于RNN的协作策略，使智能体能够交换信息并学习三个平面之间的空间关系，从而提升联合定位精度。
• 引入基于关键点的对齐模块，为智能体提供热启动，提升在噪声较多的3D超声容积中的训练稳定性和收敛性。
• 基于预测平面与真实平面之间的几何相似性（二面角与到原点的距离）和内容相似性（SSIM），定义了奖励函数。
• 使用自建的1,000个体积的3D超声数据集进行训练，并通过空间与内容指标评估性能。

实验结果

研究问题

• RQ1 多智能体强化学习框架是否能在3D超声中对多个标准平面的定位上优于单智能体方法？
• RQ2 为每个平面特定智能体引入可微分神经架构搜索（GDAS）是否能提升定位精度并降低模型复杂度？
• RQ3 基于RNN的通信机制是否能有效建模3D超声中三个解剖上不同的平面之间的空间关系？
• RQ4 与现有最先进方法相比，所提框架在3D超声平面定位中的精度、计算成本与模型规模方面表现如何？
• RQ5 该方法对子宫内装置（IUD）、妊娠及先天性异常等解剖变异是否具有鲁棒性？

主要发现

• 所提出的G-MARL-R方法在整体性能上表现最佳，冠状面定位的平均误差为5.93°/0.89 mm，正中矢状面为7.05°/2.21 mm，横断面为8.62°/2.36 mm。
• 与基于VGG的MARL基线相比，该方法将FLOPs降低了63%，参数量减少了70%，显著降低了计算成本。
• 与DARTS相比，基于GDAS的NAS方法减少了训练时间与内存使用，G-MARL-R的训练效率达到0.04天/epoch，而D-MARL-R为0.12天/epoch。
• 引入基于RNN的协作策略后，角度与距离之和（SAD）降低，Q值曲线更加稳定，表明收敛性与精度均得到改善。
• 消融实验证实，NAS与RNN协作策略在性能提升中均具有独立且协同的作用，G-MARL-R优于所有基线方法。
• 该方法展现出强大的泛化能力，在大规模内部数据集上，所有平面的SSIM均达到0.80，平均位移仅为1.82 mm。