CaTs and DAGs: Integrating Directed Acyclic Graphs with Transformers for Causally Constrained Predictions
Authors: Matthew James Vowels
Mathieu Rochat
Sina Akbari
Venue: ICLR 2026 Poster (2026)
Keywords: transformers
causal inference
causality
inductive bias
DAGs
Primary area: causal reasoning
TL;DR: Causal Transformers (CaTs) are neural networks constrained by a causal DAG, combining the power of standard ANNs with improved robustness to covariate shift, greater reliability, and interpretability for real-world applications.

Abstract: Artificial Neural Networks (ANNs), including fully-connected networks and transformers, are highly flexible and powerful function approximators, widely applied in fields like computer vision and natural language processing. However, their inability to inherently respect causal structures can limit their robustness, making them vulnerable to covariate shift and difficult to interpret/explain. This poses significant challenges for their reliability in real-world applications. In this paper, we introduce Causal Transformers (CaTs), a general model class designed to operate under predefined causal constraints, as specified by a Directed Acyclic Graph (DAG). CaTs retain the powerful function approximation abilities of traditional neural networks while adhering to the underlying structural constraints, improving robustness, reliability, and interpretability at inference time. This approach opens new avenues for deploying neural networks in more demanding, real-world scenarios where robustness and explainability is critical.

Published as a conference paper at ICLR 2026
CATS AND DAGS: INTEGRATING DIRECTED AcyclIC GRAPHS WITH TRANSFORMERS FOR CAUSALLY CON- STRAINED PREDICTIONS
Matthew J. Vowels
Kivira Health.
The Sense Innovation and Research Center, CHUV, Lausanne, Switzerland. Centre for Vision, Speech and Signal Processing, University of Surrey, U.K. matt ielkkiiiiiee vivvaaa.v. health
Mathieu Rochat
Institute of Psychology, University of Lausanne, Switzerland. mathieu. rechat@unil.ccm
Sina Akbari Statistical Laboratory, University of Cambridge, U.K.
sa23850cass ac.uk
ABSTRACT
Artificial Neural Networks (ANNs), including fully-connected networks and transformers, are highly flexible and powerful function approximators, widely ap- plied in fields like computer vision and natural language processing. However, their inability to inherently respect causal structures can limit their robustness, making them vulnerable to covariate shift and difficult interpret/explain. This poses significant challenges for their reliability in real-world applications. In this paper, we introduce Causal Transformers (CaTs), general model class designed to operate under predefined causal constraints, as specified by Directed Acyclic Graph (DAG), CaTs retain the powerful function approximation abilities of tra- ditional neural networks while adhering the underlying structural constraints, improving robustness, reliability, and interpretability at inference time. This ap- proach opens avenues for deploying neural networks in more demanding, real-world scenarios where robustness and explainability is critical.
INTRODUCTION
Machine learning has long focused on discovering powerful function approximation techniques for mapping inputs outputs. Two of the most widely used approaches are the fully-connected neural network and the transformer (Vaswani et al., 2017), both of which are Artificial Neural Networks (ANNs). These methods have achieved broad success in computer vision (Dosovitskiy al. 202T), audio (Radford et al., 2023). reinforcement learning (Vinyals et al. 2019). natural language pro- cessing (Devlin al., 2019 Brown al.) 2020), and other domains (J. et al., 202T). Despite their efficacy, ANNs rely primarily on statistical, data-driven mechanisms and do not typically incorpo- rate prior knowledge about the underlying Data Generating Process (DGP).
Models that fail to respect the DGP often show sensitivity to covariate shift (Arjovsky et al.) 2020 Rojas-Carulla et al., 2018 de Haan et al. 2019). They may exploit non-causal associations in train- ing data relationships that lack mechanistic meaning and underperform when such correlations change or are irrelevant to the outcome (see Fig. [ A well-known example classifying camels versus cows: training data often confound animals with their backgrounds (camels in sandy regions, cows in green pastures) (Arjovsky et al.] 2020). non-causal model may rely on the background rather than the animal itself, leading enrors such misclassifying camel in grassy field. In

---

Published as a conference paper at ICLR 2026
6
CFCN
CAT
5
Random Forest
MLP
Test MSF
Transformer
3
2
1
Shift
Figure 1: Illustrating the lack of covariate shift robustness associated with conventional machine learning models like random forests (Breiman, 200T), multilayer perceptrons (MLPs), and trans- formers /Vaswani et al. 2017), compared with our CaTl and CFCN.
general, non-causal ANNs entangle distinct factors, reducing robustness and making predictions highly sensitive irrelevant features.
A central goal of science, however, is to understand causality (Pearl 2009 Herman 2018 Rohrer 2018 Schookkopf 2019 Peters et al. 2017). Causal inference seeks estimate the effect of an intervention (c.g., taking a drug) even in the absence of randomized data, task made difficult by confounding. For instance, comparing outcomes between patients who select surgery versus medication is biased by age older patients are more likely to choose medication and less likely to recover. Without causal structure, models typically yield biased effect estimates (Cadei et al.) 2024 Vowels 2023a Pearl (009). To isolate true causal effects from spurious associations, explicit structural constraints are required, beyond what "purely statistical" approaches in machine learning provide.
Our main contribution a modification to the transformer architecture that constrains interactions according to Directed Acyclic Graph (DAG) representing causal structure. We call this the Causal Transformer (CaT). For comparison, we also extend the masking idea from MADE (Germain et al, 2013) to create DAG-constrained fully connected networks, which we term Causal Fully Connected Networks (CFCNs). Such constraints reduce sensitivity to covariate shift and broaden applicability causal domains including medical imaging (Castro al. 2019), policy evaluation (Kreif & DiazOrdaz 2019), medicine (Petersen et al. 2017 Deaton & Cartwright 2018), advertising (Bottou
al. 2013). and social science (Grosz
2020 Vowels 2023a
Rohrer 2018).
In summary, CaTH extends transformers to respect causal knowledge, enabling their application across empirical domains. CaTs can operate on inputs of arbitrary dimensionality, from tabular variables to high-ding sighddimensional multimodal embeddings. The goal of this work is not to provide state-of-the-are performance on specific benchmarks, to propose general modeling framework that integrates structural inductive biases into a popular neural architecture. This integration en- hances interpretability, fairness (Locatello al. 2019a), and data efficiency (Vowels) 2023b), while promoting robustness to distributional changes (Arjovsky et al. 2020). These models are intended as foundations for further innovation, allowing researchers to adapt and extend them for domain- specific applications where causal structure is essential.
RELATED WORK
The design of effective machine learning models often relies on inductive bias. Such bias can be introduced in two main ways: (a) A structural bias constrains possible interactions between variables, dimensions, or inputs, for example to disentangle generative factors (Locatello et al.)
2019h Higgins et al. 2017 Burgess al. 2018 Yang at. 2020 Bengio al. 2019) estimate
Code for CaT's and the experiments can be found https::/gittb..com//eea.
Causal Former
2

---

Published as a conference paper at ICLR 2026
the effect of an intervention (Shalit et al. [017 Louizos et al. 2017b). (b) A functional bias restricts the class of functions that can be learned. such as convolution for vision tasks (LeCun & Bengio,
1998). regularization for smoothness (Goyal & Bengio 2022). translation equivariance (Fuchs et al. 2020). Functional biases also project inputs into specialized spaces, e.g. for registration in medical images (Friedman et al. 2023) or DAG-based emotion prediction (Shen et al. 2021).
Our work falls under the first category. Structural inductive bias has been integrated into neural models for causal inference via multi-task deep networks (Xia et al.) 2021; Kocaoglu et al. 2018
Shalit et al., 2017, Shh et al., 2019), Variational Autoencoders (Kingma & Welling 2014 Louizos
al. 2017a| Zhang et al. 2021: Vowels et al., 202Ib), and variants with adversarial ( (Yon et al,, 2018 or semi- parameattt objectives (Shi et al.) [2019] Vowels et al.] 2023b). Beyond neural net-
works, Structural Equation Modeling (Kline, 2005) and non-linear extensions (Kelava et al., 2014 Umbach et al. 2017 van Kesteren & Oberski, (022) remain popular for causal effect estimation, with the added beneñt of statistical inference. Unbiased estimators can also be obtained by com- bining predictive models such MLPs or Random Forests (Breiman 2001) with careful predictor selection based on the backdoor criterion (Pearl 2009).
Most causal inference methods target tabular data, but structural bias has also been applied in vision tasks (Vowels et al., 2021a Krishnan et al.) 2015; Chen et al., 2019 Li & Mandt, (0188). where knowledge of changing versus static factors (c.g. pose identity) helps disentanglement. Our
CaT approach can similarly exploit multidimensional embeddings. Indeed, (Vafa et al. 2024) show that missing inductive biases leave generative transformers fragile, highlighting the need for explicit constraints.
Attempts make transformers "causal" have often focused on autoregressive masking, which re- stricts attention to past tokens (Vaswani et al. 2017 Yin et al. 2024 Yenduri et al. (024). This con- straint resembles Granger causality (Granger 1980) and remains weak: it cannot enforce high-level disentanglement, causal inference, or robustness to covariate shift. More targeted work includes (Melnychuk et al.] 2022). who use separate transformers for treatment, covariates, and outcomes, but this requires grouping variables by estimand. By contrast, CaT applies to any causal effect implied by a DAG without subnetworks.
Other
combinations pursue different goals. (Luo et al.) 2023) allow bidirectional links (violating causal ordering) for outcome prediction, while (Gagrani et al. 2022) predict topo- logical ordering itself. In general, while functional inductive biases (e.g. convolution, weight decay) are common, neural networks rarely enforce structural constraints. Designing such models is chal- lenging because ANNs usually facilitate maximal, not restricted, interactions.
The most comparable work to our baseline (CFCN) includes MADE (Germain et al. 2015), Graph- ical Normalizing Flows (Wehenkel & Louppe 2021 and Structured Neural Networks (Chen al. 2023). MADE masks weights so each output depends only on earlier inputs in an autoregressive order. For example, with variables A, B, C in an autoregressive structure, masks enforce C1 B, A), AIA, and Ålø. Wehenkel & Louppe (2021) combine normalizing flows with DAGs, while Chen al. (2023) express adjacency as a binary factorization optimized across layers. In contrast, our approach applies masks directly from the DAG and adapts them to layer shapes, yielding a simple and comparable architectural baseline.
3 METHODS
3.1 PROBLEM SETUP AND BACKGROUND
In this work we address two main problem setups. We consider dataset D ([Y,, D,, 1l, whereby N is the sample size, and treatment D, outcome Y. and covariates Lget L' are multidimensional embeddings (possibly of different dimensionality). CaT is concerned with this setup, but also functions in the case where all variables are univariate. CFCN, our architectural baseline, functions in the univariate case. We make no assumptions about the para- metric or distributional form of each uni- or size, |Z| to indicate the number of nodes in the DAG, and, C to indicate the feature dimensionality for each variable.
variable. We use B to indicate batch
3

---

Published as a conference paper at ICLR 2026
Input
X
Input Stage
X
CaT
Causal Black
D
FC Layer
Layer
Layer
Cancal Head
Y
Layer
RClayer
Canool Cross-Aeeriio
x 12)
D
QK
IF 5.e
Optici
A
softmmooo
1 Y 1 0
D
(8))
Layer
12 [2] x
Inpert
D LLL
skip
Output Targrts Reconstructions
12 x
Hasi's Blocks
Figure 2: A top-level depiction of how the masking and routing is applied in CaT to an input with a batch B of |Z| = C-dimensional embeddings and a corresponding causal DAG. The input parameter y, which is initially random, is recursively fed an input to the causal Heads, "extracting" information from the embeddings of the input via the causal cross-attention operation, according to the constraints imposed by the adjacency matrix A. Best viewed in color. See main text for details.
both scenarios we are concerned with building a model encoding set of conditional indepen- dencies encoded by Directed Acyclic Graph (DAG) G. The DAG represents the factorization of a joint distribution P(Y, D, Lset) of all variables using Z| corresponding nodes & E and edges (i, k) E E, where (j, k) indicates the presence of directed edge from node to node k. The acyclicity indicates the absence of any cyclic paths (such that there are no directed paths hetween a popular tool for representing causal structures or DGPs. For instance, the DAG D - M + Y tells us that D causes variable M which itself is a mediator between D and Y where here would be the outcome of the causal chain. DAGs can be represented adjacency matrix. With each row representing a possible parent (i.e. direct causal predecessor), and each column representing a possible child (i.e., variable which is affected by its parent). The adjacency matrix representation will be a key component in enforcing the causal constraints in our CFCNs and CaTs.
node and itself). DAGs
square
The end goal is to be able to reason causally, under the constraints established by the DAG. One associated task is the estimation of the average treatment effect T, which can be expressed as follows:
= = Py Y
E P [Y]
(1)
-EL >p
E EY-P (do)p((D))D)))))))
Y
E PYdd(D T.
Y
Here, the operator (Pearl 2009) denotes an intervention, simulating the case whereby = 0 and D 1, even though in the joint distribution P, D takes on values according to the underlying DGP. This quantity is therefore hypothetical, but can be estimated unbiasedly under the right conditions, and these conditions can usually be derived from the DAG itself (Pearl 2009 Peters et al.) 2017). In practice, causal discovery techniques, domain knowledge and inductive biases can be used for the specification of a DAG, although in this work we assume the DAG is known a priori. One of the key advantages of models that respect the underlying causal structure (or, more loosely, respect some underlying properties of the DGP) is that predictions are less sensitive distributional shift. Conventional machine learning models which leverage all available correlations make predictions which fluctuate according to changes in spurious or confounding factors (Arjovsky al., 2020) Vowels 2022). Causal models, on the other hand, may have less predictive power under a stable joint distribution, but maintain this predictive power under covariate shift precisely because they leverage only those relationships which are invariant under such shift.
Importantly, we note that the success of causal inference rests on a number of strong and often untestable assumptions known as ignorability, positivity, and the stable unit treatment value as- sumption, each of which we describe further in Sec S8 Interested readers are directed to surveys and introductions by (Vowels et al. 2022 Vowels 2024 2018 Peters et al.; Scholkopt 2019 for more information on causality, and Sec. SS for further preliminaries.
Glymour et al. 2019 Pearl & Mackenzie,
4

---

Published as a conference paper at ICLR 2026
Note that if relax the requirement for unbiased estimates for causal effects, can use the DAG to loosely constrain the model according to relevant factors and to make predictions for arbitrary interventions (more this below). As such, CaTs are not only useful for causal inference, but also as models which are robust to covariate shift in proportion to the strength of the constraints imposed via the DAG, and the coherence between this DAG and the true, underlying DGP.
3.2 CAUSAL TRANSFORMERS (CATs)
causal approaches (Vowels et al. 202Ib; Shi al. 2019 Louizos et al. 2017a] Shalit et al. 2017 Yoon et al., 2018) and related methods (van der Laan et al., 2007 Wager & Athey, 2018 Vowels, 2023c) is that they treat each variable as single value, restrict- ing inference to a uni diledimmnnional setting. This prevents direct use of high-dimensional embeddings such as skeleton joint sets, voice embeddings, or multimodal features from images and text. con- trast, our mransoormerbbaeed approach views the input as sequence of dddimensional embeddings, making it applicable to both tabular data (with scalar variables) and complex multimodal data. This property particularly useful in applied fields like psychology and medicine, where researchers often rely on multi-item scales or multidimensional outcomes (Hilpert et al.] [019] [Bulling et al.) 2023). For example, the NEO Personality Inventory comprises 240 items (Costa & McCrae (2008). typically averaged into five trait scores. With CaT's, each subscale can instead be represented as a vector of items, preserving information that would otherwise be lost.
One key limitation of existing
For a brief review of self- and cross-attention mechanisms in typical transformers, please see Sec. S5 A top-level depiction of the core change to the transformer attention architecture is shown in Fig. The complete process for the single causally-masked cross-attention head of CaT is given below. We provide proof that the constraints implied by the DAG are retained through the proposed architecture in S13
First, we embed the input X from shape (B x [Z] x C") into Xg with dimensionality (B x | dg). Crucially, this is done using |Z] independent linear layers. It is also important that dg is at least as large as C' and greater than In practice, we find that even if C dr, but de = 1, the network struggles to disentangle the contributions from the separate variables for all but trivial cases. This stage can also be used to unify the dimensionality of all input embeddings the case where the different variables do not have the same dimensionality (this may be the case with multimodal data, for instance).
Next, we randomly initialize learnable embedding 7, which is of size [Z] x dg. This parameter combined using cross-attention, with the embedded input Xg throughout the network and ac- cording to the interactions implied by the DAG. For cross- noss-attention, the keys, queries and values computed as follows:
K = X W K + bK:
Q = >We Q + bQ:
(2)
V = Wy + by.
Importantly, note that whilst the keys and values are functions of the embedded input Xg, the query function of the learnable embedding 7 in the first block, and a function of the output of the previous block for subsequent blocks (see Eq. 6p.
We then mask the attention between Q (which is a function of 7) and KT (which is a function of Xg). The mask is applied via Hadamard produet (indicated by o) between the transposed, topo- logically sorted, adjacency matrix A and the attention matrix. We perform the softmax operation, and matrix multiply this with V:
QK Vhg ) V.
0 softmax(A A
Importantly, and in contrast with the CFCN, we do not modify A to include self-connections / an identity diagonal after the first layer/block (see below). The self-connections are not needed in CaT. Overall, the process presented in Equation is undertaken in parallel according to the number of heads in the Causal Multi- Muad Cross-Attention (CMCA), and these outputs are concatenated whilst maintaining the independence between the input variables:
5

---

Published as a conference paper at ICLR 2026
CMCA CMCA(7,, Xg) o (Dropoutt(roo((onn((( (1, XE); 0² 0 Xg),
(4)
O " (>. x)))))),
where Conc a concatenation operation (concatenating along the last dimension corresponding with the dimensionality of each head), H is the number of heads in the CMCA, and Proj is linear projection layer which projects the output to the same dimensionality as the expected head input.
The CMCA forms part of causal block ('CBlock'), which includes residual/skip connections:
O CMCA
(5)
BN(CMCA(O X8) + O Block
Block
O Block
(6)
ANCFF(OCMCA) O CMCA):
where BN is (optional) batch normalization (loffe & Szegedy 2013), and FF is a two-layer percep- tron with dropout and Swish (Ramachandran al. (2018) activation function. At each layer, the embedded input X5 is always fed to the hsad(s), whilst the learned embedding y is used only for the head(s) in the first layer, i.c. Block y for 0. Finally, we stack number of such blocks in sequence (corresponding with the number of 'layers') and include one last linear layer before the final output.
Summary of Key Modifications: The key modifications to the transformer architecture are:
The application of the cross-attention mask to constrain the network to respect the under- lying DGP / DAG.
The initialization of the learned embedding which forms the query vector.
Feeding Xg at every layer and every block so that the downstream 7 (which becomes Of D can attend to both its prior state and the observed embedded input. This allows the model account for the surprise in the observed value compared to the current estimate. Usually, cross-attention is between two different but informative embeddings. In our case, it is between a source information Xg and an empty embedding which becomes the output it propagates through the network. Feeding in the input at each layer enables CaT to "compare' the current O Block against the input to see what it needs to attend to in the input (whilst constraint according to the DAG) to produce a useful output.
We do not include layer normalization [a et al. 2016). Despite its successful integration with transformers, layer normalization is not used as it would rescale potentially calibrated interventions.
Model
MSE (+ Std)
eATE (+ Std)
CaT False DAG 1
0.555 (# 0.003)
2,314 (+ 0.010)
CaT False DAG 2
0.557 (# 0.004)
1.382 (+ 0.147)
CaT Truc DAG
1.004 (# 0.007)
0.058 (# 0.056)
CFCN False DAG
0.622 0.200)
2.330 0.056)
CFCN False DAG
0.568 (# 0.008)
1.335 0.047)
CFCN True DAG
1.023 (# 0.014)
0.149 (# 0.309)
Rand. Forest
0.622 (= 0.004)
2.311 (# 0.012)
Transformer
0.615 0.012)
2.379 0.057)
MLP
0.559 0.006)
2.325 0.013)
Table 1: Simulation results for motivating example highlighting that models which make correct structural assumptions do necessarily yield predictions with the lowest Mean Squared Error on the test set (often incorrect causal assumptions result in higher predictive capacity, but lower robustness), and that lowest absolute error on the estimation of the Average Treatment Effect (eATE) can only be achieved with the correct causal graph.

---

Published as a conference paper at ICLR 2026
3.3 CAUSAL FULLY- CONLE-CONNECTED NETWORKS (CFCNs)
The CFCN is constructed similarly to an autoencoder, but only insofar as the inputs are also the targets (see Fig. 5). Targets cannot be predicted from themselves, but only by their structural/causal parents. The output is therefore reflective of the conditional distribution implied by the DAG. For each layer r, where 4. Pr. where pr- is the number of neurons in the layer before layer r, and Φr is the number of neurons in the current layer. We assert that the number of neurons in each layer is never less than the dimensionality of the input to avoid bottleneck issues resulting in dropout of signal along the depth of the network. Furthermore, we assume that the width of each layer is an integer multiple or factor of the preceding layer size, such that when growing or shrinking the network across its depth, the number of neurons capable of interacting with a particular input is always equal across inputs and always greater than equal 1. Taking the adjacency matrix we reshape/expand/hhire it to reflect the shape of the underlying weight matrices for each layer, whilst respecting the original dependencies encoded by the associated DAG. The masking process is applied similarly to (Germain iet al. 2015):
the total number of layers, we pre-allocate a mask M, of zeros of shape
=0(r-1 M,) + b,),
where a is an activation function, O, is the vector output, W, is the weight matrix, b, is a bias vector, and M, is a mask for given layer r, respectively. is the Hadamard product. An example is shown in Sec. [3 and Fig. 4 in the supp. mat. One key element illustrated with two examples in Fig. is the introduction of the identity diagonal to the DAG for all layers greater than 1. As such, the first layer acts as the 'barrier' to prevent inputs from attending to themselves, but once an intermediate (and causally valid) interaction has occurred, any interactions can then be passed on to the output. The identity cannot be introduced into the first layer, because this would enable the network to simply pass through the input to the output, and also violate the fact that variables without parents are exogenous. In contrast, without the introduction of the identity in the second layer, no signals could be propagated from intermediate, causally valid interactions. In Sec. SS we present CFCN 'mod', which is alternative we find to perform similarly.
3.4 TRAINING AND LOSS FUNCTIONS
We use per-variable loss functions, and users are required specify the variable type as continuous (mean-squaree error) or binary (binary cross-entropy). If embeddings are used, the variables are likely to be continuous, but for more specific tasks we provide the option to tailor the loss func- tion to the variable. Importantly, the loss functions can be determined independently of the causal constraint, and other loss functions (e.g. with custom regularizers) can be used as long as the in- dependence constraints implied by the DAG are not violated. For any endogenous node 2: if is continuous, we use mean squared error on its output head; if binary, Bernoulli cross-entropy; cate- gorical, softmax cross-entropy; if vector embedding, MSE on that vector. The experimenter is thus free to implement the loss they wish, understanding that the predictions are constrained according to the assumed DAG. A strength of CaT is that there are no special losses or regularizers.
3.5 ESTIMANDS, INTERVENTIONS, AND RECURSIVE INFERENCE
Constraining models with DAGS not only improves robustness covariate shift but also enables estimation of intervention effects. Consider the mediating graph D + M + Y (D is treatment, Y an outcome, and M a mechanism), together with D <- L + Y where L (e.g.. age) affects both treatment choice and recovery. The latter paths confound the estimation of D's effect on Y. Using do-calculus (Pearl) 2009), the interventional distribution Y do(D 1) can be identified as Y D = 1, T., since conditioning on I, removes the spurious L + D path. Importantly. M must be excluded: conditioning on M would make 1) redundant because M is more proximal cause. This is case of transitive reduction (Richardson al.) 2017]] Verma & Pcarl 1990; Vowels 2023b), where intermediate nodes can be dropped if they do not affect relevant conditional independencies. Thus the chain D MM- + Y simplifies to D > Y for estimating the total effect.
CaT supports all conditional average (CATE) and averagage (ATE) causal effect queries identifiable from the DAG via standard identification, operationalized as follows: set intervened variables to their intervention values, then recursively update descendants in topological order by re-evaluating the corresponding node heads with parents fixed to their current values, following the g-formula
7

---

Published as a conference paper at ICLR 2026
Robins 1986). Using the DAG, this algorithm iteratively updates downstream children of the inter- vention variables to propagate effects. Full details are given in Sec. [S6and and Alg. in supp. mat.
EXPERIMENTS
We provide results for five experiments: (a) simulated dataset motivating the need for CaT', (b)
causal inference benchmark datasets Jobs (oalo 1986, Smith & Todd 2005). Twins (Almond
ct al. 2005) and (c) a real-world application of CaT to psychology data during the COVID-19
pandemic (McBride et al., (02T] Vowels et al. 2023a highlighting its flexibilityy in leveraging multi- dimensional inputs. Note that owing to space constraints some of these details concerning these evaluations are presented in Sections [STO] and [SIT in supp. mat. No hyperparameter tuning was undertaken for these experiments again, the goal is not to provide a model which excels at particular task, but provide general modeling class which facilitates the integration of structural constraints across a range of tasks. We include CFCN as a baseline as it isolates the value of structural masking but without attention. Using the same DAG mask with layered identity handling, it serves as an architectural control to show that benefits are not unique to the attention mechanism of the transformer.
Twins
Jobs
Model
cATE ws
eATE os
Risk ws
Risk os
eAIT ws
eATT os
0110 # 0061
.0133 + 0078
.25 + 02
.7 00
.016 01
.86 06
CFCN
.098 0048
.0022
.25
26
06
GANITE
Your et al.
2018
0058 + 0017
.089 a. 0075
.13 01
.14 00
.01 = :.01
.06 : 03
CFRwass
Shalli el
2017
00112 0016
.0844 0012
.17
21
04+
.9
BART
Chipaan et di 2010
1206 + 0236
.1265 0234
23 00
25 02
02
08 + 00
C Forest
wager Allyy
2018
0286 + 0035
.0335 2 .0.3
.19 00
20 02
.31
= 00
CEVAR
20170
.15
.26 00
.02+.00 00
.0
TVAB
Vowels al
20216)
.66 2 00
.162.00
.01 # 00
.01 = 00
Table 2: Results + standard deviations for the Twins and Jobs benchmark datasets. eATE absolute
error on the estimation of the Average Treatment Effect; Risk Policy risk; eATT absolute error on the estimation of the Average Treatment effect on the Treated; ws within sample; os out of sample.
Toy Motivating Example: We illustrate the advantage of causal constraints with simple example adapted from (Hilpert & Vowels, 2024). Data are generated according to the graph in Fig. 3a The task estimate the total effect of D on Y, and to evaluate both the absolute error on this estimate (eATE) and the Mean Squared Error (MSE) of predicting Y. We compare three conditions: the true graph (Fig. Ba), a graph reflecting typical machine learning assumptions (Fig. 3bb,, and graph with a missing and a reversed edge (Fig. (c). CFCN and CaT are evaluated alongside a standard
transformer (Vaswani et al., 2017 Paszke et al. (017). random forest (Breiman 2001] Pedregosal
al. 2011), and multilayer perceptron. Results are shown in Table D and Fig. No hyperparameter tuning was applied to any model.
Three observations follow. First, CaT and CFCN achieve the lowest eATE under the true graph, as expected when respecting the DGP. Second, with the correct causal graph, models also attain the lowest MSE, though in general non-causal predictors can sometimes yield better accuracy than faithful causal structure. Notably, CaTT and CFCN perform comparably to other methods when the graph is misspecified. Third, Fig. [ssows that non-causal models, while potentially achieving lower MSE, are highly sensitive to covariate shift. In contrast, CaT and CFCN remain stable, reflecting robustness that extends beyond causal inference tasks.
Causal Inference Tasks: Results for the three benchmark datasets are shown in Table 2 As noted by recent authors (Vowels et al., 2023b Curth et al. 202T), such benchmarks have limitations, so our interpretations remain cautious (see also Sec. (S12). Evaluation metrics are detailed in Sec. [STT Lower scores indicate better performance. Since most comparison methods are designed and tuned specifically for causal inference, we do not expect CaT or CFCN to excel. Nevertheless, both achieve performance comparable to specialized methods. This result is notable given our minimal use of inductive bias. For the Twins and Jobs datasets, the only assumption is that all non- non-treatmeet/nnon outcome variables may be confounders. This graph is implicitly assumed by other methods, which treat all covariates equally. In contrast, TVAE, CFRWass, DML., and TL. incorporate additional
8

---

Published as a conference paper at ICLR 2026
D
L
D
D
1
Y
L
Y
L1
Y
L2
L
L
(a) True DAG (h) False DAG (c) False DAG 2
Figure 3: Motivating example. The data for the experiment are generated from DAG (a), whilst (b) and (c) represent misspecified DAGs used for estimation. Dashed curved lines in (b) emphasise possible endogeneity of D, L1 and L2, and the absence of independence in typical non-causal ma- chine learning approaches. The thick edges depict the target causal relationship of interest.
domain-specific techniques, giving them a natural advantage. By avoiding such tailoring, our ap- proach at an inherent disadvantage, yet CaT and CFCN remain competitive, underscoring their potential.
Real-World Psychology: We also conducted a secondary analysis of UK data from wave of the COVID-19 Psychological Research Consortium Study (McBride et al. (202)) collected during April/May 2020. After excluding cases with missing attachment data, 895 participants remained. The original study examined effects of attachment style compliance and mental health outcomes. Following (Vowels et al., 2023a). we estimated the causal effect of shifting attachment styles on depression, but using the full multidimensional set of questionnaire items. We compared naive bivariate linear model, a targeted learning (TL. estimator with semi-parametric techniques, and Cат.
The DAG from (Vowels al., 2023a), derived from causal discovery and domain expertise, was applied (see Sec. SII).. Competing methods had to operate on aggregated scale scores, whereas CaT used all questionnaire items. Using secure attachment as the reference, we estimated effects of fearful, anxious, and avoidant styles on depression. TL and the naive model yielded standard errors directly, while CaT's were obtained via 100 bootstrap iterations. Despite not being optimized for causal inference, CaT produced results broadly consistent with TL, without hyperparameter search. As with any real-world causal analysis, however, no ground truth exists, so these findings serve primarily demonstration of applicability.
LIMITATIONS & DISCUSSION
One principal limitation of CaT's concerns the recursive inference procedure. For DAGs which include long mediation structures, and when the intervention node(s) is(are) early in the causal ordering, the recursion can lead to compounding error as the prediction for each mediator is fed back into the model for the prediction of the next descendant. This issue can be alleviated with transitive reduction (Richardson et al., 2017 Verma & Pearl 1990, Vowels 2023b) to remove nodes which are not principally important for inference or downstream tasks. Additionally, in many application areas, it may not be possible to specify DAG because either the phenomenon is too complex, or because the DGP us unknown. However, such cases, users can either leverage causal discovery techniques (Vowels et al. 2022 Heinzee-Dee Deml et al.) 2018 Glymour et al. 2019: Spirtes & Zhang, 2016) and work with putative graph, and otherwise err on the side of caution by including more rather than fewer edges in the DAG (the removal of an edge represents stronger assumption than an inclusion). Furthermore, CaTs are flexible in that they can be used both for causal inference tasks requiring confident knowledge of the DGP, but equally for prediction, where the incorporation of inductive biases can be beneficial in building robust models.
CAT has yet to be scaled to and tested more complex computer vision tasks with pre-defined DAGs (such as autonomous driving or human-to-human interactions). Regardless, CaT is agnostic to the integration of various scaling techniques for efficient attention (incl. Flash Attention (Dao al. 2022 Dao 2024), Hedgehog & the Porcupine (Zhang et al.] 2024), or Performer (Choromanski
al. 2020

---

Published as a conference paper at ICLR 2026
Group
Naive
TL
CaT
Fearful
.13 1 .02
.05 1 .01
.01 .22
Anxious
.12 022
.05
002
.03
Avoidant
.01 01
.01 01
.03 04
Table 3: Bootstrapped effect sizes + standard emror for impact of change of attachment style on levels of depression. Compares naive (bivariate) correlation, Targeted Learning (TL), and CaT. The naive and TL based approaches use averages of the construuc//call/qquestionii CaT uses the full questionnaire without aggregation of questionnaire items.
items, whereas
CONCLUSION
Transformers powerful function approximators but lack robustness to covariate shift and cannot naturally encode structural inductive biases, limiting their reliability in scientific applications where robustness and interpretability are critical. We addressed this by introducing the Causal Transformer (CaT), which integrates domain knowledge via DAG-based constraints, and comparable baseline (CFCN). These models combine the flexibility of neural networks with explicit causal structure, while remaining agnostic to efficiency improvements such as Flash attention. Importantly, CaT and CFCN are complementary to methods like Double Machine Learning (DML.; (Chernozhukov al,, 2018 2017)) and Inverse Probability Weighting (IPW; (Rosenbaum & Rubin 1983)), and can serve as plug-in estimators. This makes them versatile addition to the empirical researcher's toolbox.
REFERENCES
D. Almond, K. Y. Chay, and D. S. Lee. The costs of low birth weight. The Quarterly Journal of Economics, 12033):1331-103) 2005.
M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz. Invariant risk minimization. arXiv: www.:19077.99995 2020.
Athey and G.W. Imbens. Recursive partitioning for heterogeneous causal effects. Proc. of the National Academy of Sciences of the USA, 113(27), 2016
J.L. Ba, J.R. Kiros, and G.E. Hinton. Layer normalization. arXiv preprint, arxiv:6007.0642 2016.
Y. Bengio, T. Deleu, N. Rahaman, N. R. Ke, S. Lachapelle, O. Bilaniuk, A. Goyal, and C. Pal. A meta-transfer objective for learning disentangle causal mechanisms. arXiv:1901.1011222. 2019.
I. Bottou, Peters, Quinonero-Candelll J., D. X. Charles, D. M. Chickering, E. Portugaly, D. Ray, P. Simard, and E. Snelson. Counterfactual reasoning and learning systems: the example of com- putational advertising. Journal of Machine Learning Research, 14, 2013.
L. Breiman. Random forests.
Machine Learning, 45(1):5-32, 2001. doi: 10.1023/A:
1010933404324.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., http:/aapeerttee firfffilesss wwt.fflle//pprr/2222222 File/ EL020/80006000001111333333222222
2020. URL
L. Bulling, R.E. Heyman, and G. Bodenmann. Bringing behavioral observation of couples into the st century. Journal of Family Psychology, 37(1):1-9, 2023. doi: 10. 103.
10

---

Published as a conference paper at ICLR 2026
C. P. Burgess, Higgins, A. Pal, L. Matthey, N. Watters, G. Desjardins, and A. Lerchner. Under- standing disentangling in Beta- VAE. arXiv:/804.0359vl1,2000
R. Cadei, L. Lindorfer, S. Cremer, C. Schmid, and F. Locatello. Smoke and mirrors in causal downstream tasks. arXiv preprint, wxxi:2405.171511 2024.
D. C. Castro, I. Walker, and B. Glocker. Causality matters in medical imaging. arXiv:1912.08142vv. 2019.
A.Q. Chen, R. Shi, X. Gao, R. Baptista, and R.G. Krishnan. Structured neural networks for density estimation and causal inference. 37th Conference on Neural Information Processing Systems (NeurIPS), 2023.
C. Chen, C. X. Lu, N. Trigoni, and A. Markham. DynaNet: Neural Kalman dynamical model for motion estimation and prediction. arXiv: 1wwp:/998../111 18v1 2019.
V. Chernozhukov, D. Chetverikov, M. Demirer, E. Duñoo C. Hansen, and W. Newey. Dou- ble/debiasee/NNeyman machine learning of treatment effects. American Economic Review, 5, 2017.
V. Chernozhukov, D. Chetverikov, M. Demirer, E. Duflo, C. Hansen, W. Newey, and J. Robins. Dou- ble/debiased machine learning for treatment and structural parameters. Econometries Journal, 21 (1):C1-C68, 2018. doi: 10.1 1 1/ectj//...// 12097.
H.A. Chipman, E.I. George, and R.E. McCulloch. BART: Bayesian additive regression trees. The Annals of Applied Statistics, 4( 4(1):266-298 2010. doi: https://doi.org// 10.1214999--AA22555
K. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, Sarlós, Hawkins. J. Davis,
A. Mohiuddin, L. Kaiser, D. Belanger, L.J. Colwell, and A. Weller. Rethinking attention with performers. CoRR, hiss22009.14944 2020. URL BRRlittps://arxe. wttp://www....////////o////////. 14794
P.T. Costa and R.R. McCrae. The SAGE handbook of personality theory and assessment, chapter The revised neo personality inventory (neo-pi-r), pp. 179-198 198. SAGE Publications, 2008
A. Curth, D. Svensson, J. Weatherall, and M. van der Schaar. Really doing great estimating CATE? critical look ML benchmarking practices treatment effect estimation. 35th Confer- ence onf Neural Information Processing Systems (NeurIPS 2021), 2021.
T. Dao. FlashAttention-22: Faster attention with better parallelism and work partitioning. In Interna- tional Conference on Learning Representations (ICLR), 2024.
T. Dao, Daniel Y. Fu, S. Ermon, A. Rudra, and C. Ré. FlashAttention: Fast and memory- exact attention with 10-awareness. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
de Haan, D. Jayaraman, and S. Levine. Causal confusion in imitation learning. 33rd Conference on Neural Information Processing Systems, 2019.
A. Deaton and N. Cartwright. Understanding and misundertstanding randomized controlled trials. Social Science and Medicine, 210:2-21, 2018. doi: 10, (0.101//.socscimeel 2oc.cimedd.001..12.000
R. H. Dehejia and S. Wahba. Propensity score-matching methods for nonexperimental causal stud- ies. Review of Economics and Statistics, 84(1):151-161, 2002.
Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional trans- formers for language understanding. arXiv preprint,
2019.
V. Dorie. Narmmerttss for causal inference.
2016.
A. Dosovitskiy, Beyey L., A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An image is worth 16x words: Transformers for image recognition scale. ICLR, 2021.

---

Published as a conference paper at ICLR 2026
S.F. Friedman, G.E. Moran, M. Rakic, and A. Phillipakis. Genetic architectures of medical images revealed by registration and fusion of multiple modalities. bioRxiv preprint, 2023. doi: https: //doi.org/10.1 101.01110122002..27..555555
B. Fuchs, D.E. Worrall, V. Fischer, and M. Welling. SEa3-TTransfommrs:: 3D roto-translation equivariant attention networks. 34th Conference on Neural Information Processing Systems, 2020.
M. Gagrani, C. Rainone, Y. Yang, H. Teague, W. Joen, R. Bondesan, H. van Hoof, Lott, W. Zeng. and P. Zappi. Neural topological ordering for computation graphs. Advances in Neural Informa- tion Processing Systems, 35, 2022.
M. Germain, K. Gregor, I. Murray, and H. Larochelle. MADE: Masked autoencoder for distribution estimation. Proceedings of the 32nd International Conference on Machine Learning, 2015.
C. Glymour, K. Zhang, and Spirtes. Review of causal discovery methods based on graphical models. Frontiers in Genefics, 10, 2019.
A. Goyal and Y. Bengio. Inductive biases for deep learning of higher-level cognition. Proceedings of the Royal Society A, 478(2266), 2022. doi: https://doi.org/10.. http:/oorrg....00//ssaaa.a.000 10w.1000.0....20.0000 .0068.
C.W.J. Granger. Testing for causality: A personal viewpoint. Journal of Economic Dynamics and Control, 2(1):329-3522 1980.
M.P. Grosz, J.M. Rohrer, and F. Thoemmes. The taboo against explicit causal inference in nonexperimental psychology. Perspectives on Psychological Science, pp. 1-13, 2020. doi: (T0177//1745691222215551
R. Guo, J. Li, and H. Liu. Learning individual causal effects from networked observational data. Association for Computing Machinery, 2020.
N. Hassanpour and R. Greiner. Learning disentangled representations for counterfactual regression. ICLR, 2020.
C. Heinze-Deml M.H. Maathuis, and N. Meinshausen. Causal structure learning. Annual Review of Statistics and Its Application, 5, 2018.
M. Hernan. The c-word: scientific euphemisms do not improve causal inference from observational data. American Journal of Public Health, 108(5):625-62) 2018. doi: 10.2105/AJPH.2018 304337.
I. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, and A. Lerchner. Beta-VAE: Learning basic visual concepts with constrained variational framework. ICLR, 2017.
Hilpert and MJJ. Vowels. A critical evaluation of theories and their empirical testing methods. OSF psyarxiv preprint, 2024. doi: https://osf.ioffchbb.
Hilpert, T. Brick, C. Flueckiger, M. J. Vowels, Ceuleman, P. Kuppens, and L. Sels. What can be learned from couple research: Examining emotional co-regulation processes in face-to- face interactions. Journal of Counseling Psychology, 2019.
G.W. Imbens and D.B. Rubin. Causal inference for statistics, social, and biomedical sciences. An Introduction. Cambridge University Press, New York, 2015.
S. Ioffe and C. Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. 1rXiv:1502.0331673,, 2015.
Jumper J., R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneberger, and D. Hassabis. Highly accurate protein structure prediction with alphafold. Nature, 596:583-589, 2021. doi: https: //doi.org/10. 0038//41586-011-9355
A. Kelava, B. Nagengast, and H. Brandt. A nonlinear structural equation mixture modeling approach for nonnormally distributed latent predictor variables. Structural Equation Modeling: A Multidis- ciplinary Journal, 21(3):468 481446888811118 2014. doi: https://doi.org/10.. 1080/ 01080/107055111 1.2014.915379.
12

---

Published as a conference paper at ICLR 2026
D. P. Kingma and M. Welling. Auto-encoding variational Bayes. arX1v:131..6114v10,
R.B. Kline. Principles and practice of structural equation modeling. Guilford Press, 2005.
M. Kocaoglu, C. Snyder, A.G. Dimakis, and S. Vishwanath. CausalGAN: Learning causal implicit generative models with adversarial training. ICLR, 2018.
N. Kreif and K. DiazOrdaz. Machine learning in policy evaluation: new tools for causal inference. arXiv:1903.000 00402... 2019.
R. G. Krishnan, U. Shalit, and Sontag. Deep kalman filters. arXiv:/511.05/2v 2015.
K. Kroenke, R.I., Spitzer, and J.B. Williams. The phq- 15: validity new measure for evaluating the severity of somatic symptoms. Psychosomatic Medicine, 64(2):258- 2002. doi: 10.1097/
00006842-200203300000000
L. Laffers and G. Mellace. Identification of the average treatment effect when SUTVA is violated. Discussion Papers on Business and Economics, 3, 2020.
R.J. LaLonde. Evaluating the econometric evaluations of training programs with experimental data. The American Economic Review, pp. 604-620, 1986.
Y. LeCun and Y. Bengio. Convolutional networks for images, speech, and time-series. In M. A. Arbib (ed.), The handbook of brain theory and neural networks. MIT Press, Cambridge, Mas- sachusetts, 1998.
Y. Li and S. Mandt. Disentangled sequential autoencoder. Proceedings on the 35th International Conference on Machine Learning, 2018.
Locatello, Abbati, T. Rainforth, T. Bauer, S. Bauer, Scholkopf, and O. Bachem. On the faimess of disentangled representations. MxXi::905.13662v/ 2019a.
F. Locatello, S. Bauer, M. Lucie, G. Ratsch, S. Gelly, B. Scholkopf, and Bachem O. Chal- lenging common assumptions in the unsupervised learning of disentangled representations. urXiv::811 I. 12359vv3, 2019b.
I. Loshchilov and F. Hutter. Decoupled weight decay regularization. ICLR, 2019.
C. Louizos, U. Shalit, Mooij, D. Sontag, R. Zemel, and M. Welling. Causal effect inference with deep latent-variable models. 31st Conference on Neural Information Processing Systems, 2017a.
C. Louizos, K. Swersky, Y. Li, M. Welling, and R. Zemel. The variational fair autoencoder. 1511:5030000000 2017b.
Y. Luo, V. Thost, and L. Shi. Transformers over Directed Acyelic Graphs. 37th Conference on Neural Information Processing Systems (NeurIPS), 2023.
D. Malinsky, I. Shpitser, and T. Richardson. A potential outcomes calculus for identifying ditional path-specific effects. Proceedings of the 22nd International Conference on Artificial Intelligence and Statistics (AISTATS), 2019.
O. McBride, Murphy, M. Shevlin, Gibson-Miller. T.K. Hartman, Hyland, L. Levita, L. Ma- son, A.P. Martinez, R. McKay, T.V.A. Stocks, K. Bennett, F. Vallieres, T. Karatzias, C. Valiente, C. Vazquez, and R.P. Bentall. Monitoring the psychological, social, and economic impact of the covid- 19 pandemic in the population: Context, design and conduct of the longitudinal covid-19 -99 psychological research consortium (c19pre) study. International Journal of Methods Psychi- atric Research, 30(1), 2021. doi: 10. 1000//pr/r 18611 861.
V. Melnychuk, D. Frauen, and S. Feurerriegel. Causal transformer for estimating counterfactual outcomes. Proc. 39th International Conference on Machine Learning, 2022.
A. Paszke, S. Gross, S. Chintala, G. Chanan, E. Yang, Z. DeVito, Z. Lin, A. Desmaison, Antiga, and A. Lerer. Automatic differentiation in PyTorch. NeurIPS Workshop, 2017.
J. Pearl. Causality. Cambridge University Press, Cambridge, 2009.
13

---

Published as a conference paper at ICLR 2026
J. Pearl and D. Mackenzie. The book of why. Penguin Books, 2018.
Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, and B. al. Thirion. Scikit-learn: Machine learning in Python. JMLR, 12:2825-2830, 2011.
J. Peters, D. Janzing, and B. Scholkopf. Elements of Causal Inference. MIT Press, Cambridge, Massachusetts, 2017.
M. Petersen, L.. Balzer, D. Kwarsimma, N. Sang, G. Chamie, J. Ayieko, J. Kabami, A. Owaraganise, T. Liegier, Mwangwa, and K. Kadede. Association of implementation of universal testing and treatment intervention with HIV diagnosis, receipt of antiretroviral therapy, and viral suppression in East Africa. Journal of American Medical Association, 317(2, 317(21):2196-2266 2017. doi: 10. 1001/jama.2017.576.
M.L. Petersen, K.E. Porter, S. Gruber, Y. Wang, and MJJ. van der Laan. Diagnosing and responding to violations in the positivity assumption. Statistical Methods in Medical Research, 21(1):31-54, 2012. doi: 10.1 10.117709922800113382222
Probst, M.N. Wright, and A.-L. Boulesteix. Hyperparameters and tuning strategies for random forest. Wires Data Mining and Knowledge Discovery, 2018. doi: 10. 10.100//wddm.3000
A. Radford, J.W. Kim, I. Xu, G. Brockman, C. Mcleavey, and I. Sutskever. Robust speech recog- nition via scrge weak supervision. Proceedings of the 40th International Conference on Machine Learning, PMLR, 202:28492-285111 2023.
Ramachandran, Zoph, and Q.V. Le. Searching for activation functions. ICLR Workshop, 2018.
T.S. Richardson, R.J. Evans, J.M. Robins, and Shpitser. Nested Markov properties for Acyclic Directed Mixed Graphs. arXiv preprint, arXiv: arXiv:1701..000 1ee.17700..666622 2017.
Robins. new approach causal inference in mortality studies with sustained exposure period application to control of the healthy worker survivor effect. Mathematical Modelling, 7:1393- 1512, 1986. doi: 10.1016/0270- (288770000558880088888888
J.M. Rohrer. Thinking clearly about correlations and causation: Graphical causal models for ob- servational data. Advances in Methods and Practices in Psychological Science, 2018. doi:
10.117/551525911445699
M. Rojas-Carulla, B. Scholkopf, R. Turner, and J. Peters. Invariant models for causal transfer learn- ing. Journal of Machine Learning Research, 19:1-34, 2018.
R. Rosenbaum and D. B. Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-55, 1983.
D. B. Rubin. Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American Statistical Association, 100(469):322-331, 2005. doi: 10. 10.1198/0162144440000011111
D.B. Rubin. Comment on 'randomization analysis of experimental data in the Fisher randomization test' by D. Basu. Journal of the American Statistical Association, 75:591-593, 1980.
B. Scholkopf. Causality for machine learning. arXiv:1911.10500%1 2019.
S. Schwartz, N.M. Gatto, and U.B. Campbell. Extending the sufficient component cause model to describe the stable unit treatment value assumption (SUTVA). Epidemiologie Perspectives and Innovations, 9(3), 2012. doi: 10. 10.11866 0e.11886/7722.0 11118//7742-53-9-33
U. Shalit, F. D. Johansson, and D. Sontag. Estimating individual treatment effect: generalization bounds and algorithms. ICML, 2017.
Shen, S. Wu, Y. Yang, and X. Quan. Directed Acyclic Graph Network for conversational emotion recognition. 59th Annual Meeting for the Association for Computational Linguistics, 2021.
C. Shi, D. M. Blei, and V. Veitch. Adapting neural networks for the estimation of treatment effects. 33rd Conference Neural Information Processing Systems, 2019.
14

---

Published as a conference paper at ICLR 2026
Iya Shpitser, Thomas S. Richardson, and James M. Robins. Multivariate Counterfactual Systems and Causal Graphical Models. pp. 813- 852. Association for Computing Machinery, New York, NY, USA, 1 edition, 2022. ISBN 9781450395861. URL
3501714 3501757
J. A. Smith and P. E. Todd. Does matching overcome LaLonde's critique of nonexperimental esti- mators? Journal of Econometrics, (25(1): (25(3)305-335, 2005.
P. Spirtes and K. Zhang. Causal discovery and inference: concepts and recent methodological advances. Applied Informatics, 3(3), 2016. doi: 11/86//40535-01--0000
R.L. Spitzer, K. Kroenke, J.B.W. Williams, and B. Löwe. A brief measure for assessing generalized anxiety disorder: The gad-7. Archives of Internal Medicine, 166(10):0022-009,, 2006. doi: 10. 10.1001/architte.16 10..10..0.... 109.
N. Umbach, K. Naumann, H. Brandt, and A. Kelava. Fitting nonlinear structural equation models in with package nlsem. Journal of Statistical Software, 77(7), 2017. doi: https://oi.org/10 18637/jjs.v007..ll v077.i07.
B. Uria, I. Murray, and H. Larochelle. deep and tractable density estimator. 31st International Conference on Machine Learning (ICML), pp. 467-475, 2014.
K. Vafa, J.Y. Chen, J. Kleinberg, Mullainathan, and A. Rambachan. Evaluating the world model implicit in generative model. arXiv preprint, arXiv:2406.0366 036669 2024.
M. J. van der Laan and R. J. C. M. Starmans. Entering the era of data science: targeted learning and the integration statistics and computational data analysis. Advances in Statistics, 2014.
M.J. van der Laan, E.C. Polley, and A.E. Hubbard. Super Learner. Statistical Applications of Genetics and Molecular Biology, 6(25), 2007. doi: 00.2202/554-6111513300
E.J. van Kesteren and D.I.. Oberski. Fleixble extensions to structural equation models using compu- tation graphs. Structural Equation Modeling: A Mulridisciplinary Journal, 29(2):233-247, 2022. doi: https://doi.org// 10. 1.200/00005511222..111
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A.N. Gomez, L. Kaiser, and I. Polo- sukhin. Attention is all you need. 3stt Conference on Neural Information Processing Systems, 2017.
Verma and Pearl. Equivalence and synthesis of causal models. Proc. 6th Conf. on Uncertainty in Artificial Intelligence, 1990.
O. Vinyals, Babuschkin, W.M. Czarnecki, M. Mathieu, Dudzik, J. Chung, D.H. Choi, and D. Silver. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575: 350-354, 2019. doi: https://doi.org/100 wwt.:///......11....eoo 018660019-11 1724-7.
L.M. Vowels, M.J. Vowels, K.B. Carnelley, A. Millings, and Gibson-Miller. Toward a causal link between attachment styles and mental health during the COVID-19 pandemic. British Journal of Clinical Psychology, 02(3))605-620, 2023a. doi: https://doi.org/I1 1 12428
M. J. Vowels. Misspecification and unreliable interpretations in psychology and social science. Psychological Methods, 2023a. doi: https://doi.org/100 httg//00..03//mmm000049999
M. J. Vowels. Prespecification of structure for increasing research transparency and for the opti- mization data collection. Collabra: Psychology, 9(1), 2023b. doi: https://doi.org// 1555/ collabra.71300.
M. J. Vowels, N. C. Camgoz, and R. Bowden. Vdsm: Unsupervised video disentanglement with state-space modeling and deep mixtures of experts. IEEE Conference on Computer Vision and Pattern Recognition, 2021a.
M. J. Vowels, N.C. Camgoz, and R. Bowden. Targeted VAE: Structured inference and targeted learning for causal parameter estimation. IEEE SMDS, 2021b.
15

---

Published as a conference paper at ICLR 2026
M.J. Vowels. Trying to outrun causality with machine learning: Limitations of model explainability techniques for identifying predictive variables. arXiv preprint, arXiv:2202.09875, 2022.
M.J. Vowels. SLEM: Machine learning for path modeling and causal inference with super learner equation modeling. arXiv preprint, arXiv:2308. 04366,, 2023c.
M.J. Vowels. A causal research pipeline and tutorial for psychologists and social scientists. Psycho- logical Methods, Advanced Online Publication, 2024.
M.J. Vowels, N.C. Camgoz, and R. Bowden. D' y'ya like DAGs? A survey structure learning and causal discovery. ACM Comput. Surv., 2022. doi: 10. 10.1145/3527154.
M.J. Vowels, S. Akbari, N.C. Camgoz, and R. Bowden. A free lunch with influence functions? an empirical evaluation of influence functions for average treatment effect estimation. Transactions on Machine Learning Research, 2023b. doi:
S. Wager and S. Athey. Estimation and inference on heterogeneous treatment effects using random forests. Journal of the American Statistical Association, 1136623);1288-122,, 2018. doi: 1080001621559.217.1133 13119899..
A. Wehenkel and G. Louppe. Graphical normalizing flows. Proceedings of the 24th International Conference on Artificial Intelligence and Statistics (AISTATS), 130, 2021.
K. Xia, K.Z. Lee, Y. Bengio, and E. Bareinboim. The causal-neural connection: Expressive- ness, learnability, and inference. 35th Conference Neural Information Processing Systems (NeurIPS), 2021.
M. Yang, F. Liu, Z. Chen, X. Shen, J. Hao, and J. Wang. CausalVAE: disentangled representation learning via neural structural causal models. arXiv:2004.08697v4 2020.
L. Yao, S. Li, Y. Li, M. Huai, Gao, and A. Zhang. Representation learning for treatment effect estimation from observational data. 32nd Conference Neural Information Processing Systems (NeurIPS), 2018.
L. Yao. Z. Chu, S. Li, L., Gao, and A. Zhang, A survey causal inference. ACM Transactions on Knowledge Discovery from Data, 15(5):1-46, 2020. doi: 10.1145/3444944.
G. Yenduri, M. Ramalingam, G.C. Selvi, Y. Supriya, G. Srivastava, P.K.R. Maddikunta, and T.R. Gadekallu. Generative pre-trained transformer: A comprehensive review enabling technolo- gies, potential applications, emerging challenges, and future directions. IEEE Access, 12:54608- 54649, 2024. doi: https://doi.org/ 10. 109/ACCESS 109/A//CESS.2224.3333533 2024.3389497..
Q. Yin, X. He, X. Zhang, Y. Zhao, J. Yao, X. Shen, and Q. Zhang. Stablemask: Refining causal masking in decoder-only transformer. arXiv preprint, arXiv:2402.0477991. 2024.
J. Yoon, J. Jordan, and M. van der Schaar. GANITE: Estimation of individualized treatment effects using generative adversarial nets. ICLR, 2018.
M. Zhang, K. Bhatia, H. Kumbong, and C. Re. The hedgehog and the porcupine: Expressive linear attentions with softmax mimicry, 2024. URL htt.tttss//eww...rrr/bbbb////
04347
W. Zhang, L. Liu, and J. Li. Treatment effect estimation with disentangled latent factors. Proceed- ings of the Thirty-Fifth AAAI Conference on Artificial Intelligence, 2021.
16

---

Published as a conference paper at ICLR 2026
L
1.
L
6x3
X
X
X
0.
1.
3x3
1.
1.
1.
3x6
X
0.
0.
0.
0.
0.
1
1.
1.
1.
0.
1
1.
X
0.
0.
1.
0.
0.
0.
0.
0.
0.
0.
1.
1.
0.
1.
1.
X
0.
0.
0.
0.
1.
Layer Mask
Layer 2 Mask (neurons=6)
0.
1.
Adj. Matrix
(neurons=3)
Layer Mask
(neurons=3)
Figure 4: Simple example sequence of matrix transformations used to generate masks for CFCN in three variable, three layer case, where the DAG includes X - X - X mediation as well as direct path X + X and the number of neurons in each layer is 3, 6, and 3. The transition from Layer to Layer 2 includes the introduction of the diagonal "pass-through" which is absent in the first layer.
SUPPLEMENTARY MATERIAL
S1. OVERVIEW
This supplementary material accompanies the paper CaT' and DAGs and full code, including exper- iments, is attached. An open repository containing the code will also be released upon acceptance. We begin this supplementary material by providing list of notation used throughout, in Section [S2 We then provide some additional illustrations and information concerning the masking mechanisms used CFCN in Section S3 and describe variation of CFCN called CFCN 'mod' in Section S4 overview of the typical self- and cross-attention mechanisms used in transformers in Section S5 and present the algorithm used for recursive inference in Ca''s and CFCNs in Sec- tion S6 We describe approach called Isomorphic Shuffing in Section S7 although this was not found to yield good results in practice.
We provide
In Section 58 we present some preliminaries relating causal inference including the notion of one- step-ahead potential outcomes (which motivates the recursive inference algorithm in Section S6).. Then, in Section S9 provide the hyperparameters used for the evaluations. In Section [STO
provide additional information and results for the motivating simulations which were presented in the main text. In Section [SIT we provide additional experimental results for other causal infer- ence tasks involving the benchmark Jobs and Twins benchmark datasets well as the real-world psychology dataset.
In Section S12 we present a curious result for how the same scikit-leamn (Pedregosa et al.) (011)) random forest (Breiman] 2001) and transformer used in the motivating experiments perform on the causal inference benchmarks.
Finally, Section S13 we provide theoretical proof of why the architectural design ensures that the conditional independencies implied by the DAG hold in the output of the network.
S2. NOTATION
D dataset
sample size
z node set of nodes in graph
Y univariate outcome
Y multivariate outcome
L univariate covariate
L = multivariate covariate
Lyet set of multivariate covariates
L number of covariates
17

---

Published as a conference paper at ICLR 2026
D univariate treatment
joint observational distribution
P(Y D, L) conditional observational distribution
T = treatment effect
G topologically sorted DAG
adjacency matrix for DAG
= set of nodes in DAG
Z number of nodes in DAG
O causal ordering for each variable in topologically sorted DAG
causal descendants of intervention nodes
D' dataset updated with interventions and predictions of estimated effects
CFCN = Causal Fully Connected Network
layer block number in CFCN CaT, respectively
total number layers or blocks in CFCN CaT, respectively
W weight/parameter matrix layer
Ø activation function
b, bias vector for layer
O, = vector output of CFCN for layer
M, = CFCN parameter mask for layer
CaT = Causal Transformer
query embeddings in CaT and conventional transformer
key embeddings in CaT and conventional transformer
= value embeddings in CaT and conventional transformer
S 'similarities' in conventional transformer
X input to CaT
X g 're-embedded' input X CaT
C = embedding dimension for CaT input before 're-embedding'
d dimensionality after 're-embedding"
Y learnable embedding
number parallel heads in each block
hg = head size
CMCA Causal Multi-head Cross Attention
OCMCA output of CMCA
O'....O' output of parallel heads
Or Block output from CaT layer/block
P permutation matrix
I set of all intervention value and variable pairs (H, 2)
(n, z) an intervention value and variable pair, respectively
M Trained CaT or CFCN model
Q empirical estimator (e.g. random forest CaTT)
18

---

Published as a conference paper at ICLR 2026
11
12
1/
Ex.1
X1
X1
X2
X2 X1
X3
Ng X1
11
12
1/
Ex.2
X1
X2
X2 X1, X3
X;
X3
X3
X4
X4 Xg. X2
Figure 5: Two examples demonstrating how the delayed introduction of the identity diagonal at layer and onwards (a) prevents inputs attending to themselves the first layer and (b) allows other intermediate predictions of these inputs to be used for the prediction other variables. Without the introduction of identity, signals from e.g. X1 used for predicting X2 would be blocked, and X2 itself could not be predicted. Applies to both CaT and CFCN.
$3. CFCN MASKING
The masking process in CFCN is applied similarly to (Germain et al. 2015):
o, (o M,) + b,),
where is an activation function, is the vector output, W, is the weight matrix, b, is bias vector, and M, a mask for given layer r, respectively. is the Hadamard product. An example is shown in Figure A visualization of this process across different layers is provided in Fig Note that, while not strictly necessary for respecting the DAG, links between paths (such as the diagonal links visible in 5 improve network expressivity by enabling variables like and 22 to interact more deeply across layers, c.g. recursively processing (æ1, x22 together to influcnce Ø3.
Twins
Jobs
Model
cATE ws
eATE os
Risk ws
Risk os
eATT ws
eAIT
CFCN
0098 .0488
01022 0092
.5 02
26 06
016 01
.84 = 00
CFCN mod
.0099 .0050
.0065 : .006
.24 4 02
= .6
.011 =
.82 =
Table 4: Results + standard deviations for the CFCN 'mod' approach described in Sec. S4 for the Twins and Jobs benchmark datasets. eATE absolute ment Effect; Risk Policy risk; eATT absolute error on the estimation of the Average Treatment effect on the Treated; within sample;
the estimation of the Average Treat-
of sample.
S4. CFCN MOD
The CFCN Mod is variation the previously introduced CFCN described in Sec. 3.3 wherein each layer has direct access to the input through a linear connection. The computation of the layers given by
Yo = (x WW o Mo) + ba),
y, o W M,) x(W, M,) + b,),
where x is the input vector, Yr an output vector, W,. W: are weight matrices for layer r, M,, M are the corresponding masks, b, is bias vector, and the Hadamard product.
key distinction from the CFCN architecture lies in the handling of masks for the weight matri- ces. Specifiiallly, the masks M include the identity diagonal, ensuring that each layer maintains
19

---

Published as a conference paper at ICLR 2026
self-connections. Conversely, the mask for M, does not include this diagonal, preventing direct propagation of the imput to the output, and enforcing the DaG constraints.
These modifications ensure that each intermediate layer has access to the imput. thus mitigating the risk of imput signal degradation throughout the network. Additionally, this design facilitates a mean- ingful comparison between the intermediate output features and the original input. For instance, if we consider a model predicting weight from height, the difference between the predicted weight at an intermediate layer and the true weight (derived from the input) could serve as a valuable indicator of body composition, such as identifying obesity. We hypothesize that this structure enhances the model's s ability retain essential input information while allowing richer feature interactions across layers.
Table show the performance of the CFCN and CFCN mod on the Twins and Jobs benchmark datasets. Whilst we believe CFCN 'mod' is a valuable alternative to CFCN, our current evaluations indicate similar performance across the two variations.
S5. CONVENTIONAL MULTI-HEAD SELF- AND CROSS-ATTENTION
The core components of transformers Vaswani et al. 2017) concern the way inputs are processed by the multi-head attention (MHA). Consider an input of 7 3, C-diiensiioall embeddings a, b, and e (boldfont to indicate vectors). Together these variables constitute an input matrix X, which B |Z] X C. Each of the three constituent inputs in X are processed independently, and by three linear transformations, yielding the 'query" Q, 'key" K, and 'value" V embeddings. Self- attention, derived via S
[QK]
where h2 is the dimensionality of the embeddings, creates a 3x3
interaction matrix, describing the extent which each pair from the embeddings of a, b and similar to each other. This attention matrix represents the key opportunity for causal masking with an adjacency matrix, and we discuss how this is done below. The attention matrix is then matrix multiplied with V such that the different components of V interact with a strength proportional to the similarities encoded in
In addition to self-attention, transformers may also employ cross-attention mechanisms, which can be useful in contexts where the input and targets differ. atosstntt processes two distinct sets of embeddings: query matrix Q from the target sequence and key-value matrices K and V from the input sequence. Suppose we have a target matrix of size B Z x C. where | is the number of target embeddings. The query Q is obtained from T. while the key K and value V are derived from the input matrix X. Cross- Pross-attention then computes the similarity between the target queries and the resulting in |Z|| |Z| interaction matrix. This interaction matrix represents how much each target embedding in T attends to the input embeddings in X. As with self-attention, S is then used to reweight V, allowing information from the input to be incorporated
input keys using
QK
into the target sequence based on the learned interactions.
S6. INFERENCE ALGORITHM
The algorithm for performing recursive inference with CaT CFCN is shown in Alg. [ This algorithm can be read as an implementation of recursive substitution (see Section S8)
S7. ISOMORPHIC SHUFFLING
We develop means to shuffle the ordering of the variables and the masks such that the topology of causal structure remains the same whilst the order the adjacency matrix is permuted randomly during training. This approach is inspired by the order agnostic training presented by (Germain et al.) 2013). These permutations are isomorphic with respect to the DAG used to specify the adjacency matrix, and to the weights/parameters of the CFCN and CaT which are not shuffled. The function approximation process is thereby forced to yield a solution which respects the constraints of the underlying DAG, whilst nonetheless being forced to fulfil this role under permutations of this DAG.
20

---

Published as a conference paper at ICLR 2026
Algorithm Estimating the effects of interventions.
1: function FORWARD(D, I, A, O)
2: input: D dataset
3: input: I intervention value and variable pairs (u, 2)
4: input: .M- Trained CaT or CFCN model
5: input: A topologically sorted adjacency matrix for DAG
6: input: O - causal ordering for each variable in DAG
7: output: D' interventions and estimated effects
8:
D' cD)
Create copy of the data b Intervention nodes
9:
+ {1 (n,z) I
10:
for u. z in I do
11:
('())= <-
D Apply intervention
12:
end for
13:
F F Desc(1, I
> Get nontrivial causal descendants > Topological sort
14:
F FF
15:
for pua in F do
b Iterate through sorted descendants b Get model predictions b Recursively update dataset.
16:
preds +- M(D')
17:
('aa) preds
18:
end for
19:
return 2
20:
end function
More formally, assume that A is the adjacency matrix for the DAG representing the DGP, A is [Z] x |Z] if the DAG has [Z] nodes, where A, 1 if there is an edge between vertex and vertex We generate a random ordering p from 1 to Z and use this to construct a permutation matrix P. which also / x [Z]. The ith row of
assigned a the jth column if, according to the random ordering p. vertex should be mapped to vertex j in the permutation. The permuted (but isomorphic to the original graph) adjacency matrix A' PAPT.
Note that in practice, we find that this shuffling worsened both predictive metrics (like MSE) and widened the standard errors for the ATE estimates. However, found positive application in (Ger-) main et al. 2015) and (Uria et al. 2014). so it may be found to be useful for specific tasks.
S8. CAUSAL PRELIMINARIES
We consider the set of variables Zset. where Y, C Zset denote the (possibly multi-dimensional and/or multi-variate) outcome and treatment, respectively, and Lset = Zset \ Y.D D are covariates. We assume causal ordering O among these variables, and that the causal relations between them represented by a DAG G. For any Z C Zees we denote by Paz the set of vertices (corresponding to variables) that have directed edge to Z in & (direct causes, parents). G conforms to the ordering O in the sense that Pa, precede 2 in the ordering O for any E Zue this indeed guarantees that G a DAG and no cycles are formed.
For ease of notation, we will use the same notation for a graph vertex as for its value. Following (Mlinssky al. 2019: Shpptser et al. 2022), we let the one-step-ahead potential outcome variables Z (Paz) denote the value of Z given its parents values Paz, possibly contrary to fact. The latter can also be defined as the output of structural equation model (SEM) where function f(-) maps the values of Paz along with an exogenous noise to values of Z. We assume the existence of all one-step-ahead potential outcomes. Given these one-step-ahead potential outcomes, for any subset D Zset= the potential outcome Z(D), is defined through recursive substitution as follows:
Z(D) = z ((Paz D, Paz \ D (D)).
(10)
When predicting the values of Z(D) in Alg. [ we apply recursive substitution by first predicting the values of Paz (D) and plugging them into Eq. equation 10 for each variable Z.
Equivalent to what we described above, Z(D) can be defined as the random variable induced by a modified SEM where the functions f(-) for those variables in D are replaced by constant functions setting their value to d. Authors denote this modified SEM using do(D = d) following Pearl's do notation (Pearl 2009). In particular, the distribution of Z(D) is shown by P doDD = d)).
21

---

Published as a conference paper at ICLR 2026
For outcome Y and treatment D, the average treatment effect is defined
== PYY - (Y-D(((do(D=))) M.
For any subset I sub Let such that ( is independent of D conditioned sub (equivalently an adjustment set in Pearl's terminology
we have
-= Er. sub
By BY-POYyde(D= [Y]
sub
P
-E 0) sub [Y]
(11)
sub ~p
PY sub (Y]
- -YY-PLYDD- sub M
In practice, we might estimate this quantity using an empirical estimator (like a random forest) Q as follows:
1
+(Q; sub) =
@(D 1, = subb,
N
(12)
@D = 0, sub = sub,
where the hat notation indicates value is estimated empirically rather than population/true quan- tity.
Finally, and importantly, we note that the success of causal inference rests on a number of strong
and often untestable assumptions Yao et al.) 2020, Guo et al. 2020 Rubin 2005 Imbens & Ru-
bin 2015 Grosz et al. 22000; Petersen et al. 2012; Rubin| 1980 Hernan 2018): (1) Ignorabil-
Exchangeability this assumes that there confounders, meaning that individuals with identical covariates should have the same probability of receiving treatment, and their potential outcomes should also be identical given the same latent covariates. (2) Positivity this requires that treatment assignment probabilities are strictly positive and not deterministic, ensuring that P(D d L,) >> 0,: V; L. In other words, for every possible subset of the data, there must be a non-zero probability of receiving any given treatment. (3) Stable Unit Treatment Value Assumption (SUTVA) (Rubin) 1980 Laffers & Mellace [220]] Schwartz et al.] 2012) this assumption states that an individual's 's potential outcomes are not influ- enced by the treatment assignments of others, meaning there are no spillover effects, and that the effect of treatment remains consistent regardless of how an individual came to receive it (Schwartz) ctt al. 2012).
ity/ Cecoonfoundedness//oonniiiii
unobserved
S9. EXPERIMENTAL SETUP AND HYPERPARAMETERS
The hyperparameters for CaT and CFCN (and, identically, CFCN mod) were the same for the mo- tivating experiment, Twins, Jobs, and the real-world dataset experiments presented in the main text, and are shown in Ta
S10. MOTIVATING EXAMPLE
We generate data from three Data Generating Processes (DGPs), the Directed Acyelic Graphs (DAGs) for which shown in Figure in the main text. The structural equations for the DGPs for all three graphs are the same (the data are generated according to the correct graph). In contrast, no graphs are given to the transformer, random forest (RF), multilayer perceptron (MLP). These do
Note that such subset always exists, e.g., Lub Pap-
22

---

Published as a conference paper at ICLR 2026
Method CFCN
Hyperparameter
Value
Dropout
0.0
Neurons per layer
[s. 2*s, 2 2*s. 2
Activation
Swish
CaT
Dropout
0.0
Num Heads
Num Layers
2
Head Size
6
Input Embedding Dim
Projection Dim
Training
No. Iters CaT
20k
No. Iters CFCN
6k
Batch Size
100
Learning Rate
5e-3
Optimizer
AdamW (Loshchloov & Hutter 2019
Table 5: Hyperparameters for CFCN, CaT, and Training Settings. Note that 's' the input size for CFCN which is determined by the dimensionality of the input for the dataset being used.
not accept graph constraints, whereas CFCN and CaT are evaluated under the three different sets of constraints. The structural equations are as follows:
UD NO 1), Us 1),
U1, ~ WNO,,) 1), Uy ~ NO 1),
(13)
D = UD, L1 0.8D Uh, == 0.8D + 0.4L, + U,, L2 = 0.7Y +0.6D + Un.
and as such, the total average treatment effect (ATE) of D on Y is (0.8 (0.8 0.4) = 1.12.
The RF and MLP are implemented in scikit-learn Pedrreosa et al. 2011). where the transformer utilizes the TransformerEncoder module from pytorch (Paszke et al. 2017), and is given the same hyperparametters as CaT. Given the simplicity of the DGP (4 unidimensional variables in total, linear functional form, Gaussian error) we use the default hyperparameters in the MI.P and RF. Indeed, the default parameters in the seikit-learn RF have been shown to work well across a wide range of tasks (Probst et al., 2018). The hyperparameters for the transformer and the CaT' are as follows: batch size 300, training iterations 20k, number of heads number of layers 3, head size 4, input embedding dimension 4, 'internal' embedding dimension 4, and learning rate of 0.0002. For CFCN, the batch size was 50, we used 4, 8, and 4 neurons in each of the hidden layers, and applied a learning rate of 0.001. Both CFCN and CaT were trained with Adam W optimizer (Loshchilov & Hutter 2019) with learning rate scheduling.
The sample size is set to 10,000 to avoid issues with random sampling variation, and generate data and undertake the analysis for 30 simulations of these data. The MSEs presented in the main text are computed over the validation set which represented 20% split of the data, and the ATES were computed over the full training-validation combination, once the models had been trained on thet training data. This paradigm (training on the training set but evaluating the ATE across the full dataset) is normal in the context of causal inference, where the optimization of the network during training does not represent the same task as in conventional supervised learning. any case, very similar results are obtained if only the validation split is used for evaluating the ATE estimation.
The results for the error on the ATE estimation across methods are shown in Fig 6 This plot highlights how the lowest error on the estimation of the ATE is provided by CaT and CFCN when the correct DAG is used, whereas all other approaches fail to yield good estimates. At best, when partially correct graph is used, such as DAG 2 in the simulation ("partially" insofar it more closely aligns with the true DGP than DAG in terms of the relationships crucial for estimating the target ATE), the results are closer than with no DAG all, which is the case for the typical machine learning approaches seen with the use an random forest, transformer, or MLP. The key takeaway from this is that specifying graph, even if it is incorrect, provides transparent means represent
23

---

Published as a conference paper at ICLR 2026
Absolute for Difierent Models
2.-
1.0
0.0
Cat Dao
car 2
cete DAD
crce Fatt DAD
arce the DAG
Redd Fersst
Tansforna
MIF
CaT The
Figure 6: Illustrating the absolute error on the Average Treatment Effect estimation associated with conventional machine learning models, random forest (Breiman) 2001) (RF), multilayer perceptron (MLP), transformer (Vaswani et al. 2017), and our CFCN and CaT networks. We show the results for CFCN and CaT under three different graphs (one correct, and two false).
1.0
0.0
GiT ralse Dac
GaT Faite 2
GaT The DMG
CPeN Falte D4D
crea Palte D46
crai Rue DMG
Rand Forest
Transforher
Figure 7: Illustrating the -at-ol-sample mean squared error for the main outcome associated with conventional machine learning models, random forest (Breiman) 2001) (RF), multilayer perceptron (MLP), transformer (Vaswani et al. 2017). and our CFCN and CaT networks. We show the results for CFCN and CaT under three different graphs (one correct, and two false). One can see that the correct causal graph has the highest MSE, whereas under all other conditions / with all other approaches the MSEs are comparable.
the assumptions concerning the underlying DGP, and may yield meaningful results than failing to specify any constraints at all.
Finally, the complete results for the MSE the estimation the outcome variable Y in the valida- tion set are shown in Fig. 7 One can see that the use the correct causal graph actually results the worst estimation performance (despite having the best robustness and ATE estimation performance).
24

---

Published as a conference paper at ICLR 2026
S11. CAUSAL INFERENCE EVALUATION
S0.1
JOBS
Following descriptions in (LaLonde 1986 Vowels et al.) 202 Louizos et al. 200 Shalit al. 2017) we also use the job outcomes dataset (Jobs) (LaLonde 1986; Smith Todd 2005) (available
httpp:: /tt///usesa nber org. / rdeeeeeeaaaaaeee /ttee/aa/aaaa///eaa/aa/aaaaaa data/ www.ataa. ht.ml J. Unlike the IHDP dataset, Jobs is real-world data with a binary outcome. We follow a similar approach to (Shalit al. 2017], who used the Dehejia-Wan Wariia (Dehejia & Wahba 2002) sample along with PSID comparison group. This dataset includes 260 treated samples and 185 control samples, in addition to the 2,490 samples from the PSID comparison group. It combines both observational and randomized control trial (RCT) data. Following the methodology of (Louizos et al.) (011 (halit et al. (0177) we apply 56/24/20 train/validation/test split and conduct 100 runs with different random split assignments to estimate the average performance and standard error. Note that we use the same random seed across models for both initialization and dataset splitting, ensuring that the variance due to these factors is consistent across experiments. As in (Louizos et al., 201 7a Shalit et al.) 2017 Yao et al.) 2018), we evaluate our network on the Jobs dataset (where only partial effect supervision is available) by measuring the Average Treatment Effect on the Treated (ATT) error:
ATT = -1D-1-1 Σx - Σ Y,
RDD
DEDo
(14)
-11 11 Σ
Q1, Q( sett,)
iED
where, in slight abuse of notation D = D1 Do constitutes all individuals in the RCT, and the subscripts denote whether or not those individuals were in the treatment (subscript 1) or control groups (subscript 0). The first terms in Eq. comppris the true ATT, and the third term the estimated ATT. We may use the policy risk as proxy for the Precision in the Estimation of the Heterogeneous Effect (PEHE- see Eq. 18:
Reol 1 EDI (T = 1] PD (TLett) 1)
(15)
- EDolY
= 0] - PD (P((TT (Lset) = 0)
where TT (Iset,,) = is the policy to treat when (1( Y,(0)= > a, and x ([sset.)) = 0 is the policy not to treat otherwise (Yao et al. 20188, Shalit et al. 2017). a treatment threshold. This threshold can be varied to understand how treatment inclusion rates affect the policy risk. We set 0, as per (Shalit et al. 2017 Louizos et al., 2017a). The subscripts for the expectation operator and probability p indicate the specific population to which these notations apply.
SO.2 TWINS
We access the Twins (Almond et al. 2005) dataset from this package repository: https:// github. com/ br http///adyaaarrralaaaae
The dataset is based births of twins in the USA between 1989 and 1991. We designate the heavier twin as the treatment group (t= 1) and the lighter twin the control (t 0). The outcome measured is 1-year mortality. For each pair of twins, 30 features were collected, including details about the parents, the pregnancy, and the birth, such as marital status, race, residence, previous births, pregnancy risk factors, care quality during pregnancy, and the number of gestational weeks.
For the Twins dataset, provide estimate for the average error on the estimated Average Treat- ment Effect (cATE) both within- and out-of out-of-sample:
cATE |T 17-fi fl
(16)
or
(1)(1) Po())
eATE =
T
(17)
i=1
25

---

Published as a conference paper at ICLR 2026
gender
chronicil
age
keyworker
enxiety
loneliness
aresson
attachment
children
cial dist
Income
covid anx
nok1month
Figure 8: The DAG used in the real-world psychology example reconstructed from the causal discovery and domain expertise results presented in (Vowels et al. (2023)). Treatment is attachment style 'attachment' (also highlighted orange) and the two outcomes of interest at the measures of depression (highlighted in green).
where 7 is the true average treatment effect and y, and Vo are the estimated potential outcomes under treatment and treatment, respectively.
Table 6: Number of items dimensionality of each construct used in the wollw data analysis.
Construct
Dimensionality
Attachment Style
Loneliness
3
Chronic IlIness
Relationship Status
Social Distancing
16
Depression
Anxiety
Age
Gender
Keyworker Status
Adults in Household
Children Household
Change in Income
Covid Anxiety
Risk Taking in the Past Month
SO.3 WORLD DATASET
We performed secondary analysis of UK data from wave 2 of the COVID-19 Psychological Re- search Consortium Study (C19PRCS), longitudinal, internet- based survey. Comprehensive details about the methodology on the OSF (attps: //osf..... 10 /2//22ur/). In brief, the UK portion of the C19PRC Study was conducted in April/May 2020 for Wave 2. During this wave, strict social distancing measures
be found in (MeBride et al.) 202I)), and the data are openly available
26

---

Published as a conference paper at ICLR 2026
enforced. Quota sampling ensured the recruitment of nationally representative panel of UK adults, based on age, gender, and household income. Participants had to be at least 18 years old, UK resi- dents, and capable of completing the survey in English. After giving consent, participants completed the survey online and were compensated by Qualtries for their time. Ethical approval was granted by UK university psychology department (Reference: 033759). In Wave 2, 1406 individuals took part, but participants who did not report their attachment styles were excluded from the analysis, and after listwise deletion of rows with other missing values, we arrived at a sample size for this wave of 895 participants.
The motivation for these particular analyses of these data stems from how the COVID-19 pandemic posed numerous challenges, including balancing health protection with maintaining a fulfilling life. For some individuals, the mental health impact of the pandemic has been more severe than for others. The original study aimed to build a theoretical causal model to better understand how individual differences attachment styles affect both social distancing compliance and mental health outcomes (loneliness, depression, and anxiety) during the pandemic.
Twins
Jobs
Model
eATE ws
eATE os
Risk
Risk
eATT ws
eATT os
GANITE
.0058 = .0177
.0089 = .0755
.13 + .01
.14 + .00
.01 + .01
.06 + .03 .09 :: .03
CFRwass
.0112 + 0016
.0284 + .032
.17 : .00
.21 + .00
.04 : .01
Transformer
.0078 = 0056
.0072 = .0885
.2352 + .0171
.2505 + .0612
.0165 + .0118
.0813 # .0585
Random Forest
.0071 = .0032
.0084 = .0061
.0704 + .0144
.2518 + .0553
.0337 + .0164
.0909 + .0654
Table 7: Results + standard deviations for the Twins and Jobs benchmark datasets comparing with non-causal estimators. cATE absolute error the estimation of the Average Treatment Effect; Risk Policy risk; eATT absolute error on the estimation of the Average Treatment effect on the Treated; ws within sample; os out of sample. Methods included are GANTTE (Yoon al. 2018). CFRWass (Shalit et al., 2017). a simple transformer [Vaswani et al. 2017). and a random forest
Breiman 2001).
We follow closely the process in (Vowels et al., 2023a) for estimating the causal effect of shifting from one category of attachment style to another on depression. We also report the results for a subset of their analyses in Table 3 which use 'naive" estimator (comprising the bivariate linear model between the categorical treatment 'attachment' and the two outcomes), a targeted learning estimator specialized for causal inference which incorporates semi-parametric techniques (van der Laan & Starmans 2014 Vowels et al. 2023b), and our results using CAT. Note that there may be some minor diffcrences in their data preprocessing which we were not able to reproduce. In partic- ular, for each node in the DAG, the original authors reduced the dimensionality of the construct to be uni- mii-dinensional by taking the sum of the scores for each of the individual items. In contrast, padded all input variables so that they were the same dimensionality as the node with the highest di- mensionality. For instance, social distancing 'social dist' was found have 16 items, loneliness, which has only 3 items, was zero gero-padded to have 16 dimensions. The enables us to use all available information in the input. The dimensionalities / number of items for each construct are shown in Table wtttubrrrwwwrrdddmmeniioooo
We also use the DAG presented in (Vowels et al. (023a) which was the result of causal discovery process alongside domain expertise, this DAG is reproduced in Figure 8
For estimating the causal effect of a change in attachment style, we treat secure attachment as the reference group, and the three other categories (fearful, anxious, and avoidant) as the comparison categories. We estimate the effect of changing groups by setting the value of attachment to each of the 4 categories and record the shift in depression. Note the anxiety measure has 7-items (GAD-7; (Spitzer et al. 2006)) and the depression measure has (PHQ-9; (Kroenke et al. (002)) we take the average over the 9 dimensions estimate the final effect on depression.
Whilst standard errors can be readily obtained from the targeted learning ('TL') and bivariate linear modeling ( ('naive") approaches, the estimates from CaT. To do this, we repeated the estimation process 100 times using a sampling- with-replacement process and sampling fraction of 0.9 of the full dataset.
used boostrapping to obtain estimates of the standard errors for
27

---

Published as a conference paper at ICLR 2026
In Table we see that, despite CaT not being specialized to causal inference, the results are reason- ably close to those the TL estimator. Note that whilst our method uses the full dimensionality of each construct, the comparison methods (naive and TL) use the aggregated scales they have no other option but to do so. The standard errors are notably wider, which may be related to the use of bootstrapping for deriving them; in contrast to TI., which errors, the CaT is initialized and optimized from fresh for each bootstrap simulation, and trained on smaller subsample of the full dataset than TL. As such, we expect that this method for estimating the effect size would yield less tight intervals than TL. Furthermore, and as mentioned in the limitations section of the main text, the presence of 'loneliness' in the graph maybe also contribute estimation variance as the recursive inference process has to provide intermediate predictions for this mediator before using these predictions to, in turn, predict the outcome. Finally, the TL estimator includes the use of a Super Learner (van der Laan et al.] 2007), which involves cross-validation approach to finding the optimal set of candidate learners for the blug-n-estimmto prediction task, whereas we did not undertake any equivalent form of hyperparameter search for CaT, making TL challeng- ing approach to compete with. Regardless, it is difficult to make a comparison with this real-world application, because no ground truth is available.
be fit once and used to derive standard
Model
PEНE
PEНE
RF (Breiman
2001
+
1.79 = 0.03
CT Athey & Imbens
2016
4.81 0.18
0.21
CF Wagerr & Athey 2018
0.19 0.16
BART (Chipan al. 2010
2.18
0.15
CFR (Shalit et al. (201
2.05
0.20
DR-CFR (Hassam pour & Greiner 2020
0.24
0.26
GANITE Yoon et al. 2018)
0.56
0.64
CEVAE
Louizos al.
2017a
0.22
0.35
TEDVAE Zhang et al.
2021
1.75
0.17
CaT
+.22
4.26
0.23
CFCN
4.14 0.13
4.17 0.10
Table 8: Comparison of methods based on within- sathin-sample (ws) PEHE and out-of-sample (os) PEHE for the ACIC 2016 causal inference challenge (Dorie 2016) including results presented in ZZhang et al. 2021
S12. A CURIOUS RESULT
Table 7 presents a curious result for how the same scikit-learn (Pedregosa et al. 2011 random for- est (RF) (Breiman 2001) and transformer (Vaswani et al. 2017) used in the motivating experiments perform on the two causal inference benchmarks used in this work. Specifically, the RF and trans- former (using default/untuned hyperparameters) perform very competitively with set of methods specifically designed for causal inference. Further work is required to explore this result, but echos message from (Curth et al. 2021 that these benchmarks for causal inference should be reconsidered.
We also provide results for the Atlantic Causal Inference Challenge (ACIC) 2016 (Dorie) 2016) which is available hcre: ltttps://githen wwtp///www/w.cccceooeoooo
wwt.rriiiaccccoop///ooo rre htt///treeeematteeeeeeeeeee /ww/ter///rr//////e////ommaa
2016 6 and follow the procedure and present relevant results from (Zhang al. 2021). This dataset comprises 77 different data generating processes, each with 58 variables and 4802 observations. We undertake 60/30/10% train/validation/test split, although for hyperparameter tuning as such tuning is unrealistic in real-world causal inference. The results for the Precision in the Estimation of the Heterogeneous Efffct (PEHE) are given in Table S This table include results from the Causal Tree (CT: (Athey & Imbens) 2016)), Disentangled Counter- factual Regression (DR- CFR: (Hassanpour & Greiner (2020) and Treatment Effect Estimation with Disentangled Variational Autoencoder (TEDVAE; (Zhang et al. 2021)). The PEHE is defined as:
discard the 30% validation set
'PEHE =
n
(Ts
(18)
1=1
28

---

Published as a conference paper at ICLR 2026
where fq is the effect for observation i from the dataset estimated by the model.
For these results, the hyperparameters for CaT were: the number layers Tlayers = 3, input em- bedding dimension = 20, head size 10, number of attention heads 5, feed-forwart embedding / projection dimension 10, and dropout = 0.0. We used the same hyperparameters for CFCN as in Table
Once again, the RF also performs well on the ACIC 2016 benchmark as can be seen from Ta- ble(8 As successful causal inference requires to know the underlying DAG, the causal inference benchmarks make unbiased estimation possible with a set of default traditional estimators (like the random forest), because all of them assume the same graph (i.e. treatment variable, a set of possible confounders or precision variables, and an outcome). At least based on this preliminary result, we recommend researchers to least double-check the performance of standard estimators on these evaluation datasets, to validate that progress is indeed being made.
Future benchmarks might, for instance, include an element of causal discovery or causal feature selection, in order to really distinguish the challenge of causal inference from a regular supervised learning task. One the advantages of CaT and CFCN that makes the inference question flexible one can estimate the causal effect between any pair of nodes in the graph. Furthermore, this graph is specified explicitly and could, for example, represent the output of upstream causal discovery.
Based on these results, we are not confident that causal inference benchmark datasets are really high- lighting the key performance differences between proposed methods. We once again re-iterate that CaT and CFCN are not tuned for causal inference, although this represents a promising direction. standard causal inference benchmarks which assume basic
When evaluating CaT and CFCN DAG structure constituting a set of covariates, a treatment, and an outcome, the estimation problem is reduced simple choice of estimator, and using CaT is similar to using typical transformer.
S13. THEORETICAL FOUNDATION FOR CAT'S ARCHITECTURE
Theorem shows that, under our architectural assumptions, CaT defines Causal Bayesian Network which has conditional factorization which matches that implied by the supplied DAG and that inter- ventions implemented through the masking and embedding mechanism coincide with the standard g-forula.. Theorem and Corollary formalize the covariate shift and transportability claims: the causal mechanisms are invariant across domains, CaT's interventional queries (such as ATE) are invariant across domains, and predictions for target node are provably insensitive to shifts in variables outside its ancestor set.
ARCHITECTURAL ASSUMPTIONS AND CAUSAL SEMANTICS
Review of terminology and notation. The tensor X Bx|2IXC contains, for each mini-batch element, the C-dimensional raw features for every node in the DAG. The slice X; = X/:, is the raw input for node Z. Embeddings are produced by |Z| independent linear maps give X E R Bx where dg 2 and dg > (required for disentanglement) The learnable query embedding y IR 2|x The topologically sorted adjacency A E (0, | with Aji if and only if E pa(i).
Cross-attention structure. Unlike standard self-attention, CaT employs causally-masked cross- attention (CMCA) where queries Q >W Q bę derive from the learnable embedding in the first block (then from previous block outputs), while keys K = W + by and values V X Wy by always derive from the embedded input X This enables iterative refinement by extracting parental information at each layer.
Assumption (Hard parental mask). At every block, attention from token can only read from its parents pa(i). This is enforced through
QKT
0 = softmax
A T
V,
Vhg
29

---

Published as a conference paper at ICLR 2026
where (AT), if and only if pa(i). Since A has the diagonal (no self- self-loopss), (A )a2 0, preventing self-attention. Thus attention weights from any S - Z pa(i) to i are zero. Assumption (Per-node embedding). The embedding stage uses |2| independent linear layers, so the embedded token E, depends only on X, This preserves causal independence the input level and enables heterogeneous feature dimensions across nodes.
Assumption 3 (Tokenwise operations), Residual connections and the positionwise feed-forward FF)) act independently across tokens. Each CBlock preserves this structure:
OCMaa CMCA
--1 Block
X
E) + O
(19)
Block:
F (( OCMOA
(20)
Block
where O OBtock - 7 and batch normalization optional. No layer normalization is used it would rescale potentially calibrated interventions.
LAYERWISE DEPENDENCE AND LOCAL MARKOV PROPERTY
Let (n(") denote token i's hidden state after block &; by construction (:4)) feeds the prediction head for node i.
Proposition 1 (Parental sufficiency). Under Assumptions for any block e 2 1 and node i E Z,
there exists measurable function F(O) such that
h 10
(1 (pa(i))
Consequently, for any S Z pa(i).
(8)
11 X Xpa(t)-
Intuition. Token i aggregates only from its parents through causally-masked cross-attention and cannot attend to itself due to the zero diagonal allows the evolving query (from O to extract relevant parental information (keys/values from Xg). The residual connections maintain and refine token i's representation across layers without requiring self-attention. The hidden state at thus depends only parental inputs Xpa(i).
Since X is fed at every layer, each block
Remark (Strict causal dependency through masking). The architecture enforces strict causal de- pendency: token never directly accesses X, through attention. The DAG mask with zero diagonal ensures h (6) is constructed purely from parental information.
Let X, be the output of the final linear prediction head for node computed from n.( O Black,, after L CBlocks.
Lemma 1 (Local Markov property). Under Assumptions 73 with the final linear layer acting
independenlly per loken,
x, ( = go( F(L) (xp(i)))-
Since the mask structurally prevents self-attention, the prediction is necessarily a function of Xpa(t)) alone, yielding
Ped(2, xz) = =Pca(( x-p(s))-
The final prediction head reads token i's representation after all CBlocks. Since the DAG mask has zero diagonal, token i never accesses X; during forward passes. This architectural constraint enforces the local Markov property of Causal Bayesian Network with respect to g.
Intuition.
Theorem (Structural soundness and DAG factorization). Under Assumptions f dictive distribution defined by CaT factorizes according to the DAG G:
the joint pre-
Pcar(ai...... x22) = II poa(2; æi =paat)).
iEZ
In particular, CaT implements a Causal Bayesian Network over whose conditional distributions respect the parental sets in G.
30

---

Published as a conference paper at ICLR 2026
Intuition. The CMCA mask guarantees that each node's representation depends only its par- ents. The parallel attention computation thus yields a model that satisfies the local Markov property for all nodes. The standard result from graphical models then implies the global factorization with respect to g.
IDENTIFIABLE INTERVENTIONS VIA MASKED ATTENTION
We now present how CaT implements the truncated factorization (g-formula) used in causal infer- ence (Robins 1986
Definition 1. For intervention set I {(1s, =) (} with S Z, CaT' implements do(Xs (4) by:
Replacing embeddings for intervened nodes: X Embed, for S.
do E.j
2. Processing _ through all CBlocks with unchanged CMCA masks.
3. Reading outputs where intervened nodes retain their set values:
X
14 [8
j E S
Block
j $
Non-intervened nodes automatically incorporate upstream interventions since their representations depend only on (possibly intervened) parents. Proposition (Truncated product via CMCA). Under Theorem the CMCA architecture with Definition 7 directly implements the g-formula [Robins) 1986)
Pca(#2 xz dd(xs))= II Pearlzy ppa(i))
xpa(i) + + pals)
JEZ\S
Treatment effects like T == E dD= B do do(D=)o) are therefore computed by forward passes with different intervention values.
Intuition. Interventions work naturally in CaT: setting Xg.j to intervention values affects only downstream nodes (children of j), since node itself cannot self-attend. The mask ensures each node reads only from parents, so intervention effects propagate causally downstream in single forward pass. The resulting interventional distribution coincides with the truncated product of Causal Bayesian Network over G.
COVARIATE SHIFT, STRUCTURAL ROBUSTNESS, AND TRANSPORTABILITY
We now formalize how CaT behaves under covariate shift and connect this to standard assumptions in causal transportability. The key result is structural robustness property: assuming that the DAG correct and the causal mechanisms that are not causal ancestors of the target shift between domains. The usual transportability of interventional queries then follows a consequence.
invariant, CaT's predictions do not change when covariates
Consider two domains (or environments) E (8,4} with the same DAG G and the same structural conditionals
(x; Fpa(s) p(#: =pa(() for all E Z,
but possibly different observational distributions (Izz due to different exogenous noise or co- variate distributions.
Assumption (Consistent estimation of conditional distributions). CaT is trained in domain s with sufficient capacity and data that for all i E Z,
(8) Pear (x: Tip((t) N-xoc
p(a; -=p(())-
The architecture is then applied in domain without retraining.
31

---

Published as a conference paper at ICLR 2026
Theorem 2 (Structural robustness under covariate shift). Let Y be target node, and suppose do- main shift affects only a subset S _ Z \ An(Y) that contains no ancestors Y. Under Assump- tions] 7 and 4 for any fixed configuration of parental variables ppalY) the Cal predictor for Y is invariant across such covariate shift::
(e) Par (y [ Pear (y ppa())) for all paa))
(8)
In other words, CaT is robust to covariate shift in variables that not causal ancestors of Y.
Intuition. Due to the CMCA mask, the representation for can only depend on its parents and their ancestors. Variables outside AY()) never influence Y's token representation. contrast, non- causal transformers/network. are vulnerable to this: in generic neural network, the representation of Y free depend on any input variable, so non-ancestors can spuriously influence predictions and make them sensitive to distributional changes that do not alter the underlying mechanisms.
Corollary (Transportability of identified causal queries). Under Assumptions D+ and f any
causal query that is identified by the g-formula (Robins, 1986) in G, such as the average treatment
effect
E dD<D==1) E E
do do(D=Ho)).
has Cal estimate that is invariant across domains and in the carge-sample limit:
CaT
(s)
N
and
Tall
(9) N-000
Equivalently, as long as the mechanisms p(a; Fpplit) are shared across domains, Cal's interven- tional estimates are transportable.
32


Reviews:
[{"review_id": "fehIX7jQrx", "reviewer": "Reviewer_E5kE", "summary": "The paper introduces Causal Transformers (CaTs), which incorporates causal graphs as masks inside the attention mechanism of transformers, which improves their performance in terms of robustness and interpretability. CaTs are capable of estimating causal queries such as treatment effects, and they are shown experimentally to outperform methods that do not take causality into account, while being able to compete with methods that do.", "questions": "8. What makes CFCN an interesting baseline, as opposed to just another method developed by the paper? It seems like an oddly specific inclusion to the paper.\n\n9. The paper claims in Sec. 4.2 that it is not expected for CaT and CFCN to excel compared to models specifically tuned for causal inference tasks. I would have listed this as a weakness, but I believe that perhaps one strength of these models is their generative capabilities across many different queries. Could the authors confirm if this is the case (e.g., could CaTs produce novel samples of both the observational and interventional data distributions)? If so, how does this model compare to common approaches used in the causal generative modeling literature such as [1][2][3]?\n\n[1] Kocaoglu, et al., “CausalGAN: Learning Causal Implicit Generative Models with Adversarial Training” (2017).\n\n[2] Xia, et al., “The Causal-Neural Connection: Expressiveness, Learnability, and Inference” (2021).\n\n[3] Saremi, “Neural Network Parameter-optimization of Gaussian Pre-marginalized Directed Acyclic Graphs” (2023).", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. Incorporating graphs into the attention mechanism ensures that the causal constraints encoded in the graphs (and therefore the resulting causal inferences) are guaranteed to hold. This is a strong advantage over methods that attempt to achieve causality through regularization. The adjustments such as the omission of layer norm are done carefully to preserve this causal integrity.\n\n2. Assumptions are stated clearly.\n\n3. The provided experimental results seem promising.", "weaknesses": "4. The contributions of the paper are largely empirical and leave much to be desired from the theoretical aspects of the models. Most notably, I believe it would be a core contribution to include a theorem that guarantees that causal constraints (such as those from Causal Bayesian Networks) are correctly enforced in the CaT architecture, and therefore, the estimated queries will be correct. Another important contribution would be to support the claim that causal models like CaT can resolve issues related to covariate shift, so it would be fundamental to include a theorem that shows what kinds of guarantees CaT can provide across domains (perhaps using causal transportability theory).\n\n5. The training details of CaT are somewhat unclear. In particular, it would be great to see the exact details of how the loss function is implemented, and how causality plays a role in it.\n\n6. It is unclear what kinds of interventional queries CaTs are capable of estimating, and how each of those queries are particularly calculated through the CaT architecture. It seems that one weakness is that CaTs are not capable of estimating ALL interventional queries (perhaps only the causal effect of treatment on target).\n\n7. Notation is not introduced properly. It is not clear what $B$, $Z$ and $C$ are.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "I6sPIR381w", "reviewer": "Reviewer_Ac6L", "summary": "This article focuses on embedding causal knowledge in the form of a Directed Acyclic Graph (DAG) into a transformer architecture. The resulting model is designed to perform causal inference tasks. To achieve this, the authors introduce a causal transformer that enables this integration. They also implement a fully connected neural network that respect DAG structure as an additional baseline. Using a toy example, the authors demonstrate the importance of having a causal structure and a model that respects this causal structure. The work is further validated through real-world experiments.", "questions": "* Is there a way to amortize this approach across different DAGs or datasets?\n* Are you confident that the differences between CAT and CCFCN are not due to insufficient hyperparameter tuning?\n* Is it possible to create synthetic datasets for these models that include high-dimensional data?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The authors provide clear motivation and demonstrate it with a compelling example that compares predictive power and treatment effect estimation.\n* They propose a novel modification of the attention mechanism for embedding DAGs into the model.\n* They introduce a formulation for incorporating DAG structure into neural networks.\n* They benchmark their approach against other causal inference methods.", "weaknesses": "* The architecture needs to be trained for each DAG on each dataset, which can limit the method from being applied to larger datasets.\n\n* While the authors describe the ability to handle multidimensional embeddings as an advantage, the experiments are conducted on tabular datasets. It would be valuable to see whether the transformer architecture can excel in such settings.\n\n* The authors quantitatively justify not using normalization; however, it is not clear whether the pros would outweigh the cons.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "DD3SLjh7Bi", "reviewer": "Reviewer_5yaS", "summary": "This paper introduces causal transformers (CaTs) for causal prediction. CaTs encode the information of a given causal graph in the architecture, offering a strong approximation power for a given causal model. Detailed experiments show the superior performance of CaTs in term of ATE estimation.", "questions": "Clarify questions:\n\n1. Does CaTs apply to all kinds of causal graphs or just the graph in Figure 2? If it applies to all causal graphs, the identification formula (1) seems to be incorrect since $L_{\\text{set}}$ may contain colliders. \n\n2. How do you model hidden confounders in Figure 3 (b)?\n\n3. Can the authors explain the use of learnable embedding $\\gamma$? In particular, why does $Q$ depend on $\\gamma$ but $K,V$ do not? \n\n4. The authors mention that they do not include normalization in CaTs because it may impact the capacity for precise estimation. In a standard transformer, LayerNorm is applied to each variable (token) independently across its feature dimension. It does not typically \"contaminate\" information between variables. Can the authors provide a more detailed explanation?\n\nExperiment questions:\n\n5. The authors include a standard transformer in Table 1, but do not include it Table 2. Can you include the results too?\n\n6. In Table 1, can you explain why misspecified DAGs can yield low MSE? This is quite counterintuitive.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 2, "strengths": "1. A key advantage highlighted is that by constraining the model to a (correct) causal DAG, it learns to ignore spurious correlations in the training data. This makes the model more robust and stable when faced with distributional shifts (i.e., covariate shift), a common failure point for traditional machine learning models.\n\n2. The experiments show superior performance over the traditional random forest models even if the DAG is misspecified. \n\n3. Unlike many causal inference methods designed for single-value, tabular data, the Causal Transformer (CaT) architecture is built to handle high-dimensional inputs. It can accept a sequence of embeddings, making it applicable to complex, multimodal data like text, images, or multi-item questionnaires.", "weaknesses": "1. The literature review part is not complete. There have been many works in the literature that consider using neural networks for counterfactual estimation. For instance, [1] establishes the connection between neural nets and causal models and constructs feed-forward networks for discrete causal models. [2] establishes similar results for causal models with mixed variables (continuous and discrete) and uses it for partial identification. Similar to this work, [3] proposes causal transformers for counterfactual estimation. There are many other works on this topic as well. Please include those works in the related work part. \n\n2. Some details are unclear in the paper. Please see the question part. \n\n3. In the experiment, the author compared the ATE given by different models. However, it is well-known that the Double Robust (DR) estimator provides a more robust way for estimation. It will provide more practical implications if the author can provide a comparison (using the same propensity model).  \n\n[1] Xia, Kevin, et al. \"The causal-neural connection: Expressiveness, learnability, and inference.\" Advances in Neural Information Processing Systems 34 (2021): 10823-10836.\n\n[2] Tan, Jiyuan, Jose Blanchet, and Vasilis Syrgkanis. \"Consistency of neural causal partial identification.\" Advances in Neural Information Processing Systems 37 (2024): 68956-68999.\n\n[3] Melnychuk, Valentyn, Dennis Frauen, and Stefan Feuerriegel. \"Causal transformer for estimating counterfactual outcomes.\" International conference on machine learning. PMLR, 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)