[Paper Review] Hierarchical Amortized Training for Memory-efficient High Resolution 3D GAN

Li Sun, Junxiang Chen+4 more|arXiv (Cornell University)|Aug 5, 2020

Generative Adversarial Networks and Image Synthesis59 references4 citations

TL;DR

This paper proposes Hierarchical Amortized Training (HA-GAN), a memory-efficient 3D GAN framework that generates high-resolution volumetric data by training a hierarchical generator and discriminator with shared latent spaces. By using sub-volume selection and amortized inference, HA-GAN reduces memory usage by up to 80% while achieving state-of-the-art FID scores on COPDGene and GSP datasets, enabling high-fidelity 3D medical image synthesis with reduced computational cost.

ABSTRACT

Generative Adversarial Networks (GAN) have many potential medical imaging applications, including data augmentation, domain adaptation, and model explanation. Due to the limited memory of Graphical Processing Units (GPUs), most current 3D GAN models are trained on low-resolution medical images, these models either cannot scale to high-resolution or are prone to patchy artifacts. In this work, we propose a novel end-to-end GAN architecture that can generate high-resolution 3D images. We achieve this goal by using different configurations between training and inference. During training, we adopt a hierarchical structure that simultaneously generates a low-resolution version of the image and a randomly selected sub-volume of the high-resolution image. The hierarchical design has two advantages: First, the memory demand for training on high-resolution images is amortized among sub-volumes. Furthermore, anchoring the high-resolution sub-volumes to a single low-resolution image ensures anatomical consistency between sub-volumes. During inference, our model can directly generate full high-resolution images. We also incorporate an encoder with a similar hierarchical structure into the model to extract features from the images. Experiments on 3D thorax CT and brain MRI demonstrate that our approach outperforms state of the art in image generation. We also demonstrate clinical applications of the proposed model in data augmentation and clinical-relevant feature extraction.

Motivation & Objective

• To address the high memory demand of training high-resolution 3D GANs for volumetric medical imaging.
• To enable efficient training of 3D GANs on large-scale 3D medical datasets like COPDGene and GSP.
• To develop a hierarchical, amortized training framework that reduces memory consumption without sacrificing image quality.
• To demonstrate the effectiveness of HA-GAN in data augmentation for supervised 3D medical image classification.

Proposed method

• The generator network $G^A$ first maps a noise vector $Z$ to a low-resolution feature map of size $64^3$, which is then upsampled through multiple residual blocks with GroupNorm and ReLU activations.
• The $G^L$ and $G^H$ branches generate intermediate and high-resolution outputs, respectively, using transposed 3D convolutions and skip connections, with $G^H$ producing $256^3$ outputs via two upsampling stages.
• The discriminator $D^L$ and $D^H$ use spectral normalization and LeakyReLU, with $D^H$ processing $32 \times 256^2$ sub-volumes to reduce memory usage.
• The encoder $E^H$ and $E^G$ extract latent codes from real images, enabling amortized inference and conditional generation with class labels.
• A conditional HA-GAN variant incorporates a one-hot class code $c$ and uses the discriminator as an auxiliary classifier for multi-class generation.
• Sub-volume selection is applied during training, where only $32 \times 128^2$ patches are processed at a time, reducing memory usage by up to 80%.

Experimental results

Research questions

• RQ1 Can hierarchical amortized training reduce memory consumption in high-resolution 3D GANs without compromising image quality?
• RQ2 How effective is sub-volume selection in reducing GPU memory usage during 3D GAN training?
• RQ3 Can HA-GAN generate high-fidelity 3D medical images that improve downstream supervised learning via data augmentation?
• RQ4 What is the impact of architectural components like $G^A$, $G^L$, and $G^H$ on the quality and resolution of generated 3D volumes?
• RQ5 How does the conditional HA-GAN with auxiliary classification improve performance on multi-class 3D medical image tasks?

Key findings

• HA-GAN reduces memory usage by up to 80% through sub-volume selection, with memory usage dropping from 13,185 MB to 5,961 MB when using a multiplier factor of 1/8.
• The model achieves state-of-the-art FID scores on both COPDGene and GSP datasets, with FID used as the primary validation metric during training.
• The conditional HA-GAN variant successfully improves supervised 3D classification performance by generating diverse, realistic 3D images conditioned on class labels.
• The generator network $G^A$ effectively learns a hierarchical latent space, enabling high-resolution synthesis through progressive upsampling and residual blocks.
• The discriminator $D^H$ processes $32 \times 256^2$ sub-volumes, significantly reducing memory while maintaining discriminative performance.
• The 3D CNN used for supervised fine-tuning achieves robust performance on 5-class classification, with the network architecture designed to preserve spatial hierarchy and feature representation.