让多个模型协作生成时,共享中间状态,省下计算和存储开销。
Streamlining the Collaborative Chain of Models into A Single Forward Pass in Generation-Based Tasks
- 通过修改输入和注意力掩码,实现多模型间键值隐藏状态共享。
- 在四个任务上保持原有性能,推理速度提升显著。
- 适合需要高效多模型协作的生成类应用。
在检索增强生成(RAG)和基于代理的框架中,通常采用“模型链”方法,即多个专用模型按顺序处理不同子任务。该方法虽有效,但需独立部署各模型,资源开销大。近期研究尝试通过提示调优(prompt tuning)让共享基础模型以少量参数变化适应多任务。然而关键挑战在于:模型间传递的中间输出为纯文本,导致推理时需重新计算隐藏状态(如Transformer中的键值对(KV)状态)。本文提出FTHSS,一种新型提示调优方法,使模型可共享KV隐藏状态,避免冗余前向传播并减少KV缓存存储。通过训练时修改输入和注意力掩码,FTHSS使模型在单轮与多轮场景下均能有效利用前序模型的KV状态。四组任务的实验证明,FTHSS在保持传统模型链性能的同时显著提升推理效率。
原文摘要 · Abstract (English)
In Retrieval-Augmented Generation (RAG) and agent-based frameworks, the "Chain of Models" approach is widely used, where multiple specialized models work sequentially on distinct sub-tasks. This approach is effective but increases resource demands as each model must be deployed separately. Recent advancements attempt to address this by applying prompt tuning, which allows a shared base model to adapt to multiple tasks with minimal parameter changes. However, a key challenge remains: intermediate outputs, passed between models as plain text, require recomputation of hidden states (i.e., Key and Value (KV) states in Transformers) during inference. In this paper, we introduce FTHSS, a novel prompt-tuning method that enables models to share KV hidden states, eliminating redundant forward passes and reducing KV cache storage. By modifying input and attention masks during training, FTHSS allows models to effectively utilize KV hidden states from prior models in both single- and multi-round scenarios. Empirical results on four tasks show that FTHSS matches the performance of traditional model chains while improving inference efficiency.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。