[论文解读] Automatic Data Augmentation via Deep Reinforcement Learning for Effective Kidney Tumor Segmentation

Tiexin Qin, Ziyuan Wang+4 more|arXiv (Cornell University)|Feb 22, 2020

Advanced Neural Network Applications参考文献 11被引用 4

一句话总结

本文提出了一种基于深度强化学习（DRL）的自动数据增强方法，用于CT扫描中的肾肿瘤分割，其中智能体通过在验证集上最大化Dice分数提升，端到端地学习最优增强策略。该方法优于传统和基于GAN的增强技术，实现了75.8%的mIoU和84.0%的DSC，表现出在数据有限情况下的优越边界定位能力和鲁棒性。

ABSTRACT

Conventional data augmentation realized by performing simple pre-processing operations (\eg, rotation, crop, \etc) has been validated for its advantage in enhancing the performance for medical image segmentation. However, the data generated by these conventional augmentation methods are random and sometimes harmful to the subsequent segmentation. In this paper, we developed a novel automatic learning-based data augmentation method for medical image segmentation which models the augmentation task as a trial-and-error procedure using deep reinforcement learning (DRL). In our method, we innovatively combine the data augmentation module and the subsequent segmentation module in an end-to-end training manner with a consistent loss. Specifically, the best sequential combination of different basic operations is automatically learned by directly maximizing the performance improvement (\ie, Dice ratio) on the available validation set. We extensively evaluated our method on CT kidney tumor segmentation which validated the promising results of our method.

研究动机与目标

• 为解决因标注训练数据有限而导致的医学图像分割性能下降问题。
• 克服传统数据增强方法的局限性，后者依赖于固定且随机的操作，缺乏对分割性能的反馈。
• 开发一种端到端学习框架，通过强化学习发现图像特定的增强策略，联合优化数据增强与分割。
• 在增强策略学习过程中直接最大化分割性能（Dice分数），确保生成的数据能提升下游任务的准确性。
• 在真实世界的CT肾肿瘤分割数据集上验证该方法，并与最先进增强基线方法进行比较。

提出的方法

• 该方法将数据增强建模为马尔可夫决策过程（MDP），其中智能体选择图像变换（如翻转、裁剪、旋转）以提升分割性能。
• 采用双-stream深度Q网络（DQN）学习最优的增强操作序列，奖励函数定义为验证集上Dice分数的提升。
• 使用分割模型作为替代策略，在DRL训练期间提供反馈，实现增强与分割的端到端联合优化。
• 智能体与环境（CT图像及掩码）交互，选择动作对输入图像进行变换，并根据最终分割质量获得奖励。
• 最终的增强数据集通过混合原始图像与DRL生成的增强图像构成，使训练集规模翻倍。
• 使用Adam优化器进行训练，所有方法共享相同的超参数以确保公平比较。

实验结果

研究问题

• RQ1 强化学习智能体能否学习到针对特定数据的增强序列，从而在有限的医学影像数据上最大化分割性能？
• RQ2 与单独训练增强和分割模块相比，增强与分割的端到端联合训练是否能带来更好的性能？
• RQ3 所提出的基于DRL的增强方法在肾肿瘤分割中与传统数据增强和基于GAN的数据生成方法相比表现如何？
• RQ4 智能体最常选择的增强操作类型（如翻转、裁剪）是什么？这些操作是否与肿瘤边界分割的改善相关？
• RQ5 该方法能否在标注有限的真实临床数据集上实现泛化，并达到最先进性能？

主要发现

• 所提方法实现了75.8%的平均交并比（mIoU），显著优于传统数据增强方法（74.3%）及其他基线方法。
• 该方法实现了84.0%的Dice相似系数（DSC），超过次优方法（传统增强方法为83.2%），并展现出最高的敏感性（94.0%）和阳性预测值（PPV: 80.2%）。
• 质心距离（CD）降低至8.86 mm，豪斯多夫距离（HD）降低至8.52 mm，平均表面距离（ASD）降低至4.01 mm，表明边界定位精度更优。
• 智能体频繁选择水平翻转和裁剪操作，这些操作增强了肿瘤的可见性，并改善了对小尺寸或不规则形状肿瘤的泛化能力。
• 可视化结果证实，该方法在低对比度或不规则形状区域中能生成更精确的肿瘤边界预测。
• 消融实验表明，通过分割模块反馈实现的端到端联合学习，相比孤立增强或固定操作集，能带来更优的策略学习效果。