Anchoring Instruction Outside Mask: Exact Reference Caching for Efficient In-Context Diffusion Transformers
Proposed a new diffusion transformer architecture using static text anchors for efficient multi-reference image editing, achieving a 3.92x speedup.
Key Findings
Methodology
The paper introduces a novel diffusion transformer architecture using static text anchors to connect instructions to the reference branch, preserving exact K and V reuse without adding parameters. Teacher-forced velocity distillation and a short on-policy stage recover generation quality.
Key Results
- In OmniContext benchmarks, with five reference images, the generation process accelerated by 3.92x, matching full-attention model quality.
- On GEdit-Bench and ImgEdit-Bench, the text-anchor model scored 4.351 and 7.540, close to the full-attention baseline.
- In multi-reference scenarios, the text-anchor method improved reference fidelity while maintaining instruction following.
Significance
This study addresses the computational bottleneck in multi-reference image editing by precisely caching reference image representations, significantly enhancing generation efficiency. The method offers substantial acceleration without sacrificing quality, with important academic and industrial applications.
Technical Contribution
Introduced a parameter-free static text anchor design that makes reference image K and V instruction-aware while preserving exact reuse across denoising steps. First use of on-policy distillation for architectural recovery in diffusion models.
Novelty
This is the first use of on-policy distillation for architectural recovery in diffusion models. Compared to existing sparse attention methods, this design resolves the information flow issue between instructions and reference images through static text anchors.
Limitations
- In certain complex scenarios, it may not fully recover the generation quality of full-attention models.
- Requires additional precomputation steps to build reference caches.
Future Work
Future research can explore combining other acceleration techniques like temporal caching and sparse attention to further enhance diffusion model efficiency and performance.
AI Executive Summary
In the field of image editing, existing diffusion transformer models face significant computational challenges when handling multiple reference images, leading to inefficiency. This paper proposes a new architecture that uses static text anchors to connect instructions to the reference image branch, addressing this issue.
The method achieves efficient generation for multi-reference image editing by precisely caching the K and V representations of reference images. Experimental results show that this method significantly speeds up the generation process while maintaining quality across various benchmarks.
While the method achieves breakthroughs in efficiency, it still faces challenges in maintaining generation quality in complex scenarios. Future research can explore combining other acceleration techniques to further enhance model performance and application scope.
Deep Analysis
Background
Diffusion transformers have made significant progress in image generation and editing, capable of handling text instructions and visual references. However, as the number of reference images increases, computation grows rapidly, becoming a major bottleneck. Existing methods reduce computation through sparse attention but limit interactions between references and targets.
Core Problem
The computational bottleneck in multi-reference image editing. Each reference image introduces thousands of tokens, leading to rapid computation growth. Existing sparse attention methods reduce computation but hinder information flow between instructions and reference images.
Innovation
Proposed a new diffusion transformer architecture using static text anchors to connect instructions to the reference image branch. This design preserves exact K and V reuse while resolving the information flow issue between instructions and reference images.
Methodology
- �� Use static text anchors at t=0 to transfer instruction information to the reference image branch.
- �� Recover generation quality through teacher-forced velocity distillation.
- �� In a short on-policy stage, the teacher supervises student-visited states to correct errors.
Experiments
Evaluated generation quality on OmniContext, GEdit-Bench, and ImgEdit-Bench benchmarks. With five reference images, the generation process accelerated by 3.92x. Compared performance with full-attention models and isolated cache models.
Results
In OmniContext, the text-anchor method achieved an overall score of 8.185, matching the full-attention model. The generation process accelerated by 3.92x. On GEdit-Bench and ImgEdit-Bench, the text-anchor model scored 4.351 and 7.540.
Applications
The method is suitable for scenarios requiring rapid processing of multiple reference images, such as real-time image editing and generation. It significantly improves generation efficiency and reduces computational costs.
Limitations & Outlook
While achieving breakthroughs in efficiency, it still faces challenges in maintaining generation quality in complex scenarios. Requires additional precomputation steps to build reference caches.
Plain Language Accessible to non-experts
Imagine you're cooking in a kitchen. You have a recipe (text instruction) and several ingredients (reference images). Traditional methods require checking each ingredient every time, which is inefficient. This method is like preparing all ingredients in advance, needing only one check. It's like pre-chopping all vegetables in the kitchen, saving time while ensuring every dish follows the recipe.
ELI14 Explained like you're 14
Hey there, imagine you're playing a super cool game with a mission to complete. You have some hints (text instructions) and some tools (reference images). Usually, you need to check these tools every time, wasting time. But this new method is like having all tools ready in advance, needing only one check to complete the mission! Isn't that awesome?
Glossary
Diffusion Transformer
A model used for generating and editing images, capable of handling text instructions and visual references.
Used in this paper for efficient multi-reference image editing.
Static Text Anchors
A method to connect instructions to the reference image branch, preserving exact K and V reuse.
Used to resolve the information flow issue between instructions and reference images.
Teacher-forced Velocity Distillation
A method to recover generation quality by supervising student-visited states.
Used to recover generation quality lost due to architectural conversion.
On-policy Distillation
A method to train student models by supervising states visited by the student itself.
Used to correct errors in the student's inference trajectory.
Sparse Attention
A method to reduce computation by selecting only a subset of token interactions.
Existing methods use this to reduce computation but limit interactions.
Open Questions Unanswered questions from this research
- 1 How to further improve generation quality in complex scenarios?
- 2 How to combine other acceleration techniques to enhance efficiency?
Applications
Immediate Applications
Real-time Image Editing
Suitable for scenarios requiring rapid processing of multiple reference images, such as real-time image editing and generation. Significantly improves generation efficiency and reduces computational costs.
Long-term Vision
Multi-modal Generation Systems
Combine other acceleration techniques to develop a unified generative model capable of handling text, images, audio, and video.
Abstract
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.