Mask-Guided Attention Regulation for Anatomically Consistent Counterfactual CXR Synthesis

Zichun Zhang; Weizhi Nie; Honglin Guo; Yuting Su

cs.CVMar 4, 2026

Mask-Guided Attention Regulation for Anatomically Consistent Counterfactual CXR Synthesis

Authors: Zichun Zhang, Weizhi Nie, Honglin Guo, Yuting Su

Organizations: Tianjin University

Abstract

Counterfactual generation for chest X-rays (CXR) aims to simulate plausible pathological changes while preserving patient-specific anatomy. However, diffusion-based editing methods often suffer from structural drift, where stable anatomical semantics propagate globally through attention and distort non-target regions, and unstable pathology expression, since subtle and localized lesions induce weak and noisy conditioning signals. We present an inference-time attention regulation framework for reliable counterfactual CXR synthesis. An anatomy-aware attention regularization module gates self-attention and anatomy-token cross-attention with organ masks, confining structural interactions to anatomical ROIs and reducing unintended distortions. A pathology-guided module enhances pathology-token cross-attention within target lung regions during early denoising and performs lightweight latent corrections driven by an attention-concentration energy, enabling controllable lesion localization and extent. Extensive evaluations on CXR datasets show improved anatomical consistency and more precise, controllable pathological edits compared with standard diffusion editing, supporting localized counterfactual analysis and data augmentation for downstream tasks.

Figures & tables

Figure 1: Overview of our inference-time attention regulation framework for counterfactual CXR generation. (a) The input image is encoded by VAE [ 15 ] , noised, and denoised by a conditional diffusion model to generate a counterfactual image. (b) Anatomy-aware self-attention gating with $M_{\text{anat}}$ preserves structural consistency, while (c) pathology-guided cross-attention reweighting with a mask-derived prior $\Omega$ localizes pathological edits.

Figure 2: Qualitative comparison with state-of-the-art counterfactual CXR generation methods. Given the target prompt and the original image, we compare our results with SD-inpainting, PIE, BiomedJourney, and ProgEmu on representative cases.

Method	Conf $\uparrow$	CLIP-I $\uparrow$	LPIPS [ 26 ] $\downarrow$	FID [ 12 ] $\downarrow$
SD-inpainting [ 19 ]	0.662	0.827	0.23	37.6
Biomedjourney [ 10 ]	0.693	0.868	0.19	30.8
ProgEmu [ 17 ]	0.701	0.866	0.17	28.4
PIE [ 16 ]	0.691	0.842	0.18	29.7
Ours	0.709	0.870	0.18	29.0

Table 1: Performance comparison with state-of-the-art methods on the counterfactual CXR generation task.

Variant	Conf $\uparrow$	CLIP-I $\uparrow$	SSIM $\uparrow$
Full (Ours)	0.71	0.87	0.80
w/o Anatomy self-attn regularization (Eq. 5 )	0.69	0.85	0.76
w/o Pathology cross-attn regulation (Eq. 6 )	0.66	0.82	0.80
w/o Latent correction (Eq. 9 )	0.70	0.86	0.79

Table 2: Ablation study of inference-time regulation modules on the counterfactual CXR generation task.

Medical Image Synthesis

Medical image synthesis creates, translates, reconstructs, and edits clinical images to fill missing data and support analysis.

See all 298 papers

Positional Segmentor-Guided Counterfactual Fine-Tuning for Spatially Localized Image Synthesis
Mar 22, 2026Tian Xia, Matthew Sinclair, Andreas Schuh +8Counterfactual Image Generation Image Generation
Department of Computing, Imperial College London, UK · HeartFlow, Inc., Mountain View, CA, USA
Generating Chest X-Ray Counterfactuals by Specialising Foundation Image Models
Sep 21, 2026Xiaodan Xing, Rajat R. Rasal, Julia A. Meister +3Medical Imaging Foundation Models Causal Counterfactual Generation
GSK · Imperial College London
Spatially Gated Diffusion for Localized Counterfactual Chest Radiograph Editing
Sep 30, 2026Kamran Ullah Afaq, Basit RazaMedical Imaging Text-Guided Image Editing
COMSATS University Islamabad, Pakistan

Mar 22, 2026J/K move · Enter open · S save

Positional Segmentor-Guided Counterfactual Fine-Tuning for Spatially Localized Image Synthesis

Tian Xia, Matthew Sinclair, Andreas Schuh, Fabio De Sousa Ribeiro+7

Department of Computing, Imperial College London, UK · HeartFlow, Inc., Mountain View, CA, USA

Counterfactual Image Generation Image Generation Counterfactual Data Augmentation Medical Image Synthesis Conditional Image Generation

Open paper AskPDF arXiv

Counterfactual image generation enables controlled data augmentation, bias mitigation, and disease modeling. However, existing methods guided by external classifiers or regressors are limited to subject-level factors (e.g., age) and fail to produce localized structural changes, often resulting in global artifacts. Pixel-level guidance using segmentation masks has been explored, but requires user-defined counterfactual masks, which are tedious and impractical. Segmentor-guided Counterfactual Fine-Tuning (Seg-CFT) addressed this by using segmentation-derived measurements to supervise structure-specific variables, yet it remains restricted to global interventions. We propose Positional Seg-CFT, which subdivides each structure into regional segments and derives independent measurements per region, enabling spatially localized and anatomically coherent counterfactuals. Experiments on coronary CT angiography show that Pos-Seg-CFT generates realistic, region-specific modifications, providing finer spatial control for modeling disease progression.

Abstract

Medical Image Synthesis

Explore similar work

Positional Segmentor-Guided Counterfactual Fine-Tuning for Spatially Localized Image Synthesis

Generating Chest X-Ray Counterfactuals by Specialising Foundation Image Models

Spatially Gated Diffusion for Localized Counterfactual Chest Radiograph Editing