Stochastic softmax tricks and Gumbel-Rao

The argmax has zero gradient almost everywhere, and a continuous relaxation restores the pathwise derivative. As a function of $\theta$, $X=\arg\max_{x\in\mathcal X}(\theta+G)^\top x$ is piecewise constant: it is differentiable almost everywhere with derivative 0, so $\mathbb E[dL(X)/d\theta]=0$ while $\nabla_\theta\mathbb E[L(X)]$ generally is not 0. All the dependence on $\theta$ sits in the jumps, which break the exchange of derivative and expectation. A relaxation that is continuous in the noise, such as the softmax, makes the exchange valid for the relaxed objective. The temperature sets the balance between bias and variance: as $\tau\to0$ the relaxed sample approaches $X$, so the bias falls, while the softmax Jacobian $(\mathrm{diag}(s)-ss^\top)/\tau$ at $s=X_\tau$ is tiny except near ties between utilities, where it grows like $1/\tau$, so the variance rises, as the Gumbel-Rao paper reports.

A linear program with random utilities defines the model, and its solver draws the sample. A linear function on $P$ is maximized at a vertex, and a unique maximizer over $\mathcal X$ is the unique maximizer over $P$, so drawing $X$ means solving a linear program; the sample is unique almost surely when $U$ gives no mass to any hyperplane $\{U^\top a=0\}$. Fast solvers exist for many structures: the top $k$ entries of $U$ by sorting for $k$-subsets, Kruskal's algorithm for a maximum spanning tree, Chu-Liu-Edmonds for an arborescence, the Hungarian method for matchings. Only a solver is needed, so sampling is cheap for sets far too large to enumerate. In exchange $p_\theta$ is defined only implicitly: for structured $\mathcal X$ the law of $X$ under Gumbel utilities is in general not $p_\theta(x)\propto e^{x^\top\theta}$ and is hard to characterize, so the paper treats $U$ as a modeling choice and warns against defaulting to $\theta+G$ or $\mathcal N(\theta,\sigma^2I)$ without analysis.

Strong convexity makes the regularized argmax the gradient of a smooth convex function. Write $\delta_P$ for the convex indicator of $P$ (0 on $P$, $+\infty$ outside), $g_\tau=\tau(f-\min f)+\delta_P$, and $g_\tau^*$ for its convex conjugate$g^*(u)=\sup_x\,u^\top x-g(x)$; when the supremum is attained at a unique point, $\nabla g^*(u)$ is that point. By conjugate duality$$\begin{aligned}X_\tau&=\nabla g_\tau^*(U),\\g_\tau^*(u)&=\sup_x\,u^\top x-g_\tau(x).\end{aligned}$$The conjugate of a strongly convex function is differentiable with a Lipschitz gradient, so $X_\tau$ is unique and continuous in $U$, and Rademacher's theorema Lipschitz function on $\mathbb R^n$ is differentiable almost everywhere makes it differentiable almost everywhere: Proposition 2. As $\tau\to0$, $g_\tau^*$ increases to the support function$u\mapsto\max_{x\in P}u^\top x$, the value of the linear program over $P$ of $P$. Where that linear program has a unique solution, the support function is differentiable on a neighborhood of $u$, and pointwise convergence of convex functions then carries over to their gradients, so $X_\tau\to X$: Proposition 1. Because $\partial X_\tau/\partial U$ is the Hessian of $g_\tau^*$, it is symmetric, so the vector-Jacobian product that backpropagation needs is a directional derivative of $X_\tau$, computable for a small step $\epsilon$ with two extra solves and no extra loss evaluation:$$\begin{aligned}\frac{dL(X_\tau)}{dU}&\approx\frac1{2\epsilon}\big(X_\tau(U+\epsilon v)\\&\qquad-X_\tau(U-\epsilon v)\big),\\v&=\frac{\partial L}{\partial X_\tau}.\end{aligned}$$

Projections, entropies and marginal inference give computable relaxations. With $f=\|x\|^2/2$, $X_\tau$ is the Euclidean projection of $U/\tau$ onto $P$. On the simplex the categorical entropy $f=\sum_ix_i\log x_i$ gives the softmax, and for independent bits the binary entropy gives the sigmoid. The most general choice comes from an exponential familythe distributions $p(x)\propto e^{\eta^\top x}$ on $\mathcal X$, indexed by the natural parameter $\eta\in\mathbb R^n$: for $p(x)\propto e^{u^\top x/\tau}$ on $\mathcal X$, the mean of $x$ solves a problem of exactly the SST form,$$\begin{aligned}\mu_\tau(u)&=\sum_{x\in\mathcal X}x\,\frac{e^{u^\top x/\tau}}{\sum_{y\in\mathcal X}e^{u^\top y/\tau}}\\&=\arg\max_{x\in P}\,u^\top x-\tau A^*(x),\end{aligned}$$with $A^*$the convex conjugate of the log-partition function $A(\eta)=\log\sum_{x\in\mathcal X}e^{\eta^\top x}$; at a mean $\mu$ it is minus the entropy of the family member with mean $\mu$ the conjugate of the log-partition function, a negative entropy; then $X_\tau=\mu_\tau(U)$. Computing it is marginal inferencecomputing the marginal probabilities, here the mean vector $\mu_\tau$, of a distribution over structures: for spanning trees the edge marginals follow from Kirchhoff's matrix-tree theoremthe sum of edge-weight products over all spanning trees of a graph equals a determinant of its reduced weighted Laplacian; edge marginals follow from that determinant's derivatives (a directed version for arborescences), and for a chain with a cardinality constraint from dynamic programming. As $\tau\to0$ an SST tends to a random draw $X$ from the SMT, because $U$ is random; annealing marginal inference toward a single most probable structure is a different limit.

Matching the true structure improves recovery and fit, and the noise distribution matters more than the regularizer. Hyperparameters, including the fixed training temperature, were chosen by random search, and models were trained with $X_\tau$ and evaluated with $X$. In Neural Relational Inferencea graph neural network whose encoder samples a latent interaction graph from observed particle trajectories and whose decoder predicts the motion by passing messages over that graph, trained on a variational lower bound on particles whose latent interaction graph is a spanning tree over 10 vertices, the spanning-tree SST reached an ELBOevidence lower bound: the variational lower bound on $\log p(\text{data})$ used as the training objective of $-1080\pm110$ and edge precision and recall of 91% from 10 observed time steps, against $-1370\pm20$ and 48% precision for the original independent-edge encoder, and from 20 steps it recovered 99% of edges. On a simplified ListOpssequences of nested prefix arithmetic such as max[ 3 min[ 8 2 ]], each labeled by its value from 0 to 9; the syntax defines a parse tree, where the true parse is a directed tree rooted at the first token, the arborescence SST with Gumbel utilities reached 95.0% task accuracy with 75.5% edge precision and 71.9% recall, against 90.1%, 13.0% and 56.4% for independent directed edges and 92.1% accuracy for an LSTM, and the same arborescence with negative exponential utilities reached only 71.5%. In L2Xlearning to explain: a model must predict a label from $k$ input words it selects, so the selected words explain the prediction on the aroma rating of beer reviews, the correlated $k$-subset SST, whose utilities link neighboring words, raised the precision of the selected words from 28.3% with L2X's own relaxation to 63.1% at $k=5$ with the one-layer convolutional encoder and lowered the mean squared error from 0.036 to 0.029, while the four regularizers for plain top-$k$ differed little. In an appendix, score-function estimators on the graph task (REINFORCE with moving-average, batch-mean or multi-sample baselines, and NVILneural variational inference and learning (Mnih and Gregor, 2014): REINFORCE with an input-dependent baseline learned by a network, all applied to $U$) struggled wherever the SSTs struggled, REINFORCE also failed with spanning trees, and NVIL with spanning trees learned some structure, less and with higher variance than the SST.

Rao-Blackwellizing ST-GS over the noise behind the outcome keeps its mean and lowers its mean squared error. ST-GS has two sources of randomness, the outcome $X$ and the Gumbel noise given $X$, and the second enters only through the point at which the softmax Jacobian is evaluated. Conditioning removes it. By the tower rule $\mathbb E[\htmlData{gr=gn13}{\hat g_{\rm GR}}]=\mathbb E[\hat g_{\rm ST}]$, so the bias is unchanged, and conditional Jensen applied to the squared error gives$$\begin{aligned}&\big\|\mathbb E[\hat g_{\rm ST}\mid X]-\nabla\big\|^2\\&\quad\le\mathbb E\big[\|\hat g_{\rm ST}-\nabla\|^2\,\big|\,X\big],\end{aligned}$$whose expectation is Proposition 1, strict whenever $\mathrm{var}[\hat g_{\rm ST}\mid X]\gt0$, which holds for finite $\tau$ and $\theta$. The mean squared error falls at every temperature, so training can move to lower temperatures, where the bias is smaller and plain ST-GS was too noisy. In a three-category example (my check: logits $(0.5,0,-0.5)$, $\tau=0.5$, a quadratic loss), ST-GS and Gumbel-Rao with 100 conditional draws had the same mean, both away from the true gradient, and the trace of the covariance fell from 0.55 to 0.32.

The Gumbels behind a given outcome have a closed form, so the conditional average costs $O(nK)$. Given $X=e_i$, the one-hot vector of the winning category $i$, write $Z=\sum_je^{\theta_j}$ and let $E_1,\dots,E_n$ be independent standard exponentialexponentially distributed with mean 1, $\mathbb P(E\gt s)=e^{-s}$ variables. Then, in distribution, with $j$ ranging over the other categories,$$\begin{aligned}\theta_i+G_i&=\log Z-\log E_i,\\\theta_j+G_j&=-\log\big(E_je^{-\theta_j}+E_i/Z\big).\end{aligned}$$The winning perturbed logit is Gumbel with location $\log Z$, independent of which index wins, and each other one is Gumbel with location $\theta_j$ truncated to lie below the winner (my check: these draws match rejection sampling from unconditional Gumbels within Monte Carlo error). Averaging the softmax Jacobian over $K$ such draws gives $\hat g_K$ ($K=1$ is ST-GS), and Jensen applied to the $K$-sample average gives Proposition 2. The extra work is $K$ softmax Jacobians of size $n$, done in parallel. Proposition 3 follows from the law of total variance: $K$ divides only the first term, which the batch size $B$ divides too, so the gain is largest at small batches and returns from larger $K$ diminish.

Gumbel-Rao trains better at low temperatures and small batches, by up to 2 nats at batch size 20. On a quadratic objective over the probability simplex in three dimensions, the variance falls at every point, most at low temperature and near the corners and edges. On ListOps with a tree-LSTM parser that its authors (Choi et al., 2017) trained with ST-GS, test accuracy for sequences up to length 10 rose at $\tau=0.01$ from 38.8% to 66.5% ($K=1000$; 66.4% with $K=10$) and at $\tau=1$ from 65.8% to 67.0%, so most of the gain is that low temperatures become usable. In a variational autoencoderan encoder draws latent variables given the input and a decoder reconstructs the input from them; both are trained on a lower bound on the log-likelihood on binarized MNIST with 16-category latent variables and batch size 20, the test bound on the negative log-likelihood improved from 98.8 to 96.8 natsunits of natural-log likelihood; the numbers are upper bounds on the negative log-likelihood, so lower is better, and with batch size 200 from 92.6 to 92.2 (an appendix table shows up to 2.9 nats at batch size 5); the best ST-GS models trained at an average temperature of 0.65 and the best Gumbel-Rao models at 0.35, and gains saturated near $K=100$. REINFORCE with a running-mean baseline reached only 169 to 204 nats at batch size 20. The relaxed Gumbel-Softmax, which changes the forward pass, was best for binary variables (98.2 against 100.5 nats at batch size 20) and matched or trailed Gumbel-Rao with 16 categories (96.8 against 96.8 at batch size 20, 94.3 against 92.2 at 200).