Research map / Large language models / Alignment and preferences
Collab: Controlled Decoding using Mixture of Agents for LLM Alignment
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Collab, a training-free decoding method that switches token-by-token among multiple off-the-shelf aligned LLMs using an implicit Q-function to select the best agent for a target reward.
Problem
Single-agent controlled decoding struggles to adapt to diverse or conflicting tasks; aligning LLMs to a new target reward without retraining billions of parameters requires a principled way to combine existing aligned models at inference time.
Approach
At each token, select the next token and the agent (from a pool of aligned LLM policies) that maximizes an implicit Q-function (approximate target-reward value minus KL to reference), so decoding switches between agents dynamically; theory bounds suboptimality by the reward gap between target and best agent plus KL terms.
Key results
- Collab achieves up to 1.56x improvement in average reward over the current SoTA decoding strategy.
- Collab reaches up to 71.89% GPT-4 win-tie rate against baseline decoding strategies.
- Switching between two diverse agents gives significantly higher reward than either individual agent; gains increase with number of diverse agents.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Switch between diverse agents using the implicit Q-function instead of decoding with Agent-I alone | Collab · Berkeley Nectar and HH-RLHF | average reward: up to 1.56x | Helped |
| Use diverse agents rather than two similar agents for switching (ablation) | Collab · Average reward across evaluated setups | average reward | Helped |
| Increase the number of diverse agents in the mixture (ablation) | Collab · Average reward across evaluated setups | average reward | Helped |
| Switch between agents using the implicit Q-function instead of using other decoding baselines | Collab · Evaluation-1 to Evaluation-3 | diversity and coherence | Helped |
Surprising findings
- Switching between non-diverse agents yields poor improvement, indicating agent diversity matters.
Limitations
- Theoretical guarantee depends on the best agent's reward being close to the target reward; gap persists otherwise.
- Evaluation limited to 7B-scale open models and two datasets with a reward model-based target.
- Sampling top-p candidates from every agent at each token increases inference cost.
Experimental setup
Datasets: Berkeley Nectar HH-RLHF
Benchmarks: Berkeley Nectar HH-RLHF
Metrics: average reward GPT-4 win-tie rate diversity (n-gram repetition) coherence (SimCSE cosine similarity)
Baselines: Agent-I (single aligned model) Agent-II (single aligned model) BoN sampling Controlled Decoding (CD) TQ* principled decoding
Builds on: controlled decoding principled decoding (TQ*) KL-regularized RL best-of-N sampling
Closest papers by meaning
- Decoding-Time Language Model Alignment with Multiple Objectives (NeurIPS 2024 · similarity 0.90)
- Controlled Decoding from Language Models (ICML 2024 · similarity 0.90)
- Aligning Large Language Models with Representation Editing: A Control Perspective (NeurIPS 2024 · similarity 0.90)
- Improving Model Alignment Through Collective Intelligence of Open-Source Models (ICML 2025 · similarity 0.89)
- The Unlocking Spell on Base LLMs: Rethinking Alignment via In-Context Learning (ICLR 2024 · similarity 0.89)
- On-the-fly Preference Alignment via Principle-Guided Decoding (ICLR 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.