LoRA-PT: Low-Rank Adapting UNETR for Hippocampus Segmentation Using Principal Tensor Singular Values and Vectors

2024-07-16

Artif. Intell. Medicine

AI Summary

To address the challenges of scarce annotated data and high computational cost in fine-tuning pre-trained UNetR for hippocampal segmentation in medical imaging, this paper proposes a novel low-rank adaptation method based on tensor singular value decomposition (t-SVD). Specifically, Transformer parameter matrices are reconstructed as third-order tensors according to their architectural scale; t-SVD is then applied to separate dominant components from residuals, enabling updates solely to the principal singular values and corresponding singular vectors. This approach falls under the parameter-efficient fine-tuning (PEFT) paradigm. Evaluated on three public hippocampal datasets, it achieves Dice score improvements of 1.2–2.8% over state-of-the-art PEFT methods, while requiring only 0.08% trainable parameters. Consequently, it significantly reduces GPU memory consumption and training time. The method establishes an efficient, scalable, and highly accurate paradigm for few-shot medical image segmentation.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Calibration & Uncertainty Quantification
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

The hippocampus is an important brain structure involved in various psychiatric disorders, and its automatic and accurate segmentation is vital for studying these diseases. Recently, deep learning-based methods have made significant progress in hippocampus segmentation. However, training deep neural network models requires substantial computational resources, time, and a large amount of labeled training data, which is frequently scarce in medical image segmentation. To address these issues, we propose LoRA-PT, a novel parameter-efficient fine-tuning (PEFT) method that transfers the pre-trained UNETR model from the BraTS2021 dataset to the hippocampus segmentation task. Specifically, LoRA-PT divides the parameter matrix of the transformer structure into three distinct sizes, yielding three third-order tensors. These tensors are decomposed using tensor singular value decomposition to generate low-rank tensors consisting of the principal singular values and vectors, with the remaining singular values and vectors forming the residual tensor. During fine-tuning, only the low-rank tensors (i.e., the principal tensor singular values and vectors) are updated, while the residual tensors remain unchanged. We validated the proposed method on three public hippocampus datasets, and the experimental results show that LoRA-PT outperformed state-of-the-art PEFT methods in segmentation accuracy while significantly reducing the number of parameter updates. Our source code is available at https://github.com/WangangCheng/LoRA-PT/tree/LoRA-PT.

Problem

• Efficiently adapts pre-trained models for hippocampus segmentation
• Reduces computational resources and parameter updates required
• Addresses limited labeled data in medical image segmentation

Innovation

• Uses tensor singular value decomposition for efficiency
• Updates only principal singular values and vectors
• Fine-tunes pre-trained UNETR with fewer parameter updates