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

Summary

The paper introduces CORY, a novel reinforcement learning (RL) technique for fine-tuning language models that casts a multi-agent framework on the trained language model by duplicating it at initialization and then using the two copies to improve one another. Essentially, a "pioneer" gives the first guess, used as a reference to an "observer" which gives the final answer, then both agents are rewarded for both of their outputs in a cooperative fashion and trained independently with an underlying RL algorithm, here PPO. The authors claim that CORY improves on single-agent PPO in terms of downstream performance, training stability, and robustness to distribution collapse.

Strengths

The paper is easy to follow and includes enough details about the experiments. Casting the problem into a multi-agent framework to change the optimization dynamics is original and will likely influence other work to build on it. The observations made by the authors about the multi-agent framework changing the dynamics training dynamics to make the observer more easily optimize for the reward while maintaining a low KL are interesting. The method can potentially be applied on top of any RL algorithm, either with independent learners as shown in the paper or by introducing centralized training critic network sharing, etc.

Weaknesses

Overall, I find the contribution of the paper to be original and likely to have impact, however, I also find a number major issues in the work that make it not ready for publication at this venue. Unclear contextualization and missing baselines: - One main contribution of the paper is that it finds a policy that has a better Pareto frontier while still using the sum of task reward and KL divergence, but the KL divergence is an artificial reward specific to current RL fintuning methods in language modeling. If the paper's method is to be highlighted from a language modeling perspective then I would expect it to compare to methods like [1] which also claim to obtain a better Pareto frontier. And more than that, I would expect the method to show its benefits in downstream tasks for example with AlpacaFarm [2] as although a better frontier could be seen a a proxy to robustness to model collapse, it isn't clear if it translates to better language modeling capabilities. Otherwise if the method is to be highlighted in general multi-objective RL fine-tuning tasks, then I would expect comparison to multi-objective RL methods. - (minor) The method can also be seen a self-improvement method, so the pepper would benefit from discussing and contrasting methods from the same family. An insufficient number of benchmarks: - The authors claim to evaluate their method "across diverse reward function landscapes, encompassing both subjective and objective categories", but only one dataset per category is tested. Judging by the scope of the claims the authors make and that this is the only form of evidence provided, I would expect at least 3 or 5 datasets per category. Hyperparameter choice: - The authors use a single set of hyperprameters, with no apparent justification. This is not enough to make claims about a method being better than another. - (minor) PPO with 1 epoch and 1 batch per epoch (as in the hyperprameter table) becomes just a policy gradient, so the authors seem to be effectively using a vanilla actor-critic with a GAE value estimator and (usually useless) value clipping. Presentation: - (minor) The wording of the claims made in the abstract and introduction "policy optimality, "resilience to distribution collapse", and "robustness during training" is not reused in the experiments section which makes it hard to connect evidence to claims. - (minor) The $\geq$ symbol in equation 6 between vector-values objectives is undefined. One would expect a definition of Pareto dominance. [1] Noukhovitch, Michael, et al. "Language model alignment with elastic reset." Advances in Neural Information Processing Systems 36 (2024). [2] Dubois, Yann, et al. "Alpacafarm: A simulation framework for methods that learn from human feedback." Advances in Neural Information Processing Systems 36 (2024).

Questions

How did the authors select the hyperprameters for their methods?

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

Although in the appendix, the authors adequately state the limitations of their work and its broader impacts.

Summary

This paper presents a Reinforcement Learning (RL) methodology to fine-tune LLMs based on Multi-Agent RL agents. One agent acts as the observer, and the other acts as the pioneer. They share knowledge through two interactions: transfer learning and role-switching. They named this methodology CORY (Coevolving with the Other You). They tested their framework with out-the-shelf LLMs and proved that learning benefits from using CORY. They compared against proximal policy optimization (PPO) in the Multi-Agent setup. Ablation studies showed the effect of model size, knowledge transfer, and role exchange. They provided theory related to RL in the context of this text-generation task and qualitative studies.

Strengths

The paper is easy to follow for a person knowledgeable in RL who understands LLMs at a basic level. Their methodology provides a viable method to improve LLM fine-tuning using RL. They used modern and relevant benchmarks not only on the side of RL (PPO) but also on the side of models.

Weaknesses

The Limitations section would've been a good addition to the main corpus of the paper as it addresses a big possible concern I imagined with your methodology: you require twice the amount of resources to train both agents. The average reader would enjoy this change, as some don't jump in appendices. Also, I think your claim about Moore's Law might be arguable, as current requirements for computing power are growing beyond expectation with the surge in LLMs.

Questions

I think I lack intuition about how the reward signals work after the N steps. Also, I am unclear about the influence of the KL divergence and how far it is relevant not to diverge from the reference policy. Could you please give me your take on that?

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors added a Limitations section in the appendix.

Summary

This paper presents CORY, a novel approach for fine-tuning large language models (LLMs) using a sequential cooperative multi-agent reinforcement learning framework. Traditional methods, primarily based on PPO and its variants, often show suboptimal performance and risk distribution collapse. CORY addresses these issues by duplicating the LLM into two agents: a pioneer and an observer. These agents generate responses based on queries and each other’s outputs, exchanging roles periodically to foster cooperation. Experiments with GPT-2 and Llama-2 on the IMDB Review and GSM8K datasets demonstrate that CORY outperforms PPO in terms of policy optimality, resistance to distribution collapse, and training robustness. This highlights its potential as a superior methodology for refining LLMs in real-world applications.

Strengths

1. The introduction of multi-agent reinforcement learning into LLM fine-tuning is both novel and compelling, with a well-justified motivation behind the designed role exchange. 2. The presentation is well-organized and clear. The comprehensive explanations and illustrative figures in section 3.2 enhance the understanding of the proposed method and its effectiveness. 3. The empirical experiments are adequate and effectively support the proposed method and its underlying rationale.

Weaknesses

1. The paper lacks a theoretical analysis of the introduced framework. 2. The baseline comparison is primarily with single-agent PPO. However, there are many more advanced algorithms for LLM fine-tuning. It would be beneficial if the authors conducted additional experiments using other baseline fine-tuning algorithms.

Questions

1. Can the authors provide a more detailed discussion on the relationship between emergent communication in multi-agent reinforcement learning (MARL) and the method in this paper? Since the authors mention that this work is inspired by the concept that "Languages develop through agent interactions and are shaped by societal and cultural influences," it would be beneficial to explore this connection further. 2. The paper suggests employing a collective task reward to facilitate cooperation, as outlined in Eq. 4. What would happen if the task reward is not the sum of individual task rewards? Can the authors provide empirical results or theoretical analysis to show the influence of the reward design? 3. Currently, the framework includes only two agents: the pioneer and the observer. Is it possible to introduce additional agents, and would doing so enhance performance? I hope the authors can discuss this possibility further.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations have been discussed in the paper.

Summary

In this paper, the authors study a multi-agent organization for LLM learning. Specifically, they have two LLMs, with the second one responding to the same query given the query itself and the response generated by another LLM. They author shows this method can achieve a better tradeoff of the task reward and the KL penalty.

Strengths

The writing is clear and the question and proposed method is interesting.

Weaknesses

Since this is an empirical paper, the expectation of the experiments, from its design to its implementation and results, would be higher. In the questions below, the reviewer mentions several points that might improve the experiments. If the authors address the concerns about the experiments, the reviewer will consider rising their score.

Questions

1. As far as the reviewer is concerned, this paper is missing several important baselines to make their experiments more convincing. (1.1) The dual-LLM setup in the LLM debating literature. This line of research is highly related to the paper. There, two LLMs response in turn to the previous generation of the other part, as well as to the original query. As the authors are explicitly considering a cooperative setting, the important question then is: which one is more effective, cooperation or competition, in terms of multi-agent LLM learning. (1.2) Mixture of experts. It appears a bit unfair to compare with a single-LLM PPO baseline--after all the authors are using two LLMs. The question is whether this pioneer-observer structure is more effective than a simpler mechanism where opinions from multiple LLMs are aggregated together by recent research in MOE. 2. Fig. 2 needs to be further explained. The major drawback is the lack of legend regarding the eta values. The reviewer is asking for the which two points (one on the CORY frontier and another on the PPO frontier) corresponding which value of eta. A minor question: For the sub-optimal frontiers of CORY and PPO, it is unclear what is shown is the training KL and reward or that of the testing time.

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

The reviewer lists some questions that make them curious, expecting some clarifications. it seem that some analysis or experiments can improve the understanding of the proposed method. (1) In the current setup, the pioneer and the observer possess the same architecture, ie, they are the same LLMs. Does this homogeneity matter? If the two LLMs are different, can the proposed method still work?

Discussion period

The discussion period is almost over, so both authors and reviewers please respond to any unaddressed questions. Reviewers, be sure that you have all of the information you need from the authors, since after the 13th, they won't be able to respond.