Apple explores how to keep AI from losing the thread when editing images

On June 1, 2026, a paper with Apple's participation proposed training image-editing models for a chain of decisions, not just an isolated command. Editing an image with natural language rarely ends with a single instruction: a person may ask to remove an object, then add text, and finally change the material of a specific area. The manuscript is signed by Apple researchers together with colleagues from the University of California, Los Angeles, the University of Texas at Austin, and Lambda, with Jiahui Huang and Yasi Zhang as equal first authors.

The method, called MT-EditFlow, applies reinforcement learning to editing models based on flow matching — the family of generative models that builds the image by transforming noise into pixels through a learned continuous flow, and the foundation of the FLUX series the study uses as its test bench. Its goal is not only that the image responds to the latest instruction, but that it preserves the content the user did not ask to change across several turns. That difference matters: an editor can succeed at recoloring a wall and, in doing so, alter a person, an object, or the style the user wanted to keep.

The problem of editing on top of your own output

Many editors are trained on image–instruction pairs: they receive an original image, apply one change, and get evaluated once. In a real conversation, the second change is made on top of the first one's output. The paper calls this an "all or nothing" demand: a single deviation can ruin the whole sequence. Exposure bias appears too, because the model must work with its own imperfect outputs instead of the ideal examples seen during training.

The scale of the problem is in the study's own table. Without sequence training, FLUX.1-Kontext-dev follows the instruction 59.97% of the time on the first turn, 32.69% on the second, and only 16.61% on the third. That is, literally, losing the thread: each edit on top of the previous result multiplies the ways to fail. Consistency of the unedited content holds up better — from 95.32% to 90.22% between the first and third turns — but the ability to follow orders collapses.

MT-EditFlow changes the training signal to observe the full trajectory. It uses two rewards: instruction following, judged by an open vision-language model (Qwen3-VL-8B) that compares reference, edits, and instructions; and content consistency, measured with EdiVal-CC as semantic distance in the regions that were not supposed to change. Both signals are converted into group-relative advantages and fused with weights — the authors fix instruction following at 1 and find the best balance at 0.3 for consistency — and the resulting trajectory-level advantage is broadcast equally to all three turns of the chain, so that local decisions answer to global success.

The researchers also compared ways of scoring. A strict binary evaluation, which only grants success if every turn goes well, resembles the user's experience but produces extreme reward sparsity: almost no example passes it and the gradient turns into nothing or noise. A binary per-turn average improves somewhat. The winner is a fine-grained per-turn score, on a 1-to-5 scale, averaged: it gives partial credit and keeps a dense signal even when most turns fail, without rewarding the model for solving one order at the expense of the rest of the image. They also test two modes for the judge: asking it to reason step by step before scoring reduces its bias but increases the signal's variance; even so, the thinking mode ends up improving the final result.

Results over three turns

The work evaluates the method on EdiVal-Bench, a sequential editing suite. For training it builds its chains with EdiVal-Agent — the benchmark's pipeline, which generates object-centric, context-consistent instruction sequences — and GPT-4o: about 2,319 reference images sourced from OpenImages, via the Pico-Banana-400K dataset, yielding 6,957 turn-level steps for reinforcement fine-tuning, always in three-turn sequences and spread across nine editing categories and twelve object categories.

In the main table, FLUX.1-Kontext-dev rises from 38.71 to 45.56 on the third-turn overall score with the MT-EditFlow-NFT variant: a 6.85-point difference, which is only 1.59 on the first turn and 5.13 on the second — the gain grows exactly where the problem worsens. The overall score combines both metrics — obedience and consistency — at each turn, and consistency rises too, if slightly: from 90.22% to 90.62% on the third turn. The method does not buy obedience at the price of fidelity. The authors place that result above Qwen-Image-Edit, which reaches 41.93 in the same column, and also report a 2.90-point improvement when applying the method to a second model, FLUX.2-klein-base-9B. The manuscript tests two reinforcement variants: one extends Flow-GRPO to multi-turn trajectories, and the other — the NFT behind the best results — builds on DiffusionNFT, which contrasts positive and negative samples without explicit likelihood estimation, and trains faster: 21.4 hours versus 34.3 on B200 GPUs for about 700 effective steps, per the paper's own indicative figures.

It is a benchmark result, not a promise that any image will be edited reliably. The instruction-following evaluation uses Qwen3-VL-8B as a visual reward model, while content consistency is measured with EdiVal-CC. The authors themselves note that vision-language models can fail at spatial reasoning, at detecting small details, or at judging artifacts and aesthetics, and that their hallucinations inject noise into the reinforcement signal even when prompted to reason step by step. They add two more limits: using large evaluators makes training slower and more expensive, and the whole evaluation happens on three-turn sequences — what happens in longer conversations remains unexplored.

What counts as a good edit

A revealing decision in the work is leaving pure visual quality out of the reward. The authors argue that an aesthetically polished image can still be a bad edit if it changes the identity, the style, or the elements that were supposed to remain. That is why they prioritize following the instruction and preserving what was not edited. The asymmetry in the table is itself a finding: what degrades across turns is not so much the fidelity of what goes untouched as the ability to execute what is actually requested. It is a reasonable choice for studying interactive editing, but it does not cover everything a user may value in a final image.

The most interesting advance is not a specific visual effect but treating the editing conversation as a sequence. Image systems are increasingly used as iterative collaborators: they try, receive corrections, and try again. For that interaction to be useful, executing isolated orders is not enough; the model must keep the thread of what it already changed and of what was deliberately left untouched. MT-EditFlow offers an experimental way to train that short-range visual memory — and its tables leave the reader with a yardstick that will outlive this paper: when facing any instruction-based image editor, ask how it performs on the third turn, not the first.