Prompt Reinjection: Alleviating Prompt Forgetting in Multimodal Diffusion Transformers

Abstract: Multimodal Diffusion Transformers (MMDiTs) for text-to-image generation maintain separate text and image branches, with bidirectional information flow between text tokens and visual latents throughout denoising. In this setting, we observe a prompt forgetting phenomenon: the semantics of the prompt representation in the text branch is progressively forgotten as depth increases. We further verify this effect on three representative MMDiTs--SD3, SD3.5, and FLUX.1 by probing linguistic attributes of the representations over the layers in the text branch. Motivated by these findings, we introduce a training-free approach, prompt reinjection, which reinjects prompt representations from early layers into later layers to alleviate this forgetting. Experiments on GenEval, DPG, and T2I-CompBench++ show consistent gains in instruction-following capability, along with improvements on metrics capturing preference, aesthetics, and overall text--image generation quality.

• The paper identifies prompt forgetting, where text-token semantics progressively degrade across MMDiT layers because denoising loss directly supervises image tokens but not text tokens.
• The paper introduces Prompt Reinjection, a training-free method that aligns shallow text features using normalization and orthogonal Procrustes rotation before reinserting them into deeper blocks during inference.
• The paper improves GenEval overall scores by up to 6.48% across SD3, SD3.5, FLUX, and Qwen-Image, with especially strong gains in spatial reasoning while preserving preference and aesthetic metrics.

Overview

This paper identifies and characterizes a failure mode of Multimodal Diffusion Transformers (MMDiTs) that the authors term prompt forgetting: the semantic content of text-token representations in the text branch degrades progressively with network depth during denoising. The diagnosis rests on a supervisory asymmetry inherent to MMDiT training—the denoising loss is defined exclusively in the image latent space, so visual tokens receive direct gradient supervision while textual tokens are updated only indirectly through joint attention. The paper then proposes Prompt Reinjection, a training-free, inference-time intervention that reintroduces aligned shallow-layer text features into deeper blocks. Across SD3-medium, SD3.5-large, FLUX.1-dev, and Qwen-Image, the method yields consistent gains on GenEval, DPG-Bench, and T2I-CompBench++ without degrading preference or aesthetic quality metrics.

Diagnosing prompt forgetting

The empirical case for prompt forgetting proceeds in two stages. First, an observational analysis measures preservation of local semantic structure via Conditional K-Nearest Neighbor Alignment (CKNNA) between the text-encoder output $T^{(0)}$ and each layer's features $T^{(l)}$, complemented by shared PCA projections of token distributions across depth. CKNNA declines monotonically in all studied models, indicating that semantically similar tokens at the input progressively cease to be neighbors deeper in the stack. The PCA projections show a pronounced collapse: most tokens concentrate into a small region of the latent space with only sparse outliers remaining separable.

Second, a functional quantification trains layer-specific MLP probes to decode five token-level linguistic categories (noun, adjective, spatial-relation, numeral, others) from intermediate representations, using identical probe capacity and training protocols so that accuracy differences reflect representational information content rather than probe confounds. Probing accuracy drops monotonically with depth on SD3, SD3.5, and FLUX.1-dev. The degradation is non-uniform: spatial-relation tokens suffer the steepest decline, which the authors connect to the consistently weak Position scores these models exhibit on GenEval. This correlation between where information is lost and where instruction following fails is one of the paper's more persuasive pieces of evidence.

The Prompt Reinjection method

A pilot study establishes two prerequisites. Using minimal prompt pairs $(P_A, P_B)$ differing in a single attribute, injecting a small scaled residual $w \cdot T^{(0)}_B$ ($w \in [0.01, 0.1]$) of prompt B's shallow features into deep blocks during generation with prompt A shifts outputs toward B's attribute—confirming that shallow features retain transferable semantics and that residual addition can carry them.

Naive cross-layer addition, however, is impeded by distributional (scale/shift) and geometric (rotational) mismatches between depths. The full method addresses both:

• Distribution anchoring and restoration: both origin and target features are Layer Normalized before fusion, and the fused result is rescaled back to the target layer's token-wise mean $\mu_{\text{tgt}}$ and standard deviation $\sigma_{\text{tgt}}$, keeping modified representations within the numerical regime downstream blocks expect.
• Geometry alignment via orthogonal Procrustes: a one-time calibration on COCO-5K solves $\min_R \|\mathbf{X}R - \mathbf{Y}\|_F^2$ subject to $R^\top R = I$, with closed-form solution $R = UV^\top$ from SVD of $\mathbf{X}^\top\mathbf{Y}$; origin features are rotated by $R$ before injection.

The origin layer is chosen just after the sharp initial distributional transition visible in the PCA analysis (layer 1 for SD3, layer 2 for FLUX), balancing semantic fidelity against cross-layer compatibility. Injection targets all subsequent blocks, matching the monotonic depth range over which recoverability degrades.

Main results

On GenEval, Prompt Reinjection improves overall scores by 6.48% for SD3.5 (0.7179 → 0.7644) and 5.64% for FLUX (0.6613 → 0.6986), with smaller but consistent gains for SD3 (0.6793 → 0.7059) and Qwen-Image (0.8756 → 0.8933). Position tasks show the clearest improvements across all four models—for example, FLUX rises from 0.1950 to 0.2525 and SD3.5 from 0.2575 to 0.3200—directly corroborating the probing finding that spatial semantics are lost most aggressively. Notably, even the 20B-parameter Qwen-Image benefits on harder subtasks (Numeracy +3.0%, Position +6.3%), indicating that scale alone does not eliminate depth-wise feature degradation.

Gains extend to DPG-Bench (e.g., SD3 85.43 → 87.12 overall) and T2I-CompBench++, particularly on color, shape, and texture binding. Crucially, human-preference proxies (HPSv2, ImageReward, PickScore) and CLIP score are maintained or marginally improved, supporting the claim that instruction adherence is improved without sacrificing generative fidelity. With reinjection enabled, layer-wise probing accuracy remains stable near shallow-layer levels rather than declining, directly demonstrating that the intervention counteracts the diagnosed forgetting. Against TACA—a LoRA-based competitor operating on the same phenomenon—the method achieves a higher GenEval overall score on FLUX (0.6986 vs. 0.6793), though TACA attains a higher Position score (0.2723 vs. 0.2525).

Ablations and cost

Several ablations refine the design picture:

Design choice	Finding
Origin layer / weight	Shallow source ($l_{\text{ori}} \in \{1,2,4\}$) with low weight ($w < 0.1$) is optimal; large weights destabilize generation severely (e.g., SD3 at $l_{\text{ori}}=0$, $w=0.1$ collapses to 0.2433)
Target coverage	Injecting into all subsequent blocks outperforms range-restricted or strided injection
Alignment components	Both anchoring and rotation help; rotation contributes more; even unaligned injection beats baseline
CFG scale	Gains are stable across CFG 4.0–8.0
Calibration data	COCO-5K works best, but BLIP3o/Echo-4o subsets perform comparably

Per-block overhead is modest: relative FLOPs of 1.0883× an SD3 block with full alignment, ~2.29 ms latency per block versus 2.118 ms native, and ~7.8 MB additional memory. The efficiency–effectiveness trade-off is favorable for an inference-time method.

Limitations and open questions

The paper is candid about constraints. Cross-depth discrepancies in feature distribution and geometry persist despite alignment, so reinjection weights cannot be made reliably large, and origin-layer choice is constrained by cross-layer compatibility—meaning new MMDiTs may require manual tuning of weight and origin layer. For Qwen-Image, the search was not as exhaustive as for smaller models, and a mid-stack origin layer (layer 30) was adopted empirically without full justification, leaving open why optimal origin depth differs so substantially across architectures. The probing analysis uses a single denoising step per model, so timestep dependence of forgetting is not characterized. Open questions include whether learned layer- or timestep-dependent injection weights would outperform the shared scalar, whether fine-tuning with reinjection enabled permits stronger injection, and whether adding direct supervision to the text branch during training (e.g., a text reconstruction loss) could address the root cause rather than the symptom.

Conclusion

The paper provides a rigorous, probe-based characterization of progressive prompt-information loss in MMDiT text branches, attributes it to the supervisory asymmetry of the diffusion objective, and shows that a lightweight, training-free reinjection of aligned shallow features measurably restores recoverability and improves instruction following across four mainstream models and three benchmarks while preserving image quality. Its principal contribution is less the intervention itself than the causal link it draws between layer-wise semantic degradation and compositional failures such as spatial non-adherence.