Heterogeneity-Adaptive Diffusion Schrödinger Bridge for PET-Guided Whole-Body MRI Translation Chengbo Wang1† , Jiacheng Yu1† , Linjie Bian2† , Ming Qi2 , Xiaosheng Liu2 , Tongtong Che3 , Jichang Zhang3 , Shuyu Li3 , and Shaoli Song2* and Xiuying Wang1* The University of Sydney, Sydney, Australia Department of Nuclear Medicine, Fudan University Shanghai Cancer Center, Shanghai, China 3 State Key Laboratory of Cognitive Neuroscience and Learning, Beijing Normal University, Beijing, China 1
arXiv:2607.07401v1 [cs.CV] 8 Jul 2026
2
. Abstract. While whole-body multimodal medical imaging scanners have been increasingly recognized for more effective medical applications, the excessive long acquisition time in PET–MR scanning is a major obstacle in more efficient clinical practice. Deep learning–based MRI translation provides a potential solution to reduce scan duration. However, current models often focus on specific anatomical regions and face challenges for whole-body scans that consists of highly heterogeneous feature distributions mainly due to (1) different anatomical regions across whole-body, and (2) lesions or pathological tissues. This paper tackles the challenges through a novel Heterogeneity-Adaptive Diffusion Schrödinger Bridge (HA-DSB) framework. By explicitly modeling translation as stochastic transport between source and target distributions, HA-DSB incorporates region context embeddings derived from a vision-language model (VLM) to enable region-specific modeling. To enhance fidelity of the pathological tissue, lesion-aware metabolic prior from PET is integrated directly into the bridge dynamics through a dual-stage guidance mechanism. Specifically, a PET-guided noise modulation module adaptively scales spatial diffusion perturbations during the forward process, while PET features are leveraged during the reverse process to selectively amplify lesion-relevant structures via an attention mechanism. Experiments demonstrate the superiority of our method across different body regions in whole-body MRI translation and show improved translation quality in lesion areas under PET guidance. Our code is available at https://github.com/xyw-medical-research/HADSB. Keywords: Medical Image Translation · PET/MR Imaging · Diffusion Schrödinger Bridge · Whole-body MR Translation contributed equally to this work. [email protected], [email protected], [email protected] * corresponding authors. †
2
C. Wang et al.
1
Introduction
Hybrid positron emission tomography combined with magnetic resonance imaging (PET/MR) is an emerging imaging modality that enables simultaneous acquisition of metabolic and high-resolution anatomical information, producing intrinsically co-registered PET and multi-sequence MR images while reducing radiation exposure compared with PET/CT [1]. Whole-body PET/MR supports comprehensive multi-bed imaging from the head and neck to the mid-thigh and has become increasingly important in clinical applications such as cancer metastasis evaluation [2]. However, whole-body PET/MR scan time typically ranges from 60–90 minutes, largely due to multiple MR sequences, which limits clinical efficiency and increases patient discomfort. Deep learning–based MR translation offers a practical solution by synthesizing time-consuming modalities from fastacquired sequences, and recent diffusion-based approaches have demonstrated promising performance in MR image translation [3–6]. However, most MR translation focus on specific anatomical regions such as the brain in public datasets [7], and extending these models to whole-body MR translation in integrated PET/MR settings remains challenging. Wholebody scans exhibit strong cross-region heterogeneity, with substantial variations in tissue composition and signal characteristics across anatomical regions [8]. Such non-stationary and region-dependent intensity distributions significantly increase the difficulty of learning a unified modality mapping [9, 10]. For example, GANs are known to suffer from mode collapse in multi-modal distributions [11, 12], while diffusion models may be biased toward dominant feature modes [13]. Without explicit region conditioning, these models may struggle to capture region-specific structures in the whole body MR translation. Another critical yet underexplored challenge in MR translation is the accurate reconstruction of pathological regions [14, 15], which are often of primary clinical importance. Compared to normal tissues, lesions exhibit highly heterogeneous appearances and complex cross-modality signal relationships. Most existing models do not explicitly differentiate pathological changes from normal tissues, causing lesion signals to be biased toward the dominant healthy distribution and leading to degraded pathological fidelity. In this work, we propose a Heterogeneity-Adaptive Diffusion Schrödinger Bridge (HA-DSB) with PET guidance for whole-body MR translation (Fig. 1). Unlike conventional diffusion-based translation methods that rely on implicit noise prediction, HA-DSB explicitly optimizes a stochastic transport plan between source and target MR distributions [3, 16]. This explicit transport formulation is better suited for whole-body MR translation, where strong cross-region domain shifts make implicit alignment prone to over-smoothing and modality ambiguity. To address the pronounced region-dependent distribution shifts in whole-body MR translation, we further introduce region context embeddings derived from structured descriptions generated by the large vision-language model and fuse them with diffusion time embeddings, enabling fine-grained, regionspecific bridge modeling and consistently improving translation performance across anatomical regions. Furthermore, to improve pathological fidelity, we ex-
Heterogeneity-Adaptive Schrödinger Bridge
3
ploit PET as a lesion-aware prior, which is naturally available in integrated PET/MR due to intrinsic co-registration. Since lesions typically manifest as focal high-uptake signals in PET, we introduce a PET-guided noise modulation mechanism that applies spatially adaptive corruption in the forward process. This spatially adaptive perturbation redistributes the stochastic transport difficulty, encouraging the model to learn pathology-sensitive mappings rather than being biased toward dominant healthy structures. A multi-scale PET-guided attention module is further incorporated into the reverse process to amplify lesionrelevant features, improving translation fidelity in pathological areas. Our main contributions include: – We propose a Heterogeneity-Adaptive Diffusion Schrödinger Bridge (HADSB) that models whole-body MR translation as an explicit distribution transport problem, achieving more accurate and stable cross-modality mapping across anatomically diverse regions compared to conventional noiseprediction-based diffusion models. – We introduce region context embeddings generated by a large vision-language model and fuse them with diffusion time embeddings as conditioning signals, enabling the bridge to adapt its region-specific transport dynamics. – We design a PET-guided lesion enhancement mechanism in which the forward corruption process is spatially modulated by PET uptake and region context, while multi-scale PET-aware attention improves lesion-relevant feature recovery during reverse denoising.
2
Methodology
2.1
Diffusion Schrödinger Bridge.
Diffusion Schrödinger bridge (DSB) [17, 18] constructs a stochastic process that transports samples between two endpoint distributions. In image-to-image translation, the bridge provides an endpoint-to-endpoint alternative to the usual Gaussian-prior diffusion models. The forward and backward process can be described as: dXt = f + g 2 ∇x log Ψ (t, Xt ) dt + g dWt ,
h i dXt = f − g 2 ∇x log Ψ̂ (t, Xt ) dt + g dWt . (1)
where X0 and X1 denote the source and endpoint states, and Ψ and Ψ̂ are timedependent Schrödinger potentials [17]. We adopt the I2 SB instantiation proposed in [18], assuming paired samples (X0 , X1 ) drawn from the joint distribution p(X0 , X1 ) = pA (X0 )pB (X1 | X0 ) and setting the drift term f = 0. Under this formulation, the Schrödinger Bridge can be efficiently trained and inferred within the score-based generative modeling (SGM) framework. The analytic posterior of DSB given a boundary pair (X0 , X1 ) can be expressed as q(Xt | X0 , X1 ) = N Xt ; µt (X0 , X1 ), Σt . µt =
σ2 σ̄t2 X0 + 2 t 2 X1 , 2 2 σ̄t + σt σ̄t + σt
Σt =
σt2 σ̄t2 I σt2 + σ̄t2
(2)
4
C. Wang et al.
Fig. 1. Overview of HA-DSB. (a) Whole-body MR translation setting. (b) Diffusion Schrödinger bridge from LAVA to T2 with noise modulation jointly conditioned on PET and region context creg . (c) UNet translation backbone with region-conditioned residual blocks and multi-scale PET attention. (d) Construction of the region context embedding creg .
Rt R1 where σt2 := 0 βτ dτ and σ̄t2 := t βτ dτ denote the variances accumulated on the two sides of time t. Unlike conventional image-to-image Schrödinger Bridge models, our method adopts a heterogeneity-adaptive conditional formulation that integrates region context embedding and PET-derived cues. The final training objective is defined as follows, where c represents the conditioning information. h i 2 0 L(θ) = E εθ (Xt , c, t) − Xtσ−X . (3) 2 t 2.2
Region context embedding guidance.
To address anatomical heterogeneity in whole-body translation, we introduce a region context embedding creg ∈ Rd that adapts the bridge dynamics to regiondependent structural variability. For each slice, we assign a body-location label lbp that partitions the body into K = 11 regions. A frozen PubMedBERT [19] encodes these labels into semantic features, which are linearly projected and combined with learnable embeddings to obtain the body-location representation ebp ∈ Rd . To capture intra-region variability, each slice is additionally associated with a set of organ labels, which are encoded by the same frozen encoder to produce i d organ embeddings {eiorg }M i=1 , where eorg ∈ R and M is the number of organs
Heterogeneity-Adaptive Schrödinger Bridge
5
present in the slice. We fuse body-location and organ information via multi-head cross-attention, where ebp serves as the query and organ embeddings act as the keys and values. The organ information is then added to the body-location representation as a residual correction, forming the final region context embedding creg : αi = Softmax
(WQ ebp )⊤ (WK eiorg ) √ dh
! ,
eattn =
M X
αi WV eiorg ,
(4)
i=1
creg = ebp + MLP(LayerNorm(eattn )) ,
(5)
where WQ , WK , WV ∈ R are learnable projections, with dh = d/H and H = 4. The resulting creg encodes both coarse body-location and fine-grained organ information. It is added to the diffusion timestep embedding et ∈ Rd and passed through an MLP to generate the scale and shift parameters in each residual block. To avoid manual annotation, we use Google Gemini 3 Pro to generate bodylocation and organ labels. The reliability of body-location labels was empirically verified, achieving approximately 96% accuracy under professional validation. For organ labels, we constrain the VLM outputs to a predefined vocabulary of major MRI-visible organs to improve reliability and avoid unconstrained predictions. dh ×d
2.3
PET-guided Noise Modulation
Standard diffusion bridges inject spatially uniform noise across all spatial locations [18], which is suboptimal for whole-body PET/MR translation because anatomical complexity and lesion-related variability are highly heterogeneous across regions. To address this, we introduce a PET-guided noise modulation mechanism that jointly conditions the forward corruption process on PET uptake and the region context embedding creg . Specifically, we use a lightweight noise modulator Mϕ to predict a spatial scaling map S ∈ [smin , smax ]H×W . The modulator extracts spatial features from PET using a two-layer convolutional encoder, incorporates anatomical context via FiLM [20] conditioning with creg , and outputs per-pixel scaling factors. We then replace the standard uniformnoise bridge forward process 1/2
xt = µt (X0 , X1 ) + Σt
ϵ ∼ N (0, I),
ϵ,
(6)
with the spatially modulated form, where ⊙ denotes element-wise multiplication. 1/2
xt = µt (X0 , X1 ) + Σt
(S ⊙ ϵ),
(7)
This yields a heterogeneity-adaptive forward corruption process, allowing anatomically and pathologically complex regions to receive different perturbation strengths.
6
C. Wang et al.
Table 1. Per-region quantitative results on the testset (mean ± std). Each cell shows SSIM (%, top) and PSNR (dB, bottom). Method
Head/Neck
Thorax
n=275
n=498
Abdomen Pelvis/Hips n=617
n=472
Thighs
Avg
n=225
n=2,087
General translation models 86.5 ± 6.3 79.9 ± 4.5 80.9 ± 5.5 83.5 ± 5.0 85.7 ± 3.7 82.5 ± 5.6 Pix2Pix[22] 22.77 ± 1.64 22.31 ± 1.89 21.58 ± 1.75 22.01 ± 2.12 23.52 ± 1.93 22.22 ± 1.97 89.6 ± 6.9 85.4 ± 4.0 85.6 ± 5.3 88.1 ± 5.1 89.8 ± 4.2 87.1 ± 5.4 Palette[23] 25.83 ± 4.50 23.83 ± 2.86 23.61 ± 2.90 24.59 ± 3.09 25.44 ± 4.03 24.37 ± 3.41 88.0 ± 6.2 82.2 ± 3.8 84.0 ± 3.6 85.5 ± 2.9 85.7 ± 2.7 84.6 ± 4.3 I 2 SB [18] 25.81 ± 1.91 21.67 ± 1.74 21.50 ± 1.31 23.06 ± 1.27 24.38 ± 2.13 22.77 ± 2.20 Medical translation models 89.6 ± 4.4 81.1 ± 3.5 79.5 ± 4.1 79.1 ± 3.2 76.5 ± 2.8 80.8 ± 5.2 PPT[24] 23.91 ± 1.44 21.42 ± 1.21 20.43 ± 1.60 20.93 ± 1.59 22.37 ± 2.09 21.45 ± 1.92 84.9 ± 7.2 84.6 ± 7.1 84.7 ± 6.6 84.9 ± 6.5 85.7 ± 6.4 84.9 ± 6.8 SelfRDB[6] 22.12 ± 2.59 22.02 ± 2.70 21.95 ± 2.63 22.10 ± 2.65 22.27 ± 2.06 22.06 ± 2.59 91.0 ± 5.5 85.4 ± 3.8 85.3 ± 4.1 87.3 ± 4.5 86.0 ± 3.3 86.6 ± 4.7 ResViT[5] 26.35 ± 2.13 23.76 ± 2.06 22.89 ± 1.93 23.74 ± 2.42 24.83 ± 2.16 23.95 ± 2.39 Our method
94.0† ± 3.1 89.0† ± 4.3 88.1† ± 4.8 93.1† ± 3.9 93.3† ± 3.0 90.8† ± 4.8 28.59† ± 2.36 24.87† ± 2.44 24.06† ± 2.81 27.42† ± 2.95 28.51† ± 2.44 26.09† ± 3.24 94.5 ± 3.5 89.5 ± 4.2 88.5 ± 4.8 93.2 ± 4.0 93.6 ± 3.2 91.1 ± 4.8 HA-DSB(PET) 28.88 ± 2.48 24.92 ± 2.46 24.15 ± 2.71 27.59 ± 3.23 28.89 ± 2.49 26.25 ± 3.35 HA-DSB(no PET)
†
Second-best result (SSIM and PSNR evaluated separately).
2.4
Multi-Scale PET Attention
Because the forward process applies spatially non-uniform corruption guided by PET and anatomical context, the reverse model must recover lesion-relevant details under the same setting. To this end, we introduce PET-aware self-attention at multiple UNet resolutions. For each feature map x ∈ RB×C×H×W , the PET image is aligned by adaptive average pooling and projected by a 1×1 convolution to obtain p′ . We concatenate x and p′ to form z = [x; p′ ], and apply PET-aware self-attention: Q = WQ z̃, K = WK z̃, V = WV z̃,
QK⊤ xout = x + Wout Softmax √ V , (8) dh
where z̃ ∈ RB×HW ×2C is the flattened form of z. We use a zero-initialized Wout so that the block starts as an identity mapping and gradually incorporates PETguided lesion cues during training.
3
Experiments
Dataset. All PET/MR examinations were performed on an integrated hybrid PET/MR system (SIGNA™ PET/MR, GE HealthCare). This platform integrates a 3.0-T MRI scanner with time-of-flight PET capability. We prospectively collect whole-body imaging data from 246 subjects at an anonymous hospital.
Heterogeneity-Adaptive Schrödinger Bridge