Balcony:生成式语言模型动态推理的轻量级方法
计算与语言
2025-03-12 v2 人工智能
机器学习
摘要
在真实应用中部署大型语言模型(LLM)常常受到严格的计算和延迟约束。虽然动态推理提供了根据不同资源预算调整模型行为的灵活性,但现有方法经常受限于硬件效率低下或性能退化。在本文中,我们提出了 Balcony,一个简单而高效的基于深度的动态推理框架。通过冻结预训练的 LLM 并在选定的退出点插入额外的 transformer 层,Balcony 在保持完整模型性能的同时,实现了对不同计算预算的实时适应。这些额外层通过简单的自蒸馏损失进行训练,使子模型输出与完整模型输出对齐。与先前方法相比,该方法需要的训练令牌和可调参数显著减少,从而大幅降低计算成本。将 Balcony 应用于 LLaMA3-8B 模型时,仅使用原始预训练数据的 0.2%,Balcony 在性能退化极小的情况下实现了显著加速。值得注意的是,我们证明 Balcony 在多个模型和多种规模上的多个基准测试中,优于 Flextron 和 Layerskip 等最先进方法以及其他主流压缩技术。
引用
@article{arxiv.2503.05005,
title = {Balcony: A Lightweight Approach to Dynamic Inference of Generative Language Models},
author = {Benyamin Jamialahmadi and Parsa Kavehzadeh and Mehdi Rezagholizadeh and Parsa Farinneya and Hossein Rajabzadeh and Aref Jafari and Boxing Chen and Marzieh S. Tahaei},
journal= {arXiv preprint arXiv:2503.05005},
year = {2025}
}