[Paper Review] Spatial Broadcast Decoder: A Simple Architecture for Learning Disentangled Representations in VAEs

Nicholas Watters, Löıc Matthey+2 more|arXiv (Cornell University)|Jan 21, 2019

Digital Media Forensic DetectionComputer Science36 references76 citations

TL;DR

The Spatial Broadcast decoder is a simple VAE decoder that tiles latent variables across space and appends fixed coordinate channels, improving disentangling and reconstruction, especially for small objects, and generalizes well.

ABSTRACT

We present a simple neural rendering architecture that helps variational autoencoders (VAEs) learn disentangled representations. Instead of the deconvolutional network typically used in the decoder of VAEs, we tile (broadcast) the latent vector across space, concatenate fixed X- and Y-"coordinate" channels, and apply a fully convolutional network with 1x1 stride. This provides an architectural prior for dissociating positional from non-positional features in the latent distribution of VAEs, yet without providing any explicit supervision to this effect. We show that this architecture, which we term the Spatial Broadcast decoder, improves disentangling, reconstruction accuracy, and generalization to held-out regions in data space. It provides a particularly dramatic benefit when applied to datasets with small objects. We also emphasize a method for visualizing learned latent spaces that helped us diagnose our models and may prove useful for others aiming to assess data representations. Finally, we show the Spatial Broadcast Decoder is complementary to state-of-the-art (SOTA) disentangling techniques and when incorporated improves their performance.

Motivation & Objective

• Motivate learning of compositional, disentangled representations in VAEs without supervision.
• Propose a Spatial Broadcast decoder as an alternative to standard DeConv decoders.
• Show that the architecture improves disentangling, reconstruction accuracy, and generalization across datasets, especially with small objects.
• Demonstrate compatibility and complementary gains when combined with state-of-the-art disentangling methods.

Proposed method

• Tile the latent vector across image space to form z_sb and concatenate fixed x- and y-coordinate channels.
• Use an unstrided convolutional decoder on the broadcast latent with coordinate channels.
• Avoid upsampling deconvolutions to reduce optimization artifacts and simplify rendering of position.
• Provide visualization tools for latent space to diagnose representations.
• Evaluate on datasets with and without positional variation to assess disentangling and generalization.
• Compare to DeConv decoders and to FactorVAE / beta-VAE while showing complementary gains.

Experimental results

Research questions

• RQ1 Does the Spatial Broadcast decoder improve disentangling and reconstruction in VAEs compared to standard decoders?
• RQ2 How does the Spatial Broadcast approach affect generalization to held-out regions of data space and to datasets with small objects?
• RQ3 Is the Spatial Broadcast decoder complementary to existing disentangling methods like beta-VAE and FactorVAE?
• RQ4 Can latent space visualizations provide more insight into representation quality beyond scalar metrics?

Key findings

• The Spatial Broadcast VAE outperforms the DeConv VAE on MIG and latent traversals, showing more interpretable disentangled factors.
• It yields better rate-distortion trade-offs and improved MIG scores across tested models, including beta-VAE and FactorVAE variants.
• Disentangling and reconstruction improvements are especially dramatic for datasets with small objects.
• Latent space visualizations reveal near-linear, well-disentangled structures under Spatial Broadcast, contrasting with entangled DeConv representations.
• The method generalizes well to held-out regions in generative factor space and to interpolation/extrapolation in pixel space.
• The Spatial Broadcast decoder is complementary to SOTA disentangling techniques and can boost their performance when combined.