# Discriminative Hamiltonian Variational Autoencoder for Accurate Tumor Segmentation in Data-Scarce Regimes

arXiv: https://arxiv.org/abs/2406.11659 · HTML: https://mlsanity.com/paper/2406.11659/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 46 of its 62 references that are in the map.
- First posted on arXiv 17 Jun 2024.

## Builds on (most cited first, 30 of 46)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Image-to-Image Translation with Conditional Adversarial Networks](https://mlsanity.com/paper/1611.07004/) (2016)
- [Going Deeper with Convolutions](https://mlsanity.com/paper/1409.4842/) (2014)
- [Perceptual Losses for Real-Time Style Transfer and Super-Resolution](https://mlsanity.com/paper/1603.08155/) (2016)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Attention U-Net: Learning Where to Look for the Pancreas](https://mlsanity.com/paper/1804.03999/) (2018)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [Least Squares Generative Adversarial Networks](https://mlsanity.com/paper/1611.04076/) (2016)
- [UNETR: Transformers for 3D Medical Image Segmentation](https://mlsanity.com/paper/2103.10504/) (2021)
- [Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images](https://mlsanity.com/paper/2201.01266/) (2022)
- [Which Training Methods for GANs do actually Converge?](https://mlsanity.com/paper/1801.04406/) (2018)
- [Improving Variational Inference with Inverse Autoregressive Flow](https://mlsanity.com/paper/1606.04934/) (2016)
- [The RSNA-ASNR-MICCAI BraTS 2021 Benchmark on Brain Tumor Segmentation and Radiogenomic Classification](https://mlsanity.com/paper/2107.02314/) (2021)
- [Tackling the Generative Learning Trilemma with Denoising Diffusion GANs](https://mlsanity.com/paper/2112.07804/) (2021)
- [nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation](https://mlsanity.com/paper/1809.10486/) (2018)
- [Deep Generative Modelling: A Comparative Review of VAEs, GANs, Normalizing Flows, Energy-Based and Autoregressive Models](https://mlsanity.com/paper/2103.04922/) (2021)
- [VAE with a VampPrior](https://mlsanity.com/paper/1705.07120/) (2017)
- [Segment Anything in Medical Images](https://mlsanity.com/paper/2304.12306/) (2023)
- [Deep Unsupervised Clustering with Gaussian Mixture Variational Autoencoders](https://mlsanity.com/paper/1611.02648/) (2016)
- [Hyperspherical Variational Auto-Encoders](https://mlsanity.com/paper/1804.00891/) (2018)
- [An overview of deep learning in medical imaging focusing on MRI](https://mlsanity.com/paper/1811.10052/) (2018)
- [Semi-Supervised Learning with Generative Adversarial Networks](https://mlsanity.com/paper/1606.01583/) (2016)
- [Deep Learning Approaches for Data Augmentation in Medical Imaging: A Review](https://mlsanity.com/paper/2307.13125/) (2023)
- [Hamiltonian Variational Auto-Encoder](https://mlsanity.com/paper/1805.11328/) (2018)

## Built on by (most cited first, 3 of 3)

- [Generative AI Enables Medical Image Segmentation in Ultra Low-Data Regimes](https://mlsanity.com/paper/2408.17421/) (2024)
- [Multi-Task Diffusion Approach For Prediction of Glioma Tumor Progression](https://mlsanity.com/paper/2509.10824/) (2025)
- [Black box behavioural modelling: Predicting human activity schedules with a deep conditional generative approach](https://mlsanity.com/paper/2512.04223/) (2025)

## Abstract (arXiv)

Deep learning has gained significant attention in medical image segmentation. However, the limited availability of annotated training data presents a challenge to achieving accurate results. In efforts to overcome this challenge, data augmentation techniques have been proposed. However, the majority of these approaches primarily focus on image generation. For segmentation tasks, providing both images and their corresponding target masks is crucial, and the generation of diverse and realistic samples remains a complex task, especially when working with limited training datasets. To this end, we propose a new end-to-end hybrid architecture based on Hamiltonian Variational Autoencoders (HVAE) and a discriminative regularization to improve the quality of generated images. Our method provides an accuracte estimation of the joint distribution of the images and masks, resulting in the generation of realistic medical images with reduced artifacts and off-distribution instances. As generating 3D volumes requires substantial time and memory, our architecture operates on a slice-by-slice basis to segment 3D volumes, capitilizing on the richly augmented dataset. Experiments conducted on two public datasets, BRATS (MRI modality) and HECKTOR (PET modality), demonstrate the efficacy of our proposed method on different medical imaging modalities with limited data.