Source-linked AI summary
GET: Generative Embedding Translation for Medical Image Segmentation
Md Maklachur Rahman, Md Hasan Al Banna, Saraf Anjum, Mahmudul Hasan, Tracy Hammond
TL;DR
Medical image segmentation needs latent-space translation that preserves target structure while remaining efficient. GET progressively translates image embeddings into mask embeddings within a frozen SD-VAE using a compact ETN. Across five datasets, it outperforms generative, CNN, and Transformer baselines, including GMS, and shows larger gains under BUS-BUSI domain shift.
Problem
Medical image segmentation requires image-to-mask translation that preserves target location, extent, shape, and boundaries while avoiding limitations of pixel-wise and computationally costly approaches.
Method
GET progressively translates image embeddings into mask embeddings across multiple resolutions within a frozen SD-VAE using a multi-stage ETN with local, global, and multi-scale modeling.
Results
Across five benchmarks, GET consistently outperforms generative, CNN, and Transformer baselines; versus GMS, it improves average Dice and IoU by 0.93% and 1.26% and reduces HD95 by 0.81 pixels.
Takeaways & Limitations
GET provides an efficient structured embedding-translation framework for medical segmentation, with stronger results under bidirectional BUS-BUSI domain shift.
Takeaways & Limitations
GET assumes that a frozen SD-VAE pretrained on large-scale natural images provides a high-fidelity latent manifold transferable to medical imagery.
Abstract
from arXiv · showhide
Generative segmentation provides an alternative to direct pixel-wise prediction by operating on learned latent representations, but effective image-to-mask translation must preserve target structure while remaining computationally efficient. We propose Generative Embedding Translation (GET), a structured embedding-translation framework that progressively transforms image embeddings into mask embeddings within the frozen latent space of a Stable Diffusion VAE. GET uses a U-Net-style Embedding Translation Network with 1.07M trainable parameters, combining Mobile Bottleneck Convolutions, Subsampled Self-Attention, and Multi-scale Feature Enrichment for local modeling, global context, and multi-scale refinement. Across five medical segmentation datasets, GET outperforms generative, CNN, and Transformer baselines. Compared with the strongest generative baseline, GMS, GET improves average Dice and IoU by 0.93% and 1.26%, reduces HD95 by 0.81 pixels, and uses 31.41% fewer trainable parameters. Under bidirectional BUS-BUSI domain shift, GET further improves Dice and IoU by 3.51% and 3.39%, while reducing HD95 by 27.37 pixels. Our code is available at: https://github.com/maklachur/GET.
1 Introduction
GET addresses the need for medical segmentation methods that capture long-range structure without the computational burden of full self-attention or iterative generative inference. It progressively translates image embeddings into mask embeddings in a frozen SD-VAE space using a compact ETN with local, global, and multi-scale modeling.
- Medical image segmentation supports diagnosis and treatment, but CNNs struggle with long-range dependencies needed for complex anatomy and ambiguous boundaries.
- Transformer models capture long-range interactions but can incur substantial computational cost and depend on large annotated datasets.
- Generative segmentation translates learned latent representations, but accurate image-to-mask translation must preserve location, extent, shape, and boundaries.
- GET progressively maps image embeddings to mask embeddings across multiple resolutions within a frozen SD-VAE space using a compact trainable translator.
- GET’s five-stage ETN combines MBC for local features, SSA for bottleneck global context, and MFE for multi-scale decoder refinement, with latent- and mask-space supervision.
- Across five benchmarks, GET improves average Dice and IoU over GMS by 0.93% and 1.26%, reduces HD95 by 0.81 pixels, and uses 31.41% fewer trainable parameters.
2 Methodology
GET performs structured image-to-mask embedding translation in a frozen SD-VAE space using a lightweight U-Net-style ETN. The ETN combines efficient local extraction, bottleneck global context, and multi-scale decoder refinement with latent and mask-space supervision.
- Latent Embedding Formulation: GET maps image latents to segmentation latents through a trainable ETN while a frozen SD-VAE encodes inputs and decodes predicted masks.During inference, the ETN predicts a segmentation latent, which the frozen decoder reconstructs and projects into a single-channel segmentation map.
- Embedding Translation Network (ETN): The five-stage U-Net-style ETN progressively translates representations across resolutions, using MBC blocks in the encoder and decoder with SSA at the bottleneck.The decoder upsamples and fuses encoder features before producing the predicted segmentation latent.
- Multi-Scale Feature Enrichment: MFE enriches decoder features through parallel dilated depthwise-convolution branches with dilation rates d ∈ {1,6,12}.The branches capture fine details and broader semantic information before feature fusion and residual refinement.
- Subsampled Self-Attention: SSA captures global context efficiently by using full-resolution queries with spatially subsampled keys and values, reducing attention complexity by about r^2.The method uses average pooling with r = 2 to form the subsampled keys and values.
- Training Objective and Loss Function: Mask-space supervision combines Dice and Focal-Tversky losses, with Focal-Tversky focusing learning on hard, imbalanced pixels.The stated Focal-Tversky parameters are (α, β, γ) = (0.7, 0.3, 0.75).
3 Experiments and Results
GET is evaluated across five medical segmentation datasets, cross-domain settings, efficiency comparisons, and ablations. It consistently improves segmentation quality while using a lightweight architecture and structured latent-space supervision.
- Datasets and evaluation: GET is evaluated on BUS, BUSI, GlaS, HAM10000, and Kvasir–Instrument using established train–test splits and comparisons with CNN, Transformer, and generative methods.Results are reported using Dice, IoU, and HD95, with comparisons averaged over five runs where specified.
- Comparison with SOTA methods: 90.76% and 83.03% DSC on BUS and BUSI exceed GMS by 2.19 and 1.70 percentage points, respectively.GET also reduces HD95 from 6.59 to 4.84 pixels on BUS and from 19.66 to 16.99 pixels on BUSI.
- Cross-domain generalization: 64.63% DSC for BUS→BUSI and 84.39% DSC for BUSI→BUS demonstrate gains over GMS under bidirectional domain shift.For BUSI→BUS, GET improves IoU by 6.34 percentage points and reduces HD95 by 6.75 pixels.
- Ablation study: The full MBC + MFE + SSA combination achieves the highest DSCs of 90.76% on BUS, 83.03% on BUSI, and 89.24% on GlaS.It also yields the lowest reported HD95 values of 4.84, 16.99, and 15.62 pixels, respectively.
- Ablation study: Three-branch dilation with d = 1, 6, 12 combines fine, intermediate, and broad receptive fields, achieving 90.76% DSC and 4.84 HD95 on BUS.The same configuration achieves 89.24% DSC and 15.62 HD95 on GlaS.
- Ablation study: The full latent- and mask-space objective performs best, while removing mask-space supervision reduces DSC to 89.60%.Adding latent reconstruction reduces BUS HD95 from 8.62 to 6.40 pixels, whereas latent alignment alone provides a smaller gain.
- Efficiency: GET requires 341.14 GFLOPs, achieves 28.07 samples/s, and uses 1.07M trainable parameters.It reduces GFLOPs by 64.9% and achieves approximately 45× higher throughput than MedSegDiff-V2.
4 Conclusion
GET is presented as a structured embedding-translation framework for medical image segmentation in a frozen SD-VAE space. Across five benchmarks and under BUS–BUSI domain shift, it improves accuracy, robustness, and efficiency relative to the reported baselines.
- Conclusion: GET improves average Dice and IoU over GMS by 0.93% and 1.26%, reduces HD95 by 0.81 pixels, and uses 31.41% fewer trainable parameters.Under BUS–BUSI domain shift, it further improves Dice and IoU by 3.51% and 3.39% and reduces HD95 by 27.37 pixels.
Supplementary Material
The supplementary material adds methodological and evaluation details supporting GET’s design choices and comparative analysis.
- Supplementary Material: Supplementary material covers related work, evaluation metrics, extensive ablations, and an efficiency–accuracy ranking plot.The ranking plot compares GET with existing state-of-the-art methods.
5 Related Work
Prior medical segmentation work spans CNNs, Transformers, and generative latent-space methods. These approaches balance local feature extraction, long-range modeling, generative priors, and computational cost in different ways.
- CNN-based methods: CNN architectures such as UNet and its variants remain central because they capture hierarchical local features through encoder–decoder designs and skip connections.The related work also describes lightweight CNN variants and improvements in feature aggregation.
- Transformer-based methods: Transformer-based methods improve long-range reasoning but often incur high computational overhead and require large annotated datasets.Hybrid and fully transformer-based designs use global or windowed attention to expand receptive fields.
- Generative methods: VAEs and VQ-VAEs provide structured latent spaces that can support downstream segmentation through smoother and more semantically organized representations.Generative latent representations are presented as a complement to direct pixel-wise prediction.
- Generative methods: Diffusion methods can produce high-quality masks through iterative denoising but are computationally expensive and slow at inference.GSS and GMS instead operate directly on pretrained latent representations to avoid iterative denoising.
6 Experiments and Results
The section defines the evaluation metrics and training objective, then compares segmentation accuracy against model complexity across five medical imaging datasets. GET is reported to outperform CNN-, Transformer-, and diffusion-based generative approaches while reducing parameter overhead.
- Evaluation Metrics: DSC, IoU, and HD95 jointly measure overlap accuracy, spatial accuracy, and boundary precision, respectively.Higher DSC and IoU are better, whereas lower HD95 indicates better boundary precision.
- Evaluation Metrics: HD95 measures boundary alignment as the 95th percentile of Euclidean distances between predicted and ground-truth boundary points.Using the 95th percentile makes the metric robust to outliers.
- Training Loss: The training objective combines latent reconstruction, L2 latent alignment, and mask-space segmentation losses.The segmentation term uses Dice and Focal-Tversky losses to address class imbalance and sharpen boundary predictions.
- Comparative Results: Figure 3 plots average DSC, IoU, and HD95 against trainable parameters on a log scale across BUS, BUSI, GlaS, HAM10000, and Kvasir–Instrument.The plots evaluate the trade-off between segmentation quality and model complexity.
- Comparative Results: GET achieves higher DSC and IoU and lower HD95 than CNN-, Transformer-, and diffusion-based generative methods while using fewer parameters.Models closer to the top-left for DSC and IoU, or bottom-left for HD95, represent better accuracy-complexity trade-offs.
7 Additional Ablation Study
The ablation study tests SSA subsampling, deep-supervision weighting, auxiliary-head weighting, and latent-alignment weighting on BUS and GlaS. Across these settings, moderate or monotonically decaying choices provide the strongest reported accuracy and boundary precision.
- SSA Subsampling: r = 2 in SSA achieves the best efficiency-accuracy balance, with peak DSC of 90.76% on BUS and 89.24% on GlaS.r = 4 loses spatial detail, while r = 1 adds computational complexity without improving results.
- Deep Supervision: Monotonically decaying deep-supervision weights wk = {1.0, 0.4, 0.3, 0.2} outperform uniform and linear weighting schemes.They achieve the highest DSC values and lowest HD95 on BUS and GlaS.
- Auxiliary Supervision: λ = 0.5 for the auxiliary segmentation head yields the highest DSC and lowest boundary errors on BUS and GlaS.The reported HD95 values are 4.84 on BUS and 15.62 on GlaS.
- Latent Alignment: λla = 0.01 provides the best overall latent-alignment performance, reaching 90.76% DSC on BUS and 89.24% on GlaS.The corresponding HD95 values are 4.84 and 15.62 pixels, respectively.