Leap-of-Thought: Accelerating Transformers via Dynamic Token Routing
Authors: Yeachan Kim
Junho Kim
Jun-Hyung Park
Mingyu Lee
SangKeun Lee
Venue: EMNLP 2023 Main (2023)
Keywords: transformer
language models
token routing
token pruning
input length reduction
TL;DR: We introduce Leap-of-Thought (LoT), a novel token reduction approach that dynamically routes tokens within transformer layers.

Abstract: Computational inefficiency in transformers has been a long-standing challenge, hindering the deployment in resource-constrained or real-time applications. One promising approach to mitigate this limitation is to progressively remove less significant tokens, given that the sequence length strongly contributes to the inefficiency. However, this approach entails a potential risk of losing crucial information due to the irrevocable nature of token removal. In this paper, we introduce Leap-of-Thought (LoT), a novel token reduction approach that dynamically routes tokens within layers. Unlike previous work that irrevocably discards tokens, LoT enables tokens to `leap' across layers. This ensures that all tokens remain accessible in subsequent layers while reducing the number of tokens processed within layers. We achieve this by pairing the transformer with dynamic token routers, which learn to selectively process tokens essential for the task. Evaluation results clearly show that LoT achieves a substantial improvement in computational efficiency. Specifically, LoT attains up to 25x faster inference time without a significant loss in accuracy

Leap-of-Thought: Accelerating Transformers via Dynamic Token Routing
Yeachan Kim Junho Kim', Jun-Hyung Park?, Mingyu Lee', SangKeun Lee
*Department of Artificial Intelligence, Korea University, Seoul, South Korea BK21 FOUR R&E Center for Artificial Intelligence, Korea University, Seoul, South Korea Department of Computer Science and Engineering, Korea University, Scoul, South Korea (yeachan, monocrat, irttiiiisso irishe7, decon9201, yakophyekoorea ac.
Abstract
transformer layers (Goyal al., 2020; Ye et al., 2021; Guan et al., 2022). The rationales behind this approach are two-fold: (i) not all tokens are equally significant to the task (Dai et al., 2020), and (ii) all token representations gradually become sim- ilar over layers (Abnar and Zuidema, 2020; Phang et al., 2021). Based on these rationales, previous studies progressively removes the less significant or redundant tokens (Figure la), selected either ran- domly (Hou et al., 2022) based on the attention scores (Wang et al., 2019). However, the permanent token removal in the previous work could entail a risk of discarding crucial information in pursuit of efficiency, which can potentially degrade perfor- mance by hindering the fine- fine-grained understanding of the input. Moreover, since the token reduction space with the permanent removal is proportionally constrained with the number of remaining tokens. sub-optimal to explore the diverse reduction strategies that potentially offer greater efficiency. These limitations suggest that there is still room for further improvement.
Computational inefficiency in transformers has been long-standing challenge, hindering the deployment in coource-constrained or real- time applications. One promising approach mitigate this limitation is progressively remove less significant tokens, given that the sequence length strongly contributes to the in- efficiency. However, this approach entails po- tential risk of losing crucial information due to the irrevocable nature of token removal. In this paper, we introduce a novel token reduction approach that dynami- cally routes tokens within layers. Unlike previ- ous work that irrevocably discards tokens, LoT enables tokens to 'leap' across layers. This ensures that all tokens remain accessible in sub- sequent layers while reducing the number of tokens processed within layers. We achieve this by pairing the transformer with dynamic token routers, which learn to selectively process to- kens essential for the task. Evaluation results clearly show that LoT achieves substantial im- provement on computational efficiency. Specif- ically, LoT attains up 25x faster inference time without a significant loss in accuracy'.
-hhoughht (LoT),
In this paper, we propose Seap-of-Thoughte (LoT)?, a novel token reduction approach that en- ables the dynamic routing of tokens across lay- ers. In contrast to permanent removal strategies, LoT allows the tokens to 'leap' over each layer, thereby retaining access all original tokens in the subsequent layers while reducing the number of to- kens processed within each layer of the transformer (Figure Ic). We achieve this by coupling the trans- former layers with dynamic token routers, which learns to decide whether the given token should be processed at the current layer or leaped forward to the next layer. Moreover, in order to steer the token router towards making informed and efficient decisions, introduce mradientgguide training
1 Introduction
The advent of Transformer (Vaswani et al., 2017) has spurred paradigm shift, most notably in natural language processing (Brown et al., 2020; Chowdhery et al., 2022), but also extending to com- puter vision (Dosovitskiy et al., 2021; Liu al., 2021). However, the impressive capabilities of the transformer typically come with non-trivial com- putational costs, which scale quadratically with the length of the input sequence. This computa- tional burden poses significant challenge when de- ploying the transformer-batt models in resource- constrained real-time systems (Sun et al., 2020).
One typical approach to tackle this challenge is reduce the number of tokens processed within
'The name comes from the LoT behavior where each token (corresponding meaning or thought) leaps over layers. LoT is not related to the Chain-of- Thought (Wai et al., 2022), which introduces the concept of continuous prompting.
OOur cnde is available at fttps://githut com/ yeachan-kr/lot.

---

W1
2
W3
W1
W2
Wg
W1
W2
(a) Width-wise approach
(h) Depth-wise approach
(c) Leap--f--houuueeeeoe (Ours)
Figure I: Reduction strategies width-wiss wiie (token reduction), depth-wise (layer reduction) and -fhoughh (ours) to achieve computational efficiency. The tokens with the dashed lines indicate the unused tokens in each layer.
that informs each router of which tokens do signifi- cantly contribute. Consequently, the router learns to identify which tokens are crucial for the task and where these tokens should be processed within layers in order achieve greater efficiency.
Related Work
In this section, we mainly review the methods that adaptively control the computation in pre-trained language models. Recent approaches can be clas- sified into two categories: width-wiss and depth- wise approaches. The former focuses on reducing the number of tokens processed by transformers, while the latter aims decrease the number of com- putational layers. Figure 1 illustrates the distinct behaviors of these approaches, including LoT.
LoT offers several advantages compared to the permanent removal. Primarily, LoT has the poten- tial to mitigate the risk losing crucial informa- tion related to the task, given that the decisions for each token are recoverable in subsequent lay- ers. In addition, LoT provides higher degree of freedom in token reduction, thereby facilitating the exploration of diverse search space for greater efficiency, which is similarly observed in network compression (Mao et al., 2017; Park et al., 2023). To substantiate the efficacy of LoT, we perform evaluations across extensive experiments. Compre- hensive results demonstrate that the model employ- ing LoT reveals substantial speedup gains without a significant loss in task accuracy. Additionally, through the analysis of LoT, provide justifica- tion for the efficacy of the dynamic token routing mechanism and illustrate how LoT achieves greater efficiency. In summary, the contributions of the pa- per include the followings:
2.1 Width-wise Reduction on Transformer
Given that the computational costs of the trans- former are heavily influenced by the length the input sequence (Tay et al., 2023), recent re- search has endeavored to minimize the sequence length by progressively removing less significant tokens from the input (Figure la). For instance. POWER-BERT (Goyal al., 2020) have initially explored the removal of tokens that receive the least attention from other words, operating on the premise that tokens with less attention are less sig- nificant. However, several studies on the trans- former interpretability have shown that the atten- tion scores might not be reliable indicators of the actual token contribution (Jain and Wallace, 2019; Abnar and Zuidema. 2020; Meister al., 2021). As such, TR-BERT (Ye al., 2021) and Tran- skimmer (Guan al., 2022) have suggested ken removal strategies that can be learned during training, by using reinforcement learning and re- parameterization tricks, respectively. Subsequently, AdapLeR (Modarressi et al., 2022) have proposed saliency- strategy that eliminates tokens by estimating the saliency of tokens via the gradients of the input embeddings with respect to the predic- tions.
We introduce Leap-of- Thought, novel to- ken reduction approach that enables dynamic token routing within the transformer, which re- duces the processed tokens within each layer while preserving crucial information.
We propose gradient-guides training steer the dynamic token router towards making more informed decisions about whether the tokens should be processed or leaped over.
We demonstrate the efficacy of LoT through extensive experiments and analysis on various benchmarks, establishing LoT as promising approach for the token reduction.
While these methods have demonstrated effi- ciency downstream tasks, they irrevocably dis-

---

Classifier
Transformer Block #1.
Gradient-guide Router Training
Fwyow,
Contribuzion
Transformer Black
Position-wise Food-forward Layer
Sclf-atiention Layer
Layer-wisc aggregation
W2
Transformer Black
Token Importancc
Position-was ard Layer
actentitn Layer
W3
Leep (w,, MS)
Leep (H5.)9))
P
Token Representation
Forward Pass
W1
W2
W3
W4
Token Rouler
Backoward Pass
Figure 2: Overview of Leap-of- Thouuhh (LoT). Starting from the embedding layer. the token routers located between every transformer layers. For the words that are decided to be leaped forward to the next layer, their representations are merged into one pseudo token to provide the minimal information unused tokens.
card input tokens, which might lead to poten- tial loss of crucial information. Moreover, the scarch space for token removal is proportionally constrained by the number of remaining tokens, thereby restricting flexibility in optimizing reduc- tion strategies. In contrast, since LoT allows the model revisit all tokens, the crucial information can be better preserved within the transformer. Be- sides, the ability to revisit tokens endows LoT with a higher degree of flexibility in exploring diverse reduction space that potentially offers greater effi- ciency.
Instead of implementing an exit strategy, Layer- Drop (Fan et al., 2020) has demonstrated its effi- ciency by randomly dropping the layers, and BERT- of- Theseus (Xu et al., 2020) has learned to replace the subsequent layers with a single layer.
While these works allow adaptive computation on different inputs to achieve efficiency, the level of granularity in the depth-wise approach is con- strained by the number of layers. This could result in the sub-optimal efficiency and difficulty in signing fine-grained computations diverse set of samples.
2.2 Depth-wise Reduction on Transformer
Leap-of- Leap-of-Thoughtt Dynamic Token Routing for Accelerating Transformer
The principle behind depth-wise approach is to al- locate minimal layer computations easy samples while dedicating more layer computations diffi- cult samples (Figuree lb). The distinction between different works lies in the criteria for difficulty. PABEE (Zhou et al., 2020) has proposed patience- based exit strategy that halts the forward at an intermediate layer only when the pre-defined number of subsequent layers yield the same pre- dictions. Similarly, DeeBERT (Xin al., 2020) and FastBERT (Liu al., 2020) have employed the predictive entropy to replace the patience, and PCEE-BERG (Zhang et al., 2022) has combined both patience and confidence for the exit criteria.
In this section, we elaborate on Leap of Thought (LoT), which dynamically routes tokens across lay- ers to improve computational efficiency. To this end, we introduce a dynamic token router in learn- ing to decide which token should be processed in the current layer lcaped forward to the subse- quent layer (Section 3.1). To ensure that the token router makes well-informed decisions, each token router trained by gradient-based token impor- tance (Section 3.2). The overall process of LoT is illustrated in Figure 2.

---

3.1 Learning to Leap Transformer Layer order to enable tokens to leap across transformer layers, we introduce a dynamic token routing mech- anism that adaptively selects tokens for utilizing in the current layer, while pushing the unused tokens forward to subsequent layers for potential use.
Notably, we employ the [CLS] token (i.e., up as the context vector, given that it serves as a contex- tual memory, being retained throughout all layers. However, training the router in an end-to-end manner is non-trivial due to the non-differentiable nature of the routing function. To circumvent this, utilize the Gumbell-oftmax reparameterizztioon (Jang et al., 2017) to approximate the discrete de- cisions during training. Specifically, we introduce a continuous relaxation for the discrete decision- making process. During the forward pass, we sam- ple from Gumbel distribution and apply the soft- max function to approximate the discrete decisions
Dynamic Token Router. To initiate the routing mechanism, we start by the definition of a dynamic token router, a lightweight module located between every transformer layers. Each router takes token representations as the imput (i.e., embedding or out- puts from the previous layer) and learns to produce a binary decision for each token: "1" denotes that it is processed at the current layer, and "O" denotes that it leaps to the next layer. The dynamic token router is formulated follows:
u (w) = softmax
(aaaaxx(((og) ((log(u(w) + g) /T), (4)
where is sample from Gumbel distribution, and 7 is the temperature parameter controlling the smoothness of the approximation. During the back- ward pass, we replace the gradient of the non- differentiable function with that of the Gumbel- softmax using straight- taattttthouuhh--iihaae estiimator (Bengio et al., 2013). This allows gradients to flow through the router, enabling end-to-end optimization of the entire model.
u(w)= = o2 (22 (Wi(L((w)) + bi) + b2) (1)
where w is a token representation, W and b denote the weights and biases for linear transformation, respectively, and indicate the GeLU activa- tion and softmax function, respectively, and L.((-) denotes the layer normalization (Ba et al., 2016) We then derive the routing decision based on the prediction of the router.
Token Merging, While the routing ability allows the model to preserve crucial information, main- taining the minimal information of unused tokens can be beneficial. We thus introduce token merging mechanism. Formally, the merged token is con- structed as follows:
R(w)= = { if (pnices(1)) > Meeup(w)
(2)
otherwise,
where the subscript of u(w) represents the prob- ability for each actions (i.e., process or leap the layer).
1 R-1
(1) mmrgge = 1I + c0)) = 0)
Routing Tokens. Once the token router is estab- lished, the routing decision is applied to all tokens before they are fed into the transformer computa- tion. Formally, let the token representations in the I-th layer be denoted as w() w|
(5)
where 11 [z]] is the indicator function that returns one if the statement x is true; otherwise zero, and m is the number of tokens to be leaped. The merged token is appended to the input and only utilized in the self-attention layer. In the next layer, the token is replaced with new merged token based the new routing results (i.e., ((+1)).
(1) w 1 where n is the length of an input. The routing decision is made for each token by applying the following gating function.
( w 0 (R!(( (m) ()) 10
3.2 Gradient-guided Router Training
where (-) is the routing function on the 1-th layer, indicates the Hadamard product, and c(I) is the context vector used to make the routing deci- sion by considering the current context information.
To steer the token router towards making informed decisions, we also introduce a mradient-guide router training, which directly provides the super- vision of the significant tokens to the routers.
We empirically observed that applying layer normaliza- tion makes the training of LoT more stable.
Guidance Derivation. As guidance for the router, the gradients of the token representations are leveraged, given that the gradient information
"Note that the [CI.S token is not forwarded to the router to perform the classification task correctly.

---

can encode the sensitivity of the output to the in- put tokens, providing insight into which tokens are being more influential for prediction (Jain and Wal- lace, 2019; Abnar and Zuidema, 2020). Inspired by Grad-CAM (gradient-weighted class activation map) (Selvaraju et al., 2017) which uses gradient information flowing into the final layer of convolu- tional neural networks, we derive class activation tokens each each layer. Formally, let y be the prediction for the bround-truth, the class activation tokens (CAT) can be derived as follows:
Experiment
In this section, evaluate the proposed method on a series of downstream tasks. We specifically demonstrate that introducing the leap action results in a more favorable computational efficiency com- pared to the prior methods.
4.1 Experimental Setup
4.1.1 Datasets
We perform diverse tasks to verify the general ap- plicability. These tasks involve scenarios where the model needs to comprehend single sequence, well as cases that requires understanding the seman- tic relationship between multiple sequences. For the single input tasks, we use SST-2 (Socher al., 2013) and IMDB (Maas et al., 2011) for sentiment analysis, AG's news (Zhang et al., 2015) for topic classification, DBpedia (Lehmann et al., 2015) for ontology classification, and HateXplain (Mathew et al., 2021) for hate speech detection. For the mul- tiple input tasks, we perform paraphrasing tasks on MRPC (Dolan and Brockett, 2005) and natural language inference tasks MNLI (Williams al., 2018) and QNLI (Rajpurkar et al., 2016).
dy
CAT
(
(6)
(1)
aw
Based on the gradient-weighted token representa- tions, we derive the importance by the magnitude of each CAT. Specifically, we aggregate the token importance from all layers since it can provide better identification for the important tokens (Qiang al., 2022) (detailed in Section 5.1):
4-1 CAT; - - CCAT 1=0
(7)
Lastly, we need to identify which range of token importance should be considered as significant. To this end, we simply select the tokens whose cumula- tive sum of their sorted and normalized importance scores falls below pre-defined threshold similar to the candidate set of nucleus sampling (Holtzman et al., 2020).
4.1.2 Baselines
Following the prior work, we use the pre-trained BERT bave (Devlin et al., 2019) backbone net- work We then compare with six baselines includ- ing the backbone model: POWER-BERT (Goyal et al., 2020) which utilizes the attention maps to eliminate the tokens; TR-BERT (Ye et al., 2021) that adopts reinforcement learning to learn a re- moval strategy; AdapLeR (Modarressi et al., 2022) that utilizes the saliency maps of the input words to remove tokens. Additionally, we also compare LoT with different direction of reduction approaches. We compare PCEE-BERT (Zhang al., 2022), which adaptively exits from the transformer by con- sidering both the confidence and patience, and Dis- tilBERT (Sanh et al., 2019), which is the resultant model from knowledge distillation.
Training Objective. The dynamic token routers are trained to process only the significant tokens which are selected from the above procedure. Let tे be the selection decision for the i-th token given the selected tokens with a value of one otherwise zero, the objective for the router is formulated as follows:
Σ ws log Uprocess (i(
+(1 wh)) log( Utcap ww (1)))) )),
(8)
The overall objective function for the downstream task can be formulated as follows:
4.1.3 Evaluation Metries
L-1
Following the recent prior work (Modarressi et al., 2022), we evaluate each method using both the task accuracy and the number of (FLOPs). Given that FLOPs are independent of
A
Σ
== Lrask +
(9)
L
1=0
where Ltask is the task-specific loss function (e.g.. cross entropy for the classification). and harmony coefficient À to balance the two loss terms.
In Appendix, we included the experiments on different architectures (TinyBERT, BERT

---

Table 1: Evaluation results of test accuracy (%) and speedup ratio on the single input tasks. The speedup ratio (denoted Speed) is computed by comparing the FLOPs of each baseline with the backbone. The best and second best results are highlighted in boldface and underlined, respectively.
SST-2
IMDB
HateXplain
AG's news
DBpedia
Method
Acc.
Speed
Acc.
Speed
Acc.
Speed
Acc.
Speed
Acc.
Speed
Baseline
92.7
1.00x
93.8
1.00x
68.3
1.00x
94.4
1.00x
99.3
1.00x
DistiBBERT (Sanh et al., 2019)
92.2
x.00x
92.4
2.00x
68.4
x.00x
94.2
2.00x
99.0
2.00
PCEE- BERT (Zhang et al., 2022)
91.9
1.56x
92.3
2.63x
67.9
3.09
93.4
5.54x
99.0
5.80
DWWER-BERT (Goyal et al., 2020)
92.1
1.18x
92.3
1.70x
66.9
2.69
92.1
12.50x
98.1
14.80
TR-BERT (Ye et al., 2021)
92.1
1,46x
93.2
2.90x
67.9
2.23x
93.2
10.20x
98.9
10.01
AdapLeR (Modarressi al., 2022)
92.3
x.49x
91.7
3.21x
68.6
4.73x
92.5
17.10x
98.9
22.23x
LoT (ours)
92.9
2
92.4
3.84x
68.8
5.21x
92.4
25.10x
99.1
19.76x
Table Evaluation results of test accuracy (%) on mul-
AdapLeR
PCEE-BERT
LoT (ours)
tiple imput tasks. The best and second best results highlighted in boldface and underlined, respectively.
Test Accuracy (%)
92
DistiBERT
MRPC
MNLJ
QNLI
90
Method
Speed
Ace.
Speed
Speed
88
Baseline
87.5
1.00x
84.2
1.00x
90.3
1.00
DistiBERT
87.7
2.00x
82.0
2.00x
87.9
x.00
1.5
2.0
2.5
PCEE-BERT
87.2
1.34x
82.5
1.10x
90.4
1.31
SpeedUp
BAWER-BERT
88.0
1.07x
82.9
1.10x
89.7
1.23x
(a) SST-2
TR-BERT
81.9
1.16x
84.8
1.00x
89.0
1.09
AdapLeR
87.6
x.27x
82.9
1.42x
89.3
1.47x
AdapLeR
PCEE BERT
LoT (ours)
LoT (ours)
88.4
3.29 x
83.1
2.53
90.2
2.74x
Test Accuracy
89
hardware, this allows us to evaluate the accelera- tion of the model without taking into account the operating environment'. Furthermore, following the recent practice (Guan et al., 2022; Modarressi et al., 2022), we compute FLOPs for single in- ference, which enables us to evaluate per-example inference and avoid pseudo speed-up resulting from the elimination of padding tokens. In the experi- ments, we present the relative speed-up compared to that of the BERT have.
88
DONIBERT
87-
1.5
2.0
2.5
3.0
SpeedUp
(b) QNLI
Figure 3: Trade-off curve between task accuracy and speedup on SST and QNLJ datasets.
accuracy. For example, LoT achieves speedup of 25.10x 10x on the AG's news dataset without a signif- icant loss in task accuracy. Such a speedup with comparable performance to other baselines demon- strates that LoT better preserves the crucial infor- mation related to tasks. This result clearly supports the significance of the leap action in the token duction approach.
4.1.4 Training Details
We implement the proposed method using PyTorch. For the
associated with LoT (i.e., threshold p in Eq. (8), and À in Eq (9)), we search the best practice parameters on the validation sets. are listed in the Appendix.
The
4.2 Main Results
Multiple input tasks. We also highlight the sults on the tasks that involve pairs of distinct sen- tences in Table 2. Given that these tasks require the model to comprehend the semantic relationships between multiple sentences, removing tokens could lead to the loss of important information needed for
Singe input tasks. In Table 1, we first present the evaluation results on the single input tasks along with strong baselines. Notably, LoT achieves sub- stantial speedup gains without compromising the "We also calculated wall-clock inference time in Appendix.

---

Table 3: Ablation study LoT on SST-2 and MRPC. and "wlo" indicates the model without the corresponding component. The token merging is related to Eq. (5), and the layer-wist wise aggregation of CAT related Eq. (7).
w/ guidance (SpeedUp: Accuracy 92.9%)
w/o guidance (SpeedUp: 1.79x, Accuracy, 91.9%)
Remaining Tokens (%)
100 80
SST-2
MRPC
Method
Acc.
SpeedUp
F1.
SpeedUp
(ours)
92.9
x.30x
88.4
x.29x
w/o token merging
92.4
2.16x
88.0
2
Layer Index
w/o layer-wise CAT
92.3
1.96x
88.6
2.46x
(a) SST-2
guidance (SpeedUp: 5.21x, Accuracy; 68.8%)
68.8%)
SST-2
MRPC QNLI
AG's news DBpedia
w/o guidance (SpeedUp: 2.31 Accuracy. 67.9%)
HateXplain
Remaining Tokens (%)
100
Remaining Tokens (%%
100
80
40-
60
2
5
6
10
12
20
Layer Index
(b) HateXplain
12
Layer Index
Figure 5: Comparison with LoT variant that learns the leap action in an end-to-end manner.
Figure 4: Remaining token distribution across various layers and datasets, excluding [PAD] tokens that may lead to pseudo speedup.
initially observe that merging tokens enhances per- formance on both metrics. This indicates that pro- viding abstractive information about tokens to be leaped enables the model to process more reduced tokens in each layer, thereby leading to the im- proved efficiency. On the other hand, in contrast GradCAM (Selvaraju al., 2017), which utilizes gradient information from the input space, LoT ag- gregates gradient information from all layers. To evaluate the impact the aggregation, we compare the performance when extracting gradient informa- tion solely from the input. The ablation shows that aggregating CAT from all layers substantially im- proves the computational efficiency, indicating that the aggregation allows the model to better identify tokens that are crucial to the task at hand.
understanding these relationships. Hence the ex- isting methods, which permanently remove the to- kens, reveal the low speedup gains these datasets. Nevertheless, LoT achieves greater speedup gains with comparable accuracy, which shows the versa- tility of LoT on diverse tasks. The overall experi- mental results verify that LoT can be successfully applied into real-world applications that demand both accuracy and efficiency.
Trade-off. To confirm the better computational efficiency of LoT, we show the trade-off curves between task accuracy and speedup gains on two representative datasets in Figure 3. This shows that LoT maintains the higher accuracy over wide range of speedups, clearly demonstrating the better trade-off of LoT compared other approaches.
5.2 Routing Distribution on Different Layers
We also analyze the routing distribution across dif- ferent layers for various datasets. Figure 4 shows the ratio of remaining tokens in each layer. Inter- estingly, we observe distinctive patterns of LoT for the reduction strategy. In the majority of datasets, less than half of the tokens are processed in the first layer. Subsequently, the behavior of LoT varies depending on the complexity of the tasks. For the simpler tasks (i.e., achieved higher speedup gains), the number the processed tokens tend to
5 Analysis
this section, we analyze the behavior of LoT in detail. We specifically focus on how LoT achieves a greater efficiency than other baselines.
5.1 Ablation Study
Table 3, we present an ablation study to dis- sect the contributions components in LoT. We

---

12
Layer Index
Dall
backes
(897)
Rast
Jaayy
boo!
SST-2
QNLI
Figure 6: Illustration of the token routing on examples. The darker block each layer indicates the of corresponding token in the layer while the lighter block denotes the leap action.
consistently decrease over the subsequent layers. Conversely, for more challenging tasks (e.g.. SST, MRPC, QNLI), the model tends to make use of larger number of tokens in few subsequent layers. These patterns indicate that LoT is capable of adap- tively executing optimized reduction strategies for different datasets, underscoring the flexibility and diversified search space of LoT.
ples on two different tasks, SST-2 and QNLI It evident that the routing maps are irregular and sparse, which demonstrates the flcxibility of LoT to reach greater efficiency. Moreover, the impor- tant tokens related to tasks (e.g.. dull, lifeless and amateur in SST-2) tend to be utilized in the deeper layers, whereas less contributing tokens are often bypassed. Interestingly, those important tokens are not consistently used in all layers. For example, the sentiment-related words SST-2 are temporally used in the earlier layers and then reused after sev- eral layers. This result highlights LoT's distinctive strategy in optimizing computational efficiency by selectively engaging with task-relevan tokens as needed. We provide the additional case study in Section F of Appendix.
5.3 Significance of Router Guidance
The token routers in LoT are supervised directly from the aggregated gradient information. To verify the significance of the supervised router training, we compare LoT with an alternative version that learns to decide the leap action without the guid- ance. To implement this baseline, we replace the guidance loss (i.e., Eq. (8)) with a regularization loss that enforces leap actions for all tokens, which is similarly used in previous work (Ye et al., 2021; Guan et al., 2022; Modarressi et al., 2022). Since this regularization term conflicts with the task ob- jective, the router could learn retain only the contributing tokens while bypassing the less con- tributing tokens. Figure shows the comparison result. We observe that the model without guidance can also achieve computational efficiency by learn- ing to leap tokens. However, explicit supervision on the important tokens enables the router to by- pass greater number of tokens, especially in the earlier layers. This phenomenon can be attributed the low learning capacities of the earlier layers identifying significant tokens without an explicit guidance. The overall results empirically justify the significance of the gradient-guided router training.
6 Conclusion
In this work, we have proposed Leap-of-Thouglt (LoT), a novel token reduction strategy that en- ables the dynamic routing of tokens within the transformer layers. Unlike the previous works that permanently remove tokens, LoT learns to decide whether the given token should be processed in the current layer leaped forward to the next layer. This ensures that all tokens remain accessible in subsequent layers while reducing the number of tokens processed within layers. Through the guid- ance from the gradient information, each router learns to process only the significant tokens to the task while bypassing the less contributing tokens. The comprehensive evaluations have convincingly supported the superiority of the proposed method by showing that LoT achieves substantial speedup gains over state-of-the-an the-rar methods with the compa- rable task accuracy. The analysis also have strongly supported that introducing the leap action leads to the substantially improved efficiency.
5.4 Case Study of Leap-of- Thoought
Lastly, we examine the behavior of LoT through case studies. Figure 6 exhibits the routing exam-

---

Limitations
Acknowledgment
While the proposed method allows transformer- based pre-trained models to achieve greater com- putational efficiency, there are a few potential limi- tations.
This work was supported by the Basic Research Program through the National Research Founda- tion of Korea (NRF) grant funded by the Korea government (MSIT) (2021R (12333102310 and In- stitute of Information & Communications Technol- ogy Planning & Evaluation (ITTP) grant funded by the Korea government (MSIT) (No.2019-0-00079, Artificial Intelligence Graduate School Program (Korea University)).
Interpretability Several existing methods for interpretability, such as layer-wise analysis (Ten- ney et al., 2019), might not be compatible with our method, given that LoT benefits from the ir- regularity. As an alternative, we believe that ag- gregating routing results across layers can serve reliable indicator of interpretability to a cer- tain extent. For example, in the case study (Figure 6), while there are irregularity on individual layers. the model tends frequently use the task-related tokens across layers, such as dull, lifeless, and am- ateur in the first sample (sentiment analysis), and who, most, player (in the question part of the pair). Bart, Starr, MVP (in the answer part of the pair) in the second example (natural language inference). Such comprehensive analysis across layers can provide some degree of interpretability for the pre- dictions.
References
Samira Abnar and Willem H. Zuidema. 2020. Quanti- fying attention flow transformers. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 4190-4197. Asso- ciation for Computational Linguistics.
Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. 2016. Layer normalization. CoRR, abs/1607.0644
Yoshua Bengio, Nicholas Léonard, and Aaron C. Courville. 2013. Estimating or propagating gradients through stochastic neurons for conditional computa- tion. CoRR, bos//308.3432.
Router Overhead In comparison to the vanilla backbone, LoT employs token routers to perform the dynamic computation, which imposes extra model parameters and computation overhead, simi- lar to other baselines (Yc et al., 2021; Modarressi al., 2022: Zhang al., 2022). This is the rea- son why we have carefully designed the routers as lightweight modules that only account for 2.15 1% (0.17% for each router) and 2.2% (0.18% for each router) of the memory and computational overhead of the entire model, respectively. To understand the effect of the router capacity, we analyze the trade- off between the computational overhead and total performance in Section D of Appendix. The result shows that such a lightweight router is sufficient to achieve significant speedup gains without compro- mising task accuracy. Nevertheless, in this paper, we confirm the applicability of such a lightweight router only in the natural language understanding tasks. Therefore, designing routers for natural lan- guage generation tasks (c.g., summarization. ma- chine translation) can be promising avenue for future research.
Tom B. Brown, Benjamin Mann. Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss. Gretchen Krueger. Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alce Radford, llya Sutskever, and Dario Amodei. 2020. Language models few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham. Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vin- odkumar Prabhakaran, Emily Reif, Nan Du Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya. Sanjay Ghemawat, Sunipa Dev. Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito. David Luan, Hycontack Lim, Barret Zoph, Alexander Spiridonov. Ryan Sepassi. David Dohan, Shivani Agrawal, Mark Omernick, An- drew M. Dai, Thanumalayan Sankaranarayana Pil- lai, Marie Pellat, Aitor Lewkowycz, Erica Moreira. Rewon Child, Oleksandr Polozov, Katherine Lee,

---

Zongwei Zhou, Xuezhi Wang. Brennan Sacta, Mark Diaz. Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstem, Douglas Eek, Jeff Dean, Slav Petrov, and Noah Fiedel. 2022. Palm: Scaling language mod- eling with pathways. CoRR, abs/2204.023111.
Le Hou, Richard Yuanzhe Pang, Tianyi Zhou, Yuexin Wu, Xinying Song. Xiaodan Song. and Deany Zhou. 2022. Token dropping for efficient BERT pretrain- ing. In Proceedings of the 60th Anmual Meeting of the Association for Computational Linguistics, pages 7774-3784. Association for Computational Linguis- tics.
Zihang Dai, Guokun Lai, Yiming Yang, and Quoe Le. 2020. fuamel-transformmr: Filtering out sequential redundancy for efficient language processing. In Advances in Neural Information Processing Systems, volume 33, pages 4271 42822 Curran Associates, Inc.
Sarthak Jain and Byron C. Wallace. 2019, Attention is not explanation. In Proceedings of the 57th An- nual Meeting of the Association for Computational Linguistics, pages 3543-3556, Association for Com- putational Linguistics.
Jacob Devlin, Ming-War Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: pre-training of deep bidirectional transformers for language under- standing. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Tech- nologies, pages 4171-4186. Association for Compu- tational Linguistics.
Eric Jang, Shixiang Gu, and Ben Poole. 2017. Cate- gorical reparameterization with gunbel-softmax. In Proceedings of the Fifth International Conference on Learning Representations. OpenReview.net.
Xiaoqi Jiao, Yichun Yin. Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang. and Qun Liu. 2020. Tinybert: Distilling BERT for natural language under- standing. In Findings of the Association for Computa- tional Linguisties: EMNLP 2020, pages 4163 4174. Association for Computational Linguistics.
William B. Dolan and Chris Brockett. 2005, Automati- cally constructing corpus of sentential paraphrases. In Proceedings of the Third International Workshop on Paraphrasing, pages 9-16. Asian Federation of Natural Language Processing.
Jens Lehmann, Robert Isele, Max Jakob, Anja Jentzsch, Dimitris Kontokostas, Pablo N. Mendes, Sebastian Hellmann, Mohamed Morsey, Patrick van Kleef, Sören Auer, and Christian Bizer. 2015, Dbpedia A large-scale, multilingual knowledge base extracted from wikipedia. In Semantic Web, volume 6, pages 167-195. IOS Press.
Alexey Dosovitskiy, Lucas Beyer. Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer. Georg Heigold. Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. 2021. An image is worth 16x16 words: Transformers for image recognition at scale. In Proceedings of the Ninth In- ternational Conference on Learning Representations. OpenReview.net.
Weijie Liu, Peng Zhou, Zhiruo Wang, Zhe Zhao, Haotang Deng, and Qi Ju. 2020. FastBERT: self- distilling BERT with adaptive inference time. In Proceedings of the 58th Annual Meeting of the Asso- ciation for Computational Linguistics, pages 6035- 6044. Association for Computational Linguisties.
Angela Fan, Edouard Grave, and Armand Joulin. 2020. Reducing transformer depth on demand with struc- tured dropout. In Proceedings of the Fighth Inter- national Conference on Learning Representations. OpenReview.net.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. 2021. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE International Conference on Computer Vision, pages 9902- IEEE Computer Society.
Saurabh Goyal, Anamitra Roy Choudhury, Saurabh Raje, Venkatesan T. Chakaravarthy, Yogish Sabhar- wal, and Ashish Verma. 2020. Power-bert: Accel- crating BERT inference via progressive word-vector climination. In Proceedings of the 37th International Conference on Machine Learning, volume pages 3690-3699. PMLR.
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. 201 I. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Longuage Technologies, pages 142-150. Association for Computational Linguistics.
Yue Guan, Zhengyi Li. Jingwen Leng, Zhouhan Lin, and Minyi Guo. 2022. Transkimmer: Transformer learns to layer-wise skim. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics, pages 7275-7286. Association for Com- putational Linguistics.
Huizi Mao, Song Han, Jeff Pool, Wenshuo Li, Xingyu Liu, Yu Wang, and William J. Dally. 2017. Exploring the gramularity of sparsity in convolutional neural net- works. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Work- shops, pages 1927-1934. IEEE Computer Society.
Ari Holtzman, Jan Buys, Li Du. Maxwell Forbes, and Yejin Choi, 2020. The curious case of neural text degeneration. In Proceedings of the Eighth Inter- national Conference on Learning Representations. OpenReview.net.
Binny Mathew, Punyajoy Saha, Seid Muhie Yimam, Chris Biemann, Pawan Goyal, and Animesh Mukher- jee. 2021. Hatexplain: benchmark dataset for ex- plainable hate speech detection. In Proceedings of

---

the AAAI Conference on Artificial Intelligence, vol- ume 35, pages 14867-14875. AAAI Press.
Zhiqing Sun, Hongkun Yu, Xiaodan Song. Renjie Liu, Yiming Yang. and Denny Zhou. 2020. Mobilebert: a compact task-agnostic BERT for resource-limited de- vices. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2158 2170. Association for Computational Linguis- tics.
Clara Meister, Stefan Lazov, Isabelle Augenstein, and Ryan Cotterell. 2021. Is sparse attention more inter- pretable? In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing, pages 122- 129. Association for Computational Linguistics.
Yi Tay, Mostafa Dehghani. Dara Bahri, and Donald Metzler. 2023, Efficient transformers: A survey. In ACM Computing Survey, volume 55, pages 109:1- 109:28. Association for Computing Machinery.
Ali Modarressi, Hosein Mohebbi. and Moham- mad Taher Pilehvar. 2022. Adapler: Speeding up inference by adaptive length reduction. In Proceed- ings of the 60th Annual Meeting of the Association for Computational Linguistics, pages 1-15. Association for Computational Linguistics.
Ian Tenney, Dipanjan Das, and Ellic Pavlick. 2019. BERT rediscovers the classical NLP pipeline. In Proceedings of the 57th Annual Meeting of the Assø- ciation for Computational Linguistics, pages 4593- 4601. Association for Computational Linguistics.
Jun- Hyung Park, Yeachan Kim. Junho Kim, Joon- Young Choi, and SangKeun Lee. 2023. Dynamic structure pruning for compressing cnns. In Proceed- ings of the AAAI Conference on Artificial Intelligence, volume 37, pages
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Pro- cessing Systems, volume 30, pages 5998-6008. Cur- ran Associates, Inc.
AAAI Press.
Jason Phang, Haokun Liu, and Samuel R. Bowman. 2021. Fine-tuned transformers show clusters of sim- ilar representations acToss layers. In Proceedings of the Fourth BlackboxNLP Workshop on Analyzing and Interpreting Neural Networks for NLP, pages 529-538. Association for Computational Linguistics.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. 2019. GLUE: A multi-task benchmark and analysis plat- form for natural language understanding. In Pro- ceedings of the Seventh International Conference on Learning Representations. OpenReview.net.
Yao Qiang. Deng Pan, Chengyin Li, Xin Li, Rhongho Jang. and Dongxiao Zhu. 2022. Atteat: Explaining transformers via attentive class activation tokens. In Advances in Neural Information Processing Systems, volume 35, pages 5052-5064. Curran Associates, Inc.
Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed H. Chi. Quoe Le. and Denny Zhou. 2022. Chain-of-thought prompt- ing elicits reasoning in large language models. In Advances in Neural Information Processing Systems, volume 35, pages 24824-24837 24837. Curran Associates. Inc.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. 2016. Squad: 100, 000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natu- ral Language Processing, pages 2383-2392, Associ- ation for Computational Linguistics.
Adina Williams, Nikita Nangia, and Samuel R. Bow- man. 2018. A broad-coverage challenge corpus for sentence understanding through inference. In Pro- ceedings of the 2018 Conference of the North Amer- ican Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1112-1122. Association for Computational Linguis- tics.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. 2019. Distilbert, distilled version of BERT: smaller, faster, cheaper and lighter. CoRR, abs/1910.01108.
Ramprasaath R. Selvaraju, Michael Cogswell, Ab- hishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. 2017. Grad-cam: Visual explana- tions from deep networks via gradient-based localiza- tion. In Proceedings of the IEEE International Con- ference on Computer Vision, pages 618-626. IEEE Computer Society.
Xin, Raphael Tang, Jaejun Lee, Yaoliang Yu, and Jimmy Lin. 2020. DeeBERT: Dynamic early exit- ing for accelerating BERT inference. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2246-2251. Asso- ciation for Computational Linguistics.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Y. Ng. and Christopher Potts. 2013. Recursive deep models for semantic compositionality over sentiment tree- bank. In Proceedings of the 51st Annual Meeting of the Association for Computational Linguistics, pages 1631-1642. Association for Computational Linguis-
Canwen Xu, Wangchunshu Zhou, Tao Ge. Furu Wei, and Ming Zhou 2020. of Compressing BERT by progressive module replacing. In Proceed- ings of the 2020 Conference on Empirical Methods in Natural Language Processing, pages 7859-7869. Association for Computational Linguistics.
tics.

---

Deming Ye, Yankai Lin, Yufei Huang, and Maosong Sun. 2021. TR- BERT: dynamic token reduction for accelerating BERT inference. In Proceedings of the 2021 Conference of the North American Chapter of
Appendix
A Dataset statisties
the Association for Computational Linguistics: Hu- mun Language Technologies, NAACL-FILT 2021, On- line, June 6-11, pages 5798-5809. Association for Computational Linguistics.
We provide the statistics of the dataset in Table A.
Table 4; Statistics of the datasets used evaluations.
Xiang Zhang, Junbo Jake Zhao, and Yann LeCun. 2015. Character-level convolutional networks for text clas- sification. In Advances Neural Information Pro- cessing Systems, NeurlPS 2015, volume 28, pages 649-657. Curran Associates, Inc.
Average length
Number of train/test data
Number of classes
Dataset
SST-2
14
70k 1.8k
2
IMDB
275
25k 25k
2
Single Inpur
HateXplain
15.4k 1.9k
3
AG's news
120k 7.6k
4
Zhen Zhang, Wei Zhu, Jinfan Zhang, Peng Wang, Rize Jin, and Tae-Sun Chung. 2022. PCEE-BERE -EEE-BERT: Ac- celerating BERT inference via patient and confident early exiting. In Findings the Association for Com- putational Linguistics: NAACL 2022. pages 327-338. Association for Computational I Linguistics.
DBpedia
64
560k 70k
14
MRPC
3.6k 1.7k
Multiple
MNLI
Inpur
QNLI
50
104k / 5.4K
2
Wangchunshu Zhou, Canwen Xu, Tao Ge, Julian McAuley, Ke Xu, and Furu Wei. 2020. Bert loses patience: Fast and robust inference with early exit. Advances in Newral Information Processing Systems, volume 33, pages 18330 18341. Curran Associates, Inc.
B Selected
In table 5, we present the selected hyper-parameters for each dataset. When selecting the threshold p. we start to assign lower values ranging from 0.0 to 0.9 with the step size 0.05. The rationale be- hind the lowest-to-highest search is that the lowest threshold (even the threshold of 0.0) can still pro- vide efficiency Section 5.3) by transforming the supervision loss into the regularization loss, similar in previous works (Ye et al., 2021; Modarressi et al., 2022). Additionally, to prevent the router from leaping at the beginning of the training, we initialize the last layer of the routers to favor the non-leap action by setting large biases against the leap action
extent (as in the analysis of
Table 5: Hyper- prrpeepprameters LoT used each dataset.
Threshold
Balance term
Temperature
Dataset
X
SST-2
0.2
2.0
IMDB
0.5
1.0
HateXplain
1.0
AG's news
0.05
4.0
DBpedia
0.05
4.0
MRPC
0.4
1.5
MNLI
0.5
0.9
1.0
QNLI
0.9
C Additional Case Study of LoT
In Figure 7, we additionally provide the case study for the two datasets, AG's news and SST-2.

---

Layer Index
Layer lindax
laeer
dingood
hammme
asad
Booer
off
Iew
med
AG's
(b) SST-2
Figure 7: Illustration of the token routing on two examples. The darker block in each layer indicates the use of corresponding token in the layer while the lighter block denotes the leap action.
D Computational Overhead
BERT (Jiao et al., 2020)) and larger model (i.e.. BERT large) than the model used in the main paper. Table 6 shows the evaluation results of different scales SST-2 dataset. The results verify that the proposed method can boost the inference speed on the different scales of PLMs, demonstrating the scalability LoT.
Since LoT requires the dynamic token routers in the transformer, it imposes additional computation cost on our method. This is why we design the router to be lightweight module, which takes only 2% of the FLOPs from the entire model. Here, we analyze the trade-off between the capacity of the router and total speed-up. Specifically, we set the target performance as fixed and evaluate the total speedup gains with the varying capacity" of the router. Figure 8 shows the evaluation results for the trade-off. Notably, increasing the capacity of the router from 0.5% 2% leads the improved speedup both datasets. However, we observe that increasing the computation the routers to 6% ad- ditional FLOPss does not bring speedup gains. This result indicates the router requiring 2% additional FLOPs is enough to achieve reasonable efficiency.
Table Evaluation results of test accuracy (%) and speedup ratio on the SST-2 dataset.
Method
Accuracy
SpeedUp
BERT large
93.5
1.00x
BERT BE LoTT (ours)
93.1
2.13x
TinyBERT
89.7
1.00x
BinyBERTLLTT (ours)
89.5
2.22x
F Wall-clock Inference Time
SST
ONU
2.74x
To assess the speedup gains on specific computa- tional environments, we measured the inference time on single NVIDIA V100 GPU. As result, we observed that the real-time speedup gains (2.2x. Base: 37ms, LoT: 17ms) consist with the gains in FLOPPs (2.3x). This observation aligns with the pre- vious finding (Ye al., 2021), which suggest that theoretical speedups (i.e., FLOPs) often closely match the actual speedup gains.
2.12x
2.38
2.30
2.25
SpeedUp
2-
0.5%
2%
6%
Additional FLOPs from Router
Figure Speedup ratio on the different capacity of dynamic token routers. Note that the router consists of two linear layers.
LoT on Different Architectures
To verify the scalability of LoT, we performed the additional experiments on smaller model (i.e., Tiny-
For the capacity variation, we adjust the dimension of hidden layers of the routers.


Reviews:
[{"review_id": "zpQGAPlv4C", "reviewer": "Reviewer_9vaB", "summary": "", "questions": "- I'm curious about the background behind the design choice of having two layers in the Token Router. How are the dimensions of W1 and W2 determined in equation 1? (From what I gather, W2 appears to be (N, 2) to decide whether to leap). Additionally, is there a specific reason for including LayerNorm in between?\n- In line 247, what's the purpose of incorporating an additional context vector? The explanation in the main text seems to fall short in clarifying this choice.\n- While the benefits of token merging were validated through an ablation study, considering the concept of self-attention, I wonder if this is truly the best approach. When thinking about the self-attention mechanism which measures relative importance between tokens, is merging multiple token representations (even if they've leaped tokens) and including them as a single token in the attention operation truly the optimal way to utilize the information of the leaped tokens? I'd like to hear your thoughts on this.\n- As the hidden dimension of the target model increases, I anticipate the training overhead for LoT would also grow more substantial. I'm interested in information regarding the LoT training overhead in relation to varying encoder model sizes.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer 9vaB for valuable comments and pointing out the positive aspects of our work. We hope our response can address many of the suggestions and concerns.\n\n### [Diversity of the applied models]\nThank you for the insightful comments to improve our paper, especially concerning LoT's applicability. We initially used the BERT base to compare LoT with existing baselines that had evaluated their methods only on the BERT base model. Following the suggestions, we have performed additional experiments on the two suggested models, which are BERT Large and TinyBERT. The experimental results on SST-5 are as follows:\n| Method                   | Accuracy | SpeedUp |\n|--------------------------|:----------:|:---------:|\n| BERT(Large)              | 93.5     | 1.00x   |\n| BERT(Large) + LoT (ours) | 93.1     | 2.13x   |\n| TinyBERT                 | 89.7     | 1.00x   |\n| TinyBERT + LoT (ours)    | 89.5     | 2.22x   |\n\nWe have observed that applying LoT to the two models produces similar speedup to thBERT base + LoT. We believe these results will help readers better understand the essence of LoT. Therefore, we will incorporate these observations into the updated version of the paper, specifically in Section 5.\n\n### [Real inference speeds in specific computational environments]\nThank you for the suggestion regarding the practical efficiency. To assess the speedup gains on specific computational environments, we measured the inference time on a single NVIDIA V100 GPU. The average inference time on SST-5 as follows.\n\n| Method               | Avg. Inference time |\n|----------------------|:---------------------:|\n| BERT(base)           | 37 ms |\n| BERT(base) + LoT(ours) | 17 ms|\n\nWe observed that the real-time speedup gains (2.2x) consist with the gains in FLOPs (2.3x). This observation aligns with previous findings [1,2], which suggest that theoretical speedups (i.e., FLOPs) often closely match the actual speedup gains. We will include these actual speedup results for specific computational environments in Section 4 of the revised paper.\n\n### [Design choice about token router]\nThank you for the detailed attention on the proposed router. The number of layers and the dimension of the token routers was determined through the empirical search. Specifically, we adjust the dimension and the layers by carefully incrementing the capacity of the router, by considering that having more layers and a larger dimension can lead to the increased overheads you mentioned. While the optimal settings can vary based on the difficulty of the tasks, we observed that 2 layers with 256 hidden dimension (i.e., W1=(D, 256), W2=(256, 2) where D is the hidden dimension of the encoder model) generally works well without introducing significant overheads (i.e., 2.1% additional costs for the BERT base and 1.8% for the BERT large). Besides, regarding the inclusion of layer normalization, it was also chosen based on our empirical analysis, where we observed an improved stability and convergence speed during training when including the layer normalization. In the updated paper, we will include the background behind the design choice of the token routers.\n\n### [Utilizing context vectors in routing decision]\nThe context vector was introduced to provide the router with broader context information when deciding on the importance of tokens (we hypothesize that the token importance is determined by its context). Indeed, we empirically observed that making decisions with context vectors results in slightly better accuracy on a few tasks (e.g., removing context vectors in MNLI results in 0.3%p loss in accuracy with similar speedup gains). We will explain it more clearly in the updated version with empirical results.\n \n### [Design choice of merging tokens]\nWe thank the reviewer for detailed attention to the design choice in merging leaped tokens. In the initial version of LoT, the merging mechanism was indeed designed to compute a weighted average of leaped tokens based on their preceding cumulative self-attention scores. However, we found that the performance difference between the self-attention-based average and a simple average was negligible. And that We suspect that such outcome might be due to the fact that, in most cases, the leaped tokens tend to represent less informative portions of the sequence. As a result, their attention scores in relation to each other are relatively uniform, thereby resulting in the simple average during the merging process.\n \n### [Overhead in Token Router]\nSince the token router uses the hidden states as input, the increased dimension of the encoder model may yet require substantial overhead in the LoT. However, when considering the overhead as a fraction of the overall computations for the encoder, the relative overhead becomes rather smaller, given that the router's overhead increases linearly, whereas the overall computations increase exponentially due to the self-attention and positional feedforward networks. For instance, with a hidden dimension of 256 set for the router (as in our experiment), the computational overhead of the router in the BERT Large model (which has a 1,024 hidden dimension and 24 layers) constitutes only 1.8% of the total computations. This percentage is comparatively lower than the 2.1% overhead in the BERT base model (768 hidden dimension and 12 layers). We appreciate your insights regarding potential limitations tied to the varying sizes of the encoder model. We will include this discussion into the revised version of our paper (Limitation section).\n \n### [Typos]\nThank you for pointing out the typos. They will be corrected in the revised paper.\n\n### References\n[1] Kim et al., Learned token pruning for transformers, KDD 2022\n\n[2] Ye et al., Tr-bert: Dynamic token reduction for accelerating bert inference, NAACL 2021", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The paper is well-structured overall and provides a detailed representation of the proposed methodology, making it easy to follow. The inclusion of descriptions regarding prior research directions that are comparable is commendable, as it facilitated a better comparison of the efficiency of the proposed method.\n- By presenting additional analyses, the author provides a deeper understanding of how this method operates in actual natural language understanding scenarios.\n- The approach of continuously retaining tokens without discarding them, and making decisions during training on whether to leap over a particular token, is an innovative method compared to related work.", "reasons_to_reject": "- The diversity of the applied models is limited. It would be beneficial to see how the proposed method performs not only on BERT-base but also on encoders of different sizes (e.g., bigger model: BERT-large, RoBERTa, smaller model: TinyBERT, SkipBERT). Understanding the tendencies across various models (both bigger and smaller) would greatly aid in evaluating and grasping the essence of this methodology.\n- While the Speed Up using FLOPS is provided, having comparative data on inference speeds in specific computational environments would offer a clearer picture of the method's practical efficiency.", "questions_to_authors": "- I'm curious about the background behind the design choice of having two layers in the Token Router. How are the dimensions of W1 and W2 determined in equation 1? (From what I gather, W2 appears to be (N, 2) to decide whether to leap). Additionally, is there a specific reason for including LayerNorm in between?\n- In line 247, what's the purpose of incorporating an additional context vector? The explanation in the main text seems to fall short in clarifying this choice.\n- While the benefits of token merging were validated through an ablation study, considering the concept of self-attention, I wonder if this is truly the best approach. When thinking about the self-attention mechanism which measures relative importance between tokens, is merging multiple token representations (even if they've leaped tokens) and including them as a single token in the attention operation truly the optimal way to utilize the information of the leaped tokens? I'd like to hear your thoughts on this.\n- As the hidden dimension of the target model increases, I anticipate the training overhead for LoT would also grow more substantial. I'm interested in information regarding the LoT training overhead in relation to varying encoder model sizes.", "paper_topic_and_main_contributions": "This paper presents a novel methodology focused on dynamic token routing, a technique devised to identify redundant tokens in the process of sequence understanding inference. Compared to the approaches in prior research, this methodology shows increased flexibility in pinpointing unnecessary tokens, which subsequently enhances efficiency during the inference phase. \n\nThis method operates by retaining all input tokens, but under the guidance of a learned router, it adaptively utilizes necessary tokens. This not only allows for a more extensive saving of tokens required for inference but also yields a higher accuracy. An ablation study and a detailed analysis of this approach robustly demonstrate its efficacy.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "- L274 of of -> of\n- L433 Table -> Figure"}, "extra_text": {}}, {"review_id": "epKZAdHLAq", "reviewer": "Reviewer_iScN", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer iScN for valuable comments and helpful feedback on our work. We hope our response can address many of the suggestions and concerns.\n\n \n### [Additional analysis on the components and hyper-parameters of LoT]\nWe appreciate the Reviewer's suggestion for a more in-depth analysis of the components and hyper-parameters of LoT.\n- **Threshold p**: the threshold p represents the cumulative ratio of important tokens over the sequence. Hence the lower the p value forces the model to decide the greater number of tokens to be leaped to the next layer, leading to a higher speedup gain. In contrast, the higher the p value encourages the model to process the greater number of tokens in each transformer block, resulting in a lower speedup gain but higher performance. This hyper-parameter, therefore, controls the trade-off between speed and accuracy. In Section B of the Appendix, we have described the searching process of the hyper-parameter.\n- **Coefficient $\\lambda$**: this is the balancing coefficient between the task-specific loss (i.e., cross-entropy) and the token router loss. Its value is empirically determined for each specific task. In our experiments, we initially set a value for a specific threshold and gradually increment the lambda value until we observe a decline in task accuracy.\n- **Merging tokens and context vectors**: The ablation study concerning token merging is already presented in Table 3. As for the context vector, we will add the results (SST-2: 92.7 (2.30x), MRPC: 88.1 (3.11x)) to the same table in the updated paper, indicating a slight performance drop for both tasks. Thank you for bringing attention to our missing ablation study.\nWe will include this discussion into Section B of the Appendix, which has described the hyper-parameters in LoT.\n\n\n### [Evaluations on NLG tasks and LLMs]\nWe understand the importance of exploring our method across a broad spectrum of tasks. However, we would like to clarify that the application of LoT to NLG tasks and auto-regressive LLMs, especially for those typically relying on decoder models, is beyond the scope of this work. The nature of generation tasks involves a different set of challenges and would require significant modifications and additional processes to LoT (e.g., cascading attention mechanism, multiple use of token routers in the decoding process). We appreciate your suggestions and plan to investigate this in separate, focused future studies.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The proposed strategy allows the model to preserve all tokens in subsequent layers while reducing the number of tokens processed within each layer, which can mitigate the risk of losing crucial information and facilitate the exploration of diverse reduction space.\n2. The paper proposes a gradient-guided router training, which directly informs each router of which tokens are more influential for the prediction. This can steer the router towards making more informed and efficient decisions about whether to process or leap tokens, as well as provide some degree of interpretability for the model.\n3. Authors demonstrate the efficacy of LoT through extensive experiments and analysis on various natural language understanding benchmarks, covering both single input and multiple input tasks. The paper shows that LoT achieves substantial speedup gains over state-of-the-art methods with comparable task accuracy, indicating that LoT can be successfully applied to real-world applications that demand both accuracy and efficiency.\nIn summary, I recommend this paper should be accepted into Findings.", "reasons_to_reject": "1. The paper does not provide enough analysis or ablation study on the impact of different components or hyperparameters of LoT, such as the token merging mechanism, the context vector, the threshold p, and the harmony coefficient λ. It would be helpful to understand how these components or hyperparameters affect the performance and behavior of LoT.\n2. The paper does not evaluate LoT on natural language generation tasks or LLMs, which could also benefit from improved computational efficiency. It would be valuable to see how LoT performs on these tasks and whether it can preserve the quality and diversity of the generated outputs.", "questions_to_authors": "", "paper_topic_and_main_contributions": "This paper proposes a token reduction strategy for transformer-based models, which aims to improve the computational efficiency and preserve the crucial information for downstream tasks. The main contributions of this paper are: 1. introducing Leap-of-Thought (LoT), a method that enables the dynamic routing of tokens within the transformer layers. LoT learns to decide whether a given token should be processed in the current layer or leaped forward to the next layer, based on a lightweight token router module. 2. The paper proposes a gradient-guided router training, which leverages the gradient information of the token representations to provide supervision for the router. This helps the router to identify and process only the significant tokens for the task at hand. 3. The paper demonstrates the efficacy of LoT through extensive experiments and analysis on various natural language understanding benchmarks, showing that LoT achieves substantial speedup gains without a significant loss in task accuracy.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "pgbNujBXBV", "reviewer": "Reviewer_K4aR", "summary": "", "questions": "1. The Gradient-guided Router is used token-wise instead of token-wise and layer-wise(Equation 7,8 use CAT_i instead of CAT_i^l), is this the reason why the gradient-guidance better than regularization loss enforcement?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer K4aR for valuable comments and helpful feedback on our work. We hope our response can address many of the suggestions and concerns.\n\n### [Actual throughputs on specific hardware]\nThank you for the suggestion regarding the actual throughputs. We followed the evaluation metrics of the previous SOTA methods (i.e., AdapLeR) that evaluate their methods on FLOPs. To provide additional clarity on the practical efficiency, we measured the inference time on a single NVIDIA V100 GPU. The average inference time on SST-5 as follows.\n\n| Method               | Avg. Inference time |\n|----------------------|:---------------------:|\n| BERT(base)           | 37 ms |\n| BERT(base) + LoT(ours) | 17 ms|\n\nWe observed that the real-time speedup gains (2.2x) consist with the gains in FLOPs (2.3x). This observation aligns with previous findings [1,2], which suggest that theoretical speedups (i.e., FLOPs) often closely match the actual speedup gains. We will include these actual speedup results for specific computational environments in Section 4 of the revised paper.\n\n### [Different model size]\nThanks for the insightful comment to improve our paper regarding LoT’s applicability. We initially used the BERT base to compare LoT with existing baselines that had evaluated only this model. Following the suggestions, we have performed the additional experiments on the two different models, which are BERT Large and TinyBERT. The experimental results are as follows:\n\n| Method                   | Accuracy | SpeedUp |\n|--------------------------|:----------:|:---------:|\n| BERT(Large)              | 93.5     | 1.00x   |\n| BERT(Large) + LoT (ours) | 93.1     | 2.13x   |\n| TinyBERT                 | 89.7     | 1.00x   |\n| TinyBERT + LoT (ours)    | 89.5     | 2.22x   |\n\nWe have observed that applying LoT to two models produces the similar speedup to the BERT base + LoT. We believe these results will help readers better understand the essence of LoT. Therefore, we will incorporate these findings into the updated version of the paper, specifically in Section 5.\n\n### [Token-wise vs. Token-wise and Layer-wise for Gradient-guided Router]\nThank you for the attention on the mechanism of the token router. First, we would like to clarify that the router is used both token-wise and layer-wise, similar to positional feed-forward networks in the transformer block. However, as you noticed, the guidance is aggregated from all layers. The rationale behind this decision lies in previous findings on transformer interpretability, which suggest that aggregating gradient information from the entire layers can highlight the important tokens in the input sequence more effectively [3,4]. We believe that this carefully designed guidance produces better performance than the regularization loss, which lacks information about the significance of each token. In the revised version of our paper, we will provide a more detailed background on the rationale behind the token-wise guidance for the router.\n\n### References \n\n[1] Kim et al., Learned token pruning for transformers, KDD 2022\n\n[2] Ye et al., Tr-bert: Dynamic token reduction for accelerating bert inference, NAACL 2021\n \n[3] Qiang et al., Explaining transformers via attentive class activation tokens, NeurIPS 2022\n\n[4] Barkan et al., Grad-SAM: Explaining Transformers via Gradient Self-Attention Maps, CIKM 2021", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1.The gradient-guided router proposed in this paper shows superiority than regularization loss(Figure 5).\n2.The result is better than previous token reduction methods.\n", "reasons_to_reject": "1.The actual throughputs on specific hardware may need to be examined and compared.\n2.The paper only shows the results with BERT base, the performance and speedup with different model size are also considered important considering the model scaling up.\n\n", "questions_to_authors": "1. The Gradient-guided Router is used token-wise instead of token-wise and layer-wise(Equation 7,8 use CAT_i instead of CAT_i^l), is this the reason why the gradient-guidance better than regularization loss enforcement?", "paper_topic_and_main_contributions": "1.The paper introduces Leap-of-Thought(LoT), a novel token reduction approach to improve transformer computational efficiency, LoT enables dynamic token routing to reduce the processed token within each layer in transformer, and proposes a gradient-guided training to steer the dynamic router.\n2.Comprehensive evaluation results on SST-2, IMDB, HateXplain, AG’s news, DBpedia, MRPC, MNLI, QNLI show that LoT achieves substantial improvement on computational efficiency.\n", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers agree that the paper is sound and moderately exciting. The routing approach presented is novel and yields good results. The experiments, although limited to a set of few models, are relatively extensive and include all necessary ablations. I recommend acceptance to either the main conference of findings.
Decision: Accept-Main