Straight-through Gumbel-Softmax estimator
Mathematical Methods AIMech:StraightThroughGumbelSoftmax DRAFT
Aliases: ST Gumbel-Softmax, ST Gumbel estimator
A surrogate-gradient estimator that uses a hard one-hot argmax of a Gumbel-Softmax sample in the forward computation and substitutes the corresponding continuous softmax relaxation’s derivative in the backward computation.
Identity Notes & Concept Scope
This is Jang and colleagues’ straight-through Gumbel estimator, which retains a discrete forward sample while using a temperature-controlled continuous backward surrogate. It is distinct from using soft Gumbel-Softmax vectors in both passes, from exact differentiation of argmax, from a score-function estimator, and from the broader family of straight-through estimators for other stochastic or quantized units. No normalized identity match was found across all 442 pinned AIO labels/scoped synonyms and 331 preceding native IDs/names/aliases. Full-checkout fixed-string searches included ignored and hidden files in worktree and main, returning 11 and 11 lines respectively. All 10 unique hit contents were inspected: neighboring Gumbel-Top-k contrasts, saved review and ledger references, new source leads, and incidental concrete-distribution wording in an older random-search report, not existing candidate records. Exact checkout counts and differences are retained in the ledger. Nine exact AIO CLASS nodes with scoped synonyms/comments and outgoing edges and six full native neighboring records were inspected. Absence is bounded to this snapshot and semantic neighborhood.
Claims & Grounded Evidence
2 Claims
Mechanism
#method
Draw independent Gumbel noise and compute y=softmax((log(alpha)+G)/tau) with positive finite weights and tau>0. In the forward pass, use z as the one-hot vector at argmax(y), equivalently argmax(log(alpha)+G). In the backward pass, approximate the derivative of z with respect to its parameters by the derivative of y, while the downstream computation has received z. Thus the downstream gradient evaluated at the hard forward state is multiplied by the relaxed sample’s Jacobian. For log-weight coordinates ell and fixed noise/temperature, that Jacobian has entries y_i(delta_ij−y_j)/tau. Hard output remains one-hot even at large positive temperature, and its categorical sampling law is independent of tau; temperature changes the backward surrogate. This is a deliberately biased surrogate, not the ordinary derivative of the hard argmax or generally the gradient obtained by evaluating the whole loss on y.
Scope
#hard-versus-soft
The source discretizes the relaxed vector by argmax for the forward pass and uses the continuous approximation for backpropagation. Its hard samples remain sparse at high temperature, unlike the positive-temperature relaxed vectors themselves. This separates the forward representation from the surrogate derivative.
⚠️ Curated Limitations & Negative Results
- Exact categorical forward sampling does not establish an unbiased gradient of the discrete expected loss. The forward/backward mismatch can persist even when the noise distribution is sampled exactly.
- The downstream computation must supply a usable derivative with respect to the hard vector, typically through a differentiable extension. A purely discrete black-box reward does not acquire such a derivative automatically.
- The derivative of hard argmax is zero almost everywhere and undefined at ties; substituting the soft Jacobian changes that derivative rather than proving argmax differentiable.
- Positive temperature and positive finite weights are required by the stated relaxed Jacobian. Temperature annealing, zero-mass categories, finite-precision ties and stability require explicit choices; there is no universal unbiasedness or variance guarantee.
- The estimator does not equal all straight-through rules or all Concrete training objectives. No author implementation, learned model or empirical speed/quality result is certified.
Artificial Intelligence Ontology (AIO) Grounding
Mapped AIO Classes
| AIO Target Class | Mapping Relation | Mapping Rationale |
|---|---|---|
AIO:SoftmaxFunction
Softmax Function
|
related_match | The estimator substitutes the Jacobian of a temperature-scaled softmax relaxation for the derivative of its hard categorical forward sample. The exact AIO node supplies this smooth backward component, not the hard forward estimator’s identity. This is a method-to-component related match, not equivalence or is-a; neither a neural layer nor a differentiable LLM is asserted. |
Evidence Sources & Primary Literature
1 Source| Source ID | Publication / Document | Type / Access | Reference & Checksum |
|---|---|---|---|
method-source
|
Eric Jang, Shixiang Gu, Ben Poole
|
primary paper
full_text
|
ID: arXiv:1611.01144v5
sha256:e424dbd5b4e8…
|
Curation Provenance
Curator: Codex
Method: ai_assisted
Model: GPT-6
Created: 2026-10-11T16:30:45Z
Updated: 2026-10-11T16:30:45Z
Source File:
data/curated/StraightThroughGumbelSoftmax.yaml