arXiv:2505.22135cs.CLcs.LG2025-05被引 9

用自预测解码识别冗余层,提升混合模型效率与性能

RAD: Redundancy-Aware Distillation for Hybrid Models via Self-Speculative Decoding

  • 通过自预测解码定位Transformer中的冗余注意力层
  • 用状态空间模型替换冗余层,实现2倍更快收敛
  • 适合追求高效推理的模型压缩与优化研究者

结合Transformer与状态空间模型(SSMs)的混合模型在性能与效率间具有潜力。然而,如何优化这类模型,特别是解决Transformer组件中固有的冗余问题,仍是挑战。本文提出RAD(冗余感知蒸馏),利用自预测解码作为诊断工具,识别模型中冗余的注意力层,并将其选择性替换为SSM组件,随后进行针对性(自)蒸馏。RAD聚焦于冗余组件的知识迁移,考虑架构变化与特定权重初始化策略。实验表明,使用RAD的自蒸馏显著优于原始基线模型,在数学与编程任务上表现更优。此外,在标准蒸馏设置下,RAD相比基线方法实现约2倍加速收敛。例如,使用较小的Llama-3.1 8B教师模型时,RAD在GSM8K上达到71.27分,CRUX上达28.25分,远超基线模型(46.17和22.75分)所用的70B教师模型。RAD为混合模型的蒸馏优化提供了新路径。

原文摘要 · Abstract (English)

Hybrid models combining Transformers and State Space Models (SSMs) are promising for balancing performance and efficiency. However, optimizing these hybrid models, particularly by addressing the potential redundancy inherent within the Transformer components, remains a significant challenge. In this paper, we propose RAD (Redundancy-Aware Distillation), a novel framework that uses self-speculative decoding as a diagnostic tool to identify redundant attention layers within the model. These identified layers are then selectively replaced with SSM components, followed by targeted (self-)distillation. Specifically, RAD focuses knowledge transfer on the components identified as redundant, considering architectural changes and specific weight initialization strategies. We experimentally demonstrate that self-distillation using RAD significantly surpasses the performance of the original base model on mathematical and coding tasks. Furthermore, RAD is also effective in standard knowledge distillation settings, achieving up to approximately 2x faster convergence compared to baseline methods. Notably, while a baseline model distilled from a Llama-3.1 70B teacher achieves scores of 46.17 on GSM8K and 22.75 on CRUX, RAD achieves significantly higher scores of 71.27 on GSM8K and 28.25 on CRUX, even when using a much smaller Llama-3.1 8B teacher. RAD offers a new pathway for efficient optimization and performance enhancement in the distillation of hybrid models.

模型蒸馏混合架构效率优化自预测解码

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。