Layerwise Progressive Freezing Enables STE-Free Training of Deep Binary Neural Networks

We investigate progressive freezing as an alternative to straight-through estimators (STE) for training binary networks from scratch. Under controlled training conditions, we find that while global progressive freezing works for binary-weight networks, it fails for full binary neural networks due to activation-induced gradient blockades. We introduce StoMPP (Stochastic Masked Partial Progressive Binarization), which uses layerwise stochastic masking to progressively replace differentiable clipped weights/activations with hard binary step functions, while only backpropagating through the unfrozen (clipped) subset (i.e., no straight-through estimator). Under a matched minimal training recipe, StoMPP improves accuracy over a BinaryConnect-style STE baseline, with gains that increase with depth (e.g., for ResNet-50 BNN: +18.0 on CIFAR-10, +13.5 on CIFAR-100, and +3.8 on ImageNet; for ResNet-18: +3.1, +4.7, and +1.3). For binary-weight networks, StoMPP achieves 91.2% accuracy on CIFAR-10 and 69.5% on CIFAR-100 with ResNet-50. We analyze training dynamics under progressive freezing, revealing non-monotonic convergence and improved depth scaling under binarization constraints.

💡 Research Summary

The paper tackles a fundamental limitation of Straight‑Through Estimators (STE) in training binary neural networks (BNNs). STE works by applying a non‑differentiable sign function in the forward pass while using a surrogate gradient in the backward pass, creating a mismatch between the forward operator and the gradient used for learning. This mismatch becomes more severe as networks deepen, leading to instability and accuracy loss. The authors explore whether a progressive freezing strategy—gradually fixing parameters to their discrete values—can replace STE. Their initial experiments show that a global freezing schedule works for binary‑weight networks (BWNs) but fails for full BNNs because once a binary activation is frozen, its derivative is zero almost everywhere, blocking gradient flow to earlier layers (the “activation‑induced gradient blockade”).

To overcome this, they propose StoMPP (Stochastic Masked Partial Progressive Binarization). StoMPP combines two ideas: (1) layer‑wise progressive freezing, where layers are binarized sequentially from input to output, ensuring that at any time there is a continuous suffix providing a gradient path to the currently transitioning layer; and (2) stochastic masking with a soft‑refresh mechanism, where each layer maintains a binary mask M indicating frozen (1) or unfrozen (0) entries. During training, a fraction 1/r of the mask entries is resampled each step from a Bernoulli(p) distribution, where p follows a monotonic schedule (cubic by default). This soft‑refresh prevents premature commitment to a particular frozen configuration while still moving toward a fully binary layer.

In the forward pass, frozen entries use the hard sign function, unfrozen entries use a smooth proxy (clip for activations, identity for weights). In the backward pass, gradients are computed only through the smooth proxy; frozen entries receive zero gradient because ∂sign/∂x = 0 almost everywhere. Thus StoMPP is truly STE‑free: it never substitutes a surrogate gradient for the discrete operator.

Algorithmically, all masks start at zero. Training proceeds epoch by epoch; each epoch focuses on a single “transition layer” ℓ. Layers before ℓ are fully frozen (M=1), layer ℓ follows the stochastic masking schedule, and layers after ℓ remain fully continuous (M=0). This ensures the transition layer always receives a valid learning signal from downstream continuous layers, avoiding the gradient blockade observed with global masking.

The authors evaluate StoMPP on CIFAR‑10, CIFAR‑100, and ImageNet using ResNet‑18 and ResNet‑50 architectures, under a minimal recipe (fixed learning rate, no weight decay, no LR schedule) to isolate algorithmic effects. Compared to a BinaryConnect‑style STE baseline, StoMPP yields substantial gains that increase with depth: for ResNet‑50 BNN, improvements of +18.0 % (CIFAR‑10), +13.5 % (CIFAR‑100), and +3.8 % (ImageNet); for ResNet‑18, +3.1 %, +4.7 %, and +1.3 % respectively. In binary‑weight mode, StoMPP reaches 91.2 % on CIFAR‑10 and 69.5 % on CIFAR‑100 with ResNet‑50. Training dynamics show non‑monotonic loss behavior and that the transitioning layer’s accuracy rises in stages, confirming the effectiveness of the progressive schedule. Ablation studies on the freezing schedule p(t) and refresh rate r reveal a sweet spot (cubic p, r = 100) balancing stability and exploration; too aggressive freezing or too frequent mask changes degrade performance.

StoMPP adds negligible computational overhead (sampling O(n/r) indices per step and element‑wise operations). It is architecture‑agnostic and can be combined with existing BNN improvements such as Bi‑Real Net. The work demonstrates that STE is not a necessary component for training deep BNNs; by carefully controlling when and where discretization occurs, one can achieve better scaling, stability, and accuracy. The authors suggest future work on non‑uniform layer schedules, extension to multi‑bit quantization, and integration with advanced BNN architectures.