IB-GAN: Disentangled Representation Learning with Information Bottleneck Generative Adversarial Networks

AI summary

Overview

Research area: Unsupervised disentangled representation learning with generative adversarial networks (computer vision; deep generative modeling).

Technical level: Advanced. The paper derives a new GAN objective through the variational inference formulation of Information Bottleneck theory, and readers need familiarity with mutual information, variational bounds, VAEs, and GAN training.

Scope: The paper proposes IB-GAN, a GAN-based unsupervised model that adds a mutual-information minimization term from Information Bottleneck theory to the InfoGAN objective, and evaluates it for disentanglement quality and sample quality on dSprites, Color-dSprites, CelebA, and 3D Chairs.

What This Paper Is About

Unsupervised models that learn "disentangled" representations — where each latent direction corresponds to a single independent factor of variation such as shape, scale, or rotation — are dominated by β-VAE and InfoGAN. β-VAE tends to produce blurry reconstructions, while InfoGAN struggles to disentangle on complicated datasets such as CelebA or 3D Chairs. The authors observe that InfoGAN's objective is missing a mutual information minimization term that the Information Bottleneck objective contains, and build a new GAN model around that missing term.

Key Contributions

1. A new GAN-based model for unsupervised disentangled representation learning. IB-GAN is presented as an extension of InfoGAN that supplies the information-constraining mechanism InfoGAN lacks from the perspective of Information Bottleneck theory.
2. A trainable latent distribution via an intermediate latent encoder. The architecture derived from the variational inference formulation of the IB framework places a stochastic layer inside the generator, between the input and the generated data, which is trained jointly with the generator end-to-end.
3. A new variational upper bound on generative mutual information. Rather than approximating the generator marginal directly (which the authors argue is difficult and risks degrading sample quality), they introduce an intermediate stochastic representation and use the Markov property (Z → R → X) to obtain a bound that constrains mutual information without directly regularizing the output image.
4. Empirical validation across four datasets. IB-GAN achieves disentanglement scores competitive with state-of-the-art β-VAEs and outperforming InfoGAN on dSprites and Color-dSprites, with FID scores on CelebA and 3D Chairs that are lower than those of VAE baselines and comparable to GAN baselines.

Main Findings

• Disentanglement scores (Kim and Mnih 2018 metric). IB-GAN reaches 0.80 ± 0.07 on dSprites and 0.79 ± 0.05 on Color-dSprites. The optimal average scores of 0.80 and 0.79 are obtained at β = 0.141 and β = 0.071, respectively. Scores are averaged over 10 random seeds.

• Comparison against baselines. On dSprites / Color-dSprites respectively: GAN 0.40 ± 0.05 / 0.35 ± 0.04; InfoGAN 0.61 ± 0.03 / 0.55 ± 0.08; VAE 0.61 ± 0.04 / 0.59 ± 0.06; β-VAE 0.69 ± 0.09 / 0.74 ± 0.06; FactorVAE 0.81 ± 0.07 / 0.82 ± 0.06; β-TCVAE 0.79 ± 0.06 / 0.80 ± 0.07. IB-GAN exceeds GAN, VAE, and InfoGAN, and is comparable to β-VAE, FactorVAE, and β-TCVAE.

• Latent traversal on dSprites. IB-GAN learns 5 out of 5 ground-truth factors (Y and X positions, scales, rotations, and shapes), aligning with the KL scores in the convergence figure.

• Latent traversal on Color-dSprites. IB-GAN captures 6 out of 6 ground-truth factors, including the color factor.

• FID scores (lower is better). CelebA / 3D Chairs: VAE 129.7 / 56.2; β-VAE 131.0 / 91.3; FactorVAE 109.7 / 44.7; β-TCVAE 125.0 / 57.3; GAN 8.4 / 27.9; InfoGAN 9.3 / 25.6; IB-GAN 7.4 / 25.5. IB-GAN's FID is substantially lower than the VAE-family baselines and comparable to the GAN baselines.

• Convergence behavior. On dSprites at β = 0.212, the KL term for each of the 10 dimensions of r increases at different points over training, indicating the representation encoder slowly adapts to capture independent factors. A similar behavior is reported in β-VAE.

• Effect of β. Experiments use a hyperparameter search over β ∈ [0, 1], taking a median over 150K training iterations per trial and averaging over 10 trials. At β = 0 the upper mutual information bound disappears and the representation diverges from the prior; a small β > 0 drastically reduces this divergence; larger β gradually lowers the upper bound and consequently the lower bound. Optimal disentanglement scores are achieved for β in the range [0.071, 0.212].

• Qualitative traversal on natural images. IB-GAN discovers human-recognizable attributes on CelebA including azimuth, gender, and skin tone, and on 3D Chairs disentangles azimuth, scales, and leg types. The paper states these attributes are hardly captured by the original InfoGAN.

• Special cases of the objective. Setting β = 0 reduces the IB-GAN objective to the InfoGAN objective. Removing the mutual information lower-bound as well reduces it to a vanilla GAN with input drawn from the prior m(r).

Methodology in Plain English

The starting point is InfoGAN, which trains a generator to produce images from a latent code while maximizing the mutual information between that code and the generated image, so the code stays informative about the output. The authors notice that the Information Bottleneck objective also contains a term that minimizes mutual information — a compression term — and that InfoGAN has no equivalent. They add it.

Directly adding an upper bound on mutual information turns out to be awkward: the natural formulation requires a model of the generator's marginal image distribution, which acts like an image prior, and a bad choice there would visibly degrade generated samples. The authors' workaround is architectural. Instead of feeding the latent code straight to the generator, they insert a stochastic "representation encoder" that maps the code z to an intermediate representation r, and the generator produces images from r. Because the chain Z → R → X is Markov, mutual information between z and the image is bounded above by the mutual information between z and r, and that quantity is easy to bound using a simple prior on r (a standard Gaussian, as in VAEs).

The final objective combines three pieces: the usual GAN adversarial loss, a reconstruction term that trains a separate reconstructor to recover z from generated images (this is the InfoGAN-style lower bound that keeps the code informative), and a KL divergence term that pushes the intermediate representation toward the Gaussian prior, with β controlling the trade-off. Gradients flow to the stochastic encoder through the reparameterization trick.

Implementation-wise, the generator and discriminator follow DCGAN with batch normalization, the reconstructor shares front-end features with the discriminator, and an MLP-based representation encoder sits before the generator. Optimization uses RMSProp with momentum 0.9, batch size 64, and images normalized to [-1, 1]. The same architecture is used across experiments except for channel sizes.

To get a representation back from a real image (since the model does not directly map x to r), a two-step procedure is needed: sample z from the learned reconstructor q(z|x), then pass it through the representation encoder e(r|z).

Why This Matters

The paper connects two lines of work that had largely been separate: the theoretical analysis showing that β-VAE's KL term corresponds to mutual information minimization under Information Bottleneck, and GAN-based disentanglement. It shows that the compression mechanism responsible for β-VAE's disentanglement can be transplanted into a GAN, potentially getting disentanglement without β-VAE's blurry reconstructions. It also contributes a new way of upper-bounding generative mutual information using the causal structure of the network rather than a direct approximation of the data marginal.

Potential real-world applications (inferred from the paper's framing of disentangled representations as valuable for downstream tasks; the paper does not itself demonstrate these applications):

• Controllable image generation and editing, where a user manipulates a single semantic attribute without changing others.
• Face attribute manipulation, given the azimuth, gender, and skin tone factors the model captures on CelebA.
• 3D shape and object design tools, given the azimuth, scale, and leg-type factors captured on 3D Chairs.
• Downstream transfer learning, where an unsupervised representation with independent factors may be reused in tasks with limited labels.

Industry relevance: Disentangled representations are directly relevant to creative and media tooling, synthetic data generation, and any pipeline where controllable, interpretable latent variables reduce the cost of labeled supervision. The paper's practical claim is narrower than its theoretical framing: it delivers β-VAE-level disentanglement with GAN-level sample quality, which matters for applications where both fidelity and control are required.

Future Directions

• Direct encoder from real images. The paper explicitly acknowledges that IB-GAN does not map the representation encoder directly from real data x to r — recovering representations requires the two-step sampling procedure through q(z|x) and e(r|z). Learning q(r|x) directly is a natural next step.
• Dataset scale and complexity. Evaluation of disentanglement relies on synthetic datasets with known ground-truth factors (dSprites and Color-dSprites); extending quantitative disentanglement evaluation to more complex real-world data remains open.
• Choosing β principledly. The optimal β differs between datasets (0.141 on dSprites, 0.071 on Color-dSprites) and required a search over [0, 1]; a principled way to select β would be valuable.
• Wider comparison with disentanglement baselines. The paper notes that many state-of-the-art models are evaluated in the Locatello et al. setting, and points toward tighter integration of the upper-bound formulation with other information-theoretic objectives.

Target Audience

Researchers and graduate students working on deep generative models, representation learning, and information-theoretic machine learning. It is most useful to readers already comfortable with GANs, VAEs, and variational inference who want to understand how Information Bottleneck theory can be applied to adversarial training. Practitioners interested in controllable image generation may benefit from the empirical results and architecture, but the derivation of the objective requires a theoretical background to follow.

Authors’ abstract

We propose a new GAN-based unsupervised model for disentangled representation learning. The new model is discovered in an attempt to utilize the Information Bottleneck (IB) framework to the optimization of GAN, thereby named IB-GAN. The architecture of IB-GAN is partially similar to that of InfoGAN but has a critical difference; an intermediate layer of the generator is leveraged to constrain the mutual information between the input and the generated output. The intermediate stochastic layer can serve as a learnable latent distribution that is trained with the generator jointly in an end-to-end fashion. As a result, the generator of IB-GAN can harness the latent space in a disentangled and interpretable manner. With the experiments on dSprites and Color-dSprites dataset, we demonstrate that IB-GAN achieves competitive disentanglement scores to those of state-of-the-art \b{eta}-VAEs and outperforms InfoGAN. Moreover, the visual quality and the diversity of samples generated by IB-GAN are often better than those by \b{eta}-VAEs and Info-GAN in terms of FID score on CelebA and 3D Chairs dataset.