RefAdapt-DiT: Adaptive Joint Attention for Reference-Conditioned Diffusion Transformers

Jian Tang; Jiawei Fan; Qiannan Zhou; Qingbin Liu; Jiang Bian; Zang Li

cs.CVSep 26, 2026

RefAdapt-DiT: Adaptive Joint Attention for Reference-Conditioned Diffusion Transformers

Authors: Jian Tang, Jiawei Fan, Qiannan Zhou, Qingbin Liu, Jiang Bian, Zang Li

Organizations: Platform and Content Group, Tencent

Abstract

Diffusion Transformers (DiTs) have become the standard backbone for high-quality generative modeling, yet deploying them in conditional generation tasks remains computationally prohibitive because bidirectional joint attention repeatedly processes large reference streams. While existing optimization schemes mitigate generic temporal redundancy, they typically rely on coarse-grained static reuse and overlook the distinct dynamics of references and targets. Specifically, we observe that reference representations often evolve slowly along the generation trajectory, while the target often assigns little attention mass to them; reference drift and this target-to-reference exposure jointly shape how strongly stale reference states affect the target. To exploit these patterns, we introduce \RefAdapt, a training-free framework for adaptive control of joint attention between references and targets. Instead of rigid static strategies, \RefAdapt combines consecutive target-Q change with previously observed target-to-reference attention mass to control reference computation adaptively at block granularity. Under ultra-few-step settings, \RefAdapt enables speedups of up to $2.097\times$ on 4-step MiniMax H3 and $3.54\times$ on 8-step Qwen Image Edit, while maintaining comparable visual quality.

Figures & tables

Figure 1: Matched reference–target K/V dynamics. Heatmaps show transition (horizontal), block (vertical), and prompt-averaged $\rho$ (color; negative means slower Ref motion). Right profiles show mean Ref drift horizontally and block vertically; pink spans prompts. H3 uses the nine shared blocks.

Model (steps $\times$ blocks)	Ref mass	Tgt mass	Tgt/Ref
MiniMax H3 ( $49\times 50$ )	9.51%	90.49%	$9.52\times$
Qwen Image Edit ( $40\times 60$ )	14.66%	85.34%	$5.82\times$

Table 1: Trajectory-wide directional attention allocation in full-fresh joint attention.

Figure 2: Adjacent-step directional Ref-mass stability. Horizontal: block; vertical: mean change $\bar{\Delta}_{b}$ in percentage points (shared scale). Mass is exposure, not causal effect or a standalone reuse rule.

Panel of Figure 3 — Figure 3: Exposure-conditioned endpoint score versus local perturbation. Each point is one $(t,b,h)$ comparison: horizontal $E^{KV}M^{\mathrm{Y}\rightarrow\mathrm{R}}$ , vertical $I^{KV}$ , and color current per-head exposure $M$ to the operational Ref set defined above. This diagnostic does not define the deployed controller score.

Figure 4: RefAdapt-DiT at one step. Per-head target-Q change

D

and lagged reference exposure

M

form

D\times M

; the first block whose head mean reaches

\theta

sets the depth cutoff. Earlier blocks reuse , the boundary reactivate s, and deeper blocks compute .

Method	Type	Efficiency		MultiBanana Evaluation
Method	Type	Speedup $\uparrow$	Time $\downarrow$	IA $\uparrow$	RC $\uparrow$	BSM $\uparrow$	PR $\uparrow$	VQ $\uparrow$	MB Score $\uparrow$
Full computation	–	1.00 $\times$	60.39	1.286	1.735	1.408	1.878	2.163	14.512
TeaCache ( Liu et al., 2025a )	C	1.13 $\times$	53.30	1.367 [-1pt] ( $+0.081$ )	1.571 [-1pt] ( $-0.164$ )	1.306 [-1pt] ( $-0.102$ )	1.816 [-1pt] ( $-0.062$ )	2.041 [-1pt] ( $-0.122$ )	13.977 [-1pt] ( $-0.535$ )
SeaCache ( Chung et al., 2026 )	C	1.60 $\times$	37.63	1.347 [-1pt] ( $+0.061$ )	1.959 [-1pt] ( $+0.224$ )	1.551 [-1pt] ( $+0.143$ )	1.776 [-1pt] ( $-0.102$ )	2.122 [-1pt] ( $-0.041$ )	15.367 [-1pt] ( $+0.855$ )
TaylorSeer ( Liu et al., 2025b ) ( $I{=}2$ )	F	1.61 $\times$	37.58	1.396 [-1pt] ( $+0.110$ )	1.875 [-1pt] ( $+0.140$ )	1.438 [-1pt] ( $+0.030$ )	2.083 [-1pt] ( $+0.205$ )	2.271 [-1pt] ( $+0.108$ )	15.605 [-1pt] ( $+1.093$ )
ToMe ( Bolya & Hoffman, 2023 )	T	1.06 $\times$	56.74	1.200 [-1pt] ( $-0.086$ )	1.180 [-1pt] ( $-0.555$ )	1.120 [-1pt] ( $-0.288$ )	1.280 [-1pt] ( $-0.598$ )	1.180 [-1pt] ( $-0.983$ )	10.720 [-1pt] ( $-3.792$ )

Table 2: Qwen Image Edit on 100 MultiBanana cases (8 steps). MB Score is the official aggregate; time is in seconds. ResComp denotes orthogonal resolution compression.

Figure 5: Qwen qualitative comparison for two multi-reference cases. Columns show full computation and four accelerated methods.

Figure 6: H3 qualitative comparison over sampled frames. Rows show full computation, SeaCache, and RefAdapt-DiT .

Method	Type	Efficiency		VLM Eval.	Video Qual.	Text Align.		Temporal Cons.
Method	Type	Speedup $\uparrow$	Time $\downarrow$	EditQ $\uparrow$	Pick $\uparrow$	Frame $\uparrow$	Video $\uparrow$	CLIP $\uparrow$	DINO $\uparrow$
Full computation	–	1.000 $\times$	200.80	7.280	20.528	28.243	24.994	0.990	0.990
TeaCache ( Liu et al., 2025a )	C	1.341 $\times$	149.69	7.340 [-1pt] ( $+0.060$ )	20.509 [-1pt] ( $-0.019$ )	28.294 [-1pt] ( $+0.051$ )	25.051 [-1pt] ( $+0.057$ )	0.990 [-1pt] ( $+0.000$ )	0.990 [-1pt] ( $+0.000$ )
SeaCache ( Chung et al., 2026 )	C	1.334 $\times$	150.52	7.060 [-1pt] ( $-0.220$ )	20.389 [-1pt] ( $-0.139$ )	28.197 [-1pt] ( $-0.046$ )	24.800 [-1pt] ( $-0.194$ )	0.988 [-1pt] ( $-0.002$ )	0.986 [-1pt] ( $-0.004$ )
TaylorSeer ( Liu et al., 2025b )	F	1.333 $\times$	150.67	7.361 [-1pt] ( $+0.081$ )	20.506 [-1pt] ( $-0.022$ )	28.305 [-1pt] ( $+0.062$ )	25.062 [-1pt] ( $+0.068$ )	0.990 [-1pt] ( $+0.000$ )	0.990 [-1pt] ( $+0.000$ )
ToMe ( Bolya & Hoffman, 2023 )	T	1.207 $\times$	166.37	5.028 [-1pt] ( $-2.252$ )	19.851 [-1pt] ( $-0.677$ )	27.755 [-1pt] ( $-0.488$ )	24.614 [-1pt] ( $-0.380$ )	0.988 [-1pt] ( $-0.002$ )	0.985 [-1pt] ( $-0.005$ )

Table 3: MiniMax H3 on 100 EditVerseBench cases (4 steps). Quality follows the official protocol; time is in seconds. ResComp denotes orthogonal resolution compression.

Configuration	Controller cue	Speedup $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$
Full computation	–	1.000 $\times$	–	1.0000
Fixed R3	fixed interval	1.464 $\times$	18.00	0.8823
Target-Q	consecutive target-Q	1.459 $\times$	29.81	0.9879
Exposure	directional exposure	1.459 $\times$	32.24	0.9958
RefAdapt-DiT	target-Q $\times$ exposure	1.482 $\times$	31.24	0.9937
Anchor-Q $\times$ Exposure	anchor-Q $\times$ exposure	1.376 $\times$	29.78	0.9864

Table 4: MiniMax H3 controller ablation on a separate 50-step Ref2VA trajectory. Adaptive variants share the executor and differ only in trigger cue; Fixed R3 is non-adaptive. Speedup uses the depth-wise executor; PSNR/SSIM are same-seed diagnostics from the signal-matched closed-loop replay and indicate trajectory fidelity, not absolute perceptual quality.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Method	3 refs	4 refs	5 refs	6 refs	7 refs	8 refs	Overall
ToMe	1.09 $\times$	1.08 $\times$	1.07 $\times$	1.07 $\times$	1.06 $\times$	1.07 $\times$	1.06 $\times$
TeaCache	1.13 $\times$	1.14 $\times$	1.13 $\times$	1.14 $\times$	1.13 $\times$	1.13 $\times$	1.13 $\times$
SeaCache	1.59 $\times$	1.60 $\times$	1.60 $\times$	1.60 $\times$	1.57 $\times$	1.62 $\times$	1.60 $\times$
TaylorSeer ( $I{=}2$ )	1.59 $\times$	1.60 $\times$	1.60 $\times$	1.60 $\times$	1.60 $\times$	1.61 $\times$	1.61 $\times$
RefAdapt-DiT	2.09 $\times$	2.26 $\times$	2.38 $\times$	2.49 $\times$	2.57 $\times$	2.62 $\times$	2.58 $\times$

Appendix

Table 5: Qwen generation speedup by number of reference images. TaylorSeer uses interval 2. Speedup is relative to full computation at the same reference count.

Diagnostic	H3 Ref2VA	Qwen multi-image
Raw endpoint vs. $I^{KV}$	0.564	0.708
Endpoint $\times$ exposure vs. $I^{KV}$	0.826	0.905
Target-Q change vs. $I^{KV}_{t+1}$	0.607	0.710
Target-Q $\times$ runtime exposure vs. $I^{KV}_{t+1}$	0.869	0.846
Sensitivity blocks improved by exposure	9/9	7/7

Appendix

Table 6: Mechanism statistics for the current canonical H3 and Qwen cases. Spearman correlations are within-case diagnostics and are reported separately for each setting.

Figure 7: Retained left-half panels from the full-block Qwen reference/target dynamics diagnostics. These four panels report target and mean-reference consecutive drift for K and V across blocks and denoising steps; the dropped right-half panels contained the corresponding log-ratio and per-block timing-coupling summaries.

Figure 8: Complete H3 raw-partition attention diagnostics. The main text keeps target exposure to Ref and aggregates the remaining mass as Tgt; the remaining panels report implementation-level allocation and directional asymmetry. Every arrow in the panel titles runs from query role to key role, not along causal flow.

Signal / target	Pearson $\uparrow$	Spearman $\uparrow$
Runtime stale-reference exposure vs. $I^{KV}_{t+1}$	0.479	0.789
Target-Q anchor change vs. $I^{KV}_{t+1}$	0.390	0.418
Anchor change $\times$ exposure vs. $I^{KV}_{t+1}$	0.589	0.545

Appendix

Table 7: Additional H3 runtime-signal diagnostics from one complete Ref2VA trajectory. Correlations are within-case screening statistics, not confidence estimates over independent inputs.

Panel of Figure 9 — Figure 9: Block-wise support for exposure conditioning. In every measured sensitivity block, endpoint drift weighted by target exposure has higher Spearman association with local stale-KV impact than raw endpoint drift (9/9 H3 blocks and 7/7 Qwen blocks).

Model	Case	Prompt
MiniMax H3	canonical	Reimagine the animal from the reference video as a giant panda ambling through the snowy forest, same motion and camera path, breath misting, soft wind and distant birdsong
MiniMax H3	red_fox_autumn	Reimagine the animal from the reference video as a red fox trotting through an autumn maple forest, preserve the original motion and camera path, rustling leaves, warm sunset light, and natural forest ambience
MiniMax H3	robot_dog_neon	Reimagine the animal from the reference video as a small silver robot dog exploring a rain-soaked neon city street, preserve the original motion and camera path, reflected lights, light rain, and distant traffic ambience
Qwen	canonical	the woman and man are hugging together
Qwen	outdoor_cafe	Place the woman and the man seated across from each other at a small outdoor cafe table, preserving both identities and realistic facial details
Qwen	umbrella_neon	Place the woman and the man standing back-to-back under one transparent umbrella on a rainy neon-lit street, preserving both identities and realistic facial details

Appendix

Table 8: Exact prompts used to construct the three-prompt reference-drift profiles in Figure 1 .

Geometry-Conditioned Image Generation

See all 11 papers

Beyond Attention Masks: Instruction Anchoring for Efficient In-Context Diffusion Generation
Aug 21, 2026Yangshuai Liu, Zheming Li, Jiaao Li +4Diffusion Transformer Diffusion Model Caching
Harbin Institute of Technology · KlingAI Research
AReA: Adaptive Receptive-Field Attention for Efficient High-Resolution Diffusion Transformers
Oct 12, 2026Javad Rajabi, Koorosh Roohi, David B. Lindell +1Diffusion Transformer Diffusion-Based Image Generation
University of Toronto · Vector Institute
RefRoute: Decoupling Conditioning Cost from References via Compact Residual Conditioning and Spatial Routing
Oct 6, 2026Wanning He, Yuyao Zhang, Yu-Wing TaiDiffusion Transformer Image Generation
Dartmouth College

Aug 21, 2026J/K move · Enter open · S save

Beyond Attention Masks: Instruction Anchoring for Efficient In-Context Diffusion Generation

Yangshuai Liu, Zheming Li, Jiaao Li, Kang He+3

Harbin Institute of Technology · KlingAI Research

Diffusion Transformer Diffusion Model Caching Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

In-context diffusion transformers concatenate instruction, target, and reference tokens into a single sequence for joint attention. Reference-side computation must therefore be repeated at every denoising step, with the cost growing rapidly as more references are added. Decoupling reference tokens from the target enables exact key-value reuse across denoising steps, but prevents the references from attending to the instruction, degrading instruction following and reference fidelity. This trade-off cannot be resolved through attention-mask design alone. We introduce AnchorCache, a parameter-free token-layout and attention-mask co-design that inserts static text anchors. These anchors condition the reference representations on the instruction during cache construction, after which the resulting reference keys and values can be reused exactly across denoising steps. To recover the quality initially lost through this structural conversion, we apply teacher-forced velocity distillation followed by a short on-policy stage that queries the teacher at student-visited states. To our knowledge, this is the first use of on-policy distillation for architectural recovery in diffusion models. Across benchmarks spanning image, speech, and video generation, AnchorCache matches full-attention quality. Its efficiency gains increase with the reference-context size, reaching a 6.40x speedup in diffusion transformer inference.

Abstract

Geometry-Conditioned Image Generation

Explore similar work

Beyond Attention Masks: Instruction Anchoring for Efficient In-Context Diffusion Generation

AReA: Adaptive Receptive-Field Attention for Efficient High-Resolution Diffusion Transformers

RefRoute: Decoupling Conditioning Cost from References via Compact Residual Conditioning and Spatial Routing