The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables (1611.00712v3)

Abstract: The reparameterization trick enables optimizing large scale stochastic computation graphs via gradient descent. The essence of the trick is to refactor each stochastic node into a differentiable function of its parameters and a random variable with fixed distribution. After refactoring, the gradients of the loss propagated by the chain rule through the graph are low variance unbiased estimators of the gradients of the expected loss. While many continuous random variables have such reparameterizations, discrete random variables lack useful reparameterizations due to the discontinuous nature of discrete states. In this work we introduce Concrete random variables---continuous relaxations of discrete random variables. The Concrete distribution is a new family of distributions with closed form densities and a simple reparameterization. Whenever a discrete stochastic node of a computation graph can be refactored into a one-hot bit representation that is treated continuously, Concrete stochastic nodes can be used with automatic differentiation to produce low-variance biased gradients of objectives (including objectives that depend on the log-probability of latent stochastic nodes) on the corresponding discrete graph. We demonstrate the effectiveness of Concrete relaxations on density estimation and structured prediction tasks using neural networks.

• The paper introduces the Concrete distribution as a differentiable relaxation that enables backpropagation in models with discrete variables.
• It leverages a softmax transformation with Gumbel noise to approximate discrete outcomes while reducing gradient variance compared to traditional methods.
• Experimental results on density estimation and structured prediction tasks show improved performance over state-of-the-art estimators like VIMCO.

The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables

Abstract and Core Contribution:

The paper introduces the Concrete distribution, a new family of continuous distributions designed to approximate discrete random variables in the context of stochastic computation graphs (SCGs). This is achieved through a relaxation technique that allows for efficient gradient-based optimization. The central contribution lies in the utility of Concrete random variables to facilitate backpropagation in models with discrete stochastic nodes.

Background and Motivation:

Traditional optimization techniques for SCGs using automatic differentiation libraries, such as TensorFlow and Theano, face significant challenges when encountering discrete nodes, primarily due to the non-differentiable nature of these nodes. The reparameterization trick, common for continuous distributions, becomes infeasible for discrete distributions. While unbiased estimators for gradients do exist, they often suffer from high variance and require complex implementations, such as REINFORCE computational variants. Therein lies the motivation for the introduction of the Concrete distribution: a differentiable, continuous approximation that allows for efficient gradient-based optimization absent the aforementioned high-variance issues.

The Concrete Distribution and Its Properties:

The Concrete distribution, denoted as $\Concrete(\alpha, \lambda)$, is parameterized by a vector $\alpha \in (0, \infty)^n$ and a positive temperature $\lambda$. The distribution operates over the $(n-1)$-simplex. The sampling procedure involves a transformation based on the softmax function applied to perturbed logits with Gumbel noise, ensuring that the samples approximate discrete outcomes closely as the temperature approaches zero. This continuous relaxation allows gradients to flow through the SCGs, thus enabling backpropagation to effectively optimize parameters even in the presence of discrete nodes.

Formal properties of the Concrete distribution include closed-form density functions and the fact that its zero-temperature limit neatly corresponds to the exact discrete distribution it approximates. Additionally, from a practical perspective, ensuring convexity in the density of the distribution at low temperatures helps maintain useful gradient properties for optimization tasks.

Application and Implementation:

The practical utility of the Concrete distribution is demonstrated through applications in variational autoencoders (VAEs) and structured prediction tasks. In the context of VAEs, the objective involves replacing discrete latent variables with Concrete random variables during training. The KL divergence between the relaxed posterior and prior distributions is managed efficiently using the properties of the Concrete nodes.

Key considerations in implementation involve the careful treatment of log-probabilities to avoid numerical underflow and over-reliance on low-variance properties of gradients. The introduction of the ExpConcrete representation, focusing on log-space values, serves as a practical enhancement to avoid numerical issues during optimization.

Experimental Evaluation:

The empirical evaluation encompasses two primary tasks: density estimation and structured output prediction on image datasets (MNIST and Omniglot). The experiments compare the performance of Concrete relaxations with state-of-the-art score function estimators such as VIMCO and NVIL. Noteworthy findings include that Concrete relaxations outperform VIMCO in non-linear models for density estimation tasks and structured prediction scenarios, suggesting their robustness in handling complex architectures. However, the results also indicate the importance of careful temperature tuning, with the Concrete distribution often favoring higher temperatures to maintain desirable gradient properties.

Implications and Future Directions:

The use of Concrete distributions represents a significant advancement in the optimization of SCGs involving discrete variables, particularly in the context of deep learning models requiring efficient and scalable solutions. The improved gradient properties and the straightforward applicability in automatic differentiation frameworks underscore the practical benefits of this approach.

Future developments may explore temperature annealing strategies and hybrid models that combine Concrete relaxations with traditional gradient estimation techniques to further enhance performance. Additionally, extending this approach to other types of combinatorial structures and exploring its implications in reinforcement learning environments provide fertile ground for further research.

The introduction of the Concrete distribution bridges the gap between continuous optimization techniques and discrete stochastic models, providing a valuable tool for advancing the capabilities and efficacy of neural network-based SCGs.