PaperlayerView exact revision on arXiv

Source-linked AI summary

Spatial Broadcast Decoder: A Simple Architecture for Learning Disentangled Representations in VAEs

Nicholas Watters, Loic Matthey, Christopher P. Burgess, Alexander Lerchner

arXiv:1901.07017v2cs.LGcs.CVstat.ML

TL;DR

VAEs often struggle to learn disentangled representations because standard architectures and disentangling objectives can be difficult to optimize and sensitive to hyperparameters. The paper introduces the Spatial Broadcast decoder, which tiles latent vectors across space, adds fixed coordinate channels, and uses an unstrided convolutional decoder. It reports improved disentangling, reconstruction, generalization, and performance alongside state-of-the-art disentangling techniques, while noting data-dependent limitations.

  • Problem

    VAE representations can be highly entangled, while existing objective variations introduce hyperparameters that are extremely sensitive; disentanglement metrics also have serious shortcomings.

  • Method

    The Spatial Broadcast decoder replaces upsampling deconvolutions by tiling the latent vector across space, appending fixed coordinate channels, and applying an unstrided fully convolutional network.

  • Results

    The Spatial Broadcast decoder improves disentangling, reconstruction accuracy, latent-space generalization, and state-of-the-art disentangling models across the evaluated VAE settings.

  • Takeaways & Limitations

    The architecture is a simple, broadly compatible step toward robust compositional visual representation learning, supported by latent-space visualization and evaluation across datasets and models.

  • Takeaways & Limitations

    The decoder may hurt when data does not benefit from absolute coordinates, and the paper evaluates its application only to VAE models.

Abstract

from arXiv · show

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.

1 Introduction

The paper introduces the Spatial Broadcast decoder as a simple VAE architecture for improving disentangled, compositional visual representations. It reports gains in disentangling, reconstruction, generalization, and compatibility with existing disentangling methods.

  • Motivation: Compositional representations are motivated as a route to recombining learned components and improving generalization.The paper relates feature compositionality to disentangled representations, while noting that disentanglement lacks a consensus definition.
  • Contribution: The Spatial Broadcast decoder replaces the standard MLP/deconvolutional decoder with an alternative architecture intended to improve representation quality and hyperparameter robustness.The authors present it as a modification applicable across VAE-based models.
  • Results: The decoder improves disentangling and reconstruction accuracy on datasets of simple objects.This is one of the paper’s stated benefits of using the Spatial Broadcast decoder for image representation learning.
  • Results: The decoder complements and improves state-of-the-art disentangling techniques.The paper explicitly reports this as a benefit beyond vanilla VAE models.
  • Results: Benefits are particularly significant for datasets containing small objects.The authors identify this as a regime that is difficult for standard VAE architectures.
  • Results: The architecture improves representational generalization to held-out regions involving interpolation and extrapolation in latent space.The paper also introduces latent-space visualization as a tool for qualitatively assessing learned representations.

2 Spatial Broadcast Decoder

The Spatial Broadcast decoder addresses limitations of standard VAE decoders by making spatial information explicit through broadcasting and coordinate channels. This simplifies rendering and improves performance, while remaining subject to data-dependent and model-scope limitations.

  • Limitations of standard decoders: Standard VAE decoders use an MLP followed by an upsampling deconvolutional network, whose learned representations can become highly entangled.VAE objective variants address this problem but introduce hyperparameters that can be highly sensitive.
  • Limitations of standard decoders: Upsampling deconvolutional networks can create checkerboard artifacts and spatial discontinuities that complicate optimization.These effects are described as potentially problematic for representation learning in the latent space.
  • Architecture: The Spatial Broadcast decoder tiles the latent vector across image space, appends fixed coordinate channels, and applies an unstrided fully convolutional network.This gives the decoder explicit spatial information instead of requiring it to propagate spatial asymmetries through upsampling layers.
  • Architecture: Making object placement a simple function eases optimization, and the Spatial Broadcast decoder improves performance across varied VAE-based models.The decoder still retains capacity for more complex datasets.
  • Architecture: The decoder adds no supervision: the model must still encode spatial information in its latent space for accurate reconstruction.Its role is to let the generative model use encoded spatial information efficiently.
  • Limitations: Reconstructions can improve with shallower, lower-parameter networks, but the decoder may hurt when absolute coordinates are not useful.A standard DeConv decoder may then better place patterns relative to one another or capture extended spatial correlations.
  • Limitations: The paper explores the Spatial Broadcast decoder only in VAE models, although it can be applied to other generative image-rendering models.This limits the scope of the reported evaluation.

3 Related Work

The Spatial Broadcast decoder builds on prior uses of coordinate information, positional embeddings, and spatial rendering mechanisms. Its distinction is to tile a learned latent embedding and render an entire image rather than local patches.

  • Coordinate information: Coordinate channels have been used to improve positional generalization and were named CoordConv in related work.The paper notes that earlier work had already used the technique and that its origin is unclear.
  • Coordinate information: Language models commonly combine fixed or learned position embeddings with sequence inputs to compute context-dependent translations.This provides a related use of explicit positional information in another modeling domain.
  • Spatial rendering: Other generative models learn where to write image patches, whereas Spatial Broadcast tiles a learned latent embedding and renders an entire image.The paper also contrasts this with Laplacian Pyramid VAE arrangements that distribute effects across global and fine scales.

4 Results

Across controlled image datasets, the Spatial Broadcast decoder improves disentangling and often reconstruction, with especially strong benefits for small objects and held-out factor regions. Latent-space visualizations complement scalar metrics by revealing geometry that MIG can miss.

  • 4.1 Performance on colored sprites: Spatial Broadcast VAE traversals separate all 8 colored-sprite factors, while DeConv traversals remain entangled.Hyperparameters were selected to minimize model error rather than explicitly optimize disentangling.
  • 4.1 Performance on colored sprites: Spatial Broadcast improves FactorVAE disentangling according to both MIG and traversal visualizations.The decoder is therefore complementary to an existing state-of-the-art disentangling objective.
  • 4.1 Performance on colored sprites: Spatial Broadcast β-VAE achieves better rate-distortion and higher MIG than DeConv β-VAE across β values on colored sprites.The comparison sweeps β from 0.4 to 5.4 using 10 replicas per value.
  • 4.2 Datasets without positional variation: Spatial Broadcast VAE representations look well-disentangled on Chairs and 3D Object-in-Room datasets without positional factors.These results use the standard VAE objective, including for factors affecting broad, non-local image regions.
  • 4.2 Datasets without positional variation: For small hue-varying circles, DeConv learns entangled, discontinuous representations, whereas Spatial Broadcast disentangles them well.The positional architectural bias is particularly useful as object size decreases.
  • 4.4 Latent space geometry visualization: Spatial Broadcast produces near-linear latent geometry and nearly perfect held-out-region generalization, while DeConv remains entangled.The latent geometry analysis visualizes a 2-dimensional factor subspace; MIG can miss this contrast because of rotation sensitivity.

5 Conclusion

The Spatial Broadcast decoder improves latent representations across datasets and VAE-based models, especially for positional variation, while supporting generalization and compositional visual representation learning.

  • The Spatial Broadcast decoder improves learned latent representations, with the largest benefits on datasets containing objects that vary in position.
  • It improves disentangling and reconstruction accuracy across a wide variety of datasets and models.
  • The decoder improves representational generalization in latent space and boosts SOTA models on disentangling and reconstruction accuracy.
  • The work positions the decoder as one step toward robust compositional visual representation learning for improved flexibility and generalization.

Supplementary material

The supplementary experiments use standardized VAE training choices and document architecture differences across vanilla, β-VAE, CoordConv, FactorVAE, and Spatial Broadcast models.

  • All VAE models use Bernoulli decoder logits for reconstruction-error computation, while fixed-variance Gaussian decoders produce qualitatively similar results and improve stability.
  • The networks use ReLU activations, truncated-normal weight initialization, zero biases, Adam optimization, and no BatchNorm or dropout.
  • Except for β-VAE in Figure 4, VAE models use the standard VAE loss with KL coefficient β = 1; FactorVAE also uses β = 1.
  • The Spatial Broadcast decoder uses fewer parameters than DeConv but requires about 50% more memory for weights, and needs three extra deconvolutional layers for 3D Object-in-Room.
  • Training uses learning rate 3 · 10^-4 and batch size 16, with same padding for convolutional and deconvolutional layers.
  • Table 1 specifies architectures for Vanilla VAE, β-VAE, CoordConv VAE, and the ablation study, with additional Spatial Broadcast layers for 3D Object-in-Room.
  • Table 2 specifies FactorVAE architectures, using the published encoder and DeConv decoder and the Spatial Broadcast decoder architecture from Table 1.

A.3 Datasets

The experiments use 64 × 64 normalized images from standard and procedurally generated datasets, varying object appearance, position, size, and other factors with controlled training and held-out distributions.

  • All datasets contain 64 × 64 images normalized to pixel values in [0, 1].
  • Colored Sprites extends dSprites with HSV colors sampled continuously, making the effective dataset size infinite.
  • Chairs contains 86,366 single-channel images, while 3D Object-in-Room contains six factors including camera angle, object size, shape, and colors.
  • The Circles datasets were created with Spriteworld’s factor distributions and renderer, using maximum episode length one and dummy-agent observations.
  • The generated Circles datasets vary controlled subsets of X-position, Y-position, Size, and Color, producing X-Y, X-H, R-G, X-Y-H Small, and X-Y-H Tiny datasets.
  • Some datasets define color ranges in RGB space and others in HSV space, as summarized by Table 3.
  • Dependent-factor datasets hold out one quarter of the data from the intersection of half-ranges of two factors, with the held-out region centered or positioned at an edge.
  • Each dataset uses 500,000 randomly sampled training images, with model training steps selected so models generally converged.

B Ablation Study

The ablation study tests whether the ordered coordinate channels matter and reports that spatially permuting them reduces both disentangling and reconstruction accuracy, including for tiny objects.

  • The Spatial Broadcast decoder concatenates constant coordinate channels to its tiled latent input, and the ablation tests whether their organization matters.
  • The ablation randomly permutes spatial coordinate entries while preserving each coordinate pair and keeping the permutation fixed during training.
  • Both disentangling and reconstruction accuracy are significantly reduced when coordinate channels are spatially permuted.
  • For small-circle datasets, circles measure 0.1 or 0.075 times the frame width, with similar disentangling results at the smaller size.
  • The lower object-size limit at which the Spatial Broadcast decoder can disentangle representations was not explored.

D CoordConv VAE

CoordConv preserves positional continuity but remains more entangled than Spatial Broadcast, while coordinate-channel shuffling harms Spatial Broadcast performance.

  • CoordConv VAE appends coordinate channels to every deconvolutional feature layer without spatially tiling the latent vector.
  • CoordConv produces lower MIG scores and more entangled traversals than Spatial Broadcast VAE on colored sprites.
  • Randomly shuffling coordinate channels damages Spatial Broadcast reconstructions and traversals, indicating that their linear organization matters.

E Extra traversals for datasets without positional variation

Additional traversal seeds provide a broader check of disentanglement for the Chairs and 3D objects datasets. The authors use them because one traversal plot only assesses local behavior at one latent-space location.

  • A single latent traversal plot shows disentanglement only locally around one point in latent space.
  • The authors therefore show traversals around a second seed for the same models used on the Chairs and 3D objects datasets.
  • These additional traversals support the claim that Spatial Broadcast VAE disentangles both datasets.

F Architecture Hyperparameters

Architecture performance depends on decoder depth, upsampling, and pre-broadcast MLPs. Deeper networks and upscaling generally reduce latent capacity or representation quality, although an MLP can help on 3D Object-in-Room.

  • Hyperparameter selection: Hyperparameters were selected using ELBO rather than disentanglement information, with large-scale decoder sweeps on colored sprites.
  • Metrics: ELBO combines negative log likelihood and KL divergence; NLL measures reconstruction accuracy, while KL measures latent information relative to the Gaussian prior.
  • ConvNet depth: Deeper Spatial Broadcast ConvNets move toward lower rate and higher distortion, reducing latent information, reconstruction accuracy, and represented factors.
  • Pre-broadcast MLP: A pre-broadcast MLP hurts colored-sprites performance but can improve 3D Object-in-Room reconstructions and latent usage.
  • Decoder continuum: The Spatial Broadcast and DeConv decoders form a continuum obtained by replacing stride-1 convolutions with stride-2 deconvolutions and reducing the tiling resolution.
  • Upscaling and DeConv depth: Upscaling deconvolutions incrementally reduce Spatial Broadcast performance, while deeper DeConv networks similarly represent fewer factors.

G Latent Space Geometry Analysis for Circle Datasets

Across circle datasets, Spatial Broadcast generally improves disentangling and latent-space geometry, especially when positional variation is strong or regions are held out. Benefits are smaller without positional variation but can remain evident in generalization settings and blank-image mixtures.

  • Overall pattern: Across X-Y, X-H, and R-G factor pairs, Spatial Broadcast nearly always improves disentangling most strongly with positional variation and least without it.
  • FactorVAE: FactorVAE results also improve with Spatial Broadcast, although FactorVAE can contract held-out holes while sacrificing latent-space geometry.
  • X-Y datasets: On dependent X-Y datasets, Spatial Broadcast yields nearly linear representations through held-out regions and helps with extrapolation in both pixel and factor space.
  • X-H datasets: On dependent X-H datasets, Spatial Broadcast improves generalization across held-out holes, including cases requiring extrapolation in pixel and generative-factor space.
  • R-G datasets: Without positional variation, Spatial Broadcast benefits are less clear for independent R-G data but clearer for dependent R-G generalization experiments.
  • Blank-image data: For X-H data containing half blank images, Spatial Broadcast improves latent-space geometry according to both MIG and traversal visualizations.
Loading 1901.07017v2…