[论文解读] MT-TransUNet: Mediating Multi-Task Tokens in Transformers for Skin Lesion Segmentation and Classification

Jingye Chen, Jieneng Chen+4 more|arXiv (Cornell University)|Dec 3, 2021

Cutaneous Melanoma Detection and Management被引用 15

一句话总结

MT-TransUNet 提出了一种基于 Transformer 的统一模型，通过协调特定任务的标记并强制注意力图与预测之间的一致性，联合执行皮肤病变分割与分类。该模型在 ISIC-2017 和 PH2 数据集上实现了最先进性能，参数量为 48M，推理时间为 0.17 秒，相较于先前方法在准确率和效率方面均有提升，同时利用了像素级与图像级标注。

ABSTRACT

Recent advances in automated skin cancer diagnosis have yielded performance on par with board-certified dermatologists. However, these approaches formulated skin cancer diagnosis as a simple classification task, dismissing the potential benefit from lesion segmentation. We argue that an accurate lesion segmentation can supplement the classification task with additive lesion information, such as asymmetry, border, intensity, and physical size; in turn, a faithful lesion classification can support the segmentation task with discriminant lesion features. To this end, this paper proposes a new multi-task framework, named MT-TransUNet, which is capable of segmenting and classifying skin lesions collaboratively by mediating multi-task tokens in Transformers. Furthermore, we have introduced dual-task and attended region consistency losses to take advantage of those images without pixel-level annotation, ensuring the model's robustness when it encounters the same image with an account of augmentation. Our MT-TransUNet exceeds the previous state of the art for lesion segmentation and classification tasks in ISIC-2017 and PH2; more importantly, it preserves compelling computational efficiency regarding model parameters (48M~vs.~130M) and inference speed (0.17s~vs.~2.02s per image). Code will be available at https://github.com/JingyeChen/MT-TransUNet.

研究动机与目标

• 解决先前方法将皮肤病变分类视为独立任务的局限性，忽略了病变分割带来的益处。
• 开发一种联合学习框架，利用分割与分类之间的共享特征以提升性能。
• 通过双任务与注意力区域一致性损失，增强在缺乏像素级标注图像上的模型鲁棒性。
• 在计算效率方面相较于先前的多任务模型，实现最先进性能的同时保持高效。
• 通过在两个任务间共享单一编码器，实现数据高效与模型高效的训练与推理。

提出的方法

• 在 Transformer 编码器-解码器架构中引入特定任务的分类与分割标记，以实现跨任务特征协调。
• 采用双任务一致性损失，将类别注意力图与分割图对齐，提升弱监督数据上的泛化能力。
• 应用注意力区域一致性损失，以强化分类与分割头之间的一致性，尤其在缺乏像素级标注时更具优势。
• 在两个任务间使用共享编码器，以减少模型参数并加速推理。
• 在半监督设置中，利用水平集函数作为分割监督的替代品，以处理无真实掩码的图像。
• 设计一个多任务学习流程，通过共享表征与一致性约束，使分割与分类相互引导。

实验结果

研究问题

• RQ1 与单任务方法相比，联合学习皮肤病变分割与分类是否能提升两个任务的性能？
• RQ2 当仅提供图像级标签时，如何强制分类与分割预测之间的一致性以提升鲁棒性？
• RQ3 共享的 Transformer 架构能否在显著减小模型规模并加快推理速度的前提下，实现最先进性能，超越先前的多任务模型？
• RQ4 该方法在无需微调的情况下，对不同数据集的泛化能力如何？
• RQ5 双任务与注意力区域一致性损失在弱监督数据上的性能影响如何？

主要发现

• MT-TransUNet 在 ISIC-2017 和 PH2 数据集上实现了最先进性能，分割任务的 Jaccard 分数达到 88.5/92.2，Dice 分数达到 93.6/95.9，分类任务准确率达到 97.1%。
• 模型将单张图像的推理时间缩短至 0.17 秒，相比 MB-DCNN 的 2.02 秒提升了 91.6%。
• 仅使用 4800 万参数，MT-TransUNet 相较于 MB-DCNN 的 1.3 亿参数，将模型规模减少了 63%。
• 在 ISIC-2017 上预训练后于 PH2 数据集上测试，该模型在识别准确率上提升 3.1%，在 Jaccard 分数上提升 2.8%，优于 MB-DCNN。
• 双任务与注意力区域一致性损失显著提升了对数据增强的鲁棒性，并改善了无像素级标注图像上的性能。
• 分割与分类的联合训练带来了相互的性能增益，证明了在皮肤镜图像分析中共享表征学习的优势。