[Paper Review] BAGAN: Data Augmentation with Balancing GAN

Giovanni Mariani, Florian Scheidegger+3 more|arXiv (Cornell University)|Mar 26, 2018

Advanced Neural Network ApplicationsComputer Science22 references258 citations

TL;DR

BAGAN trains a GAN to generate minority-class images for imbalanced datasets by jointly training on all classes with autoencoder initialization and latent-space class conditioning, improving minority-class image quality and final classifier accuracy.

ABSTRACT

Image classification datasets are often imbalanced, characteristic that negatively affects the accuracy of deep-learning classifiers. In this work we propose balancing GAN (BAGAN) as an augmentation tool to restore balance in imbalanced datasets. This is challenging because the few minority-class images may not be enough to train a GAN. We overcome this issue by including during the adversarial training all available images of majority and minority classes. The generative model learns useful features from majority classes and uses these to generate images for minority classes. We apply class conditioning in the latent space to drive the generation process towards a target class. The generator in the GAN is initialized with the encoder module of an autoencoder that enables us to learn an accurate class-conditioning in the latent space. We compare the proposed methodology with state-of-the-art GANs and demonstrate that BAGAN generates images of superior quality when trained with an imbalanced dataset.

Motivation & Objective

• Motivate and address the problem of imbalanced image datasets hurting classifier performance.
• Propose BAGAN to generate minority-class images by leveraging all available data from majority and minority classes.
• Initialize GANs with an autoencoder to start from a stable solution and learn latent-space class conditioning.
• Drive the generator to produce target-class images while avoiding mode collapse and ensuring diversity.
• Evaluate BAGAN against state-of-the-art GANs across multiple datasets to show improvements in image quality and downstream classification accuracy.

Proposed method

• Train an autoencoder on all training images (majority and minority) to learn a unconditional representation.
• Initialize the GAN by transferring weights from the autoencoder (decoder to generator, encoder to discriminator) to start adversarial training near a good solution.
• Model each class conditioning in the latent space with a class-conditional latent vector generator drawn from a per-class multivariate normal distribution N_c = N(mu_c, Sigma_c) learned from E(X_c).
• Train a single-output discriminator that predicts either a problem-specific class label c or the fake label, and train the generator to produce images matching the target class c.
• During adversarial training, supply a balanced batch with 1/(n+1) fake images; use Z_c vectors to generate fake images for each class c.
• Compare BAGAN to ACGAN and a simple GAN, analyzing quality, diversity (SSIM), and downstream classifier accuracy.

Experimental results

Research questions

• RQ1 Can a GAN trained on both majority and minority classes with autoencoder initialization generate high-quality minority-class images from imbalanced data?
• RQ2 Does embedding class conditioning in the latent space with an autoencoder-informed initialization reduce mode collapse and improve diversity for minority-class generation?
• RQ3 How does BAGAN compare to ACGAN and simple GAN in terms of image quality, diversity, and downstream classifier performance on imbalanced datasets?

Key findings

• BAGAN generally produces higher-quality minority-class images than ACGAN and simple GAN across datasets.
• Generated minority-class images achieve higher accuracy when classified by a ResNet-18 trained on the augmented data, especially under strong imbalance.
• SSIM analyses show BAGAN maintains greater diversity among generated images compared with ACGAN and simple GAN, approaching diversity of real images.
• On GTSRB BAGAN achieves the best downstream classifier accuracy among the evaluated methods; on MNIST CIFAR-10 and Flowers BAGAN often matches or beats alternative augmentations when orientation-sensitive features are involved.
• Simple GAN often collapses to a single sample per class under imbalanced conditions, while BAGAN avoids this through its unified class conditioning and initialization.
• Overall BAGAN outperforms state-of-the-art GANs for generating minority-class images from imbalanced datasets, leading to higher classifier accuracy.