Technical Summary: Diagnosing and Correcting Concept Omission in Multimodal Diffusion Transformers

Problem Statement

Multimodal Diffusion Transformers (MM-DiT), such as FLUX and Stable Diffusion 3.5, have achieved significant progress in text-to-image (T2I) generation. However, they frequently suffer from concept omission, a failure mode where specific objects requested in the prompt are entirely missing from the generated image (object omission) or where their visual properties are ignored (attribute neglect). Existing solutions often require additional model training or incur high computational costs during inference. Furthermore, prior interpretability research has largely focused on visual embeddings, leaving the role of text embeddings in detecting and correcting these omissions under-explored.

Methodology

1. Diagnosis via Linear Probing

The authors hypothesize that text token embeddings within MM-DiT contain information regarding the presence or absence of the corresponding concept during the generation process. To verify this, they employ linear probing:

• Dataset Construction: They generate images using prompts containing two objects (based on GenEval). They label the internal representations based on whether the objects successfully appear in the final image.
• Representation Selection: They focus on the key vectors ($k$) of text tokens, as these are known to heavily influence image generation in MM-DiT architectures.
• Temporal Filtering: They observe that the alignment between internal representations and final generation outcomes is poor in early and late timesteps. Consequently, they restrict their analysis to intermediate timesteps, where the agreement between predicted clean images ($\hat{x}_0$) and final labels is highest (0.965).
• Findings: Linear classifiers trained on these intermediate representations achieve high accuracy (up to 91.0%) in specific attention heads, particularly in the middle layers of the model. This confirms that specific heads encode a distinct "omission signal" that distinguishes between the presence and absence of a concept.

2. Omission Signal Intervention (OSI)

Leveraging the identified omission signal, the authors propose Omission Signal Intervention (OSI), a training-free inference-time method to correct concept omission:

• Direction Extraction: They compute a steering direction vector ($\delta$) representing the difference between the mean key vectors of omitted concepts ($y=0$) and present concepts ($y=1$).
• Intervention Mechanism: During inference, they intervene on the key vectors of specific concept tokens by adding a scaled version of this direction vector:
  $$k^{(t,l,h)}_c \leftarrow k^{(t,l,h)}_c + \alpha \sigma^{(l,h)} \theta^{(l,h)}$$
  Where $\theta$ is the normalized direction, $\sigma$ is the standard deviation of projections, and $\alpha$ is a hyperparameter controlling strength.
• Targeting Strategy: To minimize unintended perturbations, the intervention is applied only to:
  1. The top-K attention heads with the highest probing accuracy.
  2. Specific time intervals (the initial steps of the generation process, $t \in [t_{stop}, 1]$), where concept formation primarily occurs.

Key Contributions

1. Empirical Verification: The paper demonstrates that text tokens in MM-DiT capture omission-related information through joint interaction with image tokens, specifically within certain attention heads at intermediate timesteps.
2. OSI Method: The introduction of Omission Signal Intervention, which amplifies the internal "omission signal" to make the model aware of missing concepts, thereby catalyzing their generation without retraining.
3. Comprehensive Validation: The method is validated across two benchmark categories (object omission and attribute neglect) on two state-of-the-art models (FLUX.1-Dev and SD3.5-Medium), showing significant improvements over strong baselines like TACA and PLADIS.

Experimental Results

• Object Omission: On the GenEval dataset (2–6 objects) and T2I-CompBench, OSI significantly outperforms baselines. For instance, on FLUX with 6-object prompts, OSI achieves an accuracy of 0.40, compared to 0.18 for the base model and 0.20 for TACA/PLADIS.
• Attribute Neglect: OSI improves performance on attribute binding tasks (color, shape, texture) across both models. Notably, applying the intervention to attribute tokens (in addition to object tokens) yields further gains, suggesting the signal generalizes beyond object tokens.
• Ablation Studies:
  • Direction: Reversing the direction or using random vectors degrades performance or yields marginal gains, confirming the semantic alignment of the learned direction.
  • Head Selection: Intervening on low-accuracy heads or random heads fails to match the performance of selecting top-K heads.
  • Token Specificity: Intervening only on omitted tokens significantly boosts their generation probability and accuracy, whereas intervening on already-present tokens has negligible effect.
  • Hyperparameters: The method is robust across a range of scaling coefficients ($\alpha$) and head counts ($K$), with optimal settings balancing alignment and image quality (measured by MANIQA).

Significance and Claims

The paper claims that OSI offers a new perspective on the mechanism of text-to-image generation by identifying that the model possesses an internal, detectable awareness of missing concepts. By amplifying this signal, the model can be steered to "hallucinate" the missing elements intentionally.

The authors position OSI as a training-free, computationally efficient solution that alleviates concept omission without the need for fine-tuning or complex constraint optimization. They emphasize that while the method improves controllability, it does not introduce new knowledge into the model but rather steers existing representations. The work suggests that understanding and manipulating internal "omission signals" is a viable path to improving the fidelity of multimodal diffusion models.