PaperlayerView exact revision on arXiv

Source-linked AI summary

TSGate: Timestep-Aware Gated Attention for Diffusion Transformers

Boyu Zhang, Yifan Liu, Shuxia Lin, Qingjian Ni, Yinfei Xu, Xu Yang

arXiv:2609.34539v1cs.CV

TL;DR

Free-form prompts can shift away from the structured-prompt training distribution, and the paper links this degradation to attention sinks and reduced early text interaction. TSGate adds a timestep-conditioned bias to attention gating; it consistently outperforms the baseline and standard gated attention across reported benchmarks, including a 9.5% raw-prompt DPG-Bench improvement over the baseline.

  • Problem

    DiTs trained on structured prompts can degrade on free-form OOD prompts, which differ from the training distribution and are common at inference.

    Source: §1 · p. 1

  • Method

    TSGate adds an explicit timestep-conditioned bias to the gate signal to adapt gating behavior across denoising steps.

    Source: §1 · p. 2

  • Results

    TSGate scores 60.749 on raw DPG-Bench prompts, 5.266 points (9.5%) above Baseline, and consistently outperforms standard baselines across experiments.

    Source: §1 · p. 2, §4.2.2 · p. 7

  • Takeaways & Limitations

    The learned gate exhibits an emergent stage-dependent allocation rather than a manually prescribed switching schedule.

    Source: §4.2.2 · p. 7

    Evaluation centers on text-to-image generation; extending TSGate to broader multimodal tasks remains to be explored.

    Source: §5 · p. 9

Abstract

from arXiv · show

Diffusion Transformers (DiTs) have emerged as the dominant architecture for high-fidelity image and video generation. Recent DiT systems increasingly use structured prompts for training, improving caption quality and prompt adherence. However, their generation quality can degrade severely under out-of-domain (OOD) prompts, including the free-form descriptions supplied by users at inference time. Although LLM-based rewriting can convert these prompts into structured formats, it does not guarantee that the rewritten prompts align with the training distribution. Our analysis links this degradation to attention sinks and reduced early-step image-to-text attention and shows that sink suppression alone is insufficient to restore generation quality. Despite effective sink suppression, models trained with standard gated attention exhibit reduced early-step image-to-text attention and suboptimal generation quality. Based on these insights, we propose Timestep-Aware Gated Attention (TSGate), which injects a timestep-conditioned bias into the gate signal so that gating behavior adapts across denoising steps. Extensive experiments show that TSGate consistently outperforms both the baseline and standard gated attention across multiple benchmarks, improving the raw-prompt DPG score by 9.5% over the baseline.

1 INTRODUCTION

DiTs trained on structured prompts can degrade on free-form user descriptions that differ from the training distribution. The analysis links this fragility to attention sinks and reduced early image-to-text attention, motivating TSGate’s timestep-conditioned gate and its reported OOD gains.

1 INTRODUCTION

  • Structured-prompt training is presented as a widely adopted way to improve text–image alignment, but free-form user prompts often depart from that training distribution.The introduction notes that prompt engineering cannot fully eliminate this distribution shift.Source: §1 · p. 1
  • Under two OOD prompts, Figure 1 depicts clearer requested details with TSGate than with the baseline and standard gated attention.The examples show a snow-capped peak and a parrot’s head; the schematic also motivates stronger text guidance during early denoising.Source: Figure 1 · p. 2, §1 · p. 2
  • OOD prompts reduce early image-to-text attention while attention sinks intensify, limiting textual guidance as noisy image tokens acquire global semantics.The authors associate sink exacerbation with the absence of familiar header tokens.Source: §1 · pp. 1–2
  • Standard gated attention suppresses sinks but retains weak early image-to-text interaction and suboptimal generation quality, so sink suppression alone is insufficient.The content-dependent gate tends to suppress the image-to-text pathway when image tokens are dominated by noise early in denoising.Source: §1 · p. 2
  • TSGate adds a timestep-conditioned bias to the gate, aiming to preserve sink suppression while strengthening early image-to-text interaction without explicit attention-allocation supervision.The passage describes the extension as adding only a small number of parameters per layer and being compatible with existing DiT architectures.Source: §1 · p. 2
  • TSGate consistently outperforms standard baselines in experiments and a first-layer-only variant achieves substantial benefits with negligible parameter overhead.The contribution passage does not provide benchmark-specific values.Source: §1 · p. 2

2 RELATED WORK

Related work traces DiTs’ architectural development and reviews attention sinks and gating across Transformer settings. The paper positions its contribution at the intersection of sink analysis and timestep-aware attention allocation.

2.1 DIFFUSION TRANSFORMERS

  • DiTs replace convolutional U-Nets with Transformer backbones, while joint attention lets text and image tokens interact in a shared sequence.The section also notes growing use of structured, template-based training captions and says their implications for OOD prompt robustness remain largely unexplored.Source: §2.1 · p. 3

2.2 ATTENTION SINKS IN TRANSFORMERS

  • Prior sink studies span language and vision Transformers, but the paper says DiT sinks’ practical effects and their connection to early-step attention allocation remain underexplored.The related-work discussion identifies this connection as the gap addressed by the paper.Source: §2.2 · p. 3

2.3 GATED ATTENTION MECHANISMS

  • Content-only gated attention derives logits from hidden representations; TSGate complements that implicit timestep path with an explicit timestep-bias branch.The subsection describes earlier gated methods in language modeling and related gating approaches in efficient and diffusion architectures.Source: §2.3 · p. 3
  • Figure 2 depicts TSGate adding a token-shared, channel-wise time bias to content logits before the sigmoid gate scales attention features.The gated features are scaled before the output projection.Source: Figure 2 · p. 4

3 METHOD

TSGate retains token-dependent content selection while adding an explicit timestep-dependent bias to gated attention. The bias modulates feature retention across denoising steps without directly changing attention probabilities or imposing a prescribed early/late schedule.

3 METHOD

  • TSGate combines token-specific feature selection with explicit stage-dependent retention by adding a timestep-conditioned bias to content-only gating.The method is presented as complementing the backbone’s existing timestep conditioning.Source: §3 · p. 4, §3.2 · p. 4

3.1 PRELIMINARIES: GATED ATTENTION

  • Content-only gated attention forms attention features, computes content logits from input features, and scales the aggregated features before output projection.The preliminaries describe element-wise or head-wise gates and use a sigmoid followed by element-wise multiplication.Source: §3.1 · p. 4

3.2 TIMESTEP-AWARE GATED ATTENTION

  • Figure 3 contrasts structured-prompt and OOD generations, with OOD examples showing degraded object fidelity, compositionality, and scene coherence.The caption describes three representative paired cases.Source: Figure 3 · p. 5
  • A separate timestep-embedding projection produces a channel-wise bias shared across tokens; adding it to content logits before the sigmoid yields the gate that scales aggregated attention features.The bias branch is described as useful when content logits are dominated by noise at high timesteps.Source: §3.2 · p. 5
  • For fixed content logits, raising a channel’s time bias increases its retention while preserving gate ordering across tokens.The gate acts on value-aggregated features, not the softmax matrix; it does not renormalize probabilities or directly favor text keys over image keys.Source: §3.2 · p. 5
  • TSGate’s pre-projection channel gate is distinct from adaLN-Zero’s post-projection residual scaling, and its stage-dependent allocation emerges through learning rather than a prescribed schedule.The paper states that the gate does not directly renormalize text-versus-image probabilities.Source: §3.2 · p. 5

4 EXPERIMENTS

Experiments trace prompt-format failures to attention sinks and weak early text interaction, then show that TSGate improves generation quality across prompt conditions and benchmarks. The results also examine its learned attention allocation, architectural variants, and timestep-bias alignment.

4.1 ATTENTION SINKS IN DIFFUSION TRANSFORMERS

  • In-domain template markers receive more than 20× average-token attention, while OOD prompts disrupt sinks and show reduced early text attention across sampled layers.Figure 4 depicts sink bands persisting across two layers and denoising steps; Figure 5 compares early attention trajectories for in-domain and OOD prompts.Source: §4.1 · pp. 5–6, Figure 4 · p. 6, Figure 5 · p. 6

4.2.1 EXPERIMENTAL SETUP

  • Models were trained from scratch on a curated 100M-image LAION-5B subset re-captioned by Qwen2.5-VL and evaluated on DPG-Bench, GenEval, and T2I-CompBench.The standard evaluation used 50 denoising steps at 320 × 320 resolution.Source: §4.2 · p. 6, §4.2.1 · p. 7

4.2.2 MAIN RESULTS

  • On raw DPG-Bench prompts, TSGate scores 60.749, 5.266 points (9.5%) above Baseline and 1.773 above Gated attention; it also leads on full-recaption, GenEval, and CompBench.Table 1 compares Baseline, Gated attention, and TSGate across prompt formats and benchmarks.Source: Table 1 · p. 7, §4.2.2 · p. 7
  • At t = 1000, TSGate’s image-to-text attention exceeds content-only gating by 15.9 percentage points (103% relative) and Baseline by 7.8 points.The early TSGate–Gated-attention gap is 16.6 points for short prompts and 15.2 for long prompts; attention later crosses below Baseline without a prescribed switching rule.Source: §4.2.2 · p. 7, Figure 6 · p. 8
  • As prompt structure is removed, TSGate remains ahead of Baseline, with the DPG-Bench gap widening from 0.7 points in-domain to 5.3 on raw OOD prompts.Removing headings leaves TSGate near its in-domain score (79.5 versus 79.4), while Baseline falls to 77.2.Source: §4.2.2 · p. 8, Figure 7 · p. 8

4.2.3 TSGATE FAMILY

  • On raw DPG-Bench, head-wise TSGate scores 55.069 below element-wise TSGate, while a dense-backbone comparison rises from 41.619 to 64.159 but is not parameter-matched.The dense-backbone results support applicability beyond MoE, particularly for raw prompts.Source: §4.2.3 · p. 9
  • Restricting the timestep branch to Layer 0 retains 93.8% of TSGate’s improvement over Baseline while reducing time-branch parameters from 85.0M to 7.1M.TSGate-L0 scores 60.4 on raw DPG-Bench versus 60.7 for all-layer TSGate.Source: §4.2.3 · p. 9

4.2.4 TIME-BIAS INTERVENTIONS

  • Reversing the learned time-bias schedule while keeping its values and average fixed reduces raw DPG-Bench by 2.6 points.The intervention changes only the assignment of the bias values to denoising stages.Source: §4.2.4 · p. 9

5 CONCLUSION

The conclusion presents TSGate as a timestep-conditioned gate modification for prompt-format degradation in DiTs, reporting gains across text-to-image benchmarks. Its evaluation remains centered on text-to-image generation, with broader multimodal use left for future exploration.

5 CONCLUSION

  • TSGate adds a timestep-conditioned bias to attention gating and improves raw-prompt DPG-Bench by 9.5% over Baseline, with gains on GenEval and CompBench.Applying the timestep branch only at the first layer retains 93.8% of the improvement with minimal parameter overhead.Source: §5 · p. 9
  • The reported prompt-format degradation is linked to attention sinks on template markers and reduced early-step image-to-text attention; standard gated attention suppresses sinks but does not restore generation quality.TSGate is described as jointly reducing sinks and restoring early text guidance.Source: §5 · p. 9
  • The evaluation centers on text-to-image generation; extending TSGate to text-to-audio and text-to-video remains unexplored.The conclusion also identifies spatial positional encodings as a possible direction for image editing.Source: §5 · p. 9

A GATE IMPLEMENTATION AND CONDITIONING

The appendix specifies how TSGate’s content and timestep branches form gates, how those gates are applied, and why a timestep bias can help when noisy inputs make content logits unreliable.

A.1 GRANULARITY, LAYER SCOPE, AND INITIALIZATION

  • Element-wise gating uses D gate channels, while head-wise gating uses H; both produce D attention channels, with one sigmoid value shared across each head’s channels in the head-wise form.In the evaluated MoE variants, TSGate adds time biases in all layers, while TSGate-L0 adds them only in Layer 0; both retain content gates in all 12 layers.Source: §A.1 · p. 14

A.2 JOINT ATTENTION LAYOUT

  • Joint attention concatenates keys and values across streams, computes standard softmax attention, then applies the gate to the attention output before the output projection.Queries and keys undergo per-head RMSNorm and RoPE; content logits undergo neither.Source: §A.2 · p. 14

A.3 SHARED BIAS AND CONTENT SELECTION

  • For fixed content logits, a shared bias preserves token ordering, shifts the threshold for gates above 1/2, and affects gates near 1/2 most strongly.These local properties do not guarantee increased text attention or require a monotone timestep schedule.Source: §A.3 · p. 14

A.4 OUTPUT GATING DOES NOT RENORMALIZE ATTENTION PROBABILITIES

  • The same gate scales every key’s contribution for a fixed query and channel, so gated coefficients sum to the gate value rather than defining a renormalized softmax distribution.Reported attention metrics use head-averaged pre-gate probabilities; the image-to-text fraction sums over text keys for image queries.Source: §A.4 · p. 15

A.5 RELATIONSHIP TO ADALN-ZERO AND RESIDUAL GATING

  • TSGate adds a timestep-conditioned offset to content logits before the sigmoid, while the resulting bounded gate can only attenuate or pass attention channels.Unlike adaLN-Zero’s affine scale and shift, the gate acts before the attention output projection; the operations are generally not interchangeable.Source: §A.5 · p. 15

A.6 SIGNAL-TO-NOISE VIEW OF THE TIME BIAS

  • At early high-noise steps, noise can dominate the layer input and contaminate content logits; the deterministic timestep bias supplies a token-shared prior that can recede as signal-to-noise ratio rises.The paper describes this as a learned curriculum, not an explicitly scheduled one.Source: §A.6 · pp. 15–16

B.1 ARCHITECTURE AND TRAINING

The experiments compare multiple MoE gate configurations and a dense baseline against TSGate, using shared training data and specified generation and benchmark protocols.

B.1 ARCHITECTURE AND TRAINING

  • The five MoE configurations vary output gating, timestep-bias scope, and gate granularity; a widened Baseline approximately matches TSGate’s parameter count but allocates capacity differently.The fixed-weight interventions use the same trained MoE TSGate checkpoint.Source: §B.1 · p. 16, Table 2 · p. 16
  • All configurations train on the same curated 100M-image LAION-5B subset, recaptioned into structured prompts, at 320 × 320 resolution.Training uses 32 NVIDIA H200 GPUs, AdamW, and a text-drop ratio of 0.1 for classifier-free guidance.Source: §B.1 · pp. 16–17

B.2 INFERENCE

  • MoE evaluation uses four seeds and reports DPG-Bench, GenEval, and T2I-CompBench scores with benchmark-specific averaging; inference uses CFG=5, 320 × 320 resolution, and 50 steps.Dense inference uses the same CFG, resolution, and sampling-step settings.Source: §B.2 · p. 17

C COMPLETE BENCHMARK RESULTS

The results section reports overall benchmark scores, category breakdowns, and paired confidence intervals for DPG-Bench, GenEval, and T2I-CompBench.

C COMPLETE BENCHMARK RESULTS

  • The reported comparisons cover five MoE variants: Baseline, Gated attention, TSGate, TSGate-L0, and TSGate-head.The passage specifies that DPG-Bench uses a 0–100 scale, while GenEval and T2I-CompBench use 0–1 scales, with higher scores better.Source: C COMPLETE BENCHMARK RESULTS · p. 17
  • Differences use source-paired scores averaged over four seeds, with 95% pointwise percentile confidence intervals from 10,000 source-bootstrap resamples.External benchmarks use within-category resampling.Source: C COMPLETE BENCHMARK RESULTS · p. 17
  • The section combines overall scores with category breakdowns for the three benchmarks.The passage identifies benchmark protocols and confidence intervals but does not state which model performs best.Source: C COMPLETE BENCHMARK RESULTS · p. 17

C.1 DPG-BENCH

Across the reported benchmarks, TSGate leads on raw and full-recaption DPG-Bench means and outperforms standard gated attention on GenEval and T2I-CompBench, though it does not lead every macro or category result.

C.1 DPG-BENCH

  • TSGate has the highest DPG-Bench mean for raw and full-recaption prompts, while paired TSGate-L0 versus TSGate intervals include zero.The table compares raw and full-recaption prompts on the same 1,065 sources; these paired checkpoint comparisons do not estimate training-seed variation.Source: §C.1 · p. 17, Table 4 · p. 18

C.3 GENEVAL: CATEGORY BREAKDOWN

  • TSGate’s GenEval category means exceed Gated attention’s in all six categories, with the largest differences in position and color attributes.The cited category passage points to Table 8; paired uncertainty for macro scores is reported separately in Table 7.Source: §C.3 · p. 18, Table 8 · p. 19

C.4 T2I-COMPBENCH: CATEGORY BREAKDOWN

  • On T2I-CompBench, the largest TSGate–Gated-attention differences are in spatial relations and color, while Gated attention leads the complex category.TSGate-head leads the overall macro score.Source: §C.4 · p. 18, Table 9 · p. 20

C.2 GENEVAL AND T2I-COMPBENCH: OVERALL SCORES

  • TSGate outperforms Gated attention on both GenEval and T2I-CompBench, has the highest GenEval mean, and ranks second on CompBench, where TSGate-head leads.The TSGate–Baseline CompBench interval includes zero.Source: §C.2 · p. 18, Table 7 · p. 19

D STRUCTURED PROMPT FORMAT

  • The training captions use five structured fields—summary, setting, subject details, visual style, and mood—written as ordinary prose.The section provides three image-caption examples.Source: D STRUCTURED PROMPT FORMAT · p. 20

E ADDITIONAL PAIRED GENERATION COMPARISONS

  • The additional paired-generation comparisons place full-recaption prompts in the top row and raw prompts in the bottom row, with Baseline, Gated attention, and TSGate in columns.Each model and prompt-form comparison within an example uses the same random seed.Source: E ADDITIONAL PAIRED GENERATION… · p. 21

F STAGE-DEPENDENT TEXT INTERACTION AND TIME-BIAS INTERVENTIONS

  • The Sink Score is concentrated at Layer 0, where sink suppression ranks Gated attention, then TSGate, then Baseline, despite TSGate outperforming Gated attention on DPG.The score is the maximum normalized column-sum of the joint-attention map and is measured before the output gate.Source: F STAGE-DEPENDENT TEXT INTERACTION AND… · p. 23, Figure 16 · p. 23

F.1 FIXED-WEIGHT INTERVENTIONS ON TIME BIASES

Fixed-weight interventions test whether TSGate’s learned bias must match the denoising stage. The results indicate that stage assignment matters and that the explicit bias branch contributes to generation quality under raw and structured inputs.

F.1 FIXED-WEIGHT INTERVENTIONS ON TIME BIASES

  • The fixed-weight tests replace TSGate’s bias sequence while keeping trained weights and backbone timestep conditioning fixed; the content branch remains active.The complete schedule has K = 50 steps; zero removes the additive bias, and reverse reorders bias values.Source: §F.1 · pp. 23–24
  • Reversing the bias schedule reduces raw DPG by 2.614 points (95% CI [−4.281, −0.994]), while removing the bias reduces it by 42.831 points.Matched-mean changes raw DPG by −0.929 points, with an interval including zero; the reverse result supports the importance of matching bias to denoising stage.Source: §F.1 · p. 24

F.2 EFFECTS ON RAW AND STRUCTURED INPUTS

  • Across raw and structured inputs, zeroing the bias produces large reductions, while matched-mean preserves generation quality substantially better than zero.The table reports nine contrasts on 200 paired sources.Source: §F.2 · p. 24, Table 10 · p. 24

G INPUT REPRESENTATIONS AND PROMPT ROBUSTNESS

The study compares three models across six prompt constructions using the same 300 source prompts and original scoring questions. TSGate scores higher than the baseline across all constructions, with its clearest advantage over gated attention on raw and repeated prompts.

G INPUT REPRESENTATIONS AND PROMPT ROBUSTNESS

  • Baseline, Gated attention, and TSGate were evaluated on 300 DPG-Bench source prompts using the same sources and original scoring questions.The six inputs include raw, structured, header-only, semantic-repeat, neutral-padding, and full-recaption prompts.Source: G INPUT REPRESENTATIONS AND PROMPT… · pp. 24–25, Figure 18 · p. 25
  • TSGate has higher mean DPG than Baseline across all six constructions; versus Gated attention, significant paired gains appear on raw and semantic-repeat prompts.The other four paired intervals include zero, and mean differences are negative for structured and header-only prompts.Source: Figure 18 · p. 25, G INPUT REPRESENTATIONS AND PROMPT… · p. 25
  • The six constructions test repetition, appended text, organization, and rewriting, rather than successive levels of degradation.Semantic repeat duplicates the original prompt, neutral padding appends low-information text, and structured prompts reorganize the original scene information.Source: G INPUT REPRESENTATIONS AND PROMPT… · p. 25

Related papers

Loading 2609.34539v1…