Train Through Discrete Decisions with the Straight-Through Estimator

Neural networks are usually trained by following gradients through a chain of differentiable operations. A hard discrete choice breaks that chain. Rounding a value, selecting a binary gate, or quantizing an activation can make the forward computation useful while leaving ordinary backpropagation with a zero or undefined derivative at the decision.

The straight-through estimator (STE) is a practical workaround. It keeps the discrete operation in the forward pass but substitutes a simpler derivative during the backward pass. That makes optimization possible, at the cost of using a gradient that is not the true derivative of the forward computation.

This article builds a precise mental model for that trade. You will see how an STE changes backpropagation, where clipping matters, how to test an implementation, and when a different estimator is a better fit.

Start with a binary gate

Consider a model that computes a scalar score a and turns it into a binary gate:

z = 1 if a >= 0
z = 0 otherwise

The rest of the network uses z, not a. A simple loss might be:

L = (z - t)^2

where t is a target value.

Suppose a = -0.2 and t = 1. The forward pass gives:

z = 0
L = (0 - 1)^2 = 1

The model clearly made an undesirable decision. The problem appears during differentiation. The step function is constant on either side of zero, so its derivative is zero almost everywhere. Applying the ordinary chain rule gives no useful signal for moving a toward the positive side.

An STE changes only the backward rule. A basic version behaves conceptually like this:

forward:  z = step(a)
backward: dz/da := 1

The symbol := is important here. The backward value is assigned as a surrogate; it is not a claim about the mathematical derivative of step.

With that surrogate, the loss gradient can pass through the gate:

dL/da ~= dL/dz

For the example above:

dL/dz = 2(z - t) = -2

so the surrogate gradient pushes a upward. If repeated updates move a across zero, the forward gate changes from 0 to 1.

That is the core idea: use the discrete value for model behavior, but use a deliberately chosen smooth path for optimization.

The forward function and backward rule are different objects

It is easy to describe an STE as if it somehow differentiates rounding or thresholding. It does not.

For a hard threshold,

z = step(a)

the actual derivative is zero almost everywhere and is not defined at the threshold itself. An STE replaces that derivative with a surrogate such as:

dz/da ~= 1

or a clipped form:

dz/da ~= 1 when |a| <= c
dz/da ~= 0 otherwise

for some chosen range c.

The distinction has practical consequences. Standard gradient descent normally follows derivatives of the objective produced by the forward computation. With an STE, the optimizer follows a modified vector field. That field can still produce useful parameters, but familiar guarantees tied to exact gradients do not automatically transfer.

Treat the surrogate derivative as part of the model’s training design, not as an implementation trick that can be ignored during review.

A common implementation pattern

Automatic-differentiation frameworks can express an STE by combining a hard forward value with a soft gradient path. The following pseudocode shows the pattern without depending on a particular library:

soft = a
hard = round(a)

z = soft + stop_gradient(hard - soft)

During the forward pass, stop_gradient returns its input unchanged:

z = soft + (hard - soft) = hard

During the backward pass, the stopped term contributes no derivative, so:

dz/da = dsoft/da = 1

This compact construction separates the two roles cleanly:

value used by forward computation -> hard
local derivative used by backward computation -> soft path

Frameworks use different names for the operation that blocks gradients, such as a detach or stop-gradient primitive. Check the framework contract rather than assuming identical syntax or aliasing behavior.

Clipping can keep the surrogate local

Passing a derivative of 1 everywhere is simple, but it can send a strong update through a unit even when its pre-activation is far from the decision boundary.

Consider a binary gate with threshold zero. A value of a = -0.01 is close to switching state. A value of a = -50 is not. Treating both as having the same local derivative may be a poor approximation for the intended training behavior.

A clipped STE can restrict the surrogate gradient to a region around the threshold:

forward:
    z = step(a)

backward surrogate:
    dz/da = 1 if -1 <= a <= 1
    dz/da = 0 otherwise

The exact interval is a design choice, not a universal constant. Its useful scale depends on parameterization, normalization, and the surrounding network.

Clipping introduces another trade-off. A narrow active region makes the surrogate more local, but units outside that region receive no gradient through the discrete operation. If many units start far outside it, optimization can stall. A wider region sends signal to more units but makes the surrogate less tied to the switching boundary.

Quantization is a natural use case

Quantized networks often need discrete values in the forward path. For example, a simplified weight quantizer might map a real-valued parameter to a small integer grid:

q(w) = round(w / s) * s

where s is the quantization step.

The true derivative of round is zero almost everywhere. If training used that derivative directly, upstream weight updates through the quantizer would vanish.

An STE can instead use:

forward:  q(w) = round(w / s) * s
backward: dq/dw ~= 1 within an allowed range

The model therefore experiences quantized weights during the forward computation while an underlying real-valued parameter can still move during optimization.

This teaching example omits details that matter in production quantization, including scale estimation, zero points, saturation ranges, per-channel parameters, integer deployment kernels, and calibration. The STE addresses the gradient break at a discrete operation; it does not by itself define a complete quantization scheme.

The estimator is biased

A gradient estimator is called biased when its expected value does not equal the gradient of the target objective under consideration. The basic STE is deliberately biased because its surrogate derivative differs from the derivative of the hard forward operation.

That does not make it useless. It means the bias must be treated as an optimization trade rather than hidden behind the word “gradient.”

The surrogate can point in a direction that improves the discrete model, but it can also:

• push parameters according to a poor local approximation;
• create updates whose magnitude depends strongly on the chosen surrogate;
• interact badly with saturation or normalization;
• settle on parameters that work under the training surrogate but are fragile near discrete boundaries.

For this reason, validation must measure the actual hard model. Do not report only a soft proxy that was convenient for backpropagation.

Check the hard path during evaluation

A common mistake is to train with a hard forward operation but accidentally evaluate a continuous relaxation. That can make offline metrics look better than the model that will actually run.

Keep the contract explicit:

training forward path   -> hard decision
training backward path  -> surrogate derivative
evaluation forward path -> hard decision
production path         -> same intended discrete rule

If deployment uses integer quantization, binary gates, or another discrete representation, evaluation should reproduce the relevant behavior closely enough to expose boundary errors and saturation effects.

Also inspect stability near thresholds. A parameter barely above zero and one far above zero can produce the same binary output, but small perturbations affect them differently. Margin statistics can therefore reveal fragility that task accuracy alone hides.

Test the custom gradient, not just the output

A forward-unit test is insufficient for an STE because the unusual behavior exists in the backward pass.

A small test should verify both sides independently. For a rounding STE:

input:   0.2, 0.8, -0.3
forward: 0,   1,    0

Then backpropagate a simple scalar function whose upstream gradient is known. If the chosen surrogate is the identity derivative, confirm that the expected upstream values pass through unchanged.

For a clipped estimator, test values inside, outside, and exactly at the clipping boundary. Decide the boundary convention explicitly so an equality case does not depend on an accidental comparison operator.

Finite-difference gradient checks need special care. They approximate the derivative of the hard forward function, so they should not be expected to match an intentionally different STE surrogate. Instead, test the surrogate rule directly and separately test the hard forward behavior.

Compare against continuous relaxations

An STE is not the only way to train around a discrete choice. Sometimes the model can use a smooth relaxation during optimization and become discrete later.

For a binary gate, a sigmoid with temperature can provide a continuous value:

z_soft = sigmoid(a / temperature)

As the temperature decreases, the transition becomes sharper. This gives an actual derivative of the relaxed forward function, unlike the hard-forward STE.

The choice changes the training problem:

STE
forward uses hard values
backward uses a surrogate
train-deploy behavior can be closely aligned

continuous relaxation
forward uses soft values during training
gradient matches the relaxed computation
train-deploy mismatch may appear when values become hard

A relaxation is attractive when soft intermediate values are meaningful and the deployment gap can be controlled. An STE is attractive when exposing the model to hard decisions during training is important and a biased surrogate is acceptable.

For stochastic discrete variables, other estimators may be more appropriate. Score-function estimators can target gradients of expected objectives without differentiating through the discrete sample itself, but their variance can be high. Reparameterized continuous relaxations offer another route when their assumptions fit the model. There is no single estimator that dominates across bias, variance, compute, and fidelity to the deployed operation.

Common failure modes

The first failure is forgetting that the surrogate is part of the algorithm. Changing an identity STE to a clipped or smooth surrogate can materially change optimization even when the forward outputs are identical.

The second is letting latent real-valued parameters drift far beyond the useful discrete range. If the forward operation saturates but the surrogate keeps passing gradients, parameter magnitude can grow without adding representational value. Clamping, regularization, or a bounded parameterization may help, depending on the model.

The third is using an STE to conceal a discrete search problem that gradient methods handle poorly. If the decision space is small enough for enumeration, dynamic programming, beam search, or another direct method, a biased surrogate may add complexity without a clear benefit.

The fourth is comparing estimators only by training loss. The deployed hard model is the object that matters. Compare task metrics, boundary stability, convergence behavior, runtime, and sensitivity to initialization across several runs when randomness is material.

When an STE is a good fit

An STE is most compelling when three conditions hold together: the forward path genuinely needs a hard operation, exact differentiation through that operation is unavailable, and the surrounding model can tolerate a surrogate gradient.

Typical examples include quantization-aware training, binary or low-bit neural components, and architectures with hard gates whose discrete behavior must appear during training.

Use more caution when the discrete choice controls a long sequence of later decisions. A local surrogate can then ignore large downstream changes caused by flipping one choice. The mismatch between the hard objective and surrogate update may become too severe.

Also consider simpler alternatives before adding a custom gradient. If a continuous model satisfies the product requirement, keeping the computation differentiable removes an entire source of optimization mismatch.

Keep the approximation visible

The straight-through estimator is useful because it separates two needs that conflict: a model may require discrete behavior in its forward computation while gradient-based optimization requires a usable backward signal.

The clean engineering approach is to make that separation explicit. Define the hard operation, define the surrogate derivative, test both paths, evaluate the hard model, and monitor behavior near discrete boundaries. If the surrogate produces unstable or misleading updates, change the estimator rather than treating the backward rule as fixed infrastructure.

An STE is not a derivative of a discrete decision. It is a chosen optimization signal. Keeping that distinction visible makes the method easier to reason about, debug, and replace when the model demands a different trade-off.