SURGE tackles the gradient-mismatch problem in training binary neural networks (BNNs), where weights and activations are constrained to ±1. The standard Straight-Through Estimator (STE) approximates the gradient of the non-differentiable sign function with a surrogate, but this causes gradient mismatch and information loss. SURGE introduces a learnable gradient-compensation framework: for each binarized layer, it adds a parallel full-precision auxiliary branch and decomposes the output so gradients flow through both paths. The full-precision branch estimates higher-order gradient terms beyond STE's first-order approximation, and an adaptive scaler balances the two gradient streams. Experiments on image classification, object detection, and language tasks show SURGE outperforms state-of-the-art BNN training methods.
Main takeaways:
- BNN training relies on gradient approximations because the sign function is non-differentiable; STE and its variants suffer from gradient mismatch and fixed-range clipping.
- SURGE adds a parallel full-precision branch for each binarized layer, allowing gradients to flow through both binarized and full-precision paths.
- The Dual-Path Gradient Compensator (DPGC) decomposes outputs to decouple gradient flow; the full-precision branch captures components STE's first-order approximation misses.
- An Adaptive Gradient Scaler (AGS) dynamically balances gradient contributions from the two branches via norm-based scaling.
- SURGE achieves state-of-the-art results on image classification, object detection, and language-understanding benchmarks.