.

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 fast-acquired 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. Whole-body 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, region-specific bridge modeling and consistently improving translation performance across anatomical regions. Furthermore, to improve pathological fidelity, we exploit 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 lesion-relevant features, improving translation fidelity in pathological areas. Our main contributions include:

• We propose a Heterogeneity-Adaptive Diffusion Schrödinger Bridge (HA-DSB) 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 noise-prediction-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:

\[\small \mathrm{d}X_t = \left[f + g^2 \nabla_x \log \Psi(t, X_t)\right]\mathrm{d}t + g \, \mathrm{d}W_t, \quad \mathrm{d}X_t = \left[f - g^2 \nabla_x \log \hat{\Psi}(t, X_t)\right]\mathrm{d}t + g \, \mathrm{d}W_t.\] where \(X_0\) and \(X_1\) denote the source and endpoint states, and \(\Psi\) and \(\hat{\Psi}\) are time-dependent Schrödinger potentials [17]. We adopt the I\(^2\)SB instantiation proposed in [18], assuming paired samples \((X_0, X_1)\) drawn from the joint distribution \(p(X_0, X_1)=p_A(X_0)p_B(X_1\mid X_0)\) 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 \((X_0, X_1)\) can be expressed as \(q(X_t \mid X_0, X_1) = \mathcal{N}\!\big(X_t;\, \mu_t(X_0, X_1),\, \Sigma_t\big)\). \[\mu_t = \frac{\bar{\sigma}_t^{2}}{\bar{\sigma}_t^{2} + \sigma_t^{2}}\, X_0 + \frac{\sigma_t^{2}}{\bar{\sigma}_t^{2} + \sigma_t^{2}}\, X_1, \qquad \Sigma_t = \frac{\sigma_t^{2}\,\bar{\sigma}_t^{2}}{\sigma_t^{2} + \bar{\sigma}_t^{2}}\, \mathbf{I} \label{eq:dsb95posterior}\tag{1}\] where \(\sigma_t^{2} \mathrel{\vcenter{:}}= \int_{0}^{t} \beta_{\tau}\,\mathrm{d}\tau\) and \(\bar{\sigma}_t^{2} \mathrel{\vcenter{:}}= \int_{t}^{1} \beta_{\tau}\,\mathrm{d}\tau\) 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.

\[\mathcal{L}(\theta)= \mathbb{E}\Big[\big\|\varepsilon_\theta(X_t,c,t)-\tfrac{X_t-X_0}{\sigma_t}\big\|_2^2\Big].\]

2.2 Region context embedding guidance.↩︎

To address anatomical heterogeneity in whole-body translation, we introduce a region context embedding \(\mathbf{c}_{\text{reg}} \in \mathbb{R}^d\) that adapts the bridge dynamics to region-dependent structural variability. For each slice, we assign a body-location label \(l_{\text{bp}}\) 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 \(\mathbf{e}_{\text{bp}} \in \mathbb{R}^d\).

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 organ embeddings \(\{\mathbf{e}_{\text{org}}^i\}_{i=1}^{M}\), where \(\mathbf{e}_{\text{org}}^i \in \mathbb{R}^d\) and \(M\) is the number of organs present in the slice. We fuse body-location and organ information via multi-head cross-attention, where \(\mathbf{e}_{\text{bp}}\) 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 \(\mathbf{c}_{\text{reg}}\): \[\small \alpha_i = \mathrm{Softmax} \left( \frac{ (\mathbf{W}_Q \mathbf{e}_{\text{bp}})^\top (\mathbf{W}_K \mathbf{e}_{\text{org}}^i) }{ \sqrt{d_h} } \right), \qquad \mathbf{e}_{\text{attn}} = \sum_{i=1}^{M} \alpha_i \, \mathbf{W}_V \mathbf{e}_{\text{org}}^{i},\] \[\small \label{eq:residual95connection} \mathbf{c}_{\text{reg}} = \mathbf{e}_{\text{bp}} + \mathrm{MLP}\!\left(\mathrm{LayerNorm}(\mathbf{e}_{\text{attn}})\right),\tag{2}\] where \(\mathbf{W}_Q,\mathbf{W}_K,\mathbf{W}_V \in \mathbb{R}^{d_h \times d}\) are learnable projections, with \(d_h=d/H\) and \(H=4\). The resulting \(\mathbf{c}_{\text{reg}}\) encodes both coarse body-location and fine-grained organ information. It is added to the diffusion timestep embedding \(\mathbf{e}_t \in \mathbb{R}^d\) 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 body-location 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.

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 \(\mathbf{c}_{\text{reg}}\). Specifically, we use a lightweight noise modulator \(\mathcal{M}_\phi\) to predict a spatial scaling map \(\mathbf{S} \in [s_{\min}, s_{\max}]^{H \times W}\). The modulator extracts spatial features from PET using a two-layer convolutional encoder, incorporates anatomical context via FiLM [20] conditioning with \(\mathbf{c}_{\text{reg}}\), and outputs per-pixel scaling factors. We then replace the standard uniform-noise bridge forward process \[\mathbf{x}_t = \mu_t(X_0, X_1) + \Sigma_t^{1/2} \boldsymbol{\epsilon}, \qquad \boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}),\] with the spatially modulated form, where \(\odot\) denotes element-wise multiplication. \[\mathbf{x}_t = \mu_t(X_0, X_1) + \Sigma_t^{1/2} (\mathbf{S} \odot \boldsymbol{\epsilon}),\] This yields a heterogeneity-adaptive forward corruption process, allowing anatomically and pathologically complex regions to receive different perturbation strengths.

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 \(\mathbf{x} \in \mathbb{R}^{B \times C \times H \times W}\), the PET image is aligned by adaptive average pooling and projected by a \(1 \times 1\) convolution to obtain \(\mathbf{p}'\). We concatenate \(\mathbf{x}\) and \(\mathbf{p}'\) to form \(\mathbf{z} = [\mathbf{x}; \mathbf{p}']\), and apply PET-aware self-attention: \[\small \mathbf{Q}=W_Q\tilde{\mathbf{z}},\; \mathbf{K}=W_K\tilde{\mathbf{z}},\; \mathbf{V}=W_V\tilde{\mathbf{z}},\qquad \mathbf{x}_{\mathrm{out}} = \mathbf{x} + W_{\mathrm{out}} \!\left[ \mathrm{Softmax}\!\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_h}}\right)\mathbf{V} \right],\] where \(\tilde{\mathbf{z}} \in \mathbb{R}^{B \times HW \times 2C}\) is the flattened form of \(\mathbf{z}\). We use a zero-initialized \(W_{\mathrm{out}}\) so that the block starts as an identity mapping and gradually incorporates PET-guided lesion cues during training.

3 Experiments↩︎

1.0pt max width=1.03,center

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. Each subject contains three co-registered modalities (LAVA, T2, and PET), all resampled to 256×256 resolution. The dataset was split into a training set of 204 subjects (20,306 slices) and a held-out testset of 42 subjects (4,162 slices). The testset comprises 21 randomly sampled subjects and an additional 21 subjects with clinically confirmed lesions, allowing us to specifically assess synthesis fidelity in pathological regions.

Implementation Details. For the denoising network, we employ an ADM-style U-Net [21] featuring \(128\) base channels, channel multipliers of \([1, 1, 2, 2, 4, 4]\), and three residual blocks per stage. PET attention blocks are integrated at resolutions of \(32 \times 32\), \(16 \times 16\), and \(8 \times 8\) within both the encoder and decoder. The noise modulator \(\mathcal{M}_\phi\) has \(s_{\min} = 0.2\) and \(s_{\max} = 2.0\). The diffusion bridge uses \(T = 200\) training discretization steps with \(\beta_{\max} = 0.35\) and a symmetric variance schedule. Training was conducted on two NVIDIA RTX 4090 GPUs for \(50\text{k}\) iterations. We utilized the AdamW optimizer with a learning rate of \(10^{-4}\), a total batch size of \(16\), and an EMA decay of \(0.9998\).

Comparison with SOTA methods. We compare HA-DSB against five representative baselines from GAN-based, diffusion-based, and bridge paradigms, all trained on whole-body data. As shown in Table ¿tbl:tab:per-region32quantitative32results32on32the32testset?, HA-DSB achieves the best average SSIM and PSNR, outperforming the strongest baseline, ResViT. Compared with the basic diffusion bridge baseline I\(^2\)SB, HA-DSB achieves consistent improvements across all anatomical regions, clearly demonstrating the effectiveness of the proposed region-aware conditioning mechanism. Simpler regions, such as Head/Neck, Pelvis/Hips, and Thighs, achieve SSIM above 93%, while structurally complex regions such as the Abdomen also show clear gains. These results suggest that region guidance helps the bridge learn region-dependent transport paths. Qualitative comparisons in Fig. 2 further support this observation.

2.6pt

Effect of PET guidance.PET guidance (HA-DSB(pet) vs.HA-DSB(no pet)) brings only modest gains in overall SSIM and PSNR on the testset (21 patients, 2,087 slices). This is likely because lesion regions occupy only a small fraction of the total volume, and their improvements are diluted by the dominant healthy tissue where PET provides limited additional information. To better isolate the impact of PET guidance on pathological regions, we further evaluate performance on a separate lesion-confirmed cohort of 21 patients. A radiologist identified slices covering disease-bearing organs, resulting in 526 lesion-containing slices out of 2,075 (Table 1). On this pathological subset, PET guidance yields markedly larger improvements in both PSNR and SSIM than on the testset, indicating that the noise modulation and multi-scale attention primarily enhance lesion-related regions. Fig. 3 further illustrates two representative cases.

4 Conclusion↩︎

In this work, we proposed HA-DSB for PET-guided whole-body MR translation. By combining region context embeddings, PET-guided noise modulation, and multi-scale PET-aware attention, HA-DSB addresses anatomical and pathological heterogeneity in whole-body PET/MR. Experiments show improved translation quality across body regions, with additional gains in lesion-containing areas. These results indicate that heterogeneity-adaptive conditioning is effective for modeling region-dependent structural variability and recovering lesion-relevant details in challenging whole-body settings. More broadly, HA-DSB provides a practical bridge-based framework for robust whole-body medical image translation.

Disclosure of Interests↩︎

The authors have no competing interests to declare that are relevant to the content of this article.

References↩︎