Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Online Joint Fine-tuning of Multi-Agent Flows

Paul Mineiro

TL;DR

This work tackles online joint fine-tuning of multi-agent flows by leveraging an environment simulator to convert episode-level preferences into node-level preferences, enabling language-model components to be trained without granular supervision. It borrows Learning to Search principles and one-step deviation credit assignment to optimize greedy inference in complex, looped flows, and it uses an online DPO-style preference learner with Coin Betting optimization and LoRA adapters. The approach is demonstrated on Musique, achieving state-of-the-art results on both Musique-Answerable and Musique-Full, including reward-free and text-feedback scenarios. The discussion notes practical limitations, such as the need for simulator access and the sensitivity to initialization, and it outlines potential directions for broader evaluation and single-sample learning in production settings.

Abstract

A Flow is a collection of component models ("Agents") which constructs the solution to a complex problem via iterative communication. Flows have emerged as state of the art architectures for code generation, and are the raison d'etre for frameworks like Autogen. However, flows are currently constructed via a combination of manual prompt engineering and stagewise supervised learning techniques; the latter is limited to acyclic flows with granular node supervision. In this writeup I describe a procedure for online joint fine-tuning of an entire flow inspired by the Learning to Search framework. The approach leverages simulator access to reduce preferences over entire episodes to preferences over individual node outputs; when the components are language models the latter is a well-studied problem. The approach is applicable to reward-free settings (e.g., text feedback) if an episode evaluator model is available. I apply to the multi-hop QA dataset Musique achieving a state-of-the-art result.

Online Joint Fine-tuning of Multi-Agent Flows

TL;DR

This work tackles online joint fine-tuning of multi-agent flows by leveraging an environment simulator to convert episode-level preferences into node-level preferences, enabling language-model components to be trained without granular supervision. It borrows Learning to Search principles and one-step deviation credit assignment to optimize greedy inference in complex, looped flows, and it uses an online DPO-style preference learner with Coin Betting optimization and LoRA adapters. The approach is demonstrated on Musique, achieving state-of-the-art results on both Musique-Answerable and Musique-Full, including reward-free and text-feedback scenarios. The discussion notes practical limitations, such as the need for simulator access and the sensitivity to initialization, and it outlines potential directions for broader evaluation and single-sample learning in production settings.

Abstract

A Flow is a collection of component models ("Agents") which constructs the solution to a complex problem via iterative communication. Flows have emerged as state of the art architectures for code generation, and are the raison d'etre for frameworks like Autogen. However, flows are currently constructed via a combination of manual prompt engineering and stagewise supervised learning techniques; the latter is limited to acyclic flows with granular node supervision. In this writeup I describe a procedure for online joint fine-tuning of an entire flow inspired by the Learning to Search framework. The approach leverages simulator access to reduce preferences over entire episodes to preferences over individual node outputs; when the components are language models the latter is a well-studied problem. The approach is applicable to reward-free settings (e.g., text feedback) if an episode evaluator model is available. I apply to the multi-hop QA dataset Musique achieving a state-of-the-art result.
Paper Structure (16 sections, 3 figures, 2 tables)

This paper contains 16 sections, 3 figures, 2 tables.

Table of Contents

  1. Introduction
  2. The Approach
  3. The Challenges
  4. Addressing the Challenges
  5. How to Generate Rollouts?
  6. Technical Errata
  7. Musique Results
  8. Prompts
  9. Hyperparameter Selection
  10. Musique Specific Techniques
  11. Subsampling the Training Set
  12. Enforcing Grounded Answers
  13. Train on Dev
  14. Pairwise Inference
  15. Miscellaneous
  16. ...and 1 more sections

Figures (3)

  • Figure 1: Actual flow utilized for Musique-Answerable. Key realistic aspects are: dependent generation with loops (the entire retrieval block is looped multiple times); dynamic conversation state ("scratchpad"); learned conditional branching ("stop retrieval"); and lack of granular supervision.
  • Figure 2: Visualization of one-step deviations for inducing individual node output preferences $(O' \succ O)$ from complete episode preferences $(E' \succ E)$. At each point in the realized unrolled computation graph, an alternative output is produced and then the episode is evaluated until termination. Note this can result in a different sequence of nodes being called due to conditional execution. Furthermore, although all the node identifiers in this picture are unique, in practice the same node appears at multiple point in the unrolled computation graph due to loops. Finally, rollouts do not themselves generate rollouts, so the computational overhead of learning is quadratic in the length of the unrolled computation graph (a linear number of rollouts, each doing a linear amount of additional inference).
  • Figure 3: For the Full version of the Musique dataset, an additional node is introduced into the flow to predict sufficiency. The flow is otherwise identical.