Optimal Control Meets Flow Matching: A Principled Route to Multi-Subject Fidelity
This addresses a key limitation in text-to-image generation for users needing accurate multi-subject depictions, though it is incremental as it builds on existing flow matching and diffusion models.
The paper tackles the problem of multi-subject fidelity in text-to-image models, which often suffer from attribute leakage and subject omissions, by introducing a theoretical framework based on stochastic optimal control to steer sampling dynamics, resulting in algorithms that improve multi-subject alignment while maintaining base-model style across models like Stable Diffusion 3.5 and FLUX.
Text-to-image (T2I) models excel on single-entity prompts but struggle with multi-subject descriptions, often showing attribute leakage, identity entanglement, and subject omissions. We introduce the first theoretical framework with a principled, optimizable objective for steering sampling dynamics toward multi-subject fidelity. Viewing flow matching (FM) through stochastic optimal control (SOC), we formulate subject disentanglement as control over a trained FM sampler. This yields two architecture-agnostic algorithms: (i) a training-free test-time controller that perturbs the base velocity with a single-pass update, and (ii) Adjoint Matching, a lightweight fine-tuning rule that regresses a control network to a backward adjoint signal while preserving base-model capabilities. The same formulation unifies prior attention heuristics, extends to diffusion models via a flow-diffusion correspondence, and provides the first fine-tuning route explicitly designed for multi-subject fidelity. Empirically, on Stable Diffusion 3.5, FLUX, and Stable Diffusion XL, both algorithms consistently improve multi-subject alignment while maintaining base-model style. Test-time control runs efficiently on commodity GPUs, and fine-tuned controllers trained on limited prompts generalize to unseen ones. We further highlight FOCUS (Flow Optimal Control for Unentangled Subjects), which achieves state-of-the-art multi-subject fidelity across models.