Internalizing Agent Experience into Diffusion Model Weights via On-Policy Context Distillation

AI summary

Overview

Research area: Agentic text-to-image (T2I) generation, diffusion model training, and knowledge distillation.

Technical level: Advanced. The paper assumes familiarity with diffusion/flow-matching generative models, velocity-field formulations, LoRA fine-tuning, preference-based diffusion training (Diffusion-DPO, D-OPSD), and on-policy distillation.

Scope: The paper proposes Diffusion On-Policy Context Distillation (D-OPCD), a self-distillation algorithm that transfers prompt-mediated gains produced by an agentic T2I harness into the weights of the image generator itself, evaluated on four in-distribution benchmarks plus one out-of-distribution suite.

What This Paper Is About

Wrapping a text-to-image model in an agentic harness (memory, skills, planning, verification, and iterative prompt refinement) substantially improves image quality, but those gains live entirely outside the generator and must be re-derived on every request. This paper asks whether the gains discovered through harness evolution can be baked directly into the diffusion model's weights, so that the model captures part of the harness's benefit when it is given only the user's original query. The proposed answer is D-OPCD, which treats the harness-produced prompt as privileged teacher context and distills it into a student model conditioned only on the original query.

Key Contributions

1. D-OPCD (Diffusion On-Policy Context Distillation): An on-policy context distillation algorithm in which an EMA teacher conditioned on both the original query q_i and the harness-selected prompt p_i supervises a student conditioned only on q_i, with both evaluated along the student's own denoising trajectory. This makes prompt-mediated harness gains available from the original query alone.

2. Auto Skill Evolver (ASE): A hierarchical test-time learning mechanism that organizes agent experience into three levels (episodes, insights, and skills), consolidating mature insights into reusable skills that guide first-prompt construction. ASE supplies the query–prompt pairs that D-OPCD trains on.

3. A harness–generator co-evolution formulation for T2I: Because internalized gains no longer need to live in the harness, the harness can discard saturated skills and evolve again around the updated generator, creating a cycle in which harness and model improve each other.

4. Empirical demonstration across four benchmarks plus OOD evaluation: D-OPCD improves direct generation over the base generator and outperforms Vanilla SFT, Diffusion-DPO, and D-OPSD on the same replayed records, with additional T2I-CompBench++ out-of-distribution tests.

Main Findings

• The harness provides large gains at the generator interface. On the base generator, a skill-free harness H_empty raises the four-benchmark average from 60.52 to 79.28 (+18.76 points), including WISE from 42.86 to 78.00 and R2I-Bench from 42.87 to 68.12. ASE-evolved skills H_r* raise the average further to 81.83, with per-benchmark gains ranging from +1.29 on R2I-Bench to +4.12 on GenEval.

• Harness gains can be internalized into weights. D-OPCD raises direct-generation average from 60.52 to 65.09 (+4.57), improving all four benchmarks: GenEval 78.22 to 81.28, GenEval2 78.14 to 82.24, WISE 42.86 to 50.00, and R2I-Bench 42.87 to 46.83.

• The updated generator nearly matches the skill-equipped agent without any skills. Attaching the skill-free harness H_empty to the D-OPCD generator scores 82.33 on average, versus 81.83 for the original skill-equipped agent on the base generator, giving a signed internalization gap of δ(r*) = +0.50. On R2I-Bench the gap is δ(r*) = +4.27 (73.68 for the updated skill-free agent).

• Co-evolution continues after internalization. After clearing skills and rerunning ASE on the updated generator to obtain H_r†, the average rises to 84.16, above both the updated skill-free agent (82.33) and the original skill-equipped agent (81.83) — an additional 1.83 points over the skill-free harness. R2I-Bench instead declines slightly after skill evolution (69.45 versus 73.68).

• D-OPCD outperforms matched training baselines. On identical replayed records from H_r*, direct-generation averages were 60.52 (base), 63.38 for Vanilla SFT (+2.86), 62.18 for Diffusion-DPO (+1.66), 62.75 for D-OPSD (+2.23), and 65.09 for D-OPCD (+4.57). D-OPCD leads on all four benchmarks and is 1.71 points above Vanilla SFT.

• Task-matched teacher context matters. A teacher conditioned on [q_i; p_i] averages 65.09, versus 64.63 for p_i alone (which reaches the highest R2I-Bench score of 47.32) and 61.05 for shuffled prompts from other tasks, which reduces scores on all four benchmarks.

• More "useful" contexts are not automatically easier to internalize. An improvement filter (retaining pairs where the submitted image beats the direct-query image by verified pass count) leaves the four-benchmark mean nearly unchanged at 65.14, and a diversity adapter lowers it to 64.42. Pair counts differ (default: 276/400/490/392 across GenEval/GenEval2/WISE/R2I-Bench; filter: 43/264/319/314; adapter: 276/400/491/381).

• Out-of-distribution generalization is partial. On seven T2I-CompBench++ categories, Color rises from 77.45 (base) to 81.61 and 81.44 for GenEval- and GenEval2-trained checkpoints, and 2D-Spatial rises from 32.32 to 36.19 and 37.33. WISE- and R2I-Bench-trained checkpoints stay close to the base overall.

• Standalone internalization remains far short of the full harness. D-OPCD's direct-generation average of 65.09 is well below the 81.83 achieved by the original skill-equipped harness; the authors suggest evolved-skill prompts may be poorly aligned with the diffusion model's latent space for learning the same behavior from q_i alone.

Methodology in Plain English

The researchers keep a frozen image generator (Z-Image-Turbo) inside an agentic harness adapted from the GEMS architecture, which has five roles: Planner, Decomposer, Generator, Verifier, and Refiner. The Planner picks applicable skills and builds a first prompt; the Decomposer converts the original query into a fixed set of yes/no visual requirements; the Verifier checks each generated image against those requirements; and the Refiner revises the prompt using the image, the feedback, and task-local history until checks pass or the attempt budget runs out. The output is the prompt and image from the earliest attempt that satisfied the most requirements.

Task experience is turned into reusable guidance by ASE in three tiers: an episode is a full execution trajectory plus evaluator feedback; an insight records what the agent discovered that could have been given in advance; and a skill consolidates several insights into guidance for constructing the first prompt. Insights become "mature" once their support count reaches a threshold and exceeds their contradiction count, and a skill manager updates the library once enough mature, unabsorbed insights accumulate. Each committed library defines a harness version.

For each completed task, the system records the original query q_i and the harness-selected prompt p_i that produced the submitted image (the image is used only for selecting which records to keep). The generator is then trained with D-OPCD: a student model receives only q_i, while an exponential-moving-average teacher receives both q_i and p_i as structured context. The student rolls out its own denoising trajectory, and at each visited state the student's predicted velocity is pulled toward the teacher's (stop-gradient) prediction in squared L2 error. The student's trajectory states are treated as constants, so gradients do not flow through the sampling process. At inference, the harness prompt is hidden and the model generates from q_i alone. Training uses LoRA adapters under a matched optimization budget across all compared methods and an EMA teacher initialized from the same weights.

Harness versions are selected on a held-out set Q_hold, then replayed on Q_train to build the context dataset so that all contexts come from a single harness version operating on the generator that D-OPCD starts from. Tasks are split into disjoint Q_train, Q_hold, and Q_eval, with final results reported only on Q_eval. After training, the skill library and its episode/insight records are cleared so the harness can evolve fresh skills around the updated generator.

Why This Matters

Impact on research. The paper reframes agentic scaffolding and model weights as two ends of a single improvement loop rather than alternatives. It provides a concrete mechanism (context distillation with asymmetric conditioning) for transferring prompt-level agent experience into generative weights, and it introduces a measurable quantity — the signed internalization gap — for how much harness benefit has migrated into the model.

Real-world applications.

• Creative and design tooling: image generation products that keep improving from accumulated user-session experience without needing a heavy agent pipeline at inference time.
• Advertising and e-commerce asset generation: systems that must satisfy recurring compositional and constraint-heavy requirements while keeping per-request latency and cost low.
• Accessibility and consumer photo tools: turning repeated user refinement patterns into persistent model behavior so users get better first-attempt outputs.
• Domain-specific image pipelines (for example, where verification checklists are standard) that can replay verified prompt records as offline training data rather than re-deriving them live.

Industry relevance. Internalizing harness knowledge amortizes a one-time offline training cost over future requests and reduces the repeated reasoning, verification, and image-generation calls that agentic serving requires. That directly addresses inference cost and context bloat — the paper notes that as memories and skills accumulate, relevant guidance becomes harder to retrieve, and each added skill consumes context.

Future Directions

• Explaining the internalization shortfall: D-OPCD's standalone average of 65.09 versus the harness's 81.83 suggests evolved-skill prompts may not align well with the diffusion model's latent space; identifying which harness designs produce experience that is both immediately effective and learnable is left open.
• Longer co-evolution studies: the paper evaluates one base generator, one D-OPCD update, and one subsequent skill-evolution round on four benchmarks, and does not establish behavior across many update-and-reset cycles, across different generators, or across different task distributions. Retention of earlier gains, transfer to new tasks, and the tradeoff between offline training cost and recurring harness inference cost remain unmeasured.
• Preventing forgetting across repeated updates: the authors note that repeated model updates may require methods that avoid catastrophic forgetting while adapting the harness to each new generator.
• Extending beyond prompt-mediated gains: the current interface captures only prompt-mediated experience, not gains from verification or stochastic sample selection. Tool-using agents that search the web, retrieve reference images, or handle image-to-image tasks carry experience (tool choices, retrieved information, visual evidence) that a final text prompt cannot represent, raising the question of which capabilities should be internalized and which should stay in tools.

Target Audience

Researchers and engineers working on diffusion and flow-matching generative models, agentic image-generation systems, and distillation or preference-based fine-tuning for generative models. It is also relevant to practitioners building production T2I pipelines who care about inference cost and about converting accumulated agent experience into durable model capability. Readers without a background in diffusion sampling or on-policy distillation will find the method sections demanding, though the problem framing and headline results are accessible.

Authors’ abstract

Wrapping an image generation model in an agentic harness can effectively boost Text-to-Image task performance: the harness can leverage memory, skills, workflow orchestration, result verification, and iterative refinement to continually construct and revise prompts, thereby eliciting better images. These gains, however, remain external to the diffusion model and are realized only while the full harness runs. We propose Diffusion On-Policy Context Distillation (D-OPCD), which treats the agent-improved prompt as privileged context and distills the knowledge encoded in the agent harness into the weights of the diffusion model, so that the model retains part of the harness's benefit when conditioned on the original query alone. Using a Text-to-Image agent equipped with our proposed Auto Skill Evolver (ASE), we show that D-OPCD can internalize harness capabilities into the generator's weights, raising the average direct-generation score from 60.52 to 65.09 across four benchmarks. With this knowledge absorbed into the weights, the harness can shed its saturated skills and resume evolving: a second ASE round on the updated generator improves on a skill-free harness by additional 1.83 points, pointing toward text-to-image systems in which harness and model keep improving each other through continual co-evolution.