Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Agents of Diffusion: Enhancing Diffusion Language Models with Multi-Agent Reinforcement Learning for Structured Data Generation (Extended Version)

Aja Khanal, Kaushik T. Ranade, Rishabh Agrawal, Kalyan S. Basu, Apurva Narayan

TL;DR

Agents of Diffusion (AoD) addresses the challenge of generating high-quality, schema-conformant structured data by uniting diffusion language models (DLMs) with autoregressive LLM agents through language-mediated reinforcement learning. The core idea is to place a frozen DLM under the guidance of a Prompt Optimization Agent and a Judge, both communicating via natural language feedback, thereby avoiding fine-tuning or handcrafted constraints. The method combines diffusion-driven semantic richness with autoregressive constraints to preserve schema integrity, achieving high semantic fidelity and diverse outputs across multiple JSON benchmarks. This approach demonstrates state-of-the-art controllable generation for structured data, is model-agnostic and hardware-accessible, and opens avenues for extending diffusion-based structure-aware synthesis to broader domains like tabular data and code.

Abstract

Generating high-quality structured data such as JSON records, remains a fundamental challenge for large language models (LLMs), particularly when semantic richness must coexist with strict schema adherence. While autoregressive LLMs offer strong structural consistency, they often struggle with semantic variation and output diversity. In contrast, diffusion language models (DLMs) introduce powerful mechanisms for semantic richness and bidirectional decoding, yet lack the inductive biases needed for reliable structure preservation. We present Agents of Diffusion (AoD), a novel framework that unifies the generative flexibility of DLMs with the reasoning capabilities of autoregressive models through language-mediated reinforcement learning. AoD frames structured text generation as a multi-agent alignment process, where a prompt optimization agent collaborates with a judge agent to iteratively guide a DLM using natural language feedback. This approach enables controllable, schema-consistent generation without modifying model parameters or relying on handcrafted constraints. AoD advances the state of controllable generation by demonstrating that diffusion models, when supervised by cooperative agents, can achieve both high semantic novelty and structural fidelity. Across multiple structured data benchmarks, AoD consistently outperforms diffusion and autoregressive baselines, establishing a new path forward for structure-aware, diversity-enhanced text synthesis.

Agents of Diffusion: Enhancing Diffusion Language Models with Multi-Agent Reinforcement Learning for Structured Data Generation (Extended Version)

TL;DR

Agents of Diffusion (AoD) addresses the challenge of generating high-quality, schema-conformant structured data by uniting diffusion language models (DLMs) with autoregressive LLM agents through language-mediated reinforcement learning. The core idea is to place a frozen DLM under the guidance of a Prompt Optimization Agent and a Judge, both communicating via natural language feedback, thereby avoiding fine-tuning or handcrafted constraints. The method combines diffusion-driven semantic richness with autoregressive constraints to preserve schema integrity, achieving high semantic fidelity and diverse outputs across multiple JSON benchmarks. This approach demonstrates state-of-the-art controllable generation for structured data, is model-agnostic and hardware-accessible, and opens avenues for extending diffusion-based structure-aware synthesis to broader domains like tabular data and code.

Abstract

Generating high-quality structured data such as JSON records, remains a fundamental challenge for large language models (LLMs), particularly when semantic richness must coexist with strict schema adherence. While autoregressive LLMs offer strong structural consistency, they often struggle with semantic variation and output diversity. In contrast, diffusion language models (DLMs) introduce powerful mechanisms for semantic richness and bidirectional decoding, yet lack the inductive biases needed for reliable structure preservation. We present Agents of Diffusion (AoD), a novel framework that unifies the generative flexibility of DLMs with the reasoning capabilities of autoregressive models through language-mediated reinforcement learning. AoD frames structured text generation as a multi-agent alignment process, where a prompt optimization agent collaborates with a judge agent to iteratively guide a DLM using natural language feedback. This approach enables controllable, schema-consistent generation without modifying model parameters or relying on handcrafted constraints. AoD advances the state of controllable generation by demonstrating that diffusion models, when supervised by cooperative agents, can achieve both high semantic novelty and structural fidelity. Across multiple structured data benchmarks, AoD consistently outperforms diffusion and autoregressive baselines, establishing a new path forward for structure-aware, diversity-enhanced text synthesis.
Paper Structure (99 sections, 6 theorems, 35 equations, 3 figures, 3 tables, 1 algorithm)

This paper contains 99 sections, 6 theorems, 35 equations, 3 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Structured Textual Synthetic Data
  4. Autoregressive Language Models
  5. Diffusion Language Models
  6. Prompt Tuning and Optimization
  7. Multi-Agent Language-Based Coordination
  8. Methodology
  9. Preliminaries
  10. Problem Formulation
  11. DLM Integrated Multi-Agent RL
  12. Prompt Optimization Agent.
  13. Diffusion Language Model Agent.
  14. Judge Agent.
  15. Experiments and Results
  16. ...and 84 more sections

Key Result

Theorem 1

Let $\mathcal{T}(P)=\mathbb{E}_{x\sim g_\phi(P)}[R(J(x,\rho),x)]$ as in (eq:aod_policy_objective). If $\mathcal{T}(P)$ is locally Lipschitz and prompt edits $\Delta P_t$ sampled from $\pi_\theta(\Delta P \mid h_t)$ are bounded, then the iterative update $P_{t+1}=U(P_t,\Delta P_t)$ constitutes a cont ensuring stable convergence toward reward-aligned, schema-consistent prompts.

Figures (3)

  • Figure 1: Agents of Diffusion: Overview of the multi-agent training framework.
  • Figure 2: Comparison of normalized metrics across datasets.
  • Figure 3: Ablation of agents, averaged across datasets and Optimizer–Judge pairs. SP = Static Prompt, PO = Prompt Optimizer, FRL–AR–S = Autoregressive generator with scalar rewards, FRL–AR–NL = Autoregressive generator with natural language feedback, FRL–DLM–S = Diffusion generator with scalar rewards, FRL = Diffusion generator with natural language feedback.

Theorems & Definitions (6)

  • Theorem 1
  • Proposition 1
  • Proposition 2
  • Theorem 1: Convergence of prompt updates
  • Proposition 1
  • Proposition 2