Fit Pixels, Get Labels: Meta-learned Implicit Networks for Image Segmentation

TL;DR

This paper tackles the challenge of applying implicit neural representations (INRs) to image segmentation by introducing MetaSeg, a meta-learning framework that jointly trains an INR to predict pixel values and segmentation labels. Through a nested MAML-style procedure, MetaSeg learns optimal initial parameters $(\theta^*, \phi^*)$ for a coordinate-based network $f_{\theta}$ and a shallow segmentation head $g_{\phi}$, enabling rapid test-time fine-tuning on unseen images. On 2D and 3D brain MRI segmentation, MetaSeg achieves Dice scores comparable to strong U-Net baselines while using roughly 90% fewer parameters, and it outperforms a closely related INR baseline (NISF) while supporting fine-grained segmentation and robustness to downsampling. These results suggest that INR-based segmentation, when guided by meta-learned priors, can offer a scalable, resource-efficient alternative to conventional heavy architectures, with learned latent representations that encode meaningful anatomical structure.

Abstract

Implicit neural representations (INRs) have achieved remarkable successes in learning expressive yet compact signal representations. However, they are not naturally amenable to predictive tasks such as segmentation, where they must learn semantic structures over a distribution of signals. In this study, we introduce MetaSeg, a meta-learning framework to train INRs for medical image segmentation. MetaSeg uses an underlying INR that simultaneously predicts per pixel intensity values and class labels. It then uses a meta-learning procedure to find optimal initial parameters for this INR over a training dataset of images and segmentation maps, such that the INR can simply be fine-tuned to fit pixels of an unseen test image, and automatically decode its class labels. We evaluated MetaSeg on 2D and 3D brain MRI segmentation tasks and report Dice scores comparable to commonly used U-Net models, but with $90\%$ fewer parameters. MetaSeg offers a fresh, scalable alternative to traditional resource-heavy architectures such as U-Nets and vision transformers for medical image segmentation. Our project is available at https://kushalvyas.github.io/metaseg.html .

Figures (5)

• Figure 1: Overview of MetaSeg. (a) We use a meta-learning framework to learn optimal initial parameters $\theta^*, \phi^*$ for an INR consisting of an $L$-layer reconstruction network $f_\theta(\cdot)$ and shallow segmentation head $g_\phi(\cdot)$. (b) At test time, optimally initialized INR $f_{\theta^*}$ is iteratively fit to the pixels of an unseen test scan. After convergence, the penultimate features $f_{\theta^*}^{L-1}(x)$ are fed as input to the segmentation head $g_\phi^*(\cdot)$ to predict per-pixel class labels.
• Figure 2: Example progression of reconstruction and segmentation performance as a MetaSeg INR fits the pixels of a test image. We report PSNR (top) and Dice (bottom) scores. MetaSeg's initialization leads to rapid convergence, with a Dice score of 0.85 emerge as early as $T_f=2$ iterations. It achieves optimal segmentation (Dice=$0.95$, PSNR=$41.3$) at $T_f=100$ iterations, and eventually declines as it overfits to the image pixels at $T_f = 5000$ iterations.
• Figure 3: Coarse and fine-grained 2D brain MRI segmentation results with 24 classes, for two subjects. MetaSeg accurately segments structures and is robust to high variations across subjects for structures such as ventricles (depicted in purple), brain stem (gray), and hippocampus (yellow). MetaSeg also adapts well to structures such as the cerebral cortex (red) which are not compact and localized in space.
• Figure 4: PCA visualization of penultimate features learned by MetaSeg and a standard INR for a 2D test scan. We performed PCA separately on the features returned by each INR across spatial coordinates. We see a strong correlation for MetaSeg features with anatomical structures. For example, component #2 approximately resembles inner brain regions, #3 resembles regions like hippocampus and basal ganglia, #4 resembles ventricles, and #5 captures the cerebral cortex. On the contrary, a standard INR yields seemingly random features.
• Figure 5: MetaSeg is capable of accurately segmenting 3D MRI volumes. (a) Various cross sections of one learned volume. (b) Renderings of intermediate coronal, axial, and sagittal planes, with high agreement (Dice$\approx 0.93$) with ground truth.