Straight-Through Estimator (STE) kernel

[superposition]

Description

Straight-through estimator for gradients through discrete/quantized operations.

Operation

# Forward: apply discrete operation (argmax, quantize, threshold)
# Backward: pass gradient straight through as if identity
output = straight_through(x, discrete_fn)

# Common variants:
hard_output = ste_hardmax(logits)      # argmax forward, softmax backward  
quant_output = ste_quantize(x, levels)  # quantize forward, identity backward

Implementation

• Custom backward that ignores forward discretization
• Fused with common discrete ops (argmax, round, quantize)
• Support for Gumbel-softmax alternative

Tests

• Forward produces discrete output
• Backward gradient equals upstream gradient (identity)
• Works with argmax (one-hot output)
• Works with quantization (fixed levels)
• Gumbel-softmax temperature annealing
• gradcheck conceptually correct (gradient = identity)

Used In

VQ-VAE, discrete latent models, hard attention, binary networks