Boosted Dynamic Neural Networks

To save storage and computational resources, efficient neural networks have been widely studied in the existing works. To this end, a large number of efficient model architectures were designed in the past decade. These architectures include SqueezeNet (Iandola et al. 2016), MobileNets (Howard et al. 2017; Sandler et al. 2018), ShuffleNet (Zhang et al. 2018b), GhostNet (Han et al. 2020), and more. Another branch of efficient methods is network pruning, which removes selected network parameters or even some entire network channels to reduce model redundancy and increase inference speed (Han et al. 2015; Han, Mao, and Dally 2016; Li et al. 2016; He, Zhang, and Sun 2017; Luo, Wu, and Lin 2017; He et al. 2019; Lin et al. 2020). On the other hand, without designing or changing network architectures, the network quantization methods quantize network weights and activations from floating-point values to low-precision integers to allow fast model inference with integer arithmetic (Courbariaux et al. 2016; Rastegari et al. 2016; Zhou et al. 2016; Li, Zhang, and Liu 2016; Zhang et al. 2018a; Liu et al. 2018; Wang et al. 2019b; Yu et al. 2019; Liu et al. 2021).

As a special type of efficient neural networks, dynamic neural networks can adaptively change their inference complexity under different computational budgets, latency requirements, and prediction confidence requirements. Existing dynamic neural networks are designed in different aspects including sample-wise (Teerapittayanon, McDanel, and Kung 2016; Wang et al. 2018; Veit and Belongie 2018; Wu et al. 2018; Yu et al. 2019; Guo et al. 2019), spatial-wise (Li et al. 2017; Wang et al. 2019a, 2020; Yang et al. 2020; Wang et al. 2022a), and temporal-wise dynamism (Shen et al. 2017; Yu, Lee, and Le 2017; Wu et al. 2019; Meng et al. 2020; Wang et al. 2022b), as categorized by (Han et al. 2021). Specially, as one type of sample-wise methods, depth-wise dynamic models with early exits adaptively exit at different layer depths given different inputs (Huang et al. 2017; Li et al. 2019; McGill and Perona 2017; Jie et al. 2019; Yang et al. 2020). In (Huang et al. 2017) and (Li et al. 2019), the authors proposed multi-scale dense networks (MSDNet) to address two issues: a) interference of coarse-level and fine-level neural features, and b) feature scale inconsistency between the shallow and deep classifiers. In (McGill and Perona 2017) and (Jie et al. 2019), different inference exiting policies are proposed for better balance of accuracy and efficiency. In (Yang et al. 2020), resolution adaptive networks (RANet) take multi-scale images as input, based on the intuition that low-resolution features are suitable for high-efficiency inference while high-resolution features are suitable for high-accuracy inference. To train the depth-wise models, these methods jointly minimize the losses of all classifiers on every training batch without considering the dynamic inference property in testing stage. In this paper, we propose a new dynamic architecture BoostNet and several training techniques to mitigate this train-test mismatch problem. For more literature on dynamic neural networks, we refer the readers to (Han et al. 2021) for a comprehensive overview on the relevant papers.

As a pioneering work, AdaBoost was introduced in (Freund and Schapire 1997) to generate a high-accuracy model by combining multiple weak learners/models. The weak learners are optimized sequentially in a way that the subsequent learners pay more attention to the training samples misclassified by the previous classifiers. In (Friedman, Hastie, and Tibshirani 2000) and (Friedman 2001), AdaBoost was generalized to gradient boosting for different loss functions.

Later, gradient boosting methods were combined with neural networks. Following (Freund and Schapire 1997), Schwenk and Bengio applied AdaBoost to neural networks for character classification (Schwenk and Bengio 1998). In recent years, more investigations have been carried out to combine boosting and neural networks. In (Moghimi et al. 2016), BoostCNN was introduced to combine gradient boosting and neural networks for image classification. The model achieved outperforming performance on multiple image classification tasks. In (Cortes et al. 2017), the authors proposed AdaNet to learn neural network structure and parameters at the same time. The network was optimized over a generalization bound which is composed of a task loss and a regularization term. In (Huang et al. 2018), ResNet (He et al. 2016) was built sequentially in a gradient boosting way. The resulted BoostResNet required less computational and memory resources, and had exponentially decaying training errors as the layer depth increases. More recently, GrowNet was proposed to train neural networks with gradient boosting (Badirli et al. 2020), where shallow neural networks were regarded as “weak classifiers”. In addition, after a layer is trained, a corrective finetuning step is applied to all the previous layers to improve the model performance. Different from these existing works, our training framework is designed for dynamic network structure and adaptive inference ability. To this end, we propose several optimization techniques to train the model effectively.

An EDNN model usually consists of multiple stacked sub-models. In inference time, given an input image, we run the model until the last layer of the first classifier to get a coarse prediction. If the prediction result is confident above a predefined threshold, the inference can be terminated at just the first classifier. Otherwise, the second classifier will be executed on the features from the first classifier. It is expected to give a more confident prediction, correcting the mistake made by the first classifier. If not, we repeat the above process until either some classifier gives a confident enough prediction or the last classifier is reached.

To train the model, existing EDNN methods train all the classifiers jointly without considering the adaptive inference property described above (Huang et al. 2017; Li et al. 2019; Yang et al. 2020). Considering an $N$-classifier EDNN model, the overall task loss will be

where ${\bm{x}}$ and $y$ are an input example and its true categorical label, $F_{n}({\bm{x}})$ and $w_{n}$ are $n$-th classifier’s prediction and its contribution to the overall loss. This makes all the classifiers treat every training sample the same way. However, in inference time, the model will exit at early classifiers on easy inputs, so the deep classifiers will only be executed on difficult ones. This train-test mismatch problem will cause a data distribution gap between model training and testing stages, rendering a trained model sub-optimal.

Rethinking the inference process above, we find the intrinsic property of EDNN model is that a deeper classifier is supplementary to its previous classifiers by correcting their mistakes. This property shares a similar spirit with gradient boosting methods.

First we give a brief overview of the gradient boosting theory. Given an input ${\bm{x}}$ and its output $y$, gradient boosting is to find some function $F({\bm{x}})$ to approximate the output $y$. $F$ is in form of a sum of multiple base classifiers as

where $C$ is a constant. Gradient boosting optimizes each weak classifier sequentially:

where $\alpha$ is a scalar and $\mathcal{F}$ is some function class. Directly optimizing this problem to find the optimal $f_{n}$ is difficult given an arbitrary loss function $L_{n}$. Gradient boosting methods make a first-order approximation to Eq. 3. As a result, we have

where $\lambda$ is the learning rate.

Inspired by gradient boosting, we organize the classifiers of EDNN in a similar structure. The final output of the $n$-th block is a linear combination of the outputs of $n$-th block and all its previous blocks. Fig. 1 shows the proposed network structure. We name this architecture Boosted Dynamic Neural Networks (BoostNet).

Different from the first-order approximation in gradient boosting, we directly optimize the neural network $f_{n}$ w.r.t. parameters ${\bm{\theta}}_{n}$ with mini-batch gradient descent for multiple steps:

where $F_{n-1}({\bm{x}})$’s gradient is disabled even if it shares a subset of parameters with $f_{n}$ as shown in Fig. 1. In the ablation studies, we will compare this setting with the case of differentiable $F_{n-1}({\bm{x}})$, and find disabling $F_{n-1}({\bm{x}})$’s gradient works better.

Following traditional gradient boosting methods, one straight-forward optimization procedure is to train the $N$ classifiers sequentially, That is, train classifier $f_{n}$ until convergence, fix its classification head, and train $f_{n+1}$. However, this makes all classifiers shallower than $f_{n+1}$ deteriorate when training $f_{n+1}$ because all the classifiers share the same network backbone. On the other hand, fixing the shared parameters and only training $f_{n+1}$’s non-shared parameters will hurt its representation capacity.

Instead, we employ a mini-batch joint optimization scheme. Given a mini-batch of training data, we forward the entire BoostNet to get $N$ prediction outputs. Then we sum up the training losses over the $N$ outputs and take one back-propagation step. The overall loss is

We experimentally find $w_{n}=1$ works well.

Similar to (Li et al. 2019), we use a gradient rescaling operation to rescale the gradients from different network branches. As shown in Fig. 2, we multiply a scalar in each branch to rescale the gradients that pass through it. With this operation, the gradient of the overall loss $L$ w.r.t. block $b_{n}$ is

which is bounded even if $N$ is a large number, given the weak assumption that $\partial L_{i}/\partial b_{n}$ is bounded, which is usually true. The effectiveness of gradient rescaling will be demonstrated in the ablation studies.

When training in the above procedure, we observed a problem that the deeper classifiers do not show much performance improvement over the shallow classifiers, although the model size increases by a large margin.

We found this is caused by insufficient training data for deeper classifiers. In Fig. 3(a), we show the number of valid training samples for each classifier. A training sample is valid if its training loss is larger than a certain threshold $v$. Here, we set $v$ to be the 10-th percentile of all the training losses of the first classifier in each batch of training data, i.e., we assume the first classifier always has $90\%$ valid training samples. From the figure, we find the shallower classifiers have more valid training data while the deeper classifiers have relatively less. In Fig. 3(b), we show the batch-wise average training losses of different classifiers. Consistent with Fig. 3(a), the deeper classifiers have smaller training losses. From another perspective, this can be explained by the fact that the deeper classifiers have a larger model size and thus more representation power, but this also indicates that the deeper classifiers need more challenging training samples instead of learning the “easy” residuals of the shallower classifiers.

To address this problem, we multiply the output of $(n-1)$-th classifier by a temperature $t_{n}$ before serving as an ensemble member of $F_{n}$ as

where $F_{n-1}$ and $F_{n}$ are pre-softmax prediction logits. We set all the temperatures $\{t_{n}\}|_{n=1}^{N}$ to $0.5$ in our experiments unless specified in the ablation studies. We name this technique prediction reweighting. With this, the ensemble output of the shallow classifiers $F_{n-1}$ is equivalent to being weakened before being used in the ensemble output $F_{n}$.

With the proposed training techniques, the overall training pipeline is summarized in Algorithm 1.

To evaluate the proposed method, we apply it to two different network backbones, MSDNet (Huang et al. 2017) and RANet (Yang et al. 2020), both of which are dynamic models designed for adaptive inference. It is noted that our method is model-agnostic, and can also be applied to other network backbones. We do experiments on two datasets, CIFAR100 (Krizhevsky 2009) and ImageNet (Deng et al. 2009) , in two prediction modes, anytime prediction mode and budgeted-batch prediction mode following (Huang et al. 2017; Yang et al. 2020). In anytime prediction mode, the model runs at a fixed computational budget for all inputs. In budgeted-batch prediction mode, the model runs at an average budget, which means it can have different inference costs for different inputs as long as the overall average cost is below a threshold. To calculate the confidence threshold determining whether the model exits or not, we randomly reserve 5000 samples in CIFAR100 and 50000 samples in ImageNet as a validation set. We utilize a similar model combination strategy to (Huang et al. 2017; Yang et al. 2020). That is, we train a single dynamic model for anytime mode on both datasets. For budgeted-batch mode, we train multiple models with different model sizes for each dataset.

For MSDNet (Huang et al. 2017), we use the same model configurations as the original model except the CIFAR100-Batch models. In (Huang et al. 2017), the classifiers are located at depths $\{2\times i+4|i=1,...,k\}$, where $k$ is $6$, $8$, or $10$ for the three different models. One problem with this setting is that the three models have the same representation capacity for the shallow classifiers, which causes redundant overlap of computational complexity. Instead, we set the locations as $\{2\times i+k|i=1,...,6\}$, where $k$ is $4$, $8$, or $12$ for the three models. Thus, each model has $6$ classifiers at different depths. For RANet (Yang et al. 2020), we follow the same settings as the original paper. For prediction reweighting, we set all temperatures $t_{n}=0.5,\ n=1,...,N$.

For CIFAR100, the models are trained for $300$ epochs with SGD optimizer with momentum $0.9$ and initial learning rate $0.1$. The learning rate is decayed at epochs $150$ and $225$. We use single GPU with training batch size $64$. For ImageNet, the models are trained for $90$ epochs with SGD optimizer with momentum $0.9$ and initial learning rate $0.1$. The learning rate is decayed at epochs $30$ and $60$. We use $4$ GPUs with training batch size $64$ on each GPU. Our training framework is implemented in PyTorch (Paszke et al. 2019).

In this section, we compare our method with MSDNet and RANet in the anytime prediction mode. We also apply gradient rescaling to the baseline models. The results on CIFAR100 and ImageNet are plotted in Tab. 1 and 2. On both datasets, our model achieves better results under the same computational budget measured in Multi-Add. The improvement is even more significant under the low computational budgets on CIFAR100.

In budgeted-batch prediction mode, we set a computational complexity upper bound $\tau$ such that the average computational cost of the input images is below this value. We assume a proportion of $\frac{p}{1}$ samples exit at every classifier. Therefore, the average computational cost will be

where $N$ is the number of classifiers and $C_{n}$ is the computational complexity of layers between block $n-1$ and $n$. Given a budget threshold, we can utilize the Newton–Raphson method to find a $p$ satisfying the equation. Finally, we run the model on the hold-out evaluation set to determine the confidence threshold for each classifier given exiting probability $p$. We define the confidence of a classifier as its maximum softmax probability over all the categories.

We compare our method with MSDNet and RANet in Fig. 4.²²2Due to various reasons, e.g., different validation sets, RANet’s CIFAR100 result here may be a bit different from that in the paper. Under the same budget, our method gets superior results. On CIFAR100, the improvement over MSDNet is even more than $1\%$. We observe that the improvement decreases as the budget increases. We speculate the reason is when the model goes deeper, the effective amount of training data decreases. The proposed prediction reweighting technique mitigates this problem, but it is not entirely eliminated.

Note that the curve tails are flat in Fig. 4(a) and Fig. 4(b), i.e., a deeper classifier does not improve the performance over a shallower one. One possible reason of such a phenomenon of diminishing returns is overfitting. To avoid performance drop when computational budget increases, we do not increase confidence threshold if this change does not bring performance improvement. This process is conducted on the hold-out evaluation dataset. A similar strategy is also employed in (Yang et al. 2020).

We show visual examples that exit at various depths of the dynamic model. Fig. 5(a) and 5(b) are schipperke images that exit at classifier $1$ and $5$ respectively. Fig. 5(c) and 5(d) are titi examples exiting at classifier $1$ and $5$ respectively. We find that image instances exiting at classifier $1$ are usually “easier” than those exiting at classifier $5$. Actually, this phenomenon can be observed in all the categories in ImageNet. Generally, images with occlusion, small object size, unusual view angle, and complex background are difficult to recognize and require a deeper classifier.

In this section, we do ablation experiments using MSDNet backbone to understand the training techniques.

Prediction reweighting. In the previous section, we introduced prediction reweighting to make sure more valid training data for deeper classifiers. In this part, we study how prediction reweighting temperature $t$ affects the performance. The results are plotted in Fig. 6(a). We find that the deeper classifiers cannot learn well without enough supervision when $t$ is $1$, although the shallower classifiers achieve better performance than the other $t$ values. On the other hand, when $t$ is $0$, it basically becomes an MSDNet, which also performs inferior to the case of $t=0.5$. As a trade-off, we set $t=0.5$ in all our experiments.

Batch size. As introduced earlier, we formulated the early-exiting dynamic-depth neural network as an additive model in a way that the prediction of a deep classifier $F_{n}$ depends on all its preceding classifiers $\{F_{i}\}_{i=1}^{n-1}$. At training step $k$, $F_{n}^{k}$ is computed using $\{F_{i}^{k-1}\}_{i=1}^{n-1}$ before the new $\{F_{i}^{k}\}_{i=1}^{n-1}$ are computed. Intuitively thinking, if these shallower classifiers fluctuate too much during training, the deeper classifier cannot be well learned. Therefore, a relatively large batch size is desired for better training stability. Meanwhile, it cannot be too large to hurt the model’s generalization ability (Keskar et al. 2016). Therefore, batch size plays an important role in dynamic network training. In Fig. 6(b), we show how batch size affects the classification performance. Both too large batch sizes and too small batch sizes get worse performance. Therefore, we set the batch size to be $64$ for CIFAR100 models and $256$ for ImageNet models as a trade-off between stability and generalization.

Gradient rescaling. In Fig. 6(c), we compare the model with gradient rescaling (red curve) and that without gradient rescaling (green curve). When the budget is over $0.8\times 10^{9}$, the model with gradient rescaling gets better performance than that without it. This shows gradient rescaling can help improve the performance of the deeper classifiers.

Trainable $F_{n-1}$. As shown in Algorithm 1, we do not back-propagate gradients of $l_{n}$ along $F_{n-1}$ branch. We do this to avoid entangling deeper losses in shallower classifiers. For comparison purpose, we train a model without cutting out the gradients of $l_{n}$ w.r.t. $F_{n-1}$. The results are shown as the blue curve in Fig. 6(c). Entangling deeper losses in shallower classifiers hurts classification performance.