Rectified Straight-Through Estimator (ReSTE)

• ReSTE is a surrogate gradient method that interpolates between identity and non-differentiable functions to enable effective training of binary neural networks.
• It introduces an equilibrium perspective by quantifying estimating error and gradient instability, allowing fine-tuning of the estimator-stability trade-off.
• Empirical evaluations on CIFAR-10 and ImageNet show that ReSTE outperforms classic STE methods, achieving superior accuracy without extra auxiliary modules.

The Rectified Straight-Through Estimator (ReSTE) is a class of surrogate gradient methods designed to address fundamental limitations in training neural networks with hard discrete operations, specifically in the context of binary neural networks (BNNs) and vector-quantized models. ReSTE systematically interpolates between the classic identity-based Straight-Through Estimator (STE) and the non-differentiable target function (e.g., sign or quantization), introducing both theoretical and practical mechanisms for balancing gradient fidelity and stability (Wu et al., 2023, Huh et al., 2023).

1. Motivation and Equilibrium Perspective

Neural network binarization compresses models by forcing parameters or activations to discrete (often binary) values. A canonical example is using the sign function $\operatorname{sign}(z)$ in forward computations. However, direct optimization is infeasible due to the sign function’s zero (almost everywhere) and undefined (at zero) derivatives. The classic STE replaces the backward gradient of the non-differentiable operator with that of a smooth proxy, typically the identity, but this introduces a critical estimator inconsistency: gradients do not reflect the true discrete nature of the forward path.

ReSTE introduces a quantitative equilibrium perspective on surrogate gradient design. Two indicators are defined:

• Estimating Error ($e$): $$e = \|\,\operatorname{sign}(z)-f(z)\|_{2} = \sqrt{\sum_{i=1}^d(\operatorname{sign}(z_i)-f(z_i))^2}$$
• Gradient Instability ($s$): $$s = \operatorname{var}(|g|) = \tfrac1N \sum_j \left(\,|g_j|-\mu\right)^2$$, $\mu$ the mean absolute gradient

A decrease in estimating error (using a sharper estimator) results in increased gradient instability, risking vanishing/exploding gradients, and vice versa. Effective training requires an equilibrium between these competing factors (Wu et al., 2023).

2. From Classic STE to Power-Function-Based ReSTE

2.1 Standard STE

STE, as implemented in BinaryConnect and DoReFa, uses

• Forward: $z_b = \operatorname{sign}(z)$
• Backward: Replace $\frac{d}{dz}\operatorname{sign}(z)$ with $\mathbf{1}_{|z|\le 1}$ (identity within $[-1,1]$, zero otherwise)

This approach yields highly stable gradients but maximizes proxy error everywhere except close to $z=0$, decoupling training signals from binarization boundaries.

2.2 Rectified STE: The ReSTE Power Function

ReSTE generalizes the backward pass with a one-parameter power function:

• Hyperparameter: $o \geq 1$; controls transition sharpness
• Forward: $z_b = \operatorname{sign}(z)\cdot\beta$ with $\beta = \|z\|_1/n$ (layer-wise scaling)
• Backward:

$$f'(z) = \frac{d}{dz}\big[\operatorname{sign}(z)\,|z|^{1/o}\big] = \frac{1}{o}|z|^{\frac{1-o}{o}} \quad \text{for}\ |z|\le t$$

$f'(z)=0$ for $|z|>t$, a clipping threshold (e.g., $t=1.5$). Near zero, a finite-difference estimate is used to avoid singularities.

When $o=1$, $f(z)=z$ (standard STE); as $o\to\infty$, $f(z)\to\operatorname{sign}(z)$, recapitulating the hard sign. The shape parameter $o$ explicitly controls the trade-off between approximating the sign function (low $e$, high $s$) and preserving stable gradients (high $e$, low $s$).

3. Error–Stability Trade-off and Empirical Characterization

Empirical investigations on tasks such as CIFAR-10 with ResNet-20 demonstrate the quantitative effect of varying $o$:

• Low $o$ (e.g., 1): Low instability $s$, high error $e$, suboptimal final accuracy.
• Intermediate $o$ ($\approx 3$): Near-optimal trade-off, best top-1 accuracy observed.
• High $o$ ($>5$): Exploding $s$, training collapse due to gradient instability.

Reported accuracy curves confirm this: accuracy $\mathrm{Acc}(o)$ exhibits a single-peaked structure, maximizing at intermediate $o$. Both $e(o)$ and $s(o)$ also behave monotonically: $e$ decreases, $s$ increases with $o$ (Wu et al., 2023).

4. ReSTE in Vector-Quantized Architectures

In the context of vector quantization (VQ) layers, the STE is used to backpropagate through non-differentiable quantization:

• Continuous embedding: $z_e = F(x)$
• Quantized code: $z_q = c_k$, $k = \arg\min_j\|z_e-c_j\|^2$
• STE: $\frac{\partial z_q}{\partial z_e} \approx I$

Limitations here include gradient sparsity and index collapse due to codebook–embedding misalignment and update asymmetry.

ReSTE for VQNs comprises three mechanisms (Huh et al., 2023):

1. Affine Re-Parameterization: Codes $c_i = A\,s_i + b$ allow global moment matching, ensuring that all codes receive updates and encoder-codebook alignment improves.
2. Alternating Optimization: EM-style separation of codebook and encoder/decoder updates reduces the STE gradient gap:

$$\Delta_{\text{gap}}(h,F) := \|\nabla_{F(x)}L_{\text{task}}(G(z_e)) - \nabla_{F(x)}L_{\text{task}}(G(z_q))\| \leq K\|z_e-z_q\|$$

($K$-Lipschitz $G$).

1. Synchronized Commitment Update: Gradients of $L_{\text{task}}$ flow to code vectors each step, avoiding a one-step update lag.

5. Algorithmic Implementation

A typical ReSTE-based BNN layer requires the following steps per iteration:

• Forward: $z_b = \operatorname{sign}(z) \times \beta$, $\beta = \|z\|_1/\dim(z)$.
• Backward: For each $z_i$,
  • If $|z_i|>t$, $f_i'=0$
  • Else if $|z_i|<m$, finite-difference $f_i'$
  • Else $f_i' = \frac{1}{o}|z_i|^{(1-o)/o}$
• Gradient synthesis: $\nabla L/\nabla z_i = g_{z_b,i} \times f_i'$

Typical hyperparameters:

• Optimizer: SGD, LR = 0.1, cosine decay
• STE-clip $t=1.5$, finite-diff threshold $m=0.1$
• $o$ linearly increased from 1 to $o_{\rm end}=3$ over epochs
• Batch size and augmentations as in canonical baselines

For VQNs, alternating optimization and code re-parametrization are incorporated, following an EM-like update schedule, with improved commitment loss (Huh et al., 2023).

6. Empirical Evaluation and Comparative Results

ReSTE has been validated across standard image classification and generative modeling tasks:

BNN (CIFAR-10, ImageNet):

Backbone	Method	W/A	Auxiliary	Top-1 Acc.	Top-5 Acc.
ResNet-20	IR-Net	1/1	Module	85.40%	—
ResNet-20	LCR-BNN	1/1	Loss	86.00%	—
ResNet-20	RBNN	1/1	Module	86.50%	—
ResNet-20	ReSTE	1/1	—	86.75%	—
ResNet-18	FDA	1/1	Module	60.20%	82.30%
ResNet-18	LCR-BNN	1/1	Loss	59.60%	81.60%
ResNet-18	ReSTE	1/1	—	60.88%	82.59%

VQN (ImageNet-100, Generative/Recon):

• Affine re-param, synchronized updates, and alternating STE improve accuracy (AlexNet to 57.9%, ResNet-18 to 71.0%, ViT to 56.7%), perplexity, and FID over baselines (Huh et al., 2023).

Ablation studies confirm ReSTE’s flexibility and rational design: it outperforms STE (84.44%), DSQ (84.11%), and RBNN (85.87%) with 86.75% accuracy on CIFAR-10/ResNet-20, without requiring extra auxiliary modules or losses.

7. Limitations and Future Directions

Open challenges remain in automating the optimal selection of ReSTE’s shape parameter $o$ for diverse tasks, architectures, or data regimes. Extending the equilibrium perspective beyond single-bit quantization to multi-bit or other discrete mappings is an active research direction. Theoretical convergence properties under ReSTE have yet to be fully established (Wu et al., 2023).

A plausible implication is that ReSTE’s explicit mechanism for controlling the estimator–stability trade-off could generalize to other non-differentiable neural operators, facilitating principled surrogate gradient design across a broad spectrum of quantized and discretized architectures.