[Paper Review] Towards continual learning in medical imaging

Chaitanya Baweja, Ben Glocker+1 more|arXiv (Cornell University)|Nov 6, 2018

Radiology practices and education11 references42 citations

TL;DR

The paper evaluates Elastic Weight Consolidation (EWC) for sequentially learning two brain MRI segmentation tasks, showing EWC reduces catastrophic forgetting but leaves substantial room for improvement.

ABSTRACT

This work investigates continual learning of two segmentation tasks in brain MRI with neural networks. To explore in this context the capabilities of current methods for countering catastrophic forgetting of the first task when a new one is learned, we investigate elastic weight consolidation, a recently proposed method based on Fisher information, originally evaluated on reinforcement learning of Atari games. We use it to sequentially learn segmentation of normal brain structures and then segmentation of white matter lesions. Our findings show this recent method reduces catastrophic forgetting, while large room for improvement exists in these challenging settings for continual learning.

Motivation & Objective

• Motivate continual lifelong learning in medical imaging where previous task data may be unavailable.
• Evaluate how regularization-based approaches can mitigate forgetting when learning a new segmentation task.
• Quantify the trade-off between preserving old task performance and acquiring new-task performance in a biomedical setting.

Proposed method

• Formalize a continual learning setup with Task A (CSF/GM/WM segmentation) followed by Task B (white matter lesion segmentation).
• Adopt Elastic Weight Consolidation (EWC) to regularize parameter updates using Fisher information about Task A.
• Compare EWC against alternative baselines including fine-tuning, L2 regularization, and multi-task learning (upper bound).
• Use DeepMedic 3D CNN to perform segmentation on UK Biobank MRI data with T1 and FLAIR sequences.
• Compute task-wise Dice similarity coefficient (DSC) and examine performance on A during training for B to assess forgetting.

Experimental results

Research questions

• RQ1 Can EWC mitigate catastrophic forgetting when learning a new segmentation task in brain MRI without access to Task A data?
• RQ2 How does EWC compare to naive regularization and fine-tuning in preserving Task A performance while learning Task B?
• RQ3 What is the trade-off between learning Task B and retaining Task A across different regularization strengths?
• RQ4 Is continual learning via current methods close to an upper bound obtained by joint training on both tasks?

Key findings

• EWC reduces forgetting of Task A compared to naive fine-tuning or L2 regularization.
• With appropriate regularization strength, EWC preserves Task A performance better than simple L2 regularization while still acquiring Task B performance.
• Multi-task learning (training A and B together) yields the best performance (upper bound) but is infeasible when Task A data is unavailable during B’s training.
• Even with regularization, a gap remains between EWC and the multi-task upper bound, indicating room for improvement in continual learning for medical imaging.