Collab: Controlled Decoding using Mixture of Agents for LLM Alignment
作者: Souradip Chakraborty, Sujay Bhatt, Udari Madhushani Sehwag, Soumya Suvra Ghosal, Jiahao Qiu, Mengdi Wang, Dinesh Manocha, Furong Huang, Alec Koppel, Sumitra Ganesh
分类: cs.CL, cs.AI
发布日期: 2025-03-27
备注: Accepted to ICLR 2025
💡 一句话要点
Collab：一种基于混合Agent的受控解码方法，用于LLM对齐。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型对齐 受控解码 混合Agent 推理时对齐 强化学习 策略切换 模型协作
📋 核心要点
- 现有RLHF对齐LLM计算成本高昂，单Agent解码难以适应多样任务。
- 提出Collab，一种混合Agent解码策略，通过token级选择动态切换LLM。
- 实验表明，Collab在奖励和胜率上均显著优于现有解码策略。
📝 摘要（中文）
大型语言模型（LLM）的对齐对于在应用中安全可靠的部署至关重要。从人类反馈中进行强化学习（RLHF）已经成为一种有效的技术，可以将LLM与人类偏好和更广泛的效用对齐，但这需要更新数十亿的模型参数，计算成本很高。相比之下，受控解码提供了一种在推理时对齐模型而无需重新训练的机制。然而，由于任务固有的复杂性和可变性，单Agent解码方法通常难以适应不同的任务。为了加强测试时针对目标任务的性能，我们提出了一种基于混合Agent的解码策略，利用现有的现成对齐LLM策略。将每个先验策略视为Agent协作中的一个Agent，我们开发了一种解码方法，该方法允许通过多个Agent之间的token级选择策略进行推理时对齐。对于每个token，根据长期效用指标从模型池中动态选择最合适的LLM。这种策略切换机制确保了每一步的最佳模型选择，从而在解码过程中实现LLM之间的高效协作和对齐。我们提出的算法的理论分析建立了关于给定现成模型的目标奖励所代表的目标任务的最佳性能。我们使用各种任务和偏好上的开源对齐模型进行了全面的实证评估，这证明了该方法优于单Agent解码基线。值得注意的是，Collab超越了当前SoTA解码策略，在平均奖励方面实现了高达1.56倍的提升，在基于GPT-4的胜-平率方面实现了71.89%的提升。
🔬 方法详解
问题定义：现有的大型语言模型对齐方法，特别是基于人类反馈的强化学习（RLHF），需要大量的计算资源来更新模型的参数。此外，单Agent解码方法在面对复杂多样的任务时，难以充分利用现有对齐模型的优势，导致性能受限。
核心思路：Collab的核心思路是将多个预先对齐的LLM视为不同的Agent，通过在解码过程中动态地选择最合适的Agent来生成每个token，从而实现推理时对齐。这种混合Agent的方法能够更好地适应不同的任务和偏好，提高生成质量。
技术框架：Collab的整体框架包括以下几个主要步骤：1) 准备一组预先对齐的LLM，每个LLM代表一个Agent。2) 在解码的每个token生成步骤，每个Agent根据当前上下文生成一个候选token。3) 使用一个效用函数来评估每个候选token的长期价值，该效用函数考虑了目标任务的奖励。4) 选择具有最高效用值的token，并将其添加到生成的序列中。5) 重复步骤2-4，直到生成完整的序列。
关键创新：Collab的关键创新在于其动态的Agent选择机制。与传统的单Agent解码方法不同，Collab能够根据当前上下文和目标任务的需求，灵活地选择最合适的LLM来生成每个token。这种策略切换机制能够充分利用不同LLM的优势，提高生成质量和对齐效果。
关键设计：Collab的关键设计包括：1) 效用函数的设计，需要能够准确地评估每个候选token的长期价值。2) Agent选择策略，需要能够高效地选择最佳Agent。3) 如何平衡不同Agent之间的贡献，避免某个Agent过度主导生成过程。具体的效用函数和选择策略可能需要根据具体的任务和数据集进行调整。
🖼️ 关键图片
📊 实验亮点
Collab在多个任务和偏好上进行了实验，结果表明其性能显著优于单Agent解码基线。具体而言，Collab在平均奖励方面实现了高达1.56倍的提升，在基于GPT-4的胜-平率方面实现了71.89%的提升。这些结果表明，Collab是一种有效的LLM对齐方法，能够充分利用多个预先对齐的LLM的优势。
🎯 应用场景
Collab可应用于各种需要LLM对齐的场景，例如对话系统、文本摘要、代码生成等。通过利用多个预先对齐的LLM，Collab能够提高生成质量，减少有害或不符合人类偏好的输出，从而提升用户体验和安全性。该方法还可用于个性化LLM，使其更好地适应不同用户的需求和偏好。
📄 摘要（原文）
Alignment of Large Language models (LLMs) is crucial for safe and trustworthy deployment in applications. Reinforcement learning from human feedback (RLHF) has emerged as an effective technique to align LLMs to human preferences and broader utilities, but it requires updating billions of model parameters, which is computationally expensive. Controlled Decoding, by contrast, provides a mechanism for aligning a model at inference time without retraining. However, single-agent decoding approaches often struggle to adapt to diverse tasks due to the complexity and variability inherent in these tasks. To strengthen the test-time performance w.r.t the target task, we propose a mixture of agent-based decoding strategies leveraging the existing off-the-shelf aligned LLM policies. Treating each prior policy as an agent in the spirit of mixture of agent collaboration, we develop a decoding method that allows for inference-time alignment through a token-level selection strategy among multiple agents. For each token, the most suitable LLM is dynamically chosen from a pool of models based on a long-term utility metric. This policy-switching mechanism ensures optimal model selection at each step, enabling efficient collaboration and alignment among LLMs during decoding. Theoretical analysis of our proposed algorithm establishes optimal performance with respect to the target task represented via a target reward for the given off-the-shelf models. We conduct comprehensive empirical evaluations with open-source aligned models on diverse tasks and preferences, which demonstrates the merits of this approach over single-agent decoding baselines. Notably, Collab surpasses the current SoTA decoding strategy, achieving an improvement of up to 1.56x in average reward and 71.89% in GPT-4 based win-tie rate.