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

In plain words

Newer image and video generators read the text prompt and the picture in one combined stream, so the text and the picture influence each other inside every attention step. Older editing tricks were built for a design where the text and the picture were handled separately, and they break down on this design.

QK-Edit makes an edit in two passes. First it runs the source image or video through the model and remembers the part of the attention inputs that describes the picture's layout. Then it generates the edited version with the target prompt, but at selected steps it swaps in the remembered picture layout. The text part keeps following the new prompt, so the new content appears, while the layout of the original stays. No retraining is needed, and the change is small enough to fit into an existing model's inference code.

The authors show it on object addition, removal and replacement, background, material, colour and style changes, for images and for longer, higher-resolution videos.

TL;DR

• Training-free editing for MM-DiT models: FLUX.1-dev for images and HunyuanVideo for videos §4.
• Analysis: in the joint attention map, the text-to-vision block has little effect. The vision-to-text block decides which text word goes where, and the vision-to-vision block decides the spatial layout §3.2.
• Method: record the source vision query and key at chosen steps, then replace the target's vision query and key with them during denoising §3.3 alg. 1.
• Author-reported: best scores on all three PIE-Bench metrics against four image baselines tab. 1, and best scores on the TGVE+ generative setting tab. 2. On the real-video setting, the text claims the best result overall, but the table disagrees on one metric (see Pitfalls).
• No official code was found. Implementation details that the paper leaves open are listed under Pitfalls.

Problem

• In U-Net diffusion models, text-to-image cross-attention and image self-attention are separate modules, so an editor can change one and keep the other §3.1 eq. 1 eq. 2.
• In MM-DiT models, text and image tokens are concatenated and passed through a single self-attention, so the two kinds of attention share one map §3.1 eq. 3 eq. 4. Attention-map editing methods designed for U-Net, such as Prompt-to-Prompt and PnP, give poor results when applied directly to FLUX §3.2 fig. 2.
• The fused attention is computed by optimised kernels, such as flash-attention, that do not expose the attention map. Editing the map directly is therefore hard in practice §3.3.
• Real inputs need an inversion pass to get their starting noise and their attention inputs. Inversion quality limits the edit, most of all for complex video motion §5.

Method

Setting. The source image or video V^src, the source prompt P^src and the target prompt P^tgt are given. The goal is to change the content named in the target prompt and keep the unedited regions §3. Synthetic sources keep their query and key from generation; real sources take them from inversion §3.

Analysis of MM-DiT §3.2.

1. Skipping single blocks of HunyuanVideo changes the video very differently depending on the block. Some blocks change layout, identity or motion, and others change little fig. 3.
2. Skipping the self-attention part of a block changes the output a lot, while skipping the feed-forward part changes little. The authors conclude that the joint attention decides the spatial arrangement, and they focus on it fig. 4.
3. The joint attention map has four quadrants: text-to-text, text-to-vision (query text, key vision), vision-to-text (query vision, key text), and vision-to-vision. Zeroing each quadrant in turn gives these findings fig. 5:
   • Text-to-text: strong values, and the map is not symmetric.
   • Text-to-vision: zeroing it has a negligible effect on the image.
   • Vision-to-text: zeroing it breaks generation. Each column corresponds to the region of one text token.
   • Vision-to-vision: defines the pixel affinities that set the layout.

Editing rule §3.3.

4. Target the vision-to-text and vision-to-vision quadrants, which the analysis links to text-region matching and to layout. The paper writes this as setting these two products to the source's products, Q_vis^src K_text^tgt and Q_vis^src K_vis^src eq. 5 eq. 6.
5. Because the text-to-vision quadrant is negligible, the two replacements can be written as one substitution of the vision part of the target query and key by the source's vision query and key. The target text query and text key are kept eq. 7 eq. 8 fig. 6.
6. The value tensor is not changed alg. 1.
7. The replacement is applied only at the selected denoising steps. The source's query and key must be stored for those steps during the source pass alg. 1 §4.

Pseudocode, following the supplementary algorithm. The "vision" part is the image or video tokens; the "text" part is the prompt tokens.

# QK-Edit, part 1: record source vision queries and keys  [alg. 1]
# Input: source latent, source prompt; Output: saved (query, key) per step
def record_source(source_latent, source_prompt):
    # Run inversion (real input) or generation (synthetic input)  [§3.3]
    saved = {}
    latent = source_latent
    for step in steps:
        # Joint text-and-vision queries, keys and values  [eq. 4]
        queries, keys, values = mmdit_qkv(latent, step, source_prompt)
        # Keep only the vision tokens of query and key  [alg. 1]
        saved[step] = (queries[:, text_len:], keys[:, text_len:])
        latent = inversion_step(latent, step, queries, keys, values)
    return saved

# QK-Edit, part 2: denoise the target with the replaced vision Q and K
# Input: noise, target prompt, saved source pairs, replace_steps
# Output: edited latent  [alg. 1] [eq. 7] [eq. 8]
def edit(noise, target_prompt, saved, replace_steps):
    latent = noise
    for step in reversed(steps):
        # Target joint queries, keys and values  [eq. 4]
        queries, keys, values = mmdit_qkv(latent, step, target_prompt)
        if step in replace_steps:
            # Swap in the source vision query and key; text part stays  [eq. 7-8]
            source_q, source_k = saved[step]
            queries = concat(queries[:, :text_len], source_q)
            keys = concat(keys[:, :text_len], source_k)
        # Values are unchanged  [alg. 1]
        latent = denoise_step(latent, step, queries, keys, values)
    return latent

Implementation notes stated in the paper. The change needs only the storage of the chosen source query and key and the two substitutions in the attention call; the authors say it takes a couple of lines §3.3 §5. The supplementary pseudocode sets the replacement steps with a list T_list and stores the source pairs for every step alg. 1.

Hyperparameters

Results

Author-reported. All numbers are from the paper. Baselines are as listed in the table captions; the paper does not give metric definitions beyond names and references to earlier work §4.2.

Real image editing on PIE-Bench tab. 1. Setting: FLUX.1-dev is the base model for image editing §4. Column names are copied from the table: CLIP-T is the text-image alignment, DINO the structure similarity, PickScore the human preference model.

Generative video editing on TGVE+ tab. 2. Both methods use HunyuanVideo, the source video is generated first, and no inversion is used §4.3.

Real video editing on TGVE+ tab. 3. Setting: each baseline runs on its own inversion path, and the paper does not describe the inversion for each §4.3.

Ablation: which query and key parts to replace on video fig. 8 §4.5. The authors report this qualitatively, with no table.

• Replacing both query and key gives the best balance: it keeps the source structure and follows the target prompt.
• Replacing key and value, or value alone, keeps the structure better but sometimes fails the edit.
• Replacing only the query or only the key follows the prompt better but often loses the structure.

HiLo-Benchmark, high-resolution and long video tab. 4 §4.5. The baselines are the authors' reimplementations on HunyuanVideo: KV-concat (MasaCtrl and Tune-a-Video style), V-replace (RF-Solver style) and KV-replace (Stable Flow style).

Qualitative results fig. 7 cover the editing types listed in the text: object addition, removal and replacement, background change, colour change, style change (dawn to sunset) and material change (steel to gold) §4.4.

Paper vs code

No official code. The paper links no repository and no project page. The only URL in the paper is the flash-attention repository §3.3. In the LaTeX source, a commented-out line in the abstract says that code will be released; this line is not in the published text. I searched GitHub for "QK-Edit", for "QK-Edit" inside the bytedance organisation, and for "query-key manipulation MM-DiT editing". None of the matches belonged to the authors, so no repository was cloned.

The supplementary Algorithm 1 is the closest thing to a reference implementation alg. 1. Some of its details differ from the text:

• The text says the replacement uses the first fraction of steps; the algorithm stores query and key for every step in part 1 and replaces only the steps in T_list in part 2 alg. 1 §4.

Pitfalls

• Quadrant labels are easy to swap. The text calls text-to-vision the second quadrant and vision-to-text the third. Keep the query and key roles fixed when you implement it.
• The substitution is not a pure quadrant edit. Replacing the vision key also changes the text-to-vision quadrant, because queries of the target text are scored against the source vision keys. The paper argues this quadrant has little effect, but it does not measure the change caused by the substitution §3.2 §3.3 eq. 7 eq. 8.
• Block selection is not stated. The analysis shows that blocks differ in their effect fig. 3, but the card cannot say which blocks the replacement covers. Start with all blocks and compare, or ask the authors.
• Step counts. "First 30%" and "first 20%" are fractions. The number of denoising steps and the sampler are not stated, so the absolute step indices are unknown.
• Inversion decides real-input quality. The authors name complex real video motion as a failure source §5. Their real-video result depends on an inversion method the paper does not name §3.3.
• Memory. The source pass stores a vision query and key for each recorded step, for both image and video. The paper gives no memory figures for high-resolution video; the stated hardware is A100 GPUs §4.
• Attention kernels. The replacement changes the inputs to the attention call, so a fused kernel such as flash-attention can still be used, as long as the modified query and key are passed in §3.3 alg. 1.
• Table 3 disagrees with the text. The text says QK-Edit outperforms all baselines on real video. In the ViCLIP_dir column, Flatten has 0.204 and QK-Edit has 0.196, yet the table bolds QK-Edit there tab. 3 §4.3. The other three columns match the claim.
• Metric names differ between table and text. Table 1 writes CLIP-T; the text writes CLIP_text. The paper gives no formula or prompt set for any metric, and the number of PIE-Bench or TGVE+ samples is not stated §4.2.
• HiLo baselines are reimplementations. The comparisons in tab. 4 use the authors' own versions of earlier methods on HunyuanVideo, so they are not the published systems §4.5.
• Known failure modes. The method makes small unintended changes outside the edit region. The authors suggest masks as a fix. Results also inherit the base model's limits: text rendering, and complex motion such as gymnastics §5.
• Scope of the video results. The main videos are 1280×720 with up to 129 frames; the supplement shows portrait output at 720×1280 with 65 frames §4 app. Video Editing Results.

Lineage

• Builds on: MM-DiT models FLUX, Stable Diffusion 3 and HunyuanVideo, which combine text and vision tokens in one self-attention §1 §3.1.
• Analysis draws on U-Net attention editing: Prompt-to-Prompt and PnP §3.1.
• Compared against: LEDITS++, MasaCtrl, RF-Inversion and RF-Solver on images tab. 1; DiTCtrl on generative video tab. 2; CCEdit, Flatten and AnyV2V on real video tab. 3; Stable Flow and Tune-a-Video in the HiLo ablation tab. 4.
• Superseded by: not known.