[论文解读] AMI-Net+: A Novel Multi-Instance Neural Network for Medical Diagnosis from Incomplete and Imbalanced Data

Zeyuan Wang, Josiah Poon+1 more|arXiv (Cornell University)|Jul 3, 2019

Machine Learning in Healthcare参考文献 17被引用 7

一句话总结

AMI-Net+ 是一种新颖的多实例神经网络，通过整合焦点损失、门控注意力机制以及自适应多实例池化机制，提升了从不完整且不平衡的临床数据中进行医学诊断的性能。该模型在两个真实世界医学数据集上优于现有模型，展现出在处理片段化和偏态的症候-疾病关系方面的卓越性能。

ABSTRACT

In medical real-world study (RWS), how to fully utilize the fragmentary and scarce information in model training to generate the solid diagnosis results is a challenging task. In this work, we introduce a novel multi-instance neural network, AMI-Net+, to train and predict from the incomplete and extremely imbalanced data. It is more effective than the state-of-art method, AMI-Net. First, we also implement embedding, multi-head attention and gated attention-based multi-instance pooling to capture the relations of symptoms themselves and with the given disease. Besides, we propose var-ious improvements to AMI-Net, that the cross-entropy loss is replaced by focal loss and we propose a novel self-adaptive multi-instance pooling method on instance-level to obtain the bag representation. We validate the performance of AMI-Net+ on two real-world datasets, from two different medical domains. Results show that our approach outperforms other base-line models by a considerable margin.

研究动机与目标

• 解决在真实世界研究中，从不完整且极度不平衡的医学数据中训练鲁棒诊断模型的挑战。
• 在 AMI-Net 框架基础上，通过引入先进的注意力机制和损失函数，实现性能提升。
• 开发一种自适应多实例池化策略，动态选择关键实例以构建袋级别表征。
• 通过真实世界数据集，在多样化医学领域中验证模型的有效性。

提出的方法

• 用焦点损失替代标准交叉熵损失，以缓解训练过程中的类别不平衡问题。
• 采用多头注意力和门控注意力机制，建模症候内部关系及症候-疾病关系。
• 引入一种自适应多实例池化层，根据实例与疾病的关联性动态学习实例权重。
• 在注意力和池化操作前，使用嵌入层对症候特征进行编码。
• 采用端到端训练，优化整个网络，以实现从袋级别症候数据中进行疾病分类。
• 采用多实例学习（MIL）范式，其中每位患者被视为一个‘袋’（包含多个症候实例），疾病标签分配给该袋。

实验结果

研究问题

• RQ1 多实例神经网络能否在真实世界环境中，有效学习来自不完整且高度不平衡的医学数据？
• RQ2 在罕见病诊断背景下，焦点损失相较于标准交叉熵损失，如何提升模型性能？
• RQ3 与固定池化策略相比，自适应池化在多实例学习中在多大程度上提升了表征学习能力？
• RQ4 注意力机制的集成是否提升了模型识别关键症候以辅助诊断的能力？
• RQ5 AMI-Net+ 在不同医学领域中，面对不同数据特征时，其泛化能力如何？

主要发现

• AMI-Net+ 在两个来自不同医学领域的现实世界医学数据集上，显著优于基线模型，包括 AMI-Net。
• 采用焦点损失显著提升了少数疾病类别上的性能，减轻了类别不平衡的负面影响。
• 自适应多实例池化机制通过聚焦最相关症候，提升了模型的可解释性和准确性。
• 注意力机制，特别是门控注意力，增强了模型捕捉复杂症候-疾病依赖关系的能力。
• 在 AUC 和 F1 分数方面，该模型达到了最先进水平，相较于现有基于 MIL 的方法实现了可测量的性能提升。