Technical Summary: SURGE – Surrogate Gradient Adaptation in Binary Neural Networks

1. Problem Statement

Binary Neural Networks (BNNs) offer significant advantages in memory efficiency (theoretically 32× reduction) and computational acceleration (theoretically 58×) by representing weights and activations with 1-bit values. However, the training of BNNs faces a fundamental challenge: the binarization operation (e.g., the sign function) is non-differentiable, causing gradients to vanish during backpropagation.

To address this, the Straight-Through Estimator (STE) and its variants are commonly used to approximate gradients. However, the paper identifies two critical limitations in prevailing methods:

1. Gradient Mismatch: STE provides only a first-order approximation of the sign function's gradient. Since the true gradient is zero almost everywhere, relying on a fixed-value approximation introduces significant estimation bias and optimization instability.
2. Information Loss via Fixed-Range Clipping: To mitigate noise, STE typically employs fixed-range gradient clipping (e.g., preserving gradients only for inputs within $[-1, 1]$). This discards gradient information for values outside this range, leading to suboptimal training, particularly for activation quantization.

Existing approaches often rely on hand-crafted heuristic designs (e.g., piecewise polynomials or specific activation functions) to approximate gradients, which cannot guarantee an optimal solution and fail to fully rectify the accumulated gradient mismatch.

2. Methodology

The authors propose SURGE (Surrogate Gradient Adaptation), a learnable gradient compensation framework designed to mitigate gradient mismatch through auxiliary backpropagation without altering the forward-pass output or inference cost. The framework consists of two primary components:

2.1 Dual-Path Gradient Compensator (DPGC)

DPGC constructs a parallel, full-precision parameterized branch (the auxiliary branch) for every binarized layer (the main branch).

• Forward Pass: The output of the layer is synthesized by combining the binary output and a scaled, detached version of the auxiliary full-precision output. Specifically, the output is formulated as:
  $$\text{output} = f_b(x; W_b) - \text{detach}(f_{ao}(x; W_a)) + f_{ao}(x; W_a)$$
  where $f_b$ is the binary computation, $f_{ao}$ is the scaled full-precision computation, and the "detach" operator ensures the auxiliary branch does not alter the forward output value. This design ensures the inference behavior remains identical to standard BNNs.
• Backward Pass: The gradient flow is decoupled into two paths. The main branch receives the standard STE surrogate gradient ($g_b$), while the auxiliary branch computes a full-precision gradient ($g_a$). The total gradient is the sum of these components, allowing the auxiliary branch to capture higher-order terms and recover gradients that would otherwise be clipped or lost by STE.

2.2 Adaptive Gradient Scaler (AGS)

To prevent the large-magnitude gradients from the auxiliary path from destabilizing the convergence of the main binary branch, SURGE introduces an Adaptive Gradient Scaler.

• Mechanism: AGS dynamically balances the contributions of the binary and auxiliary gradients using a norm-based scaling factor $\lambda_{AGS}$.
• Formula:
  $$\lambda_{AGS} = \eta \frac{\|g_b\|_2}{\|g_a\|_2 + \epsilon}$$
  where $\eta$ is a base scaling coefficient and $\epsilon$ is a numerical stabilizer.
• Theoretical Basis: The scaling factor is derived from a theoretical analysis (Theorem 5.3) aiming to minimize the mean squared error between the composite gradient and an ideal, unobserved reference gradient. The AGS ensures the binary gradient dominates the update direction while the auxiliary gradient provides magnitude-aware compensation.

3. Key Contributions

The paper outlines three main contributions:

1. SURGE Framework: A novel gradient compensation strategy that employs a Dual-Path Gradient Compensator to address gradient mismatch. Crucially, it introduces no additional overhead during inference as the auxiliary branch is discarded post-training.
2. Adaptive Gradient Scaler (AGS): A mechanism that dynamically equilibrates gradient contributions from binary and auxiliary branches based on a theoretically derived optimal scaling factor, ensuring stable optimization.
3. State-of-the-Art Performance: Extensive experiments demonstrate that SURGE achieves superior performance across four standard benchmarks (image classification, object detection, and language understanding) compared to existing methods.

4. Experimental Results

The authors evaluated SURGE on CIFAR-10, ImageNet-1K, PASCAL VOC, and the GLUE benchmark suite.

• Image Classification (CIFAR-10): SURGE outperformed all compared methods (including ReCU, IR-Net, and RBNN) across ResNet-18, ResNet-20, and VGG-Small backbones. For instance, on ResNet-18, SURGE achieved 93.1% Top-1 accuracy, surpassing ReCU (92.8%).
• Image Classification (ImageNet-1K):
  • One-Stage Training: SURGE-trained Binarized ResNet-18 achieved 62.0% Top-1 accuracy, surpassing the previous SOTA (RBONN at 61.4%) by 0.6%.
  • Two-Stage Training: SURGE achieved 66.7% Top-1 accuracy, outperforming RBONN (66.5%) and ReCU (66.4%).
• Object Detection (PASCAL VOC): Using a Faster-RCNN framework with a ResNet-18 backbone, SURGE achieved 77.0% mAP, outperforming the previous SOTA IDa-Det (76.5%) while maintaining the same FLOPs and memory usage. Notably, it surpassed the raw real-valued Faster-RCNN baseline (76.4%) in mAP.
• Language Understanding (GLUE): On the GLUE benchmark using BERT-base, SURGE achieved an average score of 72.0, outperforming BiT (70.6) and BiBERT (63.2).

5. Significance and Claims

The paper claims that SURGE represents a significant advancement in BNN training by addressing the root cause of performance gaps: the gradient mismatch inherent in STE-based training.

• Theoretical Grounding: Unlike previous heuristic approaches, SURGE is grounded in theoretical analysis that derives an optimal scaling factor for gradient compensation.
• Practical Efficiency: The method is "architecture-agnostic" and applies to arbitrary binarized linear operators (including Convolutional and Transformer blocks). It maintains the extreme efficiency of BNNs during inference by discarding the auxiliary branch, ensuring no increase in model size or latency for deployment.
• Robustness: The adaptive nature of the scaler allows the model to handle varying gradient magnitudes dynamically, leading to more stable training and higher accuracy across diverse tasks and architectures.

The authors conclude that SURGE effectively rectifies the limitations of fixed-range clipping and first-order approximations, enabling BNNs to achieve accuracy levels that significantly narrow the gap with full-precision counterparts.