Cross-Task Attention Network: Improving Multi-Task Learning for Medical Imaging Applications

TL;DR

This paper introduces a Cross-Task Attention Network (CTAN) that enhances multi-task learning in medical imaging by leveraging inter-task interactions, leading to significant accuracy improvements across diverse medical imaging tasks and datasets.

Contribution

The study presents a novel attention-based MTL framework with cross-task attention mechanisms, improving information sharing and performance over existing methods in medical imaging applications.

Findings

CTAN improves accuracy by 4.67% over single-task learning.

Outperforms baseline MTL methods like HPS and MTAN.

Effective across multiple medical imaging domains and tasks.

Abstract

Multi-task learning (MTL) is a powerful approach in deep learning that leverages the information from multiple tasks during training to improve model performance. In medical imaging, MTL has shown great potential to solve various tasks. However, existing MTL architectures in medical imaging are limited in sharing information across tasks, reducing the potential performance improvements of MTL. In this study, we introduce a novel attention-based MTL framework to better leverage inter-task interactions for various tasks from pixel-level to image-level predictions. Specifically, we propose a Cross-Task Attention Network (CTAN) which utilizes cross-task attention mechanisms to incorporate information by interacting across tasks. We validated CTAN on four medical imaging datasets that span different domains and tasks including: radiation treatment planning prediction using planning CT images of two different target cancers (Prostate, OpenKBP); pigmented skin lesion segmentation and diagnosis using dermatoscopic images (HAM10000); and COVID-19 diagnosis and severity prediction using chest CT scans (STOIC). Our study demonstrates the effectiveness of CTAN in improving the accuracy of medical imaging tasks. Compared to standard single-task learning (STL), CTAN demonstrated a 4.67% improvement in performance and outperformed both widely used MTL baselines: hard parameter sharing (HPS) with an average performance improvement of 3.22%; and multi-task attention network (MTAN) with a relative decrease of 5.38%. These findings highlight the significance of our proposed MTL framework in solving medical imaging tasks and its potential to improve their accuracy across domains.