A Tutorial on Straight-Through Estimators (STEs)

The Problem

Imagine you're training a neural network with a step function:

def step(x):
    return 1.0 if x > 0 else 0.0

The gradient of this function is zero everywhere (except at x=0 where it's undefined). If you try to backpropagate through it, you get zero gradients - your parameters won't update!

This problem appears whenever we have:

• Discrete decisions (rounding, thresholding, sampling)
• Piecewise constant functions
• Non-differentiable operations

The Core Idea

A Straight-Through Estimator (STE) says: "In the forward pass, use the non-differentiable function. In the backward pass, pretend it was something differentiable."

The simplest STE is the identity STE:

# Forward: y = step(x)
# Backward: pretend dy/dx = 1

A Concrete Example

Let's implement a sign function with an STE:

import torch

class SignSTE(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        return torch.sign(x)  # Returns -1, 0, or 1
    
    @staticmethod
    def backward(ctx, grad_output):
        return grad_output    # Pretend derivative is 1

Now gradients can flow through the sign function!

How STEs Work in Practice

Consider training a network to output binary values:

# Without STE: gradients get stuck
logits = model(input)
binary_output = (logits > 0).float()  # Gradient is 0!
loss = criterion(binary_output, target)

# With STE: gradients flow
logits = model(input)
binary_output = SignSTE.apply(logits)  # Gradient passes through
loss = criterion(binary_output, target)

The JumpReLU Example

In the JumpReLU paper, they need to train through:

• $\text{JumpReLU}_\theta(z) = z \cdot H(z - \theta)$
• The L0 penalty: $\sum_i H(\pi_i(x) - \theta_i)$

Both involve the Heaviside step function $H$, which has zero gradient with respect to $\theta$.

Their STE solution uses a "soft" gradient in a small window:

# Forward: use actual step function
y = (x > theta).float()

# Backward: use soft gradient
grad_theta = -(1/ε) * rectangle((x - theta)/ε) * grad_output

This gives gradient signal when $x$ is within $\epsilon$ of $\theta$.

When Do STEs Work?

STEs are surprisingly effective when:

1. The forward function preserves information: Even though $\text{sign}(x)$ loses magnitude, it preserves direction
2. Gradients point in roughly the right direction: We don't need exact gradients, just reasonable descent directions
3. The mismatch is local: STEs work best when the difference between forward and backward is small

Common Applications

1. Quantized Neural Networks: Training networks with discrete weights/activations
2. Discrete Latent Variables: Training VAEs with categorical variables
3. Hard Attention: Training attention mechanisms that make discrete selections
4. Sparse Models: Like the JumpReLU SAEs, where features are either on or off

Key Insights

The magic of STEs is that optimization often just needs "good enough" gradients. Perfect gradients are great, but biased gradients that point roughly the right way can still find good solutions.

As the JumpReLU paper shows, you can even design custom STEs based on your problem structure - they use kernel density estimation to create gradients that estimate the true expected loss gradient!

Code Example: Threshold Learning

Here's a simplified version of how JumpReLU uses STEs:

class ThresholdSTE(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, threshold, bandwidth=0.001):
        ctx.save_for_backward(x, threshold)
        ctx.bandwidth = bandwidth
        return x * (x > threshold)
    
    @staticmethod  
    def backward(ctx, grad_output):
        x, threshold = ctx.saved_tensors
        ε = ctx.bandwidth
        
        # Gradient w.r.t x: normal gradient where x > threshold
        grad_x = grad_output * (x > threshold)
        
        # Gradient w.r.t threshold: custom STE
        in_window = ((x - threshold).abs() < ε/2)
        grad_threshold = -threshold/ε * in_window * grad_output
        
        return grad_x, grad_threshold, None

This allows learning thresholds even though the step function has zero gradient!