Harnessing Lightweight Transformer with Contextual Synergic Enhancement for Efficient 3D Medical Image Segmentation

Xinyu Liu; Zhen Chen; Wuyang Li; Chenxin Li; Yixuan Yuan

cs.CVMar 24, 2026

Harnessing Lightweight Transformer with Contextual Synergic Enhancement for Efficient 3D Medical Image Segmentation

Authors: Xinyu Liu, Zhen Chen, Wuyang Li, Chenxin Li, Yixuan Yuan

Organizations: Department of Electronic Engineering, The Chinese University of Hong Kong, Hong Kong SAR · Centre for Artificial Intelligence and Robotics (CAIR), Hong Kong Institute of Science & Innovation, Chinese Academy of Sciences, Hong Kong SAR

Abstract

Transformers have shown remarkable performance in 3D medical image segmentation, but their high computational requirements and need for large amounts of labeled data limit their applicability. To address these challenges, we consider two crucial aspects: model efficiency and data efficiency. Specifically, we propose Light-UNETR, a lightweight transformer designed to achieve model efficiency. Light-UNETR features a Lightweight Dimension Reductive Attention (LIDR) module, which reduces spatial and channel dimensions while capturing both global and local features via multi-branch attention. Additionally, we introduce a Compact Gated Linear Unit (CGLU) to selectively control channel interaction with minimal parameters. Furthermore, we introduce a Contextual Synergic Enhancement (CSE) learning strategy, which aims to boost the data efficiency of Transformers. It first leverages the extrinsic contextual information to support the learning of unlabeled data with Attention-Guided Replacement, then applies Spatial Masking Consistency that utilizes intrinsic contextual information to enhance the spatial context reasoning for unlabeled data. Extensive experiments on various benchmarks demonstrate the superiority of our approach in both performance and efficiency. For example, with only 10% labeled data on the Left Atrial Segmentation dataset, our method surpasses BCP by 1.43% Jaccard while drastically reducing the FLOPs by 90.8% and parameters by 85.8%. Code is released at https://github.com/CUHK-AIM-Group/Light-UNETR.

Figures & tables

Fig. 1: Performance comparison between the proposed CSE-Light-UNETR, CSE-VNet with previous SSMIS methods on the LA dataset [ 1 ] . The circle size represents the number of trainable parameters in different methods.

Fig. 2: Comparison between different methods for SSMIS. (a) Traditional methods [ 12 , 17 , 15 ] combine conv-based medical image segmentation models [ 20 ] with semi-supervised learning techniques such as mean-teacher [ 41 ] , which are not only computational expensive but also have poor performance. (b) Recent methods [ 34 ] jointly train Transformers [ 22 ] with conv-based models, which are limited to 2D data and rely on natural-domain pretrained models. (c) Our proposed model applies a single shared lightweight Transformer model for SSMIS on 3D medical volumes, which achieves better performance with significantly less computational cost.

Fig. 3: Illustration of the proposed (a) Overall Light-UNETR structure, (b) Light-UNETR Block (c) LIDR module, (d) CGLU module, (e) overlap patch embedding, and (f) Light downsample.

Fig. 4: Illustration of the CSE semi-supervised learning framework, which contains (a) Attention-Guided Replacement (AGR) and (b) Spatial Masking Consistency (SMC).

Method	LA dataset					Pancreas-CT dataset
Method	L/U	Dice (%) $\uparrow$	Jac (%) $\uparrow$	95HD (mm) $\downarrow$	ASD (mm) $\downarrow$	L/U	Dice (%) $\uparrow$	Jac (%) $\uparrow$	95HD (mm) $\downarrow$	ASD (mm) $\downarrow$
Light-UNETR (Ours)	4/0	75.55	61.99	17.80	6.16	6/0	59.46	44.15	22.64	7.92
UA-MT (MICCAI’19) [ 12 ]	4/76	78.34	65.11	15.42	4.69	6/56	52.41	36.80	36.34	11.78
SS-Net (MICCAI’22) [ 13 ]	4/76	79.69	67.02	12.44	3.86	6/56	51.31	36.10	22.04	7.50
CAML (MICCAI’23) [ 16 ]	4/76	79.74	66.94	14.58	4.46	6/56	57.34	40.94	28.41	8.86
BCP (CVPR’23) [ 15 ]	4/76	84.76	73.88	10.22	2.89	6/56	68.93	53.86	14.01	4.37

TABLE I: Semi-supervised segmentation performance on LA and Pancreas-CT datasets. Bold denotes the best result.

\uparrow

and

\downarrow

denote that higher or lower values are better, respectively. L/U refers to the number of labeled/unlabeled scans. All models use the proposed Light-UNETR as base segmentation model.

Methods	BraTS 2019 dataset
Methods	L/U	Dice $\uparrow$	Jac $\uparrow$	95HD $\downarrow$	ASD $\downarrow$
Light-UNETR (Ours)	25/0	75.91	64.80	12.93	3.86
UA-MT (MICCAI’19) [ 12 ]	25/225	74.13	62.45	16.23	4.42
SS-Net (MICCAI’22) [ 13 ]	25/225	78.79	67.96	13.65	3.96
CAML (MICCAI’23) [ 16 ]	25/225	78.35	67.02	13.34	3.85
BCP (CVPR’23) [ 15 ]	25/225	77.06	66.02	18.10	5.33

TABLE II: Semi-supervised segmentation performance on the BraTS 2019 dataset. Bold denotes the best result.

\uparrow

and

\downarrow

denote that higher or lower results are preferred. L/U refers to the number of labeled/unlabeled scans. All models use the proposed Light-UNETR as base segmentation model.

Fig. 5: Qualitative segmentation results comparison of different methods on the LA dataset with 5% labeled data. From left-to-right are UA-MT [ 12 ] , SS-Net [ 13 ] , CAML [ 16 ] , BCP [ 15 ] , MLRP [ 74 ] , our CSE-Light-UNETR, and ground-truth.

Fig. 6: Qualitative segmentation results comparison of different methods on the Pancreas-CT dataset with 10% labeled data. From left-to-right are UA-MT [ 12 ] , SS-Net [ 13 ] , CAML [ 16 ] , BCP [ 15 ] , MLRP [ 74 ] , our CSE-Light-UNETR, and ground-truth.

Method	LA dataset				Pancreas-CT dataset				BraTS 2019 dataset				FLOPs (G)	Params (M)
Method	Dice $\uparrow$	Jac $\uparrow$	95HD $\downarrow$	ASD $\downarrow$	Dice $\uparrow$	Jac $\uparrow$	95HD $\downarrow$	ASD $\downarrow$	Dice $\uparrow$	Jac $\uparrow$	95HD $\downarrow$	ASD $\downarrow$
Foundation Models
MedSAM2 [ 75 ]	80.44	67.67	21.59	5.28	69.01	53.22	22.23	5.31	83.08	74.02	17.21	6.36			2,259.2	38.96
Efficient-MedSAM2 [ 75 ]	78.61	65.39	21.82	6.02	68.92	53.06	18.83	5.55	82.29	71.06	11.13	6.01	1,922.0	34.10
Med-FastSAM [ 76 ]	83.16	69.53	19.02	5.10	70.30	59.71	19.72	4.88	82.06	71.58	10.84	5.03	1,784.9	14.48
Segmentation Networks

TABLE III: Fully-supervised segmentation results on LA, Pancreas and BraTS 2019 datasets. The best results are highlighted in bold , and the second-best are underlined . FLOPs are computed with spatial size (112, 112, 80).

Method	MSD Task01 dataset				AbdomenCT-1K dataset					HNC Tumor dataset			FLOPs (G)	Params (M)
Method	WT	ET	TC	Avg	Liver	Kidney	Spleen	Pancreas	Avg	GTVp	GTVn	Avg	FLOPs (G)	Params (M)
nn-UNet [ 57 ]	85.83	62.71	83.42	77.32	97.21	96.92	95.93	93.42	95.87	94.79	95.76	95.27	200.81	13.82
SegResNet [ 77 ]	88.68	62.64	83.28	78.20	97.06	96.35	95.23	91.48	95.03	95.19	95.70	95.44	202.66	10.64
TransBTS [ 27 ]	77.89	57.44	73.51	69.61	97.08	96.47	95.38	92.86	95.45	94.02	94.31	94.16	142.07	30.61
Swin UNETR [ 24 ]	90.23	63.02	84.22	79.16	97.03	96.08	94.98	93.51	95.40	96.45	96.51	96.48	331.95	61.99
UNETR++ [ 30 ]	90.07	64.87	84.52	79.82	95.64	95.56	93.99	91.68	94.22	95.39	95.74	95.56	57.51	20.33

TABLE IV: Fully-supervised segmentation dice scores on MSD Task01, AbdomenCT-1K, and HNC Tumor datasets. The best results are highlighted in bold , and the second-best are underlined . FLOPs are computed with spatial size (96, 96, 96).

Fig. 7: Evolution roadmap of the proposed Light-UNETR. The experiments are conducted on the LA dataset with 5% labeled data.

Method				Metrics
Sup Loss	Rand R.	AGR	SMC	Dice $\uparrow$	Jac $\uparrow$	95HD $\downarrow$	ASD $\downarrow$
✓				75.55	61.99	17.80	6.16
✓	✓			86.78	77.35	8.37	2.84
✓		✓		87.79	78.97	7.51	2.33
✓			✓	83.59	72.60	11.45	3.79
✓		✓	✓	89.53	81.15	6.14	1.88

TABLE V: Ablation on the components in the proposed CSE. The experiments are conducted on the LA dataset with 5% labeled data. Rand R.: Random Replacement. AGR: Attention-Guided Replacement. SMC: Spatial Masking Consistency.

$\lambda_{ext}$	$\lambda_{int}$	Dice↑	Jac↑	95HD↓	ASD↓	$\lambda_{ext}$	$\lambda_{int}$	Dice↑	Jac↑	95HD↓	ASD↓
1.0	1.0	87.68	78.26	7.87	2.98	3.0	1.0	88.86	80.12	7.25	2.62
2.0	2.0	88.45	79.45	6.86	2.39	4.0	1.0	89.53	81.15	6.14	1.88
3.0	3.0	89.41	80.94	6.37	2.32	1.0	3.0	88.49	79.52	6.61	2.13
4.0	4.0	89.41	80.93	6.28	2.31	1.0	4.0	88.31	79.22	9.60	3.02

TABLE VI: Comparison of

\lambda_{ext}

and

\lambda_{int}

settings (LA dataset, 5% labeled data). Left: Equal

\lambda_{ext}=\lambda_{int}

; Right: Different

\lambda_{ext}\neq\lambda_{int}

Fig. 8: Ablation on the hyperparameters (a) The impact of changing the Attention-Guided Replacement patch size ratio $\alpha$ . (b) The impact of changing the mask region side length $l$ . (c) The impact of changing the mask ratio $v$ .

#Branch	Kernel	FLOPs	Params	LA (8/72)		Pancreas (12/50)
#Branch	Kernel	(G)	(M)	Dice(%)	Jac(%)	Dice(%)	Jac(%)
1	3	4.29	1.34	89.39	81.82	78.04	65.52
1	5	4.30	1.43	89.33	81.83	77.69	65.30
1	7	4.32	1.64	90.24	82.51	77.32	65.02
2	3,5	4.29	1.34	90.50	82.74	78.50	65.27
2	5,7	4.31	1.56	89.54	81.92	78.46	65.81

TABLE VII: Ablation on the number of high-frequency branches and kernel size configurations.

Method	Dice $\uparrow$	Jac $\uparrow$	95HD $\downarrow$	ASD $\downarrow$	Mem (GB) $\downarrow$
Ours	89.53	81.15	6.14	1.88	2.52
Ours+EMA	86.85	76.96	8.73	2.32	3.10
EMA Only	84.53	73.70	10.21	2.82	3.10

TABLE VIII: Performance comparison between ours using shared-weight, using shared-weight and EMA simultaneously, and EMA only.

Methods	FLOPs	Params	Metrics
Methods	(G)	(M)	Dice $\uparrow$	Jac $\uparrow$	95HD $\downarrow$	ASD $\downarrow$
V-Net † [ 20 ]	46.85	9.44	82.74	71.72	13.35	3.26
UA-MT † [ 12 ]	46.85	9.44	87.79	78.39	8.68	2.12
SS-Net † [ 13 ]	50.71	9.45	88.55	79.62	7.49	1.90
CAML † [ 16 ]	119.72	19.73	89.62	81.28	8.76	2.02
BCP † [ 15 ]	46.85	9.44	89.62	81.31	6.81	1.76

TABLE IX: Segmentation performance comparison with CNN-based methods on the LA dataset with 10% labeled data. FLOPs are computed with input size (112, 112, 80). † refers to models using V-Net [ 20 ] as base segmentation model.

Methods	L/U	Liver	Kidney	Spleen	Pancreas	Avg
Light-UNETR (Ours)	80/0	83.36	84.02	83.71	68.72	79.95
UA-MT (MICCAI’19) [ 12 ]	80/720	87.19	85.76	84.53	72.87	82.59
SS-Net (MICCAI’22) [ 13 ]	80/720	88.62	86.44	85.92	74.15	83.78
CAML (MICCAI’23) [ 16 ]	80/720	90.91	88.13	87.60	78.63	86.32
BCP (CVPR’23) [ 15 ]	80/720	93.22	91.92	89.68	85.06	89.97
MLRP (MedIA’24) [ 74 ]	80/720	91.63	89.72	89.09	85.31	88.93

TABLE X: Semi-supervised segmentation performance on the AbdomenCT-1K dataset with 10% labeled data. Dice scores (%) for each organ and average performances are reported.

Model	FLOPs	Params	LA		Pancreas
Model	(G)	(M)	4/76	8/72	6/56	12/50
Vanilla [ 36 ]	4.83	2.68	89.05	90.04	73.75	78.57
Linformer [ 82 ]	4.73	2.67	89.01	89.79	73.26	78.17
Longformer [ 83 ]	4.82	2.68	88.94	90.35	73.67	78.24
Nyström [ 84 ]	4.83	2.69	89.52	90.72	73.36	78.15
LIDR	4.29	1.34	89.53	90.50	73.77	78.50

TABLE XI: Comparison of different attention mechanisms.

Dataset	L/U	Dice $\uparrow$	Jac $\uparrow$	95HD $\downarrow$	ASD $\downarrow$	Ratio
LA	4/76	89.53	81.15	6.14	1.88	5%
	8/72	90.50	82.74	5.82	1.83	10%
	16/64	91.06	83.40	5.16	1.70	20%
	32/48	91.64	84.61	4.81	1.51	40%
	Fully Supervised	91.58	84.51	4.95	1.43	100%
Pancreas-CT	6/56	73.77	59.44	11.77	3.76	10%

TABLE XII: Performance metrics on 3 different datasets with different labeled and unlabeled sample ratios.

Fig. 9: Fourier spectrum of the low-frequency branch and the high frequency branches. Top row: The Fourier spectrum of the feature in the encoder. Bottom row: The Fourier spectrum the feature in the decoder.

Model	Dice (%)	FPS $\uparrow$	Mem (GB) $\downarrow$	Params (M)
MedNeXt [ 58 ]	92.05	28.43 (1.00x)	12.50 (1.00x)	5.57
w/ FA	92.05	28.43 (1.00x)	12.50 (1.00x)	5.57
w/ AMP	92.03	35.25 (1.24x)	10.25 (0.82x)	5.57
w/ TensorRT	92.05	54.59 (1.92x)	9.38 (0.75x)	5.57
Light-UNETR	91.58	55.18 (1.00x)	2.52 (1.00x)	1.34
w/ FA	91.58	67.07 (1.22x)	2.48 (0.98x)	1.34

TABLE XIII: Compatibility of MedNeXt and Light-UNETR with hardware optimizations on fully-supervised LA dataset.

Method	Dice $\uparrow$	Jac $\uparrow$	95HD $\downarrow$	ASD $\downarrow$
TransUNet-ViT-B-16 (MIM)	61.03	47.62	73.60	21.82
TransUNet-ViT-B-16 (SMC)	62.22	48.94	72.19	20.17

TABLE XIV: Segmentation performance comparison with different self-supervised pretraining techniques. Experimental results on the Synapse multi-organ CT dataset, all models are pretrained and finetuned for 150 epochs.

Explore similar work

LightMedSeg: Lightweight 3D Medical Image Segmentation with Learned Spatial Anchors
Mar 7, 2026Kavyansh Tyagi, Vishwas Rathi, Puneet GoyalMedical Image Segmentation 3D Medical Image Segmentation
National Institute of Technology Kurukshetra, India · Indian Institute of Technology Ropar, India
RefineFormer3D: Efficient 3D Medical Image Segmentation via Adaptive Multi-Scale Transformer with Cross Attention Fusion
Feb 18, 2026Kavyansh Tyagi, Vishwas Rathi, Puneet Goyal3D Medical Image Segmentation 3D ViTs
National Institute of Technology Kurukshetra Department of Electronics and Communication Engineering · National Institute of Technology Kurukshetra Department of Computer Engineering · Indian Institute of Technology Ropar Department of Computer Science and Engineering
SegMaFormer: A Hybrid State-Space and Transformer Model for Efficient Segmentation
Mar 23, 2026Duy D. Nguyen, Phat T. Tran-Truong3D Medical Image Segmentation 3D ViTs
Department of Software Engineering, Faculty of Computer Science and Engineering, Ho Chi Minh City University of Technology (HCMUT), 268 Ly Thuong Kiet Street, Dien Hong Ward, Ho Chi Minh City, Vietnam · Vietnam National University Ho Chi Minh City, Linh Trung Ward, Ho Chi Minh City, Vietnam

Mar 7, 2026J/K move · Enter open · S save

LightMedSeg: Lightweight 3D Medical Image Segmentation with Learned Spatial Anchors

Kavyansh Tyagi, Vishwas Rathi, Puneet Goyal

National Institute of Technology Kurukshetra, India · Indian Institute of Technology Ropar, India

Medical Image Segmentation 3D Medical Image Segmentation U-Net

Open paper AskPDF arXiv

Accurate and efficient 3D medical image segmentation is essential for clinical AI, where models must remain reliable under stringent memory, latency, and data availability constraints. Transformer-based methods achieve strong accuracy but suffer from excessive parameters, high FLOPs, and limited generalization. We propose LightMedSeg, a modular UNet-style segmentation architecture that integrates anatomical priors with adaptive context modeling. Anchor-conditioned FiLM modulation enables anatomy-aware feature calibration, while a local structural prior module and texture-aware routing dynamically allocate representational capacity to boundary-rich regions. Computational redundancy is minimized through ghost and depthwise convolutions, and multi-scale features are adaptively fused via a learned skip router with anchor-relative spatial position bias. Despite requiring only 0.48M parameters and 14.64~GFLOPs, LightMedSeg achieves segmentation accuracy within a few Dice points of heavy transformer baselines. Therefore, LightMedSeg is a deployable and data-efficient solution for 3D medical image segmentation. Code will be released publicly upon acceptance.

Abstract

Explore similar work

LightMedSeg: Lightweight 3D Medical Image Segmentation with Learned Spatial Anchors

RefineFormer3D: Efficient 3D Medical Image Segmentation via Adaptive Multi-Scale Transformer with Cross Attention Fusion

SegMaFormer: A Hybrid State-Space and Transformer Model for Efficient Segmentation