SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yiwei Lu*
yiwei.lu@uwaterloo.ca
Xinlin Li
xinlin.li1@huawei.com
Vahid Partovi Nia
vahid.partitioninia@huawei.com
In neural network binarization, BinaryConnect (BC) and its variants are considered the standard. These methods apply the sign function in their forward pass and their respective gradients are backpropagated to update the weights. However, the derivative of the sign function is zero whenever defined, which consequently freezes training. Therefore, implementations of BC (e.g., BNN) usually replace the derivative of sign in the backward computation with identity or other approximate gradient alternatives. Although such practice works well empirically, it is largely a heuristic or "training trick." We aim at shedding some light on these training tricks from the optimization perspective. Building from existing theory on ProxConnect (PC, a generalization of BC), we (1) equip PC with different forward-backward quantizers and obtain ProxConnect++ (PC++) that includes existing binarization techniques as special cases; (2) derive a principled way to synthesize forward-backward quantizers with automatic theoretical guarantees; (3) illustrate our theory by proposing an enhanced binarization algorithm BNN++; (4) conduct image classification experiments on CNNs and vision transformers, and empirically verify that BNN++ generally achieves competitive results on binarizing these models.
The recent success of numerous applications in machine learning is largely fueled by training big models with billions of parameters, e.g., GPTs in large language models [7, 8], on extremely large datasets. However, as such models continue to scale up, end-to-end training or even fine-tuning becomes prohibitively expensive, due to the heavy amount of computation, memory and storage required. Moreover, even after successful training, deploying these models on resource-limited devices or environments that require real-time inference still poses significant challenges.
A common way to tackle the above problems is through model compression, such as pruning [44, 47, 51], reusing attention [6], weight sharing [57], structured factorization [49], and network quantization [16, 30, 32, 38]. Among them, network quantization (i.e., replacing full-precision weights with lower-precision ones) is a popular approach. In this work we focus on an extreme case of network quantization: binarization, i.e., constraining a subset of the weights to be only binary (i.e.,  $\pm 1$ ), with
*Work done during an internship at Huawei Noah's Ark Lab.
the benefit of much reduced memory and storage cost, as well as inference time through simpler and faster matrix-vector multiplications, which is one of the main computationally expensive steps in transformers and the recently advanced vision transformers [14, 34, 52].
For neural network binarization, BinaryConnect [BC, 11] is considered the de facto standard. BC applies the sign function to binarize the weights in the forward pass, and evaluates the gradient at the binarized weights using the Straight Through Estimator [STE, 4] $^2$ . This widely adopted training trick has been formally justified from an optimization perspective: Dockhorn et al. [13], among others, identify BC as a nonconvex counterpart of dual averaging, which itself is a special case of the generalized conditional gradient algorithm. Dockhorn et al. [13] further propose ProxConnect (PC) as an extension of BC, by allowing arbitrary proximal quantizers (with sign being a special case) in the forward pass.
However, practical implementations [e.g., 2, 12, 22] usually apply an approximate gradient of the sign function on top of STE. For example, Hubara et al. [22] employ the hard tanh function as an approximator of sign. Thus, in the backward pass, the derivative of sign is approximated by the indicator function  $\mathbf{1}_{[-1,1]}$ , the derivative of hard tanh. Later, Darabi et al. [12] consider the sign-Swish function as a more accurate and flexible approximation in the backward pass (but still employs the sign in the forward pass).
Despite their excellent performance in practice, approximate gradient approaches cannot be readily understood in the PC framework of Dockhorn et al. [13], which does not equip any quantization in the backward pass. Thus, the main goal of this work is to further generalize PC and improve our understanding of approximate gradient approaches. Specifically, we introduce  $\mathrm{PC} + +$  that comes with a pair of forward-backward proximal quantizers, and we show that most of the existing approximate gradient approaches are special cases of our proximal quantizers, and hence offering a formal justification of their empirical success from an optimization perspective. Moreover, inspired by our theoretical findings, we propose a novel binarization algorithm BNN++ that improves BNN+ [12] on both theoretical convergence properties and empirical performances. Notably, our work provides direct guidance on designing new forward-backward proximal quantizers in the PC++ family, with immediate theoretical guarantees while enabling streamlined implementation and comparison of a wide family of existing quantization algorithms.
Empirically, we benchmark existing PC++ algorithms (including the new BNN++) on image classification tasks on CNNs and vision transformers. Specifically, we perform weight (and activation) binarization on various datasets and models. Moreover, we explore the fully binarized scenario, where the dot-product accumulators are also quantized to 8-bit integers. In general, we observe that BNN++ is very competitive against existing approaches on most tasks, and achieves 30x reduction in memory and storage with a modest 5-10% accuracy drop compared to full precision training.
We summarize our main contributions in more detail:
- We generalize ProxConnect with forward-backward quantizers and introduce ProxConnect++ (PC++) that includes existing binarization techniques as special cases.
- We derive a principled way to synthesize forward-backward quantizers with theoretical guarantees. Moreover, we design a new BNN++ variant to illustrate our theoretical findings.
- We empirically compare different choices of forward-backward quantizers on image classification benchmarks, and confirm that BNN++ is competitive against existing alternatives.
In neural network quantization, we aim at minimizing the usual (nonconvex) objective function  $\ell(\mathbf{w})$  with discrete weights  $\mathbf{w}$ :
$$
\min  _ {\mathbf {w} \in Q} \ell (\mathbf {w}), \tag {1}
$$
where  $Q \subseteq \mathbb{R}^d$  is a discrete, nonconvex quantization set such as  $Q = \{\pm 1\}^d$ . The acquired discrete weights  $\mathbf{w} \in Q$  are compared directly with continuous full precision weights, which we denote
Note that we refer to STE as its original definition by Bengio et al. [4] for binarizing weights, and other variants of STE (e.g., in BNN) as approximate gradient.
as  $\mathbf{w}^*$  for clarity. While our work easily extends to most discrete set  $Q$ , we focus on  $Q = \{\pm 1\}^{d}$  since this binary setting remains most challenging and leads to the most significant savings. Existing binarization schemes can be largely divided into the following two categories.
Post-Training Binarization (PTB): we can formulate post-training binarization schemes as the following standard forward and backward pass:
$$
\mathbf {w} _ {t} = \mathbf {P} _ {Q} (\mathbf {w} _ {t} ^ {*}), \qquad \mathbf {w} _ {t + 1} ^ {*} = \mathbf {w} _ {t} ^ {*} - \eta_ {t} \widetilde {\nabla} \ell (\mathbf {w} _ {t} ^ {*}),
$$
where  $\mathbf{P}_Q$  is the projector that binarizes the continuous weights  $\mathbf{w}^*$  deterministically (e.g., the sign function) or stochastically $^3$ , and  $\widetilde{\nabla}\ell(\mathbf{w}_t^*)$  denotes a sample (sub)gradient of  $\ell$  at  $\mathbf{w}_t^*$ . We point out that PTB is merely a post-processing step, i.e., the binarized weights  $\mathbf{w}_t$  do not affect the update of the continuous weights  $\mathbf{w}_t^*$ , which are obtained through normal training. As a result, there is no guarantee that the acquired discrete weights  $\mathbf{w}_t$  is a good solution (either global or local) to eq. (1).
Binarization-Aware Training (BAT): we then recall the more difficult binarization-aware training scheme BinaryConnect (BC), first initialized by Courbariaux et al. [11]:
$$
\mathbf {w} _ {t} = \mathbf {P} _ {Q} \left(\mathbf {w} _ {t} ^ {*}\right), \quad \mathbf {w} _ {t + 1} ^ {*} = \mathbf {w} _ {t} ^ {*} - \eta_ {t} \widetilde {\nabla} \ell \left(\mathbf {w} _ {t}\right), \tag {2}
$$
where we spot that the gradient is evaluated at the binarized weights  $\mathbf{w}_t$  but used to update the continuous weights  $\mathbf{w}_t^*$ . This approach is also known as Straight Through Estimator [STE, 4]. Note that it is also possible to update the binarized weights instead, effectively performing the proximal gradient algorithm to solve (1), as shown by Bai et al. [2]:
$$
\mathbf {w} _ {t} = \mathbf {P} _ {Q} (\mathbf {w} _ {t} ^ {*}), \qquad \mathbf {w} _ {t + 1} ^ {*} = \mathbf {w} _ {t} - \eta_ {t} \widetilde {\nabla} \ell (\mathbf {w} _ {t}).
$$
This method is known as ProxQuant, and will serve as a baseline in our experiments.
2.1 ProxConnect
Dockhorn et al. [13] proposed ProxConnect (PC) as a broad generalization of BinaryConnect in (2):
$$
\mathbf {w} _ {t} = \mathbf {P} _ {\mathrm {r}} ^ {\mu_ {t}} \left(\mathbf {w} _ {t} ^ {*}\right), \quad \mathbf {w} _ {t + 1} ^ {*} = \mathbf {w} _ {t} ^ {*} - \eta_ {t} \widetilde {\nabla} \ell \left(\mathbf {w} _ {t}\right), \tag {3}
$$
where  $\mu_t \coloneqq 1 + \sum_{\tau=1}^{t-1} \eta_\tau, \eta_t > 0$  is the step size, and  $\mathbf{P}_r^{\mu_t}$  is the proximal quantizer:
$$
\mathbf {P} _ {\mathrm {r}} ^ {\mu} (\mathbf {w}) := \underset {\mathbf {z}} {\operatorname {a r g m i n}} \frac {1}{2 \mu} \| \mathbf {w} - \mathbf {z} \| _ {2} ^ {2} + \mathrm {r} (\mathbf {z}), \text {a n d}
$$
$$
\mathcal {M} _ {\mathrm {r}} ^ {\mu} (\mathbf {w}) := \min  _ {\mathbf {z}} \frac {1}{2 \mu} \left\| \mathbf {w} - \mathbf {z} \right\| _ {2} ^ {2} + r (\mathbf {z}).
$$
In particular, when the regularizer  $\mathfrak{r} = \iota_Q$  (the indicator function of  $Q$ ),  $\mathbf{P}_{\mathrm{r}}^{\mu_t} = \mathbf{P}_Q$  (for any  $\mu_t$ ) and we recover BC in (2). Dockhorn et al. [13] showed that the PC update (3) amounts to applying the generalized conditional gradient algorithm to a smoothened dual of the regularized problem:
$$
\min _ {\mathbf {w}} \left[ \ell (\mathbf {w}) + r (\mathbf {w}) \right] \approx \min _ {\mathbf {w} ^ {*}} \ell^ {*} (- \mathbf {w} ^ {*}) + \mathcal {M} _ {r ^ {*}} ^ {\mu} (\mathbf {w} ^ {*}),
$$
where  $f^{*}(\mathbf{w}^{*}) \coloneqq \max_{\mathbf{w}} \langle \mathbf{w}, \mathbf{w}^{*} \rangle - f(\mathbf{w})$  is the Fenchel conjugate of  $f$ . The theory behind PC thus formally justifies STE from an optimization perspective. We provide a number of examples of the proximal quantizer  $\mathbf{P}_{\mathrm{r}}^{\mu_t}$  in Appendix A.
Another natural cousin of PC is the reversed PC (rPC):
$$
\mathbf {w} _ {t} = \mathbf {P} _ {\mathrm {r}} ^ {\mu_ {t}} (\mathbf {w} _ {t} ^ {*}), \quad \mathbf {w} _ {t + 1} ^ {*} = \mathbf {w} _ {t} - \eta_ {t} \widetilde {\nabla} \ell (\mathbf {w} _ {t} ^ {*}),
$$
which is able to exploit the rich landscape of the loss by evaluating the gradient at the continuous weights  $\mathbf{w}_t^*$ . Thus, we also include it as a baseline in our experiments.
We further discuss other related works in Appendix B.
3We only consider deterministic binarization in this paper.
One popular heuristic to explain BC is through the following reformulation of problem (1):
$$
\min  _ {\mathbf {w} ^ {*}} \ell \big (\mathbf {P} _ {Q} (\mathbf {w} ^ {*}) \big).
$$
Applying (stochastic) "gradient" to update the continuous weights we obtain:
$$
\mathbf {w} _ {t + 1} ^ {*} = \mathbf {w} _ {t} ^ {*} - \boldsymbol {\eta} _ {t} \cdot \mathbf {P} _ {Q} ^ {\prime} (\mathbf {w} _ {t} ^ {*}) \cdot \widetilde {\nabla} \ell (\mathbf {P} _ {Q} (\mathbf {w} _ {t} ^ {*})).
$$
Unfortunately, the derivative of the projector  $\mathbf{P}_Q$  is 0 everywhere except at the origin, where the derivative actually does not exist. BC [11], see (2), simply "pretended" that  $\mathbf{P}_Q' = I$ . Later works propose to replace the troublesome  $\mathbf{P}_Q'$  by the derivative of functions that approximate  $\mathbf{P}_Q$ , e.g., the hard tanh in BNN [22] and the sign-Swish in  $\mathrm{BNN}+$  [12]. Despite their empirical success, it is not clear what is the underlying optimization problem or if it is possible to also replace the projector inside  $\widetilde{\nabla}\ell$ , i.e., allowing the algorithm to evaluate gradients at continuous weights, a clear advantage demonstrated by Bai et al. [2] and Dockhorn et al. [13]. Moreover, the theory established in PC, through a connection to the generalized conditional gradient algorithm, does not apply to these modifications yet, which is a gap that we aim to fill in this section.
3.1 ProxConnect++
To address the above-mentioned issues, we propose to study the following regularized problem:
$$
\min  _ {\mathbf {w} ^ {*}} \ell (\mathsf {T} (\mathbf {w} ^ {*})) + r (\mathbf {w} ^ {*}), \tag {4}
$$
as a relaxation of the (equivalent) reformulation of (1):
$$
\min _ {\mathbf {w} ^ {*}} \ell (\mathsf {P} _ {Q} (\mathbf {w} ^ {*})) + \iota_ {Q} (\mathbf {w} ^ {*}).
$$
In other words,  $\mathsf{T}:\mathbb{R}^d\to \mathbb{R}^d$  is some transformation that approximates  $\mathbf{P}_Q$  and the regularizer  $\mathsf{r}:\mathbb{R}^d\to \mathbb{R}$  approximates the indicator function  $\iota_{Q}$ . Directly applying ProxConnect in (3) we obtain<sup>4</sup>:
$$
\mathbf {w} _ {t} = \mathbf {P} _ {\mathrm {r}} ^ {\mu_ {t}} \left(\mathbf {w} _ {t} ^ {*}\right), \quad \mathbf {w} _ {t + 1} ^ {*} = \mathbf {w} _ {t} ^ {*} - \eta_ {t} \mathsf {T} ^ {\prime} \left(\mathbf {w} _ {t}\right) \cdot \widetilde {\nabla} \ell \left(\mathsf {T} \left(\mathbf {w} _ {t}\right)\right). \tag {5}
$$
Introducing the forward and backward proximal quantizers:
$$
F _ {r} ^ {\mu} := T \circ \mathbf {P} _ {r} ^ {\mu}, \quad B _ {r} ^ {\mu} := T ^ {\prime} \circ \mathbf {P} _ {r} ^ {\mu}, \tag {6}
$$
we can rewrite the update in (5) simply as:
$$
\mathbf {w} _ {t + 1} ^ {*} = \mathbf {w} _ {t} ^ {*} - \eta_ {t} \cdot \mathrm {B} _ {\mathrm {r}} ^ {\mu_ {t}} \left(\mathbf {w} _ {t} ^ {*}\right) \cdot \widetilde {\nabla} \ell \left(\mathrm {F} _ {\mathrm {r}} ^ {\mu_ {t}} \left(\mathbf {w} _ {t} ^ {*}\right)\right). \tag {7}
$$
It is clear that the original ProxConnect corresponds to the special choice
$$
\mathsf {F} _ {r} ^ {\mu} = \mathbf {P} _ {r} ^ {\mu}, \quad \mathsf {B} _ {r} ^ {\mu} \equiv I.
$$
Of course, one may now follow the recipe in (6) to design new forward-backward quantizers. We call this general formulation in (7) ProxConnect++ (PC++, which covers a broad family of algorithms.
Conversely, the complete characterization of proximal quantizers in Dockhorn et al. [13] allows us also to reverse engineer T and r from manually designed forward and backward quantizers. As we will see, most existing forward-backward quantizers turn out to be special cases of our proximal quantizers, and thus their empirical success can be justified from an optimization perspective. Indeed, for simplicity, let us restrict all quantizers to univariate ones that apply component-wise. Then, the following result is proven in Appendix C.
Corollary 1. A pair of forward-backward quantizers  $(\mathsf{F},\mathsf{B})$  admits the decomposition in (6) (for some smoothing parameter  $\mu$  and regularizer  $\mathbf{r}$ ) iff both  $\mathsf{F}$  and  $\mathsf{B}$  are functions of  $\mathbb{P}(w)\coloneqq \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).
<sup>4</sup>We assume throughout that T, and any function whose derivative we use, are locally Lipschitz so that their generalized derivative is always defined, see Rockafellar and Wets [50].
Figure 1: Forward and backward pass for ProxConnect++ algorithms (red/blue arrows indicate the forward/backward pass), where fp denotes full precision, bn denotes binary and back-prop denotes backpropagation.
Importantly, with forward-backward proximal quantizers, the convergence results established by Dockhorn et al. [13] for PC directly carries over to  $\mathrm{PC} + +$  (see Appendix C for details). Let us further illustrate the convenience of Corollary 1 by some examples.
Example 1 (BNN). Hubara et al. [22] proposed BNN with the choice
$$
\begin{array}{l l l} \mathsf {F} = \text {s i g n} & a n d & \mathsf {B} = \mathbf {1} _ {[ - 1, 1 ]}, \end{array}
$$
which satisfies the decomposition in (6). Indeed, let
$$
T (w) = \min  \{1, \max  \{- 1, w \} \}, \tag {8}
$$
$$
\mathbf {P} _ {\mathrm {r}} ^ {\mu} (w) = \left\{ \begin{array}{l l} \frac {1}{\mu} w + \operatorname {s i g n} (w) \left(1 - \frac {1}{\mu}\right), & i f | w | > 1 \\ \operatorname {s i g n} (w), & i f | w | \leq 1 \end{array} . \right. \tag {9}
$$
Since  $\mathsf{B}$  is constant over  $[-1,1]$ , applying Corollary 1 we deduce that the proximal quantizer  $\mathsf{P}_{\mathsf{r}}^{\mu}$ , if exists, must coincide with  $\mathsf{F}$  over the support of  $\mathsf{B}$ . Applying monotonicity of  $\mathsf{P}_{\mathsf{r}}^{\mu}$  we may complete the reverse engineering by making the choice over  $|w| > 1$  as indicated above. We can easily verify the decomposition in (6):
$$
\mathsf {F} = \operatorname {s i g n} = \mathsf {T} \circ \mathbf {P} _ {\mathbf {r}} ^ {\mu}, \mathsf {B} = \mathbf {1} _ {[ - 1, 1 ]} = \mathsf {T} ^ {\prime} \circ \mathbf {P} _ {\mathbf {r}} ^ {\mu}.
$$
Thus, BNN is exactly BinaryConnect applied to the transformed problem in (4), where the transformation  $\mathsf{T}$  is the so-called hard tanh in (8) while the regularizer  $\mathfrak{r}$  is determined (implicitly) by the proximal quantizer  $\mathbf{P}_{\mathfrak{r}}^{\mu}$  in (9).
To our best knowledge, this is the first time the (regularized) objective function that BNN aims to optimize has been identified. The convergence properties of BNN hence follow from the general result of Dockhorn et al. [13] on ProxConnect, see Appendix C.
Table 1: Variants of ProxConnect++.
Forward Quantizer	Backward Quantizer	Algorithm
identity	identity	FP
PQ	identity	BC
Lρ	identity	PC
PQ	1[-1,1]	BNN
PQ	∇SS	BNN+
SS	∇SS	BNN++
Example 2 (BNN+). Darabi et al. [12] adopted the derivative of the sign-Swish (SS) function as a backward quantizer while retaining the sign function as the forward quantizer:
$$
\mathsf {B} (\mathbf {w}) = \nabla \operatorname {S S} (\mathbf {w}) := \mu \left[ 1 - \frac {\mu \mathbf {w}}{2} \tanh  \left(\frac {\mu \mathbf {w}}{2}\right) \right] \tanh  ^ {\prime} \left(\frac {\mu \mathbf {w}}{2}\right), \mathsf {F} = \operatorname {s i g n},
$$
where  $\mu$  is a hyperparameter that controls how well SS approximates the sign. Applying Corollary 1 we find that the derivative of SS (as backward) coupled with the sign (as forward) do not admit the decomposition in (6), for any regularizer  $\mathfrak{r}$ . Thus, we are not able to find the (regularized) objective function (if it exists) underlying BNN+.
We conclude that BNN+ cannot be justified under the framework of PC++. However, it is possible to design a variant of BNN+ that does belong to the PC++ family and hence enjoys the accompanying theoretical properties:
Example 3 (BNN++) We propose that a simple fix of BNN+ would be to replace its sign forward quantizer with the sign-Swish (SS) function:
$$
\mathsf {F} (\mathbf {w}) = \operatorname {S S} (\mathbf {w}) := \frac {\mu \mathbf {w}}{2} \tanh  ^ {\prime} \left(\frac {\mu \mathbf {w}}{2}\right) + \tanh  \left(\frac {\mu \mathbf {w}}{2}\right),
$$
which is simply the primitive of  $\mathsf{B}$ . In this case, the algorithm simply reduces to  $PC++$  applied on (4) with  $\mathsf{r} = 0$  (and hence essentially stochastic gradient descent). Of course, we could also compose with a proximal quantizer to arrive at the pair  $(\mathsf{F} \circ \mathsf{P}_{\mathsf{r}}^{\mu}, \mathsf{B} \circ \mathsf{P}_{\mathsf{r}}^{\mu})$ , which effectively reduces to  $PC++$  applied on the regularized objective in (4) with a nontrivial  $\mathsf{r}$ . We call this variant BNN++.
We will demonstrate in the next section that  $\mathrm{BNN}++$  is more desirable than  $\mathrm{BNN}+$  empirically.
More generally, we have the following result on designing new forward-backward quantizers:
Corollary 2. If the forward quantizer is continuously differentiable (with bounded support), then one can simply choose the backward quantizer as the derivative of the forward quantizer.
This follows from Corollary 1 since  $\mathsf{P}(w) \equiv w$  is clearly proximal. Note that the BNN example does not follow from Corollary 2. In Appendix F, we provide additional examples of forward-backward quantizers based on existing methods, and we show that Corollary 2 consistently improves previous practices.
In summary: (1) ProxConnect++ enables us to design forward-backward quantizers with infinite many choices of T and r, (2) it also allows us to reverse engineer T and r from existing forward-backward quantizers, which helps us to better understand existing practices, (3) with our theoretical tool, we design a new BNN++ algorithm, which enjoys immediate convergence properties. Figure 1 visualizes ProxConnect++ with a variety of forward-backward quantizers.
In this section, we perform extensive experiments to benchmark PC++ on CNN backbone models and the recently advanced vision transformer architectures in three settings: (a) binarizing weights only (BW); (b) binarizing weights and activations (BWA), where we simply apply a similar forward-backward proximal quantizer to the activations; and (c) binarizing weights, activations, with 8-bit dot-product accumulators (BWAA) [43].
Table 2: Binarizing weights (BW), binarizing weights and activation (BWA) and binarizing weights, activation, with 8-bit accumulators (BWAA) on CNN backbones. We consider the fine-tuning (FT) pipeline and the end-to-end (E2E) pipeline. We compare five variants of ProxConnect++ (BC, PC, BNN, BNN+, and BNN++) with FP, PQ, and rPC in terms of test accuracy. For the end-to-end pipeline, we omit the results for BWAA due to training divergence and report the mean of five runs with different random seeds.
Dataset	Pipeline	Task	FP	PQ	rPC	ProxConnect++
BC	PC	BNN	BNN+	BNN++
CIFAR-10	FT	BW	92.01%	89.94%	89.98%	90.31%	90.31%	90.35%	90.27%	90.40%
BWA	92.01%	88.79%	83.55%	89.39%	89.95%	90.01%	89.99%	90.22%
BWAA	92.01%	85.39%	81.10%	89.11%	89.21%	89.32%	89.55%	90.01%
E2E	BW	92.01%	81.59%	81.82%	87.51%	88.05%	89.92%	89.39%	90.03%
BWA	92.01%	81.51%	81.60%	86.99%	87.26%	89.15%	89.02%	89.91%
ImageNet-1K	FT	BW	78.87%	66.77%	69.22%	71.35%	71.29%	71.41%	70.22%	72.33%
BWA	78.87%	56.21%	58.19%	65.99%	65.61%	66.02%	65.22%	68.03%
BWAA	78.87%	53.29%	55.28%	58.18%	59.21%	59.77%	59.10%	63.02%
E2E	BW	78.87%	63.23%	66.39%	67.45%	67.51%	67.49%	66.99%	68.11%
BWA	78.87%	61.19%	64.17%	65.42%	65.31%	65.29%	65.98%	66.08%
Datasets: We perform image classification on CIFAR-10/100 datasets [27] and ImageNet-1K dataset [28]. Additional details on our experimental setting can be found in Appendix D.
Backbone architectures: (1) CNNs: we evaluate CIFAR-10 classification using ResNet20 [18], and ImageNet-1K with ResNet-50 [18]. We consider both fine-tuning and end-to-end training; (2) Vision transformers: we further evaluate our algorithm on two popular vision transformer models: ViT [14] and DeiT [52]. For ViT, we consider ViT-B model and fine-tuning task across all models. For DeiT, we consider DeiT-B, DeiT-S, and DeiT-T, which consist of 12, 6, 3 building blocks and 768, 384 and 192 embedding dimensions, respectively; we consider fine-tuning task on ImageNet-1K pre-trained model for CIFAR datasets and end-to-end training on ImageNet-1K dataset.