Gapped Straight-Through Estimator

• The paper demonstrates that GST deterministically approximates discrete gradients using a gapped softmax surrogate, significantly reducing computational overhead.
• It eliminates Monte Carlo resampling by applying a push-up and gap perturbation to logits, ensuring low-variance gradient estimates during training.
• Empirical results show that GST outperforms alternatives like STGS and GR-MC variants in efficiency and performance on tasks such as MNIST-VAE and unsupervised parsing.

The Gapped Straight-Through (GST) estimator is a deterministic, low-variance gradient estimator designed for efficient training of discrete latent variable models. GST addresses the persistent challenge of high-variance gradient estimation in settings where the forward pass involves discrete random variables, such as categorical latent variables in deep generative models. Unlike traditional variance reduction techniques that rely on Monte Carlo resampling—incurring significant computational overhead—GST provides a single-pass alternative by constructing a surrogate that is deterministically conditioned on the sampled discrete variable, eliminating the need for repeated stochastic evaluations while maintaining or improving variance reduction performance (Fan et al., 2022).

1. Discrete Latent Variables and the Gradient Bottleneck

The central problem tackled by GST is the non-differentiability of models involving discrete random variables. Formally, given parameters $\theta$ and a discrete latent variable $D \in \{e_1, \ldots, e_N\}$, with mass function $p_\theta(D)$ and a downstream cost $g(D)$ (such as the negative ELBO in a VAE), the population loss is

$L(\theta) = \mathbb{E}_{D \sim p_\theta}[g(D)].$

For typical one-hot discrete variables, $\partial D / \partial \theta$ is zero almost everywhere, obstructing standard gradient backpropagation. Consequently, surrogate gradient estimators are necessary to facilitate end-to-end training in models where latent variables mediate model outputs.

2. From Straight-Through Gumbel-Softmax to GST

GST builds on insights from the Straight-Through (ST) family of estimators, particularly the Straight-Through Gumbel-Softmax (STGS) approach. The essential template is to introduce a surrogate $h(\theta, \xi)$ such that

$$D_{ST}(\theta, \xi) = D(\bar{\theta}, \xi) - h(\bar{\theta}, \xi) + h(\theta, \xi),$$

where $\bar{\theta} = \text{stop\_grad}(\theta)$. In forward evaluation, $D_{ST} = D$, whereas in the backward pass, the surrogate's gradient $D \in \{e_1, \ldots, e_N\}$0 is used.

STGS introduces reparameterization by drawing Gumbel noise $D \in \{e_1, \ldots, e_N\}$1 and computes the differentiable surrogate using softmax with temperature $D \in \{e_1, \ldots, e_N\}$2 as

$D \in \{e_1, \ldots, e_N\}$3

where $D \in \{e_1, \ldots, e_N\}$4 are unnormalized logits. STGS is then given by

$D \in \{e_1, \ldots, e_N\}$5

Fan et al. identify three core properties crucial to the efficacy of STGS estimators:

• Consistency: The surrogate maintains the same argmax index as the discrete draw.
• Zero-gradient perturbation: Surrogate randomness does not reintroduce gradients with respect to the frozen parameters.
• Strict gap: There exists a lower bound on the gap between the largest and second-largest perturbed logits, ensuring a highly peaked surrogate (Fan et al., 2022).

3. Construction of the Gapped Surrogate

GST achieves the three key properties in a deterministic fashion, thereby obviating the need for Monte Carlo samples. The surrogate $D \in \{e_1, \ldots, e_N\}$6 is constructed as follows:

• For a forward sample $D \in \{e_1, \ldots, e_N\}$7 (a one-hot vector indicating the chosen category), define:
  • A “push-up” perturbation $D \in \{e_1, \ldots, e_N\}$8, elevating the selected logit's value:

  $D \in \{e_1, \ldots, e_N\}$9 - A “gap” perturbation $p_\theta(D)$0, depressing non-selected logits by at least $p_\theta(D)$1:

  $p_\theta(D)$2

• The GST surrogate:

$p_\theta(D)$3

• The GST estimator is then

$p_\theta(D)$4

with $p_\theta(D)$5.

Conditioning on $p_\theta(D)$6, the estimator is deterministic—there is no surrogate generation randomness.

4. Implementation and Iterative Update

For a downstream loss $p_\theta(D)$7, the gradient estimate is:

$p_\theta(D)$8

In practice, this reduces to a single sample:

$p_\theta(D)$9

A high-level pseudocode for one GST forward/backward pass:

$\partial D / \partial \theta$5

For hyperparameters, $g(D)$0 is robust. For analytic exactness, $g(D)$1 can be used, but this increases complexity.

5. Comparative Variance Reduction and Computational Efficiency

Comparisons between GST and alternatives highlight the principal computational advantage of GST:

• STGS offers low variance but relies on Gumbel-Softmax noise. Additional variance reduction, such as in the GR-MC$g(D)$2 estimator (Paulus et al.), requires $g(D)$3 samples and $g(D)$4 softmax computations per step.

• GST achieves comparable or superior variance reduction in a single pass—no repeated resampling or softmax invocations. Conditioning the surrogate entirely on the observed $g(D)$5 removes all stochasticity from the estimator, except for the discrete variable sampling itself.

• Empirical results indicate that GST's gradient variance is lower than both STGS and GR-MC$g(D)$6 during training. Per-iteration cost for GST matches STGS, whereas GR-MC$g(D)$7 incurs approximately 100x higher memory/time costs due to multiple softmax operations (Fan et al., 2022).

6. Empirical Performance and Applications

On benchmark tasks for discrete generative latent variable models, GST demonstrated improved performance relative to established baselines.

• MNIST-VAE (30 categorical latents × 10 classes):

  • STGS (T=0.5): average test $g(D)$8 (std 2.51)
  • GR-MC$g(D)$9 (T=0.5): $L(\theta) = \mathbb{E}_{D \sim p_\theta}[g(D)].$0 (std 3.01)
  • GST ($L(\theta) = \mathbb{E}_{D \sim p_\theta}[g(D)].$1, $L(\theta) = \mathbb{E}_{D \sim p_\theta}[g(D)].$2): $L(\theta) = \mathbb{E}_{D \sim p_\theta}[g(D)].$3 (std 1.08)
• ListOps Unsupervised Parsing:
  • STGS (T=1.0): test accuracy $L(\theta) = \mathbb{E}_{D \sim p_\theta}[g(D)].$4 (std 0.006)
  • GR-MC$L(\theta) = \mathbb{E}_{D \sim p_\theta}[g(D)].$5 (T=1.0): $L(\theta) = \mathbb{E}_{D \sim p_\theta}[g(D)].$6 (std 0.009)
  • GST ($L(\theta) = \mathbb{E}_{D \sim p_\theta}[g(D)].$7, $L(\theta) = \mathbb{E}_{D \sim p_\theta}[g(D)].$8): $L(\theta) = \mathbb{E}_{D \sim p_\theta}[g(D)].$9 (std 0.005)

In both settings, GST attains lower test loss or higher accuracy compared to strong baselines, with the added benefit of reduced variance and unchanged computational cost.

7. Practical Considerations and Recommendations

GST can be interpreted as enforcing a “peaked” softmax surrogate by systematically elevating the chosen logit's value and depressing non-chosen logits by a fixed margin $\partial D / \partial \theta$0. The deterministic nature of the surrogate, given $\partial D / \partial \theta$1, is central to its variance reduction properties and single-pass efficiency.

• Recommended hyperparameters are $\partial D / \partial \theta$2 and $\partial D / \partial \theta$3.
• Very low temperatures ($\partial D / \partial \theta$4) may require a mixed-temperature schedule.
• GST is directly applicable to a variety of model classes, including discrete VAEs, unsupervised parsing, and discrete-action reinforcement learning, serving as a drop-in replacement for STGS (Fan et al., 2022).

With these features, the Gapped Straight-Through estimator provides a robust and computationally efficient means of training models with discrete latent variables, obviating the necessity for costly Monte Carlo variance-reduction schemes.