SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Baselines: For ProxConnect++, we consider the 6 variants in Table 1. With different choices of the forward quantizer  $\mathsf{F}_{\mathrm{r}}^{\mu}$  and the backward quantizer  $\mathsf{B}_{\mathrm{r}}^{\mu}$ , we include the full precision (FP) baseline and 5 binarization methods: BinaryConnect (BC) [11], ProxConnect (PC) [13], Binary Neural Network (BNN) [22], the original BNN+ [12], and the modified BNN++ with  $\mathsf{F}_{\mathrm{r}}^{\mu} = \mathsf{SS}$ . Note that we linearly increase  $\mu$  in BNN++ to achieve full binarization in the end. We also compare ProxConnect++ with the ProxQuant and reverseProxConnect baselines.
Hyperparameters: We apply the same training hyperparameters and fine-tune/end-to-end training for 100/300 epochs across all models. For binarization methods: (1) PQ (ProxQuant): similar to Bai et al. [2], we apply the LinearQuantizer (LQ), see (10) in Appendix A, with initial  $\rho_0 = 0.01$  and linearly increase to  $\rho_T = 10$ ; (2) rPC (reverseProxConnect): we use the same LQ for rPC; (3) ProxConnect++: for PC, we apply the same LQ; for BNN+, we choose  $\mu = 5$  (no need to increase  $\mu$  as the forward quantizer is sign); for BNN++, we choose  $\mu_0 = 5$  and linearly increase to  $\mu_T = 30$  to achieve binarization at the final step.
Across all the experiments with random initialization, we report the mean of three runs with different random seeds. Furthermore, we provide the complete results with error bars in Appendix G.
4.2 CNN as backbone
We first compare PC++ against baseline methods on various tasks employing CNNs:
(1) Binarizing weights only (BW), where we simply binarize the weights and keep the other components (i.e., activations and accumulations) in full precision.
Note that we use pre-trained models provided by Dosovitskiy et al. [14] on the ImageNet-21K/ImageNet-1K for fine-tuning ViT-B model on the ImageNet-1K/CIFAR datasets, respectively.
Table 3: Our results on binarizing vision transformers (binarizing weights only). We compare five variants of ProxConnect++ (BC, PC, BNN, BNN+, and BNN++) with FP, PQ, and rPC. End-to-end training tasks are marked as **bold** (i.e., ImageNet-1K for DeiT-T/S/B), where the results are the mean of five runs with different random seeds.
Model	Dataset	FP	PQ	rPC	ProxConnect++
BC	PC	BNN	BNN+	BNN++
ViT-B	CIFAR-10	98.13%	85.06%	86.22%	87.97%	90.12%	89.08%	88.12%	90.24%
CIFAR-100	87.14%	72.07%	73.52%	76.35%	78.13%	77.23%	77.10%	79.22%
ImageNet-1K	77.91%	57.65%	55.33%	63.24%	66.33%	65.31%	63.55%	66.33%
DeiT-T	CIFAR-10	94.86%	82.76%	82.25%	83.10%	85.15%	86.12%	85.91%	86.41%
CIFAR-100	72.37%	54.55%	55.66%	59.65%	60.15%	60.06%	59.77%	60.33%
ImageNet-1K	72.20%	61.23%	60.35%	63.22%	66.15%	65.00%	66.67%	67.34%
DeiT-S	CIFAR-10	95.10%	81.67%	80.23%	84.85%	85.13%	85.09%	85.16%	86.19%
CIFAR-100	73.19%	45.55%	46.66%	60.12%	61.59%	60.55%	60.17%	62.98%
ImageNet-1K	79.91%	69.87%	68.74%	73.16%	73.51%	73.77%	73.23%	73.53%
DeiT-B	CIFAR-10	98.72%	85.22%	86.35%	88.95%	90.53%	90.21%	89.03%	90.67%
CIFAR-100	86.65%	72.11%	73.40%	75.40%	78.55%	76.22%	76.51%	78.30%
ImageNet-1K	81.81%	72.54%	70.11%	76.55%	76.61%	75.60%	76.63%	76.74%
Table 4: Results on binarizing vision transformers (BW, BWA, and BWAA) on DeiT-T. We compare 5 variants of ProxConnect++ (BC, PC, BNN, BNN+, and BNN++) with FP, PQ, and rPC. End-to-end training tasks are marked as **bold** (i.e., ImageNet-1K), where we omit the results for BWAA due to training divergence and the reported results are the mean of five runs with different random seeds.
Dataset	Task	FP	PQ	rPC	ProxConnect++
BC	PC	BNN	BNN+	BNN++
CIFAR-10	BW	94.85%	82.76%	82.25%	83.10%	85.15%	86.12%	85.91%	86.41%
BWA	94.85%	82.56%	82.02%	82.89%	85.01%	85.99%	85.66%	86.12%
BWAA	94.85%	81.34%	80.97%	82.08%	84.31%	84.87%	84.72%	85.31%
CIFAR-100	BW	72.37%	54.55%	55.66%	59.65%	60.15%	60.06%	59.77%	60.33%
BWA	72.37%	53.77%	54.98%	59.21%	59.71%	59.66%	59.12%	59.85%
BWAA	72.37%	52.15%	54.36%	58.15%	59.01%	58.72%	58.15%	59.06%
ImageNet-1K	BW	72.20%	61.23%	60.35%	63.23%	66.15%	65.00%	66.67%	67.34%
BWA	72.20%	60.01%	58.77%	62.13%	65.29%	63.75%	65.29%	65.65%
(2) Binarizing weights and activations (BWA), while keeping accumulation in full precision. Similar to the weights, we apply the same forward-backward proximal quantizer to binarize activations.
(3) Binarizing weights, activations, with 8-bit accumulators (BWAA). BWAA is more desirable in certain cases where the network bandwidth is narrow, e.g., in homomorphic encryption. To achieve BWAA, in addition to quantizing the weights and activations, we follow the implementation of WrapNet [43] and quantize the accumulation of each layer with an additional cyclic function. In practice, we find that with 1-bit weights and activations, the lowest bits we can successfully employ to quantize accumulation is 8, while any smaller choice would raise a high overflow rate and cause the network to diverge. Moreover, BWAA highly relies on a good initialization and cannot be successfully trained end-to-end in our evaluation (and hence omitted).
Note that for the fine-tuning pipeline, we initialize the model with their corresponding pre-trained full precision weights. For the end-to-end pipeline, we utilize random initialization. We report our results in Table 2 and observe: (1) the PC family outperforms baseline methods (i.e., PQ and rPC), and achieves competitive performance on both small and larger scale datasets; (2) BNN++ performs consistently better and is more desirable among the five variants of PC++, especially on BWA and BWAA tasks. Its advantage over BNN+ further validates our theoretical guidance.
4.3 Vision transformer as backbone
Next, we perform similar experiments on the three tasks on vision transformers.
Table 5: Ablation study on the effect of the scaling factor, normalization, pre-training, and knowledge distillation. Experiments are performed on CIFAR-10 with ViT-B.
Method	Scaling	Normalization	Pre-train	KD	Accuracy
PC	X	X	X	X	0.10%
✓	X	X	X	12.81%
✓	✓	X	X	66.51%
✓	✓	✓	X	88.53%
✓	✓	✓	✓	90.13%
BNN++	X	X	X	X	1.50%
✓	X	X	X	23.55%
✓	✓	X	X	77.22%
✓	✓	✓	X	89.05%
✓	✓	✓	✓	90.22%
Implementation on vision transformers: While network binarization is popular for CNNs, its application for vision transformers is still rare<sup>6</sup>. Here we apply four protocols for implementation:
(1) We keep the mean  $s_n$  of full precision weights  $\mathbf{w}_n^*$  for each layer  $n$  as a scaling factor (can be thus absorbed into  $\mathrm{F}_r^{\mu_t}$ ) for the binary weights  $\mathbf{w}_n$ . Such an approach keeps the range of  $\mathbf{w}_n^*$  during binarization and significantly reduces training difficulty without additional computation.
(2) For binarized vision transformer models, LayerNorm is important to avoid gradient explosion. Thus, we add one more LayerNorm layer at the end of each attention block.
(3) When fine-tuning a pre-trained model (full precision), the binarized vision transformer usually suffers from a bad
initialization. Thus, a few epochs of pre-training on the binarized vision transformer is extremely helpful and can make fine-tuning much more efficient and effective.
(4) We apply the knowledge distillation technique in BiBERT [45] to boost the performance. We use full precision pre-trained models as the teacher model.
Main Results: We report the main results of binarizing vision transformers in Table 3 (BW) and Table 4 (BW, BWA, BWAA), where we compare ProxConnect++ algorithms with the FP, PQ, and rPC baselines on fine-tuning and end-to-end training tasks. We observe that: (1) ProxConnect++ variants generally outperform PQ and rPC and are able to binarize vision transformers with less than  $10\%$  accuracy degradation on the BW task. In particular, for end-to-end training, the best performing ProxConnect++ algorithms achieve  $\approx 5\%$  accuracy drop; (2) Among the five variants, we confirm BNN++ is also generally better overall for vision transformers. This provides evidence that our Corollary 1 allows practitioners to easily design many and choose the one that performs best empirically; (3) With a clear underlying optimization objective, BNN++ again outperforms BNN+ across all tasks, which empirically verifies our theoretical findings on vision transformers; (4) In general, we find that weight binarization achieves about  $30\mathrm{x}$  reduction in memory footprint, e.g., from 450 MB to 15 MB for ViT-B.
Ablation Studies: We provide further ablation studies to gain more insights and verify our binarization protocols for vision transformers.
Post-training Binarization: in Figure 2, we verify the difference between PTB (post-training bina- rization) and BAT (binarization-aware training) on the fine-tuning task on CIFAR-10 across different models. Note that we use  $\mathrm{BNN + + }$  as a demonstration of BAT. We observe that without optimization during fine-tuning, the PTB approach fails in general, thus confirming the importance of considering BAT for vision transformers.
Effect of binarizing protocols: here we show the effect of the four binarizing protocols mentioned at the beginning, including scaling the binarized weights using the mean of full precision weights
Figure 2: Comparison between Full Precision (FP) model, BNN++, and Post-training Binarization (PTB) on the fine-tuning task on CIFAR-10.
<sup>6</sup>Notably, He et al. [19] also consider binarizing vision transformers, which we compare our implementation details and experimental results against in Appendix E.
Figure 3: Results of binarizing different components (blocks) of ViT-B architecture on CIFAR-10. Warmer color indicates significant accuracy degradation after binarization.
scaling), adding additional LayerNorm layers (normalization), BAT on the full precision pre-trained models (pre-train) and knowledge distillation. We report the results in Table 5 and confirm that each protocol is essential to binarize vision transformers successfully.
Which block should one binarize: lastly, we visualize the sensitivity of each building block to binarization in vision transformers (i.e., ViT-B) on CIFAR-10 in Figure 3. We observe that binarizing blocks near the head and the tail of the architecture causes a significant accuracy drop.
In this work we study the popular approximate gradient approach in neural network binarization. By generalizing ProxConnect and proposing PC++, we provide a principled way to understand forward-backward quantizers and cover most existing binarization techniques as special cases. Furthermore, PC++ enables us to easily design the desired quantizers (e.g., the new BNN++) with automatic theoretical guarantees. We apply PC++ to CNNs and vision transformers and compare its variants in extensive experiments. We confirm empirically that PC++ overall achieves competitive results, whereas BNN++ is generally more desirable.
We anticipate our work to further enable training and deploying advanced machine learning models to resource limited devices and environments, and help reducing energy consumption and carbon footprint at large. We do not foresee any direct negative societal impact. One limitation we hope to address in the future is to build a theoretical framework that will allow practitioners to quickly evaluate different forward-backward quantizers for a variety of applications.
We thank the reviewers and the area chair for thoughtful comments that have improved our final draft. We thank Arash Ardakani, Ali Mosleh and Marzieh Tahaei for their early participation in this project. YY gratefully acknowledges NSERC and CIFAR for funding support.
[51] M. S. Ryoo, A. Piergiovanni, A. Arnab, M. Dehghani, and A. Angelova. "TokenLearner: What Can 8 Learned Tokens Do for Images and Videos?" arXiv preprint arXiv:2106.11297. 2021.
[52] H. Touvron, M. Cord, M. Douze, F. Massa, A. Sablayrolles, and H. Jégou. "Training data-efficient image transformers & distillation through attention". In: International Conference on Machine Learning. 2021.
[53] Z. Tu, X. Chen, P. Ren, and Y. Wang. "Adabin: Improving binary neural networks with adaptive binary sets". In: European conference on computer vision. Springer. 2022, pp. 379-395.
[54] P. Wang, Q. Chen, X. He, and J. Cheng. “Towards accurate post-training network quantization via bit-split and stitching”. In: International Conference on Machine Learning. 2020, pp. 9847–9856.
[55] S. Xu, Y. Li, T. Ma, B. Zeng, B. Zhang, P. Gao, and J. Lu. "TerViT: An Efficient Ternary Vision Transformer". arXiv preprint arXiv:2201.08050. 2022.
[56] Z. Xu, M. Lin, J. Liu, J. Chen, L. Shao, Y. Gao, Y. Tian, and R. Ji. "Recu: Reviving the dead weights in binary neural networks". In: Proceedings of the IEEE/CVF international conference on computer vision. 2021, pp. 5198-5208.
[57] J. Zhang, H. Peng, K. Wu, M. Liu, B. Xiao, J. Fu, and L. Yuan. "MiniViT: Compressing Vision Transformers with Weight Multiplexing". In: IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022, pp. 12145-12154.
[58] X. Zhao, Y. Wang, X. Cai, C. Liu, and L. Zhang. "Linear symmetric quantization of neural networks for low-precision integer hardware". In: International Conference on Learning Representations. 2019.
Appendix for Understanding Neural Network Binarization with Forward and Backward Proximal Quantizers
A More on Proximal Quantizers
Dockhorn et al. [13] gave a complete characterization of the proximal quantizer  $\mathbf{P}_{\mathrm{r}}$ : a (multi-valued) mapping  $\mathbf{P}$  is a proximal quantizer (of some underlying regularizer  $\mathbf{r}$ ) iff it is monotone, compact-valued and with a closed graph. We now give a few examples to illustrate the ubiquity of proximal quantizers, as well as the generality of PC:
- Identity function: apparently, choosing  $\mathbf{P}_{\mathrm{r}}^{\mu_t}$  as the identity function recovers the full precision training.
-  $\mathbf{P}_{\mathrm{r}}^{\mu_t} = \mathbf{P}_Q$ : as  $Q = \{\pm 1\}$ , this choice recovers exactly BC in (2).
-  $\mathbf{P}_{\mathrm{r}}^{\mu_t} = \mathbf{L}_{\rho}^{\varrho}$ : This is the general piecewise linear quantizer designed by Dockhorn et al. [13]. Recall that  $Q = \{q_k\}_{k=1}^2$ , where  $q_1 = -1, q_2 = +1$ , such that  $p_2 = 0$  is the middle point. By introducing two parameters  $\rho, \varrho \geq 0$ , we can define two shifts:
$$
\text {h o r i z o n t a l}: q _ {1} ^ {-} = q _ {1}, q _ {1} ^ {+} = p _ {2} \wedge (q _ {1} + \rho)
$$
$$
q _ {2} ^ {-} = p _ {2} \vee (q _ {2} - \rho), q _ {2} ^ {+} = q _ {2}
$$
vertical:  $p_2^- = q_1 \vee (p_2 - \varrho), p_2^+ = q_2 \wedge (p_2 + \varrho)$ .
Then, we define  $\mathsf{L}_{\rho}^{\rho}$  as the piece-wise linear map (that simply connects the points by straight lines):
$$
\mathbf {L} _ {\rho} ^ {\varrho} \left(w ^ {*}\right) = \left\{ \begin{array}{l l} q _ {1}, & \text {i f} q _ {1} ^ {-} \leq w ^ {*} \leq q _ {1} ^ {+} \\ q _ {1} + \left(w ^ {*} - q _ {1} ^ {+}\right) \frac {p _ {2} ^ {-} - q _ {1}}{p _ {2} - q _ {1} ^ {+}}, & \text {i f} q _ {1} ^ {+} \leq w ^ {*} <   p _ {2} \\ p _ {2} ^ {+} + \left(w ^ {*} - p _ {2}\right) \frac {q _ {2} - p _ {2} ^ {+}}{q _ {2} ^ {-} - p _ {2}} & \text {i f} p _ {2} <   w ^ {*} \leq q _ {2} ^ {-} \\ q _ {2}, & \text {i f} q _ {2} ^ {-} \leq w ^ {*} \leq q _ {2} ^ {+} \end{array} . \right. \tag {10}
$$
For the middle points,  $\mathbf{L}_{\rho}^{\varrho}(w^{*})$  can be regarded as the intermediate state between the identity function and  $\mathbf{P}_Q$  such that, where  $\mathbf{L}_{\rho}^{\varrho}(w^{*})$  may take any value within the two limits. Note that  $\rho$  controls the discretization vicinity, such that in practice,  $\rho$  is linearly increased over time to fulfill binary weights in the end. We visualize examples of  $\mathbf{L}_{\rho}^{\varrho}(w^{*})$  in Figure 4.
(a)  $\rho = 0, \varrho = 0.2$ .
(b)  $\rho = \underline{\varrho} = 0.2$
Figure 4: Different instantiations of the proximal map  $\mathsf{L}_{\rho}^{\varrho}$  in (10) for  $Q = \{-1,1\}$ .
(c)  $\rho = 0.2,\varrho = 0$
Vision Transformer. In computer vision, vision transformers have become one of the most popular backbone architectures. Dosovitskiy et al. [14] is the first to modify the transformer model to enable images as input, namely the ViT model. Specifically, Dosovitskiy et al. [14] translates an image to a sequence of flattened image patches as input, and applies a self-attention mechanism to retrieve patch-wise information in the feature representation. Touvron et al. [52] further equips ViT with knowledge distillation and proposes DeiT that generalizes well on smaller models and datasets. Liu et al. [34] further proposes Swin as a hierarchical vision transformer that computes representation with shifted windows.
In vision transformers, the main computation overhead is the multi-head attention layer, whose cost is quadratic with the length of the image patches. As a result, such models are in general expensive to train. To reduce the computational cost, different compression techniques have been explored. For instance, Pan et al. [44] performs dynamic pruning for less important patches; Bhojanapalli et al. [6] reuses attention scores computed for one layer in multiple building blocks; Hou and Kung [21] applies multi-dimensional model compression. In this paper, we focus on an alternative approach, namely network quantization.
Network Quantization. We consider two possible scenarios of network quantization:
(1) Post-training Quantization: We first discuss the easier post-training quantization methods. Such approaches usually quantize the full-precision pre-trained model and directly apply it for inference. Post-training quantization is widely used in CNNs [1, 3, 9, 29, 41, 42, 54]. Liu et al. [38] is the first to explore PTQ for vision transformers. It optimizes the quantization intervals and considers ranking information in the loss function. However, it only considers quantization to a 6-bit model without severe performance degradation. For lower-bit quantization, it is essential to leverage training.
(2) Quantization-Aware Training (QAT): different from post-training quantization, quantization-aware training leverage quantization during pre-training or fine-tuning. Thus it can be formulated as an optimization problem for learning the optimal quantized weights [5, 10, 15, 17, 23-25, 39, 58]. Compared with PTQ, QAT can obtain less accuracy drop in low-bit quantization compared to the full-precision model. Li et al. [30] and Xu et al. [55] demonstrate that QAT requires a unique design to quantize vision transformers and it is possible to perform quantization to 3 bit without severe performance degradation. [19] further performs binarization with softmax-aware binarization and information preservation.
Binarization Techniques: (1) Here we summarize existing binarization approaches that can be either justified or improved by  $\mathrm{PC}++$ : Some existing implementations [31, 37, 46] set the forward quantizer as sign and design the backward quantizer in an ad hoc fashion (based on graphic approximation of the sign function), e.g., Liu et al. [37] applies a piece-wise polynomial approximation; Lin et al. [31] improves [37] with a dynamic polynomial approximation; Qin et al. [46] designs a dynamic error decay estimator based on the tanh function. These methods cannot be justified with  $\mathrm{PC}++$ , but could be improved by designing new forward quantizers using our Corollary 2. We discuss these new variants in Appendix F and compare them with our methods. Other implementations [36, 53] applies shift transformation on both forward and backward quantizers, which belongs to the  $\mathrm{PC}++$  family.
(2) Architecture design that can be further integrated into our PC++ as future work: Xu et al. [56] designs a rectified clamp unit to address "dead weights"; Rastegari et al. [48] applies the absolute mean of weights and activations; Martinez et al. [40] uses real-to-binary attention matching and data-driven channel re-scaling; Kim et al. [26] proposes a shifted activation function.
(3) Optimization refinement, again could be integrated into our framework as future work: Liu et al. [33] utilizes state-aware gradient state; Liu et al. [35] provides a weight decay scheme; and Helwegen et al. [20] proposes Bop, a new optimizer for BNNs.
C Additional Theoretical Results
Corollary 1. A pair of forward-backward quantizers  $(\mathsf{F},\mathsf{B})$  admits the decomposition in (6) (for some smoothing parameter  $\mu$  and regularizer  $\mathfrak{r}$ ) iff both  $\mathsf{F}$  and  $\mathsf{B}$  are functions of  $\mathbf{P}(w) := \int_{-\infty}^{w}\frac{1}{\mathsf{B}(\omega)}\mathrm{d}\mathsf{F}(\omega)$ , which is proximal (i.e., monotone, compact-valued and with a closed graph).
Proof. We first recall the decomposition in (6):
$$
F _ {r} ^ {\mu} := T \circ P _ {r} ^ {\mu}, \quad B _ {r} ^ {\mu} := T ^ {\prime} \circ P _ {r} ^ {\mu}. \tag {19}
$$
Suppose first that  $(\mathsf{F},\mathsf{B})$  satisfies the above decomposition. Clearly, both  $\mathsf{F}$  and  $\mathsf{B}$  are functions of  $\mathbf{P} = \mathbf{P}_{\mathrm{r}}^{\mu}$ . Moreover,
$$
\frac {F ^ {\prime} (\omega)}{B (\omega)} = \frac {\mathbf {P} _ {r} ^ {\mu^ {\prime}} (\omega) \cdot T ^ {\prime} \circ \mathbf {P} _ {r} ^ {\mu}}{T ^ {\prime} \circ \mathbf {P} _ {r} ^ {\mu}} = \mathbf {P} _ {r} ^ {\mu^ {\prime}} (\omega)
$$
$$
\int_ {- \infty} ^ {w} \frac {1}{B (\omega)} d F (\omega) = \mathbf {P} _ {r} ^ {\mu} (w) - \mathbf {P} _ {r} ^ {\mu} (- \infty),
$$
which is clearly proximal.
Conversely, let  $\mathbf{P}(w) := \int_{-\infty}^{w} \frac{1}{\mathsf{B}(\omega)} \mathrm{d}\mathsf{F}(\omega)$  be proximal. Taking (generalized) derivative we obtain
$$
\mathbf {P} ^ {\prime} (\omega) = \frac {\mathsf {F} ^ {\prime} (\omega)}{\mathsf {B} (\omega)}.
$$
Since  $\mathsf{B}$  is a function of  $\mathbf{P}$ , say  $\mathsf{B} = \mathsf{T}' \circ \mathbf{P}$ , performing integration we obtain
$$
\mathsf {F} = \mathsf {T} \circ \mathsf {P},
$$
up to some immaterial constant (that can be absorbed into  $\mathsf{T}$ ). Thus,  $(\mathsf{F},\mathsf{B})$  satisfies the decomposition (6).
The following convergence guarantee for  $\mathrm{PC}++$  follows directly from the results in Dockhorn et al. [13]:
Theorem 1. Fix any  $\mathbf{w}$ , the iterates in (7) satisfy:
$$
\sum_ {\tau = s} ^ {t} \eta_ {\tau} \left[ \left\langle \mathbf {w} _ {\tau} - \mathbf {w}, \widetilde {\nabla} \ell (\mathsf {T} \mathbf {w} _ {\tau}) \right\rangle + \mathrm {r} (\mathbf {w} _ {\tau}) - \mathrm {r} (\mathbf {w}) \right] \leq \Delta_ {s - 1} (\mathbf {w}) - \Delta_ {t} (\mathbf {w}) + \sum_ {\tau = s} ^ {t} \Delta_ {\tau} (\mathbf {w} _ {\tau}), \tag {11}
$$
where  $\Delta_{\tau}(\mathbf{w})\coloneqq \mathsf{r}_{\tau}(\mathbf{w}) - \mathsf{r}_{\tau}(\mathbf{w}_{\tau +1}) - \left\langle \mathbf{w} - \mathbf{w}_{\tau +1},\mathbf{w}_{\tau +1}^{*}\right\rangle$  is the Bregman divergence induced by the (possibly nonconvex) function  $\mathsf{r}_{\tau}(\mathbf{w})\coloneqq \mu_{\tau +1}\cdot \mathsf{r}(\mathbf{w}) + \frac{1}{2}\| \mathbf{w}\| _2^2$  (Recall that  $\mu_t\coloneqq 1 + \sum_{\tau = 1}^{t - 1}\eta_{\tau}$ .)
The summand on the left-hand side of (11) is related to the duality gap, which is a natural measure of stationarity for the nonconvex problem (4). Indeed, it reduces to the familiar ones when convexity is present:
Theorem 2. For convex  $\ell \circ \mathsf{T}$  and any  $\mathbf{w}$ , the iterates in (7) satisfy:
$$
\min  _ {\tau = s, \dots , t} \mathbb {E} [ f (\mathbf {w} _ {\tau}) - f (\mathbf {w}) ] \leq \frac {1}{\sum_ {\tau = s} ^ {t} \eta_ {\tau}} \cdot \mathbb {E} \left[ \Delta_ {s - 1} (\mathbf {w}) - \Delta_ {t} (\mathbf {w}) + \sum_ {\tau = s} ^ {t} \Delta_ {\tau} (\mathbf {w} _ {\tau}) \right]. \tag {12}
$$
If  $r$  is also convex, then
$$
\min  _ {\tau = s, \dots , t} \mathbb {E} [ f (\mathbf {w} _ {\tau}) - f (\mathbf {w}) ] \leq \frac {1}{\sum_ {\tau = s} ^ {t} \eta_ {\tau}} \cdot \mathbb {E} \left[ \Delta_ {s - 1} (\mathbf {w}) + \sum_ {\tau = s} ^ {t} \frac {\eta_ {\tau} ^ {2}}{2} \| \widetilde {\nabla} \ell (\mathbf {w} _ {\tau}) \| _ {2} ^ {2} \right], \tag {13}
$$