A Training-Inference Consistent Framework for Early Exiting in Language Models with Parallel Decoding

With the growing complexity of Large Language Models (LLMs), there is increasing interest in efficient inference architectures. Early exiting is a popular method that improves inference efficiency by skipping layers and generating output once the model reaches a certain confidence level. Traditional early exiting methods incorporate weighted cross-entropy loss during training, ensuring accurate predictions across all internal classifiers. However, only one correct prediction is needed to speed up the processing during inference. Additionally, early exiting in LLMs also faces challenges with KV cache updates in autoregressive decoding. Current solutions approximate the KV cache of skipped layers by copying hidden states, resulting in reduced accuracy. In response, we propose ConsistentEE, an enhanced early exiting framework that maintains consistency between training and inference. Our method frames early exiting as a reinforcement learning challenge, employing policy networks to ascertain when to exit. We also introduce Memory Layer to evaluate the hardness of instances. To tackle KV cache challenges in LLMs with early exiting, we employ a parallel decoding strategy. Experimental results show that our approach consistently outperforms the baseline approach on multiple natural language classification and generation tasks.