Scholarlink

Autoregressive large language models (LLMs) suffer from high computational cost due to their large parameter size and the token-by-token generation process. Token-level early exit strategies aim to alleviate this issue by dynamically reducing the number of layers used for each token during decoding. However, existing early exit methods exhibit instability, leading to token degradation and suboptimal performance. In this work, we systematically analyze the instability in token-level early exit and identify two key issues: (1) KV cache loss, where deeper-layer tokens lack key-value (KV) cache entries from earlier shallow-exited tokens, causing degradation in later token generations; and (2) non-robust exit layer selection, where deviations from the optimal exit layer lead to performance drops. To address the first issue, we introduce MimicNet, a lightweight network based on layer distillation that approximates hidden state transformations of a full layer, leveraging residual connections to mitigate degradation. For the second issue, we apply SparseAlignment for Intermediate Layers (SAIL) to enhance the robustness of exit layer selection while maintaining efficiency. We conduct experiments on Llama-2-7B, Llama-2-13B, and Llama-3-8B. Our results demonstrate that our approach improves the robustness of early exit while reducing computational cost by approximately 30%, with minimal impact on model performance.

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

Robust and Efficient Early Exit for Large Language Models: Mitigating KV Cache Loss and Enhancing Exit Stability

  • 会议论文

  • 发表日期：

    2026年

查看全文
  • Long Meng,
  • Ruiqing Zhang,
  • Weiqiao Shan

摘要

Autoregressive large language models (LLMs) suffer from high computational cost due to their large parameter size and the token-by-token generation process. Token-level early exit strategies aim to alleviate this issue by dynamically reducing the number of layers used for each token during decoding. However, existing early exit methods exhibit instability, leading to token degradation and suboptimal performance. In this work, we systematically analyze the instability in token-level early exit and identify two key issues: (1) KV cache loss, where deeper-layer tokens lack key-value (KV) cache entries from earlier shallow-exited tokens, causing degradation in later token generations; and (2) non-robust exit layer selection, where deviations from the optimal exit layer lead to performance drops. To address the first issue, we introduce MimicNet, a lightweight network based on layer distillation that approximates hidden state transformations of a full layer, leveraging residual connections to mitigate degradation. For the second issue, we apply SparseAlignment for Intermediate Layers (SAIL) to enhance the robustness of exit layer selection while maintaining efficiency. We conduct experiments on Llama-2-7B, Llama-2-13B, and Llama-3-8B. Our results demonstrate that our approach improves the robustness of early exit while reducing computational cost by approximately 30%, with minimal impact on model performance.