QK-Edit: Revisiting Attention-based Injection in MM-DiT for Image and Video Editing

AI summary

Overview

Research area: Computer Vision / generative AI — specifically training-free, text-guided image and video editing with Multimodal Diffusion Transformers (MM-DiTs).

Technical level: Advanced. The paper assumes familiarity with diffusion model architectures, attention mechanisms (queries, keys, values), U-Net versus Transformer designs, and inversion-based editing.

One-sentence scope: The paper introduces QK-Edit, a training-free framework that edits images and videos by substituting the vision components of the query and key embeddings inside the Multimodal Self-Attention (MM-SA) of MM-DiT models, aiming to balance prompt fidelity against structural consistency.

What This Paper Is About

MM-DiTs such as FLUX and HunyuanVideo concatenate text and image/video tokens into one joint embedding and run a single self-attention over them, which makes the text–vision and vision–vision interactions intertwined in one attention map. Older editing methods (P2P, PnP) were designed for U-Net models where cross-attention and self-attention are separate modules, so applying them directly to MM-DiTs produces poor results. QK-Edit's goal is to make precise text-guided edits — adding, removing or replacing objects, changing backgrounds, materials, colors, and styles — in both images and videos without any additional training.

Key Contributions

1. A systematic analysis of MM-DiT attention behavior. The authors break the MM-SA attention map into four quadrants and show experimentally that text-to-vision attention (Q_text·K_visionᵀ) has negligible impact on generation, vision-to-text attention (Q_vision·K_textᵀ) decides the text-vision relationship, and vision-vision attention (Q_vision·K_visionᵀ) defines spatial and temporal consistency. They also show that MM-SA, not the FFN, determines the final spatial arrangement, and that skipping the 22nd single block of HunyuanVideo does not produce a visible change while bypassing other blocks causes significant changes.

2. A query-key manipulation strategy for editing. Because the 2nd quadrant barely matters, manipulating the 3rd and 4th quadrants is equivalent to swapping only the vision portion of the query and key embeddings: [Q_text^tgt, Q_vision^tgt] becomes [Q_text^tgt, Q_vision^src], and likewise for K. This sidesteps the fact that attention maps are encapsulated in optimized packages such as flash-attention.

3. A training-free method with minimal implementation cost. QK-Edit requires only two steps: save Q_vision^src and K_vision^src during source generation or inversion, then use them to replace Q_vision^tgt and K_vision^tgt during target generation.

4. A new high-resolution, long-video benchmark. The authors introduce HiLo-Benchmark, consisting of 43 carefully selected challenging videos with resolutions exceeding 1K and sequences longer than 100 frames.

Main Findings

• Real image editing (PIE-Bench): QK-Edit achieves the highest scores on all three metrics — CLIP-T 0.264, DINO 0.776, PickScore 22.33 — against LEDITS++ (0.207, 0.571, 19.88), MasaCtrl (0.235, 0.696, 20.68), RF-Inversion (0.252, 0.523, 21.31), and RF-Solver (0.257, 0.769, 21.98).

• Generative video editing (TGVE+): QK-Edit outperforms DiTCtrl across all metrics, with CLIP_frames 0.951 versus 0.904, PickScore 21.24 versus 20.01, ViCLIP_dir 0.217 versus 0.141, and ViCLIP_out 0.268 versus 0.242. Both methods are based on HunyuanVideo for fair comparison.

• Real video editing (TGVE+): QK-Edit outperforms CCEdit, Flatten, and AnyV2V, scoring ViCLIP_dir 0.196, ViCLIP_out 0.258, CLIP_frames 0.921, and PickScore 20.52 (compared with CCEdit 0.159/0.244/0.877/19.73, Flatten 0.204/0.236/0.887/19.91, and AnyV2V 0.117/0.228/0.906/20.28).

• HiLo-Benchmark: QK-Edit's QK_replace variant scores CLIP_frames 0.908 and PickScore 20.08, compared with KV_concat 0.896/19.87, V_replace 0.900/19.94, and KV_replace 0.901/19.92.

• QKV combination ablation: Replacing QK works best at preserving source structure while following the target prompt; replacing KV or V preserves source structure better but sometimes fails to follow the prompt; replacing Q or K follows the prompt better but often fails to preserve source structure.

• Attention quadrant behavior: The 1st quadrant (Q_text·K_textᵀ) controls relationships among text tokens, is not symmetrical, and its attention values are usually much larger (≫) than the three vision-related quadrants. Setting the 2nd quadrant to zero has a negligible effect, while setting the 3rd quadrant to zero causes FLUX to fail generating a proper image.

• Architecture probing: Each MM-DiT block contains an MM-SA and a Feed-Forward Network (FFN) connected by residual connections; skipping the FFN has negligible effect, while MM-SA decides the final spatial arrangement.

• Editing task coverage: Demonstrated edits include object addition (adding a snowman), object removal (removing boxing gloves), object replacement (woman to panda), changing background (beach to hotel swimming pool), changing color (red to green), style transformation (dawn to sunset), and changing material (steel to solid gold).

• Implementation settings: For image editing, query-key manipulation is applied during the first 30% of diffusion inference steps. For video editing, manipulation is applied during the first 20% of inference steps. Video resolutions range from 320×320 with 65 frames to 1280×720 with up to 129 frames. All experiments run on NVIDIA A100 GPUs.

Methodology in Plain English

The authors started by asking why attention-editing tricks that worked on U-Net diffusion models fail on MM-DiTs. They probed a video model (HunyuanVideo) by skipping one transformer block at a time and observed that some blocks barely matter and others matter a lot. Within a block, they skipped the attention module and then the feed-forward network separately, finding that attention drives spatial and temporal layout while the feed-forward network is nearly irrelevant.

They then split the joint attention matrix into four quadrants defined by which tokens act as queries and which act as keys: text-to-text, text-to-vision, vision-to-text, and vision-to-vision. By zeroing out each quadrant one at a time and visualizing the attention maps, they identified which quadrants actually steer the generation. From this they concluded that only the vision-to-text and vision-to-vision quadrants need to be controlled.

Their editing recipe is then simple: run the model on the source image or video (or invert a real one to noise) while recording the vision queries and keys. When denoising toward the target prompt, keep the new text queries and keys but swap in the saved source vision queries and keys. Because the text-to-vision quadrant is unimportant, this substitution reproduces the effect of surgically rewriting the 3rd and 4th quadrants, without touching the encapsulated attention implementation.

Why This Matters

Impact on research: The paper provides an empirical dissection of how MM-DiT attention actually behaves, correcting the assumption that U-Net editing recipes transfer to Transformer-based diffusion models. It shows that a large class of editing capability is already latent in these models and can be unlocked without fine-tuning, and it contributes HiLo-Benchmark for evaluating high-resolution, long-duration video editing.

Real-world applications:

• Post-production and video editing tools that restyle or retouch footage while preserving motion and identity.
• Advertising and e-commerce content creation, where products need material, color, or background changes (e.g., steel to solid gold, beach to hotel swimming pool).
• Creative tools for photographers and designers performing object addition, removal, or replacement in stills.
• Social media and short-form video production at landscape or portrait aspect ratios with high resolution.

Industry relevance: The work comes from a collaboration spanning UC Merced, ByteDance Seed, ByteDance Intelligent Creation, and CUHK, and targets production-grade base models (FLUX.1-dev, HunyuanVideo). Its training-free, few-lines-of-code nature lowers the cost of adding editing features to existing generative pipelines, and its support for up to 129 frames at 1280×720 addresses scalable multimodal content creation.

Future Directions

• Improving real video inversion accuracy, which the authors note remains challenging in complex scenes with intricate motions and can compromise downstream editing.
• Incorporating masks to prevent the minor unintended modifications QK-Edit introduces in regions that should remain unaltered.
• Mitigating inherited base-model weaknesses such as suboptimal text rendering and structural deficiencies in complex motions (e.g., gymnastic videos) by adopting more advanced generative models.
• Extending the analysis and editing approach to other MM-DiT-based models beyond FLUX.1-dev and HunyuanVideo, and to further editing task types.

Target Audience

Researchers and engineers working on diffusion-based image and video generation and editing, particularly those building on MM-DiT architectures such as FLUX and HunyuanVideo. It also suits practitioners who want a low-cost, training-free editing layer for existing models, and students studying attention manipulation techniques who want a careful architectural comparison between U-Net and Transformer diffusion backbones.

Authors’ abstract

Multimodal Diffusion Transformers (MM-DiTs) have recently emerged as a powerful framework for unified text-vision synthesis, surpassing traditional U-Net architectures in generative tasks. One key innovation lies in its Multimodal Self-Attention (MM-SA) interaction, where image and text tokens are concatenated and processed via self-attention. However, this mechanism poses significant challenges for editing, rendering conventional U-Net-based attention manipulation methods ineffective. To address this limitation, we propose QK-Edit, a training-free framework that exploits the unique attention dynamics of MM-DiTs for precise text-guided image and video editing. By introducing a novel query-key manipulation strategy, our method isolates and adjusts critical attention components to achieve an optimal balance between prompt fidelity and structural consistency. This enables seamless edits across various tasks, including object addition, object removal, object replacement, changing background, changing material, changing color, and style transformation. Notably, it can be easily implemented with feature replacement in inference. QK-Edit demonstrates superior editing performance on state-of-the-art models, such as FLUX and HunyuanVideo, effectively bridging the gap between generative power and editable flexibility in MM-DiTs, and paving the way for scalable multimodal content creation.