Source-linked AI summary
Latent-to-Latent Flow for Volumetric Stochastic Segmentation
Omar Todd, Sooha Kim, Raghav Mehta, Katherine Mackay, David Bernstein, Alexandra Taylor, Fabio De Sousa Ribeiro, Ben Glocker
TL;DR
Inter-observer variation makes medical segmentation uncertain, while limited multiple annotations and volumetric computational complexity constrain stochastic segmentation research. The paper introduces L2L-Flow, which transports between latent representations of volume-conditioned priors and labels. It reports up to 14× faster inference than Flow-SSNs while retaining highly competitive performance, and clinically meaningful samples with lower memory and computation than the full-resolution counterpart.
Problem
Stochastic segmentation must represent inter-observer variation, but multiple expert annotations are scarce and volumetric probabilistic modelling adds computational complexity.
Source: §1 · pp. 1–2Method
L2L-Flow learns flow transport from a latent volume-conditional prior to a latent label distribution using encoded image and label representations.
Source: §3.3 · p. 4Results
Takeaways & Limitations
Results across multiple datasets show that extending complex probabilistic models from 2D to 3D is non-trivial, with prior methods showing methodological and training-stability issues.
Source: §5 · p. 8GED alone does not capture the distinction relevant to downstream planning, so quantitative evaluation should be complemented with qualitative assessment.
Source: §4.3 · p. 7
Abstract
from arXiv · showhide
Uncertainty arising from inter-observer variability in medical image segmentation plays an important role in developing treatment plans. Research in this area is inhibited by the lack of multiple annotations for large-scale medical datasets, especially for volumetric data, which suffers from additional scaling and computational complexity challenges. Flow matching has emerged as a powerful framework for generative modelling and has also been demonstrated to maintain strong performance when working with latent representations of images. In this work, we introduce a latent-to-latent flow technique for stochastic segmentation of medical volumes via encoded representations of both the image and label space. We evaluate our method on two challenging applications covering delineation uncertainty for radiotherapy planning and multiple organ structure segmentation, improving efficiency up to 14x compared with full resolution models while maintaining clinically relevant performance.
1 Introduction
Medical segmentation must account for ambiguous boundaries and inter-observer variation, but stochastic segmentation research and benchmarking remain concentrated on 2D data despite the prevalence of 3D volumes. This work introduces a time-shifted Flow-SSN and L2L-Flow, evaluating them on radiotherapy planning and multi-organ segmentation.
1 Introduction
- Inter-observer variation is an important source of aleatoric uncertainty that deterministic segmentation models cannot accurately capture, motivating stochastic models of plausible segmentation maps.Incorporating uncertainty into planning is important for improving patient outcomes.Source: §1 · p. 1
- The field’s 2D focus, scarcity of multiple expert annotations, and volumetric computational complexity leave flow matching for volumetric stochastic segmentation unexplored.The computational challenge matters for time-sensitive applications such as adaptive radiotherapy planning.Source: §1 · p. 2
- The authors identify high-resolution volumetric performance degradation in Flow-SSNs and introduce a time-shifted noise schedule to improve stability and performance.They state that the schedule sets a new benchmark for volumetric stochastic segmentation.Source: §1 · p. 2
- L2L-Flow is presented as an efficient, scalable alternative to time-shifted Flow-SSN, with a controllable efficiency/performance trade-off and up to 14× faster inference than Flow-SSNs.The authors report that it retains highly competitive performance.Source: §1 · p. 2
- The flow models are evaluated on radiotherapy planning and multi-organ segmentation, with clinically relevant uncertainty estimation capabilities reported.The authors highlight generative approaches for volumetric segmentation under label ambiguity.Source: §1 · p. 2
2 Related Work
Prior stochastic segmentation methods span probabilistic latent models, low-rank distributions, and annotator-specific experts, with limitations in stability, distribution complexity, or volumetric scaling. Diffusion and flow-based approaches offer alternatives, while latent flows for stochastic segmentation must retain structural information from high-resolution images and labels.
2 Related Work
- Probabilistic U-Net, PhiSeg, SSNs, and MoSE model segmentation variability with conditional densities, hierarchical latent distributions, low-rank Gaussians, or annotator-specific distributions.SSNs can have training stability issues and limited distribution complexity; MoSE can face scaling issues when training with many samples per mode on 3D data.Source: §2 · p. 2
- Diffusion models have been successfully applied to stochastic segmentation in 2D and can generate high-quality segmentations, but DDPM-based methods are very expensive at inference time, increasingly undesirable as data scale grows.Source: §2 · p. 2
- Flow-SSNs showed strong 2D performance and greater efficiency than diffusion-based stochastic segmentation, while latent rectified flows have also been applied to cardiac segmentation for shape-completion augmentation.For stochastic segmentation, latent flows need to use high-resolution image and label structure to preserve clinical utility in a lower-dimensional space.Source: §2 · p. 3
3 Methodology
The methodology moves flow matching to encoded label representations and adjusts the Flow-SSN noise schedule for volumetric tasks. At inference, sampled volume-conditioned latents are transported by an ODE and decoded into full-resolution segmentations.
3.1 Neural ODEs and Flow Matching
- Flow matching trains continuous normalising flows along conditional probability paths, while Flow-SSNs map an image-conditional prior to discrete label maps using one-hot interpolation.Source: §3.1 · p. 3
3.2 Time-Shifted Noise Scheduling
- The proposed modified shifted noise schedule accounts for the path direction from t = 0 to t = 1; this modification unlocks significant training stability and performance improvements for Flow-SSNs in volumetric settings.Source: §3.2 · pp. 3–4
3.3 Latent-to-Latent Flow
- L2L-Flow trains a volumetric label autoencoder, then learns transport from a volume-conditioned latent prior to encoded label latents using a least-squares objective and a pretrained, frozen label encoder.Source: §3.3 · pp. 4–5
- At inference, L2L-Flow samples a volume-conditioned latent, integrates the flow ODE, and decodes the result into a full-resolution segmentation.Source: Figure 1 · p. 4, §3.3 · p. 5
- L2L-Flow accelerates flow sampling by up to 14× while maintaining volumetric segmentation performance comparable to full-resolution models.Source: §3.3 · p. 5
4 Experiments
Across radiotherapy and CURVAS, L2L-Flow offers competitive volumetric segmentation performance with substantially faster inference, while the experiments also show that clinical variation and fine-structure performance require careful evaluation. The evaluation used cross-validation on datasets with limited case counts and compared stochastic models with expert annotations and challenge baselines.
4.1 Datasets
- The radiotherapy dataset was private and the CURVAS challenge dataset public; CURVAS included three annotations for pancreas, kidney, and liver.The study evaluated radiotherapy delineation uncertainty using clinician-defined ranges and CURVAS multi-class segmentation.Source: §4.1 · p. 5
4.2 Implementation Details
- Both datasets had limited case counts—55 radiotherapy cases and 90 CURVAS cases—so evaluation used 5-fold cross-validation and GED and diversity with 16 samples.The text cautions that diversity is context-dependent because random noise can trivially achieve high diversity.Source: §4.2 · p. 5
4.3 Radiotherapy Delineation Uncertainty
- More than 14× faster inference than Flow-SSN, with highly competitive segmentation performance against full-resolution models on radiotherapy.The table reports quantitative cross-validation results; Flow-SSN with the shifted schedule had the strongest overall results among full-resolution stochastic models.Source: Table 1 · p. 6, §4.3 · p. 6
4.2 Implementation Details
- For shifted Flow-SSN, α=3 gave the best experimental results; higher α values were more prone to overfitting, and shifting the schedule improved GED in high-dimensional volumes.L2L-Flow used 4× downsampling for radiotherapy and both 4× and 2× settings for CURVAS.Source: §4.2 · p. 6, Figure 2 · p. 6
4.3 Radiotherapy Delineation Uncertainty
- L2L-Flow and Flow-SSN showed clinically relevant variation against expert contours, unlike SSN’s strongly concentric variation.For a case with ground-truth IOV from 10 clinical experts, overlaps and crossings in expert contours were reflected in flow-model samples.Source: §4.3 · p. 6, Figure 3 · p. 7
- GED alone did not capture the clinical dissimilarity of SSN contours to human IOV, motivating more holistic evaluation for clinical workflows.The paper states these SSN contours were unsuitable for downstream planning.Source: §4.3 · pp. 6–7
4.4 Multi-Organ Segmentation (CURVAS Dataset)
- 3.4% better GED for the pancreas than original Flow-SSN was reported for time-shifted Flow-SSN, while L2L-Flow at 2× downsampling remained ≈5× faster than full-resolution Flow-SSN.The 2× setting improved fine-detail capture; the 4× setting had a larger performance gap on the small pancreas structure.Source: §4.4 · p. 7, Figure 4 · p. 8
5 Conclusion
L2L-Flow is a scalable generative approach for stochastic segmentation of volumetric data, using latent image and label information to improve memory and computational efficiency. Results across datasets show that 2D probabilistic methods do not transfer straightforwardly to 3D, and clinical assessment benefits from considering qualitative and clinical evidence alongside metrics.
5 Conclusion
- L2L-Flow generates clinically meaningful volumetric segmentation samples with lower memory and computational demands than its full-resolution counterpart.The approach uses high-level image and label information in a latent rectified-flow formulation.Source: §5 · p. 8
- Results across datasets indicate that extending complex probabilistic models from 2D to 3D presents methodological and training-stability challenges.The conclusion cites prior methods’ stability issues as evidence of the difficulty.Source: §5 · p. 8
- Radiotherapy experiments show that stochastic-segmentation metrics used alone can miss clinically relevant distinctions, supporting qualitative and clinical evaluation in clinical workflows.Future work will explore more expressive image and label encoders, including foundation-model-derived encoders.Source: §5 · p. 8