Balcony: A Lightweight Approach to Dynamic Inference of Generative Language Models
作者: Benyamin Jamialahmadi, Parsa Kavehzadeh, Mehdi Rezagholizadeh, Parsa Farinneya, Hossein Rajabzadeh, Aref Jafari, Boxing Chen, Marzieh Tahaei
发布时间: 2025-03-10
来源: arxiv
研究方向: 大规模语言模型动态推理
主要内容
Balcony是一种轻量级的动态推理框架，旨在解决将大型语言模型（LLMs）部署到实际应用中的计算和延迟限制问题。它通过在预训练LLM中选择性插入额外的Transformer层来减少模型深度，从而在不牺牲性能的情况下实现实时适应不同的计算预算。
主要贡献
1. 引入了Balcony，这是一种基于深度的动态推理框架，它使用单个Transformer层在出口点，同时冻结基础模型。
2. 通过在小型数据集上使用自蒸馏损失高效地调整Balcony，与先前的方法相比，大大降低了训练成本，同时优于Flextron、LayerSkip和最先进的压缩方法。
3. 通过消融研究广泛评估了所提出框架，包括在1B参数LLM上的预训练。
研究方法
1. 冻结预训练LLM并在选定的出口点插入额外的Transformer层。
2. 使用自蒸馏损失训练Balcony层，以使Balcony层的输出与完整模型对齐。
3. 在训练过程中冻结基础模型参数和LM头，以保持基础模型性能并简化训练流程。
实验结果
Balcony在LLaMA3-8B模型上实现了最小的性能下降，同时实现了显著的加速。它优于Flextron、LayerSkip以及其他领先的压缩技术，在各种基准测试中表现优异。
未来工作
探索将Balcony扩展到更高效的架构，如MoE和Mamba等状态空间模型，以进一步提高性能和可扩展性。此外，Balcony可以用于自推测解码，以在不降低模型性能的情况下实现加速。