MedSteer: Counterfactual Endoscopic Synthesis via Training-Free Activation Steering

MedSteer is a training-free activation-steering framework that generates structurally preserved counterfactual endoscopic images by manipulating cross-attention activations in diffusion transformers, outperforming existing methods in concept editing and downstream medical detection tasks.

The Gist

Imagine you are a chef trying to teach a robot how to spot a specific ingredient in a soup. The robot is great at tasting, but it keeps getting confused by the color of the bowl or the steam rising from it. To fix this, you want to show the robot two bowls of soup: one with the ingredient and one without, but everything else must be exactly the same.

The problem is, current AI tools (called "Diffusion Models") are like a chaotic artist. If you ask them to paint a bowl with a mushroom, and then ask them to paint a bowl without a mushroom, they don't just remove the mushroom. They repaint the whole bowl, change the steam, and maybe even switch the bowl to a different table. The "background" changes, so the robot can't learn what a mushroom actually looks like.

Other tools try to fix this by taking the first painting and trying to "erase" the mushroom. But this is like trying to edit a photo by smudging the paint; it leaves messy artifacts and the background gets distorted.

Enter MedSteer.

MedSteer is a new, "training-free" tool that acts like a precise surgical scalpel for AI images. Here is how it works, using simple metaphors:

1. The "Pathology Vector" (The Recipe for Change)

Imagine the AI model has a giant library of "thoughts" (called vectors) inside its brain. When the AI thinks about "Polyps" (a type of growth in the gut), it uses a specific set of thoughts. When it thinks about "Normal Tissue," it uses a different set.

MedSteer doesn't need to retrain the AI. Instead, it does a quick "taste test" first. It asks the AI to imagine a "Polyp" and then a "Normal" version, and it measures the exact difference between those two thoughts. It creates a "Pathology Vector"—think of this as a magic instruction card that says, "To turn a Polyp into Normal Tissue, you only need to change these specific ingredients, leave everything else alone."

2. The "Steering" (The GPS for the Image)

Now, when the AI starts painting a new image from scratch (using a random noise seed, like static on an old TV), MedSteer acts as a GPS navigator.

• Without MedSteer: The AI wanders randomly. If you ask for a "Normal" image, it might wander into a different room entirely.
• With MedSteer: The AI starts the journey. As it paints, MedSteer checks its "thoughts" at every single step. If the AI starts thinking about "Polyps," MedSteer gently nudges it back toward "Normal" using that magic instruction card.

Crucially, because MedSteer is steering the AI while it paints (rather than trying to fix a finished painting), the background, the lighting, and the shape of the organ remain perfectly identical. It's like two twins walking the exact same path, but one is wearing a hat and the other isn't. The path (the anatomy) is identical; only the hat (the disease) changes.

3. The "Cosine-Similarity Gate" (The Smart Switch)

One of the coolest parts is how MedSteer knows where to apply the change. Imagine the image is made of thousands of tiny puzzle pieces (tokens).

MedSteer asks each puzzle piece: "Are you part of the Polyp?"

• If a piece is part of the background (like the wall of the intestine), the answer is "No." MedSteer leaves it alone.
• If a piece is part of the Polyp, the answer is "Yes." MedSteer applies the "Normal" instruction to that specific piece.

This is like a smart switch that only turns off the lights in the room where the party is happening, leaving the rest of the house dark and undisturbed. This gives doctors a clear visual map of exactly where the AI is making changes, which is a huge deal for trust in medical AI.

Why Does This Matter?

The paper tested this on real medical data (endoscopy images of the gut).

• Better Learning: When they used MedSteer to create "fake" training data to teach a computer how to spot polyps, the computer got much smarter (97.5% accuracy) compared to other methods.
• Removing Dyes: They even used it to "remove" blue dye from images (used in surgery) without changing the shape of the tissue underneath, something other tools failed to do.
• No Retraining: The best part? They didn't have to teach the AI anything new. They just used the AI's existing brain and gave it a better steering wheel.

In summary: MedSteer is like a precision editor that can swap a disease for healthy tissue in a medical image without blurring the background or changing the shape of the organ. It helps doctors train better AI detectors and understand exactly how the AI is making its decisions.

Technical Summary

Here is a detailed technical summary of the paper "MedSteer: Counterfactual Endoscopic Synthesis via Training-Free Activation Steering."

1. Problem Statement

In medical imaging, specifically endoscopy, training robust pathology detectors is hindered by confounding anatomical features. Detectors often learn to associate disease with specific background structures rather than the pathology itself.

• The Goal: Generate "counterfactual" image pairs (e.g., a diseased image and its healthy counterpart) where the anatomy, texture, and background are identical, and only the specific pathology concept differs.
• Limitations of Existing Methods:
  • Re-prompting: Generating two images from scratch using different text prompts (e.g., "polyp" vs. "normal") causes the model to reroll the entire generation trajectory, resulting in completely different anatomies and backgrounds.
  • Inversion-Based Editing: Methods that start with a source image and edit it using Diffusion Inversion (e.g., DDIM inversion) suffer from reconstruction errors. These errors cause "structural drift," meaning the background and non-targeted anatomy change slightly, failing to provide the strict structural preservation required for causal medical analysis.
  • Training Requirements: Many existing editing methods require fine-tuning, mask annotations, or extensive training data, which are scarce in medical domains.

2. Methodology: MedSteer

MedSteer is a training-free, inversion-free framework that generates counterfactual pairs from scratch while guaranteeing structural identity. It operates on a frozen Diffusion Transformer (DiT) backbone (specifically PixArt-α).

Core Components:

1. Offline Pathology Vector Estimation:

   • Instead of editing an image, MedSteer first identifies a "pathology vector" in the model's latent space.
   • It uses contrastive prompt pairs (e.g., "dyed lifted polyp" vs. "polyp") with varied random seeds and context phrasing.
   • It collects Cross-Attention (CA) features from the frozen model for both prompts.
   • It computes the mean difference between the positive and negative prompt features, normalizes it, and derives a unit vector ($v_{l,t}$) representing the semantic difference (the pathology) for each layer ($l$) and timestep ($t$).

2. Inference-Time Steering (Spatially Selective Pathology Steering - SSPS):

   • Shared Trajectory: Both the "source" (pathological) and "target" (healthy) images are generated from the same noise seed and the same prompt (e.g., "An endoscopic image of a dyed lifted polyp").
   • Intervention: During the generation of the target image, the system intervenes in the Cross-Attention layers.
   • Cosine-Similarity Gate (CSG): To avoid over-suppressing non-target features (like anatomy), the method calculates a per-token score based on the cosine similarity between the current activation and the pathology vector.
   • Update Rule: Only tokens positively aligned with the pathology vector are modified. The update subtracts the pathology-aligned component scaled by a steering strength ($\alpha$):
     $$h'_{l,t} = h_{l,t} - \alpha \cdot \sigma_{l,t} \cdot v_{l,t}$$
   • Result: The "steered" image retains the exact same noise trajectory and structural foundation as the "unsteered" image, differing only in the steered concept.

3. Key Contributions

1. Training-Free Activation Steering: Introduces a mechanism to steer clinical concepts by removing concept-aligned components from cross-attention tokens without fine-tuning the model or requiring source images/masks.
2. Built-in Spatial Interpretability: The per-token cosine similarity gate ($\sigma_{l,t}$) acts as a spatial map, revealing exactly where and when the model modifies the image. This provides interpretability absent in inversion-based methods.
3. Inversion-Free Counterfactual Generation: Achieves true structural preservation by generating both images from the same noise seed, eliminating the reconstruction errors inherent in DDIM inversion.
4. Disentanglement of Entangled Attributes: Successfully separates co-occurring attributes (e.g., polyp morphology vs. dye staining) even when text prompts cannot explicitly describe the separation.

4. Experimental Results

The method was evaluated on Kvasir v3 and HyperKvasir datasets across three experiments:

• Downstream Polyp Detection:

  • Augmenting training data with MedSteer-generated counterfactuals significantly improved the performance of ViT-based polyp detectors.
  • Result: MedSteer achieved an AUC of 0.9755, outperforming quantity-matched re-prompting (0.9083) and inversion-based baselines (PnP: 0.9518, h-Edit: 0.9312). This confirms that the structural consistency of the counterfactuals drives better generalization.

• Counterfactual Generation Quality:

  • Evaluated on three concept pairs (Polyp $\leftrightarrow$ Normal, Colitis $\leftrightarrow$ Normal, Esophagitis $\leftrightarrow$ Normal).
  • Concept Flip Rate: MedSteer achieved high flip rates (0.800, 0.925, and 0.950), significantly outperforming PnP and h-Edit.
  • Structural Preservation: MedSteer showed superior background preservation (higher Bg-SSIM, Bg-PSNR, and lower Bg-LPIPS) compared to inversion methods, which suffered from drift.

• Dye Disentanglement:

  • Task: Remove dye staining from "dyed lifted polyps" while keeping the polyp shape intact.
  • Result: MedSteer achieved a Dye Detection Rate (DDR) of 0.250 (meaning 75% of dye was removed), vastly outperforming h-Edit (0.900) and PnP (0.800).

• Ablation Studies:

  • Identified that layers 8–16 are the "semantic formation zone" for steering.
  • Optimal steering strength ($\alpha$) was found to be 2.5.
  • The pathology vector stabilizes after approximately 50 seeds.

5. Significance

MedSteer addresses a critical bottleneck in medical AI: the lack of high-quality, causally valid training data. By enabling the generation of perfectly matched pathological and healthy pairs without structural drift, it allows models to learn disease-specific features rather than anatomical confounders.

• Clinical Impact: It facilitates the creation of robust pathology detectors that generalize better to out-of-distribution data.
• Methodological Impact: It demonstrates that "training-free" activation steering in frozen Diffusion Transformers is a viable, superior alternative to inversion-based editing for medical image synthesis, offering both high fidelity and built-in interpretability.
• Future Directions: The authors plan to extend this to 3D volumetric data, video endoscopy, and cross-institutional deployments.