PaperlayerView exact revision on arXiv

Source-linked AI summary

Coevolving with the Other You: Fine-Tuning LLM with Sequential Cooperative Multi-Agent Reinforcement Learning

Hao Ma, Tianyi Hu, Zhiqiang Pu, Boyin Liu, Xiaolin Ai, Yanyan Liang, Min Chen

arXiv:2410.06101v2cs.AIcs.MA

TL;DR

RL fine-tuning of LLMs faces instability and distribution collapse, motivating methods that better handle the challenges of sparse rewards and large action spaces. CORY trains duplicated pioneer and observer agents cooperatively with reward sharing and periodic role exchange. Across the reported evaluations, CORY surpasses PPO in policy optimality, resilience to collapse, and training robustness; the study focuses on cooperative two-agent fine-tuning.

  • Problem

    Contemporary RL algorithms can be unstable and vulnerable to distribution collapse in LLM fine-tuning, where large discrete action spaces and sparse rewards hinder exploration.

    Source: §1 · p. 2

  • Method

    CORY duplicates an LLM into pioneer and observer agents that share task rewards, transfer responses, and periodically exchange roles during cooperative RL fine-tuning.

    Source: §1 · p. 2

  • Results

    CORY surpasses PPO in policy optimality, resilience to distribution collapse, and training robustness across subjective and objective reward evaluations.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The reported results support CORY as a potential method for improving LLMs in practical applications.

    Source: §1 · p. 2

    The study focuses on cooperation and two LLMs; competitive MARL and settings with more agents are left for future study.

    Source: §6 · p. 9

Abstract

from arXiv · show

Reinforcement learning (RL) has emerged as a pivotal technique for fine-tuning large language models (LLMs) on specific tasks. However, prevailing RL fine-tuning methods predominantly rely on PPO and its variants. Though these algorithms are effective in general RL settings, they often exhibit suboptimal performance and vulnerability to distribution collapse when applied to the fine-tuning of LLMs. In this paper, we propose CORY, extending the RL fine-tuning of LLMs to a sequential cooperative multi-agent reinforcement learning framework, to leverage the inherent coevolution and emergent capabilities of multi-agent systems. In CORY, the LLM to be fine-tuned is initially duplicated into two autonomous agents: a pioneer and an observer. The pioneer generates responses based on queries, while the observer generates responses using both the queries and the pioneer's responses. The two agents are trained together. During training, the agents exchange roles periodically, fostering cooperation and coevolution between them. Experiments evaluate CORY's performance by fine-tuning GPT-2 and Llama-2 under subjective and objective reward functions on the IMDB Review and GSM8K datasets, respectively. Results show that CORY outperforms PPO in terms of policy optimality, resistance to distribution collapse, and training robustness, thereby underscoring its potential as a superior methodology for refining LLMs in real-world applications.

1 Introduction

LLM fine-tuning with RL faces instability and distribution collapse, while cooperative multi-agent RL motivates CORY’s two-agent approach. Across subjective and objective reward settings, CORY is reported to outperform PPO on policy optimality, resilience to collapse, and training robustness.

1 Introduction

  • PPO and other contemporary RL algorithms can be unstable and vulnerable to distribution collapse when fine-tuning LLMs.The paper describes collapse as over-optimization producing highly biased behavior.Source: §1 · p. 2
  • LLM fine-tuning poses large discrete action spaces and sparse rewards, which hinder exploration and contribute to PPO instability.For Llama-2, the vocabulary can yield 32,000 action choices, while reward arrives only after the full response.Source: §1 · p. 2
  • Cooperative MARL motivates fine-tuning agents together because coevolution can yield emergent skills, and CORY instantiates this with a pioneer and observer.The observer uses both the query and the pioneer’s response; the agents share rewards and periodically exchange roles.Source: §1 · p. 2
  • CORY outperforms PPO in policy optimality, resilience to distribution collapse, and training robustness across subjective and objective reward functions.The evaluation uses IMDB reviews for subjective rewards and GSM8K for objective rewards.Source: §1 · p. 2

2 Problem Formulation

The section formulates causal next-token prediction as a sequential decision process, with a task reward assigned to a generated token sequence. It also describes token-level KL penalties as a way to constrain deviation from the initial model.

2 Problem Formulation

  • The language-augmented Markov decision process represents next-token prediction using vocabulary, states, actions, rewards, transitions, and a discount factor.The model’s vocabulary contains possible tokens, and each action is a generated token sequence.Source: §2 · pp. 2–3
  • A causal language model’s sentence-level policy is a joint policy over its sequential token-level decisions.The formulation represents prompts and generated responses as token sequences in a language-augmented Markov decision process.Source: §2 · p. 3
  • Task reward is sparse because it is obtained only after the sequence’s N token-generation steps.The reward function assigns a numerical score to a token sequence; the task reward is produced after the final token.Source: §2 · p. 3
  • Token-level KL penalties constrain the current policy’s deviation from the initial policy to mitigate distribution collapse under sparse reward.Each token receives a penalty based on the KL divergence between current and initial policies.Source: §2 · p. 3

3 Method

CORY trains two copies of an LLM as a pioneer and observer, combining their task rewards and periodically exchanging roles. Its multi-objective analysis reports that CORY’s sub-optimal frontier is closer to the Pareto frontier than PPO’s, while identifying knowledge transfer and role exchange as central mechanisms.

3.1 Coevolving with the Other You (CORY)

  • CORY duplicates the fine-tuned LLM into pioneer and observer agents that act asynchronously, with the observer using the pioneer’s response as guidance.The pioneer responds to the query alone; the observer receives both the query and the pioneer’s response.Source: §3.1 · p. 3
  • CORY is presented as a three-step extension: duplicate the LLM, combine the agents’ task rewards, and periodically exchange roles.After training, either agent can perform the task independently.Source: Figure 1 · p. 4
  • The agents receive a collective task reward formed by summing their individual task rewards, so each can benefit from the other’s improvement.The resulting collective reward is combined with a KL penalty to form the reward signal.Source: §3.1 · p. 4
  • Periodic role exchange addresses the observer’s prompt bias by exposing both agents to query-only and query-plus-response formats, allowing either to perform independently.The paper describes role exchange as helping transfer observer-learned skills back to the pioneer.Source: §3.1 · p. 4
  • CORY is described as algorithm-agnostic and compatible in principle with RL algorithms beyond PPO.The authors also characterize it as simple to implement and integrate with existing frameworks.Source: §3.1 · p. 4

3.2 Understanding CORY

  • Task reward and KL divergence can oppose one another, making RL fine-tuning a multi-objective trade-off whose ideal policies lie on a Pareto frontier.The paper notes that perfect convergence to that frontier is often intractable for natural language, so it considers sub-optimal frontiers instead.Source: §3.2 · p. 5
  • CORY’s sub-optimal frontier consistently outperforms PPO’s on GSM8K across adjusted η preferences, using task reward and negative KL divergence as the two objectives.Larger values on the plotted negative-KL axis indicate better performance.Source: §3.2 · p. 5, Figure 2 · p. 5
  • Knowledge transfer guides the observer’s optimization, while role exchange addresses prompt bias and transfers learned skills back to the pioneer.The paper says CORY achieves lower KL divergence while maintaining strong target-task performance.Source: §3.2 · pp. 5–6

4 Experiments

Experiments compare CORY with single-PPO under subjective rewards on IMDB and rule-based rewards on GSM8K, then ablate model size, knowledge transfer, and role exchange.

4.1 Subjective Rewards on IMDB Review

  • On IMDB, CORY matched PPO’s task reward with less than half its KL, while its combined reward rose and PPO’s declined.The two CORY agents also reached very similar performance levels after training.Source: Figure 3 · p. 6, §4.1 · p. 7

4.2 Objective Rewards on GSM8K

  • On GSM8K, CORY achieved 18% pass@1 and steadier task rewards, lower KL, and rising combined reward, while single-PPO was unstable.The test-set comparison reports pass@1; training curves show single-PPO’s task reward peaking around iteration 50 before declining.Source: §4.2 · pp. 7–8, Figure 4 · p. 7, Figure 5 · p. 8

4.3 Ablations

  • Doubling PPO’s model size to GPT2-XL did not improve the KL penalty, which continued rising alongside distribution collapse.The comparison uses GPT2-XL (1.5B), twice the parameters of GPT2-Large.Source: §4.3 · p. 8, Figure 6 · p. 8
  • Disabling knowledge transfer while retaining role exchange made task rewards unstable and KL divergences higher than CORY’s.The authors describe this variant as resembling PPO with individual queries and a noisier reward signal.Source: §4.3 · p. 8, Figure 6 · p. 8
  • Without role exchange, both agents achieved good task rewards but had higher KL than CORY; the observer’s KL was lower than the pioneer’s.The authors relate the observer’s lower KL to receiving the pioneer’s response and using its exploration to refine its policy.Source: §4.3 · pp. 8–9, Figure 6 · p. 8

5 Related Work

The related work situates CORY among efforts to address PPO-based RLHF challenges and prior applications of multi-agent learning. It identifies asynchronous two-agent training as an architectural precedent for CORY’s sequential cooperative setting.

RLHF

  • Prior RLHF work addresses PPO’s complexity, instability, and hyperparameter sensitivity through regularization, tuning, and reward shaping, but reports no significant improvement over vanilla PPO.This motivates exploring cooperative MARL as an alternative perspective on RL fine-tuning.Source: §5 · p. 9

MARL

  • Earlier MARL work applies joint tuning to prompts and asynchronous observation to Stackelberg games, informing CORY’s sequential cooperative training design.The cited Stackelberg framework lets the later-moving agent observe the earlier agent’s decision.Source: §5 · p. 9

6 Discussion

CORY’s cooperative multi-agent approach is presented as a way to improve RL fine-tuning robustness and avoid distribution collapse. The discussion identifies competitive methods and larger agent populations as future directions.

6 Discussion

  • Cooperative MARL fine-tuning appears to increase training robustness and prevent distribution collapse.The discussion attributes this potential benefit to coevolution and the generation of varied data through multi-agent learning.Source: §6 · p. 9
  • The discussion focuses on cooperation, while competitive MARL—especially population-based methods—is proposed as a future research direction.The authors describe an auto-curriculum driven by a natural arms race as a possible mechanism for agent learning.Source: §6 · p. 9
  • The experiments examine two LLMs, and incorporating more LLMs in MARL fine-tuning is left for future study.Source: §6 · p. 9

7 Conclusion

CORY fine-tunes two role-differentiated LLM agents through knowledge transfer and role exchange. The authors report advantages over PPO in policy optimality, distribution-collapse resilience, and training robustness.

7 Conclusion

  • CORY duplicates the pretrained LLM into two differently assigned agents and uses knowledge transfer and role exchange to support collaborative learning.Source: §7 · p. 9
  • CORY surpasses PPO in policy optimality, resilience to distribution collapse, and training robustness.The authors also report that CORY has an advantage over single-agent RL in approaching the Pareto frontier between KL divergence and task reward.Source: §7 · p. 10
  • After fine-tuning, either agent can be selected to perform the task independently.Source: §7 · p. 10
  • The analysis identifies a Pareto frontier between KL divergence and task reward.Source: §7 · p. 10

A Implementation Details

The implementation details cover computing resources, hyperparameter selection, task-specific prompts, CORY’s token-level PPO updates, and its training algorithm.

A Implementation Details

  • With one GPU, CORY fine-tunes GPT2-XL on IMDB Review in 12 hours; with four GPUs, it fine-tunes 4-bit Llama-2-7b-chat on GSM8K in four hours.The reported hardware includes two AMD EPYC 7773X CPUs and eight NVIDIA A6000 GPUs.Source: A Implementation Details · p. 13

A.1 Hyperparameters

  • PPO hyperparameters were grid-searched over α and η for stable training, while most settings were reused for CORY except its learning rate α.GPT-2 settings followed TRL defaults for IMDB; Llama-2 settings primarily followed StackLlama guidelines.Source: §A.1 · p. 13

A.2 Prompt Details

  • For IMDB and GSM8K, the observer receives the pioneer’s response as a reference before completing the review or answering the question.The pioneer receives the initial task prompt; the observer prompt also includes instructions to complete the task.Source: §A.2 · p. 14

B Token-Level Policy Update of CORY

  • Pioneer and observer share the task reward but have different Q-functions because their inputs differ; PPO updates each agent independently.A value head predicts state value, and the policy uses a clipped optimization objective with an advantage estimate.Source: B Token-Level Policy Update of CORY · p. 15

C.1 Algorithm of CORY

  • CORY initializes two role-specific agents, collects their responses and rewards, updates each agent, and periodically exchanges their roles.Either fine-tuned agent can later be selected to perform the task independently.Source: §C.1 · p. 16

C.2 Token-Level Policy Update

  • The token-level update assigns KL penalties before the final token, then adds the task reward at the final token before computing advantages and updating policy and value parameters.Source: §C.2 · p. 16

D Qualitative Analysis of Experiment Results.

On IMDB Review, CORY avoids the distribution collapse observed with PPO despite comparable task reward; on GSM8K, CORY is more stable across PPO’s stable and collapsed training cases.

D Qualitative Analysis of Experiment Results.

E.1 Robustness of CORY

CORY maintains stable training across learning rates and model sizes where PPO risks distribution collapse. Additional experiments find lower task rewards in competitive than cooperative settings and effective KL-divergence and task-reward stabilization compared with REINFORCE and Elastic Reset.

E.1 Robustness of CORY

  • Across learning rates 1e-4 and 1e-5, CORY trains stably on Llama-2-7b-chat, while PPO collapses at 1e-4 and shows rising KL divergence after 100 iterations at 1e-5.CORY's KL divergence and task reward converge around iteration 10, with KL divergence remaining relatively low.Source: §E.1 · p. 20, Figure 7 · p. 20, Figure 8 · p. 20
  • On Llama-2-13b-chat, CORY remains stable at both learning rates, while PPO rapidly collapses at 1e-4 and shows rising KL divergence at 1e-5.At 1e-5, CORY's curves stabilize even though task reward does not improve; at 1e-4, it achieves good task reward and low KL divergence.Source: Figure 9 · p. 20, §E.1 · p. 21, Figure 10 · p. 21
  • Across the training curves, CORY balances task reward and KL divergence more consistently than PPO.The authors report that PPO struggles to balance the two measures.Source: §E.1 · p. 21

E.2 Different Reward Settings

  • Competitive reward settings yield significantly lower task rewards than cooperative settings.The experiments vary the other agent's reward weight λ across {-5, -3, -1, 1, 3, 5} and normalize reward values to mitigate reward-magnitude effects.Source: §E.2 · p. 21, Figure 11 · p. 21

E.3 Additional Baselines

  • Compared with REINFORCE and Elastic Reset, CORY stabilizes KL divergence and task reward effectively.The authors report that REINFORCE is more prone to distribution collapse than PPO, while Elastic Reset can recover some performance but has volatile training that complicates decisions about when to stop and save parameters.Source: §E.3 · p. 22, Figure 12 · p. 22

F Limitations

CORY requires two copies of the LLM, doubling computational resources, though parameter sharing could alleviate this cost. The paper also discusses potential positive and negative broader impacts, including possible abuse and suggested safeguards.

F Limitations

  • CORY duplicates the LLM into two copies, doubling computational resources, while parameter sharing could alleviate this cost.Source: F Limitations · p. 22

G Broader Impacts

  • The authors state that stronger RL fine-tuning could theoretically improve specialized-task capabilities beyond human level if reward exceeds a threshold.This is presented as a theoretical possibility, not an experimental result.Source: G Broader Impacts · p. 22
  • Potential abuse includes generating misleading or harmful content; the authors suggest value alignment and real-time detection as possible safeguards.Source: G Broader Impacts · p. 22

2. Limitations

The authors report discussing the work’s limitations in Appendix F and summarize broader impacts in Appendix G. They also describe reproducibility materials and experimental details in the paper and appendices.

2. Limitations

  • The authors state that the paper discusses its limitations in Appendix F.Source: §2 · p. 23
  • The authors report open access to the data and code needed to reproduce the main experimental results and provide setup, algorithm, and model-architecture descriptions, hyperparameters, and pseudocode.Source: §2 · pp. 24–25
  • Training curves are reported as mean ± standard deviation across three random seeds, and the authors direct readers to Appendix A.1 for compute and hyperparameter details.Source: §2 · pp. 25–26

10. Broader Impacts

  • The authors state that Appendix G discusses potential positive and negative societal impacts and that existing code, data, and models are credited in Section 4.Source: §10 · pp. 26–27
Loading 2410.06101v2…