Johnson-Lindenstrauss Lemma Guided Network for Efficient 3D Medical Segmentation
Authors: Jinpeng Lu
Linghan Cai
Yinda Chen
Guo Tang
Songhan Jiang
Haoyuan Shi
Zhiwei Xiong
Venue: ICLR 2026 Poster (2026)
Keywords: Efficient Medical segmentation
multimodal learning
Knowledge Transfer
Primary area: applications to computer vision, audio, language, and other modalities
TL;DR: This paper proposes VeloxSeg, a theory-based lightweight framework that systematically alleviates the ''efficiency / robustness conflict'' in 3D medical segmentation.

Abstract: Lightweight 3D medical image segmentation remains constrained by a fundamental "efficiency / robustness conflict", particularly when processing complex anatomical structures and heterogeneous modalities. In this paper, we study how to redesign the framework based on the characteristics of high-dimensional 3D images, and explore data synergy to overcome the fragile representation of lightweight methods. Our approach, VeloxSeg, begins with a deployable and extensible dual-stream CNN-Transformer architecture composed of Paired Window Attention (PWA) and Johnson-Lindenstrauss lemma-guided convolution (JLC). For each 3D image, we invoke a "glance-and-focus" principle, where PWA rapidly retrieves multi-scale information, and JLC ensures robust local feature extraction with minimal parameters, significantly enhancing the model's ability to operate with low computational budget. Followed by an extension of the dual-stream architecture that incorporates modal interaction into the multi-scale image-retrieval process, VeloxSeg efficiently models heterogeneous modalities. Finally, Spatially Decoupled Knowledge Transfer (SDKT) via Gram matrices injects the texture prior extracted by a self-supervised network into the segmentation network, yielding stronger representations than baselines at no extra inference cost. Experimental results on multimodal benchmarks show that VeloxSeg achieves a 26\% Dice improvement, alongside increasing GPU throughput by 11x, CPU by 48x, and reducing training peak GPU memory usage by 1/20, inference by 1/24. Code is available at https://github.com/JinPLu/VeloxSeg.

Published as a conference paper at ICLR 2026
JOHNSON-LINDENSTRUUS LEMMA GUIDED NET- WORK FOR EFFICIENT 3D MEDICAL SEGMENTATION
Jinpeng Lu'; Linghan Cai: Yinda Chen'; Guo Tang", Songhan Jiang',
Haoyuan Shi!, Zhiwei Xiong
University of Science and Technology of China
2 Dresden University of Technology
3 Harbin Institute of Technology, Shenzhen
j jinplu@maiin 11 ustc. edu. en, zwxiongeustc.edu.
ABSTRACT
Lightweight 3D medical image segmentation remains constrained by funda- mental "efficiency robustness conflict", particularly when processing complex anatomical structures and heterogeneous modalities. In this paper, study how to redesign the framework based on the characteristics of high-dimensional 3D images, and explore data synergy to overcome the fragile representation of lightweight methods. Our approach, VeloxSeg, begins with deployable and ex- tensible duall-streat CNN- Tranffrasfrrm architecture composed of Paired Window Attention (PWA) and
lemma coama--uided convolution (JLC). For each 3D image, we invoke a "glance-and principle, where PWA rapidly retrieves multi-scale information, and JLC ensures robust local feature extraction with minimal parameters, significantly enhancing the model's ability to operate with low computational budget. Followed by an extension of the dual-stream ar- chitecture that incorporates modal interaction into the multi-scale image-retrieval process, VeloxSeg efficiently models heterogeneous modalities. Finally, Spatially Decoupled Knowledge Transfer (SDKT) via Gram matrices injects the texture prior extracted by self-supervised network into the segmentation network, yield- ing stronger representations than baselines at no extra inference cost. Experimen- tal results on multimodal benchmarks show that VeloxSeg achieves a 26% Dice improvement, alongside increasing GPU throughput by 11x, CPU by 48x, and reducing training peak GPU memory usage by 1/20, inference by 1/24. Code is available at https:: / giggtuhuubb.... www./oo/////.... La Ve lxSeg
INTRODUCTION
3D medical image segmentation serves as a cornerstone of contemporary clinical workflows (Wu
al. 2025 Peiris al. 2023 Wang et al. 2026), driving rapid advances semantic segmenta-
lion models (Liu et al., 2024a Shaker et al. 2024 He et al., 2025) Yu et al.) 2025a Wald et al.,
2025 Cai et al. 2023 2024). However, translating these advances into clinical practice faces signif- icant obstacles, including limited hardware resources, stringent latency requirements, and the need to achieve multi-organ generalization while handling heterogeneous multimodal data in deployment environments. These challenges have spurred the development of lightweight 3D medical segmen- tation methods, leading to lightweight approaches with fewer than 5 million parameters (Perera et al. 2024 Pang et al. 2024 Yu et al. 2025b Li et al.) 2025 Ye et al. 2025). Yet, the pursuit of smaller parameter counts and lower computational costs has revealed a fundamental and increas- ingly prominent trade-off: these lightweight models struggle to maintain both efficiency and robust performance when handling heterogeneous data and complex lesions, which we term "efficiency / robustness conflict". We address this problem from two key perspectives:
Insufficient consideration of the high-dimensional complexity of 3D data. Recent sequence
models, such as Mamba (Gu & Dao 2023 Xing al. 2025) and RWKV (Peng al. 2023 Ye
"Equal contribution.
Corresponding author.

---

Published as a conference paper at ICLR 2026
68
65
VSmiTrun
YSmilann
Nesteformmer FRomere
UNETL DeuseFormer
VelonSeg
VelonSeg
Similarity Coefficient (%4
A2r3eg
UNd
U-KAN
Similarity Coufficiunt (%)
A2F
UNE
Nestedbommrr
TIKAN
U-RWKV
88
MedNeXt
MedNee:
UNCTR
UNETR
VNat UNETR C
Slim UIETR
30
50
SuperrLightNlt
Supertieeltie
HCMA-UNet
40
S
40
SeyFommaz-DD
GPU Past Zone
CPU Available Zene
UNETR++
10U
200
500
0.05
0.1
0.2
0.5
18 15
(a) GPU Throughput (Patches/s)
(b) CPU Throughput (Patchesis)
Figure 1: Comparison of our proposed VeloxSeg with recent methods on the AutoPET-II test set. Basic models, multimodal models, lightweight models, and our model are marked with circles, squares, triangles, and stars, respectively. GPU and CPU Throughput are measured on an NVIDIA RTX3090 GPU and single-core Intel(R) Xeon(R) Gold 5320 CPU, respectively.
et al. 2025), have achieved remarkable progress in segmentation, owing to their linear complexity and long-range modeling capabilities. However, due the lack of more efficient scanning strategies suitable for 3D data, these methods have not yet supplanted CNN-Tran CNN-Transformer architectures in the domain of efficient medical segmentation. Our model is built on a dual-stream TNN--Tranformer architecture, synergizing the complementary strengths of both components: the inductive bias (Le-
Cun et al., 1989, Mansour et al., 2019) and training stability (He et al. 2016)] [loff & Szegedy) | Ulyanov et al. (0166) of convolutions with the global modeling power (Vaswani et al.) 2017) and extensibility (Lu et al. 2019; Chen et al. 2020: Dosovitskiy et (2020) of Transformers. Prun- ing is the most common approach to model lightweighting (Molchanov et al., 2019 Fang et al. 2023). but its final configuration relies on dataset-specifie importance metrics and hand-tuned spar- sity schedules, which limit generalizability and lead to expensive retraining. Therefore, developing lightweight and efficient components is essential. Constructing relationships among tokens is the core of feature modeling. In principle, self-attention (Dosovitskiy al., (021) can represent arbi- trary dependencies, but in practice it is constrained by computation and memory. Window-base attention (Hatamizadeh al., 2021; Du et al.) (2025) performs fine- une-grained relation modeling within local windows, but it relies on cascaded operations to capture cross-window interactions, leading to substantial redundancy. Axial attention (Liu et al. 2024a) and downsampled attention (Pang al,, 2024 [Perera et al. [2024) Kuang et al.) 20255) accelerate the construction of relationships between token and distant tokens by constraining attention paths or operating at lower resolutions, but they tend to weaken the representation of critical local dependencies. We propose paired window atten- tion (PWA), which builds parallel multi-scale feature streams and coordinates short- and long-range attention to capture global token relations while maintaining sufficient focus on local information, at a computational cost comparable to axial or downsampled attention. Convolution with its inductive bias remains indispensable for detailed local modeling. However, common depthwise-separable de- signs (Chollet) 2017] Ma et al.) 2018 Roy et al. 2023] Muhammad et al. 2025) suffer from a key limitation: aggressive channel decoupling disrupts the original geometric adjacency among tokens, making them harder to distinguish and fragmenting the information. This issue is particularly se- vere for complex anatomical structures and heterogeneous modalities. To address this, we introduce a Liiiisoon-Lndnnsttraut
(JL) lemma lemma-guided lightweight convolution (Lindenstrauss & Johnson, 1984), which enforces a minimum number of channels per group in each convolution layer to pre- serve geometric adjacency among tokens. This design keeps the model lightweight while ensuring that fine-grained details can be robustly captured.
Insufficient exploration of data synergy, including multimodal cooperation and data priors. Exploiting multimodal complementary information is crucial for robust model representation (Mu) et al. 2020 Zheng et al. 2025 Zou et al. 2025). However, it is often ignored by lightweight meth- ods due to the potential increase in computational cost, even when training on multimodal datasets. As discussed in Appendix B bridging multimodal information across multiple scales is vital for extracting complementary information from heterogencous modalities. Therefore, we extend our dual-stream architecture, using PWA to facilitate efficient modal interaction at the additional cost
2

---

Published as a conference paper at ICLR 2026
Modal-Fision Convolation Layer
Segmentation Deceder
JLC
PET
PaRA Elubel
Groeg Grote Gatter
Group Scatn
Neen
Add & Newr
Sep. Dennde
Sep. Doods
Lawer
Piad Stufte
PPN
Prof
Eyey
SOKT
PWA
PET
Plach Erebee
Liner Drgg
PW Minue
Add A Nomm
Nomt:
DCC Desiter
DecodGe Layer
Plad Strffe
ReSr
X
Modal-Couperrtive Transformer Layer
Self-Supervied
Teacher
Fenure Flow
Skip Connection
Supervision Signal
Add Operation
Concatemation Operation
Overview of VeloxSeg. VeloxSeg employs an encoder-decoder architecture with Paired Window Attention (PWA) and
Figure 2:
lemma-gune -guieed convolution (JLC) on the left, using I x I convolution as modal mixer. GC: group convolution; GA: multimodal grouped attention.
of only 0.27 MParams and 0.09 GFLOPs. Besides, exploring prior knowledge from existing data to enhance model's detailed representation holds practical significance for efficient segmentation methods. These methods often achieve higher efficiency by performing segmentation in a com- pressed space, which comes at the cost of exploring small lesions and complex boundaries (Perera ct al. 2024 He et al. [2233 Rahman et al. 2024] Pang et al. (022)). Although establishing cross-task knowledge transfer from reconstruction to segmentation appears to be a solution (Sun et al., [0200] Rui el al. 2025 Zhang et al. 2020 Wang et al. 20255), the significant differences in their regions of interest (ROIs) often lead negative knowledge transfer et 2023). To this end, our proposed Spatially Decoupled Knowledge Transfer (SDKT) is a simple yet effective solution, motivated by the observation that a common upsampling operation in reconstruction and super-resolution tasks, "Cxnv+PixelShuffle (Su et al.] 2025)", essentially unfolds the channel relationships at each voxel position into the spatial details of the surrounding image patch. This suggests that the guidance pro- vided by texture teacher to segmentation task should be based the channel relationships within its features. The Gram matrix, commonly used to represent style in the field of image style transfer, characterizes feature channel relationships in ppatially-invarian manner (Gatys al.. 2016 2015).
Based on it, we establish a positive knowledge transfer path from a self-supervised texture teacher the segmentation network with no inference overhead (Zhu et al. 202 Akiva al. 2022).
Inspired by the above insights, we propose VeloxSeg that systematically alleviates the "efficiency / robustness conflict" during model lightweighting. Extensive experiments thoroughly explored the rationale for the design choices and demonstrated the model's excellent clinical applicability and generalization capabilities. Figure [ shows a comparison of VeloxSeg's performance with other methods on the AutoPET-II (Gatidis S. 2022) dataset, demonstrating strong competitiveness. In summary, we develop:
A Paired Window Attention to ensemble multi-scale attention groups capturing local- global information simultaneously, improving localization capabilities with less cost, and achieving low-cost but effective modal interaction at multiple scales.
A Johnson-Lindenstraue lemma-guidee convolution that theoretically determines min- imum group size to preserve spatial adjacency, ensuring robust local feature extraction without costly and data-specific pruming.
A Spatially Decoupled Knowledge Transfer strategy that uses Gram matrices to distill rich textural details from a self-supervised teacher during training, enhancing model fidelity with zero inference overhead.
3

---

Published as a conference paper at ICLR 2026
2
METHODOLOGY
2.1
OVERVIEW OF VELOXSEG
As shown in Figure 2 VeloxSeg em- ploys [wo 4-stage encoders, a modal- fusion convolution encoder and a modal- cooperative Transformer encoder, along with a segmentation decoder and a detail texture teacher. The Paired Window At- tention (PWA), key component of the transformer encoder, is designed to cap- ture multi-scale and cross- mossl context with low enough cost. The Johnson- Lindenstrauss comma-guided Convolution (JLC), a key component of the convolu- tion encoder, consists of 3 parallel JLCs at different scales to fuse modal informa- tion and model local features. Separat- ing these two avoids parameter explo- sion as the number of modalities increases, while maximizing the advantages and par- allelism of both. training, the Spatially Decoupled Knowledge Transfer (SDKT) strategy is used to enhance texture representation, which also of great significance for super- resolution and segmentation tasks.
Global
Qin Kij Vm
Local
Small Window
X m.
Paired Window Gathering
x (1-9)Mnn
Big Window
Multimodal Grouped A ttention
X
Feature Flow
Multiscale
Paired Window Scattering
Modal Interaction
Global
Local
Paired Window Mixer
Figure 3: Overview of Paired Window Attention (PWA).
2.2 PAIRED WINDOW ATTENTION
To achieve sufficiently strong clue-capturing capabilities with minimal computational cost, PWA ensembles parallel feature streams to capture key multimodal information at multiple scales. No- tably, our approach differs significantly from conventional parallel multi-attention approaches (Liu al. 2024a Shaker et al., 2024), aiming to create faster, lower-cost, more effective, and more elegant feature stream. Given M modal features from k-th stage E m = 1,... MM, they are first projected into Q K& V As shown in Figure 3 (i) partition the features into set of big windows, collecting a salient token for each small window; (ii) synchronously expand window pairs to obtain multimodal sequences X 6 X (Q, K, V) of different scales but equal length, where i is the number of the paired window; (iii) gather all sequences and compute attention across all scales and modalities at once; and (iv) use lightweight mixer to simply and efficiently blend features from all scales. The attention A is obtained by the following formula:
A = PWA (E,
E
(I)
For more information about PWA, please see the Appendix C including PyToorch code, detailed formula flow, and complexity analysis. We also provide a detailed analysis of the necessity of multi- scale modeling of medical modalities. Notably, PWA requires only log(size), size E {H, W, D} paired windows to capture global context, while the minimum window ensures the preservation of local details. PWA achieves hear-linean complexity, with a linear coefficient of approximately 7.87% of Swin Transformer (Liu al. 2021).
2.3 JOHNSON-LINDENSTRAIIS LEMMA-GUIDED CONVOLUTION
Lemma (Johnson-Linddestraasss For any finite set X C RRd with |X| = N and 6 E (0, 1), there exists a linear map f Rd + Rdd with d' 2 c 22 N, all x, y X satisfy (1 - 5)|læ vll S f())) I (1 + +) vll2.
As shown in Figure 4 depth-wise convolution destroys the adjacency relationship between data in the feature space, making it difficult to connect the current clues with the key information of case. Inspired by Likhosherstov et al. (2021) derivation the minimum attention head size
4

---

Published as a conference paper at ICLR 2026
JL. SPA R Sgroup
via the Johnson- Lindens--nndsssssuss (JL) lemma, we build our lightest but robust convolution upon the above theory frame- work. In particular, we extend it to the 3D segmentation, exploring the lower bound on group size while preserving spatial adjacency. The JL lemma states that for N points in ligeedimensional space, we need at least Ollog N) embedding dimensions to preserve pairwise distances. The volume ratio of the input image of M modalities to the intermediate feature u, and each voxel the feature must retain information from at least input voxels. Due to anatomical constraints and the boundedness of the normalized input values, the manifold M of segmentation- related information of the input image patch can be covered by finite number of samples, with coverage count of [ v). Substituting N (MM, v) into the lemma yields the size of the convolution group:
Segmentatioo--llatdd Manifold c B MKp
d2 & x Tumort
fow
DW SPACE R!
Siegmestation-iiniee invlevant Data"
Hard to Recocniae.
Coverage of Normal
Coverage of Tumor Palzh
Coverage of Input. Patch
Palch
Intuitive difference between depth-wise (DW) convolution and guided Convolution the feature space.
Figure
(JL.)
Cgroup = d' 2 CILE log (M,
(2)
where Cgroup is the number of channels per group.
Due to the lack of N in the vision domain, we empirically approximate [[M, v) using NM,, v) = (M w)a, where is related the difficulty of the segmentation task hand. We conduct ab- lation studies on datasets with the richest modality heterogeneity and data distribution to identify the most generalizable scaling factor, which use obtain a lower bound the group size of the convolution layers in each network stage. As analyzed in the Appendix D we will use N
{Ctroup] {n, 2n, 2n, 4n) as the group size for each stage of network, where determined from the most challenging AutoPET-II (Gatidis S 2022) dataset to ensure multi-organ generalization capability. Appendix D further provides detailed derivations, including the covering- number motivated parametrization and the rationale behind empirical approximation.
2.4 SPATIALLY DECOUPLED KNOWLEDGE TRANSFER
To strengthen the representation of the lightweight model, we transfer the rich texture details ex- tracted by the self-supervised texture teacher to the segmentation network via the Gram matrix. Specifically, we start with learning self-supervised detail texture teachers Tm, m 1, M,, who are optimized by M reconstruction tasks. The Gram matrix is commonly used to represent image style and can capture feature channel relationships in a spatially invariant manner. For feature maps E CCC((WDD D) with C channels, the Gram matrix is:
GM(X) CHWD (Xx*) CCC
SDKT is implemented by matching Gram matrices, which is mathematically equivalent minimiz- ing the maximum mean difference (MMD) using a second-order polynomial kernel (Li et al. 2017] Gupta al., 2017). This naturally avoids series of issues caused by excessive ROI discrepancies between the reconstructoo/supeerresoiiiit
features and the segmentation features. Specifically, Gram-based consistency constraint serves positive knowledge transfer path between the seg- mentation features D seg and the M teacher features D Final loss L is:
(Laice + Loe) + Arefore Aedkt
GI (D7) GM (D,s
For information about L. please Appendix
E
Asegs Are and Auakt are the loss weights.
5

---

Published as a conference paper at ICLR 2026
3 EXPERIMENTS
3.1 DATASETS METRICS
> Hecktor2022 (Oriillrr et al. 2022). BraTS2021 (Baid et al. (021). and BraTS2016 (Menze et al. 2014a) (dctails in Appendix F. Unlike typical medical segmentation datasets, the modality hetero- geneity of PET/CT and the complex anatomical structures of multiple organs, and even the whole body, pose unique challenges to all models. We adopt comprehensive evaluation metrics suitable for clinical settings: Model Size (MParams), Computational Complexity (GFLOPs), Efficiency (GPU/CPU Throughput), and Segmentation Performance measured by Dice similarity coefficient (Dice) as the primary indicator, alongside 95% Hausdorff distance (HD95), Precision, and Recall.
We validate the effectiveness of VeloxSeg on four public datasets: AutoPET-II (Gatidis S 2022
3.2 IMPLEMENTATION DETAILS & BASELINES
Our implementation is based on PyTorch 2.4.1. Training is performed on an NVIDIA GeForce RTX 3090 GPU, while inference is run dardized and partitioned into training, validation, and testing subsets in a 6:2:2 ratio. For training, batch size of 4 with 1:1 positiie-to-neeaativ sample ratio. Data augmentation involves random z-axis flipping with 0.5 probability. We train the model for 300 epochs using the Adam W optimizer (Loshchilov & Hutter 2017) with an initial learning rate of 2.5e-4 and a weight decay of 0.01. The learning rate is managed by a linear warmup and cosine annealing scheduler (Liu) 2022).
Intel(R) Xeon(R) Gold 5320 CPU. All datasets are stan-
To ensure a convincing evaluation, we benchmark our method against diverse set of models, in- cluding 8 basic models, 3 multimodal models, and 5 lightweight models, which are categorized accordingly in Tables [6] Furthermore, our analysis covers five distinct architectural paradigms: CNN-based models (UNet (Çiçek et al.) 2016). VNet (Milletari et al., 2016). MedNeXt (Roy et al.) A2FSeg (Wang & Hong] 2023) ) Cra-Transfrrmer hybrids (UNETR (Hatamizadeh et al,,
2023),
2022
Nestedformer (Xing al. 2022) SuperLightNet (Yu al.) )025b)); CNN-KAN hybrids (U-
KAN (Liu et al. (2024b)); CNN-Mambe hybrids (HCMA-UNet (Li et al. 2025)); and CNN-RWKV
include 2 advanced vision foundation models: SAM-Med3D (Wang et al., 2025a), which is evaluated in a zero-shot setting, and DINOv3 (Siméoni et al., (025), for which the linear head is fine-tuned (Liu et al. 2025). Our
hybrids (U-RWKV (Ye al. 2025)). The comparison is extended
comparison conforms to the fair comparison principle outlined in Isensee et al. (2024).
3.3 CLINICAL CAPABILITIES EVALUATION
Figure provides more intuitive comparison of the trade-off's between Dice and parameter count, and between Dice and GPU throughput. Specifically, regarding segmentation performance, Table shows the segmentation performance for PET/CT. Appendix (o shows the qualitative results of all models. Detailed computational costs are provided in the Appendix G Furthermore, we report the GPU memory usage of all models on the three datasets, including training and inference, in Appendix H To release the model's potential, we train VeloxSeg on the nnUNet (Isensee et al.) 2021 Huang al. 2023) training framework and compare with the nnUNet baseline, as shown in Appendix LJ In addition, to verify the modality adaptation ability of the method, we test the performance of MRI segmentation on BraTS2021.
Comparison with Basic Models. Against established basic architectures, including CNN-based, CNN-Transformer based, and CNN-KAN based methods, VeloxSeg demonstrates superior perfor- mance, with significantly lower computational cost. On the AutoPET-II dataset, VeloxSeg achieves a 62.51% Dice. This result marginally outperforms the best basic model, VSmTrans, using only 13.30% of its parameters and 1.96% of its GFLOPs. On Hccktor2022, VeloxSeg still surpasses all other models. These demonstrate that VeloxSeg is an efficient model in medical segmentation.
Comparison with Multimodal Models. When compared to specialized multimodal architectures, VeloxSeg demonstrates its effectiveness and efficiency in cross-modke feature integration. On the AutoPET-II II dataset, VeloxSeg's 's Dice of 62.51% outperformed H- H-DenseFormer, Nestedformer, and A2FSeg by 1.01%. 1.13%. and 1.65%, respectively, while achieving GPU throughput improve- ments of x 7.75x and significant reduction of computational complexity. Furthermore,

---

Published as a conference paper at ICLR 2026
AutoPET-II
Hecktor2022
Method
Venue
Dice
HID95
Prec. 1
Rec. f
Dice
HD95 +
Prec.
Rec. t
UNet
1ICCAT'16
59.41
211.31
62.32
70.74
50.25
65.03
72.13
41.50
VNet
3DV'15
53.21
242.78
53.21
60.85
55.61
41.46
78.21
46.01
MedNeXt-S
MICCAI'2
53.94
180.83
60.63
60.25
47.22
79.82
64.89
40.38
UNETR
WACV'22
51.49
257.30
51.49
61.03
48.10
73.27
70.71
39.11
Swin UNETR
MICCAI'1
62.24
242.07
62.91
73.30
44.56
103.02
62.43
37.55
VSmTrans
MIA'24
62.46
223.88
65.19
70.92
52.91
78.03
61.91
50.97
UNETR+
TMI224
36.50
178.57
36.50
60.16
29.95
27.74
61.84
21.75
U-KAN
AAAI'25
60.67
70.91
62.03
72.94
55.89
23.48
77.72
46.89
Nestedformer
MICCAI'22
61.38
265.51
61.38
64.29
40.17
72.95
63.22
32.59
A2FSeg
MICCAI'23
60.86
131.48
60.86
76.10
40.90
32.95
77 02
30.57
H--eenseFormer
2ICCAI23
61.50
252.98
61.41
75.76
46.79
34.84
78 13.33
35.31
SAM-Med3D (CT)
TNNLS'25
13.13
101.24
19.82
16.70
27.52
84
43.94
24.46
SAM-Med3D (PET)
TNNLS'25
26.59
101.94
31.92
31.86
31.94
18.03
69.69
24.35
DINOv3-L (PET)
Arxiv'25
10.87
6.85
64.96
9.43
9.40
25.93
DINOv3-L (CT+PETT
Arxiv'25
12.17
7.50
71.16
30.86
34.99
39.98
SegFormer-3D
CVPRW"24
40.44
174.43
56.73
38.19
48.47
51.29
73.63
38.35
Slim UNETR
TMI22
52.53
310.53
53.99
66.55
49.40
56.55
69.53
41.20
SuperLightNet
CVPR'25
48.35
59.09
60.82
47.61
50.03
34.36
75.29
40.65
HCMA-UNet
ICME'25
43.40
146.11
43.32
62.46
42.06
146.11
67.68
33.18
U-RWKV
MICCAI'25
57.18
61.12
66.69
59.40
45.97
56.83
64.52
39.71
VeloxSeg
Ours
62.51
241.08
67.76
66.28
56.48
47 17.66
74.81
49.24
Due to the small object and camouflage recognition involved, DINOv3-L (CT) cannot recognize tumors. ii) means that the value is of range.
Table Comparisons of segmentation performance on PET/CT datasets. The best performance is highligted by red, followed by blue. VeloxSeg is highlighted in green
on Hecktor2022, due reduced data size, other multimodal models exhibit overfitting and overly conservative predictions, while VeloxSeg's Dice score remains stable.
Comparison with Lightweight Models. Against other lightweight methods, VeloxSeg clearly superior. It leads in Dice on both datasets by significant margin of over 5%. While some com- petitors have fewer parameters, they computationally expensive lack CPU support for clinical use. VeloxSeg offers the best balance, achieving 1.66 MParams and 1.79 GFLOPs. It also achieves high GPU throughput of 599.06 patches/s and supports CPU-only devices, making it the most clinically practical solution.
Comparison of Peak GPU Memory Usage. As shown in Appendix H VeloxSeg achieves the low- est or second lowest peak GPU memory usage among all methods. Compared to basic CNN/CNN- Transformer baseline models, VeloxSeg reduces memory usage by up 20x during training and up to 24x during inference. Even among lightweight models, VeloxSeg is close to the most compact model (Slim UNETR), reducing memory usage by up to 10x compared to other lightweight models.
Train nneNet Training Framework. As analyzed in Appendix I VeloxSeg achieves 14.2% Dice improvement with only 1.87% of nnUNet's MParams and 0.058% its GFLOPs (Isensec) al. 2021), accompanied by 4.8x improvement in GPU throughput and 52.5x CPU.
Modality Adaptation Evaluation. On the BraTS2021 dataset, which contains MRI modalities, early fusion strategy VeloxSeg-C achieves superior performance, surpassing the second-best method by 1.72% Dice. This demonstrates that our VeloxSeg can adapt to diverse multimodal segmentation tasks. Details can be found in Appendix

---

Published as a conference paper at ICLR 2026
Modules
Ablation
Hyper-Parameters
Params
FLOPs
Thr. GPU
Dice
Conv. Trans. SDKT.
(M)
(G)
(Pat.s)
(%)
X
X
Width
(32. 64, 128, 256)
2.65
5.31
145.63
48.96
X
X
(16, 64, 128)
0.73
2.41
616.53
50.10
X
X
Kernel Size
(7) (1, 3, 5)
0.73
2,41
616.53
50.10
X
x
0.66
2.30
295.02
53.65
X
x
11
0.66
2.30
295.02
53.65
X
(1, 21,2,244
0.75
2.33
291.18
53.95
X
x
Group Size
(2,4, 4,8)
0.89
2.44
284.83
54.40
X
(4, 8, 8, 16)
1.18
2,66
282.13
55.84
{8, 16, 16, 32)
1.75
3.11
279.48
55. 14
(16, 32, 64, 128)
4.76
4.18
290.72
56.20
(2, 2, 2, (1, 1,
2.37
3.07
137.87
59.56
Attention Depth
x
1.88
2.90
185.08
61.03
x
(1, 1, 4,
1.88
2.90
331.56
61.03
Expansion Ratio
(3. 0,3,,,2 2)
1.61
2.84
336.94
61.43
V
Teacher
+ Texture Teacher
1.61
2.84
336.94
59.64
Up
Unify Upsampling
1.66
1.79
390.91
59.71
V
Gram
+ Gram Supervision
1.66
1.79
390.91
62.51
Table 2: Module ablation experiments on AutoPET-II. "Conv:" convolution encoder; "Trans." transformer encoder; "SDKT.": spatially decoupled knowledge transfer. The best performance is in red and the second is in blue. Final setting is highlighted in green
3.4 MODULE ABLATION
We evaluate the performance of three model designs on AutoPET-II: JLC, PWA, and SDKT (Ta- ble When using JLC alone, Params and Dice have the lowest performance. Although the frame- work is the simplest, the PLoPs/throuhhuut is suboptimal due to the use transposed convolution for upsampling. After adding the attention mechanism, the accuracy increased by 5.59%, but the throughput decreased by 233.6 Patches/s. After changing the upsampling strategy, FLOPPs are signif- icantly reduced from 2.84 G to 1.79 G, and the GPU throughput is increased from 336.9H 599.06 Patches/s. The last three rows in the table show that it is not enough to just optimize the encoder's s detail representation after adding the texture teacher. Only through the SDKT strategy based on Gram matrices can the representation ability be improved. For more specific reasons and analysis of hyper- pyper-parammter selection, please see Appendix L
PWA Effect Evaluation. To verify the robustness of PWA, we conduct three experiments:
Reduce computational redundancy through multi-scale windows. VeloxSeg utilizes PWA to par- allelize the computation of multi-scale relationships and reduce redundancy. Experimental details of Figure c5 be found in Appendix M The inter-group differences PWA are significant and positively correlated with the window size, indicating that redundant information is reduced and to long-distance modeling is efficient. Regarding the fourth attention stage, its design is more similar attention, retaining some redundancy.
Changes in features after adding PWA. We visualize the model's decoding and its t-SNE projection, shown in Figure [6] The results indicate that PWA helps distinguish tumor regions from high- metabolic regions while producing a more compact feature distribution.
Effectiveness in handling heterogeneous modalities. We test various modal input combinations, whose details could be found in Appendix N Notably, introducing modal interaction into PWA improves the Dice score by 5.75%, significantly enhancing performance robustness without signifi- cantly increasing computational costs.
JLC Effect Evaluation. To verify the robustness of JL.C, conduct four experiments:
8

---

Published as a conference paper at ICLR 2026
Moon Anention Distrece
3
12
24
12
[a]
(4)
(d)
Figure 5: Average attention distance distribu- tion of PWA on AutoPET-II. (a)-(d) show the results for PWA across 4 stage. Y-axis: age attention distance; X-axis: big window size. Wilcoxon rank-sum test: ns (0.05 < p 1), (0.01 - 0.05), **** (p S 0.0001).
Figure 6: Visualization of model decoding and t-SNE plot. (a)-(b) show results without atten- tion; (c)-(d) with attention. In (a) and (c), "Low Met." and "High Met." represent low/high- metabolism PET regions. In (b) and (d), CT background with red tumor outline.
Methods
MParams +
Dice
Similarity Coefficient (%)
$2.66
54.14
53.55
51.06
FC.
4.78 +3.58
79.97
+0.63
50
FC. Pn.
1.27 +0.07
62.00 -1.34
55089
JLC
1.20
79.34
Depth-wise buseline:
45
Table 3: Domain generalization capa- bility comparison of the JLC and & pruning methods (Filters' Importance,
42.75
Group Size
40
(2n.2..21.2n)
39.66
(n,2n,2n,44)
2016 Fang al.] 2023) (BraTTS2021 BraT52016 TCIA). "FC." represents full convolution, and "Pn." represents pruning operation.
1
2
Figure 7: Dice performance comparison between differ- ent group size configurations.
Comparison segmentation performance between the JL setting and the standard setting with sim- ilar parameter sizes, as shown Figure [ We use uniform kernel size of 3 to ensure a gen- eral setup. The JL-guided configuration, [n, 2n, 2n, 4n), consistently surpasses the larger setup, (2n, 2n, 2n, 2n in all cases. The best performance is a 54. 14% Dice score, achieved when n = 4. This suggests the JL-guided group size arrangement enables lightweight model.
robust feature extraction
Comparison of external test performance between the JL setting and the pruned setting with similar parameter sizes, as shown in Table 3 It verifies the generalization advantage of our lightweight convolution over the pruning method. The model with l pruning (Filters' Importance 2016) on full convolution performs significantly worse than JLC on BraT52015 TCIA cases, even after cycle of training, pruning, and retraining.
Testing the segmentation performance of JLC two other datasets, detailed in the Appendix O We test the segmentation performance of pure convolution networks. The convolution with JT.- guided group sizes consistently outperforms the depth-wise convolution, achieving performance gains 6.25% on Hccktor2022 and 1.16% on BraTS.0 with only a marginal increase of 0.091 million parameters. Notably, on the Hecktor2022 dataset, the JLC even surpasses the segmentation performance of the full convolution while using 0.63 million fewer parameters.
Comparison of the t-SNE projection visualizations of JLC and depth-wise convolution is shown in Figure g We test the depth-wise wise convolution and JLC in Figure [ggure/aadd aadd Appendix O providing direct visual evidence that wist convolution disrupts the geometric adjacency between tokens.

---

Published as a conference paper at ICLR 2026
Deych Oway LSND (A (o
J.. Cams LSNK (AmAPPiTI))
Deych Coty USNI (lieckio3301))
Ju. Cleor USVE (1ecktor2025)
7
Les Moc.
Low Min. High Mrt.
Iish Met.
Tame
Tomer
-10
Figure 8: t-SNE plots of depth-wise convolution and JC convolution. "Low Met." and "High Met." represent low/high metabolism PET regions, respectively.
05
S0
014
Simiarity Codfficient (8)
Strategy
Dice 1
HD95
08
0.1
59.71
291.81
+4
1.67 -58.04
626.79
Arc 530
1.0
+Affinity
41.44-18.27
354.01
06
+Shared ROI
57.15 -2.56
397.43
OCL
+SDKT
62.51 +2.80
241.08
Table Comparison of knowledge transfer paths constructed with different losses: Sang et al. (202T) use & loss. Wang et al. (2021) 1) use Affinity loss, Qiu al. (2023) use agent loss shared
Ocn
361-9C:
50a UXE
Oen 381
56o UN
Figure Box plots of Dices at different lesion vol- umes. "MM." PWA+JLC multimodal backbone; *MM++RC.": backbone with reconstruction teacher; "Ours": backbone with the teacher and SDKT.
ROI, and we use SDKT.
Gram-Based Transfer Effeet Evaluation. Our method is the only one to demonstrate positive knowledge transfer, shown in Table This due to our method's avoidance of irrelevant features from the texture teachers, leading to better convergence. To further evaluate the effectiveness of SDKT, we analyze Dice across varying lesion volumes, as shown in Figure 9 VeloxSeg outperforms other lightweight models in segmentation mask fineness across all lesion volumes. Notably, for lesion volumes between 1000 and 5000, "MM.", "MM.+RC." and "Ours" show significant upward trend. We attribute this to the increased influence of complex textures in tumor segmentation at these sizes. The above experiments show that there is potential for optimization in small lesion segmentation. The loss weight hyperparameter experiment can be found in Appendix L
CONCLUSION
In this paper, propose VeloxSeg, lightweight, theory-based framework that systematically al- leviates the " "efficiency / robustness conftict" in 3D medical image segmentation. By extending the Johnson-Lindenstrauss lemma to the convolution setting, we derive a theoretical lower bound on the group size of convolution per stage, ensuring spatial adjacency and enabling robust detail extraction. Our paired window attention mechanism, by ensembling tumor localization team composed of attention at different scales, has near-linear complexity and more powerful modeling capabilities. Furthermore, the multimodal interaction of PWA significantly enhances model representation. Fur- thermore, our Spatially Decoupled Knowledge Transfer strategy establishes positive knowledge transfer path between the self self-supervised texture teacher and the segmentation network, enabling detailed representations that surpass baseline models without increasing inference overhead. Com- prehensive evaluation on four diverse clinical datasets demonstrates that VeloxSeg achieves strong robustness with minimal computational cost, requiring only single CPU core.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Peri Akiva, Matthew Purri, and Matthew Leotta. Self-supervised material and texture representation learning for remote sensing tasks. Proceedings the IEEE/CVF conference on computer vision and pattern recognition, pp. 9203-82155 2022.
Ujjwal Baid al. The rsma-asnr- brats 2021 benchmark on brain tumor segmentation and radiogenomic classification. arXiv preprint arXiv:2107.02311 2021.
Spyridon Bakas et al. Advancing the cancer genome atlas glioma mri collections with expert seg- mentation labels and radiomic features. Scientific data, 4(1):1-13, 2017.
Johan Bussink, Johannes HAM Kaanders, Winette TA Van Der Graaf, and Wim JG Oyen. Pet-ct for radiotherapy treatment planning and response monitoring in solid tumors. Nature Reviews Clinical Oncology, 8(4):233-25 242, 2011.
Linghan Cai, Jianhao Huang, Zihang Zhu, Jinpeng Lu, and Yonghing Zhang. localization-to: segmentation framework for automatic tumor segmentation in whole-hody pet/ct images. arXiv preprint (rXiv:2309.054444 2023.
Linghan Cai, Lijiang Chen, Jianhao Huang, Yifeng Wang, and Yongbing Zhang. Know your ori- entation: viewpoint-awe framework for polyp segmentation. Medical Image Analysis, 97: 103288, 2024.
Yen-Chun Chen, Linjie Li, Licheng Yu, Ahmed El Kholy, Faisal Ahmed, Zhe Gan, Yu Cheng, and Jiinging Liu. Uniter: Universal image-text representation learning. In European conference on computer vision, pp. 104-120. Springer, 2020.
François Chollet. Xception: Deep learning with depthwise separable convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1251-1258, 2017.
Özgün Çięek, Ahmed Abdulkadir, Soeren S Lienkamp, Thomas Brox, and Olaf Ronneberger. 3d u-net: learning dense volumetric segmentation from sparse annotation. In Medical Image Comput- ing and Computer-Assisted Intervention- 2016: 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part II 19, pp. 424-432. Springer, 2016.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 11929, 2020.
Alexey Dosovitskiy et al. An image worth 16x16 words: scale. In International Conference on Learning Representations, 2021.
Transformers for image recognition at
Hao Du, Qihua Dong, Yan Xu, and ling Liao. Tdformer: Top-down token generation for 3d medical image segmentation. IEEE Journal of Biomedical and Health Informatics, 2025.
Gongfan Fang, Xinyin Ma, Mingli Song, Michael Bi Mi, and Xinchao Wang. Depgraph: Towards any structural pruning. In Proceedings of the IEEE/VV Conference on Computer Vision and Pattern Recognition, pp. 16091-16101, 2023.
Determine Filters' Importance. 1.XY:16608.08700.. 2016.
Pruning filters for efficient convnets.
arXiv preprint
Kuestner Gatidis S. A whole-body fdg-pet/ct dataset with manually annotated tumor lesions (fdg-pet-ct-lesionns) The Cancer Imaging Archive, 226, 2022.
Leon Gatys, Alexander S Ecker, and Matthias Bethge. Texture synthesis using convolutional neural networks. Advances in neural information processing systems, 28, 2015.
Leon A Gatys, Alexander S Ecker, and Matthias Bethge. Image style transfer using convolutional neural networks. In Proceedings of the IEEE conference on computer vision and pattern recog- nition, pp. 2414-2423, 2016.
11

---

Published as a conference paper at ICLR 2026
Albert Gu and Tri Dao. Mamba: Ainear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.
Agrim Gupta, Justin Johnson, Alexandre Alahi, and Li Fei-Fei. Characterizing and improving sta- bility in neural style transfer. In Proceedings of the IEEE International Conference on Computer Vision, pp. 4067-4076, 2017.
Ali Hatamizadeh, Vishwesh Nath, Yucheng Tang, Dong Yang, Holger R Roth, and Daguang Xu. Swin unetr: Swin transformers for semantic segmentation of brain tumors in mri images. In International MICCAI Brainlesion Workshop, pp. 272-284. Springer, 2021.
Ali Hatamizadeh et al. Unetr: Transformers for 3d medical image segmentation. In Proceedings of the IEEE/CVF winter conference applications of computer vision, pp. 574-584, 2022.
Along He, Kai Wang, Tao Li, Chengkun Du, Shuang Xia, and Huazhu Fu. H2former: An efficient hierarchical hybrid transformer for medical image segmentation. IEEE Transactions on Medical Imaging, 42(9):2763 2775, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp- 770-778, 2016.
Yufan He, Pengfei Guo, Yucheng Tang, Andriy Myronenko, Vishwesh Nath, Ziyue Xu, Dong Yang, Can Zhao, Benjamin Simon, Mason Belue, et al. Vista3d: A unified segmentation foundation model for 3d medical imaging. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 20863-20873 2025.
Juha Heinonen. Lectures on analysis on metric spaces. Springer Science & Business Media, 2001.
Ziyan Huang. Haoyu Wang. Zhongying Deng. Jin Ye, Yanzhou Su, Hui Sun, Junjun He, Yun Gu, Lixu Gu, Shaoting Zhang, et al. Stu-net: Scalable and transferable medical image segmentation models empowered by large-scale supervised pre-training. arXiv preprint arXiv:2304.06716. 2023.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448-456. pmlr, 2015.
Fabian Isensee, Paul F Jaeger, Simon AA Kohl, Jens Petersen, and Klaus H Maier-Hein. net: self-configuring method for deep learning-based biomedical image segmentation. Nature methods, 2812)2003-211, 2021.
Fabian Isensee, Tassilo Wald, Constantin Ulrich, Michael Baumgartner, Saikat Roy, Klaus Maier- Hein, and Paul F Jaeger. nnu-net revisited: A call for rigorous validation in 3d medical image segmentation. In International Conference on Medical Image Computing and Computer-Assisted Intervention, pp. 488 498. Springer, 2024.
Charalambos Kaittanis et Environment-respoosivee nanophores for therapy and treatment moni- toring via molecular mri quenching. Nature communications, 5(1):3384, 2014.
Vibhu Kapoor, Barry M McCook, and Frank $ Torok. An introduction to pet-ct imaging. Radio- graphics, 24(2):523-543, 2004.
Lynn D Kramer, GE Locke, SE Byrd, and JAFAR Daryabagi. Cerebral cysticercosis: documentation of natural history with ct. Radiology, 171(2):499446 1989.
Hulin Kuang, Yahui Wang, Xianzhen Tan, Jialin Yang. Jiarui Sun, Jin Liu, Wu Qiu, Jingyang Zhang, Jiulou Zhang, Chunfeng Yang, al. Lw-ctrans: A lightweight hybrid network of enn and trans- former for 3d medical image segmentation. Medical Image Analysis, 102:103545, 2025.
Yann LeCun, Berharrd Boser. John S Denker, Donnie Henderson, Richard E Howard. Wayne Hub- bard, and Lawrence Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation, 1(4):541-551, 1989.
12

---

Published as a conference paper at ICLR 2026
Haoxuan Li, Peiwu Qin, Xi Yuan, and Zhenglin Chen. Hema-unet: A hybrid enn enn-mamba unet with axial self-attention for efficient breast cancer segmentation. arXiv preprint arXiv:2501.00551, 2025.
Yanghao Li, Naiyan Wang, Jiaying Liu, and Xiaodi Hou. Demystifying neural style transfer. 2017.
Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power of self-attention matrices. arXiv preprint wVXiv22066033764 2021.
W Johnson Lindenstrauss and J Johnson. Extensions lipschitz maps into a hilbert space. Con- temp. Math, (9-2206))): 1984.
Che Liu, Yinda Chen, Haoyuan Shi, Jinpeng Lu, Bailiang Jian, Jiazhen Pan, Linghan Cai, Jiayi Wang, Yundi Zhang, Jun Li, al. Does dinov3 set medical vision standard? arXiv preprint arXiv:2509.06467, 2025.
Tiange Liu, Qingze Bai, Drew A Torigian, Yubing Tong, and Jayaram K Udupa. Vsmtrans: A hybrid paradigm integrating self-attention and convolution for 3d medical image segmentation. Medical image analysis, 98: 103295, 2024a.
Ze Liu et Swin transformer: Hierarchical vision transformer using shifted windows. Proceed- ings of the IEEE/CVF international conference on computer vision, pp- 10012-10022, 2021.
Zhao Liu. Super convergence cosine annealing with warm-up learning rate. In CAIBDA 2022; 2nd International Conference on Artificial Intelligence, Big Data and Algorithms, pp. 1-7, 2022.
Ziming Liu et al. Kan: Kolmogorov-amoll networks. arXiv preprint arXiv:2404. 19756, 2024b.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint
arXiv:111.00101, 2017.
Jiasen Lu, Dhruv Batra, Devi Parikh, and Stefan Lee. Vilbert: Pretraining task-agnostic visiolin- guistic representations for risionaand-languag tasks. Advances in neural information processing systems, 32, 2019.
Ningning Ma, Xiangyu Zhang, Hai- Tao Zheng, and Jian Sun. Shufflenet v2: Practical guidelines for efficient enn architecture design. In ECCV, 2018
Renzo Manara et al. Brain and spine mri features of hunter disease: frequency, natural evolution and response therapy. Journal Inherited Metabolic Disease: Official Journal of the Society for the Study of Inborn Errors of Metabolism, 34(3):763-780, 2011.
Tarck Mansour et al. Deep neural networks lazy: the inductive bias of deep learning. PhD thesis, Massachusetts Institute of Technology, 2019.
Bjoern H Menze, Andras Jakab, Stefan Bauer, Jayashree Kalpathy-Cramer Keyvan Farahani, Justin Kirby, Yuliya Burren, Nicole Porz, Johannes Slotboom, Roland Wiest, The multimodal brain tumor image segmentation benchmark (brats). IEEE transactions medical imaging, 34 (10):199,-2024, 2014a.
Bjoern H Menze et al. The multimodal brain tumor image segmentation benchmark (brats). IEEE Transactions on medical imaging. 34(10):1993-200 2b.
Fausto Milletari, Nassir Navab, and Seyed-Ahmad Ahmadi. V-net: Fully convolutional neural net- works for volumetric medical image segmentation. 2016 fourth international conference on 3D vision (3DV), pp. 565-571 leee, 2016.
Pavlo Molchanov, Arun Mallya, Stephen Tyree, luri Frosio. and Jan Kautz. Importance estimation for neural network pruning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11264-11272. 2019.
Wei Mu et al. Non-invasive decision support for nscle treatment using pet/ct radiomics. Nature communications, 111)::2228 2020.
13

---

Published as a conference paper at ICLR 2026
Usman Muhammad, Jorma Laaksonen, and Lyudmila Mihaylova. Towards lightweight hyper- spectral image super-resolutien with depthwise separable dilated convolutional network. arXiv preprint 2tth::25550003374 2025.
Valentin Oreiller et al. Head and neck tumor segmentation in pet/ct: the hecktor challenge. Medical image analysis, 77:102336, 2022.
Yan Pang et al. Slim unetr: Scale hybrid transformers to efficient 3d medical image segmentation under limited computational resources. IEEE transactions on medical imaging, 43(3):994-1005, 2024.
Himashi Peiris, Munawar Hayat, Zhaolin Chen, Gary Egan, and Mchrtash Harandi. Uncertainty- guided dual-views for semi-supervised volumetric medical image segmentation. Nature Machine Intelligence, 5(7):724-738, 2023.
Bo Peng et al. Rwkv: Reinventing mnss for the transformer era. arXiv preprint 2idi::30051304.. 2023.
Shehan Perera, Pouyan Navard, and Alper Yilmaz. Segformer3d: an efficient transformer for 3d medical image segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 4981-4988, 2024.
Zhongxi Qiu, Yan Hu, Xiaoshan Chen, Dan Zeng, Qingyong Hu, and Jiang Liu. Rethinking dual- stream super-resolution semantic learning in medical image segmentation. IEEE Transactions Patlern Analysis and Machine Intelligence, 46(1):451-46 2023.
Md Mostafijur Rahman, Mustafa Munir, and Radu Marculescu. Emcad: Efficient multi-scale con- volutional attention decoding for medical image segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11769-11779, 2024.
Saikat Roy et al. Mednext: transformer-driven scaling of convnets for medical image segmentation. In International Conference on Medical Image Computing and Computer-Assisted Intervention, pp. 405-415. Springer, 2023.
Shaohao Rui, Lingzhi Chen, Zhenyu Tang. Lilong Wang. Mianxin Liu, Shaoting Zhang, and Xi- aosong Wang. Multi-modal vision pre-training for medical image analysis. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 5164-5174, 2025.
Yu Sang, Jinguang Sun, Simiao Wang, Heng Qi, and Keqiu Li. Super-resolution and infection edge detection co-guided learning for covid-19 et segmentation. In ICASSP 2021 2021 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pp. 1665-1669, 2021. doi: 10. 10.1090CAASSP39928222.... 944.4443377.
Abdelrahman M Shaker, Muhammad Maaz, Hanoona Rasheed, Salman Khan, Ming-Hsuan Yang, and Fahad Shahbaz Khan. Unetr++: delving into efficient and accurate 3d medical image seg- mentation. IEEE Transactions on Medical Imaging, 2024.
Oriane Siméoni, Huy V Vo, Maximilian Seitzer, Federico Baldassarre, Maxime Oquab, Cijo Jose, Vasil Khalidov, Mare Szafraniec, Seungeun Yi, Michaël Ramamonjisoa, al. Dinov3. arXiv preprint arXiv:2508. 01022500..0004,, 2025.
Hu Su, Ying Li, Yifan Xu, Xiang Fu, and Song Liu. A review of deep-learning resolution: From methods to applications. Pattern Recognition, 157:1 10935, 2025.
beep-eeariig-bbsed based super-
Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt. Test-time train- ing with self-supervision for generalization under distribution shifts. In International conference on machine learning, pp. 9229-9248. PMLR, 2020.
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing in- gredient for fast stylization. arXiv preprint arXiv:1607.0002. 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
14

---

Published as a conference paper at ICLR 2026
Tassilo Wald, Constantin Ulrich, Stanislav Lukyanenko, Andrei Goncharov, Alberto Paderno, Max- imilian Miller. Leander Maerkisch, Paul Jaeger, and Klaus Maier-Hein. Revisiting mae pre- training for 3d medical image segmentation. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 5 5186-5196, 2025.
Haoyu Wang, Sizheng Guo, Jin Ye, Zhongying Deng, Junlong Cheng, Tianbin Li, Jianpin Chen, Yanzhou Su, Ziyan Huang, Yiqing Shen, et al. Sam-med3d: A vision foundation model for general-purpose segmentation on volumetric medical images. IEEE Transactions on Neural Net- works and Learning Systems, 2025a.
Hongyi Wang et al. Patch-free 3d medical image segmentation driven by super-resolution technique and self-supervised guidance. In Medical Image Computing and Computer Assisted Intervention- MICCAI 2021: 24th International Conference, Strasbourg, France, September 27-October 1, 2021, Proceedings, Part I 24, pp. 131-141. Springer, 2021.
Jiawen Wang al. Masktwins: Dual-form complementary masking for domain image segmentation. In Forty-second International Conference on Machine Learning, 2025b.
Zirui Wang and Yi Hong. A2fseg: Adaptive multi-modal fusion network for medical image seg- mentation. In International Conference on Medical Image Computing and Computer-Assiste Intervention, pp. 673-681 681. Springe,,, 2023.
Ziyue Wang, Linghan Cai, Chang Han Low, Haofeng Liu, Junde Wu, Jingyu Wang, Rui Wang. Lei Song, Jiang Bian, Jingjing Fu, et al. 3dmedagent: Unified perception-to-undeer to-understandingg for 3d medical analysis. arXiv preprint arXiv:2602. 18064, 2026.
Chaoyi Wu, Xiaoman Zhang, Ya Zhang, Hui Hui, Yanfeng Wang, and Weidi Xie. Towards generalist foundation model for radiology by leveraging web-scale 2d&3d medical data. Nature Communi- cations, 16(1):7866, 2025.
Zhaohu Xing, Lequan Yu, Liang Wan, Tong Han, and Lei Zhu. Nestedformer: Nested modality- aware transformer for brain tumor segmentation. In International conference on medical image computing and computer-asee assiisted intervention, pp. 150-150. Springer, 2022.
Zhaohu Xing, Tian Ye, Yijun Yang. Du Cai, Baowen Gai, Xiao-Jian Wu, Feng Gao, and Lei Zhu. Segmamba-v2: Long-range sequential modeling mamba for general 3d medical image segmenta- tion. IEEE Transactions on Medical Imaging, 2025.
Hongbo Ye al. U-rwkv: Lightweight medical image segmentation with direction direction-adaptive rwkv. arXiv preprint arXiv:2507.1144 2025.
Bin Yu, Quan Zhou, Li Yuan, Huageng Liang, Pavel Shcherbakov, and Xuming Zhang. 3d medical image segmentation using the serial- parallll convolutional neural network and transformer based on cross windoww self-attention. CAA/ Transactions on Intelligence Technology, 10(2):337-348, 2025a.
Feng Yu, Jiacheng Cao, Li Liu, and Minghua Jiang. Superlightne: tion network for multimodal brain tumor segmentation. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 5197-5206, 2025b.
Lightweight parameter aggrega-
Linfeng Zhang, Muzhou Yu, Tong Chen, Zuoqiang Shi, Chenglong Bao, and Kaisheng Ma. Auxil- iary training: Towards accurate and robust models. In Proceedings of the IEEF/CVE conference on computer vision and pattern recognition, pp. 372-381, 2020.
Shenhai Zheng, Xin Ye, Chaohui Yang, Lei Yu, Weisheng Li, Xinbo Gao, and Yue Zhao. Asym- metric adaptive heterogeneous network for multi-modality medical image segmentation. IEEE Transactions on Medical Imaging, 2025.
Lanyun Zhu Deyi Ji, Shiping Zhu, Weihao Gan, Wei Wu, and Junjie Yan. Learning statistical texture for semantic segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12537-12546, 2021.
Jing Zou, Lanqing Liu, Qi Chen, Shujun Wang, Zhanli Hu, Xiaohan Xing, and Jing Qin. Mmr- mamba: Multi modal mri reconstruction with mamba and Medical Image Analysis, 102: 103549, 2025.
information fusion.
15

---

Published as a conference paper at ICLR 2026
APPENDIX
Table of content:
Use of Large Language Models (LLMs)
Necessity of Multi-Scale Attention
Details of PWA
Details of JL-guided Group Size
Details Loss Function
Dataset Details
Details of computational performance
Details of GPU memory usage
Results on the nnUNet training framework
Qualitative Results
K
Modality Adaptation Evaluation
Hyperparameter Analysis
M
Computation of Mean Attention Distance
PWA MultiModal Evaluation
JL-Setting Generalization Evaluation
Comparison of Different Attention Mechanisms
Comparison of Different Knowledge Transfer Strategies
Scaling Law of VeloxSeg
A USE OF LARGE LANGUAGE MODELS (LLMS)
To enhance the quality and readability of this manuscript, we use Large Language Models (LLMs) for assistance with the following tasks:
1. Table Formatting:
Improving the presentation of tables, including adjustments to spacing, typography, and alignment conform publication standards.
2. Proofreading: Identifying and correcting grammatical errors, such improper tense and word usage.
3. Language Refinement: Refining phrasing and sentence structure to improve clarity, con- ciseness, and overall flow.
B NECESSITY OF MULTI-SCALE ATTENTION
CT scans the human body using X-rays and reconstructs two-dimensional image from one- dimensional projection data. These two-dimensional images are then stacked into continuous three-dimensional image. CT imaging is characterized by high resolution, low tumor specificity, and rich structural information (Kramer al. 1989).
PET generally refers to 18F-FDG PET. Radiologists use the short-lived radionuclide 18F to label glucose. After injecting this labeled glucose into the body, they observe the accumulation of glucose, which indirectly reflects the metabolic activity of human tissues. Because tumors require large amounts of glucose to support their growth and prolifcration, tumor areas often appear bright in PET images (Mu al. 2020 Kapoor et al. 2004 Bussink et al. 2011).
Magnetic resonance imaging (MRI) provides rich complementary information for analyzing brain tumors and is routinely used in clinical practice. Specifically, for gliomas, commonly used MRI sequences include T1-weighted (TI), contrast-enhancee T1-weighted (TIGd), T2-weighted (T2), and T2 fluid-attenuated inversion recovery (T2-FLAIR) images; each sequence plays different
16

---

Published as a conference paper at ICLR 2026
PW, Gathering
ie wge x.0
Linear Proj.
Channel Split
BW, Puritiio
SW Pool
Wings
Wintes Win;
Scal- Fentures
Grouped Attention
Multimodal
Grouped
Attention
E
Puired Scattering
Linnn Pmn
Charnell Split
BW Partition
SW Pool
Paralie Multi-eate Features
Nile
Win. WineL
Win,
Winst.
Wing
Win.
Figure 10: Detailed architecture of Paired Window Attention (PWA). This figure focuses visually showing the feature flows of PWA.
role in distinguishing between the tumor, peritumoral edema, and the tumor core. For meningiomas, these sequences exhibit distinct characteristic features on T1Gd and contrast-enhanced T2-FLAIR
(FLAIR-C) MRI images (Menze et 2014b) Manara et al., 2011 Kaittanis et 2014 Bakas
et al. 2017).
This indicates that different medical modalities exhibit significant differences their regions of interest. In tumor imaging, PET imaging, characterized by high metabolic sensitivity and low res- olution, excels at localizing tumors at large scales, but its low resolution prevents clear delineation of tumor morphology. While CT imaging is less sensitive for tumors, it excels at clearly delineating tumor tissue contours at small scales. Furthermore, the four contrast types in MRI contribute differ- ently to the identification of targets three different scales: tumor, peritumoral edema, and tumor core. Therefore, multi-scale modality interaction is crucial in multimodal medical tasks.
C DETAILS OF PWA
C.I PYTORCH CODE
We've ve organized the PyTorch code and feature shape changes of PWA help readers under- stand its key operations. As shown in the Algorithm [ Nwin log H / log(r) 1, which means that we expand the large window (h,, Wh, d, by Noiin features. In the AutoPET-II dataset, we set the minimum large window size of each stage to (3, 3, 3), (6, 6, 6) (3, 3,3), (3, 3, 3), which means that after synchronous expansion of paired windows, the large window sizes of each stage are:
to obtain full-image-
First Stage: (3,3,3) (6, 6, 6). (12, 12, 12), (24, 24, 24);
Second Stage: (6, 6, 6) (12, 12, 12);
Third Stage: (3, 3, (6, 6);
Forth Stage: 33.
The settings of BraTS2021 are the same. In the Hecktor2022 dataset, the minimum maximum window size at each stage is (4, 2) (8, 8, 4) (4, 4, 2), (4, 2). The number of windows must be divisible by the number of channels of the feature map the current stage avoid extensive output channels during linear mapping. Therefore, the minimum maximum window size in the second stage is doubled.
C.2 FEATURE FLOW
As shown in Figure 10 given the m-th modal feature of the k-th encoder stage, E N
CC xH* x
need to first compute set of ordered paired window sizes {Win Jisl
17

---

Published as a conference paper at ICLR 2026
Algorithm 1: Pytorch Code of Paired Window Attention (PWA)
Input:
E:
Input tensor of shape [M, C, H,H,W,D,
B: Min big window size, [b,, Wb, da).
Min small window size, [lg. Was de)-
P: Expansion ratio for the paired windows.
C: Number of channels per window after linear projection.
Output: Attentions of all modalities A
def PWA (E. B, S. r): Nwin - [og(H (H/HHHHH ha)/ (o)(r)) 1// the number of paired windows
/x 1) Linear Projection on
Q K, < (PWC(LN(E)) for in range(3))
/+ 2) Paired Window Gathering
for X in JQ, K, Vo do
Xs // Initialize list for window features for i 4- to Nwin do
// Split channels for each window feature
X; x|;
1) // Big and small window sizes expand synchroneusly
S,, B; (- 1-1t-1. S,
// Partition with big window
X; f- rearrange X;, Chat (Nh hb) (Nw wb) (Nd db)
-> M (Nh Nw Nd Chat) hb wb db" // Collecting salient tcken for each small window
X, + Fmmax.pool3d(X,, S,)
// Flatten spatial dims and concatenate multimodal sequences
X, rearrange
X;, "M (N Chat) nh nd -> Chat (M nh nw nd)
Xs xx.ap)end(X))
end
|/ Concatenate all wi ndow features
// X: ENyaon No, C,M.C.
X + torch.cat(X.... (xs, dim = 0)
end
/+ 3) multimodal Grouped Attention x1
A + multihead.attentionn( K, V)
/* 4) Paired Window Scattering
// Inverse of gathering.
// A: [MCCHH,,DD
A < window.scattering(A A)
/* 5). Paired Window Mixer
// A: [M,C, W, DJ
A E Dropout(PWC(A))
return A
18

---

Published as a conference paper at ICLR 2026
where N win is the number of window pairs:
[inn & Main
Ji= = (B,, S Nuin
1 1i=1
=
(1" h pd
h (1-
2
N
(5)
where r E N is the expansion rate (default r 2), Bh and S.¹ represent the big window and small window, respectively. lgg wg. dž represent the height, width, and depth of the big window; fg. w de correspond to the small window. Notably, r. Bf, and Sf are closely related to the computational cost, and the specific settings for different datasets given in Appendix C.3 Since B is equal N
to (H*, W D there is no need to set N win in advance. Subsequently, the encoder features will undergo the following steps in order:
C.2.1 LINEAR PROJECTION
We do not follow the habit of linear mapping: the number of output channels is equal to the number of input channels, but based the JL-guided minimum head size C the num- ber of heads N and the number of window pairs N N the number of output channels is No Č* min (nCh minnn N N win head (nC 2 C**. where Ö* is the actual head size. The formula is as follows:
PWC (LN (E..) K PWC (LN (E*,)) V PWC (LN (E*)),
(6)
where Q. K, V represent query, key, and value respectively, LN (+) represents layer normalization, and PWC (.) represents woint-wise convolution. For convenience, we use X to represent Q, K, V in Algorithm and Figure
C.2.2 PAIRED WINDOW GATHERING
Synchronously expanding paired windows ensures that the sequence lengths of Q K and V main consistent across multiple scales, which enables parallel computation. Given Q K Vk E RC #** x where i denotes the i-th paired window, the processing pipeline is follows:
Channel Split: This operation assigns features their corresponding windows along the channel dimension. The feature shape becomes (M x Ó* x H² x W x D*).
Big Window Partition: This operation partitions the features into non- overlapping blocks based on the big window size B. The feature shape becomes
(nt, head? Óh bi, ri rd)..)). where nž is the total number of large windows.
Small Window Pooling:
This operation gathers salient tokens from each small window Sk. The feature shape becomes (ntt, N Ö , ,//h/// l/, / /www , dl. /
Flatten and Concatenation: These are feature reshaping operations. The new feature size is T Nhadd O+ MT.)). where = (ht (b)) ) (wt// w (dt) /d ) the sequence length and M is the number of modalities.
The above operations are repeated for each paired window. Thanks to the synchronous expansion, the sequence length L is guaranteed to be equal across different scales. Finally, we can summarize this process follows:
k Gather
Q(ir): R² Gather (K),
K)) V
Gather (Vi,..., V
(7)
C.2.3 MULTIMODAL GROUPED ATTENTION
The similarity matrix calculation formula as follows:
S
@*))
& K
(8)
Vo
19

---

Published as a conference paper at ICLR 2026
where represents matrix multiplication. In addition, each L L block in the similarity matrix is assigned a relative position code E to strengthen the position relationship between voxels in the window and cross-modal voxels. The remaining attention is calculated follows:
w softmax (S* +E
A W VA w is the attention weight matrix, and A k is the attention obtained for each window.
(10)
C.2.4 PAIRED WINDOW SCATTERING
After computing the attention mechanism in parallel, perform the inverse operation of Paired Window Gathering to map the multi-scale attention mechanism to the original feature space, ob- taining the window attention A for each modality, which has the same size Q K& V
A A M Scatter(A")
(11)
C.2.5 PAIRED WINDOW MIXER
The above operations obtain window attention of different scales. We will use x 1 x convolution to mix them to get the final feature E The formula is follows:
E n2
E
PWC (A m 1, . M.
(12)
C.3 COMPUTATIONAL COMPLEXITY
Let N H-W. D, B hy wy de h wy das and 1 1 the computational complexity of PWA is calculated as follows:
midn
N B
1
B
B² 53
4
22
2N
mumber of big windows
multiplication operations per big window
B No 4 1 0a +
55
NK 102 + 0
B
D DETAILS OF JL-GUIDED GROUP SIZE
D.1 EMPIRICAL PARAMETRIZATION OF COVERING NUMBERS
Motivated by the classical covering number results in Heinonen 2001 Definition 10.15 and Exer- cise 10.17), we consider a hypothesis class whose covering number satisfies
N(e) s c(+)
C21, e (0,
(13)
where > is a constant independent of the ball, e E (0, 1). To ensure generality, we do not make further assumptions about the data to obtain specific parameters, but instead empirical functions for approximation: N- (M v)
M o replaces 1 1/e, representing the coverage density required per dimension. Given the constraints of JL's logarithmic scaling and the requirement that the group size divides the input channel, we omit the possible constant term.
serves difficulty coefficient reflecting the dataset's 's intrinsic complexity. We calibrate a based on the most challenging dataset to ensure robust generalization across different tasks.
20

---

Published as a conference paper at ICLR 2026
D.2 WHY WE AVOID MORE COMPLEX POLYNOMIAL APPROXIMATIONS.
The covering-number estimate is only an intermediate step; JL lemma then applies a logarithm. Consequently, low-degree and constant terms in polynomial approximation have minimal effect after taking the log. For example, consider = 43, M = 2, a = If we add a constant term 3, (M + 3) (M log(11 + // 128) O. Besides, because the group_size must divide input channel, such a small interrupt rarely changes the final group size.
D.3 ANALYSIS OF DIFFERENT NUMBER OF MODALITIES AND THE GROUP SIZE OF 2D IMAGES
For typical lightweight 3D medical segmentation method, its network contains M E {1,2,4) modalities with volume ratios { { { j = {4", 88, 16,,32 each stage. Its complexity increases k=1
with depth, which means that the group size is:
M 1 {4.9a, 6.9a, 9.0a, 11.1a}, M = 2.
(4.2a, 6.2a, 8.3a, 10.4a).
Ggroup 2
(14)
(5.5a, 7.6a, 9.7a, 11.8a). M = 4
Considering that the group size needs to be divisible by the total number of channels and that nonlin- 1
car networks have stronger compression capabilities than linear networks, we use {C& group j=m {4a, Sox, Sa, 16a) for each stage, where is determined by the most difficult AutoPET-II dataset to ensure universality. For convenience, we replace a with n = [a/4]] E N, and the final convolution group size of each stage of the network is set to {C'rroup f k=1 {n, 2n, 2n, 4n).
For lightweight convolution settings in the natural image domain, the input image typically has M 3 channels, and the volume ratio of each stage of the network is {u*) k= 112,22, 02,,,42,55 82),
which means the group size is: (Croup) f k=1
log 3, log 12, a log 48, log 192) s (1.1a, 2.5c 3.9a, 5.3a). Considering the integer divisibility of the channels, it is recommended to use group size of {a, 2ax 4a, 4ar. This setting is similar to the depth-wise convolution setting, further demonstrating the effectiveness of wise convolution in the natural image domain.
E DETAILS OF LOsS FUNCTION
E SEGMENTATION Loss
For segmentation, we use a combination of the cross entropy loss C.e. and the foreground dice loss faiiee which can optimize the detail and global segmentation effects. Deep supervision is performed on the segmentation decoder. The formula is as follows:
HWD
Lee (P,Y)
1
Σ P log (Y,),
(15)
HWD
2 > HWD P.Y.
Ldice (P,Y)=1-
(16)
Σ
HWD P,
E Y;
where P is prediction map and Y is segmentation ground truth, subscript represents the i-th voxel.
E.2 RECONSTRUCTION Loss
The texture teacher learns without data annotation, reconstructing the original input image based on model features. The loss function is simple mean squared error, as shown the following formula:
Lre
M Σ R.,. ImIl
(17)
where M is the number of input modalities, R,, and I,, represent the reconstructed and original images of the m-th modality, respectively.
21

---

Published as a conference paper at ICLR 2026
Dataset
Modalities Region
Label Type
Image Size
Crop Size
Voxel Spacing
Malignant melanoma,
Whole
Min
Fixed (2.036, 2.036, 3)
AutoPET-II
PET, CT
lymphoma. lung
(96, 96, 96) (400, 400, 200)
Body
cancer lesions
Primary gross volume (GTVp), or lymph node gross (128, 128, 67)
Head &
Min
Median (0.98, 0.98, 3.3)
Hecktor2022 PET, CT
(128, 128, 64)
Neck
tumor volume (GTVn)
Brain tumors: whole tumor (WT), tumor Fixed core (TC), enhancing (240, 240, 155)
Fixed 1.0)
BraTS2021
MRI
Brain
(96,96,9)
tumor (ET) subregions.
Glioma segmentation
BraTS 2016 TCIA
(multi-class): necro-
Fixed tumor and (240, 240, 155)
Fixed (1.0, 1.0)
MRI
Brain
(96, (86,96,90 96)
sis/active
edema.
Table 5: Details of AutoPET-II, Hecktor2022, BraTS2021, and BraTS TCIA datasets. If image size is a variable, the minimum value is reported. If voxel spacing is variable, the median value is reported.
DETAILS OF DATASET
We evaluate our proposed VeloxSeg on four public medical image datasets: AutoPET-II Ga-
fidis S 2022). Hecktor2022 Oreiller al. 2022). BraTS2021 (Baid et al., 2021). and BraTS 2016 TCIA (Menze et al., 2014a). The AutoPET-II and Hecktor2022 datasets are multimodal PET/CT datasets for tumor segmentation. AutoPET-II contains 1,014 whole-body PET/CT seans with vari- able image dimensions and is cropped to 96 x 96 x 96 patches. Hecktor2022 comprises 524 head and neck PET/CT scans, cropped to 128 x 128 x 64 patches. The BraTS2021 and BraTS 2016 TCIA datasets multimodal (TI, TIce, T2. FLAIR) MRI datasets for brain tumor segmentation. Each patient's data registered to a common spatial resolution of 240 x 240 x 155 and undergoes skull stripping. The brain tumor region is segmented into three primary sub-regions: the enhancing tumor (ET), the tumor core (TC), and the whole tumor (WT). For training efficiency, volumes are cropped 96 x 96 x 96 patches. BraTS2021 contains 1,251 cases for training and validation, while BraTS 2016 TCIA (244 cases) serves as an external test set to evaluate domain generalization capability across different data distributions.
We use four public medical image datasets to verify the effectiveness of VeloxSeg, including AutoPET-II, Hecktor2022, Bra21 and BraTS 2016 TCIA, which used as an external test set to compare the generalization ability of the model. The first two datasets contain CT and PET images, and the latter two datasets contain 3D MRI images with four modalities. The details of the datasets are described in T5
G DETAILS OF COMPUTATIONAL PERFORMANCE
We evaluate the computational performance of VeloxSeg against other leading models on the AutoPET-II, Hecktor2022, and BraTS2001 datasets. Our analysis focus on four key metrics: the number of model parameters in millions, GFLOPs, GPU throughput, and CPU throughput. On the AutoPET-II and Hecktor2022 datasets, VeloxSeg established a new standard for efficiency. As de- tailed in Table 6 our model operates with only 1.66 million parameters and the lowest GFLOPs among all competitors, requiring just 1.79 on AutoPET-II and 2.13 on Hecktor2022. This lean profile translates to exceptional speed, where VeloxSeg recorded the highest GPU throughput and second- highest CPU throughput on both datasets. In the lightweight category, while Slim UNETR marginally smaller, VeloxSeg surpasses it in computational cost and processing speed.
22

---

Published as a conference paper at ICLR 2026
AutoPET-II
Hccktor2022
Methods
Type
MP.
GE
ThrG.
ThrC.
MP.
GF.
ThrG.
ThrC.
Basic Models
UNet
CNN
5.75
136.56
101 001.04
0.23
5.75
161.84
85.60
0.19
VNet
CNN
45.60
322.22
58.99
0.14
45.60
381.89
49.82
0.11
MedNeXt-S
CNN
5.54
57.93
27.95
0.06
5.54
68.54
23.20
0.05
UNETR
CNN-Transformer
95.76
83.61
131.6
0.40
95.76
99.09
105.78
0.35
Swin UNETR
CNN-Transformer
15.51
84.26
38.37
0.14
15.51
100.66
28.58
0.10
VSmTrans
Crn-Transf
12.48
91.44
36.56
0.14
3.12
28.79
34.13
0.16
+
TNN-Transformer
19.97
57.93
161.15
0.67
19.97
68.66
138.39
0.56
U-KAN
CNN-KAN
7.06
22.90
187.06
0.82
7.06
27.13
159.92
0.75
Multimodal Models
Nestedformer
CNN-Transformer
4.71
58.62
95.63
0.41
4.71
69.48
79.05
0.35
A2FSeg
CNN
41.32
207.97
52.02
0.17
41.32
246.48
40.60
0.13
H- D--enseFFrrmr
CNN-Transformer
3.64
71.91
123.35
0.44
3.64
85.23
102.80
0.37
Lightweight Models
SegFormer-3D
CNN-Transformer
4.50
5.11
364.24
3.31
4.50
6,06
305.01
2.67
Slim UNETR
CNN- Tranffonsfrrrer
1.77
3.83
178.33
11.40
1.77
4.53
151.85
8.78
SuperLigghNNet
CNN-Transformer
2.75
19.42
55.48
0.27
2.75
23.01
47.36
0.23
HCMA-UNet
CNN-Mamba
2.81
26.15
54.51
2.81
31.00
46.48
U-RWKV
CNN-RWKV
1.44
20.68
82,09
1.44
24.51
73.98
VeloxSeg
TNN-Transformer
1.66
1.79
390.91
6.67
1.66
2.13
319.80
5.47
Table 6: Computational performance comparison of all models on AutoPET-II and Hecktor2022
datasets. "MP.": Million Parameters; "GF": GFLOPs; "ThrG.": Throughput on GPU; "ThrC."
Throughput on CPU,
On the BraTS2021 dataset, we test early-fusion VeloxSeg due to its concentrated target distribution. absence of small lesions, and low modality heterogeneity. Table shows that VeloxSeg-C is one of the smallest models with only 1.46 million parameters, yet it achieves the lowest GFLOPs at 2.64.
Most notably, it delivered the highest GPU throughput of any model, processing 536.62 images/s, alongside the second-fastes CPU throughput. This positions VeloxSeg-C faster alternative to other lightweight models like U-RWKV and SegFormer-3D.
more efficient and
Across all three benchmarks, the VeloxSeg architecture demonstrates an excellent balance between model size, computational requirements, and processing speed, making it well-suited for deploy- ment in resource-constrained environments. Furthermore, segmentation methods based sequence models, such as Mamba and RWKV, lack CPU support, significantly limiting their application in edge devices.
H DETAILS OF GPU MEMORY USAGE
Table g shows that VeloxSeg has the second lowest memory footprint, saving more GPU memory than all non-lightweight baseline models. As shown in Figure II on the AutoPET-II dataset, the base methods' GPU memory usage is 2.5 to 20 times that of VeloxSeg, with inference memory usage reaching up to 24 times higher. Compared to other lightweight models, VeloxSeg consistently has less GPU memory usage than SegFormer-3D 3D, SuperLightNet, HCMA-UNes and U-RWKV, which, despite claiming be lightweight, have GPU memory usage that is 5.9 to 17 times higher than ours.
23

---

Published as a conference paper at ICLR 2026
Methods
Type
MP.. 4
GF.
ThrG.
ThrC. 1
Basic Models
UNet
CNN
5.75
138. 14
103.70
0.23
VNet
CNN
45.61
322.85
63.37
0.14
MedNeXt-S
CNN
5.54
57.95
27.75
0,06
UNETR
Tra-Trrnsformer
102.06
85.79
128.16
0.41
Swin UNETR
Transforser
15.51
85.53
38.11
0.13
VSmTTrans
CNN Transformer
12.48
92.72
36.18
0.13
UNETR++
UNN-Transformer
19.98
58.81
153.63
0.52
U-KAN
CNN-KAN
7.06
23.69
181.89
0.85
Multimodal Models
Nestedformer
CNN-Transformer
88.43
56.92
0.23
A2FSeg
CNN
74.55
361.18
28.92
0.08
H-DenseFormer CNN-Transformer
5.39
73.18
95.61
0.42
Lightweight Models
SegFormer-3D
Tranffonsfrrmer
4.53
5.42
355.73
2.98
Slim UNETR
CNN- Transformer
1.78
6.59
97.50
10.62
SuperL.ightNet
CNN- Tranfforsformer
2.75
19.54
55.13
0.28
HCMAA-NNet
CNN-Mamba
2.81
26.69
53.72
U-RWKV
CNN-RWKV
1.43
21.08
83.15
VeloxSeg-C
CNN-Transformer
1.46
536.62
5.23
Table 7: Computational performance BraTS2021 dataset with patch size 96 x 96 96 and modalities (T//TIc//2//LLIIR)) "MP.": Million Parameters; "GF": GFLOPs; "ThrG.": Through- put on GPU; "ThrC.": Throughput CPU.
VolonSeg (Reference)
Lightweight Madici
Muttimedal
Basic Models
Truining Peak GPU Menory (MB)
17500
15000
12500
19000
7500
5.4%
4.99
6.0x
3.77x
4,03*
2500
2.70K
844
4000
Inference Peak GPU Memory (MD)
24x
3500
17x
15H
3500
9.7x
M.3x
1000
5.9x
3.67x
500
2.80k
152
Figure 11: Visualization of memory usage for training and inference all models. Arranged from left right in reverse order of T
24

---

Published as a conference paper at ICLR 2026
AutoPET-II
Hecktor2022
BraTS2021
Methods
Тype
TM.
IM.
TM.
IM.
TM.
IM.
Basic Models
UNet
CNN
5054
2268
5942
2698
5112
2090
VNet
CNN
3820
1762
4460
2072
3886
1684
MedNeXt
CNN
17216
3678
20372
4376
17258
3734
UNETR
CNN-Transformer
4114
1788
4626
1914
4106
1460
Swin UNETR
CNN-Transformer
11164
2616
13258
3704
11 166
3466
VSmTrans
-ra-sfrsforemr
13856
1956
13318
2532
14202
1972
+NETR+
TrN-Transformer
3392
940
3906
1088
3640
984
U-KAN
CNN-KAN
2138
558
2360
640
2200
606
Multimodal Models
NestedFormer
CNN-Transformer
4428
1658
5182
1966
7274
2200
A2FSeg
CNN
7748
2004
9102
2352
15604
3074
H- DenseFormer
CNN- Tra-Tsanforrmer
3172
1256
3778
1474
3712
1094
Lightweight Models
SegFormer-3D
CNN- Transforrrrrm
2272
426
2702
436
2096
430
Slim UNETR
CNN- Transformer
488
106
562
124
602
172
SuperLightNet
ONN-Transfrrmer
7812
2510
9192
2730
7842
2266
HCMA-UNet
ONN-Mamba
7730
1480
9098
1892
7938
2020
U-RWKY
CNN-RWKV
4582
900
5388
1044
4670
936
VeloxSeg
Trannrrasfrrmer
842
152
1006
178
1392
1112
Table Peak GPU memory usage for training and inference for all models across three datasets. models were tested with fixed batch size of 2, ensuring all other experimental conditions remained the same. "TM." represents the peak GPU memory usage during training, and "IM." represents the peak GPU memory usage during inference.
Dataset
Model
MParams J. GFL.OPs
Thr.GPU 1 Thr.PPU
Dice t HD95
nnUNet
88.62
3078.83
81.13
0.127
55.85
193.54
AutoPET-II
VeloxSeg
1.66
1.79
390.91
6.67
70.05
177.51
nnUNet
88.62
4828.04
68.02
0.106
60.80
36.67
Hecktor2022
VeloxSeg
1.66
2.13
319.80
5.47
62.51
30.22
Table 9; Performance comparison between nnUNet and VeloxSeg across PET/CT datasets. Both segmentation performance and computational efficiency are evaluated.
RESULTS ON THE NNUNET TRAINING FRAMEWORK
To unleash more model potential, we placed the model in the nnUNet training framework and com- pleted the training while keeping the patch size consistent with the experimental setting. The results are shown Table can be seen that our model has achieved comprehensive transcendence. In the AutoPET-II dataset, we achieved Dice that was 14.2% higher than the nnUNet baseline with 1.87% of the parameters and 5.81e-25% GLOPs, and the GPU throughput and CPU throughput in- creased by 4.8x and 52.5 respectively. Similarly, in the Hecktor2022 dataset, VeloxSeg achieved a Dice that was 1.71 % higher than the nnUNet baseline with 1.87% of the parameters and 4.41e-2% GLOPs, and the GPU throughput and CPU throughput increased by 4.7x and 51. 1,6xx respectively.
25

---

Published as a conference paper at ICLR 2026
AutoPE
UNot
H: D.DeceFnnce
Shi UNETR
Supert Sg
Hleckorr 2012
Swie LNE.TR
Natrddaamee
A213e2
Seghemer-D
Site UNITT
SaperLighiNct
HUMAA-UN
U-RWKV
Bra 2021
UNo
Vihrt
Notedianme
AZESEZ
Sept METAAF-1
Sina
SapwrEigh/ee
Figure 12: 3D qualitative visualization of different methods the AutoPET-II, Hecktor2022 and BraTS2021 datasets. In PET-CT datasets, model predictions are shown on CT images (red indicates true positives, yellow indicates false positives, and blue indicates false negatives); In IsIs datasets, model predictions are shown on T1CE images (red represents ET, yellow represents TC, and green represents WT).
J QUALITATIVE RESULTS
Figure show the qualitative results three segmentation cases: large melanoma lesions, primary and secondary lesions of right tonsil tumors or lymphomas, and glioma lesions. can be that our method can accurately locate the tumor area, exclude the wrong attention to areas such as intracranial veins, and the prediction results are consistent with the labels.
K MODALITY ADAPTATION EVALUATION
On BraTS2021 MRI brain tumor dataset, we use an early fusion strategy (VeloxSeg-C) that does slow down as the number modalities increases, shown in Table IO Since brain tumors are large and centralized, and the slices processed are relatively fixed, almost all models achieved good results. Our model improves the Dice by 1.72% compared the state-of-the-art SuperLighttNe..
HYPERPARAMETER ANALYSIS
L.1 MODULE HYPERPARAMETER
Model optimization, detailed in Table focused on balancing segmentation performance and com- putational efficiency. convolution adjustments, including reducing model width (32, 61, 128, 256) to (16.32 32, 64, 128), improved CPU throughput, from 10.83 to 20.23, and Dice, from 48.96% to 10%. Replacing large kernel convolution (7) with parallel small kernels (1,3,5) 5) yielded a 3.55%
Dice increase, from 50.10% to 53.65%, while simultaneously reducing MParams from 0.73 to 0.66,
and GFLOPs from 2.41 to 2.30. Optimal group channel setting (4, 8, 16) achieved 55.14%
Dice. Attention depth reduction (2. 2, to (1, 1, surprisingly enhanced Dice. from 59.56%
to 61.03%. As suggested Table II reducing the FFN dilation rate of Transformer/Conver (3, 2,,2,2,2 can slightly improve Dice performance while reducing computational cost.
to
26

---

Published as a conference paper at ICLR 2026
Dice Similarity Coefficient (%)
Hausdorff Distance 95% (mm)
Method
Avg.
ET
TC
WT
Avg. +
ET + TC
WT
UNet
88.18
89.62
85.65
89.28
4.93
5.95
7.52
3.59
V-Net
88.86
90.66
86.16
89.75
5.26
5.98
8.29
4.31
MedNeXt-S
90.70
92.64
88.33
91.12
4.48
4.72
7.10
3.33
UNETR
85.44
88.12
81.49
86.71
6.68
8.51
8.19
4.75
Swin UNETR
88.52
90.19
85.73
89.63
5.07
6.37
7.29
3.44
VSmTrans
86.62
91.15
78.01
90.71
7.00
6.00
12.07
3.44
UNETR+
88.77
90.26
87.01
89.05
4.49
5.30
6.98
3.76
U-KAN
88.51
90.57
86.14
88.82
5.44
5.77
8.03
4.65
Nestedformer
88.54
89.60
86.71
89.30
4.21
5.44
7.38
3.13
A2FSeg
88.18
91.78
84.58
88.18
4.66
4.47
7.51
3.72
H--DenseFormer
89.35
90.80
86.66
90.59
5.58
5.85
8.88
3.97
SegFormer-3D
89.18
90.37
87.49
89.69
6.30
3.57
Slim UNETR
87.33
89.31
85.00
87.66
5.16
6.46
7.57
3.52
SuperLightNet
89.72
91.46
87.22
90.48
4.46
5.39
6.33
3.06
HCMA-UNet
89.53
91.63
86.24
90.72
4.79
5.05
8.12
4.15
U-RWKV
89.04
91.34
86.94
88.83
5.42
6.39
7.66
4.63
VeloxSeg-C
93.09
89.00
92.24
4.41
3.35
Table 10: Segmentation performance comparison the BraTS2021 dataset. VeloxSeg- metrics are highlighted in green. The best performance is red and the second best performance is blue.
Ablation
Are
Asdit
Dice
t
1,5
1.5
58.03 -4,41
Depth
CT
PET
Enc.
Dec.
Lre
1,0
1.5
61.53
-0.92
3.22
3.06
0.5
62.44 1.5
2.85
2.82
2.73
2.75
0.5
2.5
62.23
-0.28
2.37
2.38
2.26
2.29
0,5
2.0
62.51
4
1.82
1.85
1.74
Latgle
0.5
1.5
62.44
-0.07
0.5
1.0
61.66
-0.85
Table 11: Ratio of channels to embeddings after pruning the FFN layers of PWA and JLC.
0.5
0.5
60.55 -2.06
Baseline Dice: 69, 94%; after pruning: 68.49%.
Table 12: Hyperparameters experiments with loss weight on AutoPET-III
L.2 Loss WEIGHT
Considering the differences in the contributions of various tasks to segmentation, need adjust the relative weights between different tasks to explore the optimal parameter update process. To this end, we adopt the strategy of controlling variables and adjust the loss weights of Lre and Estyle in turn. The specific results are shown in Table 12 The final parameters each experiment are highlighted in green and red and green used to indicate the improvement and deterioration in process. Finally, the optimal weight parameters are selected as Are = 0.5, Aadit = 2.0.
M COMPUTATION OF MEAN ATTENTION DISTANCE
To analyze the locality of attention heads, we compute the Mean Attention Distance (MAD), a metric that measures the average physical distance between query voxel and the key voxels it attends to, weighted by the attention scores. We extend the Mean Attention Distance metric to 3D volumes to analyze attention patterns in volumetric data, and its computation is detailed below.
27

---

Published as a conference paper at ICLR 2026
Modality
MParams +
GFLOPs
Thr. (GPU)
Dice
CT
1.39 -0.27
1.61 -0.18
768.48
+377.57
21.43 -41.01 -13.16
PET
1.39 -0.27
1.61 -0.18
768.48 48 +377.57
PET + CT
1.39 -0.27
1.70 -0.09
694.76 +303.85
56.69 -5.75
(CT. PET)
1.66
1.79
390.91
62.51
Table 13: Modality ablation experiments performed AutoPET-II. "PET+CT" indicates early fusion strategy, and "(CT, PET)" indicates consideration of modality interaction.
Let the input 3D volume be partitioned into a grid of H x W D voxels, resulting in a total of H W voxels. The attention weight matrix for a given head is denoted by W RR where Waj is the attention weight from voxel i to voxel j. Let s e R+ be scalar representing the physical edge length of a single cubic voxel.
First, we map each voxel's flattened ID index i (where i (O...., L 1}) to its physical 3D coordinates (z4, Ve, 2,). Assuming a standard row-major flattening order where the x-axis (width) the fastest-changing dimension and the z-axis (depth) is the slowest:
2; [ (H W)]
1/s = [a mod (H WW))/W
(18)
2, mod W
Next, we compute the pairwise physical distance matrix D E RLXL, Each element Dij represents the scaled Fuclidean distance between the centers of voxel and voxel j:
Dij = V(u: 22)2 (ys 2 + (z; 2,)
(19)
Finally, the Mean Attention Distance (MAD) is formulated as the expectation of the physical dis- tance over the attention distribution. As shown Equation 20 it is computed by summing all dis- tances weighted by their corresponding attention scores, and then averaging over all query voxels:
1 1--1--1
MAD = - 1=0 3=
- Wij
(20)
smaller MAD value indicates that the attention head primarily focuses on local information, whereas larger MAD value signifies more global attention pattern across the volume. The win- dow size of each stage of PWA is shown in Appendix C
N PWA MULTIMODAL EVALUATION
To verify the effectiveness of PWA in heterogeneous modal modeling, we test various inputs, as shown in Table 13 Using only PET CT reduces model size and complexity but sacrifices seg- mentation performance. An early fusion strategy achieves a Dice 56.69%, outperforming the pure convolution framework's 's 55.84%. Crucially, introducing modal interaction in PWA improves Dice by 5.75%, significantly improving performance robustness without significantly increasing compu- tational or time costs.
O JL-SETTING GENERALIZATION EVALUATION
Results demonstrate the effectiveness of JL-guided group size configurations (4, 8, 16) on various datasets. While the smallest configuration (1, 1, 1, 1) achieves the lowest computational cost, reduc- ing segmentation performance by 0.091 MParams and 0.063 to 0.075 GFLOPs, it significantly de- grades segmentation performance, particularly on the Hecktor2022 dataset, which features heteroge- neous modality data and cross-organ distribution of targets, where the Dice drops by 6.25%. Larger configurations (16, 32, 64, 128) only slightly improve the Dice (by 1.16% on BraTS2021) but signif- icantly increase computational complexity by 0.633 MParams and 0.268 to 0.317 GFLOPs. This periment further demonstrates that JL-guided configurations strike an optimal balance, maintaining competitive performance while ensuring computational efficiency suitable for clinical deployment.
28

---

Published as a conference paper at ICLR 2026
Dataset
Configuration
MParams 4
GFLOPs
Dice t
11,1,1,1)
0.618 -0.091
2.637 -0.075
37.95-6.25 -6.25
Hecktor2022
(4, 8, 16)
0.709
2.712
44,20
(16, 32, 64, 128)
1.342 +0.633
3.029 +0.317
43.21 -0.99
11,1,
0.629 -0.091
2.377
-0.063
85.82 -1.00
BraTS2021
(4, 8, 16)
0.720
2.440
86.82
(16,32, 32, 64, 128) 1.353 +0.633
2.708 +0.268
87.98 +1.16
Table 14: Performance comparison of different group size configurations across datasets. The JL- guided configuration (4, 16) used the reference baseline.
Methods
MParams
GFLOPs
Tr. Mem.
Inf. Mem.
Thr. GPU
Dice
Window
1.51 -0.10
2,80 -0.04
678 -46
1066 +918
227 .28 +45.83
56.01 -5.42
Downsample 1.52 -0.09
2.78 -0.06
1066 +342
136
239.01 +57.56
55.18 -6.25
PWA
1.61
2.84
724
148
181.45
61.43
Table 15: Computation consumption of different attention variants.
P COMPARISON OF DIFFERENT ATTENTION MECHANISMS
Under similar computational cost constraints, we replaced PWA with other attention mechanisms, such as window-bases multimodal attention and downsampling- shown in Figure 15] Although our method is not optimal due to the larger tensor size change rate, the model performance is significantly better than the other two attention mechanisms, which further validates the effectiveness of PWA for heterogeneous modality modeling.
multimodal attention, as
Q COMPARISON OF DIFFERENT KNOWLEDGE TRANSFER STRATEGIES
The comparison results with other strategies are listed in Table with settings largely consistent with the dual-stream settings in reference (Qiu al., 2023). All comparison methods were per- formed under the same conditions. The additional training overhead is listed in Table [6 where SDKT uses only about 100 MB more memory than the baseline methods.
R SCALING LAW OF VELOXSEG
Our specific parameter configuration as follows:
S represents changing the convolution kernel from [1,3,5] to [3] under the original param- eter configuration.
B represents the original parameter configuration.
B+ represents scaling up the number attention and convolution channels from 16
L represents scaling up the depth of attention and convolution from to 2.
In VeloxSeg, the Dice value is directly proportional to the number of parameters/floating point oper- ations: for versions to B, the Dice value increases by 1.5 for every 0.47M additional parameters; for versions B to B++. the Dice value increases by 1.5 for every 360M additional parameters; while for version I., with fewer aarametersffoating-poii operations than B++, the Dice value only increases by 0.55. This indicates diminishing returns and that architectural adjustments (not just scaling) key to improving performance. Non-monotonic resource ordering (L version is smaller than B++ version) results roughly equal Dice values.
29

---

Published as a conference paper at ICLR 2026
Methods
MParams
GFLOPs
Tr. Mem.
Inf. Mem.
Thr. GPU
Dice
w/o Teacher
1.61
2.84
724
148
181.45
61.43
Teacher
1.66 +0.05
1.79-1.05
824 +100
152 +4
390.91 +209.46
59.71 -1.72
1.66 +0.05
1.79-1.05
824 +100
152 +4
390.91 +209.46
1.67 -59.76
Affinity
1.66 +0.05
1.79-1.05
894 +170
152 +4
390.91
+209.46
41.44 -99.99
+ Shared ROI
1.66 +0.05
1.79-1.05
1064 +340
152 +4
390.91 +209.46
57.15 -4.28
SDKT
1.66 +0.05
1.79 -1.05
842 +118
152 +44
390.91 +209.46
62.51 +1.08
Table 16: Computation consumption of different knowledge transfer methods.
Model
Dice
Parameters (M)
FLOPS (G)
nnUNet
55.85
88.62
3078.83
VeloxSeg S
68.56
1.19
1.41
VeloxSeg B
70.05
1.66
1.79
VeloxSeg B+
71.56
5.26
4.27
VeloxSeg L 72.11
2.65
2.45
Table 17: Accuracy results of VeloxSeg after increasing model size.
30


Reviews:
[{"review_id": "glD2q7Nx84", "reviewer": "Reviewer_kwQp", "summary": "The authors propose VeloxSeg, a dual stream convolution and transformed based medical image segmentation method. The authors introduce a Johnson-Lindenstrauss theorem guided setting for grouped convolutions for extracting local features at minimized cost, a Paired Window Attention block for extracting local and global features across modalities and a Spatially Decoupled Knowledge Transfer strategy for learning rich textual features. They demonstrate superior performance against 19 baselines on 2 datasets, alongside ablation experiments and analysis of their model performance.", "questions": "1. I recommend that the authors consider restructuring the analogy in lines 72-90 in the introduction. It is rather cumbersome to picture and a simple scientific description would be better. I would also request that they take out references to case, clues, detectives and profilers throughout the text. It is difficult to read in my opinion.\n\n2. Comparisons against architectures such as ResEncL, STUNet and MedNeXt-L are prescient especially given the fact that the paper clearly states that the JL Convolution proposed is designed to counter - as stated by the authors - the phenomenon of “depth-wise convolution destroys the adjacency relationship between data in the feature space” which architectures like MedNeXt are based on. This would also address my concerns in the above weaknesses section regarding depthwise convs.\n\n3. Please investigate my concerns regarding the possible unintentional data leakage in the domain generalization experiments. It is acceptable to simply drop the domain generalization experiment if there is or find a replacement dataset.\n\n4. As the authors state that SKDT helps with rich texture details in feature learning, maybe an idea of HD95 in Table 2 might help their case.\n\n5. The authors need to explain PWA somewhat more in the text in a better way. It is difficult to understand as it currently stands.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. The analysis of the authors on their proposed methods and against their 19 baselines is impressive, as is their ablation experiments.\n2. The authors use volumetric and surface based metrics to demonstrate the efficacy of their methods.\n3. The authors compare against a large number of baselines including categories of models such as multimodal and lightweight methods.\n4. The use of Johnson-Lindenstrauss theorem to add grouped convolutions instead of depthwise convolutions is an interesting approach.\n5. The SKDT supervision is a novel way to add to their method.", "weaknesses": "1. With research papers on advanced architectures, the results are significantly more convincing when compared against equivalent powerful architectures in the field, such as nnUNet ResEncL [4], STU-Net [2], MedNeXt-L [3] which have been compared in public benchmarks as described in [1, 4]. Architectures like MedNeXt are also dependent on depthwise convolutions, which the authors should demonstrate struggle to perform against their JL Convs.\n\n2. The Paired Window Attention is rather difficult to understand in the main text - a bit more explanation would have been useful. This includes while looking at the code in the appendix. The idea of modal features is also somewhat confusing - it is unclear if the authors mean features from different imaging modalities - It should also then be made clear to the readers that their modes range from 2-4 given their choice of datasets.\n\n3. The clever application of the Johnson-Lindenstrauss theorem is appreciated. However, I am not completely convinced of the idea of depthwise convs completely destroying spatial adjacency in feature learning. I agree that they would do so in principle if they were ever applied completely in isolation. But in execution, they are likely followed by a layer of pointwise convs. This is true for some of the authors’ own references - Shufflenet V2 and MedNeXt.\n\n4. I have some concerns regarding the domain generalization experiments by training on BraTS2021 and evaluating on MSD01 Brain Tumor. BraTS2021 reuses data from older BraTS cohorts while MSD01 reuses subsets of BraTS16 and 17. While I could not directly trace sources given the long chain of data reuse in BraTS, there is some concern that both datasets used by the authors might share old cohorts of BraTS 2012 or 2013 data. \n\nReferences:\n\n[1] Bassi, Pedro RAS, et al. \"Touchstone benchmark: Are we on the right way for evaluating ai algorithms for medical segmentation?.\" NeurIPS, 2024.\n\n[2] Huang, Ziyan, et al. \"Stu-net: Scalable and transferable medical image segmentation models empowered by large-scale supervised pre-training.\" arXiv preprint arXiv:2304.06716 (2023)..\n\n[3] Roy, Saikat, et al. \"Mednext: transformer-driven scaling of convnets for medical image segmentation.\" MICCAI, 2023.\n\n[4] Isensee, Fabian, et al. \"nnu-net revisited: A call for rigorous validation in 3d medical image segmentation.\" MICCAI, 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Rlbla4SitR", "reviewer": "Reviewer_syAZ", "summary": "This paper proposes an efficient 3D medical image segmentation framework called VeloxSeg. The framework consists of three key components working together: 1) a theory-driven group convolution (JLC) inspired by the Johnson-Lindenstrauss (JL) lemma, serving as the encoder backbone to efficiently compress and encode features; 2) a progressive window attention (PWA) mechanism in the decoder to capture multi-scale contextual information; and 3) a size-aware dynamic knowledge transfer strategy (SDKT) to improve the model's segmentation performance for lesions of different sizes through knowledge distillation. Experimental results on public datasets such as AutoPET-II and Hecktor2022 demonstrate that the proposed method achieves competitive segmentation performance while maintaining high efficiency.", "questions": "### On the Design of JLC\n\nCould you elaborate on the reasoning behind choosing the approximation  \n$ N(M, v) \\approx (M \\cdot v)^{\\alpha} $?  \nDid you explore alternative functional forms or methods to estimate the complexity of the data manifold?  \nHow sensitive is the final model’s performance to the scaling factor **α** (or the base channel number n)?  \nA more detailed sensitivity analysis would greatly help readers understand the robustness of the JLC design.\n\n---\n\n### Further Clarification on JLC’s Contribution\n\nTo better distinguish the unique theoretical contribution of JLC, did you attempt to combine a more naive lightweight convolutional module, such as a standard group convolution with empirically chosen group sizes, or a depthwise separable convolution, with PWA and SDKT?  \nThis comparative experiment is crucial. It helps determine whether the key advantage truly lies in JLC’s theory-guided design principle itself, or whether any reasonably designed lightweight encoder could achieve comparable performance when integrated with strong modules like PWA and SDKT.\n\n---\n\nOthers please see the Weaknesses section for details.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Theoretical grounding for depthwise separable convolution: The use of the JL lemma to derive a lower bound on convolution group size is a notable attempt to bridge architectural design with theoretical guarantees, particularly in preserving adjacency relationships within high-dimensional feature space.\n\n- Comprehensive ablation and visualization analysis: The paper provides detailed visualizations that help validate the contribution of each proposed component, enhancing the overall interpretability and credibility of the approach.\n\n- Good performance-efficiency balance: The most significant advantage of this paper is that it achieves state-of-the-art (SOTA) or highly competitive performance on challenging 3D medical image segmentation tasks while maintaining low computational costs. This is especially important for resource-constrained scenarios in clinical applications.", "weaknesses": "### Mismatch Between Claimed Contributions and Experimental Evidence\n\nThe paper’s title and introduction highlight **JLC** as the core and primary contribution. However, based on the ablation study (Table 2), the impact of the JLC module does not appear to be the most significant. For instance, adding the PWA module on top of the JLC encoder brings a substantial performance boost (Dice score increases from around 55% to 61%). This suggests that the model’s success heavily relies on the strong capabilities of **PWA** and **SDKT**, as well as their synergistic interaction, rather than solely on the superiority of the JLC encoder. The current narrative may overemphasize the dominance of JLC. It is recommended that the authors reconsider and more accurately position the core contribution of the paper. Perhaps as the design of an efficient and cohesive system, rather than a single JLC module.\n\n---\n\n### Theoretical vs Empirical Gap\n\nWhile the theoretical motivation behind JLC is a notable strength, the leap from the theoretical boundary  \n$ C_{group} \\geq c_{JL} \\varepsilon^{-2} \\log N(M, v) $  \nto the empirical approximation $ N(M, v) \\approx (M \\cdot v)^{\\alpha} $ is substantial. The paper acknowledges this, but in practice, the final group size configuration (e.g., {n, 2n, 2n, 4n}) is still determined empirically by tuning α or n. This weakens the claim that the design is “fully theory-driven.” It might be more accurate to describe it as theory-inspired or theory-motivated.\n\n---\n\n### System Complexity and Research Depth\n\nVeloxSeg integrates three independent and non-trivial modules, making it a complex system. Although the ablation results demonstrate overall effectiveness, such complexity makes it difficult to fully disentangle and analyze the individual contributions of each module. For a venue like ICLR, which often values deep exploration of a single core idea, a paper that focuses exclusively on JLC—examining its representational properties, limitations, and task adaptability in depth—might be equally valuable. From a practical standpoint, this complexity could also pose challenges for researchers or practitioners attempting to adapt and fine-tune the model.\n\n---\n\n### Hyperparameter Sensitivity of the PWA Module\n\nThe design of the PWA module, particularly the choices of minimum/maximum window sizes and dilation rates, appears to be dataset-dependent (Appendix D.1 notes different setups for AutoPET-II and Hecktor2022). While such dependency is common in deep learning models, providing a sensitivity analysis on these hyperparameters and offering general configuration guidelines would greatly enhance the reproducibility and generalizability of the proposed method.\n\n---\n\n### Others\n\n- PWA: Choosing and aligning paired small and large windows is nontrivial, especially for 3D data. In real-world scenarios where modalities (e.g., PET/CT or MRI/CT) have different voxel spacings or mis-alignment, the spatial correspondence between large and small windows can become ambiguous. It remains unclear whether the proposed spatial window mechanism can effectively handle such cases.\n\n- Long-range understanding: The long-range dependency modeling capability of PWA should be further justified. Including attention visualizations would strengthen the understanding of how multi-modal information is integrated through attention mechanisms.\n\n- Depthwise convolution and adjacency: In Section 2.3, the authors state that “depthwise convolution destroys the adjacency relationship between data in the feature space” (Figure 3(b)). However, it is unclear whether this claim holds true in practice. Providing theoretical justification or visual evidence would help substantiate this assertion.\n\n- Efficiency and SDKT: The teacher training cost of the Spatially Decoupled Knowledge Transfer (SDKT) module remain concerns, as it requires an additional training phase for the teacher network.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hjRkS0Oh4W", "reviewer": "Reviewer_L51F", "summary": "The paper introduces VeloxSeg, a lightweight and theoretically grounded 3D medical segmentation framework designed to resolve the efficiency–robustness conflict in medical imaging. It integrates: Paired Window Attention (PWA) for efficient multi-scale context capture, Johnson–Lindenstrauss lemma-guided convolution (JLC) to preserve spatial adjacency while reducing parameters, and Spatially Decoupled Knowledge Transfer (SDKT) using Gram matrices for knowledge distillation without inference overhead. Evaluations on four multimodal datasets (AutoPET-II, Hecktor2022, BraTS2021, MSD2019) show superior Dice scores and computational efficiency compared to state-of-the-art models.", "questions": "See the Weaknesses. I am open to discussing this further during the rebuttal and will be happy to increase my score if my concerns are addressed.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Evaluations span multiple datasets and architectures (CNN, Transformer, KAN, RWKV, and Mamba hybrids). The paper includes comprehensive ablation studies, efficiency tests (GPU/CPU throughput), and shows consistent performance gains.\n\n- Strong theoretical grounding via the Johnson–Lindenstrauss lemma. The proposed JLC introduces a sound theoretical perspective into the design of lightweight segmentation networks.", "weaknesses": "- In medical image segmentation tasks, using a fixed train/validation/test split (e.g., 6:2:2) can lead to unstable or biased results, particularly due to typically small dataset sizes and high inter-sample variability. A more robust and widely accepted approach in this domain is k-fold cross-validation, which helps assess the model’s generalizability more reliably. It is unclear whether the evaluation results reported in Table 1 are obtained using multiple folds or a single split. If only a single split is used, the reported performance may not be representative and could be sensitive to the particular partitioning. I recommend clarifying this point and, if not already done, performing evaluation using multiple folds or reporting variance across different splits.\n\n- I noticed that the baseline model (nnU-Net) is, by default, trained for 1000 epochs on the AutoPET-II dataset, whereas in your paper, all models are trained for only 300 epochs. This raises concerns about whether all models were evaluated at convergence and whether the comparisons reflect fully optimized performance for each method.\n\n- The AMOS dataset is currently one of the most widely adopted benchmarks for large-scale 3D medical image segmentation, and it also includes multimodal data (CT and MRI). It is unclear why your experiments do not include evaluation on AMOS. I recommend that the authors justify their dataset choices and discuss whether the proposed method generalizes well to other benchmarks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3mYXHdoXYW", "reviewer": "Reviewer_PE39", "summary": "The paper introduces VeloxSeg, a lightweight 3D medical segmentation model that aims to keep robustness while cutting compute. It uses Paired Window Attention to capture global context with near-linear cost, a Johnson–Lindenstrauss–guided rule to set grouped-convolution sizes so features preserve distance structure, and Spatially Decoupled Knowledge Transfer with Gram-matrix supervision to inject texture priors without extra inference cost. Across some datasets, it reports competitive Dice with markedly higher GPU/CPU throughput than common baselines.", "questions": "* What is the primary target metric (e.g., Dice vs. HD95), and do your claims hold on that metric with mean ± std over ≥3 seeds? Several baselines beat VeloxSeg on non-primary metrics—please provide confidence intervals, exact p-values, and a Pareto plot to make the trade-off explicit.\n\n* Were patch size, AMP/FP32 policy, I/O, dataloader workers, caching, pre/post-processing, and TTA identical across methods? Please re-run baselines in the same codebase and report end-to-end latency on the same hardware, with CPU thread counts fixed.\n\n* All headline results appear single-seed. Can you report 3–5 seeds for every table, include Box/Violin plots, and discuss sensitivity to random init and data order?\n\n* What are the teacher architectures, pretext tasks, and training data? Confirm there is no split leakage. Report the extra training compute and quantify SDKT’s incremental gain vs. strong alternatives under matched budgets.\n\n* At fixed GFLOPs/latency, do PWA and JL still win? Please include budgeted comparisons and scaling curves showing how accuracy changes as compute increases.\n\n\n**I am open to changing my score based on the author's responses.**", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* Creative combo of (i) Paired Window Attention for global–local context with near-linear cost, (ii) a Johnson–Lindenstrauss–guided rule to size grouped convolutions, and (iii) Gram-matrix (SDKT) supervision that adds no inference overhead. The JL connection to architectural sizing is a fresh angle for lightweight 3D segmentation.\n\n* Solid empirical scope across multiple 3D medical benchmarks with both accuracy and throughput (GPU/CPU) reported alongside params and GFLOPs. Comparisons include lightweight and multimodal baselines, giving a practical sense of trade-offs.", "weaknesses": "* **Inconsistent superiority across metrics.** The model wins on some averages but lags on others (e.g., HD95/precision/recall) with non-trivial gaps. \n\n* **Comparison fairness is under-specified.** Throughput (GPU/CPU) and accuracy comparisons may be skewed by different patch sizes, AMP/FP32, dataloader workers, I/O, pre/post-processing, and TTA.\n\n* **Statistical robustness is weak.** Results appear single-seed without CIs.\n\n* **SDKT attribution unclear.** Teacher details and potential data leakage aren’t fully specified; gains could stem from extra supervision rather than architecture.\n\n\n* **External validity.** Limited evidence under domain shift (site/scanner).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)