Query-Kontext: An Unified Multimodal Model for Image Generation and Editing

Conference: ECCV2026
Paper: ECCV paper page
Authors: Yuxin Song, Wenkai Dong, Shizun Wang, Qi Zhang, Song Xue, Tao Yuan, Hu Yang, Haocheng Feng, Hang Zhou, Xinyan Xiao, Jingdong Wang
Affiliations: Baidu Inc; National University of Singapore
Area: Image Generation
Keywords: multimodal conditioning, query tokens, image editing, subject preservation, shifted RoPE

TL;DR

Query-Kontext transfers semantic and coarse visual planning from a VLM to a diffusion model through trainable multimodal query tokens, adds low-level image conditioning for detail, and unifies generation and editing with 0.88 on GenEval using rewritten prompts and 7.66 / 7.65 on the full English / Chinese GEdit-Bench splits.

Background & Motivation

Unified multimodal models must decide both what an image should contain and how to render it realistically. The first problem includes interpreting editing instructions, grounding objects, and composing reference subjects; the second depends on a diffusion model's texture and image synthesis capabilities. In a straightforward assembly of pretrained components, a frozen VLM often serves as a static semantic feature extractor whose outputs need not capture the visual organization required for generation. Jointly training understanding and generation within a tightly integrated model can instead create competition for capacity.

Query-Kontext focuses on cooperation between components rather than a new unified backbone trained from scratch. A relatively inexpensive diffusion head first teaches the VLM to emit generation-oriented multimodal conditions. These conditions are then connected to a stronger diffusion model, so learning instructions, spatial relationships, and subject identity does not require high-fidelity rendering from the outset. Coarse conditioning still cannot preserve every input detail, motivating a low-level image pathway in the final stage.

Core Idea: first train complementary text and query pathways that express the intended output, then upgrade the renderer, and finally recover detail through image latents with task-specific roles, progressively aligning semantic decisions, rendering capacity, and pixel fidelity.

Method

Overall Architecture

The inputs are text and, when needed, a source image for editing or reference images for subject conditioning. The output is always an image, covering text-to-image generation, instruction editing, customized generation, and multi-subject composition. Qwen2.5-VL-7B processes the multimodal inputs and learnable queries, producing text tokens and a fixed-length sequence of kontext tokens. These tokens describe the intended content and use of references; they are not a readable chain-of-thought.

Text and query outputs enter the diffusion model through a two-layer MLP connector. The final model additionally places VAE-encoded source or reference image latents alongside the noisy target latents in the MMDiT context. Low-level latents preserve structure and texture, while shifted RoPE distinguishes images and their source/reference roles. The diagram follows how this architecture is built during training; inference does not sequentially invoke three separate generators.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Text and source or reference images"] --> B["Text and query pathways<br/>Lightweight diffusion head trains conditions"]
    B --> C["Large diffusion model realignment<br/>Freeze VLM and train renderer"]
    C --> D["Low-level conditioning and shifted RoPE<br/>Distinguish source and reference roles"]
    D --> E["Final MMDiT<br/>Jointly process conditions and noisy latents"]
    E --> F["Generated or edited image"]

Key Designs

1. Text and query pathways: retain instructions while learning output-oriented visual conditions

Text features alone can omit how input images should participate in generation. Compressing everything into queries, however, forces the complete instruction through the same bottleneck. Query-Kontext retains both: text tokens provide explicit semantic context, while 128 kontext tokens aggregate task-relevant information from images and text into coarse visual conditions. A lightweight connector aligns both streams with the diffusion conditioning space. The queries should therefore not be confused with discrete pixel tokens representing the generated image.

Stage 1 connects Qwen2.5-VL-7B to a lightweight MMDiT diffusion head with approximately 870M parameters, training the VLM's LoRA, the connector, the queries, and the entire diffusion head. Text-to-image supervision requires the queries to carry textual requirements; reconstruction teaches visual content; image transformation teaches changes conditioned on an input image and instruction. The claimed generative reasoning emerges through image-task supervision. The main paper does not provide separate reasoning annotations, an explicit chain-of-thought loss, or a reasoning-accuracy benchmark.

2. Large diffusion model realignment: replace the renderer without making the semantic model absorb every change

Stage 2 discards the lightweight head and substitutes a larger, internally pretrained MMDiT, described as roughly a 10-fold parameter increase. Stage 1 LoRA weights are merged into the VLM, whose full backbone is then frozen. Learnable queries, the connector, and the large diffusion model remain trainable. Thus, freezing the VLM does not freeze every learnable condition entering it, and this stage is not connector-only training.

Training uses only text-to-image generation and image reconstruction to align the new conditioning interface with the strong renderer first. The authors find that a small diffusion head is relatively easy to bridge while frozen, whereas a lightweight connector struggles to adapt a large frozen diffusion model. They consequently fine-tune the entire diffusion backbone in Stage 2. Figure 6 shows semantic layouts with blurry textures after Stage 1 and preserved layouts with improved detail after Stage 2. This supports the proposed division of labor, but does not establish complete independence between semantics and texture.

3. Low-level conditioning and shifted RoPE: distinguish preserving a source image from borrowing its identity

Stage 3 introduces VAE-encoded image latents, giving the diffusion model direct access to fine structure and texture instead of requiring fixed-length queries to remember every local detail. Text, queries, image conditions, and the current noisy target latents enter MMDiT together and interact through in-context attention. Figure 5 shows benefits for adding glasses and removing extra wings, but these are qualitative comparisons, not a source of quantified improvement.

Reusing the same position coordinates for several images can confuse an editing source whose layout should remain intact with a reference whose subject identity should be reused. The proposed scheme places the source image in the negative coordinate region and shifts the n-th reference into a positively offset region. The coordinate mappings in original Eqs. (1) and (2) are:

\[ (i_{\mathrm{ref}}^{n},j_{\mathrm{ref}}^{n})=(i+wn,j+hn),\qquad (i'_{\mathrm{src}},j'_{\mathrm{src}})=(-i,-j). \]

Here \(i\in[0,w-1]\) and \(j\in[0,h-1]\) index a latent grid of height \(h\) and width \(w\), with \(n\in[1,N]\) and at most \(N=2\) reference images in the implementation. These coordinates control shifted two-dimensional RoPE; they do not translate pixel content outside the image. Source mode orders the image latents as source followed by noise, emphasizing pixel fidelity. Reference mode uses noise followed by references, allowing layout and pose changes to satisfy the instruction.

A Worked Example

Consider the instruction to add glasses to the woman in Figure 5. The input is an editing source, not an identity reference that permits arbitrary recomposition. The VLM reads the portrait and instruction together. Text tokens retain the requested addition, and the 128 query outputs supply multimodal conditions for the intended image; the connector exposes both representations to the diffusion model.

The same image passes through the VAE to obtain low-level latents. Its source coordinates use the negative mapping above, and its latents enter MMDiT alongside the noisy target. Semantic conditions identify what should be added and to whom; low-level conditions help preserve the face and unedited regions during generation. A request to place the same subject in a new scene instead uses the reference role and positive offsets to reduce constraints from the original layout. This last step illustrates the mechanism, not an additional quantitative experiment in the paper.

Loss & Training

The main paper describes noise prediction in Stage 1 and task supervision across stages, but does not specify the full loss, noise schedule, sampler, or inference step count. A conventional noise-prediction mean-squared error or flow-matching equation is therefore not presented here as the paper's verified objective.

The default resolution in Table 3 is \(512\times512\). Stage 1 uses 72K steps, global batch size 512, and learning rate \(10^{-4}\); Stage 2 uses 420K steps, batch size 1024, and learning rate \(10^{-4}\); Stage 3 uses 30K steps, batch size 512, and learning rate \(2\times10^{-5}\). Stage 3 freezes the VLM and updates the queries, connector, and diffusion-model LoRA across generation, editing, customization, and multi-subject tasks.

VLM LoRA uses rank 128 and alpha 256; diffusion LoRA uses rank 256 and alpha 256. Resolution is subsequently increased to \(1024\times1024\) with the same mixed-task dataset for another 3,000 steps, global batch size 256, and learning rate \(10^{-5}\). These are training settings, not per-image inference costs.

Generation and reconstruction use 30M open-source English image-text pairs and 170M internal Chinese pairs. Stage 1 additionally uses 800K natural image-transformation triplets, constructed by clustering images with shared captions, filtering with SigLIP features, and generating instructions with multimodal models. Stage 3 includes public editing data, nonrigid changes extracted from video, subject customization, and multi-subject synthesis. Table 1 lists 3M public editing samples, 2M internal video samples, and 300K internal editing samples. The relationship between the customization breakdown in the prose and the table's aggregate entries is not fully transparent, so these entries are not summed into a purported deduplicated total.

Key Experimental Results

Main Results

All entries below are scores reported in the paper, with higher values preferred and no time or pixel units. DINO and CLIP-I measure visual similarity to reference images, whereas CLIP-T measures image-text alignment. Higher visual similarity does not automatically imply better instruction following.

Original table / benchmark	Metric	Query-Kontext	Comparison method	Comparison score
Table 4 / GenEval	Overall	0.88	BAGEL	0.88
Table 4 / GenEval	Position	0.85	BAGEL	0.78
Table 5 / GEdit-Bench-EN full set	G_O	7.66	Qwen-Image	7.56
Table 5 / GEdit-Bench-CN full set	G_O	7.65	Qwen-Image	7.52
Table 6 / DreamBooth, single subject	DINO / CLIP-I / CLIP-T	0.786 / 0.858 / 0.307	BAGEL	0.777 / 0.851 / 0.307
Table 7 / DreamBench, multiple subjects	DINO / CLIP-I / CLIP-T	0.532 / 0.731 / 0.336	UNO-FLUX	0.542 / 0.733 / 0.322

GenEval uses Chinese prompts rewritten with DeepSeek, and Table 4 marks methods using prompt rewriting. The result should not be treated as an unassisted, no-rewriting evaluation. GEdit-Bench uses GPT-4 to evaluate semantic consistency, G_SC, and perceptual quality, G_PQ. The verbal definition in Table 5 translates into:

\[ G_O=\frac{1}{M}\sum_{m=1}^{M}\sqrt{G_{SC,m}G_{PQ,m}}. \]

Here \(M\) is the number of evaluated samples: the geometric mean is computed per sample and then averaged, rather than calculated from the two aggregate means in the table. English G_SC / G_PQ scores are 8.36 / 7.37 for Query-Kontext and 8.00 / 7.86 for Qwen-Image; Chinese scores are 8.39 / 7.35 and 7.82 / 7.79, respectively. Higher overall scores accompany stronger semantic consistency, not superior perceptual quality.

Ablation Study

Table 9 compares conditioning interfaces with the same diffusion backbone using Stage 1 GenEval Overall scores at different training lengths. The table below selects the clearly readable 20K, 60K, and 100K columns and excludes an irregularly formatted intermediate cell in the cache.

Conditioning interface	Text / query tokens	20K steps	60K steps	100K steps
Query-Kontext	512 / 128	0.661	0.859	0.882
Text-only MLP	512 / 0	0.587	0.757	0.814
Query-only Q-Former style	0 / 640	0.359	0.689	0.751

At 100K steps, the combined interface exceeds text-only conditioning by 0.068 and query-only conditioning by 0.131, supporting complementary roles for explicit text and generation queries. However, the text-only setup has a different total token count, so the entire difference cannot be attributed to one architectural variable. Table 9 extends this ablation to 100K steps, whereas Table 3 specifies a default Stage 1 length of 72K steps; these are different training-length settings.

Table 8 compares source and reference roles on DreamBooth. Source mode achieves DINO / CLIP-I / CLIP-T of 0.865 / 0.914 / 0.289; reference mode achieves 0.786 / 0.858 / 0.307. Source mode stays closer to the input image, while reference mode better follows the new text. This is a comparison of role configurations, not a pure RoPE-versus-no-RoPE ablation.

Key Findings

• The combined interface helps both early and late in training, rather than merely adjusting the final score. Figure 4 does not supply a directly verifiable speedup factor, so no training acceleration multiplier is inferred from the plot.
• Multi-subject CLIP-T of 0.336 exceeds UNO-FLUX's 0.322, but DINO and CLIP-I are slightly lower. The unified model does not dominate every metric.
• Figure 7 illustrates unseen concepts, style transfer, and zero-shot outpainting without corresponding systematic quantitative evaluation. These are capability examples, not guarantees of generalization.

Highlights & Insights

• Queries need not replace text. Retaining text as a stable condition and using a small query set for multimodal output information reduces pressure on a single compressed interface; Table 9 provides stronger evidence than simply increasing query count.
• The training order is part of the interface design: make VLM outputs generation-oriented, adapt the strong diffusion backbone, and then add fine detail. Transferring this recipe requires checking which parameters are frozen, not merely copying three stage names.
• Source and reference images impose different fidelity requirements. Encoding their roles through position can inform systems that support both local editing and subject transfer, but task data must still teach the model what those roles mean.

Limitations & Future Work

• The authors acknowledge weaker perceptual quality and leave reinforcement learning or supervised fine-tuning specifically targeting image quality and photorealism to future work. This does not mean Stage 3 instruction tuning is absent; the missing step is dedicated quality optimization.
• The internal diffusion backbone and 170M internal Chinese image-text pairs restrict independent reproduction. The authors report 192 NVIDIA H100 GPUs with 80GB each and claim about 10% of from-scratch training compute, but provide insufficiently matched GPU-hour and baseline budgets to derive an absolute cost.
• Support for at most two reference images does not establish scalable multi-subject composition. Lightweight-connector alignment to large frozen diffusion models also remains unresolved, with connector scaling left for future study.
• The cache includes the full main text and references but not the repeatedly cited appendix. Some dataset labels, aggregation conventions, and training details remain unclear. Confidence intervals, complete inference settings, and independent tests of retained understanding ability are also missing, limiting reproduction and causal interpretation.

Related Work & Insights

• vs MetaQuery: Both use queries to connect understanding and generation. Query-Kontext first trains VLM LoRA with a lightweight head, then replaces the backbone, while retaining text and adding low-level image conditions. Its contribution is the combination of bridging, curriculum, and information pathways, not merely an extra query module.
• vs BAGEL: BAGEL integrates understanding and generation more tightly; Query-Kontext retains separately pretrained VLM and diffusion components. Their shared GenEval score of 0.88 indicates parity on that benchmark, not equal understanding or interleaved-output capabilities.
• vs FLUX.1 Kontext / Step1X-Edit: These methods also emphasize contextual image conditions or VLM instruction understanding. Query-Kontext adds specifically trained query representations alongside text and low-level conditions. The transferable distinction is among explicit instructions, target visual conditions, and original-image detail, rather than model naming.

Rating

• Novelty: 4/5. Query bridging is established, but the three-stage curriculum and role-aware low-level conditions form a concrete approach.
• Experimental Thoroughness: 4/5. Four task families and interface comparisons are covered, but complete quantitative low-level ablations and open reproduction resources are missing.
• Writing Quality: 3/5. The main pipeline is clear, while data conventions, the unavailable appendix, and some training details complicate verification.
• Value: 4/5. The paper offers a practical division of training responsibilities when reusing VLMs and strong diffusion models, although resource requirements remain high.