Table of Contents
Fetching ...

Aligning Compound AI Systems via System-level DPO

Xiangwen Wang, Yibo Jacky Zhang, Zhoujie Ding, Katherine Tsai, Haolun Wu, Sanmi Koyejo

TL;DR

The paper tackles the challenge of aligning compound AI systems by modeling them as DAGs and extending Direct Preference Optimization to system-level preferences. It introduces SysDPO with two variants (SysDPO-Direct and SysDPO-Sampling) to handle cases with or without observable intermediate outputs, and provides theoretical guarantees for beta-perfect alignment. Empirically, it validates the approach on a joint LLM+diffusion task and a two-LLM collaboration system, showing that holistic, joint alignment significantly outperforms component-wise or prompting-only baselines. The work lays a foundation for scalable, reliable coordination of complex multi-component AI workflows and outlines avenues for future efficiency and scalability improvements.

Abstract

Compound AI systems, comprising multiple interacting components such as LLMs, foundation models, and external tools, have demonstrated remarkable improvements compared to single models in various tasks. To ensure their effective deployment in real-world applications, aligning these systems with human preferences is crucial. However, aligning the compound system via policy optimization, unlike the alignment of a single model, is challenging for two main reasons: (i) non-differentiable interactions between components make end-to-end gradient-based optimization method inapplicable, and (ii) system-level preferences cannot be directly transformed into component-level preferences. To address these challenges, we first formulate compound AI systems as Directed Acyclic Graphs (DAGs), explicitly modeling both component interactions and the associated data flows. Building on this formulation, we introduce $\textbf{SysDPO}$, a framework that extends Direct Preference Optimization (DPO) to enable joint system-level alignment. We propose two variants, SysDPO-Direct and SysDPO-Sampling, tailored for scenarios depending on whether we construct a system-specific preference dataset. We empirically demonstrate the effectiveness of our approach across two applications: the joint alignment of a language model and a diffusion model, and the joint alignment of an LLM collaboration system.

Aligning Compound AI Systems via System-level DPO

TL;DR

The paper tackles the challenge of aligning compound AI systems by modeling them as DAGs and extending Direct Preference Optimization to system-level preferences. It introduces SysDPO with two variants (SysDPO-Direct and SysDPO-Sampling) to handle cases with or without observable intermediate outputs, and provides theoretical guarantees for beta-perfect alignment. Empirically, it validates the approach on a joint LLM+diffusion task and a two-LLM collaboration system, showing that holistic, joint alignment significantly outperforms component-wise or prompting-only baselines. The work lays a foundation for scalable, reliable coordination of complex multi-component AI workflows and outlines avenues for future efficiency and scalability improvements.

Abstract

Compound AI systems, comprising multiple interacting components such as LLMs, foundation models, and external tools, have demonstrated remarkable improvements compared to single models in various tasks. To ensure their effective deployment in real-world applications, aligning these systems with human preferences is crucial. However, aligning the compound system via policy optimization, unlike the alignment of a single model, is challenging for two main reasons: (i) non-differentiable interactions between components make end-to-end gradient-based optimization method inapplicable, and (ii) system-level preferences cannot be directly transformed into component-level preferences. To address these challenges, we first formulate compound AI systems as Directed Acyclic Graphs (DAGs), explicitly modeling both component interactions and the associated data flows. Building on this formulation, we introduce , a framework that extends Direct Preference Optimization (DPO) to enable joint system-level alignment. We propose two variants, SysDPO-Direct and SysDPO-Sampling, tailored for scenarios depending on whether we construct a system-specific preference dataset. We empirically demonstrate the effectiveness of our approach across two applications: the joint alignment of a language model and a diffusion model, and the joint alignment of an LLM collaboration system.

Paper Structure

This paper contains 59 sections, 2 theorems, 51 equations, 8 figures, 3 tables.

Key Result

Proposition 1

Suppose an optimal model $\theta^*\in \Theta$ achieves the maximum of the RLHF objective eq:policy or the minimum of the DPO loss function eq:dpo_loss where data preference distribution ${\mathcal{D}}$ is given by the preference oracle $\mathop{\mathrm{pref}}\nolimits(\cdot)$. Then, it follows that

Figures (8)

  • Figure 1: The figure illustrates the challenges in a compound system composed of the GPT-4 and the image generator DALL-E. Given the user prompt to GPT-4, "Generate three separate images of a cat being progressively angrier", the task is to demonstrate a clear visual progression of the specified attribute, i.e., anger. (a) shows the results from one query, and (b) represents the results from another query. The captions under each image summarize the prompts generated by GPT-4 for DALL-E (complete prompts in Appendix \ref{['sec:figure1_prompt']}), where prompts from both queries reflect progressions in anger. Similarly, DALL-E accurately generates the images following the given prompts. However, (a) fails to demonstrate a clear visual progression of anger compared to (b), highlighting GPT-4's inconsistent collaboration with DALL-E.
  • Figure 2: Corresponding DAGs of compound AI systems. (a) The user gives a prompt $x$ which is processed by the LLM $\theta_1$ to produce three captions $y_1, y_2, y_3$. The diffusion model $\theta_2$ is queried to generate images $z_i$ given $y_i$ for $i=1,\ldots,3$. (b) The user gives a prompt $x$ which is processed by the first LLM $\theta_1$ to produce an intermediate result $y$. Then, $x$ and $y$ are passed to the second model $\theta_2$ to generate the final output $z$.
  • Figure 3: An example of a success output before training.
  • Figure 4: An example of a failure case before training. The first and second images lack clear progression in ice intensity.
  • Figure 5: An example of an improved case after training. The sequence shows smooth and consistent progression in the ice intensity.
  • ...and 3 more figures

Theorems & Definitions (5)

  • Definition 1: $\beta$-Perfect Alignment
  • Proposition 1
  • Theorem 1
  • proof : Proof of Proposition \ref{['prop:DPO:palign']}
  • proof : Proof of Theorem \ref{['thm:sysdpo-direct']}