Technical Summary: Layerwise Progressive Freezing for Depth-Scalable Binary Networks

Problem Statement

Training Binary Neural Networks (BNNs), where both weights and activations are constrained to $\{-1, +1\}$, from scratch remains challenging, particularly as network depth increases. The dominant training approach relies on the Straight-Through Estimator (STE), which substitutes a non-differentiable sign function in the forward pass with a surrogate gradient (typically identity) in the backward pass. While effective for shallow networks, STE suffers from a fundamental forward/backward mismatch. As networks deepen, this mismatch leads to severe accuracy degradation.

Existing literature primarily attempts to mitigate this by refining the STE itself (e.g., better surrogate gradients, gradient shaping, or learned scaling). However, these methods retain the core structure of a hard forward operator paired with a surrogate backward pass. This paper identifies a different axis of failure: activation-induced gradient blockades. In full BNNs, once an activation is "frozen" (committed to a hard binary value), its derivative becomes zero almost everywhere. If activations are frozen at arbitrary depths during training, gradient flow to earlier layers is severed, preventing effective learning in deep architectures. This issue is specific to full BNNs; Binary Weight Networks (BWNs), where activations remain continuous, do not suffer from this blockade.

Methodology: StoMPP

The authors introduce StoMPP (Stochastic Masked Partial Progressive Binarization), a training scaffold that controls when and where binarization is enforced, rather than improving the gradient estimator itself.

Core Mechanism

StoMPP maintains real-valued underlying variables for weights and pre-activations. A binary mask determines whether each entry is treated as discrete (frozen) or continuous (unfrozen). The method employs two key strategies:

1. Layerwise Progressive Scheduling: Unlike global progressive quantization (where parameters are frozen randomly throughout the network), StoMPP freezes layers sequentially from input to output.

   • Frozen Prefix: Layers $1$ to $\ell-1$ are fully binarized.
   • Transition Layer: Layer $\ell$ is partially frozen, evolving from continuous to binary.
   • Unfrozen Suffix: Layers $\ell+1$ to $N$ remain fully continuous.
   • Rationale: This ensures that the layer currently undergoing binarization always has an unfrozen downstream suffix, preserving a valid gradient path to the loss function. This prevents the gradient blockades caused by early freezing of downstream activations.

2. Stochastic Masked Partial Freezing with Soft Refresh: Within the transitioning layer, the mask is not deterministically set. Instead, a fraction $p(\tau)$ of entries is frozen according to a schedule (e.g., cubic). Crucially, the frozen set is "soft-refreshed" at each step: a small subset of indices ($1/r$) is resampled from a Bernoulli distribution. This prevents premature commitment to a specific frozen configuration and allows the model to adapt to the emerging binary pattern.

Gradient Handling

The method is agnostic to the gradient treatment of frozen entries, supporting two regimes:

• STE-free StoMPP: Frozen entries receive zero gradient ($Surr(u) = 0$). This provides a fully STE-free training procedure.
• StoMPP+STE: Frozen entries receive a standard STE identity gradient ($Surr(u) = 1$). This combines the progression structure with surrogate gradients.

The authors conduct ablation studies in the STE-free regime to isolate the contribution of the progression order from the benefits of surrogate gradients.

Key Contributions

1. StoMPP Framework: A layerwise progressive freezing procedure with stochastic masking and soft refresh. It supports both fully STE-free training and hybrid training with STE on frozen entries.
2. Decisive Role of Progression Order: The paper establishes that for full BNNs, the order of binarization is critical.
   • Forward Layerwise (Input $\to$ Output): Succeeds on deep networks.
   • Reverse Layerwise (Output $\to$ Input): Causes catastrophic collapse to near-chance performance.
   • Global Ordering: Degrades significantly with depth.
   • BWN Insensitivity: Binary Weight Networks are largely insensitive to ordering, confirming that the failure mode is specific to binary activations blocking gradients.
3. Training Dynamics Analysis: The authors characterize the training dynamics, noting a "sawtooth" convergence pattern where accuracy drops when a new layer enters the freezing phase and recovers as the unfrozen suffix adapts. They demonstrate that layerwise scheduling is robust to schedule shape and refresh rates, unlike global scheduling.

Experimental Results

The evaluation uses a controlled protocol with a minimal training recipe (constant learning rate, no weight decay) to isolate the binarization rule's effect.

• Depth Scalability: On ResNet-50 BNNs, vanilla STE performance drops significantly as depth increases (e.g., from 41.9% on ResNet-18 to 30.8% on ResNet-50 on ImageNet).
  • StoMPP (STE-free): Improves over vanilla STE by +18.0% (CIFAR-10), +13.5% (CIFAR-100), and +3.8% (ImageNet) on ResNet-50.
  • StoMPP+STE: Achieves even higher gains, improving over vanilla STE by +27.1% (CIFAR-10), +19.8% (CIFAR-100), and +17.7% (ImageNet) on ResNet-50.
• Generalization: The pattern holds across ResNet-18/34, MobileNetV2, and BERT fine-tuning on SST-2. Notably, on SST-2, vanilla STE collapses to near-random (50.9%), while StoMPP+STE reaches 82.8%.
• Ablations:
  • Ordering: Reverse layerwise ordering causes BNNs to collapse (e.g., 8.6% on CIFAR-100 ResNet-50), while forward ordering maintains performance.
  • Hyperparameters: The method is robust to the choice of schedule (cubic, linear, etc.) and refresh rate, provided a layerwise progression is used.
  • Architectures: StoMPP transfers to binary-specific architectures like Bi-Real Net without modification, though fine-tuning interactions vary.

Significance and Claims

The paper claims that the primary bottleneck in training deep BNNs is not the quality of the surrogate gradient, but the timing and location of binarization. By enforcing a forward layerwise progression, StoMPP prevents the formation of activation-induced gradient blockades that sever learning signals in deep networks.

The authors emphasize that this is an orthogonal approach to improving STE; it changes the training scaffold rather than the gradient estimator. The results suggest that the "depth collapse" observed in BNNs is a structural consequence of gradient blockades, which can be mitigated by ensuring a continuous gradient path through unfrozen downstream layers during the binarization process. The method is presented as a training procedure that can be combined with existing architectural improvements or surrogate gradient techniques, offering a path to scalable, fully binary deep networks.