SURGE: Surrogate Gradient Adaptation in Binary Neural Networks

TL;DR

SURGE introduces a learnable gradient adaptation framework for binary neural networks, improving training stability and accuracy by addressing gradient mismatch and information loss issues.

Contribution

It proposes a novel dual-path gradient compensator and adaptive gradient scaler, providing a theoretically grounded, learnable solution for better BNN training.

Findings

SURGE outperforms state-of-the-art methods on image classification tasks.

It improves training stability and reduces gradient mismatch in BNNs.

Experiments show enhanced performance in object detection and language understanding.

Abstract

The training of Binary Neural Networks (BNNs) is fundamentally based on gradient approximation for non-differentiable binarization operations (e.g., sign function). However, prevailing methods including the Straight-Through Estimator (STE) and its improved variants, rely on hand-crafted designs that suffer from gradient mismatch problem and information loss induced by fixed-range gradient clipping. To address this, we propose SURrogate GradiEnt Adaptation (SURGE), a novel learnable gradient compensation framework with theoretical grounding. SURGE mitigates gradient mismatch through auxiliary backpropagation. Specifically, we design a Dual-Path Gradient Compensator (DPGC) that constructs a parallel full-precision auxiliary branch for each binarized layer, decoupling gradient flow via output decomposition during backpropagation. DPGC enables bias-reduced gradient estimation by leveraging the full-precision branch to estimate components beyond STE's first-order approximation. To further enhance training stability, we introduce an Adaptive Gradient Scaler (AGS) based on an optimal scale factor to dynamically balance inter-branch gradient contributions via norm-based scaling. Experiments on image classification, object detection, and language understanding tasks demonstrate that SURGE performs best over state-of-the-art methods.