Segmentor-Guided Counterfactual Fine-Tuning

• The paper introduces Seg-CFT, a method that fine-tunes a HVAE using a frozen segmentation network to generate structure-specific counterfactual images.
• It employs a segmentation-derived measurement loss to enforce agreement between target scalar interventions and predicted anatomical changes.
• Positional Seg-CFT further partitions images into anatomical regions, enabling localized and anatomically coherent counterfactual adjustments.

Searching arXiv for the cited papers to ground the article in the current literature. Segmentor-Guided Counterfactual Fine-Tuning (Seg-CFT) is a method for fine-tuning a pretrained deep structural causal model so that an intervention on a structure-specific scalar variable yields an image whose segmentation-derived measurement matches the target intervention. In the formulation reported for medical imaging, the underlying generator is a hierarchical variational autoencoder (HVAE), and the supervisory signal is supplied by a frozen segmentation network rather than by an external regressor. This design is intended to preserve the simplicity of scalar-valued interventions while producing locally coherent counterfactuals for structures such as lung fields, calcified plaque, non-calcified plaque, or lumen area. A later extension, Positional Seg-CFT, subdivides each structure into anatomical regions and derives independent measurements per region, thereby enabling spatially localized and anatomically coherent counterfactuals without requiring user-defined counterfactual masks (Xia et al., 29 Sep 2025, Xia et al., 22 Mar 2026).

1. Research context and problem setting

Counterfactual image generation has been presented as a tool for controlled data augmentation, bias mitigation, disease modeling, and anatomically interpretable causal analysis. The motivating limitation of earlier counterfactual fine-tuning schemes is that classifier- or regressor-guided objectives are well suited to subject-level interventions such as age, but are insufficient for structure-specific variables such as area measurements. In the reported setting, a frozen regressor can satisfy the intervention objective by exploiting global, spurious cues, for example overall vessel intensity or brightness, rather than by altering the target anatomy itself (Xia et al., 29 Sep 2025).

Seg-CFT addresses that limitation by replacing the external regressor with a frozen segmentor. The segmentor outputs a soft segmentation map, and scalar measurements are recovered by summing pixel-wise class probabilities. The resulting loss directly penalizes disagreement between the target scalar and the segmentation-derived scalar measured on the generated counterfactual. In this way, the fine-tuning signal is tied to the spatial support of the target structure rather than to an arbitrary global predictor (Xia et al., 29 Sep 2025).

A second limitation remained in the original framework. Because each structure was represented by a single scalar measurement, the model could increase or decrease the total area of a structure but could not specify where within the image domain the change should occur. Positional Seg-CFT addresses that issue by partitioning the image into disjoint anatomical regions and imposing region-wise measurement constraints. This suggests a progression from subject-level control, to structure-specific control, and then to within-structure positional control (Xia et al., 22 Mar 2026).

2. Causal and generative formulation

The underlying deep structural causal model is implemented as an HVAE conditioned on low-dimensional parent attributes. Let $x$ denote the observed medical image and let $\{v_1,\ldots,v_{K-1}\}$ be low-dimensional attributes, among which are structure-specific variables such as calcified plaque area (CPA), non-calcified plaque area (NCPA), and lumen area (LA). In the DSCM/HVAE formulation, counterfactual inference proceeds by abduction, action, and prediction:

$$z \sim q_\phi(z \mid x, pa_x), \qquad u_k = f_k^{-1}(v_k; pa_k),$$

followed by the intervention

$do(v_i := c),$

and then prediction through the generator

$\widetilde{x} = g_\theta(z, \widetilde{pa}_x).$

The base HVAE training objective is the ELBO,

$$\mathcal{L}_{\mathrm{ELBO}}(\theta,\phi) = \mathbb{E}_{q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)} \big[ -\log p_\theta(\mathbf{x}\mid \mathbf{z},\mathbf{pa}_x) \big] + \mathrm{KL}\big(q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)\,\|\,p(\mathbf{z})\big),$$

with $q_\phi(\mathbf{z}\mid x,pa)$ the encoder, $p_\theta(x\mid z,pa)$ the decoder, and $p(z)$ a standard normal prior (Xia et al., 29 Sep 2025).

A central empirical observation is that naively sampling $\widetilde{x}$ after intervention may ignore the instructed change. Counterfactual fine-tuning remedies this by adding a term that encourages the generated image to encode the intended intervention. In regressor-based CFT, the fine-tuning signal is

$\{v_1,\ldots,v_{K-1}\}$0

with $\{v_1,\ldots,v_{K-1}\}$1 frozen. Seg-CFT replaces that regressor with a segmentor-driven measurement (Xia et al., 29 Sep 2025).

3. Segmentor-guided objective, architecture, and training

In Seg-CFT, a frozen segmentation network $\{v_1,\ldots,v_{K-1}\}$2 receives the counterfactual image $\{v_1,\ldots,v_{K-1}\}$3 and outputs a soft mask $\{v_1,\ldots,v_{K-1}\}$4. For a structure $\{v_1,\ldots,v_{K-1}\}$5 in a set $\{v_1,\ldots,v_{K-1}\}$6 of structures of interest, the predicted area is defined as

$\{v_1,\ldots,v_{K-1}\}$7

Counterfactual fine-tuning then enforces agreement between the predicted measurement and the target counterfactual scalar. The structure-wise loss is written as

$\{v_1,\ldots,v_{K-1}\}$8

where $\{v_1,\ldots,v_{K-1}\}$9 or $$z \sim q_\phi(z \mid x, pa_x), \qquad u_k = f_k^{-1}(v_k; pa_k),$$0. In the alternative presentation of the same method, the deviation is penalized by mean squared error,

$$z \sim q_\phi(z \mid x, pa_x), \qquad u_k = f_k^{-1}(v_k; pa_k),$$1

In practice, the CFT term is added to the original ELBO, weighted by a hyper-parameter $$z \sim q_\phi(z \mid x, pa_x), \qquad u_k = f_k^{-1}(v_k; pa_k),$$2,

$$z \sim q_\phi(z \mid x, pa_x), \qquad u_k = f_k^{-1}(v_k; pa_k),$$3

Optionally, if subject-level attributes such as age or sex are also to be preserved, a regressor loss can be added in parallel (Xia et al., 29 Sep 2025).

The reported architecture comprises a generator HVAE with encoder $$z \sim q_\phi(z \mid x, pa_x), \qquad u_k = f_k^{-1}(v_k; pa_k),$$4 and decoder $$z \sim q_\phi(z \mid x, pa_x), \qquad u_k = f_k^{-1}(v_k; pa_k),$$5, together with normalizing flows $$z \sim q_\phi(z \mid x, pa_x), \qquad u_k = f_k^{-1}(v_k; pa_k),$$6 for the low-dimensional variables so that there is an explicit, invertible mapping between noise $$z \sim q_\phi(z \mid x, pa_x), \qquad u_k = f_k^{-1}(v_k; pa_k),$$7 and attribute $$z \sim q_\phi(z \mid x, pa_x), \qquad u_k = f_k^{-1}(v_k; pa_k),$$8. The segmentor $$z \sim q_\phi(z \mid x, pa_x), \qquad u_k = f_k^{-1}(v_k; pa_k),$$9 is a U-Net-style convolutional network; in one description it is Dice-trained and remains frozen during fine-tuning, and in another it is pretrained on ground-truth masks of lumen and plaque, outputting a $do(v_i := c),$0-channel softmax map of size $do(v_i := c),$1, one channel per structure (Xia et al., 29 Sep 2025).

The algorithmic procedure is straightforward. A real image $do(v_i := c),$2 is sampled, its scalar parents $do(v_i := c),$3 are extracted, and a latent code is inferred via $do(v_i := c),$4. A counterfactual parent value $do(v_i := c),$5 is then specified, the image $do(v_i := c),$6 is generated, the predicted mask $do(v_i := c),$7 is computed, and the scalar $do(v_i := c),$8 is measured. The final loss is $do(v_i := c),$9, and $\widetilde{x} = g_\theta(z, \widetilde{pa}_x).$0 are updated by gradient descent. Reported fine-tuning hyper-parameters for the original paper are batch size $\widetilde{x} = g_\theta(z, \widetilde{pa}_x).$1, learning rate $\widetilde{x} = g_\theta(z, \widetilde{pa}_x).$2, number of CFT steps $\widetilde{x} = g_\theta(z, \widetilde{pa}_x).$3, and $\widetilde{x} = g_\theta(z, \widetilde{pa}_x).$4, tuned on the validation set. The same source notes that a moderate $\widetilde{x} = g_\theta(z, \widetilde{pa}_x).$5 such as $\widetilde{x} = g_\theta(z, \widetilde{pa}_x).$6–$\widetilde{x} = g_\theta(z, \widetilde{pa}_x).$7 prevents the model from overfitting local changes at the expense of realism, and that early stopping on a held-out counterfactual effectiveness metric helps prevent drift (Xia et al., 29 Sep 2025).

4. Positional Seg-CFT and spatial localization

The key limitation of the original Seg-CFT is that each structure $\widetilde{x} = g_\theta(z, \widetilde{pa}_x).$8 yields a single scalar measurement $\widetilde{x} = g_\theta(z, \widetilde{pa}_x).$9. The model can therefore increase or decrease a global structure-specific quantity, but it cannot choose in which part of the vessel or organ the change occurs. Positional Seg-CFT resolves this by subdividing the image domain into $$\mathcal{L}_{\mathrm{ELBO}}(\theta,\phi) = \mathbb{E}_{q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)} \big[ -\log p_\theta(\mathbf{x}\mid \mathbf{z},\mathbf{pa}_x) \big] + \mathrm{KL}\big(q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)\,\|\,p(\mathbf{z})\big),$$0 disjoint anatomical regions $$\mathcal{L}_{\mathrm{ELBO}}(\theta,\phi) = \mathbb{E}_{q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)} \big[ -\log p_\theta(\mathbf{x}\mid \mathbf{z},\mathbf{pa}_x) \big] + \mathrm{KL}\big(q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)\,\|\,p(\mathbf{z})\big),$$1, for example proximal, mid, and distal artery segments. These masks are binary $$\mathcal{L}_{\mathrm{ELBO}}(\theta,\phi) = \mathbb{E}_{q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)} \big[ -\log p_\theta(\mathbf{x}\mid \mathbf{z},\mathbf{pa}_x) \big] + \mathrm{KL}\big(q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)\,\|\,p(\mathbf{z})\big),$$2 maps that sum to $$\mathcal{L}_{\mathrm{ELBO}}(\theta,\phi) = \mathbb{E}_{q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)} \big[ -\log p_\theta(\mathbf{x}\mid \mathbf{z},\mathbf{pa}_x) \big] + \mathrm{KL}\big(q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)\,\|\,p(\mathbf{z})\big),$$3 across $$\mathcal{L}_{\mathrm{ELBO}}(\theta,\phi) = \mathbb{E}_{q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)} \big[ -\log p_\theta(\mathbf{x}\mid \mathbf{z},\mathbf{pa}_x) \big] + \mathrm{KL}\big(q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)\,\|\,p(\mathbf{z})\big),$$4 at each pixel (Xia et al., 22 Mar 2026).

Given a predicted segmentation map $$\mathcal{L}_{\mathrm{ELBO}}(\theta,\phi) = \mathbb{E}_{q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)} \big[ -\log p_\theta(\mathbf{x}\mid \mathbf{z},\mathbf{pa}_x) \big] + \mathrm{KL}\big(q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)\,\|\,p(\mathbf{z})\big),$$5 for structure $$\mathcal{L}_{\mathrm{ELBO}}(\theta,\phi) = \mathbb{E}_{q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)} \big[ -\log p_\theta(\mathbf{x}\mid \mathbf{z},\mathbf{pa}_x) \big] + \mathrm{KL}\big(q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)\,\|\,p(\mathbf{z})\big),$$6 on the generated image $$\mathcal{L}_{\mathrm{ELBO}}(\theta,\phi) = \mathbb{E}_{q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)} \big[ -\log p_\theta(\mathbf{x}\mid \mathbf{z},\mathbf{pa}_x) \big] + \mathrm{KL}\big(q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)\,\|\,p(\mathbf{z})\big),$$7, the region-specific measurement is

$$\mathcal{L}_{\mathrm{ELBO}}(\theta,\phi) = \mathbb{E}_{q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)} \big[ -\log p_\theta(\mathbf{x}\mid \mathbf{z},\mathbf{pa}_x) \big] + \mathrm{KL}\big(q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)\,\|\,p(\mathbf{z})\big),$$8

Instead of a single scalar per structure, one obtains $$\mathcal{L}_{\mathrm{ELBO}}(\theta,\phi) = \mathbb{E}_{q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)} \big[ -\log p_\theta(\mathbf{x}\mid \mathbf{z},\mathbf{pa}_x) \big] + \mathrm{KL}\big(q_\phi(\mathbf{z}\mid \mathbf{x},\mathbf{pa}_x)\,\|\,p(\mathbf{z})\big),$$9 independent scalar measurements, one per region. For an intervention of the form $q_\phi(\mathbf{z}\mid x,pa)$0, the positional loss becomes

$q_\phi(\mathbf{z}\mid x,pa)$1

Here $q_\phi(\mathbf{z}\mid x,pa)$2 is the target counterfactual area for structure $q_\phi(\mathbf{z}\mid x,pa)$3 in region $q_\phi(\mathbf{z}\mid x,pa)$4, while all other regions’ targets are kept at their factual values. By driving each region’s predicted area toward its target, the generator is encouraged to confine changes to the specified anatomical segment (Xia et al., 22 Mar 2026).

A notable feature of the extension is that no change is required to the HVAE or to the segmentor’s weights. The modification is confined to replacing the global-measurement loss by the region-wise loss, while at training time interventions are sampled independently per region. The overall optimization still jointly updates $q_\phi(\mathbf{z}\mid x,pa)$5 via gradient descent on $q_\phi(\mathbf{z}\mid x,pa)$6 plus the original HVAE term (Xia et al., 22 Mar 2026).

5. Empirical evidence

The original Seg-CFT paper reports experiments on chest radiographs and coronary CT angiography (CCTA). For PadChest, the dataset split is $q_\phi(\mathbf{z}\mid x,pa)$7 train, $q_\phi(\mathbf{z}\mid x,pa)$8 validation, and $q_\phi(\mathbf{z}\mid x,pa)$9 test, with images resized to $p_\theta(x\mid z,pa)$0. The target structures are left lung area (LLA), right lung area (RLA), and heart area (HA), and segmentation masks are obtained from TorchXRayVision pre-trained U-Nets. For CCTA straightened curvilinear planar reformations, the split is $p_\theta(x\mid z,pa)$1 train, $p_\theta(x\mid z,pa)$2 validation, and $p_\theta(x\mid z,pa)$3 test at $p_\theta(x\mid z,pa)$4 resolution, with structures NCPA, CPA, and LA, and masks generated by rasterizing segmented 3D meshes. Pretraining uses standard ELBO training for the HVAE for $p_\theta(x\mid z,pa)$5k steps with Adam learning rate $p_\theta(x\mid z,pa)$6; segmentors are trained for $p_\theta(x\mid z,pa)$7 epochs with Adam learning rate $p_\theta(x\mid z,pa)$8 and Dice loss; the Reg-CFT baseline uses a ResNet-18 regressor trained with MSE area regression for $p_\theta(x\mid z,pa)$9 epochs at learning rate $p(z)$0 (Xia et al., 29 Sep 2025).

Effectiveness in the original Seg-CFT experiments is measured by the deviation between the intervened parent $p(z)$1 and the predicted parent $p(z)$2 on the generated counterfactual. For chest radiographs, the paper reports mean absolute percentage error. Base DSCM obtains $p(z)$3 for LLA, $p(z)$4 for RLA, and $p(z)$5 for HA; Reg-CFT reduces these to $p(z)$6, $p(z)$7, and $p(z)$8; Seg-CFT reports $p(z)$9, $\widetilde{x}$0, and $\widetilde{x}$1, respectively. For CCTA plaque and lumen measurements, the reported MAE in mm$\widetilde{x}$2 is $\widetilde{x}$3, $\widetilde{x}$4, and $\widetilde{x}$5 for Base DSCM on NCPA, CPA, and LA; $\widetilde{x}$6, $\widetilde{x}$7, and $\widetilde{x}$8 for Reg-CFT; and $\widetilde{x}$9, $\{v_1,\ldots,v_{K-1}\}$00, and $\{v_1,\ldots,v_{K-1}\}$01 for Seg-CFT (Xia et al., 29 Sep 2025).

The positional extension is evaluated on an internal CCTA dataset of straightened curvilinear planar reformations of the left anterior descending artery, at $\{v_1,\ldots,v_{K-1}\}$02 mm resolution and cropped to $\{v_1,\ldots,v_{K-1}\}$03 pixels. Ground-truth masks for lumen, CPA, and NCPA are obtained by rasterizing 3D meshes. The split is $\{v_1,\ldots,v_{K-1}\}$04 training, $\{v_1,\ldots,v_{K-1}\}$05 validation, and $\{v_1,\ldots,v_{K-1}\}$06 testing images, and each image is automatically partitioned into three equal vessel regions—proximal, mid, and distal—along its long axis. Following Monteiro et al. (2023), effectiveness is measured as the absolute error between the target and the segmentor’s predicted area in each region,

$\{v_1,\ldots,v_{K-1}\}$07

A well-localized intervention is therefore expected to yield low error in the intervened region and minimal change in the non-target regions (Xia et al., 22 Mar 2026).

Representative quantitative results indicate improved localization. For $\{v_1,\ldots,v_{K-1}\}$08, the reported mid-region error is $\{v_1,\ldots,v_{K-1}\}$09 for No-CFT, $\{v_1,\ldots,v_{K-1}\}$10 for Reg-CFT, and $\{v_1,\ldots,v_{K-1}\}$11 for Pos-Seg-CFT, while the distal off-target error is $\{v_1,\ldots,v_{K-1}\}$12, $\{v_1,\ldots,v_{K-1}\}$13, and $\{v_1,\ldots,v_{K-1}\}$14, respectively. For $\{v_1,\ldots,v_{K-1}\}$15, the distal error is $\{v_1,\ldots,v_{K-1}\}$16 for No-CFT, $\{v_1,\ldots,v_{K-1}\}$17 for Reg-CFT, and $\{v_1,\ldots,v_{K-1}\}$18 for Pos-Seg-CFT; the proximal and mid off-target errors are also reduced to $\{v_1,\ldots,v_{K-1}\}$19 and $\{v_1,\ldots,v_{K-1}\}$20, compared with $\{v_1,\ldots,v_{K-1}\}$21 and $\{v_1,\ldots,v_{K-1}\}$22 for Reg-CFT. The accompanying qualitative comparison for a $\{v_1,\ldots,v_{K-1}\}$23 mm$\{v_1,\ldots,v_{K-1}\}$24 CPA intervention reports that Reg-CFT exhibits strong unwanted expansion of calcified plaque in non-target segments and diffuse global brightening, whereas Pos-Seg-CFT confines the added plaque to the intended vessel segment with sharply localized change and no visible artifact elsewhere (Xia et al., 22 Mar 2026).

6. Applications, limitations, and relation to adjacent methods

The reported applications of Seg-CFT include data augmentation, de-biasing datasets, modeling disease, and explainability through direct visualization of anatomical “what if” scenarios. The positional extension adds disease-progression modeling of focal stenoses, targeted augmentation for underrepresented lesion locations, and anatomically interpretable causal analyses centered on localized structural change (Xia et al., 29 Sep 2025, Xia et al., 22 Mar 2026).

Several limitations are explicit. In the original Seg-CFT formulation, the current method intervenes only on a single numeric property, namely area; shape, position, or texture interventions remain future work. The method also assumes independence between multiple structure-specific parents such as NCPA and CPA, and a full causal graph over masks is described as an open challenge. In the positional variant, region masks must be defined a priori, for example along vessel segments, and extremely fine subdivisions, such as fewer than $\{v_1,\ldots,v_{K-1}\}$25 pixels, may challenge the accuracy of segmentation-derived areas (Xia et al., 29 Sep 2025, Xia et al., 22 Mar 2026).

A recurrent misconception is that anatomy-aware counterfactual editing necessarily requires a user to supply hypothetical pixel-level masks. Seg-CFT is specifically framed as avoiding that requirement: the user specifies scalar targets, while the frozen segmentor provides the measurement used in the loss. Positional Seg-CFT preserves that scalar-based design; the regional masks are fixed anatomical partitions used to aggregate the segmentor output, not bespoke counterfactual masks. Another misconception is that regressor-guided fine-tuning is sufficient for structure-specific interventions. The reported evidence argues otherwise: regressors can rely on global spurious cues, whereas segmentation-derived supervision better couples the intervention objective to the target anatomy (Xia et al., 29 Sep 2025, Xia et al., 22 Mar 2026).