3640 - Continual Whole-Body Organ Segmentation in 3D CT Scans Using Vision Transformer and Low-Rank Adaptation

Purpose/Objective(s): Deep segmentation networks achieve high performance when trained on specific datasets. However, in clinical practice, patient privacy and data storage constraints often require that pretrained models are capable of adapting dynamically, allowing for the segmentation of new organs without access to the previous training datasets. This clinically preferred process can be viewed as a continual semantic segmentation (CSS) problem, which is a non-trivial task & understudied in medical imaging. Previous CSS works would either experience catastrophic forgetting or lead to unaffordable memory costs as models expand. In this study, we introduce a novel continual whole-body organ segmentation model using a lightweight low-rank adaptation mechanism (LoRA), which effectively segments 121 organs without catastrophic forgetting and meanwhile maintaining a low parameter increasing rate (PIR, %).

Materials/Methods: For the model development and validation, we train a 3D Pyramid Vision Transformer (PVT) based segmentation model on the public TotalSegmentator dataset (1204 CTs with 103 labeled organs). Subsequently, for each new learning task, we freeze the already trained network parameters and introduce additional lightweight trainable LoRA parameters to continually segment new head-neck organs (9 organs, 244 CTs) and new chest organs (9 organs, 153 CTs) over 2 in-house datasets. In each continual step, trainable LoRA parameters are incorporated into the patch-embedding, multi-head attention, and feed-forward layers. For validation, 20% of each dataset is reserved as test set. We report the Dice score (DSC, %) & 95% Hausdorff Distance (HD95, mm) as evaluation metrics.

Results: Our proposed model exhibits segmentation performance comparable to the PVT model separately developed on each dataset (achieving an average DSC of 89.8% and HD95 of 4.1mm), while only requiring about 37% of the parameters. Other CSS models either experience catastrophic forgetting (yielding very low performance, e.g., <25% DSC for [PLOP, MiB]) or have exploding PIR (32% for [SUN]). In comparison, our model avoids the knowledge forgetting (segmentation performance very similar to the upper bound) and meanwhile has the PIR of only 5.6%.