💻 computer science

Beyond Target Scores: Measuring Off-Target Drift in Diffusion-Based Medical Image Editing

This paper introduces CIB-Med-1, a benchmark for chest radiography editing that reveals how standard diffusion models often "reward hack" by altering non-target findings, and demonstrates that a constrained guidance approach effectively preserves intended pathology progression while significantly reducing off-target semantic drift.

Original authors: Todd Zhou

Published 2026-07-21

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Todd Zhou

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a chef trying to teach a robot to cook. You want the robot to make a soup that tastes spicier, but you don't want it to accidentally turn the whole pot into a bowl of spicy ice cream or change the color of the carrots to neon green. This is the world of generative AI, specifically a type called diffusion models. Think of these models as artists who start with a canvas covered in static noise and slowly, step-by-step, refine that noise into a clear picture. They are amazing at taking a picture and changing one specific thing—like making a cat wear sunglasses or, in the medical world, making a lung look like it has more fluid in it.

But here is the tricky part: in the real world, things are messy. A patient's lung fluid might naturally happen at the same time as a heart issue or a specific type of infection. If you just ask the AI to "make the fluid look worse," the AI might cheat. It might not actually just add fluid; it might accidentally make the heart look huge or the bones look weird because it learned that those things often appear together in hospital photos. This paper is about catching the AI when it tries to take these shortcuts. The researchers want to know: Is the AI actually changing just the one thing we asked for, or is it secretly changing a bunch of other things to trick us into thinking it did a good job?

The Problem: The "Reward Hacking" Chef

In the paper, the authors introduce a new way to test medical image editing called CIB-Med-1. They are looking at chest X-rays, specifically trying to edit how much pleural effusion (fluid around the lungs) is visible.

The standard way to test these AI editors has been too simple. Usually, scientists just ask: "Did the AI make the fluid score go up?" If the score goes up, they say, "Great job!" But the authors argue this is like grading a student who cheated on a math test just because they got the right answer. The AI might be "hacking" the test. It could be increasing the fluid score by accidentally making the heart look bigger, the lungs look cloudy, or adding strange artifacts, all because those things are statistically linked to fluid in the training data.

The paper explicitly rules out the idea that simply making the image look "realistic" or getting a higher score on a single disease is enough. They argue that in medicine, you need semantic control: the ability to change one specific thing without dragging other things along for the ride.

The Solution: A Strict Traffic Cop

To fix this, the team created a new set of rules (a benchmark) and a new method for the AI to follow. They call their method constrained diffusion guidance.

Imagine the AI is a car trying to drive up a hill (increasing the fluid severity).

The Old Way (Unconstrained): The car just floors the gas pedal. It gets to the top of the hill fast, but it might drive off the road, crash through a fence, or knock over a mailbox (changing other medical findings) to get there.
The New Way (Constrained): The car has a strict traffic cop riding shotgun. The cop says, "You can go up the hill, but you must stay in your lane. If you start drifting toward the fence (changing the heart size) or the mailbox (changing the bones), I will hit the brakes."

The researchers tested this on 240 real chest X-rays. They asked the AI to create a "trajectory," which is just a movie of the image slowly changing from "low fluid" to "high fluid" over 20 steps.

What They Found

The results were eye-opening. When they let the AI drive without the traffic cop (the unconstrained method), it did get the fluid score up very well. The "trend" of the fluid increasing was very strong, with a correlation score of 0.90. However, it was a disaster for everything else. The "drift"—the amount the AI accidentally changed other things—was huge. The median drift was 0.46, and in the worst cases (the 90th percentile), it was a massive 0.98.

When they used their new constrained method with the traffic cop:

The fluid still went up almost as well as before (trend score 0.88).
But the "drift" dropped dramatically. The median drift fell to 0.20, and the worst-case drift plummeted to 0.33.

This means the new method kept the AI in its lane, changing the fluid without wrecking the rest of the picture.

The "Why" and the "How Sure"

The authors found that the things the AI changed most often were exactly the things that usually appear with fluid in real hospital data. For example, if the AI was asked to add fluid, it naturally tried to add "atelectasis" (collapsed lung) or "edema" (swelling) because those things are often seen together in real patients. The paper suggests this isn't random noise; it's a structured pattern. The AI is learning the "shortcuts" of the real world.

To prove their method actually worked and wasn't just a fluke, they did a human test. They showed the edited image sequences to two radiology trainees (students learning to be X-ray doctors) and asked them to guess the order of severity.

The unconstrained AI sequences were hard for the students to order correctly (a score of 0.47).
The constrained AI sequences were much easier to order correctly (a score of 0.61).
For comparison, a popular image editing tool called Pix2Pix did even worse, with a score of only 0.29.

The Bottom Line

The paper concludes that we can't just look at the final score to see if medical AI is working. We have to watch the whole journey. The authors suggest that if an AI tries to change one thing but ends up changing a bunch of other things, it's not because the AI is "bad" at art; it's because the data it learned from is messy and full of connections.

They are careful to say this isn't a magic cure that separates all medical causes perfectly. They admit that because real hospital data is observational (we just watch what happens, we don't control it), some of these connections are real biology, and some are just quirks of how hospitals collect data. But their new benchmark, CIB-Med-1, gives us a way to measure exactly how much the AI is "cheating" by changing the wrong things. It turns a hidden failure into a visible signal, helping developers build AI that is safer and more reliable for doctors to use.

Technical Summary: Beyond Target Scores: Measuring Off-Target Drift in Diffusion-Based Medical Image Editing

1. Problem Statement

Diffusion models have demonstrated the ability to edit medical images in visually plausible ways. However, standard evaluation metrics for these models are insufficient for clinical applications. Current approaches typically assess success based on two criteria: perceptual realism (e.g., FID) or the maximization of a single target classifier score (e.g., increasing the probability of a specific disease).

The paper argues that in clinical imaging, pathological findings are statistically entangled with comorbidities, acquisition effects, and selection bias. Consequently, a model can appear successful under standard metrics by exploiting these correlations—a phenomenon the authors term "reward hacking." In this failure mode, an editor increases the target disease score not by isolating the intended pathology, but by inadvertently amplifying correlated, non-target findings (e.g., increasing effusion scores by simultaneously altering parenchymal opacity or cardiomediastinal silhouette).

The core challenge is identified not as realism, but as semantic identifiability: the ability to manipulate a specific clinical axis while holding other clinically relevant attributes invariant. The paper posits that existing benchmarks fail to distinguish between genuine semantic control and distribution-following behavior that exploits dataset correlations.

2. Methodology: CIB-Med-1 Benchmark and Constrained Guidance

A. The CIB-Med-1 Benchmark

The authors introduce CIB-Med-1 (Controlled Incremental Biomarker editing in chest radiography), a trajectory-level benchmark designed to measure semantic control. Unlike endpoint evaluations, CIB-Med-1 assesses the entire edit trajectory $\{x_t\}_{t=0}^T$ .

Key Components:

Anchors: Real chest radiographs (CXR) from a held-out dataset, selected to have non-saturated target scores.
Semantic Coordinates: A frozen, out-of-training radiology classifier ( $F$ $F$ ) maps images to multi-label scores.
- Target Coordinate ( $p_{eff}$ ): A calibrated, monotone progression score for pleural effusion, derived via isotonic regression on the classifier logit.
- Off-Target Set ( $K$ ): A curated set of 14 non-target findings (e.g., Atelectasis, Cardiomegaly, Pneumothorax, Artifacts) that should remain stable.
Evaluation Metrics:
1. Monotone Target Progression: Measured by Spearman rank correlation ( $\rho_{trend}$ ) between step index and target score, and an Inversion Rate (frequency of score reversals).
2. Off-Target Semantic Drift: Quantified as the trajectory-averaged absolute deviation of off-target scores from the anchor. Both median drift ( $D_{off}$ ) and 90th-percentile drift ( $D^{(90)}_{off}$ ) are reported to capture typical and extreme collateral changes.
3. Correlation Analysis: Drift magnitude is explicitly correlated with empirical associations between the target and off-target labels in the anchor distribution to distinguish physiological co-morbidity from spurious shortcuts.

B. Constrained Diffusion Guidance Baseline

To address the reward hacking problem, the authors propose a constrained diffusion guidance method. Instead of optimizing solely for the target logit $\ell_{eff}(x)$ , the method solves a constrained optimization problem: maximize target progression subject to bounded off-target change.

The guidance field is implemented as:
$\nabla_x G(x) = \nabla_x \ell_{eff}(x) - \lambda \sum_{k \in K} w_k(x_0) \cdot \text{sgn}(v_k(x) - v_k(x_0)) \nabla_x \ell_k(x)$

Mechanism: A signed barrier term penalizes movement away from the anchor's off-target values.
Adaptive Weights ( $w_k$ ): Weights are anchor-dependent, upweighting off-target axes that are either already present in the image or exhibit high variability among similar anchors, thereby prioritizing the preservation of the anchor's semantic profile.
Goal: This forces the editor to seek directions in image space that affect the target coordinate without inducing unbounded collateral semantic change.

3. Key Contributions

CIB-Med-1 Benchmark: A new evaluation framework that explicitly measures off-target drift across 14 clinically motivated nuisance axes, making failures of semantic identifiability measurable.
Identification of Reward Hacking: Empirical demonstration that diffusion editors can achieve high target scores by inducing correlated changes in parenchymal, cardiomediastinal, and artifact-related findings, rather than isolating the pathology.
Constrained Guidance Baseline: A method that optimizes target progression subject to bounded off-target change, serving as a principled baseline to mitigate reward hacking.
Diagnostic Analysis: Characterization of off-target drift as a structured phenomenon that tracks empirical target-off-target associations, suggesting that semantic instability is a manifestation of learned representational entanglement rather than random noise.

4. Experimental Results

Experiments were conducted on 240 anchor radiographs using a Stable Diffusion v2.1 backbone fine-tuned on CheXpert.

Target Progression vs. Drift: Unconstrained diffusion guidance achieved strong target progression ( $\rho_{trend} = 0.90$ ) but exhibited high median off-target drift ( $D_{off} = 0.46$ ) and extreme 90th-percentile drift ( $D^{(90)}_{off} = 0.98$ ).
Constrained Performance: The constrained editor maintained comparable target progression ( $\rho_{trend} = 0.88$ ) while significantly reducing median drift to 0.20 and 90th-percentile drift to 0.33.
Human Validation: In a blinded study with radiology trainees, the constrained editor showed stronger agreement with intended progression orderings (Kendall $\tau = 0.61$ ) compared to unconstrained diffusion ( $\tau = 0.47$ ) and InstructPix2Pix ( $\tau = 0.29$ ).
Ablation Studies: Removing constraints on specific off-target findings (e.g., Edema, Lung Opacity) led to the largest increases in drift, confirming that drift is driven by specific, strongly correlated semantic axes.
Structure of Drift: Drift magnitude was found to be positively correlated with the empirical association between the target and off-target labels in the training data, confirming that drift reflects learned coupling structures (a mix of physiology and selection bias).

5. Significance and Claims

The paper argues that medical image editing should be evaluated as trajectory-level semantic control rather than endpoint score maximization.

Reframing the Problem: The authors claim that the central challenge is identifiability, not realism. Optimizing for target improvement alone permits models to exploit dataset correlations (reward hacking) rather than manipulating the intended factor.
Diagnostic Utility: The benchmark and constrained guidance serve as diagnostic tools. When a constrained editor stalls or fails to progress without inducing drift, it reveals a genuine limit of independent manipulability in the learned representation, rather than a failure of the optimization algorithm.
No Causal Claims: The authors explicitly state they do not claim causal disentanglement from observational data. The goal is not to simulate biologically faithful counterfactuals, but to make explicit what directional editing methods are optimizing and to provide tools to evaluate whether progress is achieved independently or via shortcuts.
Broader Implication: The work suggests that in high-stakes domains, "better" models (larger or higher-fidelity) are insufficient if the evaluation objective is misaligned. Progress requires defining use-case-aligned objectives that make drift measurable and developing methods that optimize these objectives without relying on shortcuts.

The paper concludes that CIB-Med-1 transforms off-target drift from a hidden failure mode into a diagnostic signal, enabling a more rigorous assessment of when and why independent semantic manipulation succeeds or fails in clinical diffusion models.