Source-linked AI summary
Semantic Routing: Exploring Multi-Layer LLM Feature Weighting for Diffusion Transformers
Bozhou Li, Yushuo Guan, Haolin Li, Bohan Zeng, Yiyan Ji, Yue Ding, Pengfei Wan, Kun Gai, Yuanxing Zhang, Wentao Zhang
TL;DR
DiT text conditioning often uses static or single-layer representations despite hierarchical LLM semantics and non-stationary denoising. The paper introduces normalized convex multi-layer fusion with lightweight time-wise, depth-wise, and joint gates, finding depth-wise routing most robust while attributing time-wise degradation to train–inference trajectory mismatch.
Problem
Static conditioning and single-layer use do not reflect hierarchical LLM semantics or changing conditioning needs across diffusion time and DiT depth.
Method
A unified lightweight framework uses normalized convex fusion and gates to compare time-wise, depth-wise, and joint routing of multi-layer LLM features.
Results
Multi-layer fusion consistently outperforms single-layer conditioning, while depth-wise adaptive fusion delivers the most robust improvements; it gains +9.97 on GenAI-Bench Counting over the penultimate-layer baseline.
Takeaways & Limitations
Depth-wise semantic routing is a strong baseline, whereas robust time-dependent conditioning requires trajectory-aware signals.
Takeaways & Limitations
Stronger alignment may reproduce societal biases more consistently and lower barriers to harmful or infringing content.
Abstract
from arXiv · showhide
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.
1. Introduction
The paper argues that static text conditioning overlooks semantic hierarchy in LLM layers and changing conditioning needs across diffusion time and DiT depth. It introduces a unified framework for adaptively routing multi-layer features and finds depth-wise fusion strongest, while time-wise fusion suffers from trajectory mismatch.
- Static conditioning fails to account for evolving needs across diffusion time and network depth.
- LLM layers encode different semantic granularities, while DiT stages differ in structural and textural functions.
- The paper asks how to route hierarchical LLM signals by diffusion timestep and DiT depth to improve generation.
- A unified lightweight framework compares time-wise, depth-wise, and joint fusion under matched settings.
- +9.97 on GenAI-Bench Counting over the penultimate-layer baseline establishes depth-wise fusion as the strongest strategy.+5.45 is reported over uniform averaging.
- Nominal timesteps fail to track effective SNR during sampling, causing semantic misalignment in time-wise fusion.The paper identifies this train–inference trajectory mismatch as a failure mode motivating trajectory-aware conditioning.
2. Related Work
Related work shows a shift toward decoder-only LLMs for stronger diffusion conditioning, alongside evidence that LLM layers encode different semantic information. However, existing approaches generally use static fusion and lack controlled isolation of adaptive-routing effects.
- Diffusion text conditioning has progressed from U-Net cross-attention to DiT adaLN-Zero, cross-attention, and joint self-attention designs.
- Recent diffusion systems increasingly use decoder-only LLMs instead of encoder-only CLIP or T5 encoders.
- Shallow LLM layers emphasize lexical semantics, whereas deeper layers capture increasingly abstract concepts and task-specific information.
- Multi-layer conditioning has shown advantages, but prior deep-fusion methods lack rigorous controls or compare mainly against the final LLM layer.
- Diffusion denoising follows a coarse-to-fine trajectory from global structures toward high-frequency details.
3. Method
The method represents text conditioning as a normalized convex fusion of multi-layer LLM features, with routing indexed by diffusion time and DiT depth. Lightweight gates instantiate static, time-wise, depth-wise, and joint parameterizations for controlled comparison.
- Flow matching transports samples from a base distribution to the target distribution by integrating a text-conditioned vector field over continuous time.
- The conditioning site is indexed by DiT block depth d, where fused text states enter each block’s cross-attention module.
- The LLM supplies hidden-state sequences H(l) from all layers, each with text length N and hidden dimension C.
- The framework uses LayerNorm and softmax-normalized convex weights to combine layer features within their convex hull.
- The evaluation compares penultimate-layer conditioning, uniform normalized averaging, and static learnable fusion.
- A time-conditioned gate maps sinusoidally embedded time through an MLP to produce fusion logits for time-wise and joint strategies.
- Time-wise fusion shares one gate across blocks, whereas depth-wise fusion learns block-specific logits that remain static over time.
- Joint fusion assigns each DiT block its own time-conditioned gate, producing weights dependent on both timestep and depth.
4. Experiments
Experiments compare multi-layer feature-fusion strategies under matched architectures and training settings. Depth-wise routing provides the strongest reported gains, while static and purely time-wise alternatives expose limitations.
- Experimental Setup: All conditioning variants share the same DiT backbone and differ only in how multi-layer text features are fused.The backbone uses 24 Transformer blocks and approximately 2.24B parameters.
- Experimental Setup: Models are trained on approximately 30 million LAION-400M image-text pairs with dense synthetic captions.Images are resized to 256 × 256, and prompts support up to 512 tokens.
- Main Results: Multi-layer aggregation consistently outperforms penultimate-layer conditioning, while learnable static fusion does not surpass uniform averaging.The result suggests complementary semantic signals are underused by single-layer conditioning, but fixed learned weights lack sufficient adaptivity.
- Main Results: +9.97 on GenAI-Bench Counting over the penultimate-layer baseline and +5.45 over uniform averaging are achieved by depth-wise fusion.Its gains are especially pronounced in Advanced categories and support selective routing across DiT depths.
- Main Results: Purely time-wise fusion often degrades generation quality through blurriness and loss of fine details, whereas joint fusion avoids this characteristic degradation.Joint fusion remains competitive but is slightly less effective than depth-wise fusion.
5. Analysis
The analysis examines learned multi-layer fusion weights, trajectory mismatch in time-wise fusion, and computational overhead. Depth-wise routing shows strong compositional gains, while time-wise fusion can degrade generation because nominal timesteps diverge from inference dynamics.
- Weight Dynamics: Learned weights selectively emphasize internal LLM layers, with the penultimate layer dominant early before fading during later generation.The pattern is consistent with early structural layout needs differing from later texture-refinement needs.
- Weight Dynamics: Intermediate-layer weights show local peaks with neighboring suppression, especially under joint fusion, indicating selective redundancy reduction.Joint fusion produces sharper peaks and higher local contrast than decoupled time-wise and depth-wise settings.
- Weight Dynamics: Fusion weights vary across timesteps and DiT depth, while neighboring timesteps and blocks exhibit structured local smoothness.This emergent locality is presented as evidence that depth-only routing reflects coherent semantic signals rather than stochastic noise.
- Trajectory Misalignment: Time-wise fusion degrades generation quality because nominal timesteps fail to track effective SNR along the CFG inference trajectory.Iterative inference advances faster than the static training schedule, making nominal-timestep conditioning semantically mistimed.
- Trajectory Misalignment: Manual timestep advancement recovers performance across metrics, including +0.24 on GenEval, but cannot fully correct nonlinear trajectory deviations.The result supports the mismatch diagnosis while leaving robust time-aware fusion as an open challenge.
- Compute Overhead: S2 adds essentially no extra parameters and increases end-to-end latency by only ∼8%, indicating negligible computational overhead for the best depth-wise strategy.The paper describes gating as not being a computational bottleneck.
- Compute Overhead: FuseDiT’s lower FLOPs coincide with degraded generative quality, whereas modular routing preserves semantic expressiveness under comparable efficiency.The paper hypothesizes that reusing LLM K/V states restricts conditioning expressiveness.
6. Conclusion
The paper studies multi-layer LLM conditioning for DiT models through a unified framework covering time-wise, depth-wise, and joint fusion. Experiments favor depth-wise adaptive fusion, while time-wise fusion can hurt because nominal timesteps mismatch inference denoising dynamics.
- Conclusion: The unified formulation enables controlled comparison of time-wise, depth-wise, and joint adaptive fusion for multi-layer LLM conditioning.The framework uses lightweight adaptive mechanisms within a single formulation.
- Conclusion: Multi-layer fusion consistently outperforms single-layer conditioning, with depth-wise adaptive fusion delivering the most robust and substantial improvements.This is the paper’s principal comparative conclusion across the studied strategies.
- Conclusion: Purely time-wise fusion can hurt performance because nominal timesteps mismatch inference-time denoising dynamics.The learned time-wise weights remain structured, suggesting trajectory-aligned signals may support more effective time-adaptive conditioning.
7. Impact Statements
The paper frames semantic routing as a lightweight, interpretable approach for improving text-image alignment and compositional instruction following. It also discusses potential misuse, bias, and infringement risks alongside existing mitigation practices.
- Semantic routing fuses multi-layer LLM hidden states to better match DiT generation across network depth and optional temporal dynamics.
- The proposed conditioning may improve controllability and instruction consistency in text-to-image generation, supporting applications such as content editing and human-AI co-creation.
- The approach could be misused for deceptive imagery, reproduce societal biases, or lower barriers to harmful or infringing content.
- Existing safeguards such as content filtering, sensitive-concept blocking, bias evaluation, robustness testing, and watermarking can mitigate these risks.
A. TCFG Implementation Details
TCFG uses a lightweight timestep-conditioned MLP to generate layer-fusion logits, while normalized convex aggregation combines the resulting weights with LLM hidden states. Zero initialization starts training from uniform layer averaging.
- TCFG maps the diffusion timestep through a lightweight MLP to adaptively aggregate multi-layer LLM features.
- The continuous timestep is represented with a sinusoidal embedding, using dimension D_t = 128 in the experiments.
- A two-layer MLP transforms the timestep embedding into L fusion logits, with an input projection, SiLU activation, and output projection.
- Zero-initializing the final projection produces uniform softmax weights at training start, corresponding to an average over all layers.
- LayerNorm is applied before fusion, and the aggregated feature is a convex combination of normalized layer representations.
B. Image Examples
The image examples compare the baseline and fusion strategies across multiple prompts under identical sampling settings. The examples cover object counts, attributes, spatial relations, and multi-object compositions.
- Figure 5 compares B1, B2, B3, FuseDiT, and three fusion strategies across multiple prompts.
- The examples include prompts involving balloons, teddy bears, dogs and cats, pillows, cups and a fork, puppies, roses, and a pilot.
- Additional prompts depict balloons against a blue sky and a pizza with pepperoni on one half and mushrooms on the other.
C.1. Detailed Visualization of Fusion Weights
Figures 6 and 7 visualize how text-encoder fusion weights are distributed across layers and diffusion timesteps for depth-wise and joint routing.
- Figures 6 and 7 show fusion-weight distributions for the text encoder.
- The visualizations organize weights across different text-encoder layers and diffusion timesteps.
- The figures compare the depth-wise and joint strategies.
C.2. Weight Evolution and Trend Analysis
The analysis treats fusion weights as distributions over LLM-layer semantic hierarchy and tracks their evolution across diffusion time and network depth. Joint fusion remains broadly distributed and temporally stable, whereas time-wise fusion exhibits stronger drift; joint fusion also changes more smoothly between timesteps.
- Statistical representation: Fusion weights are modeled as distributions over uniformly spaced semantic-level locations, enabling mean and variance summaries.The mean indicates the semantic center, while variance measures dispersion across encoder layers.
- Statistical representation: Joint fusion forms a two-dimensional weight map over diffusion timestep and encoder layer, then marginalizes one axis to obtain interpretable one-dimensional trends.Semantic center and dispersion are computed on each normalized marginal distribution.
- Observed trends: Joint fusion keeps a largely stable, moderately deep semantic center with broad multi-layer coverage, shifting only slightly toward shallower layers at later timesteps.The wide dispersion band indicates that weights do not collapse onto only a few layers.
- Observed trends: Time-wise fusion shows pronounced monotonic drift from deeper toward shallower layers as timesteps progress, while depth-wise fusion has smoother, smaller fluctuations.These patterns indicate less stable time-dependent allocation for time-wise fusion and more consistent allocation for depth-wise fusion.
- Temporal variation: Joint fusion produces smaller consecutive-timestep 1-Wasserstein distances than time-wise fusion, indicating smoother temporal evolution of weight allocation.For joint fusion, the distance is computed after marginalizing over layers.