Beyond Discreteness: Finite-Sample Analysis of Straight-Through Estimator for Quantization

TL;DR

This work provides the first finite-sample convergence analysis of the straight-through estimator (STE) for quantization in neural networks. By studying a two-layer network with binary weights and activations under Gaussian inputs and possible label noise, it derives explicit sample-size thresholds that guarantee ergodic convergence of STE-based optimization to the binary optimum and, separately, conditions under which the last iterate also converges, with a notable recurrence behavior in the presence of noise. The analysis leverages concentration bounds in the $\ell_\infty$ norm and connects to 1-bit compressed sensing through RAIC-inspired techniques and occupation-time (dynamical systems) methods. These results reveal the critical role of data size in STE performance and provide quantitative guidance for quantization-aware training, including the counterintuitive benefit of iterates recurrences to avoid stagnation. Empirical tests corroborate the quadratic scaling with dimension and illustrate the recurrence phenomenon under noise, highlighting practical implications for finite-sample quantization in resource-constrained settings.

Abstract

Training quantized neural networks requires addressing the non-differentiable and discrete nature of the underlying optimization problem. To tackle this challenge, the straight-through estimator (STE) has become the most widely adopted heuristic, allowing backpropagation through discrete operations by introducing surrogate gradients. However, its theoretical properties remain largely unexplored, with few existing works simplifying the analysis by assuming an infinite amount of training data. In contrast, this work presents the first finite-sample analysis of STE in the context of neural network quantization. Our theoretical results highlight the critical role of sample size in the success of STE, a key insight absent from existing studies. Specifically, by analyzing the quantization-aware training of a two-layer neural network with binary weights and activations, we derive the sample complexity bound in terms of the data dimensionality that guarantees the convergence of STE-based optimization to the global minimum. Moreover, in the presence of label noises, we uncover an intriguing recurrence property of STE-gradient method, where the iterate repeatedly escape from and return to the optimal binary weights. Our analysis leverages tools from compressed sensing and dynamical systems theory.

Figures (1)

• Figure 1: Left: Recovery rate for the ergodic averaged iterate $\mathcal{Q}(\overline{\boldsymbol{w}}^T)$ in the noiseless case. Middle: Recovery rate for the last-iterate $\boldsymbol{w}^T$ in the noiseless case. Right: Curves for $\|\boldsymbol{w}^t-\boldsymbol{w}^*\|$ vs. $t$ and $L(\boldsymbol{w}^t)$ vs. $t$ in the noisy case with $m=128, n = 25, N = 140$.

Theorems & Definitions (32)

• Theorem 1: Informal
• Theorem 2: Informal
• Theorem 3: Informal
• Theorem 4
• Theorem 5
• Corollary 1
• proof : Proof of Corollary \ref{['cor:noiseless_part_bound']}
• proof : Proof of Theorem \ref{['thm:RAIC_single_term']}
• Lemma 1
• proof : Proof of Lemma \ref{['lem: Concentration of main term 1']}