NCAdapt: Dynamic adaptation with domain-specific Neural Cellular Automata for continual hippocampus segmentation

TL;DR

This work introduces NCAdapt, a Neural Cellular Automata (NCA) based method designed to address continuous learning challenges in medical imaging with a domain-specific multi-head structure, integrating adaptable convolutional layers into the NCA backbone for each new domain encountered.

Abstract

Continual learning (CL) in medical imaging presents a unique challenge, requiring models to adapt to new domains while retaining previously acquired knowledge. We introduce NCAdapt, a Neural Cellular Automata (NCA) based method designed to address this challenge. NCAdapt features a domain-specific multi-head structure, integrating adaptable convolutional layers into the NCA backbone for each new domain encountered. After initial training, the NCA backbone is frozen, and only the newly added adaptable convolutional layers, consisting of 384 parameters, are trained along with domain-specific NCA convolutions. We evaluate NCAdapt on hippocampus segmentation tasks, benchmarking its performance against Lifelong nnU-Net and U-Net models with state-of-the-art (SOTA) CL methods. Our lightweight approach achieves SOTA performance, underscoring its effectiveness in addressing CL challenges in medical imaging. Upon acceptance, we will make our code base publicly accessible to support reproducibility and foster further advancements in medical CL.

Figures (6)

• Figure 1: Simulation of 467 training stages comparing the estimated carbon footprint between Lifelong nnU-Net and NCAdapt. To estimate per stage, the codecarbon python package is used codecarbon (Lifelong nnU-Net: 616 kWh (entire Route 66 of 3940 km); NCAdapt: 122 kWh). The route is based on the average range of a Tesla Model 3, where 1 kWh equals 6.4 km (4 miles).
• Figure 1: Comparison of BWT and FWT performance against number of parameters for NCAdapt and sequential U-Net trained on the Cardiac datasets; smaller boxes indicate superior performance. Note: Values above zero on the y-axis represents FWT and values below BWT.
• Figure 2: NCAdapt for medical image segmentation using a M3D-NCA with an M3D-NCA as a backbone while having shared layers and domain-specific convolutional layers.
• Figure 3: Comparison of BWT and FWT performance against number of parameters for NCAdapt, sequential nnU-net and U-Net; smaller boxes indicate superior performance. Note: Values above zero on the y-axis represents FWT and values below BWT.
• Figure 4: CL performance for NCAdapt, NCA with EWC and RWalk, Lifelong nnU-Net (sequential, EWC and RWalk) and U-Net (Sequential and AGem) using Dice scores with positive BWT and FWT; the larger the covered area the better the method.