[论文解读] How Can We Tame the Long-Tail of Chest X-ray Datasets?

Arsh Verma|arXiv (Cornell University)|Sep 8, 2023

COVID-19 diagnosis using AIMedicine被引用 3

一句话总结

本文提出，在目标数据集训练前，先在胸部X光（CXR）数据上微调视觉模型，可显著提升对罕见、长尾标签的性能表现——且无需使用复杂的损失函数或超参数调优。通过采用两阶段预训练策略（ImageNet后接CXR特定数据集），该方法在超罕见标签上实现最高3个百分点的性能提升，通过模型平均实现6%的mAP绝对增益，表明模型初始化质量是医学影像中长尾泛化能力的关键因素。

ABSTRACT

Chest X-rays (CXRs) are a medical imaging modality that is used to infer a large number of abnormalities. While it is hard to define an exhaustive list of these abnormalities, which may co-occur on a chest X-ray, few of them are quite commonly observed and are abundantly represented in CXR datasets used to train deep learning models for automated inference. However, it is challenging for current models to learn independent discriminatory features for labels that are rare but may be of high significance. Prior works focus on the combination of multi-label and long tail problems by introducing novel loss functions or some mechanism of re-sampling or re-weighting the data. Instead, we propose that it is possible to achieve significant performance gains merely by choosing an initialization for a model that is closer to the domain of the target dataset. This method can complement the techniques proposed in existing literature, and can easily be scaled to new labels. Finally, we also examine the veracity of synthetically generated data to augment the tail labels and analyse its contribution to improving model performance.

研究动机与目标

• 解决在胸部X光（CXR）分类任务中，对罕见、长尾标签性能表现不佳的挑战。
• 探究是否可通过领域特定的预训练进行模型初始化，在不修改损失函数或数据采样策略的前提下，提升对低频标签的泛化能力。
• 评估模型平均与合成数据增强在提升尾部标签性能方面的有效性。
• 提供一种可扩展、即插即用的方法，以最小化微调成本，提升对新出现的罕见标签的模型性能。

提出的方法

• 该方法采用两阶段预训练策略：首先使用ImageNet权重进行初始化，随后在三个开源CXR数据集（NIH CXR、PadChest、CheXpert）上进行微调，以获得领域自适应的初始化权重。
• 随后，模型在MIMIC CXR数据集上进行端到端微调，该数据集包含26个标签的长尾分布，其中5个标签为预训练数据集中未出现的独有尾部标签。
• 通过集成不同主干网络（DenseNet161与ResNeXt101）的模型预测结果，应用模型平均策略，提升对低资源标签的鲁棒性。
• 探索了合成数据增强方法：将MIMIC CXR与使用RoentGen生成的小型合成数据集（5,000个样本）联合微调，以提升尾部标签的性能。
• 所有模型均使用二元交叉熵损失函数与Sigmoid激活函数，训练20个周期，输入分辨率保持在448×448，以保留细微解剖结构信息。
• 训练流程以患者为单位进行，每个患者的所有可用视图均被使用，数据划分也以患者为单位进行，以避免数据泄露。

实验结果

研究问题

• RQ1 在不修改损失函数或数据采样策略的前提下，是否可通过在多样化CXR数据集上进行预训练，提升模型对罕见、长尾标签的性能表现？
• RQ2 相较于标准的ImageNet初始化，采用两阶段预训练策略（先ImageNet，后CXR特定数据）是否能实现对尾部标签更好的泛化能力？
• RQ3 模型平均在提升低频标签性能方面有多大作用，特别是对目标数据集中独有的标签？
• RQ4 合成数据增强是否有助于缓解尾部标签的数据稀缺问题，并在模型性能上带来可测量的增益？

主要发现

• 链式预训练策略使DenseNet161在5个独有的尾部标签（Tail-U）上性能提升3个百分点，而ResNeXt101仅提升0.5个百分点，表明在超罕见类别上取得了显著增益。
• 通过模型平均，Tail-U标签的平均平均精度（mAP）实现了6%的绝对提升（从0.1806提升至0.2430），优于所有单个模型的性能表现。
• 平均模型在验证集的26个标签中，有19个标签表现最佳，并在所有类别评估（头部、中等、尾部）中均排名第一。
• 尽管合成数据增强规模有限（5,000个样本），但其在尾部标签上表现出可测量的性能提升，尤其在尾部类别中增益最为显著，显示出未来扩展的潜力。
• 该方法具有高度可扩展性：新增标签仅需少量微调周期，使其在临床实际部署中具备实用性。
• 结果表明，模型初始化是长尾泛化中一个关键但常被低估的因素，尤其在医学影像中，罕见病变更具有临床重要性。