straight_through

→Tensor

straight_through(hard: Tensor, soft: Tensor)

source edit

Straight-through estimator — hard forward, soft backward.

Returns a tensor numerically equal to hard whose gradient is the gradient of soft:

\mathrm{ST}(h, s) = s + \mathrm{sg}\!\big[h - s\big],

where $\mathrm{sg}$ is the stop-gradient operator. The forward value is exactly $h$ , and $\partial \mathrm{ST} / \partial s$ is the identity, so a discrete or otherwise non-differentiable hard can sit in the middle of a network while the surrounding layers still train.

Parameters

hardTensor

The value the forward pass should carry — typically the output of a non-differentiable step such as a one-hot argmax or a nearest-codebook lookup.

softTensor

The differentiable surrogate whose gradient is borrowed. Must be broadcast-compatible with hard; in practice the two have the same shape.

Returns

Tensor

Equal to hard in value, differentiable with respect to whatever soft depends on.

Notes

Introduced by Bengio, Léonard, and Courville, "Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation" (arXiv:1308.3432). It is the mechanism behind gumbel_softmax with hard=True and behind lucid.nn.VectorQuantizer, which is why it lives here rather than being rewritten at each call site.

The estimator is biased — the returned gradient is not the gradient of the discrete function, which has none — but it is low-variance and empirically trains where the true gradient does not exist.

Examples

>>> import lucid
>>> import lucid.nn.functional as F
>>> soft = lucid.tensor([0.3, 0.7], requires_grad=True)
>>> hard = lucid.tensor([0.0, 1.0])
>>> out = F.straight_through(hard, soft)
>>> out.tolist()               # the forward value is hard
[0.0, 1.0]
>>> out.sum().backward()
>>> soft.grad.tolist()         # gradient passed through untouched
[1.0, 1.0]

Used by 3

>>> import lucid >>> import lucid.nn.functional as F >>> soft = lucid.tensor([0.3, 0.7], requires_grad=True) >>> hard = lucid.tensor([0.0, 1.0]) >>> out = F.straight_through(hard, soft) >>> out.tolist() # the forward value is hard [0.0, 1.0] >>> out.sum().backward() >>> soft.grad.tolist() # gradient passed through untouched [1.0, 1.0]