Diffusion Model Fine-Tuning with ART

SD3 and Wan2.2 generation models trained with DiffusionART.

Fine-tune SD3 or Wan2.2 with DiffusionART to improve text rendering using OCR rewards. ART combines rewards and clean latent endpoints within each prompt group into a shared direction, then trains against frozen velocity targets at states visited during generation. Complete the shared setup, choose an image or video recipe, and inspect the training results below.

Overview

Both recipes use LoRA with FSDP2 and deterministic ODE generation. They split each prompt’s candidates into two disjoint update windows that share targets frozen before the first update. Corrective masking suppresses updates that would move an already-drifted prediction farther outward. For the derivation and reference implementation, see the DiffusionART paper, official code, and project page.

Models

SD3 / Wan2.2

Algorithm

DiffusionART

Output

Image / Video

Reward

OCR / Video OCR

You’ll do: select a recipe → launch training → inspect OCR reward and generated samples.

Prerequisites: Shared installation, model, and dataset setup. The recorded Wan experiment used four NVIDIA A800 GPUs.

Tasks

Choose the config for the output you want to generate.

Model	Config	Task
SD3.5 Medium	sd3_art_ocr	Image OCR
Wan2.2 TI2V 5B	wan22_ti2v_5b_art_video_ocr	Text-to-video OCR

Observation and Action

The environment scores each completed candidate; video rewards are averaged before ART builds guidance for the prompt group.

Value	Meaning
Observation	A text prompt containing quoted target text.
Action	The generated image or video.
Reward	One scalar OCR similarity score in [0, 1] per candidate.

Installation and Assets

Follow the shared installation, model download, and OCR dataset download instructions. Both ART recipes set env.*.reward.use_gpu: False and use the CPU OCR installation.

Run Training

Open examples/diffusion/config/sd3_art_ocr.yaml or examples/diffusion/config/wan22_ti2v_5b_art_video_ocr.yaml before launching. In the selected YAML, set actor.model.model_path, env.train.dataset.path, and env.eval.dataset.path to your downloaded directories. Check cluster.num_nodes, cluster.component_placement, and the batch settings against your hardware and the ART constraints below.

After checking the YAML, run its config name from the repository root:

bash examples/diffusion/run_diffusion.sh sd3_art_ocr

bash examples/diffusion/run_diffusion.sh wan22_ti2v_5b_art_video_ocr

The launcher generates and scores candidates, updates LoRA parameters, and writes logs under logs/<timestamp>-<config>/.

Configure the ART Update

When changing rollout size or placement, each actor shard must receive complete prompt groups. algorithm.art_updates_per_rollout must be 1 or 2 and divide algorithm.group_size; env.train.group_size must match that group size.

Warning

Set actor.global_batch_size to env.train.total_num_envs * env.train.rollout_epoch divided by algorithm.art_updates_per_rollout. For 96 environments, 12 rollout epochs, and two updates, use actor.global_batch_size=576. Each per-rank update window must be divisible by actor.micro_batch_size. Environment counts must preserve complete groups across both env and actor ranks.

ART requires actor.training_backend=fsdp, rollout.pipeline_stage_num=1, algorithm.update_epoch=1, and runner.weight_sync_interval=1. Keep env.train.max_steps_per_rollout_epoch=actor.model.num_action_chunks so that each rollout epoch produces one complete candidate. SD3 uses actor.model.sd3.train_sampler=ode with actor.model.sd3.noise_level=0.0; Wan records ODE trajectories.

Main Update Options

Option	Effect
algorithm.art_timestep_fraction	Fraction of denoising states used for training: 0.5 for SD3 and 0.25 for Wan in the supplied recipes.
algorithm.art_guidance_gain	Scales reward guidance; both recipes use 1.0.
algorithm.art_corrective_mask_enabled	Enables the corrective movement mask; both recipes use True.
algorithm.art_mask_threshold	Movement threshold for the mask; both recipes use 1e-5 with art_mask_metric=sde_kl.

Expected Results

Watch env/video_ocr for Wan or env/avg for SD3 and inspect generated samples. actor/active_group_frac shows whether prompt groups supply guidance; actor/skipped_zero_guidance marks rounds with no usable guidance on any rank. Use Training Metrics for logging and TensorBoard setup. OCR reward alone does not measure visual quality or temporal consistency, and ART and NFT losses are not directly comparable.

The runs started on November 11, 2026 used four NVIDIA A800 80 GB GPUs, Wan2.2 TI2V 5B, LoRA (rank 32, alpha 64), FSDP2, BF16, and CPU OCR. The first table compares 40,320 generated trajectories: round 35 for ART and round 70 for NFT. For each run, reward is the mean env/video_ocr over the last 5760 trajectories at that budget. The NFT entry reports the mean and sample standard deviation across two runs; the ART entry is one run.

Wan Video OCR at 40,320 Generated Trajectories

Method	Runs	Trajectories per Round	Learning Rate	Training Reward	Elapsed Time (h)
ART	1	1152	1e-4	0.8584	10.01
NFT	2	576	4e-4	0.7565 ± 0.0544	11.29 ± 0.52

At this sample budget, ART’s reward is 13.5% higher than the NFT mean (an absolute increase of 0.1019), with 11.3% less elapsed training time. Elapsed time comes from the training-loop log and excludes model initialization. ART first reached a reward of 0.8 over the same 5760-trajectory window at 21,888 trajectories (5.48 h); neither NFT run reached that windowed threshold within 40,320 trajectories.

The NFT run with the same trajectories per round and learning rate as ART recorded 29 rounds before it was stopped. Compare both at 33,408 generated trajectories, using the same last-5760-trajectory reward window. Each entry below is one run.

Matched Trajectories per Round and Learning Rate

Method	Trajectories per Round	Learning Rate	Training Reward	Elapsed Time (h)
ART	1152	1e-4	0.8446	8.31
NFT	1152	1e-4	0.1076	8.23

With matched trajectories per round and learning rate, ART’s reward is 0.7370 higher (685.2% relative improvement), with similar elapsed time. This comparison covers only the first 29 rounds, not final convergence. These are training rewards, not held-out evaluation scores. The 576-trajectory NFT configuration in the first table also changes learning rate and reward reduction; that comparison does not isolate the algorithm’s effect.

To reproduce the configurations, use wan22_ti2v_5b_art_video_ocr for ART or wan22_ti2v_5b_nft_video_ocr for NFT with the model and data paths prepared above. Set cluster.component_placement to {"actor,env,rollout": "0,1,2,3"}, env.train.total_num_envs=96, algorithm.group_size=8, and runner.val_check_interval=-1. Set env.train.reward.use_gpu and env.eval.reward.use_gpu to False. Under actor.model.wan22_ti2v_5b, use max_generation_batch_size=2, compile_transformer_forward=True, and offload_auxiliary_modules=True; retain the recipes’ 512 × 512 resolution, five frames, ten denoising steps, and guidance scale of 1.0.

For ART, set env.train.rollout_epoch=12, actor.global_batch_size=576, actor.micro_batch_size=2, actor.optim.lr=1e-4, runner.save_interval=10, runner.max_epochs=70, and runner.max_steps=70. For NFT with 576 trajectories per round, use env.train.rollout_epoch=6, actor.global_batch_size=288, actor.micro_batch_size=1, actor.optim.lr=4e-4, runner.save_interval=-1, runner.max_epochs=70, and runner.max_steps=70. Set env.*.reward.reward_type to video_level for ART or frame_level for this NFT configuration.

For NFT with 1152 trajectories per round, use the NFT recipe with the ART overrides above, including video_level rewards, but set runner.max_epochs=200 and runner.max_steps=200; the recorded run was stopped after round 29. Both NFT configurations retain algorithm.nft_tau=[1.0,0.01,0,70].

All runs used actor.seed=42 and environment seeds of 42. These settings do not fix every global RNG used for LoRA initialization and sampling; repeat runs and record the revision, dependencies, and RNG seeds when comparing results.