Volumetric Latent Diffusion for Brain Tumour Segmentation with Missing Modalities

Introduction

Segmenting gliomas on magnetic resonance imaging (MRI) supports diagnosis, surgical planning and the assessment of treatment response. The standard protocol acquires four complementary sequences: T1, contrast-enhanced T1 (T1CE), T2 and FLAIR. No single one reveals the whole lesion: enhancement marks the active core, while oedema is clearest on FLAIR. Manual delineation is slow and varies between readers, making automated segmentation a central problem [9]. Deep encoder–decoder architectures [11, 6] perform well on public benchmarks, but the complete protocol is often unavailable in practice. Methods for incomplete multimodal segmentation [4, 3, 18, 13] maintain accuracy on partial inputs, yet share an assumption that is rarely stated: the output should be a single deterministic mask. Their evaluations report how accuracy degrades as sequences are removed, but the prediction is returned in the same form either way, giving the clinician no indication of the model’s confidence. Generative methods model a distribution over plausible masks rather than one answer [7, 2, 17], which suits a target whose boundary is genuinely ambiguous. Diffusion formulations are the strongest of these, but each sample costs many network evaluations at full resolution, so they have largely been confined to two dimensions; while brain tumour segmentation is inherently volumetric. Moving the process into a learned latent space removes that cost, and pairing it with modality dropout lets one model serve every acquisition protocol. A sampling formulation also gives predictive uncertainty for free: repeated draws under the same input disagree where the model is unsure, and evaluating each case under every subset makes that comparison paired, so lesion difficulty cancels.

In this paper, we present a three-dimensional conditional latent diffusion framework for whole-tumour segmentation from incomplete multi-sequence MRI, and use it to test this behaviour. Our contributions are:

1. Volumetric latent diffusion with modality dropout. We extend latent diffusion segmentation to 3D multi-sequence MRI, training with randomly resampled modality availability so one model serves all fifteen configurations rather than one per protocol.

2. Segmentation under every acquisition configuration. We report whole-tumour accuracy for all fifteen subsets of the four sequences on 2,601 studies from BraTS 2023 and 2024.

3. Predictive uncertainty as a by-product of sampling. The same repeated draws that form the prediction also give a voxel-wise disagreement map, which we report for every configuration and show to rise as sequences are withdrawn.

Related Work

Incomplete Multimodal Tumour Segmentation

Clinical acquisitions frequently deviate from the four-sequence protocol that benchmarks assume, and the literature responds in three ways. Synthesis generates the absent sequence and passes a completed set to a conventional network, which leaves the segmentation architecture unchanged but makes accuracy dependent on the quality of the synthesis. Modality-invariant representations avoid that dependence: HeMIS [4] pools per-modality features using statistics that are defined for any subset, and RFNet [3] adds region-aware fusion that weights sequences differently for different tumour components, since enhancement and oedema are visible in different sequences. Distillation transfers knowledge from a full-modality teacher to a student that observes only a subset; mmFormer [18] combines modality-specific encoders with inter-modality attention to model long-range dependencies across sequences. UniME [13] separates representation learning from fusion, pretraining a shared encoder with masked self-supervision before any fusion module is attached, so that fusion need not simultaneously learn what each sequence contributes and how to combine them. These methods differ in mechanism but agree in output: each returns a single deterministic mask, whatever subset it was given. Accuracy is reported as a function of the available sequences; the confidence attached to the prediction is not. How a model’s expressed confidence should respond when evidence is withdrawn is, to our knowledge, not addressed in this literature.

Generative Approaches to Segmentation

Discriminative networks such as U-Net [11] and its successors [6] learn a mapping that is one-to-one by construction, so the disagreement observed between expert annotators cannot appear in the output. Latent-variable models relax this by decoding samples of a stochastic code into different masks: the Probabilistic U-Net [7] draws a global latent from a conditional prior, and PHiSeg [2] extends the idea to a hierarchy of latents at several resolutions. Both are limited by the capacity of that latent and by the tendency of the KL term to make the decoder ignore it. Diffusion models [5] build the distribution through iterative refinement instead of a single stochastic bottleneck. Conditioning a denoising network on the input image turns generation into segmentation, and repeated sampling gives an ensemble from one set of weights; MedSegDiff [17] strengthens the conditioning with frequency-domain filtering of the guidance signal. The cost is many full-resolution network evaluations per sample, which is acceptable in two dimensions and prohibitive in three. Latent diffusion [10] removes that cost by working in a compressed space, an approach applied to medical segmentation by LSegDiff [16]. MedSegLatDiff [14] pairs two autoencoders with a latent diffusion model and shows that a foreground-weighted reconstruction loss is needed when targets are small, while LatentFM [15] substitutes flow matching to shorten sampling. These methods operate on two-dimensional images with all inputs present, and evaluate on datasets whose ambiguity is fixed by the annotations. The present work extends the line to volumetric data and to inputs that are deliberately incomplete.

Method

We propose a volumetric latent diffusion framework for segmentation with missing modalities. The model conditions on the available modalities and performs diffusion in a compact latent space, while multiple generated segmentation samples provide an estimate of predictive uncertainty. Figure 1 illustrates the overall framework.

Vector (PDF)

Problem Formulation

In this problem setting, we consider a medical image segmentation dataset consisting of paired multimodal observations and voxel-wise annotations. Let

denote the dataset, where $X^{(i)} \in\mathbb{R}^{H \times W \times D \times M}$ is a co-registered volume containing $M$ imaging modalities, with $H$, $W$, and $D$ denoting the height, width, and depth in voxels, respectively, and $S^{(i)} \in\{0,1\}^{H \times W \times D}$ is the corresponding voxel-wise expert annotation. We consider binary whole-tumour segmentation, where the target $S$ defines a single foreground region.

Modality Availability In routine practice, one or more imaging modalities may be missing due to incomplete acquisition, motion corruption, or differences in acquisition protocols. We explicitly represent the available modalities using a binary vector

where $m_j = 1$ indicates that the $j$-th modality is available and $m_j = 0$ indicates that it is missing. This gives $2^{M} - 1$ possible non-empty modality configurations. Given an availability vector $\mathbf{m}$, the observed input is defined as

where $\odot$ denotes element-wise multiplication with $\mathbf{m}$ broadcast across the spatial dimensions, such that missing modalities are represented by zero-valued volumes.

Objective We formulate segmentation as a probabilistic generative task rather than a deterministic mapping from the observed input to a single segmentation. Following our previous work [14], we model the conditional distribution

which describes the distribution of plausible segmentations given the observed input $X$. When modalities are missing, the observed input is represented by $X_{\mathrm{m}}$. We therefore learn a single conditional generative model that captures the segmentation distribution across different modality configurations:

At inference time, multiple samples can be generated for the same incomplete input, whose variation provides an estimate of predictive uncertainty under different levels of modality missingness.

Volumetric Perceptual Compression

To make volumetric diffusion computationally tractable, we perform the generative process in a compact latent space following the latent diffusion paradigm [10]. Both the multimodal input and the segmentation annotation are transformed from the original voxel space into lower-dimensional latent representations. For a volumetric input with spatial dimensions $H \times W \times D$, the corresponding latent representation has spatial dimensions $h \times w \times d$, with the spatial downsampling factor defined as

The latent representation therefore has the form $z \in\mathbb{R}^{h \times w \times d \times c}$, where $c$ denotes the number of latent channels. This spatial compression substantially reduces the number of spatial positions processed by the diffusion model, making volumetric generation computationally feasible.

Multimodal Autoencoder The multimodal encoder $\mathcal{E}_{X}$ maps the incomplete multimodal input $X_{\mathrm{m}}$ into a compact latent representation. To impose a probabilistic structure on the latent space, the encoder parameterizes the latent representation as a Gaussian posterior

where $\mu_{X}$ and $\sigma_{X}^{2}$ are the mean and variance predicted by the encoder. We adopt a standard Gaussian prior $p(z_{X}) = \mathcal{N}(0, I)$ and regularize the posterior toward this prior to obtain a smooth and well-structured latent space for subsequent diffusion modeling. A latent sample $z_{X}$ is passed through the decoder $\mathcal{G}_{X}$ to reconstruct the input. The autoencoder is trained using the ELBO-based objective

where the first term is the MSE reconstruction loss that encourages the latent representation to preserve information from the input, while the second term regularizes the posterior toward the Gaussian prior. The coefficient $\beta_X$ controls the trade-off between reconstruction fidelity and latent-space regularization. After training, the posterior mean is used as the multimodal latent, $z_X = \mu_X$, which serves as the conditioning representation for the diffusion model.

Mask Autoencoder. The mask encoder $\mathcal{E}_S$ similarly maps the binary whole-tumour segmentation mask $S$ to a Gaussian posterior

with the same standard Gaussian prior $p(z_S) = \mathcal{N}(0, I)$. A latent sample $z_S$ is decoded by $\mathcal{G}_S$ to reconstruct the segmentation mask. Since $S$ is binary with a highly imbalanced foreground and background, we use a weighted binary cross-entropy reconstruction loss instead of the MSE used for the multimodal input. The resulting ELBO-based objective is

where $\hat{S} = \mathcal{G}_S(z_S)$ is the reconstructed mask and $w_+ > 1$ assigns greater importance to foreground voxels. The coefficient $\beta_S$ controls the trade-off between reconstruction fidelity and latent-space regularization. After training, the posterior mean is used as the mask latent, $z_S = \mu_S$, which serves as the target latent for diffusion training. During inference, the pretrained and frozen decoder $\mathcal{G}_S$ decodes the generated mask latent into the corresponding segmentation mask.

Latent Diffusion with Modality Dropout

Once the two volumetric perceptual compression models are trained, we perform the generative process in their latent spaces, following the latent diffusion paradigm [10]. Specifically, the diffusion model operates on the segmentation latent $z_S = \mathcal{E}_S(S)$ and learns the conditional distribution $p(z_S \mid z_X)$, where $z_X$ provides the conditioning information derived from the available modalities.

Forward Diffusion Process. To model the conditional segmentation distribution $p(z_S \mid z_X)$ in the latent space, we adapt a denoising diffusion probabilistic model (DDPM) [5] to the latent representations. The forward diffusion process starts from the segmentation latent $z_{S,0} = z_S$ and gradually adds Gaussian noise according to

where $z_{S,t}$ denotes the noisy segmentation latent at timestep $t$, and $\alpha_t$ controls the fraction of signal retained at each diffusion step. After $T$ steps, the distribution of the noisy latent approaches an isotropic Gaussian distribution. The noisy latent at any timestep can be sampled directly from the segmentation latent as

where $\gamma_t = \prod_{i=1}^{t} \alpha_i$. This formulation avoids explicitly simulating the intermediate diffusion steps during training and allows $z_{S,t}$ to be constructed directly from $z_S$.

Modality Dropout for Conditional Denoising During training, we randomly sample a non-empty modality availability vector $\mathbf{m}$ to simulate different missing-modality configurations. Since there are $2^M - 1$ possible non-empty configurations, $\mathbf{m}$ is sampled uniformly from this set:

The resulting incomplete input $X_{\mathbf{m}}$ is encoded by the pretrained multimodal encoder to obtain the conditioning latent $z_X = \mathcal{E}_X(X_{\mathbf{m}})$. The noisy segmentation latent is then conditioned on $z_X$ through channel-wise concatenation:

where $\oplus$ denotes channel-wise concatenation. Since $z_S$ and $z_X$ share the same spatial dimensions, $z_X$ can be directly concatenated with $z_{S,t}$ along the channel dimension and provided as the input to the denoising network. The denoising network $\epsilon_\theta$ is implemented as a 3D U-Net parameterized by $\theta$ and is trained to predict the Gaussian noise added to the segmentation latent using the standard MSE objective:

Here, $\mathbf{m}$ is uniformly sampled from all non-empty modality configurations, $t \sim\mathcal{U}\{1,\ldots,T\}$, and $\epsilon\sim\mathcal{N}(0,I)$ denotes the Gaussian noise used to construct $z_{S,t}$ according to Eq. (12).

Reverse Diffusion Process At inference time, we follow DDIM sampling [12] to generate the denoised segmentation latent conditioned on the multimodal latent $z_X$. Rather than using all diffusion timesteps, DDIM performs the reverse process on a smaller subset of timesteps, reducing the number of denoising updates. Specifically, the reverse process performs $l$ updates on a subsequence $\{\tau_i\}_{i=0}^{l}$, where $\tau_0 = 0$ and $\tau_l = T$. For a uniform sampling stride $\Delta t$, the selected timesteps are given by

The reverse process starts from a Gaussian noise sample $z_{S,\tau_l} \sim\mathcal{N}(0,I)$ and proceeds from $\tau_l$ to $\tau_0$. At each timestep $\tau_i$, the denoising network $\epsilon_\theta$ predicts the noise component from the conditional latent $z^{\mathrm{cond}}_{S,\tau_i}$, which is then used to compute the latent at the previous timestep:

where $i = l,l-1,\ldots,1$. Here, the clean latent $z^{\mathrm{pred}}_{S,0}$ is obtained by inverting Eq. (12) using the predicted noise. The final denoised segmentation latent is obtained as $\hat{z}_{S,0}$ at $\tau_0$. Finally, the pretrained and frozen mask decoder then produces $\hat{S} = G_S(\hat{z}_{S,0})$.

Sampling and Predictive Uncertainty

To capture the variability of the predicted segmentation for a given incomplete input, we generate multiple segmentation samples from different initial noise realizations. Specifically, the conditioning latent $z_X$ is shared across all samples, while $K$ independent Gaussian noise samples are drawn as

Each noise sample is independently propagated through the reverse diffusion process, producing a segmentation latent $z^i_S$. The pretrained and frozen mask decoder then maps each latent to a segmentation $S^i = G_S(z^i_S)$. The final segmentation is obtained from the voxel-wise mean of the generated samples, while their variability is quantified using the voxel-wise variance. The final binary segmentation is then obtained by thresholding the mean prediction at 0.5.

Experiments

Dataset and Preprocessing

We evaluate on the BraTS 2023 and 2024 adult glioma collections [9, 1], which provide 1,251 and 1,350 studies with publicly released annotations, giving 2,601 in total. Each study provides four co-registered, skull-stripped sequences (T1, T1CE, T2 and FLAIR) at an original resolution of $240 \times240 \times155$ voxels, with a voxel-wise expert annotation. Intensities are normalised per sequence and per study by $z$-scoring over brain voxels only, so that background air does not bias the statistics and scanner-dependent ranges are placed on a common scale. Each study is then cropped to a fixed cubic window centred on the centroid of the annotation bounding box. This crop defines the input size, $H = W = D = 128$, and hence a latent grid of $h = w = d = 32$. A fixed crop allows annotation latents to be computed once and cached, and this window retains the complete tumour for 97% of studies. Studies are partitioned patient-wise into training, validation and test sets in a 70:20:10 ratio with a fixed seed, giving 1,821, 520 and 260 studies respectively. All configurations of a given study fall in the same partition, so no case is ever evaluated under one subset of sequences while having been trained on under another.

Implementation Details

Autoencoders Both use a residual convolutional encoder–decoder with two residual units per level and two downsampling stages, mapping $H \times W \times D$ inputs to $h \times w \times d$ latents. The image autoencoder uses channel widths (64, 128, 256) and $\beta_X = 10^{-4}$; the annotation autoencoder uses (32, 64, 128), $C_S = 4$ and $\beta_S = 10^{-2}$. The annotation autoencoder is trained on the single binary whole-tumour mask, with $w_+ = 4.4$. Both are trained for 200 epochs with AdamW [8] at learning rate $10^{-4}$, weight decay $10^{-4}$ and batch size 32, with early stopping after 25 validation checks without improvement, then frozen. Modality dropout is applied to the image autoencoder as well as the diffusion model, so its encoder is not shown partial inputs for the first time at inference.

Diffusion model The denoising network is a 3D U-Net with base width 64, channel multipliers (1, 2, 4) and two residual blocks per level, totalling 31.3M parameters. The image latent is projected from 256 to 64 channels before concatenation with the noisy annotation latent. Training uses $T = 1000$ timesteps with a linear schedule, $2 \times10^{5}$ steps, batch size 64 and AdamW at learning rate $10^{-4}$ and weight decay $10^{-4}$, with $m$ resampled uniformly over the 15 non-empty configurations for every example.

Inference Sampling uses DDIM [12] with 200 steps, beyond which accuracy changes little while cost grows in proportion. We draw $K = 5$ samples per case, average them, and threshold at $\tau= 0.5$; Figure 2 justifies that choice. All experiments use a fixed seed and are implemented in PyTorch; training was performed on NVIDIA H200 GPUs.

Vector (PDF)

Evaluation protocol Every test case is evaluated under all 15 modality configurations, giving a complete case $\times$ configuration grid. The whole tumour is the only target: we do not segment or report the tumour core or the enhancing tumour. Because the autoencoder is frozen, its reconstruction fidelity upper-bounds what any diffusion model in that latent space can reach, so Table 1 reports it separately and Table 2 should be read against that ceiling.

Autoencoder	Target	Dice	IoU	SSIM	PSNR (dB)
Annotation	WT	0.9729	0.9475	0.8833	24.87
Image	FLAIR	n/a	n/a	0.8687	23.55
Image	T1CE	n/a	n/a	0.8527	19.62
Image	T1	n/a	n/a	0.8977	22.73
Image	T2	n/a	n/a	0.9374	24.59

$|m|$	Available sequences	Ours	HeMIS	RFNet	UniME
4	FLAIR+T1CE+T1+T2	0.955	0.821	0.901	0.933
	FLAIR+T1CE+T1	0.889	0.767	0.880	0.891
3	FLAIR+T1CE+T2	0.922	0.832	0.880	0.917
	FLAIR+T1+T2	0.914	0.863	0.902	0.911
	T1CE+T1+T2	0.766	0.702	0.777	0.765
2	FLAIR+T1CE	0.890	0.800	0.899	0.900
	FLAIR+T1	0.880	0.782	0.812	0.902
	FLAIR+T2	0.919	0.812	0.857	0.912
	T1CE+T1	0.646	0.341	0.619	0.635
	T1CE+T2	0.630	0.582	0.593	0.627
	T1+T2	0.681	0.582	0.509	0.702
1	FLAIR	0.904	0.701	0.827	0.897
	T1CE	0.281	0.223	0.324	0.329
	T1	0.355	0.192	0.404	0.423
	T2	0.578	0.251	0.426	0.562
	Mean over all 15	0.747	0.617	0.707	0.754

Metrics For anomaly segmentation, we use the Dice coefficient and Intersection over Union (IoU), which measure the overlap between the predicted and ground-truth masks. Dice is used as the primary segmentation metric. For reconstructed image quality, we use the structural similarity index (SSIM) and peak signal-to-noise ratio (PSNR), which assess the similarity and fidelity of reconstructed images to their corresponding reference images. All four metrics are higher-is-better measures.

Results and Discussion

The frozen annotation autoencoder reconstructs the whole tumour at a Dice of 0.973 (Table 1), which is the ceiling every segmentation number below is measured against; the remaining gap is attributable to the diffusion model rather than to compression. Among the four sequences the image autoencoder reconstructs T1CE least well, at an SSIM of 0.853 against 0.937 for T2, and T1CE is also the sequence that contributes least to the whole tumour, so the two limitations compound rather than offset.

Figure 2 shows whole-tumour Dice against the number of sampling seeds on the fully observed configuration. Accuracy rises sharply from 0.834 at $K = 1$ to 0.955 at $K = 5$ and shows no further trend beyond that point, fluctuating between 0.912 and 0.955 over $5 \le K \le12$. We therefore use $K = 5$, the smallest ensemble that attains the maximum. The residual fluctuation is worth stating plainly: at roughly $\pm0.02$ Dice it is of the same order as the margins separating the methods in Table 2, so individual configurations should be read as indicative and the comparison rests on the pattern across all fifteen rather than on any single cell.

Table 2 reports whole-tumour Dice for every configuration against three incomplete-modality baselines. With all four sequences the proposed model reaches 0.955, above UniME at 0.933, RFNet at 0.901 and HeMIS at 0.821, and it is the strongest of the four in eight of the fifteen configurations. Averaged over all fifteen, however, it reaches 0.747 against 0.754 for UniME, 0.707 for RFNet and 0.617 for HeMIS: the proposed model leads wherever FLAIR is present and in the fully observed case, while UniME is stronger on the sparse T1CE- and T1-dominated subsets, so the two are separated by less than the spread across configurations. The ordering across configurations is not governed by how many sequences are supplied. FLAIR alone attains 0.904, within 0.051 of the fully observed result and above every configuration that omits it; the three sequences other than FLAIR together reach only 0.766, and T1CE alone reaches 0.281. This is expected rather than anomalous: the whole tumour is delimited by peritumoural oedema, which is what FLAIR is acquired to display [9, 1], so most of the information relevant to this target is already present in that one sequence. T1CE instead highlights the enhancing rim and carries little information about the outer boundary being scored. The same asymmetry is reported throughout the missing-modality literature [4, 3, 18]. The configurations where the proposed model trails the baselines are precisely those without FLAIR, which is also where the compression ceiling and the weaker T1CE reconstruction are most exposed.

Figure 3 shows what these differences look like on a single case under three configurations. With T2 alone the four methods produce visibly similar masks on this slice. As T1CE and then FLAIR are added, the proposed model’s mask tightens towards the annotation, while HeMIS changes comparatively little; with all four sequences it recovers the upper-left extension that HeMIS and RFNet miss and fragments the interior less than UniME. The ranking in Table 2 is over the whole test set.

Figure 4 shows, for the proposed model, how the prediction changes as sequences are withheld. The five sampling seeds agree closely when the input is informative and diverge as it is impoverished: with all four sequences they differ only along a thin band at the tumour boundary, whereas with T1CE alone they disagree over most of the lesion and two of the five recover almost none of it. The uncertainty maps follow the same progression, from a narrow rim to a broad region covering the interior. Every row is the same case and the same annotation, so the only variable is how much of the input the model sees. The figure establishes the mechanism; Table 2 carries the claim across the test set.

Conclusion

We presented a three-dimensional conditional latent diffusion framework for whole-tumour segmentation from incomplete multi-sequence MRI. Two volumetric autoencoders compress images and annotations by a factor of four per axis, which makes diffusion over whole volumes tractable, and modality dropout during training lets a single model serve all fifteen acquisition configurations.

Repeated sampling yields both a consensus segmentation and a voxel-wise map of disagreement between plausible annotations. Beyond accuracy, we treated modality availability as an experimental variable rather than a nuisance. Withholding a sequence is a controlled reduction in the evidence available, and evaluating every case under every subset makes the resulting change in uncertainty measurable within each case. Our results indicate that the predicted uncertainty follows this behaviour, which suggests it reflects genuine information loss rather than sampling noise.

Limitations and Future Work

Two aspects of the protocol limit how the numbers should be read. First, the $128^3$ window is centred on the centroid of the annotation bounding box in every split, so localisation is given rather than inferred and the absolute Dice values are correspondingly optimistic: they measure segmentation inside a window already known to contain the lesion, and are not comparable with results obtained by whole-volume inference. The comparison between methods is unaffected, since all four receive identical crops. A deployable system would derive the window from a brain mask, an intensity heuristic or a separate detector, or dispense

with it through sliding-window inference; we have not re-run the study under that protocol. Second, the number of sampling seeds $K$ was read off Figure 2, which was computed on the test split rather than on held-out validation data; the remaining hyperparameters were fixed a priori or selected by early stopping on the validation split, and any $K \ge5$ changes Dice by at most 0.043.

The autoencoders apply a fixed $4 \times$ downsampling per axis, chosen so that diffusion over whole volumes fits within memory, and the diffusion model works entirely inside that latent space, so the reconstruction fidelity of Table 1 is a ceiling it cannot exceed. We segment the whole tumour only; extending to the tumour core and the enhancing tumour is the natural next step, and because those structures are smaller and thinner it would likely depend more on the compression scheme than on the diffusion model, which would require repeating the autoencoder study at several downsampling factors. The annotation autoencoder is trained with a weighted cross-entropy alone; region-overlap losses such as soft Dice, or boundary-aware and focal variants, may preserve small structures better at the same compression factor, and we have not compared them. Finally, the predictive uncertainty is reported but not compared against other estimators, so Monte Carlo dropout, deep ensembles and the Probabilistic U-Net would place it on a common footing with established methods.

Code availability. Training, evaluation and figure code is available at https://github.com/Toan-Nguyen26/medseglatdiff-3d.