MACAW: A Causal Generative Model for Medical Imaging
Authors: Vibujithan Vigneshwaran
Erik Yuiti Ohara
Matthias Wilms
Nils Forkert
Venue: Accepted by TMLR ()

Abstract: Although deep learning techniques show promising results for many neuroimaging tasks in research settings, they have not yet found widespread use in clinical scenarios. One of the reasons for this problem is that many machine learning models only identify correlations between the input images and the outputs of interest, which can lead to many practical problems, such as encoding of uninformative biases and reduced explainability. Thus, recent research is exploring if integrating \textit{a priori} causal knowledge into deep learning models is a potential avenue to identify these problems. However, encoding causal reasoning and generating genuine counterfactuals necessitates computationally expensive invertible processes, thus restricting analyses to a small number of causal variables and rendering them infeasible for generating even 2D images. To overcome these limitations, this work introduces a new causal generative architecture named Masked Causal Flow (MACAW) for neuroimaging applications. Within this context, three main contributions are described. First, a novel approach that integrates complex causal structures into normalizing flows is proposed. Second, counterfactual prediction is performed to identify the changes in effect variables associated with a cause variable. Finally, an explicit Bayesian inference for classification is derived and implemented, providing an inherent uncertainty estimation. The feasibility of the proposed method was first evaluated using synthetic data and then using MRI brain data from more than 23000 participants of the UK biobank study. The evaluation results show that the proposed method can (1) accurately encode causal reasoning and generate counterfactuals highlighting the structural changes in the brain known to be associated with aging, (2) accurately predict a subject's age from a single 2D MRI slice, and (3) generate new samples assuming other values for subject-specific indicators such as age, sex, and body mass index.

Published in Transactions on Machine Learning Research (03/2026)
MACAW: A Causal Generative Model for Medical Imaging
Vibujithan Vignoshwaran University of Calgary, Canada
Erik Y. Ohara
crik.okarr@uccalooy..c
University of Calgary, Canada
Matthias Wilms
wilmm:wmichh.c. edu
University of Michigan, USA
Nils D. Forkert
nils. mabkforrkrt@uallgrrr..
University of Calgary, Canada
Reviewed on OpenReview: https: // /t openreview. net/ forum? id= id=eYW0370904
Abstract
Although deep learning techniques show promising results for many neuroimaging tasks in rescarch settings, they have not yct found widespread use in clinical scenarios. One of the reasons for this problem is that many machine learning models only identify correlations between the input images and the outputs of interest, which can lead to many practical problems, such as encoding of uninformative biases and reduced explainability. Thus, re- cent research exploring integrating a priori causal knowledge into deep learning models a potential avenue to identify and overcome these problems. However, encoding causal reasoning and generating genuine counterfactuals necessitates computationally expensive invertible processes, thus restricting analyses to small number of causal variables and ren- dering them infeasible for generating even 2D images. To overcome these limitations, this work introduces a new causal generative architecture named Masked Causal Flow (MACAW) for neuroimaging application: Within this context, three main contributions are described. First, a novel approach that integrates complex causal structures into normalizing flows is proposed. Second, counterfactual prediction performed to identify the changes effect variables associated with cause variable. Finally, an explicit Bayesian inference for clas- sification is derived and implemented, providing inherent uncertainty estimation. The feasibility of the proposed method was first evaluated using synthetic data and then using MRI brain data from more than 23,000 participants of the UK biobank study. The eval- uation results show that the proposed method can (1) accurately encode causal reasoning and generate counterfactuals highlighting the structural changes in the brain known to be associated with aging, (2) accurately predict a subject's 's age from single 2D MRT slice, and (3) generate new samples assuming other values for subject-specific indicators, such as age, sex, and body mass index.
1 Introduction
Recent advances in medical imaging and the emerging availability of digital health records have resulted in an abundance of data in healthcare. This wealth of data, along with the continuing increase in computing resources, has helped the field of medical image analysis enter new era deep learning. This leap holds significant promise for disease prevention, diagnosis, and treatment planning (MacEachern & Forkert (021). However, the translation of deep learning techniques from academic research into clinical deployment has
Code available httpp://gittubb.com wtwp///atta/aaaaaaaaaaa

---

Published in Transactions on Machine Learning Research (03/2026)
been slow (Winder et al., 2024). The primary challenge often lies in the mismatch between the data used for training and the data encountered in real clinical scenarios (Castro et al. 2020p. This mismatch often causes models that perform well in research settings to generalize poorly when applied in real-world clinical environments. Furthermore, the limited clinical translation of deep learning models can be partly attributed to their "black box" nature, which lacks inherent explainability for their decisions (Vercio el al. 2020
Both of these issues can be primarily related to the way these models are trained. More precisely, dis- criminative deep learning techniques designed to maximize accuracy a given dataset. This approach encourages models to exploit all possible correlations in the data, including shortcuts, to improve performance on that specific dataset. As a result, these models often capture spurious correlations and uninformative bi- ases present in imaging data (Souza et al. 2024). Thus, while they excel in research settings, they frequently fail to generalize effectively or provide meaningful explanations for their predictions when they are deployed in clinical settings.
One approach to better analyze and interpret the available medical data is to move beyond correlation- based analyses and instead utilize causally informed models. This approach has not only the potential to identify which variables (e.g., age, sex, race, scanner type, etc.) causally affect medical images but also to what extent. Within this context, cansal discovery methods address the question of which variables have an impact, while causal reasoning methods explore how they affect the images (Sanchez et al.] 2022b). While causal discovery is beyond the scope of this work, we demonstrate in this work how true causal reasoning can be achieved using medical images through generative modeling. This causal generative modeling provides a powerful framework to explore variables related to the data- aata-generating process in a transparent manner.
A straightforward approach to causal generative modeling is integrating a known structural causal model (SCM) into an existing generalive framework. Pearl (012)) describes causal model as comprising three levels of complexity, referred to as the "causal ladder". These levels, in order of increasing complexity, are association, intervention, and counterfactuals. Association pertains to correlations in the data and is solely focused on modeling the probability distribution in the dataset. Intervention relies on structural assumptions about the underlying data-generation process and involves exploring interactions with variables to observe how outcomes change on a population level. Many existing deep generative models, such as conditional VAEs, conditional GANs, and causal generative neural networks, only fulfill the requirements up to the intervention level. Lastly, Counterfucluals can be used to investigate hypothetical scenarios on an individual level, encompassing both interventional and associational inquiries. A counterfactual query essentially asks the trained causal generative model, "How would a subject's data appear if it had been acquired under different conditions?". The generation of counterfactuals not only provides causal insights into metadata and medical images but also holds signíficant potential tasks with practical applications, such as fairness, bias mitigation, data augmentation, data harmonization, and digital twins (Pawlowski et al., 2020
While the term "counterfactual" is frequently employed in literature, only few models, such as Deep Structural Causal Models (DSCMs) (Pawlowski et al. 2020), Neural Causal Models (NCMs) (Xia et al. 2022), Hicrarchical Variational Autocncoders (HVAE) (Ribciro et al. 2023). VQ-VAE and generalized lincar models [Png el al. 2024). and Diffusion SCM (Dir--CCM) (Sanchez & Tsaltaris 2022b), have the ability to produce causally-grouss grounded counterfactuals by following the Abduction Actiiionn-cttti-- Action-Prediictionnn Prediction steps as defined by Pearl 2 Other methods this domain mostly focus on generating realistic adversarial images aimed at deceiving classifiers. This scarcity of true counterfactual models for images arises from the requirement for invertible deep networks to achieve genuine counterfactual generation. To date, only the normalizing flow model is invertible by nature and has been previously used for causal modeling (Pawlowski et al. 2020 In this series of studics, each causal variable within graph is modeled by a separate conditional normalizing flow, necessitating multiple normalizing flow models to represent the complete causal framework.
Alternatively, our work demonstrates, for the first time, that a single normalizing flow model coupled with masked autoencoders is sufficient to efficiently model complex causal structures. While we only consider three causal variables in the experiments in this work, the approach presented can be easily extended to any number of causal variables by defining the corrosponding adjacency matrix. Within this context, we demonstrate that combining a standard dimensionality reduction technique, such as kernel principal component analysis (KPCA), with normalizing flows can effectively encode a given causal structure. In theory, the KPCA
2

---

Published in Transactions on Machine Learning Research (03/2026)
technique can be replaced with any dimensionality reduction technique. Within its medical context, relevant previous study is the work by Wilms et al. [2022] on counterfactual image generation and Bayesian classification. The key difference to the work presented here lies in their exclusive focus on a setup assuming independence between conditioning variables, while the method presented this work encodes a complex causal structure interactions between variables using masked autoencoders. Thus, the main contributions of this work can be summarized as follows:
causal structures of the data-generating process into a generative model. We present and evaluate a novel method named masked causal flow (MACAW) for encoding the
2. Using this model, we generate new high-resolution brain images and evaluate counterfactuals asso- ciated with brain aging.
The model's explicit density estimation enables direct Bayesian classification, eliminating the need for a separate discriminative model.
2 Background
2.1 Causal graphical models
Causal graphical models employ nodes to depict variables and edges to define their connections within directed acyclic graph (DAG) to provide an intuitive method for defining and exploring dependencies. The two important conditions, the Markov condition and faithfulness condition, ensure that conditional independence in the joint probability distribution is accurately reflected in the causal graphical model. Readers interested in these concepts can find in-depth information in the book by Peters et al. 2017).
Consider the d-dimensional random vector x [x1,...,,TT] x,]" E Rd, which follows the distribution px(x). Let causal graphical model be DAG containing nodes, each represented by x; [1,4]; the connection between nodes is defined by adjacency matrix E (0, (0.1)dxd, It essential emphasize that the adja- cency matrix of any topologically ordered DAG is always triangular. Assuming Markovianity, G is valid representation of (x(x) if and only if the probability density px(x) can be factorized follows:
d
(x II (1))((a)))
(1)
Here, T(;;) denotes the set of parents of the node 2;, where =xx;) {zy 1}. The basic network structure (adjacency matrix A) is typically established by making of existing causal knowledge in the field. In cases with no or limited existing causal knowledge about the data-generating process, causal discovery algorithms can be utilized for this purpose (Glymour et al. 2019
The typical method for modeling a probabilistic causal model to use Structural Equation Modeling (SEM) with random noise. The structural equation for each variable xy defined as S xy = (x(r)) ny), where n, represents the mutually independent exogenous noise variable of the noise distribution Pa- In this formulation, the observational distribution of variables px(x) can be conceptualized sampling from noise distribution Pn. and then applying set of structural equations S to the sampled values. This implics that the observed variables are influenced by both, the noise distribution and the causal relationships represented by the structural equations.
being generated by
The causal calculus (Pearl 2012) was created for the utilization of a causal model. Specifically, the do() operator enables an intervention in the model. When applying the do() operator, specific functions are replaced with a constant in an SEM. Similarly, in the corresponding DAG, the edges going into the target of intervention removed, but the edges exiting the target retained.
3

---

Published in Transactions on Machine Learning Research (03/2026)
2.2 Normalizing flows
Normalizing flows are used to model complex probability distributions, denoted as Px, by applying sequence of transformations T = Tio..-oT to simple density prior Ps. Transformations in T must be both invertible and diffcrentiable to allow training of the model using the change variables formula:
Px(x)= = Px (x)) det./r (xx)
Here, det.J--1 is the determinant of the Jacobian of the inverse transformations. Efficient model optimization hinges on the ease computing this determinant. As a result, Jacobians are often designed as triangular matrices, allowing for computation in (n)) time. Various techniques have been introduced in the literature achieve this triangular structure, with common approaches including the use of coupling and autoregressive functions.
Of particular relevance to this present work, briefty describe autoregressive functions next. Autoregressive functions transform variable z; using variables to 2;-1, which in turn constrains the Jacobian of the transformation to be lower triangular. This is similar to writing a multivariate density Px as a product of univariate conditional densities:
d I a18-1) i=2
px(x) = (x(
(3)
Following this step, the flow model is trained directly through maximum likelihood optimization using equation More detailed descriptions about normalizing flows are, for example, provided by Kobyzev el. al. h2021h.
2.3 Related work
This subsection outlines the role of related works in shaping the development of the proposed MACAW model. Specifically, our research built upon the foundations laid by Wehennel & Louppe (2021) their work on graphical flows and [Khemakhem al. (2021) in their study on causal flows. Theoretical underpinnings of our framework closely align with these studies, and our work builds directly these established proofs.
Wehenkel & Louppe (2021) highlighted the similarity between equation and equation 3 and argued that autoregressive transformations can be interpreted as a method to model a causal network with predeter- mined node ordering. Conversely, in the case of a specific DAG with a predefined causal relationship, the autoregressive conditioners can be selectively masked to incorporate the causal relationship into the model. Based on this new perspective, they proposed the graphical normalizing flow technique, a new invert transformation with either a prescribed or a learnable graphical structure to inject domain knowledge into normalizing flows. In similar work, Khemakhem et al. (2021) suggested that SEMs and autoregressive flows are similar and introduced framework called causal autoregressive flow (CAREFL) for causal discovery. Furthermore, they showed that normalizing flows can generate effective counterfactual queries due to their invertible nature.
Both of these studies primarily focused on identifying the causal structure or topology within a given dataset, operating with limited variables. Specifically, the counterfactual implementation only permitted coupling flows, accommodating two sets of independent variables. In contrast, the primary focus of our work is to develop method for classification and counterfactual generation for digher-dimensionall datasets, such images, and to encode complex (non-antoreggressive) causal structures into the flows. Consequently, model is required that scales effectively for larger dimensions and allows the parallel exccution of the flow. Therefore, we introeece a neural network called the causally masked autoencoder (C-MADE) in this work, which is inspired by the masked encoder developed by Germain et al. (2015). This implementation requires only a single forward pass to compute all causal dependencies and their respective conditional likelihoods, making a computationally efficient density estimator compared to existing alternatives. Subsequently, multiple C-MADFs are arranged in sequence to form the Masked Causal Flow (MACAW), akin to how
4

---

Published in Transactions on Machine Learning Research (03/2026)
Papamakarios et al. (2017) utilized stacked MADEs to create masked autoregressive flows. The following
section provides a detailed explanation of this method.
3
Methods
DAG
C-MADE
MACAW
X
X
(1,0)
Xo
Xo
zO
C-MADE
C-MADE
1
(5
X1
exp(s])x1 +b1
X1
X2
X2
exp(s2)x2+b2
X
(5,
Figure A causal DAG (left) and its respective C-MADE network (center). The MACAW architecture (right) consists of multiple C-MADE networks connected a series, thereby forming normalizing flow.
3.1 Causally-masked autoencoders (C-MADE)
We constructed C-MADE networks to efficiently represent the causal structure using a neural network for a given causal graphical model with the adjacency matrix A. For simplicity, we assume that a neural network has the same number of input units (x) and output units (z). Since the output 2, must depend only on its parents (2,) C x, there must be computational path between the output unit 2 and any non-parent input units. A convenient way of zeroing connections multiplying cach matrix element-wiss wise with binary mask matrix derived from A, whose entries that are set to zero corresponding to the connections that need to be removed. Thus, the masks are essentially responsible for satisfying the causal structure.
Furthermore, we assume that our input is D-dimonsional, i.c., x = [ri,...,Ip]] where the adjacency matrix A characterizes the causal graph. impose the causal properly, we first assign an integer each unit the input and output layers from D. After numbering all the units, causal constraints on each unit are simply imposed by masking the connections based matrix A.
While hidden layers can be added to this network, the causal parental structure has to be accurately main- tained through the hidden layers. In this setup, ith hidden neuron in jth layer, h/, is connected to its parents I (h²) h well as h which ensures that the learned hidden values propagated to the final layer. This can be done by augmenting the adjacency matrix A's diagonal units to 1 for all input-hidden layer connections and be a multiple of D to preserve the causal utructure. If hidden layer has n x D neurons, A can be duplicated n times and stacked to create the adjacency matrix for that layer. However, as the neurons in the final layer (output layer) are connected only their parents, the network becomes strictly triangular, thus preserving the initial causal DAG structure. Figure shows causal graph and its corresponding C-MADE network.
layer connections. Moreover, the number of neurons in a hidden layer must
3.2 Masked Causal Flow (MACAW)
Let us consider a normalizing flow model whose transformed probability density p(z) is given as the product of prior distributions. The prior distributions for the source variables (variables that do not have parents) should be selected to resemble their original distributions closely. This is because the distributions of these variables are not modified throughout the training process. A standard normal prior is assigned to all variables that are not sources. Thus, a non-source variable z distribution is given by p(z: m<=))) _N(00, 1). This can achieved by using the change of variable formula equation [ with a series affine transformations Tü where E [1, K]. The transformed variables are denoted as m, For notation consistency, we assign x mg and z mæ. Each affine transformation T. has scale and shift vectors, S, and b,, estimated by unique C-MADE network Ct.
5

---

Published in Transactions on Machine Learning Research (03/2026)
(8t, b.) (m
(4)
m: = cxp(ss)mmm + b
(5)
Since C-MADE constructed as a causal network, the predicted affine transformation parameters of a variable can be considered as a function of its parents. Thus, the likelihood of the transformed probability density complies with equation 2 The negative log-likelihood of p(z) was used as the loss function to optimize the C-MADE networks' weights. Following the optimization, the likelihood of data point can be readily measured using forward flow. Stacking multiple C-MADEs in sequence provides the flexibility required to transform a complex distribution to the defined distribution (Papamakarios et al. 2017).
3.2.1 Generative sampling
To generate data using the trained model, we start by sampling the source variables from their prior distri- bution. Next, we proceed to sample each of the other variables iteratively by performing a backward flow. Moreover, in the case of interventional sampling, we have the option to set a particular value for a variable and continue sampling the remaining variables as usual.
3.2.2 Counterfactual inference
Counterfactual queries aim to assess statements about hypothetical situations of already existing obser- vations. Let's assume that the MACAW was used to model the density of dataset X containing data vectors ix", ".....,,x",, x Within this context, one obscrved data vector from the datasct is denoted as x | ..I,...,IpI) For instance, if variable Vj had taken the value xj = in our observed feature vector ohe counterfactual queries can be used to determine what the value of variable 2; would have been. This is denoted as Ti, According to Pearl (2012), generating causal counterfactuals requires three steps: abduction, action, and prediction. The abduction step evaluates the probability distribution over latent vari- ables zobs given observations xobe, In our model, this can be simply done using forward flow, computing the transformation zabs = T(x The next step (action) is to intervene and fix the value of Tj to a specific value a, denoted as do(r, = a xx"),, which makes it independent of its causes x(x)). In this step, the corresponding value of change of Ti,2, is adjusted in the transformed space zore. The final step (prediction) is performed by computing an inverse transformation pass of the intervened z zas to generate the counterfact.ual feature vector xef, which is done through backward flow. The algorithm outlinees the steps for the counterfactual process (Khemakhem et al., 2021).
Algorithm Counterfactual query
Require: observed data xabs, of variable æy, and of value 1. Abduction forward flow: gobe +- T(x
2. Action change hypothetical values in the space
(a) obs t + T, (T)))) Z5 <- a))
(b) aizytta + 20t for i $
obs
3. Prediction backward flow X2,+a - T
(zabs
(
3.2.3 Bayesian classification
Let's denote our input distribution p(x) as a joint distribution of set of features f and a parent variable c that we need to classify, which can take one value from (c),...,c#t) Thus, p(x) can be written as p(c, f). When we set = ct, the forward flow of the network provides the posterior of p(c= ef f). This can be easily computed for all possible values of {ct,..., c)) simply by setting the class variable accordingly. Thus, for given class = c², the posterior can be computed using Bayes' theorem as follows:
p(c= f) = plc = e
(6)
E.p(c=effj
6

---

Published in Transactions on Machine Learning Research (03/2026)
Then, the class with the maximum a-posteriori (MAP) is chosen as the predicted class label. It is important to note that for each class variable, it is necessary to evaluate all possible posterior values and perform a forward pass with each of them.
3.3 Dimensionality reduction
As described carlier, normalizing flows are known for their resource-hungry nature. Consequently, to maintain the dimensions of the variables throughout the flows, a naive implementation of MACAW results in severe computational challenges. Therefore, when dealing with images, operating with all image pixels is impractical and even 2D images have to be projected onto low-dinensional latent space before applying the MACAW model for density estimation. In this work, we utilized Kernel Principal Components Analysis (KPCA) to reduce the dimensionality of the training images. However, this technique can be easily replaced with any dimensionality reduction technique. The basic idea of this method to project data that is not linearly separable onlo higherr-dimensional space, thereby transforming into linearly separable form. We selected this technique for two primary reasons: (1) projected features are not linearly correlated with each other, they projections onto an orthogonal basis. This property reduces the dependence between projected image fealures, leading us to assume that this makes the optimization faster. (2) The inverse (preinage) of the projected KPCA features (here: polynomial kernel with a degree of 3) can be computed efficiently using existing algorithms.
3.4 Quantitative evaluation
Assessing the effectiveness of counterfactual image generation techniques poses challenges due the absence of real-world ground truth. Monteiro et al. (2023) and then Melistas et al. (2024) proposed framework for evaluating counterfactual images generated by various generative techniques. Their framework uses struc- Aeriiion-Acttioo-PPeeee steps to assess the performance of the different methods. In line with their study, we also conduct a specific set of experiments to evaluate the quality of our counterfactuals in terms of Realism and Effectiveness.
tural causal models and employs the counterfactual inference based on Pearl's
Realism: We used the Fréchet Inception Distance (FID) to quantify the semantic similarity between the generated counterfactual images and images in the training set. Therefore, real and generated samples were passed through Inception v3 model (Szegedy et al. 2015) (pre-trained on Imagenet) to extract their semantic features and to calculate the FID between these two feature representations. A lower FID indicates that the features contain similar semantic information.
Effectiveness: Effectiveness aims to assess how well a counterfactual query performs. To quantitatively evaluate the effectiveness a particular counterfactual image, we trained a separate discriminative deep learning model on the training set to predict. the value of the intervened variable based on the image. While this traditional inference model may capture spurious correlations in the training data, it still provides information about the degree of confidence in the counterfactual generations.
4 Experiments and Results
We performed two experiments to investigate the effectiveness of MACAW. In the first experiment, we used synthetic data with a known causal structure to demonstrate that likelihood estimation, intervention, and comnteracttall analysis function as expected. In the second experiment, we utilized the UK Biobank brain MRI images to showcase the practical applications and benefits of interventional sampling, counterfactual inference, and Bayesian classification that are only possible with true causal deep learning framework.
4.1 Synthetic data
4.1.1 Data
The first experiment aimed to evaluate if the proposed MACAW model can effectively learn the given causal structure within dataset and accurately perform interventional and counterfactual queries. Therefore, we
7

---

Published in Transactions on Machine Learning Research (03/2026)
created a synthetic dataset with predefined structural equations for sample generation to investigate this in detail. The causal structure of this dataset and the specific structural equations are defined follows:
20 140
21
22 200 + .11 +
(7)
23 = 2x0 + 113
24 = 62222 + 721
Where the noise variable ny was sampled from uniform 2 and normal W distributions as follows:
no 1(0,1) 11 N(1 (0, n4 7 (00,0.)
122 (0,
(8)
We generated 10,000 samples using equation and partitioned it for training (70%) and testing (30%). Next, defined prior distributions for the latent variables as follows: 20 drawn from uniform distribution u(0, 1), 21 drawn from a normal distribution 1),, while all other latent variables followed standard Gaussian distributions.
4.1.2 Training
For this experiment, our MACAW model consisted of 10 C-MADEs, each consisting of three hidden layers with 15 neurons each. During the training process, we utilized the negative log-likelihood as the loss function and halted training when reaching the point of minimal validation loss.
4.1.3 Generative sampling
After model training, we generated 10,000 random samples using the MACAW model. The distribution from SEMs and those generated using MACAW were very similar. Table 1 outlines the mean and the variance of the variables generated. We further compared the maximum mean discrepancy (MMD) between the ground-truth and generated datasets, obtaining value of 3.187 x 10-4, which indicates only small distance between the two distributions. In addition, we have included the generated versus ground-truth data distributions in the appendix.
Table 1: Mean and the variance of generated samples using SEM and MACAW.
D2
F3
SEM
MACAW
SFM
MACAW
SEM
MACAW
Mean
4.01
3.99
2.99
2.99
73.96
74.55
Variance 5.32
5.92
0.58
0.53
2429.86
2863.50
4.1.4 Counterfactuals
To measure the accuracy of the counterfactual inference quantitatively, we estimated the counterfactual values for the action do(æ; 2 x x²) in the test set. We computed the absolute difference between the observed values and the counterfactual values and then summed them up. Fig. 2(a) shows that the counterfactual query only affected 12 and while other variables remained unaffected, as expected based on the causal structure ure depicted in equation 7 Nexi, we compared the counterfactual values with their respective ground truth value. These ground truth values were directly determined by substituting 22 = 2 in the equation x4 6xzazhe n4, where the noise term ny was computed as x: 6.rghex:'h Fig. 2 shows that the expected values of the ground truth and counterfactual values are perfectly aligned in the range of [-5,15), which is due to 99% of 2a values falling within this range during training.
8

---

Published in Transactions on Machine Learning Research (03/2026)
Cf(x2-2)
Cf vs Ground truth
(a)
(b)
Total Total change after
200
200
4400
-800
X0
x1
X2
X4
-10
-5
5
10
15
30
Variables
X2 values
Figure/ 2: (a) The sum of the absolute difference between each variable the test set after the counterfactual x (b) The expected values of the predicted counterfactual values
query for the action do(x2 and respective ground truth values.
4.1.5 Ablation Study: Causal vs. Correlation-Bant Based Conditioning
A natural question is whether the causal masking mechanism in MACAW provides any advantage over a model that simply conditions on all available preceding variables. We address this question by training two MACAW models on the same synthetic data: one with the true causal DAG, and one with fully-conmected DAG in which every variable may depend on all variables that precede it in the topological ordering. The two models identical in architecture, hyperparameters, and training procedure, and they differ only in the masks induced by their respective DAGs.
Data- getea-enerating process. We design five-variable structural causal model with a strong confounding structure:
20 N(O, 1),
21 20 + E1,
61 ~ 0.33,
32 5.0 + 62
62 0.3,,
(9)
2x +
((0,0.3)
= 3 x2 +
6 (10,00))..3
The true causal edges are {0~1, 0-+2, 1-3, 2-+4]. Crucially, Vg and 22 share the common cause To with a high correlation = = 0.988). The fully-connected DAG adds all edges {i - j::: f}, introducing six spurious connections (most importantly + 2), which allows the model to exploit the strong æ 22 correlation if it were direct causal link.
Observational fit. Both models are trained on N=10,000 samples using identical hyperparameters (3 flow layers, hidden multipliers [4, 8, 4]. 800 epochs). As expected, both achieve comparable observational fit as measured by Maximum Mean Discrepancy (MMD) between generated and real samples (causal: 0.000723, fully-connected: 0.000652). The fully-connected model has strictly greater capacity (a superset of edges), it can represent the observational distribution slightly better than the causal model. This confirms that observational metrics alone cannot. distinguish the two approaches.
Interventional predictions. The key discriminating test is do(z) v). Under the true causal structure, is not a descendant of æ1, and they are conditionally independent given Zo. so the intervention should leave unchanged at [8[22] & 0. We sweep v E [-15, 15] and plot E|zz do(z)=u)] for both models (Fig. Bh). The causal model correctly produces a flat response, while the fully-connected model shows a strong linear edge.
trend as a direct. consequence of the spurious
Counterfactual predictions. We further evaluate counterfactual queries of the form "What would 22, Ra. F4 have been if æ1 10?" Under the true causal structure, only 21 and its descendant 23 should

---

Published in Transactions on Machine Learning Research (03/2026)
Intervention Sweep: E|x2 ddoxi=vv]]
Causal E[x2]=-0.01
Fornecce model
(a)
[B4
-1
dolx1
Causal: dix1=10)
Cully-oomeeted: cflx2=10)
25
10-
320 sharggg afted
2.0
(b)
3.5
X1 A2
10 14
Figure 3: (a) When intervening on variable 21, the causal model correctly yields a near flat response, whereas the fully connected model exhibits a strong linear trend, reflecting the spurious edge ween the two variables. (b) When changing the variable I counterfactually, the causal model correctly limits changes 21 and 23, with negligible effects on the other variables. In contrast, the fully connected model introduces spurious changes in 22, which further propagate to TT.
change; No; 32, and x4 should remain at their observed values. Fig. 3 shows the total absolute change per variable after the counterfactual. The causal model correctly confines changes æ and I3, with near-zero changes in Zo, 22, and xt. The fully-connected model produces substantial spurious changes in 22 (via the + 2 edge) and consequently in xa (propagated through 2 > 4).
This ablation demonstrates that the causal masking in MACAW is not merely structural prior that aids training. but it essential for correct interventional and counterfactual inference. model with more capacity (more edges) fits the observational distribution equally well but fails precisely where causal reason- ing required, bocause it cannot distinguish genuine causal effects from spurious correlations inducod by confounders.
4.2
UK Biobank images
4.2.1 Data
The second and main experiment used neuroimaging data from the UK Biobank (UKBB) cohort (Miller et al. 20166.. This study retrieved data under application 77508: Explainable and interpretable machine learning solutions in computational medicine. The UKBB's Tl-weighted structural magnetic resonance imaging (MRI) used 3D MPRAGE sequence with 1-mm isotropic resolution and 208 256 256 mm field of view. In the first step, starting with all subjects with available Tl-weighted MRI data, participants with diagnosed brain-related disorders based on ICD10 codes (data field 41202, chapter V Mental and behavioral disorders and chapter VI -Diseases of the nervous system) were excluded. We obtained the participants' sex information from the genetic sex data field (22001) and their age from the recorded values during the imaging visit (data field: 21003-2.0). Participants younger than and older than were excluded from the analysis because there was not a sufficient number of participants in these specific age groups to perform the age-stratified train-test. split. Additionally, the body mass index (BMT) values were
10

---

Published in Transactions on Machine Learning Research (03/2026)
Age
Sex
Latents
Age
Sex
Age
Sex
BMI
BMI
Latents
Lalents
(a)
(b)
Figure 4: (a) Predefined causal graph for the UKBB dataset, incorporating age, sex, and BMI values. (b) Corresponding adjacency matrix, where filled cells indicate and unfilled cells indicate 0.
retrieved directly using the UKBB data field 21001. Subjects with NaN values for either sex, age, or BMI were excluded from this work, resulting in a sample size of 23,692 (male 11,050, female 12,642).
All Tl-weighted MRI data of the selected participants were aligned to the SRI24 atlas [Rohlfing et al, 2010), using the affine registration implemented in ANTs (Aantts et al.) 2011). From the 3D stack of data, we specifically chose the axial slice subject's image that primarily covered the lateral ventricular region, an area the brain that captures the extent of visible atrophy due to aging. Finally, the data was split into training and lesting sets of 80% and 20%, stratified by age.
4.2.2 Training
For model training, the extracted 2D images underwent center-cropping, leading to image sizes of 180 180 pixels. Subsequently, these images were projected onto 1500-dimensional subspace (latents) using KPCA. Training all 1500 latents using the same MACAW network did not converge successfully, whereas training with 60 latents yielded a better likelihood estimation. Consequently, we divided the latents into subgroups of 60 and trained a separate MACAW model for each group, optimizing the likelihood individually for each model. Each MACAW model incorporated the causal structure illustrated in Figure ! this setup, age and sex were treated as discrete distributions, while BMI was considered continuous. The prior distributions for sex and age were determined based on the distribution of the training data, adopting Bernoulli and categorical distributions, respectively. Gaussian distributions were used priors for both BMT and the latents. A 10% validation set was taken from the training set during the training phase to determine the best early stopping criteria.
4.2.3 Generative sampling
Generative sampling was conducted in an autoregressive manner using the trained models. The process began by randomly sampling age, sex, BMT, and 60 components from the first model. This process continued iteratively until 1,500 components were generated, which were then reconstructed by estimating the preimage of the KPCA components. During interventional sampling, spocific values for one or more parent variables (age, sex, BMT) were manually defined, while the remaining variables were sampled from the first model. This autoregressive sampling approach was then repeated using the subsequent models. The results from the generative and interventional sampling are presented in Figure
4.2.4 Counterfactuals
Counterfactual inference was performed on data points within the test set. Using this setup, we aimed to estimate how an image would appear if the person had different biological characteristics including age, sex, and BMI. Coutterfactuals were generated independently for each model and then reconstructed, following a process similar to the generative sampling. Figure D visually demonstrates that age has an expected impact on the ventricular and sulci regions. More precisely, increasing the age tends enlarge the ventricular area while decreasing the age variable reduces the ventricular volume. Furthermore, changes in the sulei
11

---

Published in Transactions on Machine Learning Research (03/2026)
Male, 58, BMI:22
Female, BMI:20
70,
Random
Female, 75, BMI:22
Male, 75, BMI:26
Female. 75, BMI:31
do(Age=
Figure 5: Results from two types of generative sampling: unconditional (top) and interventional (bottom). For interventionll sampling, we set the age variable to 75 years.
regions are also observed when changing the age variable. When considering sex-related counterfactuals, variations in brain size and subtle changes (Lotze et al. 2019) in the ventricular regions can be observed in the figure. Furthermore, when generating counterfactual images to illustrate sex differences, the BMI values appropriately change, with males exhibiting higher BMI values. This is attributed to the causal influence of sex on BMI. Finally, BMI changes result in alterations of the lateral and ventricular parts of the brain, which may be related accelerated brain aging (Beck al., 2021). In addition, Fig displays results from counterfactual queries simulating brain aging.
Female, 69, BMI33
Male, 53, BMI35
Male, 76, BML:25
Femalc, 79, BML26
Male, 57, BML 33
Male, 64, BMI: 24
Actual
do(Age=50)
do(Age=75)
do(Sex- Female)
do(Sex Male)
do(BM] 20)
do(BM]- 35)
Female, BM1333
Male, 75, BML34
Famale,, 76, BM1:24
Male, 79, 27
Mele, 57, BMI:
Male, 64, BMB:
Counterfactual
Difference
Figure 6: Illustration of the outcomes of counterfactual queries, where causal factors of randomly selected raw images were altered to generate new counterfactual images. Every column represents a distinct counterfactual query, identified by "do(.)" commands.
12

---

Published in Transactions on Machine Learning Research (03/2026)
Table 2:
Realism and effectiveness are measured on test, blurred, and counterfactual images. FID: Fréchet
Inception Distance, MAE: Mean Absolute Error.
Test
Blurred
CFss
C'6o
C'66s
C'Fo
CF7s
Realism (FID)
0.49
53.23
3.12
2.09
2.48
6.02
11.78
Effectiveness (MAE) +
3.63
7.25
4.82
4.71
6.72
10.23
Age:58.0
Age:81.0
Age:88.0
Age:71.0
Age:78.0
Counterfactuak
Ag.:48.0
Difference
Observation
Figure 7: Results from counterfactual queries simulating brain aging. Each counterfactual image corresponds to unique query "do(Age "²
To assess the effectiveness of counterfactual inference, we generated counterfactual age values for the entire test set, consisting of 4,739 images. Essentially, the query asks what would happen if everyone in the test set had age a, denoted as C.a,, where E 60, 60, 65, 70, 75). The following experiments conducted to assess these counterfactuals.
Realism: In the first step, the FID between the training and test sets was measured to serve as the lower bound. Subsequently, all images in the test set were blurred using Gaussian blur with a standard deviation of 1, and the FID between each original test image and smoothed test image was measured as a comparative baseline. Finally, the distance between each CF,, defined above and the training set was calculated. The results displayed in Table 2 demonstrate that the generated counterfactual images show superior realism compared to even slightly blurred images. Additionally, it was observed that realism notably high for the counterfactual images generated for the age of 65, which aligns with the centered age distribution of the data.
Effectiveness:
To assess whether the generated counterfactuals indeed contain predictive age information, we trained a separate classifier trained on the entire training set. Specifically, we utilized a popular archi- tecture for brain age estimation, the SFCN model (Peng et al.) 2021 trained for 2D slices, which achieved a mean absolute error (MAE) 3.63 for test images. Subsequently, the C.F, sets were tested for age prediction the ground-truth target value for MAE computation. Table 2 displays
using the counterfactual age these results, showing that CF6o and CFes performed reasonably well. Since the UKBB has large amount of training data for these age bins, it can be speculated that the model performs better when the queried counterfactual age is closer to the actual age (less severe changes on average). Therefore, to evaluate how the generated images change with counterfactual age differences, we computed the difference between the actual age and the counterfactual query age for all CF,, sets and subsoquently measured the MAE value for this counterfactual age gap. Figure X illustrates these results, indicating that the counterfactual query performs well within the age range of -10 to +10.
13

---

Published in Transactions on Machine Learning Research (03/2026)
Cf age gap Absolute error
25
15
Absolute
-24.0
-19.0
11.0
16.0
26.0
Cf age gap
Figure 8: The figure displays how the effectiveness (MAE) of the generated images changes with the difference between actual age and queried counterfactual age. The dotted lines show the MAE of the test set (unaltered images).
4.2.5 Classification
Next, employed the model's Bayesian classification abilities to predict the age of each participant in the test dataset as described in section 3.2.3 Therefore, we computed the posterior for each age and selected the age with the maximum a-posteriori (MAP) value as the predicted age. The overall accuracy of this prediction resulted in an MAE of 5.047 (standard deviation (std) = 0.052) when we used the first model (60 latents) for classification. Figure g illustrates the postcrior distribution for a participant and the disparity between actual and predicted values. The figures show that (1) the prediction assigns higher posterior values around the true chronological age and lower for distant ages, and (2) the error in the prediction distribution is centered around O. These two results indicate that the model identifies the causal connection between age and images quite effectively.
Max likelihood = 71, Actual age=75
Difference between actual and predicted age
Likellhood
Count
51
61
66
76
55
10
Age
Difference
Figure 9: Posterior estimated by the model for each age, with the actual age being 75 and the predicted being 71 (left); distribution the differences between the actual age and the predicted age for the whole test set (right).
5 Discussion
This work presents a novel causal generative framework called MACAW, which satisfies all three levels of the causal ladder (association, intervention, and counterfactual) and holds considerable promise for various
14

---

Published in Transactions on Machine Learning Research (03/2026)
applications in medical and clinical contexts. Specifically, counterfactual generation offers causal insights the individual level, which cannot be achieved using corvelation-based approaches. Counterfactuals on the MRI data revealed brain regions associated with aging, sex differences, and BMI changes that are all in line with current knowledge.
Disentangling spurious correlations and identifying factors that causally influence medical images has been challenging. Discriminative deep learning models prioritize achieving accuracy on the training dataset, often exploiting shortcuts and spurious correlations. This makes it difficult to model and explore how variables like age, sex, and BMI causally affect medical images. In contrast, the causal generative modeling method described in this work encodes the data-ggnerating process through a causal DAG, offering causal explana- tions at both, the population level (interventions) and the individual level (counterfactuals)) Furthermore, causal counterfactual generation offers theoretical benefits with respect to fairness (Kusner et al. 2018). bias mitigation, data augmentation, data harmonization, and digital twin development. Furthermore, is im- portant to emphasize the distinction between conditional and causal approaches. In a standard conditional generation setup, the factors influencing image features are often assumed to be independent. For instance, when conditioning on the age variable, it is typically assumed that this has no impact on the BMI variable. However, in our sctup, the model considers causal relationships within the datasct. As result, when we intervene on the age variable, it correctly influences both, the BMI and the image features, reflecting genuine causal connections.
Quantitative assessments employing the metrics related to Realism and Effectiveness confirm that the gener- ated counterfactual images effectively explain relevant causal information in the images. Furthermore, using interventional sampling, we created new samples with predefined parent values. This process can generate new samples that align with the characteristics learned from the population. Consequently, these generated samples can serve as supplementary data for training other deep learning models, particularly in cases where certain classes have insufficient data. However, potentially more relevant for knowledge discovery in the biomedical context, it is possible to use the mothod proposed in this work as the basis for digital twins.
While the primary purpose of our model is counterfactual generation, it can also perform classification using Bayes' theorem, similar to other density estimation methods. However, its classification accuracy is lower compared to CNN-based methods, like SFCN for brain age estimation (Peng et al.. 2021). One potential reason for this finding is that, uulike lask-specific models, such as the SFCN, our model's 's lower-dinensional subspace must capture features sufficient for accurate image reconstruction. However, the key advantage of our approach is that a single trained model can classify multiple variables (e.g., age, sex, BMI), whereas CNN-based methods require separate models for each variable. Although classification is not the model's primary focus, it may serve as solid foundation for developing discriminative deep learning models in the
future.
Thus far, researchers have used VAEs (Pawlowski et al. 2020 et al. 2021), and diffusion models (Sanchez & Tsaftaris, 2022a Sanchez et al. 2022a for counterfactual image generation. However, for a causal model to truly perform counterfactuals, deterministic invertibility is essential (Pearl 2012). Normalizing flows intrinsically offer this invertibility, and the experiments described in this work demonstrated that a single normalizing flow combined with masked autoencoders is effective in modeling complex causal structures. The proposed framework grounded solid mathematical foundations. In partieular, using masking networks with a DAG-indueing adjacency matrix has been demonstrated as a valid causal generative prior in several previous studies, including Graphical Normalizing Flows (Wehenkel & Louppe 2021 and Causal Autoregressive Flows- CAREFL (Khemakhem et al. 2021). The theoretical basis of our framework aligns closely with these prior works, and our approach builds directly upon their established proofs. While these previous studies primarily focus on causal structure discovery, our work extends these principles to generate image- based counterfactuals, representing a novel application in this domain.
Ribeiro et al. 2023), GANs Nemirovsky
Due to the high dimensionality of the image space, we projected images onto latent space to enable feasible model training. The dimensionality of the latent variables affects the model in two ways. First, it impacts the reconstruction of the final counterfactual image, as using fewer latent variables may lead to information loss. Second, it determines how effectively the causal parents can influence the generated images,
15

---

Published in Transactions on Machine Learning Research (03/2026)
with larger latent spaces potentially diluting this effect. The ranked structure of principal components is particularly useful in this context. For instance, the first 50 principal components capture over 90% of the variance in the images in our experiments, while the remaining components were found to contribute only minor changes to the generated counterfactuals. Consequently, even if the causal inffuence of the parents diminishes the later latent components, it has minimal effect on counterfact performance. An extensive ablation study on this topic has been conducted previously (Ohara al.) 2025).
This study has several limitations, with one of the most notable being the reliance on predefined causal graph. In our research, assumed the availability of causal graph for the UKBB dataset. However, predefined causal graph may not be readily available in many real-world scenarios. Typically, rescarchers need to estimate causal relationships from existing domain knowledge or through randomized controlled trials, which may not always be feasible or ethical. In cases where a predefined causal graph is absent, causal discovery techniques can be employed (Glymour et al. 2019). Within this context, it important to emphasize that estimating causal relationships solely from obsorvational and pross-soctional data can be an extremely challenging lask, often even impossible. Additionally, to caplure the complex structure of the brain associated with aging, it is necessary to extend our method to 3D images, which still remains our future work. Another potential limitation of the proposed model relates to the computational demands of the classification task. Specifically, for each class variable, it is necessary to evaluate all possible values and perform a forward pass with each of them. Consequently, this process incurs a computational cost several times greater (equal to the number of classes) than the typical computational load associated with a standard discriminative network. Moreover, this classification process is limited to anti-causal prediction, meaning it involves predicting a top-level parent variable based its effects.
In conclusion, the experimental results demonstrate the potential and effectiveness of MACAW in generating interventional and counterfactual images, as well as performing Bayesian classification. Future research should prioritize the development of a technique that seamlessly integrates dimensionality reduction and the normalizing flow framework into a single model. This integration, in turn, would facilitate the efficient processing of 3D images, ensuring that the model effectively captures all crucial information. The proposed technique holds potential for exploring and identifying potential new biomarkers for various diseases.
References
Brian B. Avants, Nicholas J. Tustison, Gang Song, Philip A. Cook, Arno Klein, and James C. Gee. A reproducible evaluation of ANTs similarity metric performance in brain image registration. Neurolmage, (0l::2233-2044 February 2011. ISSN 1095-9572. doi: 100.:///..eroiimee200.00..... 20..0......
Dani Beck, Ann-Marie G. de Lange, Mads L. Pedersen, Dag Alnæs, Ivan 1. Maximov, Irene Voldsbekk, Geneviève Richard, Anne-Marthe Sanders, Kristine M. Ulrichsen, Erlend S. Dørum, Knut K. Kolskår, Einar A. Høgestøl, Nils Fiel Steen, Srdjan Djurovic, Ole A. Andreassen, Jan E. Nordvik, Tobias Kaufmann, and Lars T. Westlye. Cardiometabolic risk factors associated with brain age and accolcrate brain ageing. Human Brain Mapping, (300))700 720, October 2021. ISSN 1065-9471. doi: 111100000.....
Daniel C. Castro, Jeremy Tan, Bernhard Kainz, Ender Konukoglu, and Ben Glocker Morpho-MNIST: MNIST: Quantitative Assessment and Diagnostics for Representation Learning, October 2019. URL http://
arxiv. wwwii.oor//aaa wwwp/b//g///.../....r. 1109 10780 arX1v:1800.10780 [cs, stat].
Daniel C. Castro, Ian Walker, and Ben Glocker. Causality matters in medical imaging. Nature Communi- cations, 1161)33 July 2020. ISSN 2041-1723. doi: 10.1038//411477:00--77772 w. Number: 1 Publisher: Nature Publishing Group.
Mathieu Germain, Karol Gregor, Jain Murray, and Hugo Larochelle. MADE: Masked Autoencoder for Distribution Estimation. In Proceedings of the 32nd International Conference on Machine Learning, pp. 881-889. PMLR, June 2015. ISSN: 1938-7228.
Clark Glymour, Kun Zhang, and Peter Spirtes. Review of Causal Discovery Methods Based on Graphical Models. Frontiers in Geneties, 10, 2019. ISSN 1664-8021.
16

---

Published in Transactions on Machine Learning Research (03/2026)
Ilyes Khemakhem, Ricardo Monti, Robert Leech, and Aapo Hyvarinen. Causal Autoregressive Flows. In Proceedings of The 24th International Conference on Artificial Intelligence and Stalistics, pp. 3520-3522 PMLR, March 2021. ISSN: 2640-3498.
Ivan Kobyzev, Simon J.D. Prince, and Marcus A. Brubaker. Normalizing Flows: An Introduction and Review of Current Methods. IEEE Transactions on Pattern Analysis and Machine Intelligence, 43(11): 3964 3979, November 2021. ISSN 1939-3539. doi: /Tt.://A...00.299992222222 Conference Name:
TEEE Transactions on Pattern Analysis and Machine Intelligence.
Matt J. Kusner, Joshua R. Loftus, Chris Russell, and Ricardo Silva. Counterfactual fairness, 2018.
Martin Lotze, Martin Domin, Florian H. Gerlach, Christian Gaser, Eilcen Lueders, Carsten O. Schmidt, and Nicola Neumann. Novel findings from 2,838 Adult Brains Sex Differences in Gray Matter Brain Volume. Scientáfic Reports, 9(1):1671, February 2019. ISSN 2045-2322. doi: (.1008/s41508--08--338 Number: Publisher: Nature Publishing Group.
Sarah J. MacEachern and Nils D. Forkert. Machine learning for precision medicine. Genome, 64(4):416-4) 425,
April 2021. ISSN 1480-3321. doi: 10.1139/gan 20120000111
Thomas Melistas, Nikos Spyrou, Nefeli Gkouti, Pedro Sanchez, Athanasios Vlontzos, Yannis Panagakis, Giorgos Papanastasion, and Sotirios A Tsaftaris. Benchmarking counterfactual image generation. Advances in Neural Information Processing Systems, 37:133207-133222 2024.
Karla L. Miller, Fidel Alfaro-Almagro, Neal K. Bangerter, David L. Thomas, Essa Yacoub, Junqian Xu, Andreas J. Bartsch, Saad Jbabdi, Stamatios N. Sotiropoulos, Jesper L. R. Andersson, Ludovica Griffanti, Gwenaëlle Douand, Thomas W. Okell, Peter Weale, Iulius Dragomu, Steve Garratt. Sarah Hudson, Rory Collins, Mark Jenkinson, Paul M. Matthews, and Stephen M. Smith. Multimodal population brain imaging in the UK Biobank prospective epidemiological study. Nature Neuroscience, 19(11):1523-14 -1536, November 2016. ISSN 1526-1726. doi: 10.1038/mm 4393... Number: 11 Publisher: Nature Publishing Group.
Miguel Monteiro, Fabio De Sousa Ribeiro, Nick Pawlowski, Daniel C. Castro, and Ben Glocker. Measuring axiomatic soundness of counterfactual image models, 2023.
Daniel Nemirovsky, Nicolas Thiebaut, Ye Xu, and Abhishek Gupta. CounteRGAN: Generating Realistic Counterfactuals with Residual Generative Adversarial Nets, May 2021. arXiv:2009.0511 [cs, stat].
Erik Y Ohara, Vibujithan Vigneshwaran, Raissa Souza, Finn G Vamosi, Matthias Wilms, and Nils D Forkert Dimensionality reduction in 3d causal deep learning for ncuroimage generation: an evaluation study. Journal of Medical Imaging, 12(2):024506 024506, 2025.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked Autoregressive Flow for Density Estimation. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Nick Pawlowski, Daniel Coclho de Castro, and Ben Glocker. Deep Structural Causal Models for Tractable Counterfactual Inference. In Advonces in Neural Information Processing Systems, volume 33, pp. 869- Curran Associates, Inc., 2020.
Judea Pearl. The Causal Foundations of Structural Equation Modeling: Technical report, Defense Technical Information Center, Fort Belvoir, VA, February 2012.
Han Peng, Weikang Gong, Christian F. Beckmann, Andrea Vedaldi, and Stephen M. Smith. Accurate brain age prediction with lightweight deep neural networks. Medical Image Analysis, 68:101871, February 2021. ISSN 1361-8415. doi: 10.1016/j.media.22000 2020.ia.2220111177
Wei Peng, Tian Xia, Fabio De Sousa Ribeiro, Tomas Bosschieter, Ehsan Adeli, Qingyu Zhao, Ben Glocker, and Kilian M. Pohl. Latent 3d brain mri counterfactual, 2024.
Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: foundations and learning algorithms. Adaptive computation and machine learning series. The MIT Press, Cambridge, Massachuestts, 2017. ISBN 978-0-262-03731-0..
17

---

Published in Transactions on Machine Learning Research (03/2026)
Fabio De Sousa Ribeiro, Tian Xia, Miguel Monteiro, Nick Pawlowski, and Ben Glocker. High Fidelity Image Counterfactuals with Probabilistic Causal Models. In Proceedings of the 40th International Conference on Machine Learning, pp. 7390 7425. PMLR, July 2023. ISSN: 2640-3498.
Torsten Rohlfing, Natalie M. Zahr, Edith V. Sullivan, and Adolf Pfefferbaum. The SRI24 multichannel atlas of normal adult human brain structure. Human Brain Mapping. 31(5):798-819, May 2010. ISSN 1097-0193. doi: 10.1002 10.1002/hbm.20000 20906
Pedro Sanchez and Sotirios Tsaftaris. Diffusion causal models for counterfactual estimation. arXiv preprint 10160 2022a.
Pedro Sanchez and Solirios A. Tsaftaris. Diffusion causal models for counterfactual estimation, 2022b.
Pedro Sanchez, Antanas Kascenas, Xiao Liu, Alison Q. O'Neil, and Sotirios A. Tsaftaris. What is Healthy? Generative Counterfactual Diffusion for Lesion Localization. In Anirban Mukhopadhyay, Πaka Oksuz, Sandy Engelhardt, Dajiang Zhu, and Yixuan Yuan (eds.), Deep Generative Models, Lecture Notes in Computer Science, pp. 34 44, Cham, 2022a. Springer Nature Switzerland. ISBN 978-3-031-18576-2. doi: 10.1007/978-33 011179997888--1-1877777 1817-18576-2
Pedro Sanchez, Jeremy P. Voisey, Tian Xia, Hannah I. Watson, Alison Q. O'Neil, and Sotirios A. Tsaftaris. Causal machine learning for healthcare and precision medicine. Royal Society Open Science, 9(8):220638, August 2022b. doi:
220638, Publisher: Royal Society.
Raissa Souza, Anthony Winder, Emma A. M. Stanley, Vibujithan Vigneshwaran, Milton Camacho, Richard Camicioli, Oury Monchi, Matthias Wilms, and Nils D. Forkert. Identifying biases in multicenter mri database for parkinson's 's discase classification: Is the discase classifier secret site classifier? IEEE Journal of Biomedical and Health Informatics, 28(4):2047-2000 2024. doi: N0.1109/JBIITS JBH....22..3222222222
Christian Szegedy, Vincent Vanhoucke, Sergey Toffe, Jonathon Shlens, and Zbigniew Wojna. Rethinking the Inception Architecture for Computer Vision, December 2015. arXiv:1512.0056 [cs] version: 3.
Will Taylor-Melanson, Zahra Sadeghi, and Stan Matwin. Causal generative explainers using counterfactual inference: a case study on the morpho-mnist datasot. Pattern Analysis and Applications, 27(3):89, 2024.
Lucas Lo Vercio, Kimberly Amador, Jordan J Bannister, Sebastian Crites, Alejandro Gutierrez, M. Ethan MacDonald, Jasmine Moore, Pauline Mouches, Deepthi Rajashckar, Screna Schimert, Nagesh Subbanna, Anup Tuladhar, Nanjia Wang, Matthias Wilms, Authony Winder, and Nils D Forkert. Supervised machine learning tools: tutorial for clinicians. Journal of Neural Engineering, 17(6):062001, nov 2020. doi: 10.1088/1741-25552/able
Antoine Wehenkel and Gilles Louppe. Graphical Normalizing Flows. In Proceedings of The 24th International Conference on Artificial Intelligence and Stalistics, pp. 37 45. PMLR, March 2021. ISSN: 2640-3498.
Matthias Wilms, Jordan J. Bannister, Pauline Mouches, M. Ethan MacDonald, Deepthi Rajashekar, Sönke Langner, and Nils D. Forkert. Invertible Modeling of Bidirectional Relationships in Neuroimaging With Normalizing Flows: Application to Brain Aging. IEFE Transactions on Medical Imaging, 41 (9)):3331-2344 September 2022. ISSN 1558-254X. doi: on Medical Imaging.
Conference Name: IEEE Transactions
Anthony J Winder, Emma AM Stanley, Jens Fiehler, and Nils D Forkert. Challenges and potential of artificial intelligence in neuroradiology. Clinical Neuroradiology, pp- 1 1-13, 2024.
Kevin Xia, Kai-Zhan Lee, Yoshua Bengio, and Elias Bareinboim. The causal-neural connection: Expressive- ness, learnability, and inference, 2022.
18

---

Published in Transactions on Machine Learning Research (03/2026)
A
Appendix
A.1
Additional results for 1-D experiment
0.175
0.150
0.125
Donsiyy
0.100
0.075
0.050
Real
0.025
Sampled
0.000
10
-200
200
400
x4
Figure 10: After model training, we generated 10,000 random samples using the MACAW model. As shown in the figure, the distribution from SEMs (ground truth) and those generated using MACAW were very similar.
0.20
0.15
Density
0.10
0.05
-100
100 200
x3
x4
Figure 11: We generated samples from the model by intervening with two conditions: do(æg- 1.5 x obe) and do(x; 0.5 x xb The red lines represent the expected values derived from the equations (the gold standard). It is evident that the distribution's mean closely matches these expected equation-based values.
Distributions x and x
between x gold standard
0.040
xg
0.7
0.035
K
0.6
0.030
0.5
Denstty
0.025
Dereity
0.4
0.020
0.3
0.015
0.010
0.2
0.005
0.1
0.000
0.0
-100
100
200
300
-10
10
Difference
Figure 12: The figure displays both the observational and counterfactual distributions, along with the differences between the gold standard and the predicted values.
19

---

Published in Transactions on Machine Learning Research (03/2026)
A.2 MorphoMNISS benchmarking
For direct quantitative comparison with existing counterfactual image generation methods, evaluated MACAW on the MorphoMNISS benchmark introduced by Melistas et al. (2024). This benchmark provides a standardized evaluation framework with establishod metrics, dataset splits, and corresponding bascline results for three model families: variational autoencoders (VAE), hierarchical VAEs (HVAE), and generative adversarial networks (GAN), all operating within a structural causal model (SCM) framework.
B Models
VAE (Decp-SCM) (Pawlowski et al. ):20 : A deep structural causal model that combines variational infer- ence and normalizing flows enable tractable inference of latent noise variables for counterfactual reasoning, It supports all three levels of causal analysis association, intervention, and counterfactuals and is validated synthetic and medical imaging data.
HVAE (Ribeiro et al. 2023): A causal generative framework designed for high-dimensional data (e.g., im- ages), enabling accurate counterfactuul and interventional inference. Tt intuues structured causal mecha- nisms and demonstrates strong performance estimating direct, indirect, and total causal effects.
et al., 2024): A causal generative approach for interpreting image classifiers through counterfactual explanations. It identifies influential features using Shapley-based methods and generates interpretable counterfactuals, showing improved explanation quality compared to existing tools.
GAN
B.0.1 Data
MorphoMNIST (Castro 2019) is variant of MNIST where handwritten digit images are augmented with continuous morphological attributes. Following the benchmark protocol, we used the dataset provided by Mclistas et al., which contains 60,000 training and 10,000 test images of size 28 x 28 pixels with two continuous altribules: stroke thickness (I E [0.88, 6.26)) and pixel intensily (i [66.6, 254.9)), along with the digit class label E 9). Images were normalized to 1], matching the benchmark convention.
B.0.2 Training
The 20-dimensional encoded latents, together with the standardized thickness and intensity values and one- hot encoded digit labels, formed the input vector to the MACAW model. The causal edges for the MACAW model were defined to reflect the benchmark DAG: thickness - intensity, thickness - latents, intensity + latents, labels + latents, and autoregressive connections among the latent dimensions. The MACAW model was trained using Ny = 4 normalizing flow layers with hidden multipliers [4,6, 6, 6, 4] for 200 epochs.
B.0.3 Evaluation
We followed the benchmark's evaluation protocol with three metries:
Effectiveness. For each test image, a counterfactual was generated by intervening on a single attribute (do(t), do(1), or do(y)) with a target value sampled from the training distribution, ensuring the target differod from the factual value. Three anti-causal CNN predictors (a thickness regressor, an intensity regressor, and digit classifier) were independently trained on the real training images to predict each attribute from the images. These predictors were then applied to the generated counterfactual images, and effectiveness measured the mean absolute error (MAE) between the predicted and target attribute values for continuous attributes, and classification accuracy for the digit label as shown in Table
Composition. The composition metric evaluates the stability of the encoding and then decoding under repeated null counterfactual operations (no intervention). For each test image, the cycle of latent space encoding, passing through the MACAW counterfactual mechanism with no attribute changes, decoding was repeated for K cycles. The distance between the original and cycled images were measured at cycles K (1, 10).
20

---

Published in Transactions on Machine Learning Research (03/2026)
Eifectiveness: single-attribute interventions
dec,
de0,
t=1. 1696
1-:88
in 1-787
1=122
1-230
1=173
1=127
9
8
3
9
7
8
3
1-.4
T-23
1028
126
1:26
1-3.7
9
8
3
-113
1-1
1:200
1-12
1::66
6
4
9
7
Figure 13: The figure shows the outcomes of counterfactual queries, where causal factors of selected test images (Factual) were altered to generate new images. Every row represents distinct counterfactual query, identified as "do(.)" commands.
Realism. The Fréchet Inception Distance (FID) was computed between 10,000 counterfactual images (generated by intervening on thickness with random targets) and the real test images.
B.0.4 Results
Table Effectiveness on the MorphoMNIST test set. For each single-attribute intervention, all three attributes are measured: the directly targeted attribute and the preservation of non-intervened attributes. Thickness and intensity are reported MAE (4), digit classification accuracy (t). Values for HVAE,
VAE, and GAN baselines are taken from Melistas et al. (2024).
MACAW
HVAE
VAE
GAN
do(t)
Thick. MAE
0.179
0.086
0.109
0.228
do(i) Intens. MAE 4
15.39
3.52
5.33
15.14
do(y)
Digit Acc. f
0.972
0.775
0.451
Table presents the effectiveness results. For the direct thickness intervention do(t), MACAW achieved a thickness MAE of 0.179, compared 0.086 for the H-VAE baseline. Under digit intervention do(y), the model achieved classification accuracy of 0.93 while proserving the morphological attributes.
Table 4: Composition, realism, and minimality on MorphoMNIST.. Values for HVAE, VAE, and GAN
baselines taken from Melistas et al. (2024).
MACAW
HVAE
VAE
GAN
Composition (K 1)) Composition & (K= 10) 4
0.0432
0.438
2,600
3.807
0.1275
1.550
7.698
11.697
FID 4
16.66
5.362
10.124
35.568
Table reports the remaining metrics. The composition scores indicate the stability of the encode-decode cycle, where lower values correspond to better image preservation under repeated cycling. Interestingly, the MACAW model outperformed the three baseline models, likely due to the invertible nature of the proposed network and the simple latent encoding.
It should be noted that the bascline methods (HVAE, VAE, GAN) in the benchmark employ deep genera- tive architectures specifically designed for image generation, with each model containing dedicated image
21

---

Published in Transactions on Machine Learning Research (03/2026)
decoder that directly maps latent representations to pixel space. In contrast, MACAW operates
pressed representation (d=20), with the normalizing flow modeling the causal relationships in this latent space. Despite this architectural difference, MACAW demonstrates competitive counterfactual generation quality, particularly in attribute preservation under single-variable interventions.
22


Reviews:
[{"review_id": "Ny6FZJGwBV", "reviewer": "Reviewer_MKFc", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces a novel causal generative model for medical imaging that incorporates complex causal structures into normalizing flows. This design enables counterfactual prediction and provides an explicit Bayesian inference mechanism for classification.\n\nThe model is evaluated on both a synthetic data and a  real clinical scenario. The counterfactual predictions demonstrated in the neuroimaging experiments are particularly compelling and suggest strong potential for impactful clinical applications.", "claims_and_evidence_yn": "No", "claims_explanation": "Although the empirical goal of the proposed algorithm is interesting, there are several concerns regarding the soundness of the approach.\n\n1. While the authors claim that the complex causal graph is encoded within the flow model, the encoding appears to only support information flow from parent nodes to their child. In particular, there is no mechanism for information to propagate from child nodes back to parent nodes. This restricted formulation requires further clarification. As a result, the causal graph effectively functions as a preprocessing step for removing bidirected edges caused by confounders, yielding a simplified DAG-constrained model. This departs from prior work that does not impose such directional restrictions, and the implications of this restriction should be more thoroughly discussed.\n\n2. It remains unclear whether similar performance could be achieved without explicitly incorporating causality, particularly if correlations in the data are sufficient to drive the observed results. Without a comparative ablation or justification, it is difficult to determine whether the improvements arise from causal reasoning or simply from exploiting statistical associations.", "audience_yn": "Yes", "audience_explanation": "The task is challenging, and focusing on causality rather than correlation makes the approach practically useful.", "requested_changes": "1. An ablation study comparing performance under a fully connected causal graph versus the correct, sparser causal graph would be valuable.\n\n2. Including additional relevant baselines would strengthen the empirical evaluation.\n\n3. It would be helpful to clarify how the number of latent variables influences model performance. Specifically, how a user should determine the appropriate number of latent variables in practice.\n\n4. Is there any theoretical guarantee for the resulting counterfactual predictions or for the predictions of the parent variables?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "2E9gYc1b7P", "reviewer": "Reviewer_Ry1z", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposed MACAW, a flow-based causal generative model that injects a known DAG into the architecture via C-MADE (causally-masked autoencoders). In particular, each layer’s connectivity is masked according to the adjacency matrix so that each variable depends only on its parents. Stacking C-MADEs yields a Masked Causal Flow enabling exact density modeling, interventional and counterfactual sampling, and downstream Bayesian classification. For high-dimensional image data, the method first projects slices with KPCA and learn the flow in the reduced space, using a preimage map to visualize counterfactuals. Experiments on one synthetic dataset and UK Biobank T1 MRI (2D slices) evaluate the generated counterfactuals with FID and a discriminative “effectiveness” proxy.", "claims_and_evidence_yn": "No", "claims_explanation": "**Strengths:**\n1. The masking network encodes the DAG directly in the conditioner, preserving triangular Jacobians and efficient likelihoods.\n\n2. Compared to previous works, a single normalizing flow handles the whole graph (vs one flow per node), improving scalability in principle.\n\n3. The proposed method enables full causal reasoning, including interventional and counterfactual queries.\n\n**Weaknesses:**\n1. Some notations are not clear. For example, in $T(\\boldsymbol x^{obs})$ and $T(\\alpha, \\boldsymbol x^{obs}\\_{\\pi(j)})$, the dimension of the input space varies when $j$ is not the leaf node because $\\pi(j)$ only contains the parent set. Shouldn't it be noted as the restriction of $T$ on $(\\boldsymbol x^{obs}\\_{\\pi(j)}, \\boldsymbol x_j)$? For the counterfactual action, the author used e.g. $do(x_2=2)$, shouldn't it be $do(x_2=2|\\boldsymbol x = x^{obs})$?\n\n2. The experiment on the synthetic dataset needs more extensions. In Table 1, the author only compared the mean and variance of the variables generated by the ground-truth model and the MACAW. However, since these variables are not Gaussian, these two moments may not be sufficient. I suggest the author also include, e.g., the Maximum Mean Discrepancy (MMD) metric and some other baseline methods to validate the counterfactual prediction more rigorously. In addition, it seems that the author assumed the distribution class of the latent variables is known. I wonder whether the method is robust if the distribution of the latent variables is misspecified?\n\n3. The main experiment lacks baseline methods to compare. The author only evaluates their method without including related works, e.g. Deep-SCM and Diff-SCM. In addition, some other metrics, such as counterfactual latent divergence (CLD), IM1, IM2 can be included or discussed. \n\n4. The evaluation is merely on one toy example and UKBB. \n\n5. There are no experiments on the run-time, while the author seems to claim that scalability is the strength of MACAW. For example, the training/inference time on UKBB, how the speed and performance are sensitive to the choice of the hidden layers.", "audience_yn": "Yes", "audience_explanation": "The paper proposed a novel causal generative model that has wide applications.", "requested_changes": "1. Clear notation.\n2. Sufficient baseline methods, benchmark datasets, and evaluation metrics.\n3. Run-time comparison.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "OGH0woCB3q", "reviewer": "Reviewer_ZpQ3", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces MACAW, a causal generative framework that integrates a single normalizing flow with masked autoencoders to capture causal structures in medical imaging data. The model can recover causal pathways, generate counterfactual images, and perform Bayesian classification within one unified framework. Results on synthetic and UK Biobank MRI data show biologically plausible counterfactuals and reasonable predictive performance.\n\n\n\n$\\textbf{Strengths:}$\n\n\n1. Novel integration of causal graphs and normalizing flows for image-level counterfactual generation.\n\n2. The approach supports causal reasoning, interventional sampling, and classification simultaneously.\n\n3. Demonstrations on a large real-world neuroimaging dataset highlight practical potential.\n\n\n\n$\\textbf{Weaknesses}$\n\n1. The causal structure is predetermined, not learned, limiting flexibility.\n\n2. Evaluation is limited, robustness under noise or with more variables should be tested.\n\n3. Missing comparisons with other competing causal generative methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Overall, the paper’s main claims are reasonably supported, though the evidence is partial rather than comprehensive. To make the evaluation more convincing regarding the method’s effectiveness, the authors should test the approach on more complex and diverse tasks within the UK Biobank MRI dataset or other publicly available medical image tasks.", "audience_yn": "Yes", "audience_explanation": "Yes, the paper addresses an important intersection between causal generative modeling and medical imaging, with clear relevance to explainable deep learning in the medical domain.", "requested_changes": "The paper is overall interesting and relevant, introducing a novel perspective on causal generative modeling in medical imaging. However, the experimental section could be strengthened to make the claims more convincing and to better demonstrate the robustness and practical value of the proposed method. I suggest the following improvements:\n\n1.Robustness to causal structure assumptions:\n\nIt would be informative to test how the model behaves when the assumed causal graph is partially incorrect or perturbed. In real-world medical scenarios, causal relationships among demographic and clinical variables are often uncertain. Evaluating MACAW’s sensitivity to such structural disruptions would provide valuable insights into its robustness.\n\n2.Inclusion of additional or irrelevant variables:\n\nThe current experimental setup uses only three causal factors (age, sex, BMI), which may be too simple to reflect practical conditions. Introducing more variables—some relevant, others intentionally irrelevant or latent features extracted from images—could help assess how the model scales and whether it can ignore spurious inputs.\n\n3.Comparison with competing causal generative methods:\n\nA quantitative comparison with other causal or conditional generative models (e.g., VAE-based approaches) would clarify the advantages of MACAW’s masking and flow architecture. Using similar latent representations would make such a comparison fair and highlight the specific contributions of the proposed structure.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)