Scholarlink

<p>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 <i>Memory Layer</i> 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.</p>

错误：搜索内容不能为空，请输入英文关键词
错误：关键词超出字数限制，请精简
高级检索

A training-inference consistent framework for early exiting in language models with parallel decoding

  • 期刊论文

  • 发表日期：

    2026年2月16日

查看全文
  • Ziqian Zeng,
  • Zelin Chen,
  • Huiping Zhuang

摘要

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.