Designing Instance-Level Sampling Schedules via REINFORCE with James-Stein Shrinkage

arXiv:2511.22177v2 Announce Type: replace-cross Abstract: Most post-training methods for text-to-image samplers focus on model weights: either fine-tuning the backbone for alignment or distilling it for few-step efficiency. We take a different route: rescheduling the sampling timeline of a frozen sampler. Instead of a fixed, global schedule, we learn instance-level (prompt- and noise-conditioned) schedules through a single-pass Dirichlet policy. To ensure accurate gradient estimates in high-dimensional policy learning, we introduce a novel reward baseline based on a principled James-Stein estimator; it provably achieves lower estimation errors than commonly used variants and leads to superior performance. Our rescheduled samplers consistently improve text-image alignment including text rendering and compositional control across modern Stable Diffusion and Flux model families. Additionally, a 5-step Flux-Dev sampler with our schedules can attain generation quality comparable to deliberately distilled samplers like Flux-Schnell. We thus position our scheduling framework as an emerging model-agnostic post-training lever that unlocks additional generative potential in pretrained samplers.

Related

• Reward Score Matching: Unifying Reward-based Fine-tuning for Flow and Diffusion Models
• SOAR: Self-Correction for Optimal Alignment and Refinement in Diffusion Models
• SafeMERGE: Preserving Safety Alignment in Fine-Tuned Large Language Models via Selective Layer-Wise Model Merging
• Diagnosing and Improving Diffusion Models by Estimating the Optimal Loss Value

Source: arXiv cs.CV | 2026-04-28