PaperlayerView exact revision on arXiv

Source-linked AI summary

OneReward: Unified Mask-Guided Image Generation via Multi-Task Human Preference Learning

Yuan Gong, Xionghui Wang, Jie Wu, Shiyin Wang, Yitong Wang, Xinglong Wu

arXiv:2508.21066v1cs.CV

TL;DR

Mask-guided editing tasks share a conditioning format but differ in data distributions and evaluation criteria, and existing approaches often rely on task-specific fine-tuning. OneReward uses a single VLM for task- and criterion-aware preference learning, and Seedream 3.0 Fill performs strongly across multiple editing tasks and dimensions.

  • Problem

    Mask-guided editing tasks share an input format but differ in data distributions and evaluation metrics, while existing methods often rely on task-specific SFT.

    Source: §1 · p. 1

  • Method

    OneReward uses one VLM to make task- and criterion-specific pairwise judgments and guides multi-task reinforcement learning of Seedream 3.0 Fill directly from a pre-trained model.

    Source: §1 · p. 3, §4.3 · p. 9, §4.4 · p. 11

  • Results

    Seedream 3.0 Fill achieves state-of-the-art performance across image fill, image extend, object removal, and text rendering, outperforming commercial and open-source models on most dimensions.

    Source: §6 · p. 15

  • Takeaways & Limitations

    The results support using one reward model for multi-task, multi-dimensional reinforcement learning of image-generation models.

    Source: §1 · p. 3, §6 · p. 15

    The training pipeline maintains policy, reference, and EMA models in parallel, increasing memory use and engineering complexity; low-quality reference images may also yield overly easy preference pairs.

    Source: §5.3 · p. 14

Abstract

from arXiv · show

In this paper, we introduce OneReward, a unified reinforcement learning framework that enhances the model's generative capabilities across multiple tasks under different evaluation criteria using only \textit{One Reward} model. By employing a single vision-language model (VLM) as the generative reward model, which can distinguish the winner and loser for a given task and a given evaluation criterion, it can be effectively applied to multi-task generation models, particularly in contexts with varied data and diverse task objectives. We utilize OneReward for mask-guided image generation, which can be further divided into several sub-tasks such as image fill, image extend, object removal, and text rendering, involving a binary mask as the edit area. Although these domain-specific tasks share same conditioning paradigm, they differ significantly in underlying data distributions and evaluation metrics. Existing methods often rely on task-specific supervised fine-tuning (SFT), which limits generalization and training efficiency. Building on OneReward, we develop Seedream 3.0 Fill, a mask-guided generation model trained via multi-task reinforcement learning directly on a pre-trained base model, eliminating the need for task-specific SFT. Experimental results demonstrate that our unified edit model consistently outperforms both commercial and open-source competitors, such as Ideogram, Adobe Photoshop, and FLUX Fill [Pro], across multiple evaluation dimensions. Code and model are available at: https://one-reward.github.io

1 INTRODUCTION

Mask-guided editing tasks share an input format but differ in data distributions and evaluation criteria, while existing approaches often specialize through task-specific fine-tuning. OneReward uses a single VLM for task- and criterion-aware preference judgments, and Seedream 3.0 Fill applies multi-task reinforcement learning directly to a pre-trained model.

1 INTRODUCTION

  • Mask-guided tasks share an input format but differ in conditional distributions and evaluation metrics, and existing methods often rely on task-specific SFT or LoRA.Image fill, image extend, object removal, and text rendering also have distinct task objectives.Source: §1 · p. 1
  • DPO can struggle to rank outputs when evaluation dimensions conflict, while traditional multimodal reward models may require separate models for each criterion.ReFL can face reward conflicts across tasks, and GRPO-based methods rely on group-relative advantage estimation rather than explicitly maximizing reward signals.Source: §1 · pp. 2–3
  • OneReward conditions a single VLM reward model on task category and evaluation metric to make pairwise judgments about which output is better.Seedream 3.0 Fill is optimized with reinforcement learning from a pre-trained model without task-specific SFT.Source: §1 · p. 3
  • Seedream 3.0 Fill is presented as a unified model for image fill, image extend, object removal, and text rendering, outperforming leading commercial and open-source competitors.The paper also introduces FLUX Fill [dev][OneReward], which outperforms the original model on inpainting and outpainting.Source: §1 · p. 3

2 RELATED WORK

Prior work established diffusion-based mask-guided editing and preference alignment, but models often remain specialized or lack robust generalization across editing modalities. OneReward combines task-aware multi-task reinforcement learning with a unified image-editing model.

2 RELATED WORK

  • Diffusion models became state of the art for inpainting and outpainting, while later systems added instruction-based editing data and feedback learning.ByteEdit applies separate SFT and RL procedures across sub-tasks, and FLUX Fill is a strong open-source inpainting and outpainting baseline.Source: §2 · pp. 3, 5
  • Prior editing models are often specialized or lack robust generalization across distinct modalities, motivating a framework that unifies inpainting, outpainting, object removal, and text rendering.The paper builds on these foundations with multi-task RLHF.Source: §2 · p. 5
  • ReFL trains on expert comparisons and backpropagates reward scores as preference losses; VisionReward decomposes preferences into interpretable axes, DPO directly optimizes preferred/dispreferred log-probability ratios, and GRPO aligns diffusion and flow-matching models.Source: §2 · p. 5
  • OneReward uses one VLM to provide task-aware feedback, addressing DPO's difficulty distinguishing winners when preferences vary across evaluation dimensions.The framework jointly learns multiple image-editing sub-tasks.Source: §2 · p. 5

3 PRELIMINARIES

The preliminaries describe flow matching as learning a conditioned velocity field and RLHF as learning from human comparisons. They also introduce ReFL's direct gradient-based use of reward scores for fine-tuning diffusion models.

3.1 FLOW MATCHING

  • Flow matching trains a conditioned neural velocity field to approximate a target vector field, using a conditional regression loss.The conditioning information can include text or a binary mask.Source: §3.1 · pp. 5–6
  • Rectified Flow defines the target vector field as the constant direction from a noise sample to its corresponding data sample.The paper describes this as a more sampling-efficient special case of flow matching.Source: §3.1 · p. 6

3.2 REINFORCEMENT LEARNING FROM HUMAN FEEDBACK

  • RLHF uses pairwise human preferences to train a reward model that predicts scalar scores for generated outputs.The passage introduces the Bradley-Terry model for the probability of preferring a winner over a loser.Source: §3.2 · p. 6
  • ReFL randomly selects a late denoising step, scores the partially denoised image, and backpropagates a truncated reward through that step.This provides direct gradient-based fine-tuning using a pre-trained reward model.Source: §3.2 · p. 6

4 METHODOLOGY

The methodology combines human preference data across editing tasks and criteria with a VLM reward model and multi-task reinforcement learning. The resulting policy is trained to improve reward across multiple dimensions.

4.1 OVERVIEW

  • The pipeline partitions image-condition data by task, with task identifiers and metric sets conditioning a shared reward model's comparison of policy and reference outputs.Tasks can have distinct evaluation criteria.Source: §4.1 · p. 6

4.2 HUMAN PREFERENCE DATA COLLECTION

  • Each sample is a source-image/mask/prompt triplet; fill and extend use content prompts, while removal uses a fixed generic prompt and randomized inference settings diversify candidate outputs.The mask indicates the region to edit.Source: §4.2 · p. 7
  • Annotators select best and worst outputs separately for each dimension; metric-level pairs retain disagreements rather than forcing potentially misleading holistic labels.A candidate may be preferred on one dimension and disfavored on another.Source: §4.2 · p. 8

4.3 MULTI-DIMENSIONAL PAIRWISE REWARD MODEL TRAINING

  • OneReward uses a pretrained VLM with task-and-dimension queries to compare image pairs; the original prompt is included only for prompt-image alignment.The VLM's binary preference judgment is trained from human winner/loser pairs.Source: §4.3 · pp. 9–10
  • Text Alignment exceeds 80% accuracy for image fill and image extend; Removal Quality reaches 84.93% for object removal.Table 1 reports test-set accuracy for distinguishing preferred from less-preferred outputs.Source: Table 1 · p. 10, §4.3 · p. 10

4.4 MULTI-TASK REWARD FEEDBACK LEARNING

  • Training compares a fully denoised reference with a partially denoised policy image from the same condition, then maximizes bounded preference reward averaged across task criteria.The policy image uses a randomly selected denoising timestep, and the probability of the affirmative token supplies the reward signal.Source: §4.4 · pp. 10–11
  • Reward curves rise across all dimensions despite moderate fluctuations under multi-task training.The authors also report high final reward values across dimensions.Source: Figure 6 · p. 11, §4.4 · p. 12

5 EXPERIMENTS

Seedream 3.0 Fill performs strongly across the reported editing tasks, while OneReward provides a shared reward model without task-specific SFT. The experiments also compare a dynamic EMA-reference strategy and RL-enhanced visual results.

5.1 IMPLEMENTATION DETAILS

  • The dataset pipeline selects a diverse subset with CLIP embeddings and K-means, then creates human preference pairs for reward-model and generation-model training.The reward model fine-tunes Qwen2.5-VL-7B-Instruct; image-fill, image-extend, and object-removal samples use probabilities of 50%, 25%, and 25%.Source: §5.1 · p. 12

5.2.1 EVALUATION SETTINGS

  • The benchmark covers image fill, object removal, and text-guided and text-free image extend, comparing against commercial APIs and open-source models.It contains 130 fill images, 100 removal images, and 200 extend images.Source: §5.2.1 · p. 12

5.2.2 HUMAN EVALUATION

  • Seedream 3.0 Fill leads across tasks: fill usability is 69.04% versus 52.11% for the second-best competitor, text-free extend usability is 87.54%, and removal usability and quality are 82.22% and 86.33%.Style consistency is the exception in image fill, where Ideogram has a slight advantage.Source: Table 2 · p. 13, §5.2.2 · pp. 13–14
  • Forty participants rated outputs on quality and task-specific dimensions, using binary success rates for usability, text rendering, and removal quality and 1–5 MOS for other dimensions.Higher MOS scores indicate better quality.Source: §5.2.2 · p. 13
  • Good–Same–Bad evaluation gives the OneReward variant a higher proportion of “Good” ratings than the base model in every task.The reported results indicate that the unified reward model generally shifts outputs toward preferred generations.Source: Figure 7 · p. 13, §5.2.2 · p. 14

5.3 DYNAMIC REINFORCEMENT LEARNING

  • Dynamic training reuses the EMA model as the reference, reducing memory overhead while providing a progressively stronger baseline and more stable, adaptive reward signals.The standard pipeline maintains policy, reference, and EMA models in parallel; insufficient-quality reference images may yield overly easy pairs and potentially hinder learning through reward hacking.Source: §5.3 · pp. 14–15

5.4 QUALITATIVE RESULTS

  • In Fig. 9, OneReward RL-enhanced models based on Flux Fill [dev] and trained with Alg. 1 and Alg. 2 show superior visual quality to the base model and closed-source API.Source: §5.4 · p. 15, Figure 9 · p. 15

6 CONCLUSIONS

The conclusion presents OneReward as a unified reward model for multi-task, multi-dimensional reinforcement learning and Seedream 3.0 Fill as a strong multi-task editing model. It also notes a style-consistency weakness and describes an enhanced FLUX variant released for broader use.

  • OneReward uses a VLM as a unified reward model for multi-task, multi-dimensional reinforcement learning of diffusion and flow-matching models.Source: §6 · p. 15
  • Seedream 3.0 Fill achieves strong performance on image fill, image extend, object removal, and text rendering, outperforming commercial APIs and open-source models on most dimensions.Source: §6 · p. 15
  • Style consistency remains a relative weakness, while the authors identify detailed data annotation and post-training refinement as next steps.The conclusion also says an RL-trained FLUX Fill [dev] variant will be open-sourced.Source: §6 · p. 15

A APPENDIX

The appendix outlines dynamic multi-task reinforcement learning, including task sampling, preference comparisons, policy updates, and an EMA reference. Its schematic contrasts this dynamic reference with a fixed baseline.

  • Figure 10 schematically contrasts rewards measured against a fixed reference with rewards computed against an EMA-updated reference; it is conceptual, not a depiction of actual parameters or training dynamics.Source: Figure 10 · p. 18
  • The algorithm takes multi-task image-condition datasets, task sampling probabilities, task IDs, and evaluation dimensions as inputs.Source: A APPENDIX · p. 18
  • For each iteration, it samples a task condition and noise, generates reference and evaluated images, and forms queries for each evaluation dimension.Source: A APPENDIX · p. 18
  • The reward model scores each comparison for policy updating, after which the reference is updated as an EMA of the policy.Source: A APPENDIX · p. 18
Loading 2508.21066v1…