Linearizing Vision Transformer with Test-Time Training (SD3.5-T5)

Quick answer

SD3.5-T5 replaces Softmax attention in Stable Diffusion 3.5 with a Test-Time Training (TTT) layer and reaches comparable text-to-image quality after only 3000 fine-tuning steps (about 1 hour on 4 H20 GPUs), then runs 1.32x faster at 1024x1024 and 1.47x faster at 2048x2048 on a single H20. On GenEval it scores 0.69 against 0.70 for the fine-tuned Softmax model and 0.66 for the original SD3.5; on DPG-Bench it edges the original (84.43 vs 83.83). The reason it works so cheaply: TTT’s two-layer dynamic form is structurally the same shape as Softmax attention, so the pretrained Q/K/V/O projections transfer directly instead of being thrown away. The claim is conversion economics on diffusion image generation, not that linear attention now equals Softmax in general.

Why converting Softmax to linear attention usually fails

Linear-complexity attention removes the quadratic cost that makes high-resolution diffusion slow, but training a linear model from scratch costs as much as training the original. The shortcut everyone wants is to inherit a pretrained Transformer’s weights and fine-tune briefly. That shortcut breaks because the representation a Softmax attention layer learns does not fit a linear kernel: prior conversion methods like LiT-style grafting only keep the MLP weights and discard the attention projections, throwing away the part that actually does the mixing. In the paper’s ImageNet probe, plain linear attention fine-tuned from a DeiT backbone reaches 63.30% Top-1, LiT 69.52%, CLEAR 68.66%, and Hedgehog diverges to NaN, all against the 72.05% Softmax DeiT-Tiny. The gap is the cost of the lost attention knowledge.

How TTT inherits Softmax attention weights

The genuine idea is a structural match, not a new kernel. The paper rewrites Softmax attention as a two-layer transform with a nonlinearity in the middle: a query passes through one weight matrix built from the keys, gets a nonlinear squash, then passes through a second matrix built from the values. Test-Time Training has exactly this shape. Its inner model is a small two-layer MLP whose first layer is adjusted by one gradient step at inference, so it also reads as query, nonlinear middle, output. Because the two architectures line up term for term, the pretrained Q, K, V, and O projections plug straight into the TTT layer and keep their meaning. Unlike a fixed linear kernel, the nonlinear inner MLP has enough capacity to fit the function a Softmax layer learned, which is why the inheritance sticks instead of degrading.

The ablation isolates this. On a DiT-S/2 generator, full weight inheritance gives FID 68.52; inheriting only the MLP weights gives 89.71; inheriting only the attention weights gives 93.18. Keeping both is what holds quality, and the attention weights are not optional.

The two fixes that make inheritance trainable

Architectural alignment alone is not enough because Softmax and TTT differ in two representational habits, and the paper patches each.

First, key shift-invariance. Softmax cancels any constant offset added to all keys, so a shifted key set produces the same attention. TTT has no such invariance, and pretrained keys carry a large shift: across 5K ViT activations the measured key shift ratio is about 0.5, versus about 0.07 for random init. Feeding those off-center keys into the TTT inner optimization makes its gradient explode, which is why a naive port hits NaN. Key instance normalization centers the keys across the sequence before they enter the layer, and that single change is the difference between divergence and 71.19% Top-1 on DeiT-Tiny.

Second, locality. Softmax attention in a ViT quietly behaves locally, and TTT loses that bias. A lightweight module restores it: a depthwise convolution on queries and keys lifts DeiT-Tiny from 69.25% to 71.19%, and adding a small neighborhood-attention term reaches 72.06%, which actually passes the 72.05% Softmax baseline. The locality module is cheap engineering bolted onto the real architectural insight, and the paper is honest that it is doing inductive-bias repair, not magic.

Key results

• GenEval (SD3.5): SD3.5-T5 scores 0.69 overall vs 0.70 for the fine-tuned Softmax model and 0.66 for the original SD3.5, so conversion holds prompt-following quality. Per-axis: Single 1.00, Two 0.88, Counting 0.58, Color 0.88, Position 0.25, Color-Attr 0.55.
• DPG-Bench (SD3.5): 84.43 vs 83.83 for the original SD3.5-Medium, a small win on dense prompt fidelity.
• Speedup: latency at 1024x1024 drops from 25s to 19s on a single H20, a 1.32x gain; at 2048x2048 the gain rises to 1.47x, because the quadratic term grows with sequence length and is exactly what the linear layer removes.
• Conversion cost: 3000 steps, batch size 64, on Flux-generated 1024x1024 images, about 1 hour on 4 H20 GPUs. For class-conditional DiT, the linearized DiT5-XL/2 reaches FID 2.48 vs the original 2.27 after 8 epochs against a 1400-epoch baseline, roughly 0.57% of the training cost.
• ImageNet conversion (DeiT-Tiny, 30 epochs): T5-T at 71.19% and T5+-T at 72.06% vs the 72.05% Softmax DeiT-Tiny, while plain linear attention sits at 63.30% and Hedgehog diverges.
• Weight-inheritance ablation (DiT-S/2): full inheritance FID 68.52 vs MLP-only 89.71 vs attention-only 93.18, proving the attention projections carry transferable knowledge.

How to read the headline

The speedup is real but conditional, and reading it as “linear attention now matches Softmax” overshoots. The 1.32x and 1.47x gains appear at 1K and 2K resolution because that is where the quadratic cost dominates; at low resolution or short sequence length the linear layer has little to win and can even lose, which is why the paper reports its image-generation speedups specifically at 1K and 2K. The quality claim is also “comparable,” anchored to GenEval 0.69 vs 0.70 and DPG-Bench 84.43 vs 83.83, not a clean win, and the paper reports no direct FID for SD3.5 itself. The honest framing: for a high-resolution diffusion Transformer you already have, this is a cheap way to buy a real inference speedup at near-parity quality, and the transferable insight is that picking a linear architecture whose math lines up with Softmax beats inventing a clever kernel. The same instinct of grafting pretrained Softmax knowledge into a cheaper attention scheme drives Full Attention Strikes Back, and the broader push to make diffusion Transformers cheaper at high resolution also runs through LoomVideo.

Limits and open questions

There is no dedicated limitations section, which for a conversion method with this many knobs is worth noting. The quality is comparable, not better: 0.69 vs 0.70 on GenEval is inside noise, and Position at 0.25 stays weak, so the method preserves the base model’s failures rather than fixing them. The speedup is resolution-gated; nothing here helps below 1K, and the paper does not report where the crossover sits. The TTT-specific parameters need a 20x learning-rate multiplier and key instance normalization just to avoid NaN, so the recipe is sensitive and may not transfer to other backbones without retuning. The SD3.5 results lean on GenEval and DPG-Bench with no direct FID, and the fine-tuning data is Flux-generated synthetic images, which leaves open how much of the parity reflects distilling Flux rather than preserving SD3.5. Finally, every result is image generation and ImageNet classification; nothing here tests video, long-sequence, or text settings where linear attention is most wanted.

FAQ

Does SD3.5-T5 match the fine-tuned Softmax model on text-to-image quality (GenEval comparison)?

Close, not above. SD3.5-T5 scores GenEval 0.69 against 0.70 for the fine-tuned Softmax SD3.5 and 0.66 for the original, and DPG-Bench 84.43 against 83.83 for the original. The paper’s word is “comparable,” and it reports no direct FID for SD3.5, so treat this as parity within noise rather than a quality gain. The actual payoff is the 1.32x and 1.47x inference speedup at 1K and 2K.

Why can Test-Time Training inherit pretrained attention weights when linear attention cannot (the architecture insight)?

Softmax attention can be written as a two-layer transform: a query through a key-built matrix, a nonlinearity, then a value-built matrix. TTT’s inner model is also a two-layer MLP with a nonlinear middle, so the two line up term for term and the pretrained Q/K/V/O projections transfer with their meaning intact. A plain linear kernel has no nonlinear middle and cannot represent the function Softmax learned, which is why linear-attention conversion lands at 63.30% Top-1 on DeiT-Tiny while T5 reaches 71.19%.

Where do the 1.32x and 1.47x speedups actually show up?

At high resolution. The 1.32x is at 1024x1024 (latency 25s to 19s on one H20) and 1.47x at 2048x2048. The linear layer removes the quadratic term in sequence length, so the saving grows with resolution and is small or absent at low resolution. Reading the speedup as a universal win is the main misread.

What breaks if you skip key instance normalization in SD3.5-T5?

Training diverges to NaN. Pretrained keys carry a large constant shift (measured shift ratio about 0.5 vs 0.07 for random init) that Softmax cancels but TTT does not, so the inner optimization gradient explodes. Centering the keys across the sequence with instance normalization is what turns NaN into a working 71.19% Top-1 on DeiT-Tiny. Hedgehog, which has no such fix, diverges in the same probe.

How cheap is the conversion compared to training from scratch?

For SD3.5 it is 3000 steps at batch size 64 on Flux-generated 1024x1024 images, about 1 hour on 4 H20 GPUs. For class-conditional DiT it is 8 epochs against a 1400-epoch from-scratch baseline, roughly 0.57% of the original training cost, landing FID 2.48 vs 2.27. That cost gap is the paper’s real argument, alongside related efficiency work like MiniMax sparse attention.

One line: SD3.5-T5 shows that picking a linear architecture whose two-layer math matches Softmax lets you inherit pretrained attention weights and buy a 1.32x-1.47x high-resolution speedup at near-parity quality for about an hour of fine-tuning. Read the original paper on arXiv.