Skip to content

README.md

Latest commit

 

History

History
52 lines (32 loc) · 4.14 KB

README.md

File metadata and controls

52 lines (32 loc) · 4.14 KB
Raw

CyL-Adapter: Cross-Modality Latent Normalizing Flow Adapter

This repository serves as a pure architectural reference for the CyL-Adapter framework, paired with a Unified 3D V-Net Foundation Model. It is explicitly designed for medical imaging researchers seeking to perform highly parameter-efficient, single-shot Domain Adaptation across severe domain shifts (e.g., CINE MRI to LGE MRI, or CT to MRI) where both domains share the same anatomical label space.

1. The Foundation Model Architecture

At the core of the framework is the Unified Foundation Model, implemented as a deep 3D V-Net encoder-decoder structure. This model is trained from scratch on the abundant, annotated Source Domain.

  • Encoder (encode): Utilizes powerful SqueezeExcitation3D and AttentionGate3D blocks to iteratively downsample the 3D volume, ultimately projecting the image into a highly condensed, semantic latent bottleneck space: $Z_{source} \in \mathbb{R}^{B \times 256 \times D \times H \times W}$.
  • Decoder (decode): A symmetric expanding pathway that reconstructs dense, multi-scale features using ASPP3D.
  • Classification Head: A standard $1 \times 1 \times 1$ 3D Convolution that projects the final features into segmentation logits.

Note

Once Phase 1 is complete, the Foundation Model has thoroughly mapped the anatomy of the Source Domain. It is strictly frozen for the remainder of the pipeline to prevent catastrophic forgetting.

2. The CyL-Adapter Architecture

When you introduce a new Target Domain (e.g., a new MRI sequence or a different hospital's scanner), standard fine-tuning often destroys the robust features learned by the Foundation Model, especially if you only have a single-shot (one patient) to train on.

To solve this, we introduce the CyL-Adapter. It acts as a mathematically invertible translator between the Target Domain and the Source Domain using Normalizing Flows.

  • Forward Flow ($T_{flow}$): An Invertible RealNVP-style Affine Coupling layer comprising $1 \times 1 \times 1$ 3D convolutions. It intercepts the bottleneck latent space of the Target image ($Z_{target}$) and dynamically warps it to match the anatomical distribution of the Source latent space ($Z_{source}$).

    $$ Z_{aligned} = T_{flow}(Z_{target}) $$

Because the flow preserves spatial topologies while shifting feature distributions, the aligned target latent space is perfectly understood by the frozen Foundation Decoder.

Why is there no Inverse Flow?

In traditional unsupervised cycle-consistency frameworks, an inverse flow ($T_{inv}$) is used to reconstruct the original image. However, because our framework targets supervised adaptation (where both modalities share the identical set of universal segmentation labels), we rely entirely on the Forward Flow. The aligned features flow directly through the frozen Decoder and Head, outputting the universal segmentation mask.

3. The Two-Stage Training Process

The repository exposes two highly abstracted, dataset-agnostic training modules ready to plug into your custom 2D/3D dataloaders.

Phase 1: train_foundation

Located in train_phase1_foundation.py. The Unified 3D V-Net is trained from scratch on your abundant Source Domain dataset using standard Cross-Entropy Loss. The model comprehensively learns the underlying anatomy and structures a rich $Z_{source}$ latent space.

Phase 2: train_adapter

Located in train_phase2_adapter.py.

  1. The Foundation Model is entirely frozen.
  2. The CyLAdapterModel wraps the foundation model and initializes the $T_{flow}$ layer.
  3. Target Domain images are fed in.
  4. The Magic: The adapter is trained purely using the Task Loss (Cross-Entropy Loss on the segmentation output). The error gradients flow backwards through the frozen Head, through the frozen Decoder, and straight into $T_{flow}$. The flow is mathematically forced to warp $Z_{target}$ into the exact distribution that the frozen Source components expect.

Because only the extremely lightweight Normalizing Flow is being updated, the adapter is mathematically protected from overfitting, making it uniquely powerful for Single-Shot Domain Adaptation.