SlimFormer++: Token Mixing Where It Matters for Efficient 3D Medical Image Segmentation

Yang Hong; Lei Zhang; Zhihong Liu; Jianqing Mo; Xujiong Ye

cs.CVOct 18, 2026

SlimFormer++: Token Mixing Where It Matters for Efficient 3D Medical Image Segmentation

Authors: Yang Hong, Lei Zhang, Zhihong Liu, Jianqing Mo, Xujiong Ye

Organizations: the University of Exeter, Exeter, U.K. · Guangdong University of Technology

Abstract

To determine where computation is needed in volumetric segmentation networks, we analyze block-level redundancy and stage-wise token-mixer contribution. Centered Kernel Alignment reveals high input-output representational similarity within blocks of representative convolutional, hybrid, and Transformer-style segmentation networks. Similarity-guided pruning reduces their parameter counts by 39-75% while decreasing mean Dice by less than 0.4 percentage points. To examine token mixers within blocks, we introduce Anatomical Relation Distortion, which quantifies how bypassing a mixer distorts feature relations between labeled anatomical positions. Applied to an unpruned candidate architecture, Centered Kernel Alignment indicates high representational redundancy in intermediate encoder mixers, whereas Anatomical Relation Distortion shows that they contribute little and the bottleneck mixer most. These findings motivate SlimFormer++, which uses Selective Mixing blocks with windowed Mamba mixing in the first encoder stage and only residual channel MLPs in deeper stages. At the bottleneck, Augmented Mixing injects similarity-weighted sparse long-range context into the axial taps of a depthwise convolution at cost linear in voxel count. Evaluated against 12 baselines across eight public CT, MRI, and ultrasound datasets (1,809 cases), SlimFormer++ achieves the highest mean Dice (87.84%) and lowest mean HD95 (6.50 mm) with the fewest parameters (5.98 million). These results support stage-aware token mixing as a design strategy for improving the accuracy-efficiency trade-off in volumetric segmentation. The source code is available at https://github.com/WHHHHHY822/SlimFormerPlusPlus

Figures & tables

Fig. 1: Performance–efficiency trade-off of representative 3D segmentation architectures. Dice is the unweighted mean over the eight datasets. Arrows connect each model before pruning (hollow) to its pruned variant (filled), and the percentage on each arrow is the resulting parameter reduction.

Fig. 2: Redundancy and contribution measures for token mixers. ARD compares the intact features $\mathbf{F}$ with the features $\mathbf{F}^{(-b)}$ obtained by bypassing the token mixer of a block, and CKA compares the representations $\mathbf{U}$ and $\mathbf{V}$ before and after token mixing. The right panel sketches the stage-wise pattern measured on the unpruned network (Section IV-C2 ).

Fig. 3: Overall framework of SlimFormer++. (a) Macro architecture: a convolutional Stage 0, three encoder stages (S1–S3), a bottleneck (BN), and a decoder. (b) Selective Mixing (SM) block: a causal Mamba scan within nonoverlapping windows. (c) Token mixers across stages, set by the redundancy–contribution analysis (right): active in S1, pruned in S2 and S3, where only the channel MLPs remain, and strengthened at the bottleneck (d). (d) Augmented Mixing (AM) block: remote context sampled along six axial directions is mixed into the axial taps of the depthwise convolution. (e) Residual convolution block of Stage 0. (f) Aligned upsampling module of the decoder.

Dataset	Region	Target	Modality
AbdomenCT-1K [ 36 ]	Abdomen	Multi-organ	CT
ACDC [ 37 ]	Thorax	Ventricles	MRI
AMOS [ 38 ]	Abdomen	Multi-organ	CT
CHAOS [ 39 ]	Abdomen	Liver	MRI
RESECT [ 40 ]	Head	Brain tumor	US
MMWHS [ 41 ]	Thorax	Whole heart	CT, MRI

TABLE I: Summary of the eight public segmentation datasets.

Fig. 4 : Sensitivity of the unpruned network to stage removal (left) and joint bypass of a stage’s token mixers (right). Bars show the Dice drop ( $\Delta$ Dice, symmetric log scale); lines show whole-stage CKA and ARD. BN denotes the bottleneck.

Fig. 6 : Representation redundancy is measured by stage-averaged CKA, the mean whole-stage input–output CKA over the encoder stages and the bottleneck. Higher CKA indicates greater representational similarity and potential redundancy.

Encoder: alternative mixers activated throughout S1–S3
			AbdomenCT-1K				AMOS				ACDC
Mixer	$\Delta$ P $\downarrow$	$\Delta$ G $\downarrow$	Dice $\uparrow$	Jaccard $\uparrow$	HD95 $\downarrow$	ASD $\downarrow$	Dice $\uparrow$	Jaccard $\uparrow$	HD95 $\downarrow$	ASD $\downarrow$	Dice $\uparrow$	Jaccard $\uparrow$	HD95 $\downarrow$	ASD $\downarrow$
Ours	–	–	95.62	91.83	2.84	0.67	89.76	84.14	11.36	6.38	91.56	85.05	3.70	0.69
Attention	+0.34	+37.9	$-$ 0.10	$-$ 0.17	$-$ 0.01	+0.04	$-$ 0.28	$-$ 0.40	$-$ 0.10	$-$ 0.29	$-$ 0.20	$-$ 0.34	+0.84	+0.09
GDNet	+0.75	+21.9	$-$ 0.30	$-$ 0.49	+0.69	+0.10	$-$ 0.27	$-$ 0.40	$-$ 0.70	$-$ 0.89	$-$ 0.16	+0.08	+0.24	+0.01
DWConv	+0.03	+6.8	$-$ 0.11	$-$ 0.18	+0.52	$-$ 0.15	$-$ 0.25	$-$ 0.36	$-$ 0.88	$-$ 0.87	+0.05	+0.04	+0.30	0.00

TABLE II: Token-mixer replacement in SlimFormer++. Ours gives absolute values (Dice and Jaccard in %, HD95 and ASD in mm); other rows give changes relative to Ours.

\Delta

P: parameter change (M);

\Delta

G: GFLOP change (%, averaged over datasets).

Model	Params	AbdomenCT-1K		ACDC		AMOS		CHAOS		RESECT		MMWHS		VESSEL12		Leg3DUS
Model	(M)	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$	Dice $\uparrow$	HD95 $\downarrow$
Representative models
nnUNet [ 2 ] Nat. Methods’21	30.8	94.73	4.59	90.06	4.71	83.61	21.01	79.66	17.61	68.12	6.98	87.63	6.56	99.51	0.74	81.37	13.10
TransUNet [ 3 ] MedIA’24	151.1	95.42	3.34	90.63	4.74	89.64	11.81	80.92	17.54	66.56	8.13	88.93	6.68	99.55	0.74	83.37	11.41
SlimFormer backbone [ 17 ]	69.8	94.39	76.52	88.12	4.46	84.37	17.19	79.38	18.06	64.74	5.98	88.04	7.51	98.33	10.40	76.44	24.51
nnWNet3D [ 24 ] CVPR’25	53.2	95.11	3.95	91.47	3.97	87.50	13.21	83.70	12.78	67.80	6.64	90.00	6.36	99.60	0.71	82.50	19.83

TABLE III: Quantitative comparison with state-of-the-art 3D medical image segmentation methods on eight public datasets. Dice is in % (

\uparrow

) and HD95 in mm (

\downarrow

). Bold : best; underlined : second-best distinct value for each dataset and metric.

Fig. 8: Efficiency profiles of nine models: parameters (M), GFLOPs, latency (ms), peak memory (GB), and energy (J). Values are min–max normalized in logarithmic space and inverted , so larger radii indicate lower costs within the displayed comparison set.

Fig. 9: Qualitative comparison on three representative datasets covering CT, MRI, and ultrasound (rows, ordered as in Table I ).

3D Medical Image Segmentation

3D Medical Image Segmentation identifies organs, lesions, and fine structures in volumetric scans for measurement and clinical planning.

See all 213 papers

BATS: Resource-Efficient Volumetric Segmentation with Boundary-Aware Mixed-Resolution Tokens
Jul 29, 2026David Hagerman, Roman Naeem, Fredrik KahlImage Segmentation 3D Medical Image Segmentation
Chalmers University of Technology, 412 96 Gothenburg, Sweden
SegMaFormer: A Hybrid State-Space and Transformer Model for Efficient Segmentation
Mar 23, 2026Duy D. Nguyen, Phat T. Tran-TruongHybrid Transformer-State Space Model Architectures 3D Medical Image Segmentation
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
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

Jul 29, 2026J/K move · Enter open · S save

BATS: Resource-Efficient Volumetric Segmentation with Boundary-Aware Mixed-Resolution Tokens

David Hagerman, Roman Naeem, Fredrik Kahl

Chalmers University of Technology, 412 96 Gothenburg, Sweden

Image Segmentation 3D Medical Image Segmentation Boundary-Aware Segmentation

Open paper AskPDF arXiv

Many high-performing volumetric segmentation models maintain dense multi-scale feature maps, leading to high activation memory and inference cost. We present BATS (Boundary-Aware Token Selection), a 3D medical image segmentation architecture that concentrates fine-resolution processing near predicted class boundaries. A dense boundary predictor identifies where additional resolution is needed, while a fine-first context cascade constructs an input-dependent mixed-resolution hierarchy. Homogeneous regions are represented coarsely, with finer tokens retained around boundaries, thin structures, and small targets. The sparse hierarchy is refined and rasterised into a dense segmentation. BATS predicts boundary relevance independently at every resolution level, preventing an erroneous coarse-scale decision from suppressing fine-scale evidence. Parent cluster attention further injects hierarchical ancestor tokens into local attention neighbourhoods, providing cross-scale context without dense multi-scale feature maps or cross-scale neighbour search. We evaluate BATS on five public CT and MRI datasets using the standardised nnU-Net Revisited protocol. BATS achieves the highest LiTS Dice among the compared methods and averages within 0.37 Dice points of the strongest dense baseline, MedNeXt-L, across the five datasets. Relative to MedNeXt-L, it reduces peak allocated GPU memory by more than 53% on KiTS, LiTS, and BraTS. Inference is up to 30% faster on KiTS and LiTS, which retain fewer tokens, but slower on the more token-dense BraTS. Mixed-resolution processing therefore provides consistent memory savings, while runtime and accuracy gains depend on dataset boundary density.

Abstract

3D Medical Image Segmentation

Explore similar work

BATS: Resource-Efficient Volumetric Segmentation with Boundary-Aware Mixed-Resolution Tokens

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

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