Semantic Routing: Exploring Multi-Layer LLM Feature Weighting for Diffusion Transformers

TL;DR

This work tackles the static text conditioning in DiT-based diffusion models by introducing a unified, lightweight fusion framework that can route multi-layer LLM features along diffusion time $t$ and DiT depth $d$ through a convex, LayerNorm-stabilized fusion. Among time-wise, depth-wise, and joint strategies, depth-wise routing (S2) delivers the strongest improvements in text-image alignment and compositional generation, with a notable $+9.97$ boost on GenAI-Bench Counting over the penultimate-layer baseline. The authors identify a train-inference trajectory mismatch that plagues time-wise conditioning, providing mechanistic insight and a counterfactual timestep-shift validation; joint time-and-depth fusion mitigates these issues by coupling the axes. The approach achieves better conditioning expressiveness with modest compute overhead, highlighting the importance of trajectory-aware, hierarchy-aligned conditioning for robust, controllable diffusion-based generation.

Abstract

Recent DiT-based text-to-image models increasingly adopt LLMs as text encoders, yet text conditioning remains largely static and often utilizes only a single LLM layer, despite pronounced semantic hierarchy across LLM layers and non-stationary denoising dynamics over both diffusion time and network depth. To better match the dynamic process of DiT generation and thereby enhance the diffusion model's generative capability, we introduce a unified normalized convex fusion framework equipped with lightweight gates to systematically organize multi-layer LLM hidden states via time-wise, depth-wise, and joint fusion. Experiments establish Depth-wise Semantic Routing as the superior conditioning strategy, consistently improving text-image alignment and compositional generation (e.g., +9.97 on the GenAI-Bench Counting task). Conversely, we find that purely time-wise fusion can paradoxically degrade visual generation fidelity. We attribute this to a train-inference trajectory mismatch: under classifier-free guidance, nominal timesteps fail to track the effective SNR, causing semantically mistimed feature injection during inference. Overall, our results position depth-wise routing as a strong and effective baseline and highlight the critical need for trajectory-aware signals to enable robust time-dependent conditioning.

Figures (9)

• Figure 1: Learned Fusion Weights. We observe distinct weight distributions for shallow versus deep DiT blocks, indicating that different generative stages utilize distinct semantic levels from the text encoder.
• Figure 2: Weight distributions under different fusion-weight parameterizations. The x-axis denotes the text encoder layer index $l$, and the y-axis denotes the normalized fusion weight $\alpha_{t,d}$. For time-wise fusion, we sample $t \in [0,1]$ with a step size of $0.2$. For joint fusion and depth-wise fusion, we report representative DiT block indices $d \in \{0, 11, 23\}$. Additional visualizations are provided in the appendix \ref{['app:detail_weights']}
• Figure 3: Visualizing the local smoothness of learned fusion weights. We compute pairwise JS similarity between normalized fusion-weight distributions along two axes: (a) across DiT blocks (depth) and (b) across diffusion timesteps (time).
• Figure 4: Trajectory analysis on GenAI-Bench: Latent MSE (left) and PSNR (right) evolution.
• Figure 5: Qualitative comparisons across strategies under multiple prompts. Columns: B1/B2/B3, FuseDiT baseline, and three fusion strategies (S1/S2/S3). All images use identical sampling settings for fair comparison.