Skip to content

← JAX Stochasticity step 6 of 25

Hard Primitives

Gumbel-Softmax

Jump to the task

Why this matters

A categorical sample chooses one class. That hard choice does not provide a useful ordinary gradient with respect to the logits. Gumbel-Softmax instead returns a smooth vector that can be differentiated for a fixed draw of noise and a positive temperature.

First add independent standard Gumbel noise to each logit. Taking the largest perturbed score gives a categorical sample with probabilities softmax(logits). Applying softmax to the perturbed scores divided by temperature gives the relaxed sample used here.

This is a continuous relaxation of a discrete choice. Gradients of a loss through this vector generally differ from gradients of the original discrete objective. The exercise does not use a hard forward sample or a straight-through gradient estimator.

Worked mini-example

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
logits = jnp.array([1.0, 2.0, 3.0])
temperature = 0.5
g = jax.random.gumbel(key, logits.shape)
y = jax.nn.softmax((logits + g) / temperature)
# Shape (3,), nonnegative entries, sum approximately 1.

For a fixed noise draw with a unique largest perturbed score, decreasing temperature toward zero concentrates the result on that class. Increasing temperature without bound approaches a uniform vector. A particular value such as 0.1 or 10 does not guarantee either behavior for every input.

Problem

Implement gumbel_softmax(seed, logits, temperature) using:

  1. key = jax.random.PRNGKey(int(seed)).
  2. g = jax.random.gumbel(key, logits.shape) with the default float32 noise.
  3. jax.nn.softmax((logits + g) / temperature).

seed is a nonnegative integer-valued scalar supplied as a number. logits is a nonempty finite float32 vector of shape (K,). temperature is finite and positive; assume the perturbed, scaled scores remain finite. Return a float32 vector of shape (K,) whose entries sum to approximately one.

The fixtures check this seeded draw, not just any vector summing to one. They do not by themselves verify gradients. Reusing the same key, shape and random configuration reproduces the noise. New draws in a training loop require fresh keys, usually obtained with jax.random.split.

Common pitfalls

  • Divide the entire perturbed score, including the noise, by temperature. Scaling only the logits changes the sampling distribution.
  • The returned vector is neither a one-hot draw nor the original softmax(logits) probabilities. It includes random noise.
  • Very small temperatures can produce saturated outputs and unstable or unhelpful gradients. Zero temperature is not valid.
  • The Python int(seed) conversion needs a concrete seed. If wrapping this exercise in jax.jit, mark the seed static. A reusable training function would usually accept a JAX key directly.

The chart below shows how temperature changes softmax for fixed scores. It does not generate Gumbel noise.

Loading visualization…