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

TL;DR

This paper introduces a novel instance-level sampling schedule learning method for text-to-image models, improving alignment and efficiency without retraining the model weights.

Contribution

It proposes a Dirichlet policy-based approach with a James-Stein based reward baseline for better gradient estimation in policy learning.

Findings

Improves text-image alignment and compositional control across multiple models.

Achieves comparable quality to distilled models with fewer sampling steps.

Demonstrates the method's model-agnostic applicability as a post-training enhancement.

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.