SlowGuess/ACL_2020_OCR huggingface dataset:
Luca Soldaini
Manhattan Beach, CA, USA
lssoldai@amazon.com
Alessandro Moschitti
amosch@amazon.com
Large transformer-based language models have been shown to be very effective in many classification tasks. However, their computational complexity prevents their use in applications requiring the classification of a large set of candidates. While previous works have investigated approaches to reduce model size, relatively little attention has been paid to techniques to improve batch throughput during inference. In this paper, we introduce the Cascade Transformer, a simple yet effective technique to adapt transformer-based models into a cascade of rankers. Each ranker is used to prune a subset of candidates in a batch, thus dramatically increasing throughput at inference time. Partial encodings from the transformer model are shared among rerankers, providing further speed-up. When compared to a state-of-the-art transformer model, our approach reduces computation by  $37\%$  with almost no impact on accuracy, as measured on two English Question Answering datasets.
Recent research has shown that transformer-based neural networks can greatly advance the state of the art over many natural language processing tasks. Efforts such as BERT (Devlin et al., 2019), RoBERTa (Liu et al., 2019c), XLNet (Dai et al., 2019), and others have led to major advancements in several NLP subfields. These models are able to approximate syntactic and semantic relations between words and their compounds by pre-training on copious amounts of unlabeled data (Clark et al., 2019; Jawahar et al., 2019). Then, they can easily be applied to different tasks by just fine-tuning them on training data from the target domain/task (Liu et al., 2019a; Peters et al., 2019). The impressive effectiveness of transformer-based neural networks can be partially attributed to their large number of parameters (ranging from 110 million
for "base" models to over 8 billion (Shoeybi et al., 2019); however, this also makes them rather expensive in terms of computation time and resources. Being aware of this problem, the research community has been developing techniques to prune unnecessary network parameters (Lan et al., 2019; Sanh et al., 2019) or optimize the transformer architecture (Zhang et al., 2018; Xiao et al., 2019).
In this paper, we propose a completely different approach for increasing the efficiency of transformer models, which is orthogonal to previous work, and thus can be applied in addition to any of the methods described above. Its main idea is that a large class of NLP problems requires choosing one correct candidate among many. For some applications, this often entails running the model over hundreds or thousands of instances. However, it is well-known that, in many cases, some candidates can be more easily excluded from the optimal solution (Land and Doig, 1960), i.e., they may require less computation. In the case of hierarchical transformer models, this property can be exploited by using a subset of model layers to score a significant portion of candidates, i.e., those that can be more easily excluded from search. Additionally, the hierarchical structure of transformer models intuitively enables the re-use of the computation of lower blocks to feed the upper blocks.
Following the intuition above, this work aims at studying how transformer models can be cascaded to efficiently find the max scoring elements among a large set of candidates. More specifically, the contributions of this paper are:
First, we build a sequence of rerankers  $SR_{N} = \{R_{1}, R_{2}, \dots, R_{N}\}$  of different complexity, which process the candidates in a pipeline. Each reranker at position  $i$  takes the set of candidates selected by  $(i - 1)$ -th reranker and provides top  $k_{i}$  candidates to the reranker of position  $i + 1$ . By requiring that  $k_{i} < k_{i - 1} \forall i = 1, \dots, N - 1$ , this approach
5697
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 5697-5708 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
allows us to save computation time from the more expensive rerankers by progressively reducing the number of candidates at each step. We build  $R_{i}$  using transformer networks of 4, 6, 8, 10, and 12 blocks from RoBERTa pre-trained models.
Second, we introduce a further optimization on  $SR_{N}$  to increase its efficiency based on the observation that models  $R_{i}$  in  $SR_{N}$  process their input independently. In contrast, we propose the Cascade Transformer (CT), a sequence of rerankers built on top of a single transformer model. Rerankers  $R_{1},\ldots ,R_{N}$  are obtained by adding small feedforward classification networks at different transformer block positions; therefore, the partial encodings of the transformer blocks are used as both input to reranker  $R_{i}$ , as well as to subsequent transformer encoding blocks. This allows us to efficiently re-use partial results consumed by  $R_{i}$  for rankers  $R_{i + 1},\dots ,R_{N}$ .
To enable this approach, the parameters of all rerankers must be compatible. Thus, we trained CT in a multi-task learning fashion, alternating the optimization for different  $i$ , i.e., the layers of  $R_{i}$  are affected by the back-propagation of its loss as well as by the loss of  $R_{j}$ , with  $j \leq i$ .
Finally, as a test case for CT, we target Answer Sentence Selection (AS2), a well-known task in the domain of Question Answering (QA). Given a question and a set of sentence candidates (e.g., retrieved by a search engine), this task consists in selecting sentences that correctly answer the question. We tested our approach on two different datasets: (i) ASNQ, recently made available by Garg et al. (2020); and (ii) a benchmark dataset built from a set of anonymized questions asked to Amazon Alexa. Our code, ASNQ split, and models trained on ASNQ are publicly available.
Our experimental results show that: (i) The selection of different  $k_{i}$  for  $SR_{N}$  determines different trade-off points between efficiency and accuracy. For example, it is possible to reduce the overall computation by  $10\%$  with just  $1.9\%$  decrease in accuracy. (ii) Most importantly, the CT approach largely improves over SR, reducing the cost by  $37\%$  with almost no loss in accuracy. (iii) The rerankers trained through our cascade approach achieve equivalent or better performance than transformer models trained independently. Finally, (iv) our results suggest that CT can be used with other
NLP tasks that require candidate ranking, e.g., parsing, summarization, and many other structured prediction tasks.
In this section, we first summarize related work for sequential reranking of passages and documents, then we focus on the latest methods for AS2, and finally, we discuss the latest techniques for reducing transformer complexity.
Reranking in QA and IR The approach introduced in this paper is inspired by our previous work (Matsubara et al., 2020); there, we used a fast AS2 neural model to select a subset of instances to be input to a transformer model. This reduced the computation time of the latter up to four times, preserving most accuracy.
Before our paper, the main work on sequential rankers originated from document retrieval research. For example, Wang et al. (2011) formulated and developed a cascade ranking model that improved both top-k ranked effectiveness and retrieval efficiency. Dang et al. (2013) proposed two stage approaches using a limited set of textual features and a final model trained using a larger set of query- and document-dependent features. Wang et al. (2016) focused on quickly identifying a set of good candidate documents that should be passed to the second and further cascades. Gallagher et al. (2019) presented a new general framework for learning an end-to-end cascade of rankers using back-propagation. Asadi and Lin (2013) studied effectiveness/efficiency trade-offs with three candidate generation approaches. While these methods are aligned with our approach, they target document retrieval, which is a very different setting. Further, they only used linear models or simple neural models. Agarwal et al. (2012) focused on AS2, but just applied linear models.
Answer Sentence Selection (AS2) In the last few years, several approaches have been proposed for AS2. For example, Severyn and Moschitti (2015) applied CNN to create question and answer representations, while others proposed interweighted alignment networks (Shen et al., 2017; Tran et al., 2018; Tay et al., 2018). The use of compare and aggregate architectures has also been extensively evaluated (Wang and Jiang, 2016; Bian et al., 2017; Yoon et al., 2019). This family of approaches uses a shallow attention mechanism
1https://github.com/alexa/wqa-cascade-transformers
5698
over the question and answer sentence embeddings. Finally, Tayyar Madabushi et al. (2018) exploited fine-grained question classification to further improve answer selection.
Transformer models have been fine-tuned on several tasks that are closely related to AS2. For example, they were used for machine reading (Devlin et al., 2019; Yang et al., 2019a; Wang et al., 2019), ad-hoc document retrieval (Yang et al., 2019b; MacAvaney et al., 2019), and semantic understanding (Liu et al., 2019b) tasks to obtain significant improvement over previous neural methods. Recently, Garg et al. (2020) applied transformer models, obtaining an impressive boost of the state of the art for AS2 tasks.
Reducing Transformer Complexity The high computational cost of transformer models prevents their use in many real-word applications. Some proposed solutions rely on leveraging knowledge distillation in the pre-training step, e.g., (Sanh et al., 2019), or used parameter reduction techniques (Lan et al., 2019) to reduce inference cost. However, the effectiveness of these approaches varies depending on the target task they have been applied to. Others have investigated methods to reduce inference latency by modifying how self-attention operates, either during encoding (Child et al., 2019; Guo et al., 2019b), or decoding (Xiao et al., 2019; Zhang et al., 2018). Overall, all these solutions are mostly orthogonal to our approach, as they change the architecture of transformer cells rather than efficiently re-using intermediate results.
With respect to the model architecture, our approach is similar to probing models $^2$  (Adi et al., 2017; Liu et al., 2019a; Hupkes et al., 2018; Belinkov et al., 2017), as we train classification layers based on partial encoding on the input sequence. However,  $(i)$  our intermediate classifiers are integral part of the model, rather than being trained on frozen partial encodings, and  $(ii)$  we use these classifiers not to inspect model properties, but rather to improve inference throughput.
Our approach also shares some similarities with student-teacher (ST) approaches for self-training (Yarowsky, 1995; McClosky et al., 2006). Under this setting, a model is used both as a "teacher" (which makes predictions on unlabeled data to obtain automatic labels) and as a "student" (which learns both from gold standard and automatic labels). In recent years, many variants of ST have
been proposed, including treating teacher predictions as soft labels (Hinton et al., 2015), masking part of the label (Clark et al., 2018), or use multiple modules for the teacher (Zhou and Li, 2005; Ruder and Plank, 2018). Unlike classic ST approaches, we do not aim at improving the teacher models or creating efficient students; instead, we trained models to be used as sequential ranking components. This may be seen as a generalization of the ST approach, where the student needs to learn a simpler task than the teacher. However, our approach is significantly different from the traditional ST setting, which our preliminary investigation showed to be not very effective.
3 Preliminaries and Task Definition
We first formalize the problem of selecting the most likely element in a set as a reranking problem; then, we define sequential reranking (SR); finally, we contextualize AS2 task in such framework.
3.1 Max Element Selection
In general, a large class of NLP (and other) problems can be formulated as a max element selection task: given a query  $q$  and a set of candidates  $A = \{a_{1},..,a_{n}\}$ , select  $a_{j}$  that is an optimal element for  $q$ . We can model the task as a selector function  $\pi : Q \times \mathcal{P}(A) \to A$ , defined as  $\pi(q, A) = a_{j}$ , where  $\mathcal{P}(A)$  is the powerset of  $A$ ,  $j = \operatorname{argmax}_{i} p(q, a_{i})$ , and  $p(q, a_{i})$  is the probability of  $a_{i}$  to be the required element.  $p(q, a_{i})$  can be estimated using a neural network model. In the case of transformers, said model can be optimized using a point-wise loss, i.e., we only use the target candidate to generate the selection probability. Pairwise or list-wise approaches can still be used (Bian et al., 2017), but (i) they would not change the findings of our study, and (ii) point-wise methods have been shown to achieve competitive performance in the case of transformer models.
3.2 Search with Sequential Reranking (SR)
Assuming that no heuristics are available to preselect a subset of most-likely candidates, max element selection requires evaluating each sample using a relevance estimator. Instead of a single estimator, it is often more efficient to use a sequence of rerankers to progressively reduce the number of candidates.
We define a reranker as a function  $R:Q\times \mathcal{P}(A)\to \mathcal{P}(A)$ , which takes a subset  $\Sigma \subseteq$
Also known as auxiliary or diagnostic classifiers.
5699
$A$ , and returns a set of elements,  $R(q,\Sigma) = \{a_{i1},\ldots ,a_{ik}\} \subset \Sigma$  of size  $k$ , with the highest probability to be relevant to the query. That is,  $p(q,a) > p(q,b)\quad \forall a\in \Sigma ,\quad \forall b\in A - \Sigma .$
Given a sequence of rerankers sorted in terms of computational efficiency,  $(R_{1}, R_{2}, \ldots, R_{N})$ , we assume that the ranking accuracy,  $\mathcal{A}$  (e.g., in terms of MAP and MRR), increases in reverse order of the efficiency, i.e.,  $\mathcal{A}(R_{j}) > \mathcal{A}(R_{i})$  iff  $j > i$ . Then, we define a Sequential Reranker of order  $N$  as the composition of  $N$  rerankers:  $SR_{N}(A) = R_{N} \circ R_{N-1} \circ \ldots \circ R_{1}(A)$ , where  $R_{N}$  can also be the element selector  $\pi(q, \cdot)$ . Each  $R_{i}$  is associated with a different  $k_{i} = |R_{i}(\cdot)|$ , i.e., the number of elements the reranker returns. Depending on the values of  $k_{i}$ , SR models with different trade-offs between accuracy and efficiency can be obtained.
3.3 AS2 Definition
The definition of AS2 directly follows from the definition of element selection of Section 3.1, where the query is a natural language question and the elements are answer sentence candidates retrieved with any approach, e.g., using a search engine.
4 SR with transformers
In this section, we explain how to exploit the hierarchical architecture of a traditional transformer model to build an SR model. First, we briefly recap how traditional transformer models (we refer to them as "monolithic") are used for sequence classification, and how to derive a set of sequential rerankers from a pre-trained transformer model (Section 4.1). Then, we introduce our Cascade Transformer (CT) model, a SR model that efficiently uses partial encodings of its input to build a set of sequential rerankers  $R_{i}$  (Section 4.3). Finally, we explain how such model is trained and used for inference in sections 4.3.1 and 4.3.2, respectively.
4.1 Monolithic Transformer Models
We first briefly describe the use of transformer models for sequence classification. We call them monolithic as, for all input samples, the computation flows from the first until the last of their layers.
Let  $\mathcal{T} = \{E;L_1,L_2,\dots ,L_n\}$  be a standard stacked transformer model (Vaswani et al., 2017), where  $E$  is the embedding layer, and  $L_{i}$  are the
Figure 1: A visual representation of the Cascade Transformer (CT) model proposed in this paper. Components in yellow represent layers of a traditional transformer model, while elements in purple are unique to CT; input and outputs of the model are shown in blue. In this example, drop rate  $\alpha = 0.4$  causes sample  $X^3$  to be removed by partial classifier  $C_{\rho(i)}$ .
transformer layers $^4$  generating contextualized representations for an input sequence;  $n$  is typically referred to as the depth of the encoder, i.e., the number of layers. Typical values for  $n$  range from 12 to 24, although more recent works have experimented with up to 72 layers (Shoeybi et al., 2019).  $\mathcal{T}$  can be pre-trained on large amounts of unlabeled text using a masked (Devlin et al., 2019; Liu et al., 2019c) or autoregressive (Yang et al., 2019c; Radford et al., 2019) language modeling objective.
Pre-trained language models are fine-tuned for the target tasks using additional layers and data, e.g., a fully connected layer is typically stacked on top of  $\mathcal{T}$  to obtain a sentence classifier. Formally, given a sequence of input symbols,  $X = \{x_0,x_1,\dots ,x_m\}$ , an encoding  $H =$
4That is, an entire transformer block, constituted by layers for multi-head attention, normalization, feed forward processing and positional embeddings.
For ranking tasks, the sequence of input symbols is typically a concatenation of the query  $q$  and a candidate  $a_{j}$ . In
3The design of an end-to-end algorithm to learn the optimal parameter set for a given target trade-off is left as future work.
5700
$\{h_0,h_1,\ldots ,h_m\}$  is first obtained by recursively applying  $H_{i}$  to the input:
$H_0 = E(X), H_i = L_i(H_{i-1}) \quad \forall i = 1, \ldots, n,$  where  $H = H_n$ . Then, the first symbol of the input sequence is fed into a sequence of dense feedforward layers  $D$  to obtain a final output score, i.e.,  $y = D(h_0)$ .  $D$  is fine-tuned together with the entire model on a task-specific dataset (a set of question and candidate answer pairs, in our case).
4.2 Transformer-based Sequential Reranker (SR) Models
Monolithic transformers can be easily modified or combined to build a sequence of rerankers as described in Seciton 3.2. In our case, we adapt an existing monolithic  $\mathcal{T}$  to obtain a sequence of  $N$  rerankers  $R_{i}$ . Each  $R_{i}$  consists of encoders from  $\mathcal{T}$  up to layer  $\rho (i)$ , followed by a classification layer  $D_{i}$ , i.e.,  $R_{i} = \{E;L_{1},\ldots ,L_{\rho (i)},D_{i}\}$ . For a sequence of input symbols  $X$ , all rerankers in the sequence are designed to predict  $p(q,a)$ , which we indicate as  $R_{i}(X) = y_{\rho (i)}$ . All rerankers in  $SR_{N}$  are trained independently on the target data.
In our experiments, we obtained the best performance by setting  $N = 5$  and using the following formula to determine the architecture of each reranker  $R_{i}$ :
$$
\rho (i) = 4 + 2 \cdot (i - 1) \quad \forall i = \{1, \dots , 5 \}
$$
In other words, we assemble sequential reranker  $\mathrm{SR}_5$  using five rerankers built with transformer models of 4, 6, 8, 10 and 12 layers, respectively. This choice is due to the fact that our experimental results seem to indicate that the information in layers 1 to 3 is not structured enough to achieve satisfactory classification performance for our task. This observation is in line with recent works on the effectiveness of partial encoders for semantic tasks similar to AS2 (Peters et al., 2019).
4.3 Cascade Transformer (CT) Models
During inference, monolithic transformer models evaluate a sequence  $X$  through the entire computation graph to obtain the classification scores  $Y$ .
order for the model to distinguish between the two, a special token such as "[SEP]" or  $\langle \langle s\rangle \rangle$  is used. Some models also use a second embedding layer to represent which sequence each symbol comes from.
Before being processed by a transformer model, sequences are typically prefixed by a start symbol, such as "[CLS]" or "<s)". This allows transformer models to accumulate knowledge about the entire sequence at this position without compromising token-specific representations (Devlin et al., 2019).
This means that when using  $SR_{N}$ , examples are processed multiple times by similar layers for different  $R_{i}$ , e.g., for  $i = 1$ , all  $R_{i}$  compute the same operations of the first  $\rho(i)$  transformer layers, for  $i = 2$ ,  $N - 1$  rerankers compute the same  $\rho(i) - \rho(i + 1)$ , layers and so on. A more computationally-efficient approach is to share all the common transformer blocks between the different rerankers in  $SR_{N}$ .
We speed up this computation by using one transformer encoder to implement all required  $R_{i}$ . This can be easily obtained by adding a classification layer  $C_{\rho(i)}$  after each  $\rho(i)$  layers (see Figure 1). Consequently, given a sample  $X$ , the classifiers  $C_{\rho(i)}$  produces scores  $y_{\rho(i)}$  only using a partial encoding. To build a CT model, we use each  $C_{\rho(i)}$  to build rerankers  $R_{i}$ , and select the top  $k_{i}$  candidates to score with the subsequent rerankers  $R_{i+1}$ . We use the same setting choices of  $N$  and  $\rho(i)$  described in Section 4.2.
Finally, we observed the best performance when all encodings in  $H_{\rho(i)}$  are used as input to partial classifier  $C_{\rho(i)}$ , rather than just the partial encoding of the classification token  $h_{\rho(i),0}$ . Therefore, we use their average to obtain score  $y_{\rho(i)} = C_{\rho(i)} \left( \frac{1}{m} \sum_{l=1,\dots,m} h_{\rho(i),l} \right)$ . In line with Kovaleva et al. (2019), we hypothesize that, at lower encoding layers, long dependencies might not be properly accounted in  $h_{\rho(i),0}$ . However, in our experiments, we found no benefits in further parametrizing this operation, e.g., by either using more complex networks or weighting the average operation.
4.3.1 Training CT
The training of the proposed model is conducted in a multi-task fashion. For every mini-batch, we randomly sample one of the rankers  $R_{i}$  (including the final output ranker), calculate its loss against the target labels, and back-propagate its loss throughout the entire model down to the embedding layers. We experimented with several more complex sampling strategies, including a round-robin selection process and a parametrized bias towards early rankers for the first few epochs, but we ultimately found that uniform sampling works best. We also empirically determined that, for all classifiers  $C_{\rho}(i)$ , backpropagating the loss to the input embeddings, as opposed to stopping it at layer  $\rho(i - 1)$ , is crucial to ensure convergence. A possible explanation could be: enabling each classifier to influence the input representation during backpropagation ensures that later rerankers are more robust against
variance in partial encodings, induced by early classifiers. We experimentally found that if the gradient does not flow throughout the different blocks, the development set performance for later classifiers drops when early classifiers start converging.
4.3.2 Inference
Recall that we are interested in speeding up inference for classification tasks such as answer selection, where hundreds of candidates are associated with each question. Therefore, we can assume without loss of generality that each batch of samples  $B = \{X^1,\ldots ,X^b\}$  contains candidate answers for the same question. We use our partial classifiers to throw away a fraction  $\alpha$  of candidates, to increase throughput. That is, we discard  $k_{i} = \lfloor \alpha \cdot k_{i - 1}\rfloor$  candidates, where  $\lfloor \cdot \rfloor$  rounds  $\alpha \cdot k_{i - 1}$  down to the closest integer.