PaperlayerView exact revision on arXiv

Source-linked AI summary

Anchoring Instruction Outside Mask: Exact Reference Caching for Efficient In-Context Diffusion Transformers

Yangshuai Liu, Zheming Li, Jiaao Li, Kang He, Ziliang Lai, Zhitai Liu, Chengru Song

arXiv:2608.21229v1cs.CV

TL;DR

Diffusion transformers need efficient multi-reference conditioning without sacrificing instruction-aware visual representations. The paper uses static text anchors and two-stage distillation to preserve exact reference caching while recovering quality, matching full-attention generation across three benchmarks and achieving substantial end-to-end speedups.

  • Problem

    Reference tokens make in-context diffusion computation grow rapidly, while structured sparse attention preserves exact caching by blocking visual references from accessing text instructions.

  • Method

    The method adds parameter-free static text anchors for temporary instruction access during cache construction, then applies teacher-forced and on-policy velocity distillation for architectural recovery.

  • Results

    The method matches full-attention quality across three image-editing benchmarks and delivers a 3.92× end-to-end speedup with five reference images.

  • Takeaways & Limitations

    Exact reference reuse can provide substantial acceleration while retaining full-attention-level generation quality in multi-reference image editing.

Abstract

from arXiv · show

Omnimodal generation is central to a wide range of content creation and editing applications. In-context conditioning is essential to this paradigm. It allows diffusion transformers to process text instructions and visual references in a shared attention sequence. However, each reference image introduces thousands of tokens. Computation therefore grows rapidly with the number of references. Existing methods reduce computation through structured sparse attention, which limits interactions between reference and target tokens. This structure also makes the reference K and V independent of the denoising target, allowing them to be computed once and reused across steps. However, it blocks visual references from attending to the text instruction. This substantially degrades instruction following and reference fidelity in multi-reference editing. To resolve this conflict, we jointly redesign the token sequence and attention mask. Our beyond-mask design uses static text anchors to connect the instruction to the reference branch. It preserves exact K and V reuse without adding parameters. However, this direct architectural conversion degrades generation quality. We recover the lost performance through teacher-forced velocity distillation, followed by a short on-policy stage in which the teacher supervises student-visited states. To our knowledge, this is the first use of on-policy distillation for architectural recovery in diffusion models. Across three image-editing benchmarks, our method matches full-attention generation quality. With five reference images, it accelerates the complete 40-step denoising process by 3.92x, while static text anchors introduce negligible runtime overhead; the speedup reaches 5.47x at ten references in our scaling study.

1 INTRODUCTION

In-context diffusion transformers face a trade-off between exact reference caching and instruction-aware visual references. The paper resolves this through static text anchors, regular attention, and two-stage distillation recovery.

  • Reference images add roughly 4k tokens at 1024^2 resolution, and target-dependent reference states require recomputation at every denoising step.
  • Structured sparse attention enables exact K and V reuse by isolating references, but blocks reference tokens from accessing updated instruction information.
  • The beyond-mask design preserves regular attention and compatibility with optimized kernels, enabling practical exact-caching speedups without specialized kernel development.
  • Static text anchors provide a temporary instruction pathway during reference precomputation without depending on the evolving target or adding parameters.
  • The two-stage recovery procedure combines teacher-forced velocity distillation with on-policy supervision at student-visited states to correct trajectory-specific errors.
  • The method matches full-attention quality across three image-editing benchmarks and achieves a 3.92× speedup for five references, reaching 5.47× at ten references.

2 RELATED WORK

Related work accelerates diffusion transformers through sparse attention, temporal or structural caching, and architectural recovery. On-policy distillation has broader precedents, but prior work primarily targets acceleration or generative training rather than attention-structure recovery.

  • Sparse-attention methods reduce token interactions, but dynamic and irregular patterns can complicate kernel and system design.
  • Temporal caching reuses activations across nearby denoising steps approximately, whereas structural methods enable exact K and V reuse by isolating fixed condition tokens.
  • Architectural changes often cause an initial quality drop; lightweight modifications may recover through short fine-tuning or distillation, while larger operator changes typically require longer retraining.
  • Prior on-policy distillation studies mainly address sampling acceleration or generative training rather than recovery from attention-structure modifications.

3 METHODOLOGY

The methodology combines in-context diffusion processing with instruction-aware exact reference caching and two-stage distillation-based architectural recovery. Static text anchors preserve instruction access without target-dependent reference states, while on-policy refinement addresses residual errors along the student trajectory.

  • In-context diffusion transformers: In-context diffusion transformers jointly process live text, evolving target tokens, and visual-reference tokens, whose hidden states may change through attention to the target.Reference tokens are concatenated across images, increasing sequence length as references are added.
  • Exact reference caching: Exact caching computes reference K and V once when reference states receive no information from step-dependent target states.The cached representations remain unchanged across denoising steps and can be reused throughout denoising.
  • Missing instruction path: Under structured attention, removing target dependencies also blocks the instruction path needed for reference tokens to identify instruction-relevant visual content.The resulting dependency prevents both live instruction access and exact caching within the original sequence structure.
  • Static text anchors: Static text anchors extend the sequence with a separate instruction pathway that connects to references without opening a target-to-reference dependency.The static subgraph is evaluated once, temporary anchor states are discarded, and only instruction-aware reference K and V are retained.
  • Static text anchors: The anchor design adds only a one-time cache-construction cost, with negligible runtime overhead after the anchors are discarded.The additional attention interactions are paid during precomputation rather than every denoising step.
  • Two-stage architectural recovery: Two-stage distillation recovers quality lost by architectural conversion, first using teacher-forced data-derived states and then student-visited states.Stage 1 restores most lost performance, while Stage 2 addresses residual errors along the student’s inference trajectory.

4 EXPERIMENTS

Experiments show that static text anchors preserve full-attention-level quality while retaining the efficiency of exact reference caching. Two-stage recovery, especially on-policy supervision, closes the architectural quality gap, and scaling benefits increase with more references.

  • Main Results: 3.92× speedup reduces complete 40-step latency to 94.4 seconds with eight references, while static anchors add only 0.3 seconds over the isolated cache.The isolated cache takes 94.1 seconds under the same setting.
  • Main Results: 8.185 Overall matches the full-attention base model’s 8.119 and improves the isolated cache from 7.680 at nearly identical latency.The improvement over the isolated cache includes a 0.677 gain in subject consistency.
  • Main Results: The final model closely matches full attention across OmniContext, ImgEdit-Bench, and GEdit-Bench, with Overall scores of 8.185 versus 8.119, 4.351 versus 4.324, and 7.540 versus 7.574.The comparison spans different instruction distributions, metrics, and evaluation protocols.
  • Ablation Study: 500 on-policy updates raise Overall to 8.185, whereas continued teacher-forced training remains nearly saturated at 8.016 ± 0.015.The same Stage 1 checkpoint initializes both recovery strategies, and improvement is already visible after 250 on-policy updates at 8.057.
  • Ablation Study: Full-trajectory sampling with K = 4 is selected as the default OPD configuration because increasing query count does not uniformly improve recovery.The effect of K depends on the query interval, and full-trajectory sampling performs best for both K = 1 and K = 4.
  • Scaling: Speedup increases from 1.82× with one reference image to 5.47× with ten as cached reference representations are reused throughout denoising.Full-attention latency grows rapidly with additional references, while the proposed method grows substantially more slowly.

5 CONCLUSION

The method combines static text anchors with two-stage recovery to make exact reference caching compatible with instruction-aware editing. Across three benchmarks it matches full-attention quality while accelerating 40-step generation, with larger gains as references increase.

  • Conclusion: Static text anchors provide a temporary instruction path during cache construction, while two-stage distillation recovers quality after the architectural conversion.The recovery combines teacher-forced velocity distillation with on-policy supervision at student-visited states.
  • Conclusion: 3.92× end-to-end speedup with five reference images matches full-attention quality, and speedup reaches 5.47× at ten references.Static text anchors introduce negligible overhead in the five-reference setting.

A A GRAPH STATEMENT OF THE IMPOSSIBILITY

The graph formulation proves that the original token set cannot simultaneously preserve exact reference cacheability and live instruction access under the stated connectivity assumptions. Adding a static instruction carrier changes the graph domain and restores both properties structurally.

  • Impossibility Statement: StructCacheable(R) requires that no path from the evolving target X reaches the reference R at any layer.This is a structural guarantee: a target-to-reference path removes the mask-only guarantee of step-independent reference states.
  • Impossibility Statement: LiveTextReachable(R) requires a directed path from a live-text state to the reference at a later layer.The definition captures access to contextualized live text rather than static instruction tokens.
  • Impossibility Statement: Under shared masks and preserved text–target coupling, no mask over {LT, X, R} satisfies both exact cacheability and live-text reachability.Theorem 1 formalizes the incompatibility on the original token set.
  • Construction: Adding static instruction carrier ST lets the reference and ST exchange information without either receiving information from the live target stream.This creates an instruction path to R without creating a path from X to R, preserving structural exact cacheability.
  • Scope: The theorem guarantees structural dependency properties, not numerical activity of every permitted edge or improved editing quality from text reachability alone.These claims are explicitly outside the theorem’s scope.

B WHY THE RECOVERY IS STAGED

The recovery is staged because teacher-forced distillation first restores most quality but eventually saturates, while student-visited states provide a further recovery signal.

  • Stage 1 and Stage 2: Teacher-forced recovery uses data-derived interpolants, whereas on-policy recovery queries states from the student’s own rollout distribution.The student rollout is frozen during each update, and the teacher is queried at those same visited states.
  • Why switching helps: Changing the query distribution matters because a nonzero architectural residual can be weighted differently across state-space regions.Exact pointwise teacher matching is not guaranteed after the architectural change, so different query distributions need not share the same optimum.
  • Caveat: Stage 2 is a moving-query stochastic procedure rather than minimization of one fixed objective, and the paper makes no convergence claim.The student-induced distribution changes with the student parameters across iterations.
  • Stage 1 and Stage 2: Stage 1 avoids student rollouts for each update and recovers most of the quality lost after architectural conversion before Stage 2 begins.Continued Stage 1 training saturates under the reported schedule, while switching distributions improves every evaluated Stage 2 checkpoint.
  • Why switching helps: After teacher-forced recovery plateaus, changing the query distribution provides an effective additional recovery signal.The analysis does not claim that teacher-forced training is mathematically required before every on-policy procedure.

C COST ACCOUNTING

The cost model separates one-time static precomputation from repeated live computation, explaining how reference count affects structural speedup while bounding its interpretation.

  • Cost decomposition: The attention-interaction model splits computation into one-time evaluation of [ST, R] and per-step queries from [LT, Xt].It counts one unit per query–key interaction and compares the resulting structure with full attention.
  • Scaling properties: The analytical speedup is strictly increasing with the number of reference tokens for T > 2.This monotonicity holds under the stated interaction-counting convention.
  • Scaling properties: The model predicts R(NR) < T for every finite reference count, approaching the T-fold ceiling only as NR →∞.Thus no finite reference count reaches a T-fold speedup in this model.
  • Approximation limits: The interaction model characterizes structural scaling rather than exact measured seconds because it omits hardware effects and position-wise computation.Actual caching also avoids reference-side projections and feed-forward work, while memory traffic and kernel overhead remain excluded.
  • Crossover: The crossover estimate NR ≈TS is asymptotic rather than an exact boundary because the precise crossover depends on NST.The denominator combines one-time static precomputation with repeated live computation.
  • Approximation limits: At one reference, the large-context approximation predicts 0.87× while measured end-to-end speedup is 1.82×.The exact model retains NST and should use the actual static-text token count.

D FULL EXPERIMENTAL DETAILS

The experiments use fixed recovery settings and benchmark-specific protocols, while the scoring pipeline introduces reproducibility and interpretation caveats.

  • Recovery setup: The gradient-norm clip is 0.05, and Stage 2 resumes from a Stage-1 checkpoint rather than pretrained weights.These settings make the stages a sequential recovery procedure.
  • Recovery setup: Both recovery stages train the full parameter set in bf16 with DeepSpeed ZeRO-2 and gradient checkpointing on a single node.Table 4 collects the recovery hyper-parameters.
  • Sampling protocols: OmniContext uses 40 denoising steps, 1024 × 1024 references and outputs, γ = 4, blank negative prompts, and seed 42 across 400 examples.GEdit-Bench instead follows its released 50-step protocol with per-example seeds.
  • Sampling protocols: GEdit-Bench numbers are not comparable to OmniContext tables because the benchmarks use different released sampling protocols.The stated differences include denoising steps and seed handling.
  • Scoring caveats: The OmniContext judge uses default API temperature, producing a measured 0.054 rescore spread on unchanged images.Unparseable JSON responses receive random scores in [0, 10], potentially shifting a 50-example sub-task mean by up to 0.2.
  • Scoring caveats: The paper declines to treat differences of a few hundredths as an ordering because scoring randomness and rare parse failures affect reported means.The released evaluation code is retained for comparability with published numbers.

E WHAT DOES NOT WORK

Guidance-aware recovery objectives do not improve on plain on-policy regression in the tested setting, but the result does not isolate their individual components.

  • Guidance choice: The recovery stage uses no guidance, although inference retains true classifier-free guidance at γ = 4.Teacher and student are evaluated under their conditional branches during recovery.
  • Guidance-aware objectives: The tested guided and positive-direction packages each cost about 0.16 Overall relative to plain on-policy regression.Their 0.004 difference is below the reported reproducibility scale.
  • Guidance-aware objectives: Neither deployment-aligned package beats plain on-policy regression under the reported experiments.The negative result applies to the packages as tested, not necessarily to either component alone.
  • Interpretation limit: The negative result cannot separate guided trajectories from branch-aware losses because a guided-trajectory, positive-only control was not run.The tested configurations jointly vary the two components.
  • Interpretation limit: In this setting, shared negative conditioning makes both branch errors fall together, leaving no cross-branch compensation for a branch-aware objective to remove.The remaining cost of guided rollouts is not offset by recovery gains when the residual is small.

F.1 WHERE THE BASE MODEL STANDS, UNDER ONE PROTOCOL

Table 5 positions the base model among recent instruction-based editing systems under a single evaluation protocol, while emphasizing that the comparison is contextual rather than a controlled counterpart to the paper’s recovery experiments.

  • Comparison scope: Table 5 compares recent instruction-based editing systems under one protocol, rather than serving as a controlled counterpart to the paper’s §4 experiments.The systems differ in scale, training data, and inference configuration.
  • Evaluation protocol: Every row was generated and scored with the same judge, prompt template, sampler, step count, resolution, and without per-model guidance or negative-prompt tuning.This uniform setup shifts absolute scores below published figures, including for the authors’ base model.
  • Interpretation: Only relative ordering is intended because published scores use each system’s preferred settings and are not directly reproduced by the uniform protocol.The table caption explicitly cautions that absolute values differ from published results.
  • Evaluation boundary: The table is not comparable to Table 1 or Figure 4, which use the OmniContext protocol, and it does not affect the paper’s controlled teacher-versus-mask comparison.The §4 claims vary only the mask on the same base model.
  • Qualitative examples: Figures 7 and 6 provide ten contiguous GEdit-Bench examples in a three-column layout: input, full-attention teacher, isolated condition, and Text Anchor.Instructions appear above each row, including prompts to change the background to the ocean and remove the umbrella.
Loading 2608.21229v1…