PaperlayerView exact revision on arXiv

Source-linked AI summary

DualHeads: Reward-tilted few-step diffusion made simple

Yuanzhi Zhu, Eleftherios Tsonis, Simo Alami Chehboune, Xi Wang, Stéphane Lathuilière, Vicky Kalogeiton

arXiv:2610.28858v1cs.CVcs.AIcs.LG

TL;DR

Few-step diffusion generators can produce high-quality images, but adapting them to task requirements and preferences is challenging when existing methods need signals a distilled generator may not provide. DualHeads uses reward and reference classifiers on a frozen backbone to guide generation. In four steps, it reaches OCR and GenEval scores comparable to multi-step RL methods while improving quality and preference scores over the base model.

  • Problem

    Adapting few-step generators to task requirements is difficult because existing approaches may rely on differentiable rewards, score fields, or stochastic trajectories that distilled generators may lack.

    Source: §1 · p. 2

  • Method

    DualHeads trains reward and reference classifier heads on a frozen diffusion backbone and updates the generator to maximize their weighted logits.

    Source: §1 · p. 3

  • Results

    With four steps, DualHeads reaches OCR 0.92 and GenEval 0.90, comparable to RL methods with 40–80 NFE, while improving quality and preference scores over the base model.

    Source: §5 · p. 9

  • Takeaways & Limitations

    At optimal heads, the weighted-logit objective is KL-regularized maximization of the learned reward, with the head-weight ratio setting the regularization strength.

    Source: §1 · p. 3

    The analysis assumes optimal heads and a fixed reward head, while training uses approximately optimal heads and refits the reward head; the alternating updates have no convergence guarantee.

    Source: A LIMITATIONS AND FUTURE WORK · p. 15

Abstract

from arXiv · show

Distilled diffusion models generate high-quality images in a few sampling steps, but aligning them with task requirements and human preferences is difficult. Methods developed for multi-step diffusion models rely on differentiable rewards, or on score fields and stochastic trajectories that distilled generators may lack. We ask whether a few-step generator can be adapted using only external reward evaluations and reference samples, and show that two binary classifiers suffice. Our method, DualHeads, implements them as lightweight heads on a frozen diffusion backbone. The reward head learns which samples of the same prompt score higher, the reference head separates reference images from generated ones, and the generator maximizes a weighted sum of their logits. When the heads are optimal, this objective is KL-regularized maximization of the learned reward, and the ratio of head weights sets the regularization strength. Empirically, with both rule-based and model-based rewards, a few-step DualHeads generator reaches task rewards close to those of 40-step reinforcement learning methods while preserving image quality, at no extra inference cost. We hope that this simple and elegant dual-head formulation will help advance few-step reward-tilted generation.

1 INTRODUCTION

Few-step generators offer fast, high-quality samples, but adapting them to task rewards is difficult when methods require gradients, scores, or stochastic trajectories the generator may not provide. DualHeads uses two classifiers to guide a four-step generator, reaching rewards close to 40-step RL methods while preserving quality.

1 INTRODUCTION

  • Prior few-step and multi-step reward methods often require trajectory likelihoods, score fields, differentiable rewards, or teacher scores beyond samples and black-box evaluations.Few-step stochastic policy-gradient updates can have high variance, while forward-process objectives need fields a distilled generator may not preserve.Source: §1 · p. 2
  • The central question is whether few-step generators can be tilted toward higher-reward regions while remaining close to reference images using only samples, black-box rewards, and reference examples.Distillation can yield high-quality samples in a handful of network evaluations, but those samples may still miss task requirements.Source: §1 · p. 2
  • DualHeads raises OCR accuracy from 0.59 to 0.92 and GenEval from 0.63 to 0.90 over the base model, and gives the best HPSv3 on Z-Image.The paper reports these results in a four-step setup while preserving generation quality.Source: §1 · p. 2
  • DualHeads adapts few-step generators from reward evaluations and reference samples alone using two lightweight classifier heads on a frozen diffusion backbone.The reward head learns within-prompt reward preferences; the reference head distinguishes reference images from generated predictions, and the generator maximizes their weighted logits.Source: §1 · pp. 2–3
  • With optimal heads, the weighted-logit generator objective is reverse-KL-regularized reward maximization, with the head-weight ratio setting regularization strength.This classification-based formulation replaces direct optimization of the black-box reward and the generator-to-reference density ratio.Source: §1 · pp. 2–3
  • With rule-based and model-based rewards, a four-step student reaches task rewards close to 40-step RL methods while preserving quality.The reported evaluations include OCR, GenEval, and human-preference rewards.Source: §1 · p. 3

2 RELATED WORK

Prior work addresses diffusion distillation, reward optimization, and reward-guided distillation through different supervision and fidelity mechanisms. DualHeads post-trains a distilled generator with online reward and reference heads on frozen diffusion features.

2 RELATED WORK

  • Distillation methods reduce diffusion sampling to a few large trajectory jumps or match student and teacher output distributions using score differences; DualHeads builds on a distilled model.Its reference head follows LADD’s use of heads on frozen diffusion features.Source: §2 · p. 3
  • Prior reward-optimization methods pass reward signals through differentiable rewards, policy gradients, or forward-process objectives; DualHeads instead learns an explicit classifier that guides the student through its input gradient.The cited forward-process methods use rewards as sample weights or regression targets rather than learning an explicit reward model.Source: §2 · p. 3
  • DualHeads post-trains a distilled generator and refits its reward head online from pointwise group-relative black-box labels, while its reference head keeps outputs close to the reference pool.Reward Lightning is architecturally closest, but its preference-pair-trained reward head remains frozen.Source: §2 · p. 3
  • In the generator update, the N-step rollout is followed by one differentiable prediction whose weighted reference- and reward-head logits define the utility to maximize.The backbone and both heads stay fixed during this generator update.Source: Figure 2 · p. 4

3 METHOD

DualHeads adapts few-step generators using black-box reward evaluations and reference samples, replacing unavailable reward gradients and density ratios with two classifier heads on a frozen backbone. The generator maximizes their weighted logits, which at optimal heads corresponds to a KL-regularized reward objective.

3.1 BACKGROUND

  • The target maximizes expected reward while penalizing reverse KL to the reference distribution, yielding a reference distribution reweighted by exponentiated reward.β controls reference regularization relative to reward.Source: §3.1 · p. 4

3.2 DUALHEADS

  • DualHeads uses reference and reward classifiers as surrogates because the black-box reward and generator-to-reference density ratio are not directly tractable.Both lightweight heads share a frozen diffusion backbone.Source: §3.2 · pp. 4–5
  • The reference head distinguishes noised references from generator predictions, while the reward head predicts soft optimality labels derived from rewards within each prompt group.The labels are computed from group-relative rewards; the reward head is trained with binary cross-entropy.Source: §3.2 · p. 5

3.3 GENERATOR OBJECTIVE

  • With optimal heads, the weighted-logit generator objective is KL-regularized reward maximization, with the ratio λD/λR setting regularization strength.The reference-head term penalizes reverse KL, while the reward-head term favors larger reward logits.Source: §3.3 · pp. 5–6
  • Training alternates classifier and generator updates, then discards both heads so inference uses only the fine-tuned few-step generator at no extra inference cost.The generator update ascends the weighted classifier logits while the heads are fixed.Source: §3.3 · p. 6

4 EXPERIMENTS

DualHeads is evaluated with few-step generators across text rendering, compositional generation, and preference rewards, then ablated for reward strength, reward-head warm-up, and reference-pool choice. It improves task scores while retaining image quality, with ablations showing that the reference pool can also shape properties not measured by the reward.

4.1 EXPERIMENTAL SETUP

  • Four-step training uses SD3.5-M and Z-Image with OCR, GenEval, or preference rewards, and evaluates models on OCR and GenEval test sets as well as DrawBench prompts.The default setup warms up the reward head for 1,000 steps; the main results use prompt groups of K=24 and four sampling steps.Source: §4.1 · pp. 6–7

4.2 TASK PERFORMANCE

  • 0.92 OCR accuracy and 0.90 GenEval are reached with four-step DualHeads, close to multi-step methods while improving quality and preference scores over the base model.For OCR, only TDM-R1 (0.95) and REST (0.96) score higher among few-step methods; on GenEval, DualHeads is slightly below few-step baselines at 0.92–0.94.Source: Table 1 · p. 7, §4.2 · pp. 7–8
  • 13.93 HPSv3 is reached by four-step DualHeads on Z-Image-Turbo, above TDM-R1 (12.67) and the 100-NFE base model (9.93).It also leads on PickScore (23.03) and HPSv2.1 (0.314), while Aesthetic and CLIPScore drop slightly; on SD3.5-M, the multi-reward setup improves all five metrics over baseline but trails MeanFlowNFT and RTDMD on most.Source: Table 2 · p. 8, §4.2 · p. 8
  • DualHeads renders the five OCR texts correctly and often places GenEval objects in scene context, unlike several baselines; Nano Banana makes samples more photographic.The Nano Banana pool keeps OCR at 0.92 but lowers GenEval to 0.73.Source: §4.2 · p. 8, Figure 3 · p. 8

4.3 ABLATIONS

  • 0.907 OCR at βeq = 0.1 exceeds 0.888 at 0.25 and 0.862 at 1; lower βeq increases reward weighting, with no observed collapse at 0.1.The two βeq = 0.25 runs agree within 0.015 across sixteen checkpoints despite halving both absolute head weights.Source: §4.3 · p. 9
  • A freshly initialized reward head starts 0.497 and 0.459 below a warmed-up head, but recovers by step 3000 and converges within 0.003 and 0.004.The warm-up therefore speeds training without changing the final result at either tested weight setting.Source: §4.3 · p. 9
  • With a Nano Banana reference pool, OCR stays at 0.92 and aesthetic score rises from 5.46 to 5.62, while GenEval falls from 0.90 to 0.73.ImageReward, PickScore, and HPS barely change, whereas CLIPScore drops from 28.75 to 26.86; the reference pool largely sets properties left unspecified by the reward.Source: §4.3 · p. 9

5 CONCLUSION

DualHeads uses two lightweight classifier heads on a frozen diffusion backbone to adapt few-step generators from external rewards and reference samples. The conclusion reports strong four-step OCR and GenEval results and notes that the reference pool shapes properties not specified by the reward; the appendix points to limitations, experimental details, and further results.

5 CONCLUSION

  • DualHeads converts KL-regularized reward maximization into two binary classification problems and reaches OCR 0.92 and GenEval 0.90 with four steps.The conclusion reports performance comparable to RL methods with 40–80 NFE and improvements over the base model on all quality and preference scores.Source: §5 · p. 9
  • The reward head sets what the reward measures, while the reference pool largely determines other properties, such as style..Source: §5 · p. 9
  • The authors report using generative AI tools for drafting, literature review, implementation and experiment inspection, and interpretation; they say they reviewed the material and checked citations and reported numbers.Source: AI USE STATEMENT · p. 10

Appendix for DualHeads

A LIMITATIONS AND FUTURE WORK

The paper identifies limitations in its idealized analysis, learned-reward dependence, and reference-pool fidelity, and lists convergence analysis and broader evaluation as future work.

A LIMITATIONS AND FUTURE WORK

  • The analysis assumes optimal heads and a fixed reward head, but practical heads are approximate and the refit reward head moves the target; alternating updates have no convergence guarantee.Source: A LIMITATIONS AND FUTURE WORK · p. 15
  • Training against the reward head inherits both its errors and the external reward's flaws, while the reference head cannot identify reward exploitation.Source: A LIMITATIONS AND FUTURE WORK · p. 15
  • Because fidelity is measured against the reference pool, generated samples also acquire reward-ignored properties such as style and composition.Proposed directions include richer reference pools, convergence analysis, and applications during distillation, at higher resolutions, or to few-step video.Source: A LIMITATIONS AND FUTURE WORK · p. 15

B EXPERIMENTAL DETAILS

The experimental details passage specifies inference and compute settings for the reported experiments.

B EXPERIMENTAL DETAILS

  • The experiments use EMA weights for inference; ablations run on one NVIDIA H100 GPU and main runs on four.The text specifies that its values describe the SD3.5-M OCR configuration.Source: B EXPERIMENTAL DETAILS · p. 15

B.1 TRAINING CYCLE

The training cycle scores grouped generator rollouts, pairs them with prompt-matched reference images, and updates the generator through a single recomputed prediction. The appendix details lightweight patch-logit heads, optimization settings, and reward definitions for the evaluated tasks.

B.1 TRAINING CYCLE

  • Each prompt’s four-step rollout endpoints are scored, and within-group reward labels train the reward head on re-noised samples.Reference images are paired per prompt, re-noised over the same noise range, and used as positive examples; recomputed generator predictions provide negative examples and a pathwise utility gradient.Source: §B.1 · p. 15

B.2 HEADS DETAILS

  • For SD3.5-M, both heads use features from four frozen-transformer blocks and output 32 × 32 patch-logit maps.Each level has its own convolutional head, and both binary-cross-entropy losses and the utility average over patches.Source: §B.2 · pp. 15–16

B.3 OPTIMIZATION AND HYPERPARAMETERS

  • The generator maintains an EMA for inference, while utility weights are swept to vary the reward and reference contributions.Training rollouts use raw generator weights; the appendix gives example weight pairs (1, 10) and (0.5, 5).Source: §B.3 · p. 16

B.4 REWARDS

  • OCR reward compares recognized text with the prompt’s target string, while GenEval scores object counts and, when counts match, color or position.The GenEval dense training reward averages credits, whereas its strict score requires every requirement to be met exactly.Source: §B.4 · p. 16
  • The multi-reward setup equally weights PickScore, CLIPScore, and HPSv2.1, while the Z-Image setup trains on HPSv3.The multi-reward setting uses Pick-a-Pic prompts and images; the Z-Image setting evaluates on DrawBench prompts.Source: §B.4 · p. 17, Table 4 · p. 17

C.1 MULTI-REWARD TRAINING

The multi-reward experiment trains SD3.5-M on three preference and image-scoring rewards, and the accompanying GenEval results are also broken down by category.

C.1 MULTI-REWARD TRAINING

  • SD3.5-M is trained jointly on PickScore, CLIPScore, and HPSv2.1 using Pick-a-Pic prompts and reference images, with evaluation on DrawBench prompts.Table 4 reports the results for this multi-reward setting.Source: §C.1 · p. 17

C.2 PER-CATEGORY GENEVAL RESULTS

  • Table 5 breaks down the overall GenEval scores by category.The categories include single-object, two-object, counting, colors, position, and attribute binding.Source: §C.2 · p. 17

C.3 QUALITATIVE RESULTS

Qualitative comparisons show how reward weighting and reference choice affect generated images, while two-step sampling retains reward close to four-step sampling. The section also presents qualitative comparisons and uncurated samples for Z-Image-based models.

C.3 QUALITATIVE RESULTS

  • With λD = 1 and C = 2.5, increasing λR across {2.5, 5, 10, 25} makes text larger and easier to read; at λR = 25 it tends to dominate the image, while scene layout stays similar across rows.Source: §C.3 · pp. 17–18
  • Two-step samples retain reward close to four-step samples for both OCR- and GenEval-trained generators.The generators were trained with four sampling steps; the figure compares both step counts on the same prompts.Source: §C.3 · p. 18, Figure 6 · p. 19
  • Changing the reference pool changes generator style but not the content rendered for the same prompts.The figure contrasts base-model references with Nano Banana references.Source: §C.3 · p. 18, Figure 7 · p. 19
  • The section includes a qualitative comparison of the Z-Image base model with its distilled and RL-trained variants, alongside uncurated GenEval and reward-training samples.The uncurated samples are drawn from evaluation runs with prompts sampled uniformly without replacement.Source: §C.3 · p. 18, Figure 8 · p. 19, Figure 9 · p. 20

C.4 SAMPLES FROM THE REFERENCE DATASET

Reference images come from Nano Banana or teacher-simulated SD3.5-M samples selected by reward thresholds. The sources have similar high reward among retained references but differ in prompt coverage, and low scores can reflect evaluator errors as well as genuine image failures.

C.4 SAMPLES FROM THE REFERENCE DATASET

  • Teacher references cover 93–97% of prompts, compared with 62–63% for Nano Banana, while retained references have mean reward r of 0.90–0.98.References are each prompt’s highest-reward sample above the OCR or GenEval threshold.Source: §C.4 · p. 21
  • Low-reward samples include correct images misjudged by OCR or dense GenEval, as well as teacher samples with genuine missing or illegible text.Examples include stylized text OCR fails to read and cluttered scenes whose extra detections depress the dense GenEval score.Source: §C.4 · p. 21
  • Figure 12 shows random Nano Banana samples from OCR and GenEval prompts, annotated with reward and target text or prompt.The figure distinguishes OCR training prompts from GenEval training prompts.Source: Figure 12 · p. 21

C.5 ABLATIONS OF THE HEADS AND OF THE GENERATOR LOSS TYPE

The head ablations show that removing the reference head leads to OCR-training collapse, while removing the reward head leaves samples natural but yields only a slight reward increase. Figure 16 also includes comparisons of three generator-loss types.

C.5 ABLATIONS OF THE HEADS AND OF THE GENERATOR LOSS TYPE

  • Without the reference head, OCR reward rises initially then collapses to zero after about 900 training steps.Without the reward head, samples remain natural, but the reference head alone produces only a slight training-reward increase, likely because the reference pool is high quality.Source: §C.5 · p. 21
  • Figure 16 also evaluates three classic generator-training loss types.The passage does not specify their comparative outcomes.Source: §C.5 · p. 21, Figure 16 · p. 22
  • About 29% of lowest-reward Nano Banana OCR images and 20% of teacher-simulated OCR samples receive the minimum reward of zero.The accompanying figures also show low-scoring GenEval samples; the teacher-simulated examples are limited to at most one sample per prompt.Source: Figure 14 · p. 22, Figure 15 · p. 22

Related papers

Loading 2610.28858v1…