How Far Can INRs Go? Cross-Domain Parameter-efficient INR-Based Semantic Segmentation for Brain MRI

Ziyao Shang; Pouya Sadeghi; Letian Jiang; Alexander Wong; Sirisha Rambhatla

cs.CVSep 25, 2026

How Far Can INRs Go? Cross-Domain Parameter-efficient INR-Based Semantic Segmentation for Brain MRI

Authors: Ziyao Shang, Pouya Sadeghi, Letian Jiang, Alexander Wong, Sirisha Rambhatla

Organizations: Department of Systems Design Engineering University of Waterloo · David R. Cheriton School of Computer Science University of Waterloo · Department of Management Science and Engineering University of Waterloo

Abstract

Biomedical image segmentation is central to medical image analysis, but practical deployment often faces limited annotations, memory constraints, and cross-site distribution shifts. Implicit Neural Representations (INRs) have recently emerged as a lightweight alternative for semantic segmentation, achieving competitive performance with substantially fewer parameters than conventional architectures. However, the mechanisms, scaling behavior, and domain generalization abilities of INR-based segmentation remain insufficiently understood. In this work, we study these questions in the context of cross-domain brain MRI segmentation. We analyze INR-based segmentation across low-parameter regimes, comparing it with conventional pipelines in both in-domain and out-of-domain settings. Surprisingly, we find that INR-based models do not simply improve with increasing parameter budget. Their advantage is most pronounced under low-parameter and limited-augmentation settings, while U-Net-based models benefit more from larger capacity and standard augmentation. We also investigate how INRs encode semantic information in their hidden features and show that complementary segmentation-relevant structure is distributed across multiple INR layers. Building on this insight, we introduce HierINRSeg, a hierarchical INR-based architecture that aggregates multi-layer representations for improved robustness and generalization. Extensive experiments show that HierINRSeg consistently outperforms MetaSeg, a strong recent INR-based segmentation baseline, with an average improvement of 5.6 percentage points in Dice for the in-domain test set and 8.2 percentage points out-of-domain. Overall, our analysis identifies the conditions under which INR-based segmentation is most effective, providing concrete guidance for model selection and future research.

Figures & tables

Figure 1: Scaling comparison of U-Net- and INR-based segmentation models under training augmentation and domain shift. Results are shown on KKI, the source (in-domain) site, and CMU, a target (out-of-domain) site. Left: foreground mean Dice across parameter budgets for U-Net, MetaSeg, and our HierINRSeg, demonstrating stronger performance for INR-based models in the low-parameter regime, while U-Net becomes more competitive as model capacity increases. Right: representative segmentation results at lower- and higher-capacity settings, showing a transition from INR-advantaged performance at smaller model sizes to U-Net-advantaged performance at larger model sizes. Red circles highlight typical segmentation errors. The ground truth segmentation is denoted as GT.

Figure 2: Visualizing MetaSeg INR features before and after fitting to a single image using UMAP. Earlier layers remain largely stable against fitting, while the final layer changes substantially, suggesting that early and middle layers contain more reusable structure while latter layers are more image-specific.

Figure 3: Layer-wise UMAP of MetaSeg INR features pooled across validation images after image-specific fitting. The first layer shows weak semantic organization, middle layers preserve a shared brain-like spatial structure, and latter layers become more label-compact while losing global anatomical layout.

Figure 4: The HierINRSeg architecture: instead of predicting from a single INR layer as in MetaSeg (Orange panel), our HierINRSeg (green panel) progressively fuses multi-layer INR features with the segmentation head. This allows the segmentation head to gradually combine complementary spatial and semantic information from different INR layers before producing the final label map. The gray panel shows how 2D images are extracted from 3D MRI volumes. We select a segment along the Y-axis of the brain, centered at the brain center, that covers k% of its total length. From this segment, slices in the XZ-plane are sampled to form the 2D images used in our experiments. We use k=5 for the main experiments, while the set of values for k is expanded to {5, 20, 40, 60, 80} in Section 5.2 .

Panel of Figure 5 — Figure 5: Controlled parameter-scaling comparison. All models are trained on KKI without augmentation and evaluated on KKI (ID, test split) and Caltech / CMU (OOD). We compare MetaSeg, HierINRSeg, and U-Nets across parameter budgets from 5K to 1M. Dashed lines represent the large-scale TransUNet. All runs are repeated for three different seeds, with markers denoting the mean and the surrounding shades showing ± standard deviation. The two stars represent MetaSeg and our HierINRSeg using MetaSeg’s original hidden dimension of 128. HierINRSeg consistently improves over MetaSeg across parameter budgets, while U-Net requires substantially larger capacity to become competitive, especially on the OOD test sets.

Panel of Figure 6 — Figure 6: Practical parameter-scaling comparison with strengthened training. All models are trained on KKI with augmentation and evaluated on KKI (ID, test split) and Caltech/CMU (OOD). We include both cross-entropy and focal-loss U-Nets and the parameter-efficient UNeXt, with the large-scale nnU-Net and TransUNet shown as reference baselines. All runs are repeated for three different seeds, with markers denoting the mean and the surrounding shades showing ± standard deviation. Our HierINRSeg consistently improves over MetaSeg and remains competitive with strengthened U-Net and UNeXt baselines across domains. The UNeXt parameter sweep is shown in the second row for visual clarity.

Figure 7: Qualitative results at the 20K parameter budget with training augmentation. We compare predictions from U-Net, MetaSeg, and HierINRSeg on the in-domain KKI test set and the out-of-domain Caltech and CMU test sets. Solid red circles indicate ignored regions, and dashed red circles indicate incorrect segmentations. HierINRSeg produces segmentations that more closely match the ground truth across both in-domain and out-of-domain examples.

Budget	Model	Performance (Dice) on each percentage
Budget	Model	5%	20%	40%	60%	80%
10K	MetaSeg	0.666 ± 0.020	0.648 ± 0.016	0.616 ± 0.011	0.535 ± 0.034	0.526 ± 0.016
	HierINRSeg	0.703 ± 0.013	0.724 ± 0.059	0.680 ± 0.008	0.615 ± 0.019	0.644 ± 0.024
	U-Net	0.466 ± 0.093	0.450 ± 0.061	0.451 ± 0.016	0.476 ± 0.090	0.617 ± 0.014
200K	MetaSeg	0.697 ± 0.011	0.694 ± 0.018	0.623 ± 0.007	0.512 ± 0.055	0.517 ± 0.015
	HierINRSeg	0.746 ± 0.003	0.768 ± 0.005	0.735 ± 0.048	0.666 ± 0.016	0.640 ± 0.034

Table 1: KKI in-domain performance under increasing 3D spatial coverage. Slices are sampled from progressively wider central regions of each volume (5–80%). Results are reported at representative 10K and 200K parameter budgets.

Panel of Figure 8 — Figure 8: Mean normalized augmentation effect ( $R(k)$ from eq. 1 ) for each INR layer during test-time INR fitting for HierINRSeg and MetaSeg (120k) trained without augmentation. In both models, the augmentation effect first spikes, then becomes relatively small after only a short period of fitting, indicating that as test-time INR adaptation progresses, effects of augmentation-induced variations decrease relative to the variations between different images. The same trend is observed for both HierINRSeg and MetaSeg trained with augmentation, as shown in Appendix Figure 11 .

Head config	Params	ID	OOD
Head config	Params	KKI	Caltech	CMU
W32-D1	67K + 4K	0.670 ± 0.001	0.562 ± 0.010	0.656 ± 0.013
W64-D1	67K + 8K	0.676 ± 0.012	0.560 ± 0.007	0.652 ± 0.012
W128-D1 (MetaSeg)	67K + 17K	0.685 ± 0.005	0.575 ± 0.014	0.658 ± 0.005
W256-D1	67K + 34K	0.689 ± 0.020	0.575 ± 0.026	0.659 ± 0.020
W512-D1	67K + 69K	0.687 ± 0.009	0.575 ± 0.006	0.656 ± 0.006

Table 2: Head-capacity ablation for MetaSeg with the INR kept fixed. All models are trained on KKI, and Dice scores are averaged across three seeds. Parameter counts are reported as INR + head. W denotes the width of the hidden layers of the segmentation head. D denotes the number of hidden layers.

Model	Params	Selected layers					ID	OOD
Model	Params	1	2	3	4	5	KKI	Caltech	CMU
MetaSeg	67K+17K					✓	0.685 ± 0.005	0.575 ± 0.014	0.658 ± 0.005
MetaSeg: large	67K + 594K					✓	0.672 ± 0.005	0.565 ± 0.005	0.627 ± 0.022
HierINRSeg: 1 layer	67K+3K					✓	0.700 ± 0.014	0.586 ± 0.005	0.690 ± 0.005
HierINRSeg: 2 layers	67K+36K				✓	✓	0.727 ± 0.023	0.636 ± 0.010	0.716 ± 0.010
HierINRSeg: 3 layers	67K+70K			✓	✓	✓	0.731 ± 0.006	0.638 ± 0.003	0.721 ± 0.010

Table 3: Component ablation of HierINRSeg where INR features from different sets of layers are used. All models are trained on KKI, and Dice scores are averaged across three seeds. Parameter counts are reported as INR + head.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Target Size	Feature counts	Parameters
5k	$(1,1,4,6,13)$	5,037
10k	$(1,2,5,9,18)$	10,005
20k	$(2,4,8,13,23)$	19,999
40k	$(2,5,10,18,36)$	40,011
70k	$(3,6,12,24,49)$	70,003
80k	$(4,7,12,26,52)$	79,999

Appendix

Table 4: U-Net model sizes.

Target size	$d$ (layer width)	Parameters
5k	30	4,957
10k	43	9,897
20k	62	20,157
40k	88	40,047
70k	117	70,207
80k	128	83,847

Appendix

Table 5: Metaseg model sizes.

Target size	$d$ (layer width)	Parameters
5k	18	5,101
10k	26	9,861
20k	38	19,881
40k	55	39,992
70k	74	70,677
80k	79	80,192

Appendix

Table 6: HierINRSeg model sizes.

Target size	Feature counts	Parameters
5k	$(1,3,5,6,21)$	$5{,}000$
10k	$(3,3,5,12,25)$	$10{,}000$
20k	$(2,3,5,20,34)$	$20{,}000$
40k	$(1,6,12,25,51)$	$40{,}000$
70k	$(3,7,18,31,71)$	$70{,}000$
80k	$(6,11,17,34,74)$	$80{,}000$

Appendix

Table 7: UNeXT model sizes. Feature counts specify the channel widths at the five stages.

Hyperparameter	Value
Step 1 inner-loop fitting steps	2
Validation/test time INR fitting steps	100
Validation frequency	Every 50 iterations
Focal loss:	gamma=1, no class weights
INR learning rate	$1\times 10^{-4}$
Classifier learning rate	$5\times 10^{-5}$

Appendix

Table 8: Remaining training and validation hyperparameters.

Figure 9: UMAP parameter sweep for visualizations shown in Figure 2 .

Figure 10: UMAP parameter sweep for visualizations shown in Figure 3 .

Budget	Model	Performance (Dice) on each percentage
Budget	Model	5%	20%	40%	60%	80%
10K	MetaSeg	0.542 ± 0.016	0.518 ± 0.027	0.491 ± 0.021	0.432 ± 0.048	0.426 ± 0.012
	HierINRSeg	0.596 ± 0.025	0.610 ± 0.009	0.587 ± 0.009	0.534 ± 0.014	0.555 ± 0.017
	U-Net	0.461 ± 0.027	0.419 ± 0.074	0.404 ± 0.043	0.467 ± 0.095	0.583 ± 0.011
200K	MetaSeg	0.601 ± 0.017	0.583 ± 0.022	0.522 ± 0.008	0.406 ± 0.055	0.415 ± 0.040
	HierINRSeg	0.683 ± 0.019	0.699 ± 0.002	0.663 ± 0.004	0.610 ± 0.018	0.575 ± 0.015

Appendix

Table 9: Caltech (OOD) performance under increasing 3D spatial coverage. Slices are sampled from progressively wider central regions of each volume (5–80%). Results are reported at representative 10K and 200K parameter budgets.

Budget	Model	Performance (Dice) on each percentage
Budget	Model	5%	20%	40%	60%	80%
10K	MetaSeg	0.624 ± 0.010	0.560 ± 0.030	0.536 ± 0.029	0.475 ± 0.057	0.458 ± 0.013
	HierINRSeg	0.661 ± 0.006	0.660 ± 0.010	0.624 ± 0.002	0.568 ± 0.003	0.580 ± 0.022
	U-Net	0.355 ± 0.026	0.391 ± 0.075	0.408 ± 0.032	0.457 ± 0.105	0.593 ± 0.005
200K	MetaSeg	0.678 ± 0.005	0.637 ± 0.017	0.563 ± 0.011	0.449 ± 0.054	0.454 ± 0.041
	HierINRSeg	0.738 ± 0.007	0.742 ± 0.005	0.696 ± 0.006	0.634 ± 0.008	0.592 ± 0.016

Appendix

Table 10: CMU (OOD) performance under increasing 3D spatial coverage. Slices are sampled from progressively wider central regions of each volume (5–80%). Results are reported at representative 10K and 200K parameter budgets.

Panel of Figure 11 — Figure 11: Normalized augmentation effect during test-time INR fitting for HierINRSeg and MetaSeg (120k) trained with augmentation. In both models, the relative augmentation effect becomes small after only a short period of fitting, showing the same trend as Figure 8 .

Dataset	MetaSeg	HierINRSeg	U-Net
KKI (ID)	0.688 (+0.019)	0.750 (+0.001)	0.701 (+0.016)
Caltech (OOD)	0.569 (+0.040)	0.676 (+0.013)	0.654 (+0.100)
CMU (OOD)	0.658 (+0.038)	0.735 (+0.016)	0.661 (+0.086)

Appendix

Table 11: Effect of adding augmentation to training at the 120K parameter budget. Each entry reports the Dice score for the model trained with augmentation, followed by the augmentation gain (as compared to the same model trained without augmentation).

Figure 12: The absolute extents of feature changes caused by augmentation (left) and image variation (right) for MetaSeg trained without augmentation

Figure 13: The absolute extents of feature changes caused by augmentation (left) and image variation (right) for MetaSeg trained with augmentation

Figure 14: The absolute extents of feature changes caused by augmentation (left) and image variation (right) for HierINRSeg trained without augmentation

Figure 15: The absolute extents of feature changes caused by augmentation (left) and image variation (right) for HierINRSeg trained with augmentation

Selected INR layers	Selected layers					$\Delta$ Dice
Selected INR layers	1	2	3	4	5	KKI	Caltech	CMU
1 layer					✓	+0.060	+0.046	+0.069
2 layers				✓	✓	+0.039	+0.043	+0.018
3 layers			✓	✓	✓	+0.002	-0.011	-0.004
excl. middle	✓				✓	+0.008	+0.008	+0.005
4 layers		✓	✓	✓	✓	-0.004	-0.019	-0.007

Appendix

Table 12: An alternative HierINRSeg variant where we replace the input coordinates to the segmentation head with the earliest selected INR layer. The numbers shown are the differences in Dice (embedded coordinates

-

earliest-INR feature)

Semantic Segmentation

Semantic segmentation assigns a class label to every image pixel, enabling detailed scene understanding across domains and sensors.

See all 416 papers

Disentangled Learning Improves Implicit Neural Representations for Medical Reconstruction
May 5, 2026Qing Wu, Xuanyu Tian, Chenhe Du +4Disentangled Representation Learning Implicit Neural Representations
Medical AI Lab, Ant Group · ShanghaiTech University · Shanghai Jiao Tong University
Confidence-Calibrating Regularization for Robust Brain MRI Segmentation Under Domain Shift
Sep 27, 2025Behraj Khan, Tahir Qasim Syed, Syed Ahmad Chan BukhariImage Segmentation Promptable Image Segmentation
St. John’s University New York, USA · Institute of Business Administration Karachi Pakistan
LUMINR: Intensity-conditioned Meta-learned INRs for Accelerated Medical Image Registration
Oct 16, 2026Damian Rynczak, Marek WodzinskiMedical Image Registration Deformable Image Registration
AGH University of Krakow · Sano - Centre for Computational Personalised Medicine, International Research Foundation

May 5, 2026J/K move · Enter open · S save

Disentangled Learning Improves Implicit Neural Representations for Medical Reconstruction

Qing Wu, Xuanyu Tian, Chenhe Du, Haonan Zhang+3

Medical AI Lab, Ant Group · ShanghaiTech University · Shanghai Jiao Tong University

Disentangled Representation Learning Implicit Neural Representations Image Reconstruction Medical Image Representation Learning

Open paper AskPDF arXiv

Implicit neural representations (INRs) have emerged as a powerful paradigm for medical imaging via physics-informed unsupervised learning. Classical INRs optimize an entire network from scratch for each subject, leading to inefficient training and suboptimal imaging quality. Recent initialization-based approaches attempt to inject population priors into pre-trained networks, yet they rely on high-quality images and often suffer from catastrophic forgetting during fine-tuning. We present DisINR, a novel INR framework that explicitly disentangles shared and subject-specific representations. DisINR introduces a shared encoder-decoder pair and subject-specific encoders, whose features are jointly decoded for image reconstruction. By integrating differentiable forward models, it pre-trains the shared modules directly from limited raw measurements, removing the need for pre-acquired high-quality images. During test-time adaptation, only the subject-specific encoder is optimized, while the shared pair remains frozen, effectively preserving learned priors. Extensive evaluations on three representative medical imaging tasks show that DisINR significantly outperforms state-of-the-art INRs in both reconstruction accuracy and efficiency.

Abstract

Semantic Segmentation

Explore similar work

Disentangled Learning Improves Implicit Neural Representations for Medical Reconstruction

Confidence-Calibrating Regularization for Robust Brain MRI Segmentation Under Domain Shift

LUMINR: Intensity-conditioned Meta-learned INRs for Accelerated Medical Image Registration