FFN - S kip LLM : A Hidden Gem for Autoregressive Decoding with Adaptive Feed Forward Skipping (2024)

Abstract: Autoregressive Large Language Models (e.g., LLaMa, GPTs) are omnipresent achieving remarkable success in language understanding and generation. However, such impressive capability typically comes with a substantial model size, which presents significant challenges for autoregressive token-by-token generation. To mitigate computation overload incurred during generation, several early-exit and layer-dropping strategies have been proposed. Despite some promising success due to the redundancy across LLMs layers on metrics like Rough-L/BLUE, our careful knowledge-intensive evaluation unveils issues such as generation collapse, hallucination, and noticeable performance drop even at the trivial exit ratio of ~10-15% of layers. We attribute these errors primarily to ineffective handling of the KV cache through state copying during early exit. In this work, we observe the saturation of computationally expensive feed-forward blocks of LLM layers and propose FFN-SkipLLM, which is a novel fine-grained skip strategy for autoregressive LLMs. FFN-SkipLLM leverages an input-adaptive feed-forward skipping approach that can skip ~25-30% of FFN blocks of LLMs with marginal change in performance on knowledge-intensive generation tasks without any requirement to handle the KV cache. Our extensive experiments and ablation studies across benchmarks like MT-Bench, Factoid-QA, and variable-length text summarization illustrate how our simple and easy-to-use method can facilitate faster autoregressive decoding.

## introduction autoregressive large language models (llms) have been recently show-stealers, profoundly influencing not only the landscape of nlp @xcite @xcite @xcite @xcite , but also recently buttressing numerous computer vision @xcite @xcite @xcite and graph neural networks @xcite @xcite @xcite algorithms, achieving stellar performance across various task benchmarks. however, their widespread adoption is hindered by their massive scale, characterized by billions of parameters, which demand exceedingly high computational resources and memory capacities. for instance, the gpt-175b model necessitates 325 gb of gpu memory for loading its weights and relies on a minimum of five a100 (80gb) gpus employing sophisticated parallelism techniques @xcite . this imposing computational and memory requirement presents a challenge to the broader accessibility of these models. to alleviate the demanding hardware requirements for deploying massive trained models, considerable efforts have been devoted to mitigating their high computational inference cost resulting from token-by-token generation. among several model compression techniques such as quantization @xcite @xcite @xcite , and sparse neural networks @xcite @xcite @xcite @xcite which require additional hardware support for speedup, token-level early exit or layer-skip has emerged as a promising technique to alleviate these limitations by allowing tokens to cease computation as soon as their hidden states reach saturation @xcite @xcite . these methods exploit existing redundancy across llms' layers which can be ignored during tokenby-token generation significantly saving massive computation involved within a layer (e.g., ∼ 200-300 million parameters in a single llama layer). although the proposed methods have shown some promising success, their performance is widely restricted by the issue of inappropriately handling kv caching. kv caching saves keys and values of all attention layers for previously generated tokens and accelerates sequence generation by reducing redundant computation (though at the cost of higher memory usage). given a token generated via early exiting, its kv caches in subsequent layers are incomplete which impedes the generation of future tokens beyond the exiting layer of the current token. for handling the kv cache issue, some recent works @xcite @xcite @xcite propose three main solutions: copying hidden states, pre-fixed token-level skip pattern, and kv recomputation. despite these mitigation methods, our careful knowledge-intensive investigation reveals that layer-skipping induces permanent damage due to deviation from the inference process that the model is trained to excel at, leading to significant hallucination of wrong facts and token generation collapse. figure 1 shows the comparison of the responses generated by two recent layer skipping methods, namely @xcite and shortgpt @xcite for a knowledge-intensive qa example. in their response, both shortgpt and skipdeocde fail to generate the correct answer "narendra modi", suffer from token collapse, and hallucinate misinformation. in this work, we ask an interesting unexplored question: instead of attempting to fix the kv cache, can we completely circumvent the kv cache bottleneck of layer-skipping and still avoid unnecessary computational expenses while mitigating hallucination and token generation collapse? to this end, our work is the first attempt to investigate a fine-grained layer-skipping strategy that focuses on computationally expensive feed-forward network (ffn) blocks in llms. ## layer-skipping: an knowledge-intensive evaluation recent advancements in autoregressive models @xcite @xcite have revolutionized the quality of language generation in various generative tasks, including question answering @xcite , summarization @xcite , and machine translation @xcite . however, these large transformer models face challenges in terms of high inference latency attributed to their numerous layers and the autoregressive decoding process. the sequential computation of multiple stacks of transformer layers for each token during the inference stage imposes significant computational overheads, thus limiting their real-time adaptability. to counter the computational cost of token-bytoken generation with modern gigantic llms, several works @xcite @xcite @xcite have been recently exploring token-level early exit and layer-skipping (depth-pruning) strategies. the primary challenge associated with these approaches is that if the current token exits at a higher layer, there arises a need to recalculate the key-value (kv) caches for preceding tokens. to this end, three major approaches have been explored: (1) copy the hidden states of the current token at the exiting layer to all later layers, which will be used to compute the keys and values at later attention layers, (2) pre-specify the exiting layer for each token, while ensuring that kv missing in previous tokens will not hinder the generation of later tokens; with this approach, the ability of token-wise adaptive selection of exits is inevitably lost, (3) kv recomputation which is a variant of synchronized parallel decoding and adds additional computational and memory overhead. despite some notable performance gains over some metrics (e.g., perplexity, rough-l, blue), our careful knowledge-intensive investigation reveals that the kv cache problem during layer-skip is not effectively addressed. figure 1 illustrates the responses generated by two recent layer-skipping methods skipdecode @xcite and @xcite for a given factoid-based qa task which requires answering using relevant entities and attributes ingested within llms during pre-training. interestingly, answers generated by the skipdecode agent hallucinate misinformation claiming '... does not have a prime minister ... india abolished its cabinet posts ... ' while the shortgpt agent fails to generate any factoid to answer the question. note that both agents suffer from token collapse and start generating repetitive ## experimental results baseline details: to empirically evaluate the performance gains enabled by our proposed ffn-skipllm across multiple knowledge-intensive tasks. we aim to investigate how well fnn block skipping can retain the ability to access factoid answers ingested during pretraining, perform multiturn instruction following, and in-context summarization. our baselines are: 1 full model which indicate the maximum capability of llm under consideration; 2 random skip where ffn-blocks are dropped at random without giving careful consideration of cold and non-cold regions; 3 no input adaptive where we do not track the cosine similarity per token and ffn-blocks are dropped at random from the non-cold region. our baselines are constructed to carefully validate the effect of our observations in ffn-skipllm. ## conclusion in this paper, we explore an orthogonal dimension for layer-skipping and early-exit strategies that suffer from kv cache issues leading to the hallucination of misinformation and token collapse. we propose ffn-skipllm, a novel fine-grained skip strategy of autoregressive llms which is an input-adaptive feed-forward skipping strategy that can skip ∼ 25-30% of ffn blocks of llms with marginal change in performance on knowledgeintensive tasks. fnn-skip llm is built on the core observation of monotonically increasing redundancy within the ffn blocks of llms. our future work includes exploring parameter-efficient continual fine-tuning techniques to push the performance of ffn-skipllm for high skip ratios.