Streamlining the Collaborative Chain of Models into A Single Forward Pass in Generation-Based Tasks (2025)

Abstract: 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.

## introduction in many retrieval-augmented generation (rag) and agent-based frameworks @xcite , multiple large language models (llms) often collaborate sequentially. each model focuses on a specific sub-task and passes its output as input to the next model until the task is completed @xcite . for instance, some rag post-retrieval optimization methods @xcite involve summarizing retrieved documents with a summarization model, and then generating answers with a question-answering model. these stepwise approaches leverage the strengths of individual models and have proven effective in many scenarios. as a result, the "chain of models" approach has gained popularity @xcite . deploying every specialized llm in such chains significantly increases the resources needed. to address this, researchers have explored parameterefficient fine-tuning (peft) methods, such as prompt tuning @xcite and @xcite . these techniques allow fine-tuning with a fraction of the parameters when training. during inference, a shared base model is deployed on a single device and handles multiple tasks with distinct parameter configurations. this approach merges "chain of models" workflows into a single architecture, adapting to various sub-tasks through selective parameter usage. however, a critical bottleneck remains: in the chain, the intermediate keyvalue (kv) hidden states from one model cannot be directly reused by the next model due to parameter differences. as a result, communication between models in the chain relies on passing plain text, forcing the downstream model to recompute hidden states. this practice not only adds computational overhead, but also raises kv cache storage requirements for each model in the chain, further hampering efficiency. in this paper, we argue that such recomputation is unnecessary. even with parameter differences, the kv hidden states produced by one model should only differ marginally from those recalculated by the next. particularly in prompt-tuning methods, the kv hidden states produced by the previous model are essentially conditioned on a few noisy tokens. with appropriate fine-tuning, the subsequent model can effectively interpret and utilize the kv hidden states of the previous model despite these noises, as figure 1 shows. to realize this vision, we propose fthss(fine-tuning for hidden state sharing), a prompt-tuningbased method that enables models in a chain to share kv hidden states. specifically, when finetuning the model in single-round scenarios, where each model is invoked only once, we use kv hidden states from the prior models as input rather than plain text. this training approach requires extensive storage and access to kv hidden states, which may potentially increase training time and storage demands. to mitigate this, we introduce an online optimization strategy. by modifying the input and attention mask for each layer, we recompute the prior model's kv hidden states in memory during training, thus avoiding the overhead of storage and access. in multi-round scenarios, where models in the chain are invoked repeatedly, each model must adapt to the kv hidden states of others, so all models in the chain are trained synchronously to ensure mutual adaptation. after fine-tuning, models can dynamically switch learnable prompt tokens during inference, adapting based on task requirements, while leveraging precomputed kv cache for direct generation. and since prompt-tuning-based methods enable the deployment of multiple models on a single device, communication overhead for hidden states is effectively eliminated. empirical results on four tasks, including singleround and multi-round, demonstrate that fthss leads to a comparable performance to the chain of models, while enhancing inference efficiency. technical contributions of this paper can be summarized as follows: ## methodology in this section, we begin by highlighting a key challenge: model chains rely on text-based communication, which prevents the direct transfer of kv hidden states between models. we then explore the feasibility of fine-tuning the downstream model to process kv hidden states from the upstream model, although these hidden states often include noise tokens irrelevant to the downstream task. lastly, we propose training strategies, fthss, to achieve kv hidden state sharing. these strategies include modifying inputs and attention masks to support both single-and multi-round scenarios. ## conclusion in this paper, we introduced fthss, a method that enables models in a chain to directly share kv hidden states, eliminating redundant forward passes over intermediate results and reducing kv cache storage. by reordering the input and attention masks at each layer, fthss allows downstream models to leverage kv hidden states from upstream models. our experiments demonstrate that fthss matches the performance of traditional model chains while significantly improving the inference efficiency in both single-round and multiround scenarios.