Learn2Synth: Learning Optimal Data Synthesis Using Hypergradients for Brain Image Segmentation

Conference: ICCV 2025 arXiv: 2411.16719 Code: https://github.com/HuXiaoling/Learn2Synth Area: Image Segmentation Keywords: Domain Randomization, Hypergradients, Synthetic Data Augmentation, Brain Image Segmentation, Domain Generalization

TL;DR

This paper proposes Learn2Synth, a training framework that leverages hypergradients to learn optimal synthetic data augmentation parameters, enabling segmentation networks trained exclusively on synthetic data to achieve peak performance on real data. The framework simultaneously attains high in-domain accuracy and strong out-of-domain generalization, outperforming both SynthSeg and supervised learning baselines on brain MRI segmentation tasks.

Background & Motivation

Acquiring high-quality annotated data in medical imaging is constrained by scanning costs, image noise and artifacts, and the specialized expertise and time required for annotation. This leads to the persistent problem of poor generalization in modality-specific models—an issue particularly pronounced in brain image segmentation, where contrast varies substantially across scanners, pulse sequences (MPRAGE vs. FLASH), and acquisition parameters.

Existing Approaches and Their Limitations:

Supervised Learning: Achieves peak in-domain performance but degrades rapidly out-of-domain and severely overfits on small datasets.

Domain Randomization (e.g., SynthSeg): Generates synthetic images with randomized contrast from label maps to train networks, yielding strong generalization but suffering from a "reality gap" that persistently limits in-domain accuracy compared to supervised methods.

Mixed Training (Synthetic + Real Data): The network may internalize parallel sub-networks and partially overfit to the limited real data.

Distribution Matching (GANs, Contrastive Learning, Diffusion Models): Makes synthetic images "look like" real images but introduces objectives unrelated to the segmentation task and may disrupt label–image alignment.

Key Challenge: How can one simultaneously achieve high in-domain accuracy (the advantage of supervised learning) and strong out-of-domain generalization (the advantage of domain randomization)?

Core Idea: Rather than exposing the segmentation network directly to real data, Learn2Synth trains a learnable augmentation network to "calibrate" synthetic data such that the segmentation network trained on the calibrated data achieves optimal performance on real data. The key is transmitting the loss signal from real data to the augmentation network via hypergradients (differentiating through the update step).

Method

Overall Architecture

Learn2Synth alternates between two passes: 1. Synthetic Pass: The augmentation network \(A_\boldsymbol{\theta}\) is frozen; synthetic data are augmented and fed to the segmentation network \(S_\boldsymbol{\phi}\) for training, updating \(\boldsymbol{\phi}\). 2. Real Pass: The segmentation network is frozen; real data are passed through the segmentation network to compute the loss, and the augmentation network parameters \(\boldsymbol{\theta}\) are updated via hypergradients.

A critical property is that the segmentation network never has its weights updated directly on real data, preventing overfitting to real data.

Key Designs

1. Hypergradient Mechanism:

   • Function: Establishes a gradient pathway from the augmentation network to segmentation accuracy on real data.
   • Mechanism: In the Real Pass, real data are passed through the (already updated) segmentation network to obtain the loss \(\mathcal{L}_{\text{real}} = \text{SoftDice}(S_{\boldsymbol{\phi}^*}(\mathbf{x}_{\text{real}}), \mathbf{y}_{\text{real}})\). The gradient with respect to the augmentation network is: \(\mathbf{g}_\theta = \frac{\partial \mathcal{L}_{\text{real}}}{\partial \boldsymbol{\theta}} = \frac{\partial \mathcal{L}_{\text{real}}}{\partial \boldsymbol{\phi}^*} \times \frac{\partial \boldsymbol{\phi}^*}{\partial \mathbf{g}_\phi} \times \frac{\partial^2 \mathcal{L}_{\text{synth}}}{\partial \boldsymbol{\phi} \partial \boldsymbol{\theta}^T}\)
   • Interpretation of the three components: (i) gradient of the real loss with respect to segmentation network weights; (ii) derivative of the update step (learning rate times identity under SGD); (iii) Hessian of the synthetic loss with respect to both parameter sets (computed via automatic differentiation without explicit construction).
   • Design Motivation: Rather than using distribution matching (GAN/contrastive learning), this approach adopts "improving segmentation accuracy on real data" as the sole optimization objective.

2. Parametric Augmentation Model:

   • Function: Learns optimal parameters for Gaussian noise and intensity non-uniformity (INU) in MRI images.
   • INU Model: Models the spatial profile of receive coils using multi-frequency B-spline basis functions. Three random fields at different spatial frequencies (\(K=3\)) are combined via learnable coefficients \(\mathbf{c} = [c_{\text{low}}, c_{\text{mid}}, c_{\text{high}}]\): \(\boldsymbol{\alpha} = \prod_{k=1}^K \boldsymbol{\alpha}_k^{c_k}, \quad \mathbf{x}_{\text{synth}} \leftarrow \mathbf{x}_{\text{synth}} \odot \boldsymbol{\alpha}\)
   • Gaussian Noise Model: Learns a learnable standard deviation \(\sigma\): \(\mathbf{x}_{\text{synth}} \leftarrow \mathbf{x}_{\text{synth}} + \sigma \cdot \boldsymbol{\varepsilon}\)
   • Variants: Fixed \(\sigma\) vs. randomly modulated \(\sigma\) (i.e., \(\sigma \cdot s\), \(s \sim \mathcal{N}(0,1)\)).

3. Nonparametric Augmentation Model:

   • Function: Employs a UNet to learn augmentation residuals of arbitrary form.
   • Mechanism: The synthetic image is concatenated with a single channel of Gaussian noise and fed into the UNet to learn a residual augmentation: \(\mathbf{x}_{\text{synth}} \leftarrow \mathbf{x}_{\text{synth}} + A_\boldsymbol{\theta}([\mathbf{x}_{\text{synth}}, \boldsymbol{\xi}]), \quad \boldsymbol{\xi} \sim \mathcal{N}_N(0,1)\)
   • Design Motivation: The parametric model requires predefined augmentation types, whereas the nonparametric model can automatically discover optimal augmentations. However, when a good parametric model is available, the parametric approach is preferable.

Loss & Training

• The Synthetic Pass updates the segmentation network using SoftDice Loss.
• The Real Pass updates the augmentation network via hypergradients.
• Training is based on 434 automatically segmented brain images from the OASIS dataset.
• The segmentation network and augmentation network are updated in alternation.
• Learnable parameters: for the parametric model, only \([c_1, c_2, c_3, \sigma]\) (4 scalars); for the nonparametric model, the UNet parameters.

Key Experimental Results

Main Results

Segmentation accuracy (Dice) on real brain MRI datasets (ABIDE and OASIS3):

Method	ABIDE	OASIS3	Notes
Supervised UNet	0.908	0.899	In-domain upper bound
SAMSEG	0.875	0.841	Unsupervised Bayesian
Naive SynthSeg	0.869	0.831	Standard domain randomization
Mixed SynthSeg	0.875	0.854	Mixed real + synthetic
Finetuned SynthSeg	0.871	0.847	Fine-tuned on real data
AdvChain	0.867	0.848	Adversarial augmentation baseline
Learn2Synth (nonparam)	0.879	0.881	Surpasses mixed training

Ablation Study

Cross-contrast generalization (MPRAGE training → FLASH testing):

Method	#Train	MPRAGE	FLASH 3°	FLASH 5°	FLASH 20°	FLASH 30°
SynthSeg	/	0.861	0.776	0.694	0.766	0.781
Supervised (29)	29	0.941	0.419	0.396	0.671	0.769
Supervised (5)	5	0.907	0.397	0.413	0.586	0.692
Learn2Synth (29)	29	0.895	0.804	0.789	0.785	0.797
Learn2Synth (5)	5	0.867	0.798	0.789	0.795	0.799

Parameter inference experiment (noise parameter recovery on synthetic data):

Preset \(\hat{\sigma}\)	0	0.050	0.100	0.150	[0.025,0.2]
Inferred \(\sigma^*\)	0.001	0.042	0.098	0.146	0.134

The inferred parameters closely match the preset values, validating the effectiveness of the learning mechanism.

Key Findings

• Learn2Synth achieves a Dice score of 0.881 on OASIS3, surpassing all baselines (including supervised learning, which achieves 0.899 but with poor generalization).
• Remarkable cross-contrast generalization: Supervised learning achieves only 0.419 on FLASH 3°, while Learn2Synth reaches 0.804—nearly double.
• Learn2Synth with only 5 training samples outperforms SynthSeg on all FLASH sequences, demonstrating efficient utilization of limited annotations.
• The parametric model outperforms the nonparametric model when the augmentation type is known; the nonparametric model is more flexible when augmentation types are unknown.
• As expected, Learn2Synth simultaneously improves both in-domain (MPRAGE) and out-of-domain (FLASH) performance.

Highlights & Insights

• Elegant Training Paradigm: Real data are leveraged indirectly through hypergradients; the segmentation network never directly touches real data, perfectly balancing in-domain accuracy and out-of-domain generalization.
• Single Optimization Objective: "Maximize segmentation accuracy on real data" is the sole objective, avoiding the introduction of task-irrelevant goals as in GAN-based methods.
• Interpretable Parameters: The learned augmentation parameters reveal the optimal training distribution, providing insights to guide manual hyperparameter tuning.
• Excellent Experimental Design: Synthetic experiments validate parameter recovery, real experiments validate practical utility, and cross-contrast experiments validate generalization.
• The framework is applicable to any scenario involving synthetic data training, not limited to medical imaging.

Limitations & Future Work

• Hypergradient computation requires backpropagation through the update step, incurring high computational cost due to second-order gradients.
• Validation is currently limited to 2D brain image segmentation; 3D volumetric segmentation and other organs remain to be explored.
• The parametric model requires prior knowledge of augmentation types (noise, INU); nonparametric models are needed for unknown artifact types.
• The UNet in the nonparametric model introduces additional parameters and computational overhead.
• Only SoftDice is used as the loss function; compatibility with other task-specific losses (e.g., topology-aware losses) has not been verified.

Related Work & Insights

• Hypergradients have been applied in meta-learning; this work creatively extends their use to data augmentation optimization.
• The training strategy of "never exposing the model directly to real data" is generalizable to other data-scarce scenarios.
• The comparison between parametric and nonparametric augmentation models provides guiding principles for selecting augmentation strategies.
• The combination of domain randomization and Learn2Synth can be extended to other Synth* series methods such as SynthMorph and SynthSR.

Rating

• Novelty: ⭐⭐⭐⭐⭐ Applying hypergradients to augmentation learning is a genuinely novel idea with an elegant training paradigm.
• Experimental Thoroughness: ⭐⭐⭐⭐ Synthetic and real experiments are comprehensive, though limited to 2D brain segmentation.
• Writing Quality: ⭐⭐⭐⭐⭐ Mathematical derivations are clear and rigorous; experimental design is logically progressive.
• Value: ⭐⭐⭐⭐ The training paradigm has broad applicability and carries significant value for medical image segmentation.

Related Papers

• [ICCV 2025] CLOT: Closed Loop Optimal Transport for Unsupervised Action Segmentation
• [ICCV 2025] Dynamic Dictionary Learning for Remote Sensing Image Segmentation
• [ICCV 2025] HiMTok: Learning Hierarchical Mask Tokens for Image Segmentation with Large Multimodal Model
• [ICCV 2025] UniGlyph: Unified Segmentation-Conditioned Diffusion for Precise Visual Text Synthesis
• [ICCV 2025] LEGION: Learning to Ground and Explain for Synthetic Image Detection