💻 computer science

FOCUS: Optimal Control for Multi-Entity World Modeling in Text-to-Image Generation

This paper introduces FOCUS, a principled framework that casts flow matching as stochastic optimal control to derive training-free and fine-tuning algorithms that significantly improve multi-subject fidelity and prevent attribute leakage in text-to-image generation.

Original authors: Eric Tillmann Bill, Enis Simsar, Thomas Hofmann

Published 2026-03-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Eric Tillmann Bill, Enis Simsar, Thomas Hofmann

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a master chef (the AI) who is amazing at cooking single dishes. If you ask for "a perfect steak," you get a masterpiece. But if you ask for "a steak, a salad, and a glass of wine all on one plate," the AI often gets confused. It might paint the steak with the salad's dressing, turn the wine into a red steak, or forget the wine entirely. The ingredients get "entangled" or mixed up in a messy way.

This paper, titled FOCUS, introduces a new way to teach the AI how to keep these ingredients separate and distinct, even when they are all on the same plate.

Here is the breakdown using simple analogies:

1. The Problem: The "Messy Kitchen"

Current AI image generators (like Stable Diffusion or FLUX) are great at following instructions, but they struggle with multi-entity prompts (requests with multiple distinct objects).

Attribute Leakage: You ask for a "red hat" and a "blue car," and the AI paints the car red.
Identity Entanglement: You ask for a "dog" and a "cat," and the AI draws a creature that is half-dog, half-cat.
Subject Omission: You ask for three items, and the AI only draws two.

The AI treats the whole image as one big, blurry soup rather than distinct objects with their own identities.

2. The Solution: The "Traffic Cop" (Optimal Control)

The authors realized that generating an image is like driving a car from point A (random noise) to point B (the final picture). Usually, the AI just follows the GPS (the base model) blindly.

FOCUS acts like a Traffic Cop or a Co-Pilot that steps in during the drive.

It doesn't rebuild the car (the AI model).
It doesn't need to retrain the driver for years.
Instead, it gently steers the car away from "traffic jams" (where objects mix up) and keeps the lanes clear.

They call this Optimal Control. It's a mathematical way of saying: "Make the smallest possible correction at every step to ensure the red hat stays on the head and the blue car stays blue."

3. The Two Ways to Use the Co-Pilot

The paper offers two ways to use this new steering system:

A. The "Instant Fix" (Test-Time Control)

How it works: You don't change the AI model at all. When you type a prompt, the system adds a tiny, real-time "nudge" to the AI's thinking process.
Analogy: It's like wearing noise-canceling headphones while listening to a song. The song (the image) is still the same, but the headphones (FOCUS) filter out the static (the confusion) so you hear the music clearly.
Pros: Works immediately on any computer; no waiting for training.
Cons: It takes a little longer to generate the image because the system has to do extra math on the fly.

B. The "Muscle Memory" (Fine-Tuning)

How it works: The system teaches the AI a new habit. It runs through thousands of examples, learning exactly how to keep objects separate. Once trained, the AI "remembers" this skill.
Analogy: It's like a musician practicing scales. At first, they have to think hard about every note (Test-Time). But after practice, their fingers just know where to go (Fine-Tuning). They can play the song faster and more naturally.
Pros: The image generation is just as fast as before, but the quality is permanently better.
Cons: Requires some computing power to do the initial training.

4. The Secret Sauce: "FOCUS" (The Attention Map)

How does the Traffic Cop know when to steer? It uses a new rule called FOCUS (Flow Optimal Control for Unentangled Subjects).

Imagine the AI has a set of spotlights (called "attention maps") that shine on different parts of the image based on the words in your prompt.

Without FOCUS: The spotlight for "dog" and the spotlight for "cat" might overlap, causing the AI to mix them up.
With FOCUS: The system checks the spotlights. If the "dog" light is shining on the "cat," it says, "No, move that light over here!" It mathematically forces the spotlights to stay in their own lanes.

It treats these spotlights like probability clouds. The goal is to make the "dog cloud" and the "cat cloud" distinct and non-overlapping, ensuring the AI knows exactly where each object belongs.

5. Why This Matters

This isn't just about making prettier pictures of dogs and cats. It's about World Modeling.

Storytelling: Imagine generating a comic book where a hero, a villain, and a sidekick are all in the same scene, each with their own unique clothes and weapons.
Science: Visualizing a cell with a nucleus, mitochondria, and ribosome, where each part is distinct and labeled correctly.
Design: Creating a room with a specific chair, lamp, and rug, where the lamp doesn't accidentally turn into a chair.

Summary

FOCUS is a smart steering wheel for AI image generators. It uses math to gently guide the AI, ensuring that when you ask for multiple things, the AI gives you exactly what you asked for—distinct, separate, and perfectly bound to their own descriptions—without losing the artistic style of the original model. Whether you need a quick fix or a permanent upgrade, FOCUS provides the tools to keep the AI's world organized.

1. Problem Statement

Current Text-to-Image (T2I) models, particularly those based on Flow Matching (FM) like Stable Diffusion 3.5 and FLUX.1, excel at single-entity generation but struggle with multi-subject prompts. Common failure modes include:

Attribute Leakage: Attributes (e.g., color, material) from one subject incorrectly applying to another.
Identity Entanglement: Subjects merging into a single hybrid entity.
Subject Omission: One or more subjects in the prompt failing to appear in the generated image.

Existing solutions are largely heuristic (e.g., manipulating cross-attention maps at test time) and lack a unified theoretical framework. They often require extensive hyperparameter tuning, are specific to older diffusion architectures (U-Net), and do not generalize well to modern Flow Matching models.

2. Methodology

The authors propose a principled framework that reframes multi-subject generation as a Stochastic Optimal Control (SOC) problem.

A. Theoretical Framework: Stochastic Optimal Control (SOC)

The generation process in Flow Matching is modeled as a stochastic differential equation (SDE). The authors introduce a control term $u_t$ to the drift of the SDE to minimize a cost function $f$ (representing subject entanglement) while staying close to the base model.

Objective: Minimize the expected cost of entanglement plus the energy of the control.
Derivation: Using the Hamilton-Jacobi-Bellman equation, they derive the optimal control $u^*_t = -\sigma(t)a(t)$ , where $a(t)$ is the adjoint (co-state) signal.
Approximation: Since solving the full adjoint dynamics is computationally prohibitive at test time, they approximate the adjoint locally, leading to a gradient-based guidance term:
$v^*_t \approx v_{base}(X_t, t) - \eta(t) \nabla_X f(X_t, t)$
This unifies prior attention heuristics under a single, architecture-agnostic control interface.

B. The FOCUS Objective

To instantiate the cost function $f$ , the authors introduce FOCUS (Flow Optimal Control for Unentangled Subjects).

Probabilistic Attention: Instead of treating attention maps as raw activations, FOCUS treats them as probability distributions over spatial locations.
Jensen-Shannon Divergence (JSD): The loss function uses JSD to measure:
1. Within-Subject Consistency: Minimizing dispersion within a subject's attention maps (ensuring the subject is coherent).
2. Between-Subject Separation: Maximizing the divergence between the mean attention distributions of different subjects (ensuring they do not overlap).
Normalization: The JSD is normalized to ensure comparability across varying numbers of subjects.

C. Two Instantiations

The framework supports two practical algorithms:

Test-Time Control (Training-Free):
- Computes the control signal on-the-fly during sampling.
- Perturbs the base velocity with a single-pass update.
- Requires no model retraining but increases inference latency (approx. 2x).
Adjoint Matching Fine-Tuning:
- Learns a lightweight control network (using LoRA) to approximate the optimal control signal.
- Uses Adjoint Matching, regressing the control network to a "lean adjoint" computed on frozen trajectories, avoiding backpropagation through the controlled path.
- Results in a model that generalizes to unseen prompts with zero inference-time overhead.

3. Key Contributions

FOCUS Objective: A novel, probabilistic attention-binding objective based on Jensen-Shannon divergence that effectively measures and penalizes multi-entity entanglement.
SOC Perspective: A unified theoretical interpretation of latent-update guidance as Stochastic Optimal Control, providing a single hyperparameter ( $\lambda$ ) to trade off between base fidelity and disentanglement.
Dual Algorithms:
- A training-free test-time controller that works with any off-the-shelf SDE/ODE solver.
- An Adjoint Matching fine-tuning route that learns an amortized controller, preserving base style and diversity while improving multi-subject fidelity.
Architecture Agnosticism: The method is designed for modern Flow Matching models (SD 3.5, FLUX.1) but successfully transfers to older U-Net diffusion models (SD 1.5).

4. Experimental Results

The authors evaluated their approach on Stable Diffusion 3.5 and FLUX.1 using a custom dataset of 150 prompts with 2–4 subjects.

Quantitative Metrics:
- FOCUS consistently achieved the highest Composite Scores (a macro-average of CLIP, SigLIP, BLIP, Qwen2, PickScore, and ImageReward) compared to baselines like Attend&Excite, CONFORM, and Divide&Bind.
- Fine-tuned models outperformed test-time controllers on composite scores, demonstrating the benefit of learning the control policy.
Human Preference Study:
- In a pairwise study with 50 participants (2,000 comparisons), FOCUS achieved the highest win rates and Elo ratings for both test-time control and fine-tuning scenarios.
- Users preferred FOCUS-generated images for better subject presence, correct attribute binding, and reduced leakage.
Generalization:
- Fine-tuned models generalized well to unseen prompts.
- Remarkably, fine-tuning on a dataset as small as 1 prompt yielded significant improvements, indicating the SOC controller learns a robust disentangling direction.
Efficiency:
- Test-time control runs efficiently on commodity GPUs (12GB VRAM).
- Fine-tuned models add zero overhead to inference time.

5. Significance

This paper represents a significant shift from heuristic, ad-hoc interventions to a principled, control-theoretic approach for T2I generation.

Unification: It provides a common language to understand and improve multi-subject generation across different model architectures (U-Net vs. DiT/Flow Matching).
Practicality: By offering both a zero-shot test-time solution and a lightweight fine-tuning path, it makes high-fidelity multi-entity generation accessible without requiring massive retraining or complex user layouts.
Future Direction: The SOC framework opens the door for applying advanced control theory techniques to other generative modeling challenges, such as temporal consistency in video or 3D generation.

In summary, FOCUS solves the critical problem of multi-subject entanglement in modern T2I models by leveraging optimal control theory, delivering state-of-the-art results in both automated metrics and human preference while maintaining the stylistic integrity of the base models.