Leap-of-Thought: Accelerating Transformers via Dynamic Token Routing (2023)

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

## introduction the advent of transformer @xcite has spurred a paradigm shift, most notably in natural language processing @xcite , but also extending to computer vision @xcite . however, the impressive capabilities of the transformer typically come with non-trivial computational costs, which scale quadratically with the length of the input sequence. this computational burden poses a significant challenge when deploying the transformer-based models in resourceconstrained or real-time systems @xcite . one typical approach to tackle this challenge is to reduce the number of tokens processed within transformer layers @xcite @xcite . the rationales behind this approach are two-fold: (i) not all tokens are equally significant to the task @xcite , and (ii) all token representations gradually become similar over layers @xcite . based on these rationales, previous studies progressively removes the less significant or redundant tokens (figure 1a ), selected either randomly @xcite or based on the attention scores @xcite . 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 performance by hindering the 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, it is sub-optimal to explore the diverse reduction strategies that potentially offer a greater efficiency. these limitations suggest that there is still room for further improvement. in this paper, we propose leap-of-thought (lot) foot_1 , a novel token reduction approach that enables the dynamic routing of tokens across layers. in contrast to permanent removal strategies, lot allows the tokens to 'leap' over each layer, thereby retaining access to all original tokens in the subsequent layers while reducing the number of tokens processed within each layer of the transformer (figure 1c ). we achieve this by coupling the transformer 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, we introduce a gradient-guided training that informs each router of which tokens do significantly 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 to achieve greater efficiency. lot offers several advantages compared to the permanent removal. primarily, lot has the potential to mitigate the risk of losing crucial information related to the task, given that the decisions for each token are recoverable in subsequent layers. in addition, lot provides a higher degree of freedom in token reduction, thereby facilitating the exploration of a diverse search space for greater efficiency, which is similarly observed in network compression @xcite . to substantiate the efficacy of lot, we perform evaluations across extensive experiments. comprehensive results demonstrate that the model employing lot reveals substantial speedup gains without a significant loss in task accuracy. additionally, through the analysis of lot, we provide justification for the efficacy of the dynamic token routing mechanism and illustrate how lot achieves greater efficiency. in summary, the contributions of the paper include the followings: ## related work in this section, we mainly review the methods that adaptively control the computation in pre-trained language models. recent approaches can be classified into two categories: width-wise and depthwise approaches. the former focuses on reducing the number of tokens processed by transformers, while the latter aims to decrease the number of computational layers. figure 1 illustrates the distinct behaviors of these approaches, including lot. ## experiment in this section, we 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 compared to the prior methods. ## analysis in this section, we analyze the behavior of lot in detail. we specifically focus on how lot achieves a greater efficiency than other baselines. ## conclusion in this work, we have proposed leap-of-thought (lot), a novel token reduction strategy that enables 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 or 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 guidance 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-art methods with the comparable task accuracy. the analysis also have strongly supported that introducing the leap action leads to the substantially improved efficiency.