SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Don Kurian Dennis
Anish Sevekari
Kazuhito Koishida
We study the problem of progressive ensemble distillation: Given a large, pretrained teacher model  $g$ , we seek to decompose the model into smaller, low-inference cost student models  $f_{i}$ , such that progressively evaluating additional models in this ensemble leads to improved predictions. The resulting ensemble allows for flexibly tuning accuracy vs. inference cost at runtime, which is useful for a number of applications in on-device inference. The method we propose, B-DISTIL, relies on an algorithmic procedure that uses function composition over intermediate activations to construct expressive ensembles with similar performance as  $g$ , but with smaller student models. We demonstrate the effectiveness of B-DISTIL by decomposing pretrained models across standard image, speech, and sensor datasets. We also provide theoretical guarantees in terms of convergence and generalization.
Knowledge distillation aims to transfer the knowledge of a large model into a smaller one [5, 23]. While this technique is commonly used for model compression, one downside is that the procedure is fairly rigid—resulting in a single compressed model of a fixed size. In this work, we instead consider the problem of progressive ensemble distillation: approximating a large model via an ensemble of smaller, low-latency models such that such that progressively evaluating additional models in this ensemble leads to improved predictions. The resulting decomposition is useful for many applications in on-device and low-latency inference. For example, components of the ensemble can be selectively combined to flexibly meet accuracy/latency constraints [31, 44], can enable efficient parallel inference execution schemes, and can facilitate early-exit [4, 11] or anytime inference [36, 28] applications, which are scenarios where inference may be interrupted due to variable resource availability.
More specifically, our work seeks to distill a large pretrained model,  $g$ , onto an ensemble of 'smaller' models, such that evaluating the first model produces a coarse estimate of the prediction (e.g., covering common cases), and evaluating additional models improves on this estimate (see Figure 1). There are major advantages to such an ensemble for on-device efficient inference.
Figure 1: In progressive ensemble distillation, a large teacher model is distilled into an ensemble of low inference cost models. The more student models we evaluate, the closer the ensemble's decision boundary is to that of the teacher model. Models in the ensemble are allowed to depend on previously computed features to reduce overhead and inference cost.
Corresponding author: Don Dennis <dondennis@cmu.edu>.
Concretely, (i) inference cost vs. accuracy trade-offs can be controlled on-demand at execution time, (ii) the ensemble can either be executed in parallel or in sequence, or possibly a mix of both, and (iii) we can improve upon coarse initial predictions without re-evaluation at runtime.
While traditional distillation methods are effective when transferring information to a single model of similar capacity, it has been shown that performance can degrade significantly when reducing the capacity of the student model [34, 19]. Moreover, distillation of a deep network onto a weighted sum of shallow networks rarely performs better than distillation onto a single model [9, 1].
Our insight in this work is that by composing and reusing activations and by explicitly incentivizing models to be weak learners during distillation, we can successfully find weak learners even when the capacity gap is relatively large. As long as these composition functions are resource-efficient, we are able to increase our hypothesis class capacity at roughly the same inference cost as a single model. Moreover, we show that our procedure retains the theoretical guarantees of classical boosting methods [39]. Concretely, we make the following contributions:
- We formulate progressive ensemble distillation as a two player zero-sum game, derive a weak learning condition for distillation, and present our algorithm, B-DISTIL, to approximately solve this game. To make the search for weak learners in low parameter count models feasible, we solve a log-barrier based relaxation of our weak learning condition. By allowing models to reuse computation from select intermediate layers of previously evaluated models of the ensemble, we can increase the model's capacity without significantly increasing inference cost.
- We empirically evaluate our algorithm on synthetic and real-world classification tasks from computer vision, speech, and sensor processing with models suitable for the respective domains. We show that our ensemble behaves like a decomposition, allowing a run-time trade-off between accuracy and computation, while retaining competitive performance with the teacher model.
- We provide theoretical guarantees for our algorithm in terms of in-sample convergence and generalization performance. Our framework is not architecture or task specific and can recover existing ensemble models used in efficient inference; we believe our work thus puts forth a general lens to view previous work and also to develop new, principled approaches for efficient inference.
Efficient Inference. Machine learning inference is often resource-constrained when deployed in practical applications due to memory, energy, cost, or latency constraints. This has spurred the development of numerous techniques for efficient inference. Pruning and approximations of pretrained parameter tensors through low-rank, sparse and quantized representations [22, 3, 18, 20] have been effective in reducing resource requirements. There are also architecture and task specific techniques for efficient inference [12, 45, 13]. In contrast to compressing an already trained model, algorithms have also been developed to train compressed models by incorporating resource constraints as part of their training routines [2, 7]. More recently, algorithms have been developed to search and find smaller sub-models from a single pre-trained model [46, 6].
Knowledge distillation. Knowledge distillation aims to transfer the knowledge of a larger model (or model ensemble) to a smaller one [5, 23]. Despite its popularity, performing compression via distillation has several known pitfalls. Most notably, it is well-documented that distillation performs poorly when there is a capacity gap, i.e., the teacher is significantly larger than the student [34, 19, 9, 1]. When performing distillation onto a weighted combination of ensembles, it has been observed that adding additional models into the ensemble does not dramatically improve performance over that of a single distilled model [1]. There is also a lack of theoretical work characterizing when and why distillation is effective for compression [21]. Our work aims to address these pitfalls by developing a principled approach for progressively distilling a large model onto an ensemble of smaller, low-capacity ones. We defer readers to [21] for a recent survey on varied applications of and approaches for distillation at large.
Early exits and anytime inference. Many applications stand to benefit from the output of progressive ensemble distillation, which allows for flexibly tuning accuracy vs. inference cost and executing inference in parallel. Enabling trade-offs between accuracy and inference cost is particularly useful for applications that use early exit or anytime inference schemes. In on-device continuous (online) inference settings, early exit models aim to evaluate common cases quickly in order to improve
energy efficiency and prolong battery life [11, 4]. For instance, a battery powered device continuously listening for voice commands can use early exit methods to improve battery efficiency by quickly classifying non-command speech. Many early exit methods are also applicable to anytime inference [28, 36]. In anytime inference, the aim is to produce a prediction even when inference is interrupted, e.g., due to resource contention or a scheduler decision. Unlike early exit methods where the classifier chooses when to exit, anytime inference methods have no control over when they are interrupted. We explore the effectiveness of our method, B-DISTIL, for such applications in Section 5.
Two-player games, online optimization and boosting. In this work we formulate progressive ensemble distillation as a two player zero-sum game. The importance of equilibrium of two player zero-sum games have been recognized since the foundational work of von Neumann and Morgenstern [42]. Later applications by Schapire [38] and Freund [15], Freund and Schapire [16] identified close connections between boosting and two-player games. On the basis of this result, a number of practical boosting-based learning approaches such as AdaBoost [17], gradient boosting [33], and XGboost [8] have been developed. Boosting has only recently seen success in modern deep learning applications. In particular, Suggala et al. [40] propose a generalized boosting framework to train boosting based ensembles of deep networks. Their key insight is that allowing function compositions in feature space can help boost deep neural networks. Although they focus on training and do not produce decompositions of pretrained models, we use a similar approach in our work to select intermediate layers connections between ensemble components (Section 3.3). A more general application of boosting that is similar to our setup is by Trevisan et al. [41]. They prove that given a target bounded function  $g$  (e.g., the teacher model) and class of candidate approximating functions  $f \in \mathcal{F}$ , one can iteratively approximate  $g$  arbitrarily well with respect to  $\mathcal{F}$  using ideas from online learning and boosting. However, these results depend on the ability to find a function  $f_{t}$  in iteration  $t$  that leads to at least a small constant improvement in a round-dependent approximation loss. A key contribution of our work is showing that such functions can be found for the practical application of progressive ensemble distillation by carefully selecting candidate models.
3 Progressive Ensemble Distillation with B-DISTIL
As discussed, our goal in performing progressive ensemble distillation is to approximate a large model via an ensemble of smaller, low-latency models so that we can easily trade-off between accuracy and inference-time/latency at runtime. In this section we formalize the problem of progressive ensemble distillation as a two player game and discuss our proposed algorithm, B-DISTIL.
Consider repeated plays of a general two player zero-sum game with the pure strategy sets comprising of a hypothesis class  $\mathcal{F}$  and a probability distribution  $\mathcal{P}$ . Given a loss function  $\mathcal{L}$ , we let the loss (and reward) of the players be given by  $F(f,p) = \mathbb{E}_{x\sim p}[\mathcal{L}(f,x)]$ , and the minimax value of the game is:
$$
\max  _ {p \in \mathcal {P}} \min  _ {f \in \mathcal {F}} F (f, p). \tag {1}
$$
In the context of distillation, given a training set  $\{x_{i}\}$ , we can think of the role of the max player in Equation (1) as producing distributions over the training set and the role of the min player as producing a hypothesis that minimizes the loss on this distribution. In this setting, note that  $\mathcal{P} = \left\{p\in \mathbb{R}^{N\times L}:p_{i,j}\geq 0\quad \forall j\sum_{i}p_{i,j} = 1\right\}$  is the product of simplices in  $N\times L$  dimensions, and
$$
\left(\nabla_ {f} F (f, p)\right) _ {j} = \sum_ {i} p _ {i, j} \left(f \left(x _ {i}\right) - g \left(x _ {i}\right)\right) _ {j}. \tag {2}
$$
Our goal is to produce an ensemble of predictors from the set of hypothesis classes  $\{\mathcal{F}_m\}$  to approximate  $g$  'well'. We now appeal to tools from the framework of stochastic minimax optimization to approximately attain the value in Equation (1) (see Appendix A for a more involved discussion). As is common in this setup, we assume our algorithm is provided access weak gradient vector  $h$  such that, when queried at distribution  $p\in \mathcal{P}$  and for  $\beta >0$
$$
\langle h, \nabla_ {f} F (f, p) \rangle \geq \beta . \tag {3}
$$
We perform this construction iteratively by searching for weak learners or weak gradients in the sense of Equation (3) with respect to the target  $g$  in the class  $\mathcal{F}_m$ . Conditioned on a successful search we can guarantee in-sample convergence to the minimax value in Equation (1) (Theorem 1) and bound the excess risk of the ensemble (Theorem 2). Although Equation (3) is a seemingly easier notion than the full optimization, in many problems of interest even this is challenging. In fact, in the multilabel setting that we focus on, one of the main algorithmic challenges is to construct an algorithm that can reliably find low cost weak gradients/learners (see Section 3.3).
Algorithm 1 B-DISTIL: Main algorithm
Require: Target  $g$  ,rounds  $T$  data  $\{(x_i,y_i)\}_{i = 1}^N$  learning rate  $\eta$  , model classes  $\{\mathcal{F}_m\}_{m = 1}^M$    
1:  $K_{t}^{+}(i,j),K_{t}^{-}(i,j)\gets \frac{1}{2N},\frac{1}{2N}\quad \forall (i,j)$    
2:  $F,r,t\gets \emptyset ,1,1$    
3: while  $r <   R$  and  $t <   T$  do   
4:  $f_{t} = \mathrm{FIND - WL}(K_{t}^{+},K_{t}^{-},\mathcal{F}_{r})$    
5: if  $f_{t}$  is NONE then   
6:  $r\gets r + 1$    
7: continue   
8: end if   
9: With  $l\coloneqq f_t - g$  ,update  $K_{t}^{+},K_{t}^{-}.\forall (i,j)$ $K_{t + 1}^{+}(i,j)\gets K_{t}^{+}(i,j)\exp (-\eta \cdot l(x_{i})_{j})$ $K_{t + 1}^{-}(i,j)\gets K_{t}^{-}(i,j)\exp (\eta \cdot l(x_{i})_{j})$    
10:Normalize  $K_{t}^{+},K_{t}^{-}$    
11:  $F,t\gets F\cup \{f_t\} ,t + 1$    
12:end while   
13: Return  $\frac{1}{|F|}\sum_{i = 1}^{|F|}f_i$
Algorithm 2 FIND-WL
Require: Probability matrices  $K^{+}, K^{-}$ , model class  $\mathcal{F}$  parameterized by  $\theta \in \Theta$ , hyperparameters for SGD  
1: Obtain  $\{\mathcal{F}_r\}_1^R$  by expanding  $\mathcal{F}$  (Section 3.2).  
2: for  $\mathcal{F}' \in \{\mathcal{F}_r\}_{r=1}^R$  do  
3: Initialize initial parameter  $\theta_0 \in \mathcal{F}'$ .  
4: for  $i \in \{1, \dots, \text{max-search}\}$  do  
5: Randomly initialize  $f_{\theta_i}$ .  
6: Run SGD to solve Equation (6).  
7: if  $f_{\theta_i}$  is a weak learner then  
8: Return  $f_{\theta_i}$   
9: end if  
10: end for  
11: end for  
12: Return NONE
3.2 B-DISTIL Algorithm
Concretely, at each round  $t$ , our proposed algorithm, B-DISTIL, maintains matrices  $K_{t}^{+} \in R^{N \times L}$  and  $K_{t}^{-} \in R^{N \times L}$  of probabilities (in our setting, it turns out to be easier to maintain the positive errors and the negative error separately). Note that the matrices  $K_{t}^{+}$  and  $K_{t}^{-}$  are such that for all  $j \in [L]$ ,  $\sum_{i} K_{t}^{+}(i,j) + K_{t}^{-}(i,j) = 1$ . Moreover, for all  $(i,j) \in [N] \times [L]$ ,  $0 \leq K_{t}^{+}(i,j), K_{t}^{-}(i,j) \leq 1$ . The elements  $K_{t}^{+}(i,j)$  and  $K_{t}^{-}(i,j)$  can be thought of as the weight on the residual errors  $f_{t-1}(x) - g(x)$  and  $g(x) - f_{t-1}(x)$  respectively, up-weighting large deviations from the teacher model  $g(x)$ . We formalize our notion of weak learners in this setting using Definition 1, which can be seen as a natural extension of the standard weak learning assumption in the boosting literature.
Definition 1 (Weak learning condition). Given a dataset  $\{(x_i, y_i)\}_{i=1}^N$ , a target function  $g: \mathcal{X} \to \mathbb{R}^L$  and probability matrices  $K_t^+, K_t^-$ , a function  $f_t: \mathcal{X} \to \mathbb{R}^L$  is said to satisfy the weak learning condition with respect to  $g, \forall j$ , if the following sum is strictly positive:
$$
\sum_ {i} K _ {t} ^ {+} (i, j) \left(f _ {t} \left(x _ {i}\right) - g \left(x _ {i}\right)\right) _ {j} + K _ {t} ^ {-} (i, j) \left(g \left(x _ {i}\right) - f _ {t} \left(x _ {i}\right)\right) _ {j}.
$$
At each round  $t$ , with the current probability matrices  $K_{t}^{-}, K_{t}^{+}$ , B-DISTIL performs two steps; first, it invokes a subroutine FIND-WL (discussed below) that attempts to find a classifier  $f_{t} \in \mathcal{F}_{r}$  satisfying the weak learning condition (Definition 1). If such a predictor is found, we add it to our ensemble and proceed to the second step, updating the probability matrices  $K_{t}^{-}, K_{t}^{+}$  based on errors made by  $f_{t}$ . This is similar in spirit to boosting algorithms such as AdaBoost [39] for binary classification. If no such predictor can be found, we invoke the subroutine with the next class,  $\mathcal{F}_{r + 1}$ , and repeat the search till a weak learner is found or we have no more classes to search in.
3.3 Finding Weak Learners
As mentioned above, the main difficulty in provably approximating the teacher model in this setting is in finding a single learner  $f_{t}$  at round  $t$  that satisfies our weak learning condition simultaneously for all labels  $j$ . Existing boosting methods for classification treat multi-class settings ( $L > 1$ ) as  $L$  instances of the binary classification problem (one vs. all) [39]. They typically choose  $L$  different weak learners for each instance, which is unsuitable for resource efficient on-device inference. The difficulty is further increased by the capacity gap between the student and teacher models we consider for distillation. Thus, along with controlling temperature for distillation, we employ two additional strategies: 1) we use a regularizer in the objective FIND-WL solves to promote weak learning and, 2) we efficiently reuse a limited number of stored activation outputs of previously evaluated models to increase the expressivity of the current base class.
Log-barrier regularizer. To find a weak learner, the FIND-WL method minimizes the sum of two loss terms using stochastic gradient descent. The first is standard binary/multi-class cross-entropy distillation loss [23], with temperature smoothing. The second term is defined in Equation (6):
$$
- \frac {1}{\gamma} \sum_ {i, j} I _ {i j} ^ {+} \log \left(1 + \frac {l \left(x _ {i}\right) _ {j}}{2 B}\right) + \left(1 - I _ {i j} ^ {+}\right) \log \left(1 - \frac {l \left(x _ {i}\right) _ {j}}{2 B}\right) \tag {6}
$$
Here  $I_{ij}^{+} \coloneqq I[K_{t}^{+}(i,j) > K_{t}^{-}(i,j)]$ ,  $B$  is an upper bound on the magnitude of the logits, and  $l(x_{i}) \coloneqq f(x_{i}) - g(x_{i})$ . To see the intuition behind Equation (6), assume the following holds;  $\forall (i,j)$
$$
\left. \left(K _ {t} ^ {+} (i, j) - K _ {t} ^ {-} (i, j)\right) \left(f \left(x _ {i}\right) - g \left(x _ {i}\right)\right) _ {j} > 0. \right. \tag {7}
$$
Summing over all  $x_{i}$ , we can see that this is sufficient for  $f$  to be a weak learner with respect to  $g$ . Equation (6) is a soft log-barrier version of the weak learning condition, that penalizes those  $(i,j)$  for which Equation (7) does not hold. By tuning  $\gamma$  we can increase the relative importance of the regularization objective, encouraging  $f_{t}$  to be a weak learner potentially at the expense of classification performance.
Intermediate layer connections and profiling. As discussed in Section 2, distillation onto a linear combination of low capacity student models often offers no better performance than that of any single model in the ensemble trained independently. For boosting, empirically we see that once the first weak learner has been found in some class  $\mathcal{F}_m$  of low-capacity deep networks, it is difficult to find a weak learner for the reweighed objective from the same class  $\mathcal{F}_m$ . To work around this we let our class of weak learners at round  $t$  include functions that depend on the output of intermediate layers of previous weak learners [40].
As a concrete example, consider a deep fully connected network with  $U$  layers, parameterized as  $f = W\phi_{1:U}$ . Here  $\phi_{1:u}$  can be thought of as a feature transform on  $x$  using the first  $u$  layers into  $\mathbb{R}^{d_u}$  and  $W\in R^{d_U\times L}$  is a linear transform. With two layer fully connected base model class  $\mathcal{F}_0\coloneqq \{W^{(0)}\phi_{1:2}^{(0)}\mid W^{(0)}\in \mathbb{R}^{L\times d_2}\}$  (dropping subscript  $m$  for simplicity), we define:
$$
\mathcal {F} _ {r} = \left\{W ^ {(r)} \phi_ {1: 2} ^ {(r)} (\mathrm {i d} + \phi_ {1: 2} ^ {(r - 1)}) \right\} \quad \text {a n d ,} \quad \mathcal {F} _ {r} ^ {\prime} \left\{W ^ {(r)} \phi_ {2} ^ {(r)} (\mathrm {i d} + \phi_ {1} ^ {(r)} + \phi_ {1} ^ {(r - 1)}) \right\},
$$
with  $\operatorname{id}(x) := x$ . It can be seen that  $\{\mathcal{F}_r\}$  and  $\{\mathcal{F}_r'\}$  define a variant of residual connection based networks [27]. It can be shown that classes of function  $\{\mathcal{F}_r\}$  (and  $\{\mathcal{F}_r'\}$ ) increase in capacity with  $r$ . Moreover, when evaluating sequentially the inference cost of a model from  $\mathcal{F}_r$  is roughly equal to that of  $\mathcal{F}$ , since each subsequent evaluation reuses stored activations from previous evaluations. For this reason the parameter count of each  $\mathcal{F}_r$  remains the same as that of the base class. Note that by picking the base class as dense networks at various scales and connections as dense connections our algorithm can recover MSDNets studied in [27]. Similarly, by picking the base class as root nodes, and connections as binary connections, we recover an HNE from [36].
We informally refer to the process of constructing  $\{\mathcal{F}_r\}$  given a choice of base class  $\mathcal{F}_0$ , the parameter  $R$  and the connection type as expanding  $\mathcal{F}_0$ . Note that while intermediate connections help with capacity, they often reduce parallelizability as models become mutually dependent. As a practical consequence dependencies on activation outputs of later layers are preferred, and we use the Rasley et al. [35] profiler to measure inference cost during training rounds and rank models (see Appendix C).
In this section we provide theoretical analysis and justification for our method. First, we show that the ensemble output produced by algorithm 1 converges to  $g$  at  $\mathcal{O}(1 / \sqrt{T})$  rate, provided that the procedure FIND-WL succeeds at every time  $t$ .
Theorem 1. Suppose the class  $\mathcal{F}$  satisfies that for all  $f\in \mathcal{F}$ ,  $\| f - g\|_{\infty}\leq G_{\infty}$ . Let  $F = \{f_t\}$  be the ensemble after  $T$  rounds of Algorithm 1, with the final output  $F_{t} = \frac{1}{T}\sum_{t = 1}^{T}f_{t}$ . If  $f_{t}$  satisfies eq. (7) for all  $t\le T$  then for  $T\geq \ln 2N$  and  $\eta = \frac{1}{G_{\infty}}\sqrt{\frac{\ln 2N}{T}}$ , we have for all  $j$
$$
\left\| F _ {t, j} - g _ {j} \right\| _ {\infty} \leq G _ {\infty} \sqrt {\frac {\ln 2 N}{T}}, \tag {8}
$$
where  $F_{t,j}$  and  $g_{j}$  are the  $j^{th}$  coordinates of the functions  $F_{t}$  and  $g$  respectively.
We defer the details of the proof to the Appendix B. The main idea behind the proof is to bound the rate of convergence of the algorithm towards the minimax solution. This proceeds by maintaining a potential function and keeping track of its progress through the algorithm. The bounds and techniques here are general in the sense that for various objectives and loss functions appropriate designed weak learners give similar rates of convergence to the minimax solution. Furthermore, a stronger version that shows exponential rates can be shown by additionally assuming an edge for the weak learner.
In addition to the claim above about the in-sample risk, we also show that the algorithm has a strong out-of-sample guarantee. We show this by bounding the generalization error of the algorithm in terms of the generalization error of the class  $\mathcal{F}$ . In the following theorem, we restrict to the case of binary classification for simplicity, but the general result follows along similar lines. Let  $\mathcal{C}_T$  denote the class of functions of the form
$$
F _ {T} (x) = \mathrm {s i g n} (\frac {1}{T} \sum_ {i = 1} ^ {T} f _ {t} (x)),
$$
where  $f_{t}$  are functions in class  $\mathcal{F}$ . We then have the following generalization guarantee: