The official implementation of Regularized Low-Rank Adaptation for Few-Shot Organ Segmentation.
📜 Medical Image Computing and Computer Assisted Intervention (MICCAI)
Ghassen Baklouti1,2,
Julio Silva-RodrÃguez1,
Jose Dolz1,2,
Houda Bahig2,
Ismail Ben Ayed1,2
1ÉTS Montréal ⋅ 2 Centre Hospitalier de l’Université de Montréal (CRCHUM)
| Conference | ArXiv |
Introduction
Parameter-efficient fine-tuning (PEFT) of pre-trained foundation models is increasingly attracting interest in medical imaging due to its effectiveness and computational efficiency. Among these methods, Low-Rank Adaptation (LoRA) is a notable approach based on the assumption that the adaptation inherently occurs in a low-dimensional subspace. Despite its empirical success and its ability to reduce the number of trainable parameters, LoRA still operates with a fixed rank that must be manually specified and cannot be changed throughout training. This constraint limits its flexibility, as the optimal rank often varies across downstream tasks (see Fig. 1). This limitation is exacerbated in few-shot regimes, where relying on validation data for finding the optimal configuration is unrealistic. To address this challenge, We propose ARENA, a novel Adaptive Rank Segmentation method that dynamically adjusts the intrinsic rank during adaptation. By Viewing the low-rank representation of the trainable weight matrices as a singular value decomposition (SVD), our method introduces an l1 sparsity regularizer into the loss function, and tackle it using a proximal optimizer. This regularizer acts as a penalty on the decomposition rank, enabling the model to automatically discover task-adapted ranks during training.
Figure 1: Adaptive-LoRA. We introduce a novel few-shot PEFT technique for Adaptive Rank Segmentation (ARENA) that is (a) robust to rank initialization and (b) enhances the parameter efficiency vs. performance trade-off.
ARENA is evaluated under a realistic few-shot fine-tuning setting, where it is first compared to the standard LoRA and then to several other PEFT baselines across two distinguishable segmentation tasks: base organs and novel organs. Extensive experiments demonstrate that our method yields notable performance gains, particularly when adapting foundation models to new tasks. ARENA not only enhances efficiency but also shows strong robustness to suboptimal rank initialization. Further experimental details and analysis are provided in the full paper Regularized Low-Rank Adaptation for Few-Shot Organ Segmentation .
Requirements
Environment Configuration
To set up the environment for running ARENA, create a new Python environment and install the required packages. Make sure the PyTorch version is compatible with your CUDA setup.
conda create -n arena python=3.9 -y
conda activate arena
conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.7 -c pytorch -c nvidia
pip install 'monai[all]'
pip install -r requirements.txt
Datasets Installation
To evaluate the adaptation performance of ARENA, we employed two datasets that exhibit severe domain drifts with respect to the pre-training distribution.
| Dataset | Selected Tasks |
|---|---|
| TotalSegmentator | Binary segmentation on 9 base organs |
| Parcellation of novel structures: heart, lung, and gluteus | |
| FLARE22 | Multi-class segmetation of 9 base organs |
- The employed train/test splits are located at
local_data/partitions/transferability.txt. - Check
local_data/datasets/README.mdfor an overview on how to organize these datasets.
Usage
Experimenting with ARENA
To run experiments with ARENA, use the following command:
python main_fseft.py --model_id fseft --dataset totalseg --organ organ_name --k 5 --method ARENA --seeds 3 --decoder frozen --rank 8
In this configuration, you can choose from the following base organs {spleen; kidney_left; gallbladder; esophagus; liver; pancreas; stomach; duodenum; aorta} or from the following novel organs {heart_myocardium heart_atrium_left heart_atrium_right heart_ventricle_left heart_ventricle_right}. Note that the decoder is frozen when transferring to known tasks (--decoder frozen), and entirely updated when adapting to novel organs (--decoder fine-tuned).
To experiemnt with the FLARE dataset, use:
--dataset flare` and `--organ selected
Experimenting with Other Baselines
To experiement with baseline methods, use the following command:
python main_fseft.py --model_id fseft --dataset totalseg --organ organ_name --k 5 --method method_name --seeds 3 --decoder frozen --rank 8
You can choose among the following methods:
- Generalization
--method generalization - Linear probing
--method LP - Full fine-tuning
--method FT - LoRA
--method LoRA - AdaLoRA
--method AdaLoRA - BitFit
--method Bias - Affine-LN
--method Affine
Note that, for LoRA and AdaLoRA you might specify the initial rank using the option -- rank init_rank_value. For organ selection, decoder configuration and dataset setup the same guidelines provided for ARENA are applied here.
Experimenting with Recently Released Pre-trained Foundation Models
The above experiements can also be conducted using several recently released pre-trained foundation models. To use a specific model:
- Download the corresponding weight file from the provided link.
- Save it to the following path:
models/pretrained_weights/[ID].pth - Use the appropriate
--model_idin the command when running the experiment.
The table below lists available models:
| Model Name | Architecture | Model [ID] | Repository | Weights | |
|---|---|---|---|---|---|
| Self-Supervised | 2022 | Swin-UNETR | selfsup | LINK | LINK |
| Dataset-specific (BTCV) | 2022 | Swin-UNETR | btcv | SwinUNETR | LINK |
| CLIP-Driven | 2023 | Swin-UNETR | clipdriven | CLIP-Driven | LINK |
| FSEFT | 2023 | Swin-UNETR | fseft | FSEFT | LINK |
| SuPreM | 2024 | Swin-UNETR | suprem_swinunetr | SuPreM | LINK |
The configuration of each model can be found at models/configs.py.
Note: Please, check other authors' repositories for an updated link to the weights.
Acknowledgement
- We express our sincere gratitude to the authors of FSEFT for their open-source contribution and model, which form the foundation of this repository.
- We also thank the authors of CLIP-Driven and SuPreM for making their pre-trained models publicly available for research purposes.
- The FSEFT framework on which our work is built, is itself based on the MONAI library for medical image segmentation.
Contact
For any inquiries, please contact us at at ghassen.baklouti.1@ens.etsmtl.ca or feel free to create an issue.
Citation
If you find this repository useful, please consider citing this paper:
@article{baklouti2025regularized,
title={Regularized Low-Rank Adaptation for Few-Shot Organ Segmentation},
author={Baklouti, Ghassen and Silva-Rodr{\'\i}guez, Julio and Dolz, Jose and Bahig, Houda and Ayed, Ismail Ben},
journal={arXiv preprint arXiv:2507.15793},
year={2025}
}