F i RST : Finetuning Router-Selective Transformers for Input-Adaptive Latency Reduction (2025)

Abstract: Auto-regressive Large Language Models (LLMs) demonstrate remarkable performance across different domains such as vision and language tasks. However, due to sequential processing through multiple transformer layers, autoregressive decoding faces significant computational challenges, particularly in resource-constrained environments like mobile and edge devices. Existing approaches in literature that aim to improve latency via skipping layers have two distinct flavors: (1) early exit, and (2) input-agnostic heuristics where tokens exit at pre-determined layers irrespective of input sequence. Both the above strategies have limitations, the former cannot be applied in the presence of KV caching, which is essential for speed-ups in modern inference frameworks, and the latter fails to capture variation in layer importance across tasks or, more generally, across input sequences. To address these limitations, we propose FiRST, a model-agnostic framework that reduces inference latency by using layer-specific routers to adaptively skip transformer layers during decoding, based on routing decisions made from the input prompt in the prefill stage. FiRST remains fully compatible with KV caching, enabling faster decoding while maintaining quality. Our method reveals that input adaptivity is essential: Different tasks rely on different subsets of layers to evolve meaningful representations. Extensive experiments show that FiRST significantly reduces latency while outperforming existing layer selection strategies in quality. It retains performance comparable to the base model without skipping. FiRST is thus a promising and efficient solution for LLM deployment in low-resource environments.

## problem statement our goal is to exploit the heterogeneity in inputs and tasks to selectively use llm layers in a qualityaware manner for reducing inference latency and compute for on-device constraints. ideally, we want to select an optimal sub-sequence of layers within a transformer architecture for a given input and task, such that the overall latency, as well as expended computation, are both low, while quality is comparable to the un-modified case where every input sequence passes through every layer. for ease of explanation, without loss of generality, we assume the task is same and simply consider an input sequence for describing the problem. let us consider an input sequence @xmath0, x 2 , . . . , x n } with n tokens. let there be m transformer layers in the model, where the i th transformer layer is represented as the function ϕ i (). as stated lucidly in @xcite , x is first converted to an initial latent representation h 0 = {h 1 0 , h 2 0 , . . . , h n 0 }, where h 0 @xmath1, ∀@xmath2. thereafter, every transformer layer ϕ i () operates on the latent vectors h i to generate the embedding for the i th layer as follows. for the j th token, the embedding at layer i is computed as follows: the optimality above is with respect to the minimum subsequence of layers that can help achieve the above, to minimize latency while keeping quality unaffected. note that, the optimal subsequence f op t (x ) need to obey the same autoregressive computation on previous tokens as given in equation 1. hence, any algorithm that determines the optimal subsequence, need to be compatible with kv cache handling, to avoid the re-computation of values for tokens preceding the current token. the potential number of subsequences for m layers is 2 m , hence a brute force is infeasible and also beats the purpose of such a layer selection in the first place: reducing latency and compute. in the absence of any known substructure in the behaviour of the latent layers on each input sequence, it is difficult to arrive at the optimal solution polynomially or with low additional latency or compute. we propose to learn an approximation of the optimal subsequence of layers for any input sequence with low additional latency and minimal training. ## proposed solution: first let us first understand what it entails to learn an optimal subsequence of layers for any input. consider the full transformer sequence to be f * = {ϕ 1 , ϕ 2 , . . . , ϕ m }. any optimal subsequence for an input x : f op t,x could be thought of as finding an optimal path through a binary tree of functions. formally, let every level in the binary tree correspond to a transformer layer and the 0 th layer corresponds to the initial embedding look up; i.e., at depth @xmath3, there would be 2 i nodes, each corresponding to either ϕ i or ϕ i , where the former denotes that a particular transformer layer is included in the optimal path whereas the latter denotes that it is not included. each (of the 2 i-1 nodes) ϕ i or ϕ i has two children, corresponding to the next transformer layer: ϕ i+1 and ϕ i+1 (see figure 1 ). in such a tree structure, for example, the path {ϕ i , ϕ i+1 , ϕ i+2 } indicates the subsequence of transformer layers {ϕ i , ϕ i+2 }. for any transformer layer ϕ i in this tree, let anc(ϕ i ) = k, 0 ≤ @xmath4. in the above example, anc(ϕ i+2 ) = ϕ i . consider a sequence of functions f, where for level i, anc(ϕ i ) = ϕ k . the autoregressive computations for the j th token in the input sequence (originally eq 1), would now be modified as: our problem translates to navigating this binary tree to find the optimal path f op t for an input sequence and task. since there are 2 m paths in this tree, we propose to approximate the optimal by making a decision in a greedy fashion at each node. formally, we add a (lightweight and fast) router r i before every transformer layer ϕ i in the model, that will predict whether ϕ i will be selected or not. our aim is to learn to predict the layer choice at a sequence level (not token) to maintain compatibility with the autoregressive computations and avoid re-computation of of kv cache values. moreover, we should spend minimal compute for learning the r i functions. finally, r i functions should not add any significant latency to the overall computation. first modifies any off-shelf pre-trained transformer based model by incorporating and training a router or probability function r i before every transformer layer ϕ i . the output of r i is a score ρ i denoting the probability of selecting ϕ i in the layer sequence. during inference, ρ i is rounded to determine selection of ϕ i . let ⌊ρ i ⌉ = 1 if ρ @xmath5.5, else 0. equation 1 is now modified as: this recursively approximates eq 3 for the optimal f in a probabilistic, greedy manner. we train the functions r i on datasets and tasks, and further fine tune using lora adapters to make the layer selections smooth and improve the output quality. ## first framework and algorithm in this section, we describe the training and inference frameworks for first in details. we discuss how to train routers to be adaptive to input sequences. given an off-the-shelf pre-trained llm, we propose two training phases. in the first phase, we train a router for each layer that decides whether the input sequence should skip the layer. in the second phase, to tackle the issue of unseen skipping during pre-training, we fine-tune the routeraugmented llm keeping router weights fixed to ensure the model improves performance on the target dataset without reducing the skipping level. while joint training of the router and lora modules is theoretically possible, we find it introduces optimization instability (see appendix a.4 for further discussion on this design choice and its impact). ## experiments we conduct experiments on three benchmark tasks: machine translation, text summarization, and question answering, demonstrating the robustness and scalability of first across diverse settings. we base our task selection on prior work in the field to ensure a fair and meaningful comparison @xcite @xcite . datasets: for machine translation, we use wmt development sets (2017-2020) for english-to-chinese and english-to-german tasks, evaluating performance on the wmt 2022 test set, which covers diverse domains such as news, social media, e-commerce, and conversational contexts. for summarization, we use the cnn/dailymail dataset, with 4,000 randomly selected training samples and evaluation on the standard test set of 11,490 samples. for question answering (qa), we utilize squad v1.1 and natural questions (nq). we train on 4,000 randomly selected samples from each dataset and evaluate on their respective validation sets as test set labels are unavailable. for nq, we incorporate a retrieval step before answer generation, retrieving relevant passages as context. appendix a.2 contains detailed descriptions. evaluation metrics: we employ standard metrics to assess output quality across three distinct tasks. for machine translation, we benchmark performance using bleu and comet; comet is included for its more nuanced assessment capabilities beyond the n-gram overlap measured by bleu. summarization quality is evaluated with rouge scores and bert score, the latter capturing meaning-based similarity. for question answering, we use exact match (em) and f1 score on the squad dataset, and report exact match (em) for natural questions (nq). finally, to benchmark latency, we measure time per output token (tpot) on gpu, which gauges overall decoding performance. detailed descriptions of all evaluation metrics are available in appendix a.3. hyperparameters used during training and inference are provided in appendix a.4. ## conclusion we propose a new framework, first, for layer selection that adapts to the input sequence and task, aiming to reduce latency in a quality-aware manner. this approach is sequence-dependent and operates compatibly with kv caching. with an optimal layer skipping rate of around 15%, first achieves a 10-20% reduction in latency. this speedup is achieved while remaining quality-neutral, maintaining approximately 80% or more of the base model's performance on quality metrics across multiple tasks (machine translation, summarization, question answering) and model architectures on well-known open-source datasets. furthermore, our method significantly outperforms other layer selection strategies on most quality metrics.