[论文解读] Evolution-based Fine-tuning of CNNs for Prostate Cancer Detection

Khashayar Namdar, Isha Gujrathi+2 more|arXiv (Cornell University)|Nov 4, 2019

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

一句话总结

本文提出一种基于进化计算的微调方法，采用遗传算法（GA）优化预训练卷积神经网络（CNN）的全连接（FC）层，以提升六通道扩散加权磁共振成像（DW-MRI）中前列腺癌检测的性能。通过在层级别而非单个权重上应用GA，该方法在1,334张测试切片上将测试AUC提升9.3%（从0.707提升至0.773），实现了显著的性能增益，且无需微调特征提取层。

ABSTRACT

Convolutional Neural Networks (CNNs) have been used for automated detection of prostate cancer where Area Under Receiver Operating Characteristic (ROC) curve (AUC) is usually used as the performance metric. Given that AUC is not differentiable, common practice is to train the CNN using a loss functions based on other performance metrics such as cross entropy and monitoring AUC to select the best model. In this work, we propose to fine-tune a trained CNN for prostate cancer detection using a Genetic Algorithm to achieve a higher AUC. Our dataset contained 6-channel Diffusion-Weighted MRI slices of prostate. On a cohort of 2,955 training, 1,417 validation, and 1,334 test slices, we reached test AUC of 0.773; a 9.3% improvement compared to the base CNN model.

研究动机与目标

• 解决尽管使用标准反向传播和可微损失函数，基于CNN的前列腺癌检测仍存在AUC性能不佳的问题。
• 通过在分类器头引入进化计算，克服SGD在优化AUC等不可微指标时的局限性。
• 通过在层级别而非单个权重上操作，降低基于GA的CNN优化计算成本。
• 证明对全连接层进行进化微调可显著提升AUC，且无需重新训练卷积特征提取器。
• 在包含414名患者共5,706张前列腺切片的真实DW-MRI数据集上验证该方法。

提出的方法

• 使用包含3个卷积层和3个全连接（FC）层的预训练CNN，权重通过Xavier和均匀分布初始化。
• 首先使用交叉熵损失和SGD（固定初始学习率、动量和L2正则化）对CNN进行训练。
• 仅对全连接层应用遗传算法，将每一层视为单一实体进行交叉与变异操作。
• 初始GA种群包含512个分类器：1个来自训练好的CNN，其余511个随机初始化。
• 基于训练集上的AUC评估适应度，经过多代迭代后应用选择、交叉与变异（概率为1%）。
• 当验证集与训练集AUC不再提升时停止优化，最终模型根据最高验证AUC选取。

实验结果

研究问题

• RQ1 遗传算法能否有效优化CNN的全连接层，从而提升DW-MRI中前列腺癌检测的AUC？
• RQ2 与节点级别优化相比，层级别GA的变异与交叉是否能在保持性能增益的同时降低计算成本？
• RQ3 即使特征提取层被冻结，进化微调是否仍能在AUC上超越标准SGD训练？
• RQ4 AUC的提升在训练集、验证集和独立测试集上是否具有统计显著性且具鲁棒性？
• RQ5 与常规的AUC监控训练相比，该方法在最终模型性能上表现如何？

主要发现

• 所提出的基于GA的微调方法将测试AUC从SGD训练模型的0.707提升至0.773，相对提升9.3%。
• 最佳模型仅在三轮GA迭代内达成，验证AUC峰值达0.815，测试AUC达0.773。
• 该方法仅在Geforce GTX 1060 GPU上耗时10分钟即实现最优性能，证明了其计算可行性。
• GA将训练集AUC从SGD的0.867提升至0.877，表明泛化能力更强。
• 通过进化搜索绕过反向传播约束，成功优化了不可微的AUC指标。
• 与节点级别优化相比，层级别GA方法显著降低了计算开销，使其适用于低性能硬件。