Optimal Control Meets Flow Matching: A Principled Route to Multi-Subject Fidelity

Abstract: 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.

• The paper introduces FOCUS, a method employing optimal control in flow matching models to disentangle multiple subjects in text-to-image generation.
• The test-time controller and Adjoint Matching fine-tuning reduce attribute leakage and identity entanglement, achieving up to 85% improvement over baselines.
• Empirical results on architectures like Stable Diffusion 3.5 and SDXL demonstrate robust multi-subject performance with minimal fine-tuning and computational overhead.

Optimal Control for Multi-Subject Fidelity in Flow Matching Models

Introduction

This paper presents a principled framework for improving multi-subject fidelity in text-to-image (T2I) generative models, particularly those based on flow matching (FM) and rectified-flow (RF) architectures. The authors address persistent failure modes in multi-subject prompts—attribute leakage, identity entanglement, and subject omission—by formulating subject disentanglement as a stochastic optimal control (SOC) problem over the sampling dynamics of trained FM models. The proposed approach yields two architecture-agnostic algorithms: a training-free test-time controller and a lightweight fine-tuning rule (Adjoint Matching), both of which are instantiated in the FOCUS (Flow Optimal Control for Unentangled Subjects) algorithm.

Flow Matching and Stochastic Optimal Control Formulation

Flow matching models parameterize generation as a time-dependent flow from a base distribution (e.g., Gaussian noise) to the data distribution via a learned vector field $v_\theta(x, t)$. Sampling is performed by integrating the learned ODE or SDE, producing a trajectory from noise to image. The authors leverage the FM framework to unify both modern (SD 3.5, FLUX) and classical (SDXL) architectures, enabling cross-architecture statements and interventions.

The core insight is to augment the base FM dynamics with a control $u(x, t)$, steering the trajectory to minimize a quadratic cost that penalizes both control magnitude and subject entanglement. The SOC objective is:

$$\min_{u} \mathbb{E} \left[ \int_0^1 \frac{1}{2} \|u(X_t^u, t)\|_2^2 + f(X_t^u, t) dt \right]$$

where $f$ is a differentiable disentanglement cost (e.g., FOCUS), and $X_t^u$ evolves under the controlled dynamics. The Hamiltonian formalism yields the optimal control in terms of the adjoint $a(t)$, which is approximated for efficient single-pass inference.

Algorithms: Test-Time Control and Adjoint Matching Fine-Tuning

Test-Time Controller

The test-time controller computes an instantaneous control update at each sampling step by locally approximating the adjoint. This yields a velocity reparameterization:

$$v_t^\star = v_{\text{base}}(X_t, t) - \frac{\sigma^2_{\text{mem}}(t)}{2}(1-t) \nabla_X f(X_t, t)$$

This update is compatible with any SDE/ODE solver and requires only the extraction of subject token indices and cross-attention maps. The controller is training-free and runs efficiently on commodity GPUs.

Figure 1: Optimal control makes flow matching models reliable on multi-subject prompts. Using FOCUS at test time or via fine-tuning yields faithful multi-subject compositions with correct attributes, minimal leakage, and no omissions, while preserving base style.

Adjoint Matching Fine-Tuning

Fine-tuning is performed by regressing a control network $u_\theta$ onto a lean adjoint signal computed along frozen forward trajectories, using a memoryless noise schedule to ensure generalization. The Adjoint Matching loss is:

$$\mathcal{L}_{\text{AM}}(\theta) = \frac{1}{2} \int_0^1 \|u_\theta(X_t, t) + \sigma_{\text{mem}}(t) \tilde{a}(t)\|^2 dt$$

This approach requires only text prompts and subject token indices during training, and preserves base model style and support.

FOCUS: Probabilistic Attention-Based Disentanglement

FOCUS introduces a probabilistic attention loss that treats cross-attention maps as distributions over spatial locations. The loss combines within-subject agreement (minimizing dispersion among attention maps for each subject) and between-subject separation (maximizing Jensen-Shannon divergence between subject means):

$$\text{FOCUS}(S) = \frac{1}{2} \left( \frac{1}{|S|} \sum_{s \in S} \widehat{D}_{\text{JS}}(P_s) \right) + \frac{1}{2} \left( 1 - \widehat{D}_{\text{JS}}(M) \right)$$

where $P_s$ is the set of attention maps for subject $s$, and $M$ is the set of subject means. This loss is differentiable and normalized, enabling direct optimization during sampling or fine-tuning.

Empirical Evaluation

Experiments are conducted on Stable Diffusion 3.5, FLUX, and SDXL, using a curated dataset of 150 multi-subject prompts. Metrics include image-text alignment (CLIP, SigLIP-2), caption-based faithfulness (BLIP, Qwen2-VL), and human preference studies. Both test-time control and fine-tuning consistently improve multi-subject fidelity over baselines and prior heuristics (Attend & Excite, CONFORM, Divide & Bind).

Figure 2: Effect of the control parameter $\lambda$ on test-time control with Stable Diffusion 3.5.

Fine-tuned models generalize from small prompt sets to unseen prompts, indicating that the attention-level failure mode is robust across subject categories and prompt structures. FOCUS achieves the highest composite scores and human preference win rates, with up to 85% relative improvement in composite metrics for SD 3.5.

Human Preference Study

A prompt-conditioned, pairwise preference study with 50 participants and 2,000 judgments demonstrates that FOCUS-controlled models are preferred over baselines and competing heuristics, both at test-time and after fine-tuning.

Figure 3: User interface for the prompt-conditioned, pairwise preference study.

Transfer to Classical Diffusion Models

The SOC formulation and FOCUS loss are shown to transfer to classical denoising diffusion models (SDXL) via a flow-diffusion correspondence, improving multi-subject fidelity without retraining.

Implementation and Resource Considerations

• Test-time control: Requires only subject token indices and cross-attention extraction; incurs $\sim$2x inference time overhead but is compatible with commodity GPUs (12GB VRAM).
• Fine-tuning: LoRA-based, trains $<0.1\%$ of parameters; fits within H100 VRAM; training time is 17–79 minutes depending on model.
• Generalization: Fine-tuned controllers trained on limited prompts generalize to unseen prompts and subject categories.
• Deployment: Test-time controller is suitable for plug-and-play use; fine-tuned models match base inference speed.

Implications and Future Directions

The control-theoretic framework provides a unified, optimizable objective for multi-subject fidelity, subsuming prior attention heuristics and extending to modern FM architectures. The strong generalization from minimal fine-tuning data suggests that multi-subject entanglement is an attention-level failure mode, motivating further research into annotation-free proxies and automated subject tokenization. The approach is compatible with both deterministic and stochastic sampling, and can be extended to other generative modalities.

Conclusion

The paper establishes a principled route to multi-subject fidelity in T2I models by formulating subject disentanglement as an optimal control problem over flow matching dynamics. The FOCUS algorithm, instantiated via test-time control or fine-tuning, consistently improves multi-subject alignment, reduces attribute leakage, and preserves base model style across architectures. The framework unifies and extends prior heuristics, and empirical results demonstrate strong gains in both automated metrics and human preference. Future work should explore annotation-free disentanglement objectives and further probe the attention-level mechanisms underlying multi-subject failures.