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

• Static conditioning fails to account for evolving needs across diffusion time and network depth.Source: §1 · p. 1
• LLM layers encode different semantic granularities, while DiT stages differ in structural and textural functions.Source: §1 · pp. 1–2
• The paper asks how to route hierarchical LLM signals by diffusion timestep and DiT depth to improve generation.Source: §1 · p. 2
• A unified lightweight framework compares time-wise, depth-wise, and joint fusion under matched settings.Source: §1 · p. 2
• +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.Source: §1 · p. 2
• 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.Source: §1 · p. 2

• Diffusion text conditioning has progressed from U-Net cross-attention to DiT adaLN-Zero, cross-attention, and joint self-attention designs.Source: §2.1 · p. 2
• Recent diffusion systems increasingly use decoder-only LLMs instead of encoder-only CLIP or T5 encoders.Source: §2.1 · p. 2
• Shallow LLM layers emphasize lexical semantics, whereas deeper layers capture increasingly abstract concepts and task-specific information.Source: §2.2 · p. 2
• Multi-layer conditioning has shown advantages, but prior deep-fusion methods lack rigorous controls or compare mainly against the final LLM layer.Source: §2.2 · pp. 2–3
• Diffusion denoising follows a coarse-to-fine trajectory from global structures toward high-frequency details.Source: §2.3 · p. 3

• Flow matching transports samples from a base distribution to the target distribution by integrating a text-conditioned vector field over continuous time.Source: §3.2 · p. 3
• The conditioning site is indexed by DiT block depth d, where fused text states enter each block’s cross-attention module.Source: §3.2 · p. 3
• The LLM supplies hidden-state sequences H(l) from all layers, each with text length N and hidden dimension C.Source: §3.2 · p. 3
• The framework uses LayerNorm and softmax-normalized convex weights to combine layer features within their convex hull.Source: §3.3 · pp. 3–4
• The evaluation compares penultimate-layer conditioning, uniform normalized averaging, and static learnable fusion.Source: §3.4 · p. 4
• A time-conditioned gate maps sinusoidally embedded time through an MLP to produce fusion logits for time-wise and joint strategies.Source: §3.4 · p. 4
• Time-wise fusion shares one gate across blocks, whereas depth-wise fusion learns block-specific logits that remain static over time.Source: §3.4 · p. 4
• Joint fusion assigns each DiT block its own time-conditioned gate, producing weights dependent on both timestep and depth.Source: §3.4 · p. 4

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.Source: §4.1 · p. 4
• 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.Source: §4.1 · p. 4

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.Source: §4.2 · p. 5
• +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.Source: §4.2 · p. 5
• 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.Source: §4.2 · p. 5

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.Source: §5.1 · p. 6
• 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.Source: §5.1 · p. 6
• 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.Source: §5.1 · p. 6, Figure 3 · p. 7

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.Source: §5.2 · pp. 7–8
• 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.Source: §5.2 · p. 8

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.Source: §5.3 · p. 8
• 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.Source: §5.3 · p. 8

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.Source: §6 · p. 8
• 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.Source: §6 · p. 8
• 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.Source: §6 · p. 8

• Semantic routing fuses multi-layer LLM hidden states to better match DiT generation across network depth and optional temporal dynamics.Source: §7 · p. 9
• 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.Source: §7 · p. 9
• The approach could be misused for deceptive imagery, reproduce societal biases, or lower barriers to harmful or infringing content.Source: §7 · p. 9
• Existing safeguards such as content filtering, sensitive-concept blocking, bias evaluation, robustness testing, and watermarking can mitigate these risks.Source: §7 · p. 9

• TCFG maps the diffusion timestep through a lightweight MLP to adaptively aggregate multi-layer LLM features.Source: A. TCFG Implementation Details · p. 12
• The continuous timestep is represented with a sinusoidal embedding, using dimension D_t = 128 in the experiments.Source: A. TCFG Implementation Details · p. 12
• A two-layer MLP transforms the timestep embedding into L fusion logits, with an input projection, SiLU activation, and output projection.Source: A. TCFG Implementation Details · p. 12
• Zero-initializing the final projection produces uniform softmax weights at training start, corresponding to an average over all layers.Source: A. TCFG Implementation Details · p. 12
• LayerNorm is applied before fusion, and the aggregated feature is a convex combination of normalized layer representations.Source: A. TCFG Implementation Details · p. 12

• Figure 5 compares B1, B2, B3, FuseDiT, and three fusion strategies across multiple prompts.Source: Figure 5
• The examples include prompts involving balloons, teddy bears, dogs and cats, pillows, cups and a fork, puppies, roses, and a pilot.Source: B. Image Examples · pp. 13–14
• Additional prompts depict balloons against a blue sky and a pizza with pepperoni on one half and mushrooms on the other.Source: B. Image Examples · p. 13

• Figures 6 and 7 show fusion-weight distributions for the text encoder.Source: §C.1 · p. 14
• The visualizations organize weights across different text-encoder layers and diffusion timesteps.Source: §C.1 · p. 14
• The figures compare the depth-wise and joint strategies.Source: §C.1 · p. 14

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.Source: §C.2 · pp. 14–15
• 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.Source: §C.2 · p. 15

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.Source: §C.2 · p. 15
• 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.Source: §C.2 · p. 15

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.Source: §C.2 · pp. 17–18