Skip to content

Heterogeneity-Adaptive Diffusion Schrodinger Bridge for PET-Guided Whole-Body MRI Translation

🕒 Published (v1): 2026-07-08 13:35 UTC · Source: Arxiv · Venue: MICCAI 2026 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

HA-DSB is a Diffusion Schrödinger Bridge framework for translating LAVA MRI to T2-weighted MRI across the whole body in integrated PET/MR scanners. It addresses cross-region anatomical heterogeneity via VLM-derived region context embeddings and improves pathological fidelity via PET-guided noise modulation and multi-scale attention. It outperforms all baselines across five body regions on a 246-subject clinical dataset.

Problem

Existing MRI translation models target single anatomical regions (predominantly brain) and fail on whole-body scans due to two compounding heterogeneities: (1) large cross-region variation in tissue composition and intensity distributions, causing standard diffusion models to be biased toward dominant feature modes; and (2) lesion regions exhibiting highly atypical signal characteristics that get suppressed toward the healthy-tissue distribution, degrading clinical utility in pathological cases.

Method

HA-DSB builds on the I²SB (Image-to-Image Schrödinger Bridge) formulation, which optimizes a stochastic transport plan between paired source (\(X_0\)) and target (\(X_1\)) distributions rather than denoising from Gaussian noise. Three key components are added:

  1. Region context embeddings (\(c_\text{reg}\)): Google Gemini 3 Pro generates body-location (K=11 regions) and organ labels per slice. A frozen PubMedBERT encodes these; body-location embedding \(e_\text{bp}\) and organ embeddings \(\{e^i_\text{org}\}\) are fused via multi-head cross-attention (H=4 heads), then \(c_\text{reg} = e_\text{bp} + \text{MLP}(\text{LayerNorm}(e_\text{attn}))\). This embedding is summed with the diffusion timestep embedding and used to generate FiLM scale/shift parameters in each residual block.

  2. PET-guided noise modulation (forward process): A lightweight convolutional modulator \(M_\phi\) predicts a per-pixel spatial scaling map \(S \in [s_\text{min}, s_\text{max}]^{H \times W}\) conditioned on PET uptake and \(c_\text{reg}\) via FiLM. The standard bridge forward sample \(x_t = \mu_t(X_0, X_1) + \Sigma_t^{1/2}\epsilon\) is replaced by \(x_t = \mu_t(X_0, X_1) + \Sigma_t^{1/2}(S \odot \epsilon)\), concentrating stochastic difficulty on lesion-rich and anatomically complex areas.

  3. Multi-scale PET attention (reverse process): At UNet resolutions \(32\times32\), \(16\times16\), \(8\times8\), PET is pooled to match feature map size and concatenated with features \(x\) to form \(z = [x; p']\). Self-attention is applied over \(z\) with a zero-initialized output projection \(W_\text{out}\), so the block initializes as an identity and progressively learns PET-guided lesion cue incorporation.

Key Contributions

  • Frames whole-body MRI translation as an explicit distribution-transport Schrödinger Bridge, avoiding over-smoothing from implicit noise-prediction approaches under strong cross-region domain shifts.
  • Region context embeddings from a frozen VLM (PubMedBERT-encoded labels generated by Gemini 3 Pro) injected into bridge dynamics via FiLM, enabling fine-grained region-specific transport.
  • PET-guided spatially adaptive noise modulation in the forward process to redistribute stochastic difficulty toward pathologically complex regions.
  • Multi-scale PET-aware self-attention in the reverse process with zero-initialized output gating to enhance lesion-relevant feature recovery.

Results

  • Overall (n=2,087 test slices, 42 subjects): HA-DSB(PET) achieves mean SSIM 91.1% / PSNR 26.25 dB, vs. best baseline ResViT at 86.6% / 23.95 dB; second best is Palette at 87.1% / 24.37 dB.
  • Per-region SSIM: Head/Neck 94.5%, Thorax 89.5%, Abdomen 88.5%, Pelvis/Hips 93.2%, Thighs 93.6% — all best in class.
  • Lesion-confirmed cohort (526 lesion-containing slices from 21 subjects): PET guidance yields \(\Delta\)PSNR = +0.80 dB and \(\Delta\)SSIM = +1.7% over no-PET variant, compared to only +0.16 dB / +0.3% on the full test set, confirming lesion-specific benefit.
  • Baselines compared: Pix2Pix, Palette, I²SB (general), PPT, SelfRDB, ResViT.

Limitations

  • Dataset is single-center (246 subjects, one PET/MR scanner model: SIGNAâ„¢ GE), limiting generalizability.
  • VLM-generated organ labels are constrained to a predefined vocabulary, which may miss atypical or rare organ configurations.
  • Body-location label accuracy was verified at ~96%, meaning ~4% of slices carry erroneous region conditioning.
  • PET guidance provides only modest aggregate improvement because lesion regions are a small fraction of total slice volume; per-region gains on non-lesion slices are negligible.
  • Evaluation is restricted to LAVA→T2 translation; generalization to other MR sequence pairs is not demonstrated.
  • No perceptual or radiologist reader study for clinical acceptability; evaluation relies solely on SSIM and PSNR.

Relevance to Vision-Language Models

VLMs enter this work as a structured annotation oracle: Google Gemini 3 Pro generates body-location and organ labels from MRI slices, which are then encoded by a frozen biomedical language model (PubMedBERT) to produce semantic conditioning signals for the diffusion bridge. This represents a practical deployment pattern where a large VLM replaces manual expert annotation at ~96% label accuracy, feeding a downstream generative model rather than acting as the generator itself. For researchers tracking VLMs, this paper is evidence that frozen VLM-derived embeddings can function as reliable spatial priors in medical imaging pipelines where region semantics are structurally meaningful but prohibitively expensive to annotate at scale.