SphereFlow: Missing Modality Imputation via Geometric Transport on Hypersphere

AI summary

Overview

• Research area: Computer vision / medical image generation — specifically cross-modal (missing modality) imputation, latent representation design, and flow matching on Riemannian manifolds.
• Technical level: Intermediate to Advanced. The core ideas are describable in plain language, but the method relies on flow matching, representation autoencoders, and a geometric proof about off-manifold deviation.
• Scope: SphereFlow proposes projecting all imaging modalities onto a single shared hypersphere via a frozen DINOv2 encoder, so that imputing a missing modality becomes a short geometric transport between nearby latent points rather than a long jump between separate manifolds. Experiments cover BraTS multi-modal MRI synthesis, SynthRAD2023 CT–MRI translation, and downstream tumor segmentation.

What This Paper Is About

Missing modality imputation means synthesizing an image modality that was never acquired (for example, generating a T2 MRI scan when only T1 was scanned), using the modalities that were observed. The authors argue the real bottleneck is not the generator but the latent representation: existing methods encode each modality into its own separate latent space (pixel space or per-modality VAE), so the model must learn a long, unstructured mapping for every source–target pair. The paper's goal is to show that if all modalities are instead mapped into one shared, geometrically clean latent space — a hypersphere — then imputation reduces to a short, well-behaved movement between nearby points, and that flow matching can be adapted to exploit this.

Key Contributions

1. Reframing imputation as geometric transport on a hypersphere. Through a pilot study on three tokenizers (VAE, MAE, DINOv2), the authors identify two properties a shared latent space needs — cross-modal compactness (paired modalities of the same subject closer than different subjects of the same modality) and geometric regularity (latents concentrated on a known manifold) — and show that a DINOv2-based Representation Autoencoder uniquely satisfies both.

2. Identifying and fixing two "geometric conflicts." Applying standard flow matching on the sphere fails because (i) the Gaussian noise source has norm ≈ √d while data has norm 1, and (ii) linear interpolation between two sphere points dips inside the sphere, where the decoder has never been trained. The fixes are using the observed-modality latent as the flow source (cross-modal flow matching) and a time-weighted expansion loss that trains decoders to reconstruct from off-manifold interpolants.

3. A theoretical bound on off-manifold deviation. Proposition 1 shows the worst-case deviation of the linear interpolant is Δ = 1 − cos(α/2) at τ = 1/2, depending only on the angular separation α between source and target latents — not on the latent dimension d — compared to Θ(√d) for Gaussian-source flow matching.

4. State-of-the-art results on BraTS 1-to-1 (all 12 directions) and N-to-N synthesis, SynthRAD2023 CT↔MRI translation, and downstream brain tumor segmentation with imputed data.

Main Findings

• Tokenizer pilot study: On 500 paired BraTS T1/T1ce test slices, VAE gives a cross-modal-to-intra-modal cosine distance ratio of ≈7.96 (modalities land in unrelated regions); MAE reduces it to 0.47 but only by collapsing all latents into a tiny region; DINOv2 achieves the lowest ratio (0.41) while preserving discriminability. DINOv2 latents also concentrate sharply at ‖z‖ ≈ 1 with near-zero coefficient of variation.

• 1-to-1 BraTS translation: SphereFlow achieves the best PSNR on every one of the 12 directions, averaging 34.83 dB (+1.11 dB over the runner-up), and reduces average rFID from 29.99 to 19.93 (33.6%). The largest gains are on hard pairs, e.g. T2→T1ce improves by +1.24 dB.

• N-to-N BraTS synthesis: Averaged over 14 input combinations, SphereFlow reaches PSNR 35.88 (+0.75 dB over the runner-up), SSIM 0.9530, and rFID 19.68 (−1.5% vs. Zhang). Simple averaging of source latents suffices for multi-source fusion because all modalities share one hypersphere.

• CT–MRI translation (SynthRAD2023): On MRI→CT, SphereFlow reaches 27.82 dB PSNR, SSIM 0.9084, rFID 13.99. On CT→MRI, Zhang slightly leads PSNR/SSIM (25.51 / 0.8635 vs. 24.70 / 0.8452), but SphereFlow roughly halves rFID from 30.91 to 13.65.

• Downstream segmentation: With a 3D U-Net trained on full-modality BraTS and tested on imputed inputs, SphereFlow reaches average Dice 0.7473 (+15.4 over runner-up CodeBrain) and reduces HD95 from 46.27 to 19.02 (−59%). It lands within 0.108 of the full-modality oracle Dice of 0.8554 — the abstract describes this as closing roughly 80% of the gap. In single-source cases, baselines collapse below Dice 0.10–0.20 while SphereFlow stays above 0.45.

• Encoder ablation (T1c, Stage-1 reconstruction): MedMAE wins PSNR (31.23) but has 3.6× higher iFID (144.45 vs. 40.64); DINOv2 gives up 0.66 dB of PSNR for the best rFID/IS/iFID (30.57 dB PSNR, rFID 35.29, iFID 40.64), which the authors read as geometric regularity mattering more than pixel fidelity for transport.

• Stage-wise decoder ablation (T1/T1ce): Step 1 (single-modality warm-up) gives on-manifold 33.14 dB / rFID 28.29 and chord 32.04 dB / rFID 41.66. Adding Step 2 (manifold expansion) gives +0.69 dB on-manifold and +1.25 dB on the chord, though chord rFID slightly worsens (28.29→30.92). Adding Step 3 (shared-encoder alignment) is decisive: on-manifold PSNR jumps to 44.06, chord to 41.40, and the on-manifold-vs-chord rFID gap collapses from 13.4 to 0.3.

Methodology in Plain English

The approach has two independently trained stages.

Stage 1 — building the shared latent space. A DINOv2 ViT-B/14 image encoder is used to turn every modality into a feature map. Because of its final LayerNorm, these features naturally sit on a unit hypersphere. Each modality gets its own decoder (ViT-G/14) that maps a latent back to an image. Decoders are trained in three steps: first reconstruct their own modality from on-sphere latents; second, an expansion step where they are asked to reconstruct from interpolated points between a source and target latent, with weights τ and (1−τ) so responsibility shifts smoothly toward whichever endpoint the interpolant is closer to; third, the encoder is unfrozen and fine-tuned jointly with all decoders, anchored to a frozen copy of the pretrained encoder by a cosine + global-average-pooling alignment loss so it adapts to medical images without drifting from the pretrained geometry.

Stage 2 — learning the transport. With the tokenizer fixed, a Diffusion Transformer (DiTDH-XL) learns a velocity field that maps a source-modality latent to a target-modality latent. Instead of starting from Gaussian noise, the flow starts at the observed latent z_s and the target velocity is simply the constant displacement z_t − z_s, conditioned on a modality indicator so one model handles all directions. At inference, the ODE is integrated starting from z_s (2 Euler steps on BraTS, 4 on SynthRAD) and the result is decoded by the target-modality decoder.

The two "fixes" are what make this work: using the real observed latent as the source removes a huge, irrelevant radial contraction the network would otherwise have to learn, and training decoders on interpolated interior points means the decoder is valid everywhere the flow actually travels.

Why This Matters

Impact on research. The paper shifts attention in cross-modal generation from generator architecture to latent geometry, and gives a concrete, provable criterion (dimension-independent off-manifold deviation controlled by inter-modality similarity) for why a data-as-source flow formulation is preferable to noise-as-source on a manifold. It also offers a reusable three-tokenizer diagnostic (compactness ratio + coefficient of variation of latent norms) for choosing a shared representation. The decoupled two-stage design means the expensive tokenizer is trained once while the flow model can be retrained per modality pair.

Real-world applications:

• Clinical imaging under incomplete protocols — generating a missing MRI contrast (e.g. T1ce) when scanner time, contrast-agent tolerance, or protocol constraints prevent acquisition.
• Radiotherapy planning — SynthRAD-style CT↔MRI translation, where MRI gives soft-tissue contrast but CT is needed for dose calculation.
• Tumor segmentation pipelines — the paper directly demonstrates that imputed modalities feed a 3D U-Net and recover most of the full-modality Dice (0.7473 vs. 0.8554 oracle).
• Remote sensing and multi-sensor fusion — the introduction names remote sensing and autonomous driving as domains where acquiring all modalities is infeasible.

Industry relevance. Any deployed multi-sensor system that must degrade gracefully when a sensor drops out stands to benefit, and the frozen-encoder-plus-lightweight-flow design is attractive for teams that already have strong self-supervised backbones and want to avoid training large paired generators from scratch.

Future Directions

• Modalities without shared visual semantics. The stated limitation: the method relies on DINOv2 providing cross-modal compactness, so pairs like text and imaging would need a domain-adapted encoder or a dedicated cross-modal pretraining stage.
• Extending the geometry-matching principle beyond imaging, which the authors suggest for image-guided speech synthesis and satellite multi-spectral fusion.
• Better nonlinear paths on the sphere. The current interpolant is a straight chord that provably dips inside the sphere; the paper notes Riemannian flow matching with SLERP fixes the path but still departs from noise, so a data-as-source spherical path remains an open combination.
• Generalization beyond 2-modality-pair training — the paper handles N-to-N by averaging available source latents, which invites questions about smarter fusion when many sources of varying reliability are present, and about behavior on modality pairs or anatomies outside BraTS and SynthRAD2023.

Target Audience

Researchers and graduate students working on medical image synthesis, cross-modal generation, or flow/diffusion models on non-Euclidean latent spaces; practitioners building clinical imaging pipelines that must tolerate missing acquisitions; and anyone interested in representation-learning arguments for why a shared, geometrically regular latent space beats per-modality encoders for conditional generation. Readers should be comfortable with diffusion/flow matching terminology and basic linear algebra; the appendix proof is short and accessible.

Authors’ abstract

Missing modality imputation aims to synthesize unobserved modalities from observed ones. Existing methods treat each modality as living on its own latent manifold, forcing the generator to learn a long, unstructured jump between them. We present SphereFlow, which takes a fundamentally different view: by projecting all modalities into a shared hyperspherical latent space via a frozen self-supervised encoder, cross-modal imputation reduces to short-range geometric transport between nearby points on the sphere. Realizing this idea with flow matching, however, exposes two geometric conflicts: the standard Gaussian noise source is catastrophically mismatched with the unit-norm data manifold, and the linear interpolation path departs the sphere surface into a region where the decoder has no training signal. We resolve both with simple, geometry-aware modifications: replacing the noise source with the observed-modality latent so that the velocity field learns only the cross-modal displacement, and training the decoder with a time-weighted loss that extends its domain to the near-surface region traversed by the flow. We prove that this data-as-source formulation reduces the worst-case off-manifold deviation from a quantity that grows with the latent dimension to a small, dimension-independent bound governed by inter-modality similarity. Experiments on BraTS multi-modal MRI imputation and CT-MRI translation demonstrate state-of-the-art generation fidelity across all settings, while downstream brain tumor segmentation with our imputed data closes roughly eighty percent of the gap to the full-modality oracle. Code is available at https://github.com/WTHH031230/SphereFlow.