PaperlayerView exact revision on arXiv

Source-linked AI summary

Categorical Reparameterization with Gumbel-Softmax

Eric Jang, Shixiang Gu, Ben Poole

arXiv:1611.01144v5stat.MLcs.LG

TL;DR

Training stochastic neural networks with categorical latent variables is difficult because categorical samples are non-differentiable. The paper introduces Gumbel-Softmax, a reparameterizable continuous approximation that outperforms existing stochastic gradient estimators and enables dramatic speedups for discrete latent-variable inference.

  • Problem

    Existing gradient estimators were not formulated specifically for categorical variables, making discrete stochastic networks difficult to train with backpropagation.

  • Method

    Gumbel-Softmax provides a differentiable continuous approximation to categorical samples whose parameter gradients are computed through reparameterization.

  • Results

    Gumbel-Softmax outperforms existing stochastic gradient estimators across structured prediction and variational autoencoder tasks while enabling dramatic speedups for discrete latent-variable inference.

  • Takeaways & Limitations

    The estimator offers a simple differentiable sampling mechanism that integrates categorical variables into neural networks trained with standard backpropagation.

  • Takeaways & Limitations

    At non-zero temperature, Gumbel-Softmax samples differ from categorical samples, creating a tradeoff between sample discreteness and gradient variance.

Abstract

from arXiv · show

Categorical variables are a natural choice for representing discrete structure in the world. However, stochastic neural networks rarely use categorical latent variables due to the inability to backpropagate through samples. In this work, we present an efficient gradient estimator that replaces the non-differentiable sample from a categorical distribution with a differentiable sample from a novel Gumbel-Softmax distribution. This distribution has the essential property that it can be smoothly annealed into a categorical distribution. We show that our Gumbel-Softmax estimator outperforms state-of-the-art gradient estimators on structured output prediction and unsupervised generative modeling tasks with categorical latent variables, and enables large speedups on semi-supervised classification.

1 INTRODUCTION

The paper introduces Gumbel-Softmax, a differentiable continuous distribution that approximates categorical samples and supports reparameterized gradients. It enables standard backpropagation for categorical variables and efficient semi-supervised training without costly marginalization.

  • Motivation: Discrete random variables represent distributions in unsupervised learning, language modeling, attention mechanisms, and reinforcement learning.They can encode semantic classes, image regions, and memory locations.
  • Challenge: Backpropagation cannot train stochastic networks containing non-differentiable discrete-variable layers.Prior gradient estimation focused on score-function estimators with variance reduction or biased path-derivative methods.
  • Method: Gumbel-Softmax is a continuous simplex distribution that approximates categorical samples and provides parameter gradients through the reparameterization trick.This supplies a differentiable approximate sampling mechanism for categorical variables.
  • Empirical results: Gumbel-Softmax outperforms all single-sample gradient estimators on Bernoulli and categorical variables.The claim is experimental and covers both variable types.
  • Applications: The estimator efficiently trains semi-supervised models without costly marginalization over unobserved categorical latent variables.The paper gives Kingma et al. (2014) as an example of such a model.

2 THE GUMBEL-SOFTMAX DISTRIBUTION

The Gumbel-Softmax is a continuous distribution over the simplex that approximates categorical one-hot samples while remaining differentiable for positive temperature. Annealing temperature toward zero recovers categorical behavior, with practical tradeoffs between sample discreteness and gradient variance.

  • Definition: The Gumbel-Softmax distribution places continuous samples on the simplex to approximate categorical variables encoded as one-hot vectors.Categorical samples lie on the corners of the (k −1)-dimensional simplex, while the relaxation occupies the simplex continuously.
  • Construction: Gumbel noise combined with a softmax provides a differentiable approximation to the categorical sampling arg max.The construction generates k-dimensional sample vectors y ∈∆k−1 from independently sampled Gumbel variables.
  • Temperature: As τ approaches 0, Gumbel-Softmax samples become one-hot and the distribution becomes identical to the categorical distribution.For positive temperature, the distribution remains smooth and has a well-defined gradient with respect to π, enabling backpropagation.
  • Temperature: Small temperatures produce nearly one-hot samples but high gradient variance, whereas large temperatures produce smoother samples with lower gradient variance.In practice, training starts at a high temperature and anneals to a small but non-zero temperature.
  • Straight-Through Estimator: When discrete outputs are required, the Straight-Through Gumbel Estimator uses arg max in the forward pass and the continuous approximation for the backward pass.It approximates ∇θz ≈∇θy for constrained discrete sampling scenarios.

3 RELATED WORK

This section reviews gradient estimators for discrete stochastic variables, contrasting unbiased score-function methods and control variates with biased pathwise proxies. It also describes the computational cost of marginalizing categorical variables in semi-supervised learning and motivates differentiable Gumbel-Softmax samples.

  • Pathwise and proxy estimators: Pathwise estimators reparameterize samples as z = g(θ, ϵ), enabling gradients through deterministic transformations; Straight-Through instead uses a differentiable proxy and can create forward-backward discrepancies.For k = 2, ST Gumbel-Softmax resembles slope-annealed Straight-Through but uses a softmax.
  • Gumbel-Softmax: Gumbel-Softmax avoids Straight-Through’s sample-independent proxy issue because each y is a differentiable proxy for its corresponding discrete z.The estimator is constructed using the reparameterization trick applied to continuous samples.
  • Score-function estimators: Score-function estimators are unbiased and require only continuity of pθ(z), but their variance scales linearly with sample-vector dimensionality, hindering categorical models.They do not require backpropagation through f or z, yet their high variance slows convergence.
  • Control variates: Control variates reduce score-function variance by subtracting b(z) from f and adding its analytical expectation µb, preserving unbiasedness.The estimator uses (f(z) − b(z))∇θ log pθ(z) + µb.
  • Semi-supervised learning: Marginalizing all k classes in semi-supervised classification costs O(D + k(I + G)) per forward/backward step, becoming prohibitively expensive as the number of classes grows.Gumbel-Softmax instead permits backpropagation through a single sample y ∼ qφ(y|x).

4 EXPERIMENTAL RESULTS

Across structured prediction, generative modeling, and semi-supervised classification, Gumbel-Softmax generally outperformed competing estimators while enabling faster single-sample inference. The experiments used MNIST benchmarks and compared Gumbel-Softmax and straight-through variants with established stochastic gradient estimators.

  • Experimental setup: The experiments compared Gumbel-Softmax and ST Gumbel-Softmax with Score-Function, DARN, MuProp, Straight-Through, and Slope-Annealed ST on structured prediction and generative modeling.MNIST used fixed binarization, with learning rates selected on validation data and test performance reported.
  • Structured output prediction: Gumbel-Softmax outperformed other estimators for both Bernoulli and categorical variables in structured output prediction, while ST Gumbel-Softmax improved specifically on categorical variables.The softmax temperature was fixed at τ = 1 for this task.
  • Variational autoencoders: Gumbel-Softmax drastically outperformed other estimators for Bernoulli and categorical latent variables in variational autoencoder training, while ST Gumbel-Softmax outperformed on categorical variables.Temperature annealing was used, and evaluation employed a multi-sample bound with m = 1000.
  • Semi-supervised classification: Gumbel-Softmax enabled direct backpropagation through single samples from qφ(y, z|x), avoiding marginalization over y without compromising generative or classification performance.Marginalization-based inference has computational cost scaling linearly with the number of classes.
  • Semi-supervised classification: 2× as fast for 10 classes and 9.9× as fast for 100 classes was training with Gumbel-Softmax compared with marginalization.The speed comparison varied the number of categorical classes using MNIST images with randomly generated labels.

5 DISCUSSION

The discussion identifies the reparameterizable Gumbel-Softmax distribution as the work’s primary contribution, providing low-variance path-derivative gradients for categorical variables and effective performance across structured prediction and variational autoencoder tasks.

  • Contribution: The primary contribution is a reparameterizable Gumbel-Softmax distribution for categorical variables.Its corresponding estimator affords low-variance path derivative gradients.
  • Empirical effectiveness: Gumbel-Softmax and Straight-Through Gumbel-Softmax are effective on structured output prediction and variational autoencoder tasks.
  • Empirical effectiveness: They outperform existing stochastic gradient estimators for both Bernoulli and categorical latent variables.

A SEMI-SUPERVISED CLASSIFICATION MODEL

The semi-supervised classification experiments use the architecture described in Figures 6 and 7, with a generative model and inference model involving Gaussian style and categorical class variables.

  • A SEMI-SUPERVISED CLASSIFICATION MODEL: The experimental architecture for semi-supervised classification is described in Figures 6 and 7.These figures present the model used in the experiments for semi-supervised classification.
  • A SEMI-SUPERVISED CLASSIFICATION MODEL: The generative model pθ(x|y, z) synthesizes images from Gaussian style variable z and categorical class variable y, while qφ(y, z|x) infers them from x.Gaussian z is differentiable with respect to its parameters because it is reparameterizable; when y is unobserved, prior VAE training requires marginalizing over all y values.

B DERIVING THE DENSITY OF THE GUMBEL-SOFTMAX DISTRIBUTION

This section derives the probability density function of the Gumbel-Softmax distribution with probabilities π1, ..., πk and temperature τ. It begins by defining logits from the probabilities and introducing Gumbel noise used to compute a sample.

  • The derivation targets the Gumbel-Softmax probability density function for probabilities π1, ..., πk and temperature τ.
  • The logits are defined as xi = log πi for the categorical probabilities.

B.1 CENTERED GUMBEL DENSITY

This section addresses the non-invertibility introduced by softmax normalization by defining an equivalent centered sampling process and deriving its density through marginalization and a change of variables.

  • Centered sampling: Softmax normalization removes one degree of freedom, so the mapping from Gumbel samples g to Gumbel-Softmax samples y is not invertible.The equivalent process subtracts (x_k + g_k)/τ before applying softmax.
  • Centered sampling: The centered variables are defined as u_i = x_i + g_i − (x_k + g_k) for i = 1, ..., k − 1, with g_i ∼ Gumbel(0, 1).The centered multivariate Gumbel density is derived for this equivalent sampling process.
  • Density derivation: The density is computed by marginalizing out the last Gumbel sample, g_k, using the Gumbel probability density.For scale β = 1 and mean μ, the density at z is given as f(z, μ) = e^(μ−z)−e^(μ−z).
  • Density derivation: A change of variables v = e^(−g_k) gives dg_k = dv/v, while setting u_k = 0 simplifies the notation.The substitution uses dv = −e^(−g_k)dg_k.

B.2 TRANSFORMING TO A GUMBEL-SOFTMAX

The section derives the Gumbel-Softmax density by transforming k−1 centered Gumbel samples, with the final coordinate determined by the others. It applies the change-of-variables formula using the transformation’s inverse and Jacobian determinant.

  • Transformation: A deterministic transformation h maps k−1 centered Gumbel samples to the first k−1 Gumbel-Softmax coordinates.The construction begins with samples u1, ..., uk−1 from the centered Gumbel distribution.
  • Transformation: The final coordinate probability yk is fixed once the first k−1 coordinates are given.This reduces the relevant degrees of freedom to the first k−1 variables.
  • Density derivation: The Gumbel-Softmax sample probability is computed with the change-of-variables formula applied only to the first k−1 variables.The derivation uses the lower-dimensional transformation because the final coordinate is determined by the others.
  • Density derivation: The derivation requires the inverse of h and its Jacobian determinant.These two components are computed before substitution into the density formula.
  • Density derivation: The resulting expression substitutes the centered Gumbel density, h’s inverse, and the Jacobian determinant into the change-of-variables formula.The determinant calculation uses identities for products, diagonal matrices, and det(I + uvT) = 1 + uT v.
Loading 1611.01144v5…