Normalized Narrow Jump To Conclusions: Normalized Narrow Shortcuts for Parameter Efficient Early Exit Transformer Prediction (2024)

Abstract: With the size and cost of large transformer-based language models growing, recently, there has been interest in shortcut casting of early transformer hidden-representations to final-representations for cheaper model inference. In particular, shortcutting pre-trained transformers with linear transformations over early layers has been shown to improve precision in early inference. However, for large language models, even this becomes computationally expensive. In this work, we propose Narrow Jump to Conclusions (NJTC) and Normalized Narrow Jump to Conclusions (N-NJTC) - parameter efficient alternatives to standard linear shortcutting that reduces shortcut parameter count by over 97%. We show that N-NJTC reliably outperforms Identity shortcuts at early stages and offers stable precision from all transformer block levels for GPT-2-XL, Phi3-Mini and Llama2-7B transformer models, demonstrating the viability of more parameter efficient short-cutting approaches.

## introduction transformer based large language models @xcite stack blocks made up of multi-headed self-attention and feed forward layers sequentially. modern sophisticated language models stack upwards of 30 such blocks. for example, phi3-mini @xcite stacks 32 transformer blocks, gpt2-xl @xcite stacks 48 blocks, and deeper models like llama-2 70b @xcite and gpt-3 @xcite stack as many as 80 and 96 sequential blocks. however while such stacking typically improves model performance, it also increases the computational costs of inference. more gpu memory is required to store the additional stacked transformer blocks and more time is required to forward-pass inputs sequentially. there have been attempts to reduce the computational costs of such large language models by short-cutting transformers during inference. in short-cutting, one makes intermediate-predictions from an approximation of the final transformer output -that can be cheaply inferred from intermediate transformer-block outputs at each stage in the forward pass. a decision to 'early-exit' from the forward pass is made once the confidence of these intermediate-predictions reaches a certain pre-set confidence level λ. initial methods to short-cutting @xcite approximated the final transformer output directly by intermediate representations (identity shortcuts (id) : figure 1 ). more recently, @xcite proposed jump-to-conclusions (jtc) shortcuts that demonstrate that significant gains in this early-exit inference can be achieved by using a simple linear transformation over token representations to approximate the final transformer output. however, the jtc linear transformation proposed by @xcite adds an additional h x h parameters for each short-cut inference (for a transformer hidden dimension of h). for deep language models with large hidden dimensions, this too becomes very computationally expensive. for example, if we are to shortcut a phi3-mini model, we use a transformer hidden dimension of 3072. so each jtc shortcut roughly requires 9.43 million new parameters. with 32 transformer-blocks, having a shortcut inference option from each block requires that we train and store 9.43 * 31 > 292 million new parameters. for larger and deeper models like llama 2 70b, this number grows to over 5 billion new shortcut parameters. clearly, with the size and depth of large transformer models increasing, there is a need to develop more parameter efficient alternatives for short-cutting than the jtc method. independent of short-cutting methods, there has also been interest in matrix-decomposition for pre-trained model compression @xcite and success in low-rank fine-tuning of transformer weights @xcite . @xcite in particular has gained traction as a reliable method for parameter efficient fine-tuning -demonstrating that the weight update matrix can be decomposed into low-rank representations to reduce the total number of traninable parameters and reduce costs. however, these prior works focus on increasing efficiency within individual transformer blocks without skipping any block-computations. they do not consider any applications to transformer short-cutting -that skips multiple transformer block computations at a time, and is in itself a method for extreme efficiency. taking inspiration from (noach and goldberg, 2020) and @xcite , in this work we address the parameter inefficiency of jtc shortcuts @xcite . as mentioned in section 1, @xcite was the first to propose shortcuts for early exit transformer prediction. however they make the assumption that all transformer block outputs operate in the same space and use direct identity shortcuts for prediction. @xcite consider the output of different transformer blocks to operate in different representational spaces and recently demonstrated that linear transformation shortcuts significantly improve the performance of early-exit prediction. however, as discussed, they use full h × h linear matrices and are not very parameter efficient. independent of these works @xcite ; noach and goldberg, 2020) considered matrix-decomposition for pre-trained model compression and @xcite demonstrated that training low rank matrix decompositions of a weight update matrix approximate good fine-tuning results for transformers with parameter-efficiency. however as mentioned earlier these methods focus on efficiency within fixed transformer-blocks and do not consider any applications to transformer short-cutting -that skips computations of entire blocks at a time, and is itself a method for extreme efficiency. ## method 3.1 narrow jump to conclusions (njtc) given a transformer model with hidden dimension h, to approximate a short-cut between its block-outputs at any two levels l and m, given a set of n input sentences s, we forward pass each sentence s @xmath0, h m i )} n @xmath1. we then fit a simple 2 layer linear neural network made up of matrices a : h × h 100 and b : h 100 × h that takes as input h l i and approximates ĥm i . we note in particular that while other model informed-choices for low-rank short-cutting dimensions may be possible, in eq.1 to standardize our approach to diverse transformers with potentially different hidden dimensions, we consider a fixed low-rank reduction to 1% of the transformer hidden dimension size ( h 100 ). we fit the two matrices a and b jointly using gradient descent to minimize the mean squared error loss l lm between the approximated representations ĥm i and the transformer block outputs h m i . the hidden representation of each token in a sentence at level l is passed through a and b to obtain approximations of the hidden representations at level m. as a result of this low rank matrix decomposition, each njtc shortcut uses 2 * (h × h 100 ) = 0.02h 2 parameters: only 2% the number of parameters of a jtc shortcut.