Predicting Probabilities of Error to Combine Quantization and Early Exiting: QuEE

Our methodology involves learning to predict the performance we could achieve given more computational power. This would motivate using $PE(f_{\pi}|x)$ as the target values during training, but we do not have access to these quantities. Therefore, we must rely on approximations to set the learning targets for $\widehat{pe}_{\theta}(f_{\pi}|x)$. There is existing work on estimating these quantities [14, 12], but the proposed procedures tend to be only reliable for high-performing classifiers (we verified this through experiment). The intermediate classifiers at lower blocks are often relatively inaccurate.

We base our approximations on the assumption that $PE(f_{\pi}|x_{i})$ is smooth over some transformation of the space $t_{\pi}(x)$. This allows us to discretize the space into $K$ partitions $\{Q_{1},\dots,Q_{K}\}$ for which, under our assumption, the resulting distances $|PE(f_{\pi}|x_{i})-\frac{1}{p_{Q}}\int_{x\in Q}PE(f_{\pi}|x)P(dx)|$ are small for all $x_{i}\in Q$. Here we have introduced $p_{Q}\triangleq\int_{Q}P(dx)$.

We can then compute the empirical estimator of $\frac{1}{p_{Q}}\int_{Q}PE(f_{\pi}|x)P(dx)$ for each partition using a validation set. Assume that we have $m$ samples $\{x_{j},y_{j}\}_{j=1}^{m}$ in the validation set and denote by $m_{Q}$ the number of samples in $Q$, i.e. $m_{Q}=\sum_{j=1}^{m}\mathbbm{1}[x_{j}\in Q]$. Then:

	$\displaystyle PE(f_{\pi}|x_{i})$	$\displaystyle\approx\frac{1}{\Pr(Q)}\int_{Q}PE(f_{\pi}|x)P(dx),\quad\mathrm{for}\,\,x_{i}\in Q\,,$		(10)
		$\displaystyle\approx\frac{1}{m_{Q}}\sum_{j=1}^{m}\mathbbm{1}[f_{\pi}(x_{j})\neq y_{j},x_{j}\in Q],$		(11)
		$\displaystyle\triangleq\widetilde{pe}(f_{\pi}|x_{i})\text{ our approximation used to train }\widehat{pe}_{\theta}(f_{\pi}|x).$		(12)

In practice, we discretize the space using a clustering algorithm and we use the predicted probability vectors as the clustering features ($t_{\pi}(x)=\hat{\mathbf{p}}_{\pi}$). We provide empirical results and qualitative analysis to motivate this choice in Appendix 7.2 and in the result section.

The learning task of $h^{j}_{\theta}(\mathcal{I}_{j}(x),\pi^{\prime})$ is to predict the ground truth approximation $\widetilde{pe}(f_{\pi}|x)$ of future classifiers $f_{\pi^{\prime}}$. Each classifier is encoded by its associated path $\pi^{\prime}$, and we have access to past predicted probability vectors $\{\hat{\mathbf{p}}_{\pi_{i}[:1]},\hat{\mathbf{p}}_{\pi_{i}[:2]},\dots\hat{\mathbf{p}}_{\pi_{i}[:j]}\}$ through $\mathcal{I}_{j}(x_{i})$. We cannot employ a complex architecture for $h^{j}_{\theta}(\mathcal{I}_{j}(x),\pi^{\prime})$ because the goal is to reduce computational overhead.

We base the decision on the current probability vector $\hat{\mathbf{p}}_{\pi_{i}[:j]}$ and the previous one $\hat{\mathbf{p}}_{\pi_{i}[:j-1]}$. We additionally extract helpful statistics such as entropy $H(\hat{\mathbf{p}}_{\pi_{i}[:j]}),H(\hat{\mathbf{p}}_{\pi_{i}[:j-1]})$ and maximum probability, $\max(\hat{\mathbf{p}}_{\pi_{i}[:j]}),\max(\hat{\mathbf{p}}_{\pi_{i}[:j-1]})$. These are often used in early exiting algorithms [2, 18, 23, 49]. The final vector $u_{i}\in{}^{2|\mathcal{Y}|+4}$ that encodes all the input features is:

$\displaystyle u_{i}$

$\displaystyle=\hat{\mathbf{p}}_{\pi_{i}[:j]}\,||\,\hat{\mathbf{p}}_{\pi_{i}[:j-1]}\,||\,H(\hat{\mathbf{p}}_{\pi_{i}[:j]})\,||\,H(\hat{\mathbf{p}}_{\pi_{i}[:j-1]})\,||\,\max(\hat{\mathbf{p}}_{\pi_{i}[:j-1]})\,||\,\max(\hat{\mathbf{p}}_{\pi_{i}[:j]})\,.$

(13)

As for the path $\pi^{\prime}$, we encode it in a fixed length vector $p_{i}$ of size $E{+}1$ by padding shorter paths with 0s, and add the number of layers evaluated at the end:

$\displaystyle p_{i}$

$\displaystyle=[b_{1},b_{2},\dots,b_{e},0,0,e]\in{}^{E+1}.$

(14)

Our decision architecture is a simple 2-layer network:

$\displaystyle\widehat{pe}_{\theta}(f_{\pi^{\prime}}|x)=\mathrm{sigmoid}(NN(NN(u_{i}||p_{i}))),$

(15)

and we train each predictor module $h^{j}_{\theta}(\mathcal{I}_{j}(x),\pi^{\prime})$ with a standard MSE loss:

$\displaystyle\mathcal{L}(x,\pi^{\prime})=||h^{j}_{\theta}(\mathcal{I}_{j}(x),\pi^{\prime})-\widetilde{pe}(f_{\pi^{\prime}}|x))||_{2}.$

(16)

We start by augmenting the pre-trained backbone with randomly initialized classifiers at predefined layers (See App. 7.1 for details about exit placement). The classifiers have the same architecture as the final inference head. We train intermediate classifiers until convergence, keeping the backbone frozen. We then quantize the network at various bit-widths corresponding to our choices of computation levels, $\{b_{1},\dots,b_{B}\}$, as follows. The highest bit-width in $b_{B}$ is chosen to be the lowest value that maintains the accuracy of the final classifier (usually around 8 bits). We then add decreasing bit-widths to $b_{B}$ as long as the accuracy of the final classifier is within 5% of the original performance. Typically, we obtained at most $B=4$. We perform the quantization at each bit-width using PTQ4ViT [51], a state-of-art PTQ algorithm for transformers (See Appendix 7.1.2 for an in-depth description). The network obtained is thus a multi-quantization architecture comprised of L blocks augmented with an inference head, where each block consists of the backbone layer quantized at the B levels.

QuEE components We finally augment the multi-quantization network with gates $\{g_{\theta}^{j}\}_{j\in E}$ at every exit. Each gate $g_{\theta}^{j}$ contains a small learnable predictor module $h_{\theta}^{j}$. The first gate $g_{\theta}^{1}$ is placed at the start of the network and always routes to the highest computation level.

Before training the predictor modules, we fit a K-means algorithm for every path in the network in order to obtain discretized predictor targets, where $K$ is a hyperparameter. In practice, we observe that paths that move from a low bit resolution to a higher bit resolution result in subpar performance. We thus do not consider these paths during training and inference. This effectively reduces the set of paths $\mathcal{P}$ by a factor of two. If the resulting $|\mathcal{P}|>50$, we randomly sample 50 paths, both during training and at inference. For each valid path $\pi$, we obtain $K$ clusters by clustering the predicted probability vectors $\hat{\mathbf{p}}_{\pi}\in[0,1]^{|\mathcal{Y}|}$ using the validation set. For each cluster $Q$ we compute $\widetilde{pe}(f_{\pi,Q}|x_{i})$ using equation 11. At training, we obtain the discretized targets by predicting the appropriate cluster $Q^{i}$ for $\hat{\mathbf{p}}_{\pi^{\prime}}$ and using $\widetilde{pe}(f_{\pi,Q^{i}}|x_{i})$ as the target in the MSE loss (equation 16). This is optimized with the Adam optimizer and early stopping (using the validation set). The complete set of optimization hyperparameters is provided in Appendix 7.1.

At inference, when the sample $x$ reaches $g^{2}$, the probability of future error $\widehat{pe}_{\theta}(f_{\pi^{\prime}}|x)$ for every path $\pi^{\prime}\in\mathcal{P}_{\pi[1]}$ is computed and $x$ is forwarded along the optimal path, as stated in equation 9. The cost $c_{\pi^{\prime}}$ of each path in equation 9 is the normalized cost in BitOPS [35, 43, 26], computed as follows. For each layer $l$ in $\pi^{\prime}=b_{1}\rightarrow\dots\rightarrow b_{e}$, we compute the layer cost $c_{l}=\text{BitOPS}(l)=\text{FLOPS}(l)\times b_{l}$. The unnormalized cost of the path is $c_{\pi^{\prime},\text{unnorm}}=\sum_{l=1}^{e}c_{l}$. We normalize all path costs by dividing them by the cost of the most costly path, $\pi_{\text{max}}=b_{\text{max}}\rightarrow b_{\text{max}}\dots$, which is incurred when evaluating the entire network using the highest bit-width $b_{\text{max}}$. Therefore:

$$c_{\pi}=\frac{c_{\pi^{\prime},\text{unnorm}}}{c_{\pi_{\text{max}}}}=\frac{\sum_{l=1}^{e}\text{FLOPS}(l)\times b_{l}}{\sum_{l=1}^{L}\text{FLOPS}(l)\times b_{\text{max}}}\,.$$

(17)

At inference, we can obtain different cost-accuracy operating points without retraining QuEE by incorporating a cost-importance hyper-parameter $\lambda\in\mathbb{R}^{+}$ in equation 9. This allows us to prioritize efficiency or accuracy:

$$\hat{\pi}_{\text{adaptive}}^{*}(x)=\operatorname*{arg\,min}_{\pi^{\prime}\in\mathcal{P}_{\pi[:j-1]}}\,\big{\{}c_{\pi^{\prime}}\lambda+\widehat{pe}_{\theta}(f_{\pi^{\prime}}|x)\big{\}}\,.$$

(18)

As our proposal combines early exiting and quantization, we include baselines from both settings.

• •

  DQNet-gate [35] is a data-adaptive mixed-precision quantization architecture. It is a modified version of DQNet [35], where the parameters of the backbone are fixed instead of being trained alongside the bit-selector module. We achieve different accuracy/cost points by changing the target bit-width and cost importance parameter $\alpha$.

• •

  PTQ4ViT [51] is a SOTA PTQ method for ViT models. We report the accuracy obtained a various levels of quantization.

• •

  JEIDNN[40] is a SOTA EE method for frozen backbones. As a backbone model, we use the backbone quantized at the highest precision bit-width we consider.

• •

  Thresholding applies the most popular gating decision used in the EE literature, where we exit if the max probability exceeds a certain threshold.

The performance ordering of each baseline is maintained across datasets and backbone architectures, as presented in Figure 3. Overall, QuEE outperforms the baselines, especially for lower-cost regimes, with the exception of the SVHN dataset, which we will discuss later. JEIDNN closely follows, performing particularly well in higher-cost regimes, where it outperforms both QuEE and the initial accuracy of the backbone due to its ability to train the IMs jointly with the gating mechanism. DQNet is less competitive, as it is limited in the range of costs it can reach, since it cannot early exit. It generally outperforms simple quantization (PTQ4ViT), maintaining the same accuracy at a lower cost. We emphasize that both JEIDNN and DQNet require retraining, with adjustment of the trade-off hyperparameters, to obtain each point along the operating curve. This leads to instabilities that are apparent in the results, making these methods less reliable. In practice, this means that in order to obtain a specific target cost or performance, the models would need to be retrained many times until the desired values are attained. In contrast, QuEE only has to be trained once to obtain a full curve, because the cost can be modified at inference to obtain a different trade-off. This makes QuEE more efficient, stable, and flexible, and therefore more practical.

As we stated in our methodology, the strength of QuEE lies in its ability to combine quantization with early exit. For this to be an advantageous strategy, both computation reduction methods need to perform relatively well independently — there is no benefit in combining a weak method with a strong one. As a result, we expect QuEE to perform the best when both quantization and early exit are approximately equally effective. This is confirmed by our experimental results. Figure 3C) compares both computation reduction methods for the various experiments. When the methods are on par with each other, as is the case for CIFAR-10, ImageNet with T2T-ViT, and CIFAR-100, QuEE obtains its best performance. In contrast, for SVHN, the quantization methods are significantly outperformed by early exit, and the early exit method JEIDNN outperforms QuEE on this dataset. SVHN is a relatively simple dataset, and the T2T-ViT-7 architecture is much deeper than necessary for this task, making it particularly suitable for early exit. Conversely, for ImageNet with ViT pre-trained with Dino, we observe the inverse trend, and quantization methods outperform early exit. This indicates that the intermediate representations of ViT are not well suited to be used as input to perform inference, which could be a side effect of the self-supervised objective in DinoV2 [39]. This could also explain the instabilities of JEIDNN, because it jointly trains the inference heads with the gating mechanisms.