Adjustment for Confounding using Pre-Trained Representations
Authors: Rickmer Schulte
David Rügamer
Thomas Nagler
Venue: ICML 2025 poster (2025)
Keywords: ATE
Pre-trained Models
Doubly Robust Estimation
Intrinsic Dimensions
Primary area: general_machine_learning->causality

Abstract: There is growing interest in extending average treatment effect (ATE) estimation to incorporate non-tabular data, such as images and text, which may act as sources of confounding. Neglecting these effects risks biased results and flawed scientific conclusions. However, incorporating non-tabular data necessitates sophisticated feature extractors, often in combination with ideas of transfer learning. In this work, we investigate how latent features from pre-trained neural networks can be leveraged to adjust for sources of confounding. We formalize conditions under which these latent features enable valid adjustment and statistical inference in ATE estimation, demonstrating results along the example of double machine learning. We discuss critical challenges inherent to latent feature learning and downstream parameter estimation arising from the high dimensionality and non-identifiability of representations. Common structural assumptions for obtaining fast convergence rates with additive or sparse linear models are shown to be unrealistic for latent features. We argue, however, that neural networks are largely insensitive to these issues. In particular, we show that neural networks can achieve fast convergence rates by adapting to intrinsic notions of sparsity and dimension of the learning problem.

Adjustment for Confounding using Representations
Rickmer Schulte!! David Rügamer" 12 Thomas Nagler'
Abstract
and in some cases even impossible obtain, either due to ethical or economic reasons. This often implies relying on observational data, typically subject to (unmeasured) confounding- ((hideen factors that affect both the exposure and the outcome. To overcome this issue of confounding and to obtain unbiased estimates, several inferential methods have been developed to properly adjust the ATE estimation for confounders. One approach that has garnered significant attention in recent years is the debiased/double machine learning (DML) framework (Chernozhukov et al., 2017; 2018), which allows the incorporation of machine learning methods to adjust for non-linear or complex confounding effects in the ATE estimation. DML is usually applied in the context of tabular features and was introduced for ML methods tailored to such features. However, confounding information might only be present in non-tabular data, such as images or text.
There is growing interest in extending average treatment effect (ATE) estimation to incorporate non-tabular data, such images and text, which may act sources of confounding. Neglecting these effects risks biased results and flawed sci- entific conclusions. However, incorporating non- tabular data necessitates sophisticated feature ex- tractors, often in combination with ideas of trans- fer learning. In this work, we investigate how la- tent features from pre- eraitriieed neural networks be leveraged to adjust for sources of confounding. We formalize conditions under which these latent features enable valid adjustment and statistical in- ference in ATE estimation, demonstrating results along the example of double machine learning. We discuss critical challenges inherent to latent feature learning and downstream parameter esti- mation arising from the high dimensionality and non-identiiiability of representations. Common structural assumptions for obtaining fast conver- gence rates with additive or sparse linear models are shown to be unrealistic for latent features. We argue, however, that neural networks are largely insensitive these issues. In particular, that neural networks rates by adapting to intrinsic notions of sparsity and dimension of the learning problem.
Non-tabular Data as Sources of Confounding Espe- cially in medical domains, imaging is a key component of the diagnostic process. Frequently, CT scans X-rays are the basis to infer a diagnosis and a suitable treatment for a patient. However, as the information in such medical images often also affects the outcome of the therapy, the information in the image acts as a confounder. Similarly, treatment and health outcomes are often both related to pa- tient's files, which are typically in text form. Consequently, ATE estimation based on such observational data will likely be biased if the confounder is not adequately accounted for. Typical examples would be the severity of a disease or fracture. The extent of a fracture impacts the likelihood of surgical or conservative therapy, and the severity of a disease may impact the decision for palliative or chemother- apy. In both cases, the severity will likely also impact the outcome of interest, e.g., the patient's recovery rate. An- other famous example is the Simpson's 's Paradox observed in the kidney stone treatment study of Charig et al. (1986). The size of the stone (information inferred from imaging) impacts both the treatment decision and the outcome, which leads to flawed conclusions about the effectiveness of the treatment if confounding is not accounted for (Julious & Mullee, 1994).
show achieve fast convergence
1. Introduction
Causal inference often involves estimating the average treat- ment effect (ATE), which represents the causal impact of an exposure on outcome. Under controlled study se- tups of randomized controlled trials (RCTs), valid inference methods for ATE estimation are well established (Deaton & Cartwright, 2018). However, RCT data is usually scarce
'Department of Statistics, LMU Munich, Munich, Germany "Munich Center for Machine Learning (MCML), Munich, Ger- many. Correspondence to: Rickmer Schulte <schulte@ una- muenchen.de>.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Contemporary Applications While the DML framework provides a solution for non-linear confounding, previous

---

Adjustment for Confounding using Pre- Tre-TTrined Representation
Maximum Likelihood Estimation or the DML approach es- timate both nuisance functions 9 and m. These methods thus only require the product of their estimation errors to converge at Vn-rate (Robins & Rotmitzyy, 1995; van der Laan & Rubin, 2006; van der Laan & Rose, 201 1; Cher- nozhukov et al., 2017; 2018). However, even this can be hard to achieve, given the curse of dimensionality when considering the high-dimensionality of non-tabulas data W such as images. Especially given the often limited number of samples available in many medical studies involving im- ages, estimating m and as a function of W, e.g., via neural networks, might not be feasible or overfit easily. To cope with such issues, a common approach is to adopt ideas from transfer learning and use pre-trained neural networks.
W!
T
Y
Figure 1. Schematic (left) and DAG visualization (right) of the effect of a treatment T on outcome Y that is confounded by non- tabular data W (c.g. information from medical imaging).
examples demonstrate that modern data applications require extending ATE estimation to incorporate non-tabe non-tabular data. In contrast to traditional statistical methods and classical machine learning approaches, information in non-tabular data usually requires additional feature extraction mecha- nisms to condense high-da -ighesniensi inputs to the relevant information in the data. This is usually done by employing neural network-bascd based approaches such as foundation mod- els or other pre-traine pre-trained neural networks. While it may seem straightforward use such feature extractors extract la- tent features from non-tabular data and use the resulting information in classical DML approaches, we show that this necessitates special caution. In particular, incorporat- ing such features into ATE estimation requires overcoming previously unaddressed theoretical and practical challenges, including non-identifiability, high dimensionality, and the resulting limitations of standard assumptions like sparsity.
Our Contributions In this paper, we discuss under what conditions pre-trained representations Z == sp(W) obtained from pre-trained neural networks o can replace W in the estimation of nuisance functions (2) and (3). Although the dimensionality of Z is usually drastically reduced compared to W, one major obstacle from theoretical point of view is that representations can only be learned up to invertible lin- ear transformations (e.g.. rotations). We argue that common assumptions allowing fast convergence rates, e.g., sparsity or additivity of the nuisance function, are no longer reason- able in such settings. In contrast, we build on the idea of low intrinsic dimensionality of the pre- pre-trained representations. Combining invariance of intrinsic dimensions and functional smoothness with structural sparsity, we establish conditions that allow for sufficiently fast convergence rates of nuisance function estimation and, thus, valid ATE estimation and inference. Our work, therefore, not only advances the the- oretical understanding of causal inference in this context but also provides practical insights for integrating modern machine learning tools into ATE estimation.
Given n independent and identically dis- tributed (i.i.d.) observations of (T, W, Y), are interested in estimating the ATE of a binary variable T E {0. some outcome of interest Y E R while adjusting for some source of confounding W e W (cf. Figure 1). W is pre- treatment data from some potentially complex sampling space W that is assumed to be sufficient for adjustment. The definition of sufficiency will be formalized in Section 3.1. Under positivity and consistency assumption- the standard assumptions in causality- the target parameter of interest can be identified as
Problem Setup
2. Related Work
The DML framework was initially proposed for tabular fea- tures in combination with classical machine learning meth- ods (Chernozhukov et al., 2017: 2018). Several theoretical and practical extensions to incorporate neural networks have been made with a focus on tabular data (Shi et al., 2019; Farrell et al., 2021; Chernozhukov et al., 2022; Zhang & Bradic, 2024). Additionally, there is a growing body of research that aims to incorporate non-tabular data as adjust- ment into DML (Veitch et al., 2019; 2020: Klaassen et al., 2024). While the latter directly incorporates the non-tabular data in the estimation, none of them discuss conditions that would theoretically justify fast convergence rates necessary for valid inference. A different strand of research instead uses either derived predictions (Zhang et al., 2023; Battaglia et al., 2024; Jerzak et al., 2022; 2023a;b) or proxy variables (Kuroki & Pearl, 2014; Kallus al., 2018; Miao et al., 2018;
ATE: = BE
EY 0, W]]. (L)
While there are many well-known ATE estimators, most require to estimate either the outcome regression function
g(t, :) = 1, W w]
or the propensity score
(t|/w)= == P W w)
(3)
at parametric rate Vn. Doubly robust estimators such as the Augmented Inverse Probability Weighted, the Targeted
2

---

Adjustment for Confounding using Pre- Pre-Trained Representation
Mastouri et al., 2021; Dhawan et al., 2024) in downstream estimation. In contrast to these proposals, consider the particularly broad setup of using pre-trained representations for confounding adjustment. Given the increasing popular- ity of pre-trained models, Dai et al. (2022) and Christgau & Hansen (2024) establish theoretical conditions justify- ing the use of derived representations downstream tasks, which we will review in the next section. The idea of a low intrinsic dimensionality of non-tabular data and its latent representations to explain the superior performance of deep neural networks in non-tabular data domains has been ex- plored and validated both empirically (Gong et al., 2019; Ansuini et al., 2019; Pope et al., 2021; Konz & Mazurowski, 2024) and theoretically (Chen et al., 2019; Schmidt-Hieber 2019; Nakada & Imaizumi, 2020). By connecting several of those theoretical ideas and empirical findings, our work establishes a set of novel theoretical results and conditions that allow to obtain valid inference when using pre-trained representations in adjustment for confounding.
(iii) P-ODS (Outcome Distribution Sufficient) if:
Y Lp Z Z.
Remark 3.2. If 2 is P-ODS, it is also called sufficient embedding in the literature (Dai et al., 2022).
The three conditions in Definition 3.1 place different restric- tions on the nuisance functions (2) and (3). While P-ODS is most restrictive (followed by P-OMS) and thus guaran- tees valid downstream inference more generally, the strictly weaker condition of P-validity is already sufficient (and in fact necessary) to guarantee that Z E Z is valid adjust- ment set in the ATE estimation (Christgau & Hansen, 2024), Thus, any pre-trained representation Z considered in the following is assumed to be at least P-valid.
3. Properties of Pre-Trained Representations
Given the high dimensional nature of non-tabular data, to- gether with the often limited number of samples available (especially in medical domains), training feature extractors such as deep neural networks from scratch is often infeasi- ble. This makes the use of latent features from pre-trained neural networks a popular alternative (Erhan et al., 2010). In order to use pre-trained representations for adjustment in the considered ATE setup, certain conditions regarding the representations are required.
W
pp())
Z
2(Z)
Figure 2. Schematic visualization of pre-trained neural network (--) and representations (pW)).
3.2. Non-Identifiability under ILTs
In practice, the representation 7 = (W) is extracted from some layer of pre-trained neural network yp This infor- mation does not change under bijective transformations of Z, so the representation Z itself is not identifiable. We ar- gue that, in this context, non-identifiability with respect to invertible linear transformations (ILTs) is most important. Suppose Z = (WW) is extracted from a deep network's fth layer. During pre-training the network further processes Z through model head (ZZ), schematically depicted in Figure 2. The model head usually has the form (><(AZ b) where A, b are the weights and biases of the éth layer, and Φ>t summarizes all following computations. Due to this structure, any bijective linear transformation Z + QZ can be reversed by the weights A + À AQ-1 so that the networks (A +6) and &> +-) have the same output.
3.1. Sufficiency of Pre-Trained Representations
Given any pre-trained model yp, trained independently of W on another dataset, we denote the learned (last-layer) rep- resentations Z (W)). Due the non-identifiability of 7 up to certain orthogonal transformations, further dis- cussed in Section 3.2, we define the following conditions for the induced equivalence class of representations 2 fol- lowing Christgau & Hansen (2024). For this, we abstract the adjustment as conditioning on information in the ATE estimation, namely conditioning on the uniquely identifiable information contained in the sigma-algebra a (Z) generated by any Z E Z (see also Appendix A.I for a special case).
Definition 3.1. [Christgau & Hansen (2024)) Given the joint distribution P of (T, W, Y), sigma-algebra (2) of Z, and t E {0, 1}, we say that any Z E Z is
Definition 3.3 (Invariance to ILTss). Given a latent rep- resentation Z, we say that model (head) Φ with pa- rameters & E E is non-identifiable up to invertible lin- ear transformations if for any invertible matrix Q Rdd a e E: (Q() [ dz(Z).
(i) P-valid if:
Ep Ey = EP W]]
(ii) P-OMS (Outcome Mean Sufficient) if:
Important examples of ILTs are rotations, permutations, and scalings of the feature space as well as compositions thereof.
=, Ep (Pass))
3

---

Adjustment for Confounding using Pre-Trained Representation
Smoothness
+ Additivity
Sparsity & Linearity
Intrinsic Dimension
Stone (1982)
Stone (1985)
Raskuttii et al. (2009)
Bickel & Li (2007)
0 2++0
O(n 2++1
O O()plog(d)/n)) < d
O 2TTM ) dM <d
Table 1. Assumptions and related minimax convergence rates of the estimation error
4. Estimation using Pre-Trained Representations
assumed due to its invariance to ILTs. If the feature dimen- sion is large, however, unrealistic amount of smoothness would be required to guarantee fast convergence rates (e.g.. of order n n-1/4). This necessitates additional structural or distributional assumptions.
The previous section discussed sufficient and necessary (in- formation theoretic) conditions for pre-trained representa- tions, justiffing their usage for adjustment in downstream tasks. The following section will discuss aspects of the func- tional estimation in such adjustments. Valid statistical infer- ence in downstream tasks usually requires fast convergence of nuisance function estimators. However, obtaining fast convergence rates in high- -igeesiienii estimation problems is particularly difficult. We argue that some commonly made assumptions are unreasonable due to the ooniddetifiability of representations. We discuss this in the general setting of nonparametric estimation as described in the following.
4.2. Structural Assumptions II: Additivity & Sparsity
The common structural assumption is that f is additive, f(x) Σ fz(z)), i.e., the sum of univariate s-smooth functions. In this case, the minimax convergence rate re- duces to n žनन (Stone, 1985). Another common approach is to rely on the idea of sparsity Assuming that f is p-sparse implies that it only depends on p < min d) features. In case one further assumes the univariate functions to be linear in each feature, i.e. f(x) Σj= Byæz with coef- ficient 33 E R. the optimal convergence rate reduces to vp log (d/p//)// (Raskutti et al., 2009).
The Curse of Dimensionality The general problem in nonparametric regression is to estimate some function f in the regression model
It can easily be shown that the previously discussed condi- tions are both preserved under permutation and scaling. But as the following lemma shows, sparsity and additivity of f are (almost surely) not preserved under generic ILTs such as rotations.
Y = f(x)+ +E
(4)
with outcome Y E R. features X Rd, and error E~ (O, d²). The minimax rate for estimating Lipschitz functions is known to be n - (Stone, 1982). This rate becomes very slow for increasing d, known the curse of dimensionality. Several additional structural and distri- butional assumptions are commonly encountered to obtain faster convergence rates in high dimensions.
Lemma 4.2 (Non- Inv-Invariance of Additivity and Sparsity under IL.Ts). Let f Rdd -> TR be a function of x E RRd, We distinguish between two cases:
(i) Additive: f(x) = -E (I()) with univariate func- tions fj :RR + R. and at least one f, being non non-linear.
4.1. Structural Assumption I: Smoothness
(ii) Sparse Linear: [xx = Σ5= Byzj, where }3 E R and at least one (but not all) 3, =0.
A common structural assumption is the smoothness of the function / in (4), i.e., the existence of x bounded and contin- uous derivatives. Most convergence rate results assume at least some level of smoothness (see Table 1). The following lemma verifies that this condition is also preserved under ILTs.
Then, for almost every Q drawn from the Haar measure on the set of ILTs, it holds:
(i) If f is additive, then h = f is not additive. (ii) If f is sparse linear, then h = f Q-1 is not sparse.
Lemma 4.1 (Smoothness Invariance under ILTss). Let D C IR. be open set, f D - R.. be an s-smooth--uuction on D, and Q by any ILT. Then h = f o 1 Q(D) + R is also s-smooth on the transformed domain (D).
Given the non-identifiability of representations with respect to ILT's and the non-invariance result of Lemma 4.2, any additivity or sparsity assumption about the target function if of the latent features seems unjustified. An example of this rotational non-invariance of sparsity is given in Fig- ure 3. This also implies that learners such as the lasso (with
The proof of Lemma 4.1 and subsequent lemmas of this section are given in Appendix A.
The lemma shows that a certain level of smoothness of a function defined on latent representations may reasonably be
4

---

Adjustment for Confounding using Pre- Pra-Trained Representation
tion Z but for the entire equivalence class of latent represen- tations induced by the class of ILTs. This is crucial given the non-identifiability of latent representations, highlighting the importance of low intrinsic dimensions (IDs).
Non- 2000--000-- contiiiints
Lstrrr Pattrr Dinencer
X0
700
Deep Networks Can Adapt to Intrinsic Dimensions Re- cently, several theoretical works have shown that DNNs can adapt to the low intrinsic dimension of the data and thereby attain the optimal rate of 2х+dx (Chen et al., 2019; 9; Schmidt-Hieber, 2019; Nakada & Imaizumi, 2020; Kohler et al., 2023). In Section we present a new conver- gence rate result that builds on the ideas of low ID and a hierarchical composition of functions particularly suited for DNNs.
random feature motaiins
Figure 3. Non-zero coefficients of linear classifier on latent fea- tures, showing that sparsity is lost with an increasing number of random feature rotations.
underlying sparsity assumption), tree-based methods that based axis-aligned splits (including corresponding boosting methods), and most feature selection al algorithms are not ILT-invariaat, Further examples can be found in Ng (2004).
5. Downstream Inference
The manifold assumption alone, however, cannot guarantee sufficient approximation rates in our setting. Even the manifold dimension d' much smaller than the ambient dimension d (for example, dM 30), unreasonably high degree of smoothness would need to be assumed to allow for convergence rates below n-1/1, In what follows, we give more realistic assumption to achieve such rates, In particular, we combine the low-dimensional manifold structure in the feature space with a structural smoothness and sparsity assumption on the target function.
4.3. Distributional Assumption: Intrinsic Dimension
While the previous conditions are structural assumptions regarding the function f itself, faster convergence rates can also be achieved by making distribution assumptions about the support of f. A popular belief that the d-dimensional data X Id lie on close to a low-dimensional manifold M with intrinsic dimension d. This relates to the famous manifold hypothesis that many high-dimensional data con- centrate low-dimensional manifolds (Fefferman et al., 2016, e.). There is strong empirical support for this as- sumption, especially for con-larular modalities such as text and images, see Appendix B. Given that dm < d, and much again assuming f to be s-smooth, this can lead faster convergence rate of n अनM (Bickel & Li, 2007), as is independent of the dimension d of the ambient space.
5.1. Structural Sparsity on the Manifold
Kohler Langer (2021) recently derived convergence rates based on the following assumption. Definition 5.1 (Hicrarchical composition model, HCM).
Similarly to Lemma 4.1, the following lemma shows the invariance of the intrinsic dimension of manifold with respect to any ILT of the coordinates in the d-dimensional ambient space.
(a) We say that f: Rd -> IR satisfies a HCM of level 0, if [(x) = zi for some E (I.....d)..
(b) We say that f satisfies HCM of level k 2 1, if there is a s-smooth function h: R - R. such that
Lemma 4.3 (Intrinsic Dimension Invariance under ILTs).
Let M C Idd be a smooth manifold of dimension d S d. For any ILT Q, the transformed set
f(x) = (a(ma(t),....hh(((()).
where hi...., hp Rd + R are HCMs of level k 1.
Q(M) (Qx x
is also smooth manifold of dimension d.M Remark 4.4. Put differently, in case the latent representa- tions Z E ddd lie on a diy-dimmnsionnl smooth manifold M, then the IL-transformed representations ((Z) also lie on a smooth manifold Q(M) dimension d.m-
The collection of all pairs (s,p) E R appearing in the specification is called the constraint of the HCM.
An illustration of Definition 5.1 given Appendix B.2. The assumption includes the case of sparse linear and (gen- eralized) additive models more general. Kohler & Langer (2021) and Schmidt-Hieber (2020) exploit such a structure to show that neural networks can approximate the target function at a rate that is only determined by the worst-case pair (s.p) p) appearing in the constraint set. It already follows from Lemma 4.2 that the
special case but is much
Summarizing previous results, the structural and distribution assumptions of smoothness and low intrinsic dimensionality are invariant with respect to any ILT of the features. Hence, as opposed to additivity or sparsity, the two conditions hold not only for a particular instantiation of latent representa-
5

---

Adjustment for Confounding using Pre- Trained Representation
constraint set of such model is not invariant to ILTs of the input space. Furthermore, the assumption does not ploit the potentially low intrinsic dimensionality of the input space. To overcome these limitations, we propose a new assumption combining the input space's 's manifold structure with the hierarchical composition model.
The result shows that the convergence rate of the neural networks is only determined by the worst-case pair (s, p) appearing in the constraint set of the HCM and the embed- ding map v. The theorem extends the results of Kohler & Langer (2021) in two ways. First, it allows for more gen- eral loss functions than the square loss. This is important since classification methods are often used to adjust for con- founding effects. Second, it explicitly exploits the manifold structure of the input space, which may lead to much sparser HCM specifications and dramatically improved rates.
Assumption 5.2. The target function fo can be decom- posed as fo = fo vh; where M is a smooth, compact, dm-diimnsioonal manifold, eh: .M- + R2 is sy -s--mhhhh and is a HCM of level & E N with constraint set P.
Whitney's embedding theorem (e.g., Lee & Lee, 2012, Chap- ter 6) allows any smooth manifold to smoothly embedded into g2dx, This corresponds mapping v with Sy oo and p 2d in the assumption above. If not all informa- tion in the pre-trained representation Z is relevant, however, can be much smaller. Importantly, Assumption 5.2 is not affected by ILTs.
5.3. Validity of DML Inference
In the previous sections, explored plausible conditions under which the ATE is identifiable, and DNNs can estimate the nuisance functions with fast rates. We now combine our findings to give a general result for the validity of DML from pre-trained representations.
For binary treatment T E {0, 1} and pre-trained representa- tions Z, we define the outcome regression function
Lemma 5.3 (Invariance of Assumption 5.2 under ILTs). Let Q be any ILT. If fo satisfies Assumption 5.2 for a given P and (ss, dm), ) then fo = fo o Q-1 satisfies Assumption 5.2 with the same P and (sp, d.)..
g 2) = B| t, Z 2],
and the propensity score
5.2. Convergence Rate of DNNs
m(z) = P[ 1|2 = 2].
We now show that DNNs can efficiently exploit this struc- ture. Let (Yi 1 be i.i.d. observations and & be a loss function. Define
Suppose we are given an i.i.d. sample (Y,, Z, T.) DML estimators of the ATE are typically based on cross-fitting procedure. Specifically, let U 4 {1..... n be partition of the sample indices such that // - 1/K.
fo = arg min B(((f(Z),))) f: R
Let g(()) and (k) denote estimators of g and m computed only from the samples (Y,, Zi, Tı) Ti)sgIe Defining
1 Σ (((()))))
f = arg min fEF(Im))))
ATE (k)
1 Σ
p(T,, Y,, Y,
m (*))
[4N
where F(L, v) is the set of feed-forward neural networks with L layers and v neurons per layer. Let Z ~ Pz and define the dP(). We make the following assumption on the loss function L.
with orthogonalized score
of function f as II (P2))) =
P(T,YY, Z45 g. m) = g(( 2n) - g(0, 2,)
T,(Y, =g(1, + (1 T.)(Y, - g(0, Zz))
+ m(Z))
Assumption 5.4. There a, b E (0, oc) such that
1- 1-m(2) Z.)
[L(((Z), B[((fo(Z), Y)] E [a, 6).
the final DML. estimate of ATE is given by
[| TollZ,(PP)
1
(R)
ATE
ATE
K
Assumption 5.4 is satisfied for the squared and logistic loss, among others (e.g., Farrell et al., 2021, Lemma 8).
We need the following additional conditions. Assumption 5.6. It holds
Theorem 5.5. Suppose Assumption 5.2 and Assumption 5.4 hold. There are sequences Ln Vn and a corresponding sequence of neural network architectures F(In, pn) such that (up to log factors)
max B[]g((, 2) <0, E|||*] < xo, E (0.1)
E
0, Pr(m(Z) E (e. - -))]= =1,
fl - follzg(()) = =Op
max ((.))EPU())ddM)
n 20t0
for some & > 0.
6

---

Adjustment for Confounding using Pre- Pre-Trained Representation
The first two conditions ensure that the tails of Y and g(l, Z) not too heavy. The second two conditions required for the ATE to be identifiable.
Salmnnr ATE
Trea
Theorem 5.7. Suppose the pre- pre-trained representation is P-- valid, Assumption 5.6 holds, and the outcome regression and propensity score functions g and m satisfy Assumption 5.2 with constraints Pg (8y, d and Pmn (s's, d,),, respec- tively. Suppose further
12
Oracle
S-Leamer
S-Leamar
SLeamee (Losso)
DML
DML (Lasso)
x min s' 1 (sp)EPyulsyddu P (5 (sts aMd p
[Lneer)
(RF)
(Linsar)
(RF)
min
(5)
4'
Figure 4. Label Confounding: Comparison of ATE estimators on the IMDb dataset. DMI. and S-Learner pre-trained represen- tations. Point estimates and 95% CIs are depicted.
and the estimators g(k) and ( are DNNs as specified in Theorem 5.5 with the restriction that in is clipped away from 0 and 1. Then
obtained by passing the images through a pre-trained convo- lutional neural network and extracting the 1024-dimensional last hidden layer features of the model. We use the pre- trained Densenet-121 model from the TorchXRay Vision library (Cohen et al., 2022), which was trained on several publicly available chest X-ray datasets (Cohen et al., 2020). Further details on the datasets and pre-trained models used in our experiments are provided in Appendix C.I.
Vi(ATE- ATE)- - O'
where o² Elp(T; Y,, m)
Condition (5) is our primary regularity condition, ensuring sufficiently fast convergence for valid DML inference. It characterizes the necessary trade-off between smoothness and dimensionality of the components in the HCM. In par- ticular, it is satisfied when each component function in the model has input dimension less than twice its smoothness.
Confounding Setup For both data applications, we sim- ulate treatment and outcome variables while inducing con- founding based on the labels. As example, for the modi- fied image dataset, children with pneumonia have higher chance of receiving treatment compared to healthy children. In contrast, pneumonia negatively impacts the outcome vari- able. The same confounding is present in our modified text dataset. Hence, the label creates a negative bias in both ATE settings not properly accounted for. Further details about the confounding setups provided in Appendix C.2.
6. Experiments
In the following, we will complement our theoretical results from the previous section with empirical evidence from several experiments. The experiments include both images and text as non-tabular data, which act as the source of confounding in the ATE setting. Further experiments can be found in Appendix D.
ATE Estimators We compare the performance of DML using three types of muisance estimators: linear models with and without L1 de-penalizttion (Lasso/ Linear), as well as ran- dom forest (RF). For comparison, we also include another common causal estimator, called S-Leaner, which only esti- mates the outcome function (2) (details in Appendix C.3). In each of the simulations, estimators facilitate the information contained in the non-tabular data to adjust for confounding by using the latent features from the pre-trained models in the estimation. As a benchmark, we compare the estimate to the ones of a Naive estimator (unadjusted estimation) and the Oracle estimator (adjusts for the true label).
6.1. Validity of ATE Inference from Pre-Trained Representations
Text Data We utilize the IMDb Movie Reviews dataset from Lhoest et al. (2021) consisting of 50,000 movie views labeled for sentiment analysis. The latent features Z as representations of the movie reviews are computed using the last hidden layer of the pre-trained Transformer- based model BERT (Devlin et al., 2019). More specifically, each review results in a -i8-dimensional latent variable Z by extracting the [CLS] token that summarizes the entire sequence. For this, each review is tokenized using BERT's subword tokenizer pert-base-uncasedd truncated to maxi- mum length of 128 tokens, and padded where necessary.
Label Confounding Results The results for the IMDb experiment over 5 simulations are depicted in Figure 4. As expected, the naive estimator shows a strong negative bias. The same can be observed for the S-Learner (for all nuisance estimators) and for DML using lasso or ran- dom forest. In contrast, DML, using linear nuisance es-
Image Data We further use the dataset from Kermany et al. (2018) that contains 5,863 chest X-ray images of chil- dren. Each image is labeled according to whether the lung disease pneumonia present or not. The latent features are
7

---

Adjustment for Confounding using Fre-Trained Representation
Intrrric Dimension Estinates
15.0
12.5
Estmaled
5.0
Gog00
26
9000
Naive
S S-Leamer (NN)
Leamer (RF)
DML (NN)
DML (RF)
CheXpert All
NIH
RSNA
Pre-Trained Representations from Model
Figure Complex Confounding: Comparison of ATE extimators on the X-ray dataset. DML and S-rrer pre-trained repre- sentations. Point estimates and 95% Cls are depicted.
Figure 6. Different Intrinsic Dimension (ID) estimates of pre- trained representations obtained from different pre-trained models. Representations are bosed the X-Ray dataset.
timators (without sparsity-inducing penalty) yields unbi- ased estimates with good coverage, can be seen by the confidence intervals (CIs). First, these results indicate that DML seems to benefit from the doubly robust estimation. Second, DML fails when using ILT non-invariant nuisance estimators such as lasso or random forest. This is because neither of the two can achieve sufficiently fast convergence rates without structural assumptions, such as sparsity or additivity. The latter being unlikely to hold given that rep- resentations were shown to be identifiable only up to ILTs. The results for image-based experiment are given in Ap- pendix D.I, where the same phenomenon can be observed.
estimate is strongly biased similar to the random forest- based estimators. contrast, the neural network-based estimators exhibit much less bias. While the S-Learner's confidence intervals shows high coverage and is therefore the only estimator that enables valid inference. The results for the IMDb dataset with complex confounding are given in Appendix D.1.
optimistic, the DML estimator
Low Intrinsic Dimension We also investigate the low intrinsic dimension hypothesis about pre-trained represen- tations. Using different intrinsic dimension (ID) estima- tors such as the Maximum Likelihood (MLE) (Levina Bickel, 2004), the Expected Simplex Skewness (ESS), and the local Principal Component Analysis (IPCA) we esti- mate the ID of different pre-trained representations of the X-ray dataset obtained from different pre-trained models from the TorchXRay Vision library (Cohen et al., 2022). The results in Figure 6 indicate that the intrinsic dimension of the pre-trained representations is much smaller than the di- mension of the ambient space (1024). A finding that is in line with previous research, which further discussed in Appendix B. Additional information the experiment and the estimators used can be found in Appendix C.4.
6.2. Neural Networks Adapt to Functions on Low Dimensional Manifolds
In second line of experiments, we investigate the ability of neural network-based nuisance estimation to adapt to low intrinsic dimensions. The features in our data sets already concentrate on a low-dimensional manifold. For example, Figure 6 shows that the intrinsic dimension of the X-ray im- ages is around dyy= 12, whereas the ambient dimension is d = 1024. To simulate complex confounding with structural smoothness and sparsity, we first train an autoencoder (AE) with 5-dimensional latent space on the pre-trained represen- tations. These oiw-dimesionnl encodings from the AE are then used simulate confounding. Due this construc- tion of confounding, the true nuisance functions correspond to encoder-ther then-lineat functions, which multi-layered hierarchical compositions and therefore align with Assump- tion 5.2. We refer to this as complex confounding.
6.3. The Power of Pre-Training for Estimation
In another line of experiments, explore the benefits of pre-training in our setup. In particular, we are investigating whether pre-trained neural feature extractors actually out- perform non-pre- non-pre-trained feature extractors in the nuisance estimation of DML-based ATE estimation. We conduct the experiments in the context of the previously introduced image-based Label Confounding setup. To adjust for con- founding in this setup, nuisance estimators must extract the relevant information from the X-rays. For this pur- pose, we compare DML using pre-trained feature extractors against DML using neural feature extractors that are trained on downstream data from scratch. While the former uses the same pre-trained Densenet-122 model that was used in previous image-based experiments, the latter incorpo-
Complex Confounding Results We again compare DML and the S-Learner with different nuisance estimators. In con- trast to the previous section, we now use a neural network (with ReLU activation, 100 hidden layers with 50 neurons each) instead of a linear model in the outcome regression nuisance estimation. The results are depicted Figure 5. Similar to the previous experiments, we find that the naive
8

---

Adjustment for Confounding using Pre-Trained Representation
establish conditions for fast convergence rates in nuisance estimation. This ensures valid ATE estimation and inference, contributing both theoretical insights and practical guidance for integrating machine learning into causal inference.
2.6
Edtimaledd ATE
2.0
AueATe
Limitations and Future Research In this work, we focus on a single source confounding from a non-tabular data modality. A potential future research direction is to study the influence of multiple modalities on ATE estimation. In particular, having multiple modalities requires further causal and structural assumptions the interplay of the modalities. For example, this could mean that each modality is best processed by a separate network or that the confounding information can only be extracted through a joint network that correctly fuses modalities at some point. We note, however, that this is more of technical aspect and a matter of domain knowledge, and thus being of minor relevance for the discussion and theoretical contributions of our study.
1.0
Naiv
Oracle
DML (Pre-trained)
DML (CNN)
Figure 7. Comparison of DML using pre-trained representations
"DML (Pre-trained)" and DML without pre-training "DML (CNN)"
for ATE estimation. Experiment is based on the X-Ray dataset. Point estimates and 95% Cls are depicted.
rates Convolutional Neural Networks (CNNs) as nuisance estimators into the DML ATE estimation routine. The fol- lowing experiment is based on 500 sampled images from the X-Ray dataset, where five-layer CNNs are used in the con-pre-trainge DML version. Further details about the training and architecture of the utilized CNNs can be found in Appendix C.5.
Moreover, we focused on the estimation of the ATE in this paper, given its popularity in both theory and practice. How- ever, our approach could also be extended to cover other target parameters such as the average treatment effect on the treated (ATT) or the conditional ATE (CATE). While each of these would require a dedicated discussion of the necessary assumptions, we believe that many of the core ideas and results presented here-such as the convergence rates for neural network- -baed estimation-coule also be transferred and used a theoretical investigation in those settings.
The results are depicted in Figure 7. poses, we also show the estimates of the Naive and Oracle estimators, which match those of previous experiments. The key finding of Figure 7 is that DML using pre-trained fea- ture extractors (DML (Pre-trained)) yields unbiased ATE estimates and well-calibrated confidence intervals, while DML without pre-training (DML. (CNN)) does not. The same phenomenon can be observed in experiments with varying sample sizes and CNN architecture. These experi- ments are discussed in Appendix D.2. Overall, the results emphasize the benefits of using DML in combination with pre-trained models when utilizing non-tabular data such as images, for confounding adjustment in ATE estimation.
For illustrative pur-
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Further Experiments Further experiments on the asymp- totic normality of DML-based ATE estimation as well as the role of the HCM structure of the nuisance functions are given and discussed in Appendix D.3 and D.4.
References
Aghajanyan, A., Zettlemoyer, I.,, and Gupta, S. Intrin- sic dimensionality explains the effectiveness of language model fine-tuning. arXiv preprint arXiv:2012.13255, 2020.
7. Discussion
In this work, we explore ATE estimation under confounding induced by non-tabular data. We investigate conditions un- der which pre-trained neural representations can effectively be used adjust for such kind of confounding. While the representations typically have lower dimensionality, their in- variance under orthogonal transformations challenges com- mon assumptions to obtain fast nuisance function conver- gence rates, like sparsity and additivity. Instead, the study leverages the concept of low intrinsic dimensionality, com- bining it with invariance properties and structural sparsity to
Ansuini, A., Laio, A., Macke, J. H., and Zoccolan, D. In- trinsic dimension of data representations in deep neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Bach, P., Chermozhukov, V., Kurz, M. S., and Spindler, M. DoubleML-an object-oriented implementation of double machine learning in python. Journal of Machine Learning Research, 23(53):1-6, 2022.
Bartlett, P. L.,, Harvey, N., Liaw, C., and Mehrabian, A.

---

Adjustment for Confounding using Pre- Pre-nTrined Representation
Nearly-tight VC-dimension and pseudodimension bounds for piecewise linear neural networks. Journal of Machine Learning Research, 10(63):1-17, 2019.
Dai, B., Shen, X., and Wang, J. Embedding learning. Jour- nal of the American Statistical Association, 117(537)) 307-319, 2022.
Battaglia, L., Christensen, T., Hansen, S., and Sacher, S. Inference for regression with variables generated from unstructured data. 2024.
Deaton, A. and Cartwright, N. Understanding and misun- derstanding randomized controlled trials. Social science & medicine, 210:2-21, 2018.
Bickel, P. and Li, B. Local polynomial regression on unknown manifolds. Lecture Moooogoho Series,
Deng, J., Dong, W., Socher, R., Li, L.JJ,, Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255, Ieee, 2009.
pp. 177-186, 2007.
Charig, C. R., Webb, D. R., Payne, S. R., and Wickham, E. Comparison of treatment of renal calculi by open surgery, percutaneous nephrolithotomy, and extracorpo- real shockwave lithotripsy. British Medical Journal (Clin- ical research ed.), 29246524):879-882 1986.
Devlin, J., Chang, M.- W., Lee, K., and Toutanova, K. BERT: Pre-training of deep bidirectional transformers for lan- guage understanding. In Burstein,
Doran, C., and Solorio, T. (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Com- putational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. 1171-4186, Min- neapolis, Minnesota, June 2019. Association for Compu- tational Linguistics.
Chen, H., Harinen, T., Lee, J.-Y., Yung, M., and Zhao, Z. CausalML: Python Package for Causal Machine Learning. 2020.
Chen, M., Jiang, H., Liao, W., and Zhao, T. Efficient ap- proximation of deep relu networks for functions on low dimensional manifolds. Advances in Neural Information Processing Systems, 32, 2019.
Dhawan, N., Cotta, L., Ullrich, K., Krishnan, R. G., and Maddison, J. End-To-End Causal Effect Estimation from Unstructured Natural Language Data. In Advances in Neural Information Processing Systems, volume 37, pp. 77165-77199, 2024.
Chernozhukov, V., Chetverikov, D., Demirer, M., Duflo, E., Hansen, C., and Newey, W. Double/äehinseddeymmnn ma- chine learning of treatment effects. American Economic Review, 10265):211 2017.
Erhan, D., Courville, A., Bengio. Y., and Vincent, P. Why does unsupervised pre-training help deep learning? In Proceedings of the thirteenth international conference on artificial intelligence and statisties, pp. 201-208. JMLR Workshop and Conference Proceedings, 2010.
Cheroozhukov, V., Chetverikov, D., Demirer, M., Duñlo, E., Hansen, C., Newey, W., and Robins, J. Double/debiased machine learning for treatment and structural parameters. The Econometrics Journal, 21(1), 2018
Farrell, M. H., Liang, T., and Misra, S. Deep neural net- works for estimation and inference. Econometrica, 89(1): 181-213, 2021.
Chernozhukov, V., Newey, W., Quintas-Martuee V. M., and Syrgkanis, V. Riesznet and forestriesz: Automatic debiased machine learning with neural nets and random forests. In International Conference on Machine Learn- ing. pp. 3901-93333 3901-3914. PMLR, 2022.
Fefferman, C., Mitter, S., and Narayanan, H. Testing the manifold hypothesis. Journal of the American Mathemat- ical Society, 29(4):983-1049 2016.
Christgau, A. M. and Hansen, N. R. Efficient adjustment for complex covariates: Gaining efficiency with DOPE. arXiv preprint arXiv:2402. 12980, 2024.
Fukunaga, K. and Olsen, D. R. An Algorithm for Finding Intrinsic Dimensionality of Data. IEEE Transactions computers, 100(2):176-18) 1971.
Cohen, J. P., Hashir, M., Brooks, R., and Bertrand, H. On
the limits of cross-domain generalization in automated X- ray prediction. In Medical Imaging with Deep Learning, pp. 136-136-5 155. PMLR, 2020.
Gong, S., Boddeti, V. N., and Jain, A. K. On the intrinsic dimensionality of image representations. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3987-3996, 2019.
Cohen, J. P., Viviano, J. D., Bertin, P., Morrison, P., Tora- bian, P., Guarrera, M., Lungren, M. P., Chaudhari, A., Brooks, R., Hashir, M., et al. TorchXRayVision A li- brary of chest X-ray datasets and models. International Conference on Medical Imaging with Deep Learning, pp- 231-249 PMLR, 2022.
Huang, G., Liu, Z., Van Der Maaten, L., and Weinberger, K. Q. Densely connected convolutional networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4700 2017.
10

---

Adjustment for Confounding using Pre- Prainrrined Representation
Jerzak, C. T.,, Johansson, F., and Daoud, A. Estimat- ing causal effects under image confounding bias with application to poverty in Africa. arXiv preprint orXiv:2206.06410. 2022.
Levina, E. and Bickel, P. Maximum likelihood estimation of intrinsic dimension. Advances in Neural Information Processing Systems, 17, 2004.
Lhoest, Q. Villanova del Moral, A., Jernite, Y., Thakur, A., von Platen, P., Patil, S., Chaumond, J., Drame, M., Plu, J., Tunstall, L., Davison, J., Saško, M., Chhablani, G., Malik, B., Brandeis, S., Le Scao, T., Sanh, V., Xu, C., Patry, N., McMilla-MMajor, A., Schmid, P., Gugger, S., Delangue, C.,, Matussière, T., Debut, L., Bekman, S., Cistac, P., Goehringer, T., Mustar, V., Lagunas, F., Rush, A., and Wolf. T. Datasets: A Community Library for Natural Language Processing. In Proceedings of the 1 Conference on Empirical Methods in Natural Language Processing:
Jerzak, C. T., Johansson, F., and Daoud, A. Integrating earth observation data into causal inference: challenges and opportunities, arXiv preprint arXiv:2301.12965 2023a.
Jerzak, T., Johansson, D., and Daoud, A. Image-based Treatment Effect Heterogeneity. In Proceedings of the Second Conference on Causal Learning and Reasoning, pp. 531-553122 PMLR, 2023b.
Johnsson, K., Soneson, C., and Fontes, M. Low Bias Local Intrinsic Dimension Estimation from Expected Simplex Skewness. IEEE transactions on pattern analysis and machine intelligence, 37(1):196 202, 2015.
System Demonstrations, pp. 175-184, 2021.
Mastouri, A., Zhu, Y., Gultchin, L.., Korba, A., Silva, R., Kusner, M., Gretton, A., and Muandet, K. Proximal Causal Learning with Kernels: Avo-Stage Estimation and Moment Restriction. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of P Proceedings of Machine Learning Research, pp. 7512-7523, PMLR, 2021.
Julious, S. A. and Mullee, M. A. Confounding and Simp- son's paradox. BMJ, (0M,)330966667)144881444883 1994.
Kallus, N., Mao, X., and Udell, M. Causal inference with noisy and missing covariates via matrix factorization. Advances in Neural Information Processing Systems, 31, 2018.
Miao, W., Geng, Z., and Tchetgen Tchetgen, E. J. Identify- ing causal effects with proxy variables of an unmeasured confounder. Biometrika, 105(():987-99)) 993779999,971 2018.
Kermany, D., Zhang, K., and Goldbaum, M. Labeled Op- tical Coherence Tomography (OCT) and Chest X-Ray Images for Classification. Mendeley Data, 2018.
Nakada, R. and Imaizumi, M. Adaptive approximation and generalization of deep neural network with intrinsic dimensionality. Journal of Machine Learning Research, 21 21(174):1-38, 2020.
Klaassen, S., Kiichert-Klugee J., Bach, P., Chernozhukov, V., Spindler, M., and Vijaykumar, S. DoubleMI.Deep. Esti- mation of Causal Effects with Multimodal Data. arXiv preprint arXiv:2402.07780 2024.
Ng, A. Y. Feature selection, T, vs. T.2 regularization, and rotational invariance. In Proceedings of the twenty-first International Conference on Machine Learning, pp. 78, 2004.
Kohler, M. and Langer, S. On the rate of convergence of fully connected deep neural network regression estimates. The Annals of Statistics, (((4):2211-2249) 2021.
Pope, P., Zhu, C., Abdelkader, A., Goldblum, M., and Gold- stein, T. The intrinsic dimension of images and its impact on learning. In International Conference on Learning Representations, 2021.
Kohler, M., Langer, S., and Reif, U. Estimation of gression function on manifold by fully connected deep neural networks. Journal of Statistical Planning and Inference, 222:160-188 2023. ISSN 0378-3758. doi: http://bb.ooo///00000..jjp......
Raskutti, G., Yu, B., and Wainwright, M. Lower bounds on minimax rates for nonparametric regression with addi- tive sparsity and smoothness. Advances in Neural Infor- mation Processing Systems, 22, 2009.
Konz, N. and Mazurowski, M. A. The effect of intrinsic dataset properties on generalization: Unraveling learn- ing differences between natural and medical images. In International Conference Learning Representations, 2024.
Robins, J. M. and Rotnitzky, A. Semiparametric efficiency in multivariate regression models with missing data. Jour- nal of the American Statistical Association, 122 129, 1995.
Kuroki, M. and Pearl, J. Measurement bias and effect restoration in causal inference. Biometrika, 101(2):423- 437, 2014.
Robinson, P. M. Root-N-consistent semiparametric regres- sion. Econometrica: Journal of the Econometric Society, pp. 931 1988.
Lee, M. and Lee, M. Smooth manifolds. Springer, 2012.

---

Adjustment for Confounding using Pre- Trainra Representation
Scheidt-Hieber, J. Deep ReLu network approxima- tion of functions manifold. arXiv:1908.00000 2019.
arXiv preprint
Schmidt- J. Nonparametrie regression using deep neural networks with ReLU activation function. The Annals of Statistics, 2020.
Shi, C., Blei, D., and Veitch, V. Adapting Neural Networks for the Estimation of Treatment Effects. Advances Neural Information Processing Systems, 32, 2019.
Stone, C. J. Optimal global rates of convergence for non- parametric regression. The Annals of Statistics, pp. 1040- 1053, 1982.
Stone, C. J. Additive regression and other nomparametric models. The Annals of Statistics, 13(2):689-705, 1985.
der Laan, M. J. and Rose, S. Targeted Learning: Causal Inference for Observational and Experimental Data, vol- ume 4. Springer, 2011.
van der Laan, M. J. and Rubin, D. Targeted Maximum Likelihood Learning. The International Journal of Bio- statistics, 2(1), 2006.
van der Vaart, A. W. Asymptotic Statistics. Cambridge Series Statistical and Probabilistic Mathematics. Cam- bridge University Press, 1998.
Van der Vaart, A. W. and Wellner, J. A. Weak Convergence and Empirical Processes: With Applications to Statistics. Springer Nature, 2023.
Veitch, V., Wang, Y., and Blei, D. Using embeddings to cor- rect for unobserved confounding in networks. Advances in Neural Information Processing Systems, 32, 2019.
Veitch, V., Sridhar, D., and Blei, D. Adapting text embed- dings for causal inference. In Conference on Uncertainty in Artificial Intelligence, pp. 919-928. PMLR, 2020.
Young, E. H. and Shah, R. D. ROSE Random Forests for Robust Semiparametric Efficient Estimation. arXiv preprint arXiv:2410.03471,, 2024.
Zhang, Xue, W., Yu, Y., and Tan, Y. Debiasing Machine- Learning-or Al-Generated Regressors in Partial Linear Models. Available SSRN, 2023.
Zhang, Y. and Bradic, Causal inference through multi- stage learning and doubly robust deep neural networks. arXiv preprint ar arXiv:2407.088 2024.
12

---

Adjustment for Confounding using Pre- Pre-Trained Representation
A. Proofs and Additional Results
A.1. Equivalence Class of Representations
Lemma A.1 (Equivalence Class of Representations), Let F.,,, be a probability space, and let : Q - R2 be a measurable map (a random representation). Then for each ILT Q the random variable ((Z) satisfies
a (Q(Z)) oZ),
where (Z) denotes the o-algebra generated by the random variable Z. Consequently.
Z = {Q(Z) Q 2}
forms an equivalence class of representations that are indistingrishable from the viewpoint of measurable information.
Proof. Each Q E Q an invertible linear transformation. Consequently, Q is Borel measurable bijection with Borel measurable inverse. To show (()) o (Z), consider any Borel set B S Rd.. We have
{w Ω Q(Z(w)) B} (w Ω (w) (B)) Since QB is Borel (as Q is a Borel isomorphism), the pre-image fw Z(w) E (B) belongs to a (Z). Similarly, for any Borel set A C Rd,
(w E Ω: 7(w) E A) = [ Q(Z(W) E {(((()))((()) {{)), which belongs to (Q(2)) Therefore, (Q(2)) = =Z)).
A.2. Proof of Lemma 4.1
Proof. We consider f being C* on the open domain D C IRd so by definition, all partial derivatives of f up to order exist and are continuous on D. Further, we consider any invetible matrix Q. Such linear transformations are known to be infinitely smooth (as all their partial derivatives of any order exist and are constant, hence continuous). Hence, the function = f Q-1 the composition of a C function f with linear (and thus (o map Q-1
Applying the multivariate chain rule, of f and the linear transformation Q-1 Specifically, since Q-1 is O and f is C* their composition h retains the C smoothness. Lastly, the (transformed) domain Q(D) is also open as linear (and thus continuous) transformations preserve the openness of sets in Rd. Therefore, h is well-defined and C on Q(D).
easily verify that the differentiability properties of h inherited from those
A.3. Proof of Lemma 4.2
Proof. Suppose that Q is an invertible matrix representing the linear map + Q(=). Denote by Q = =Q-1 its inverse and its rows by .......04 Id.
(i) Additivity
Assume that f A - IR is additive, where X C Rd, such that
f(x) = (ZT))
1=1
and suppose that at least one fj is nonlinear. Now consider the transformed input space X = Q(x) = (Qx x E X), induced by the invertible linear transformation Q. Let h X + I be given by h(ā) == f(Q-tt). Then represents the
same mapping as ff but expressed in the transformed coordinate system Y. In particular, h(ā) = f(z), Vā E X. Further,
we have
h(t) = ES 3=1
Assume without loss of generality that fi is nonlinear. The set of invertible matrices where タ equals a multiple of a standard basis vector has Haar measure 0. Hence, fi(ğī ă) is almost everywhere a nonlinear function all coordinates of ä, implying that is not additive.
13

---

Adjustment for Confounding using Pre-Trained Representation
(ii) Sparsity
Assume f x R. where X' S R24 is sparse linear of the form f(x) = B with 1 S ||81lo < d. We again consider the transformed input space X = Q()) {Qu æ E 2), induced by the invertible linear transformation Q, and define h : X + In given by h(ā)= f(Q 2). Then we have h(ā)) f(Q- 3 Q 2 =: B z. While the map h is still linear, the set of matrices Q such that || = || Q F d has Haar measure zero. Hence, h is almost everywhere not sparse.
A.4. Proof of Lemma 4.3
Proof. As in the previous proof in Appendix A.2, it is essential to note that ILTs Q are linear, invertible maps that are C700 (infinitely differentiable) with inverses that are likewise C Specifically, Q serves as a global diffcomorphism on IRd, ensuring that both Q and Q are smooth (Cxc functions.
Given that M is a d dimensional smooth manifold, for each point x on the manifold (x E M), there exists neighborhood U E M and smooth chart yp U + IR that a diffeomorphism onto its image. Applying the orthogonal transformation Q to M results in the set Q(M), and correspondingly, the image Q(U) - Q(M). To construct smooth chart for (M),
we can consider the map
Q(U) - RdM, (()()) = =p(x).
where E U. Since Q is a diffcomorphism, the composition ; yp Q-1 restricted to Q(U) remains a smooth diffeomorphism onto its image. Hence, this defines a valid smooth chart for (MM) Covering ((M) with such transformed charts derived from those of M ensures that Q(M) inherits a smooth manifold structure. Each chart e smoothly maps an open subset of (M) to an open subset of Rd preserving the intrinsic dimension. Therefore, the intrinsic dimension d. of the manifold M is preserved under any orthogonal transformation Q, and Q(M) remains dyy-dimensional smooth manifold in Rd
A.5. Proof of Lemma 5.3
Proof. Recall that Q is an invertible linear map, fo = fo bb M + R. and fo fo v Q- : - R. Write f= é with so Q ((M) + R. Since M is smooth manifold, Q(M) is smooth manifold with the same intrinsic dimension d by Lemma 4.3. Since : - Q-1 is continuous and M is compact, Q(M) is also compact. Next, since do is sp-Smooth by assumption, vi is also Sm-Shooth by Lemma 4.1. Finally, the HCM part f in the two models fo and fo is the same, they share the same constraint set P. This concludes the proof.
A.6. Proof of Theorem 5.5
We will use Theorem 3.4.1 of Van der Vaart & Wellner (2023) to show that the neural network ff converges at the rate stated in the theorem. For ease of reference, we restate slightly simplified version of the theorem adapted to the notation used in our paper. Here and in the following, we write S b to indicate S Cb for constant C' E (0, oo) not depending on n. Proposition A.2. Let Fa be a sequence of function classes, E be some loss function, So the estimation target, and
f = arg min JEF, ni i=1
E ((F(Z)).Y..
Define Fn,s = ( E Fn: - Volloo S 0) and suppose that for every > it holds inf Ele(f(Z), E((C))))))) 2))
2 8²,
(A.2.1)
/ and, writing 8z(2.1)= = =(((())))) E(folz), v), that
1 Σ (,( -(1,)),) <
2n()) Vn
E
sup JEF..,
(A.2.2)
for functions on (5) such that + (,)68/ 22 is decreasing for some = > 0. there are Fo Fn and en 2 such that =(2 -(((6(())) (n(En) S Vns
(A.2.3)
(A.2.4)
14

---

Adjustment for Confounding using Pre- Pre-Trained Representation
it holds IIʃ Joll (p(sn)-
Proof of Theorem 5.5. Define (s*, ,d") arg (a) s/p and denote the targeted rate of convergence by
= ====)==nn n कد (logn)*
2 (log n)
En (U.p)EPU(
max
We now check the conditions of Proposition A.2.
Condition (A.2.1):
Follows from Assumption 5.4, since
inf JeF V
alss - Sollin(P2) 2 4 B(e(((2)))) - -[(e(6(2)) fer JeFna
inf
00
Condition (A.2.2): Let [,, L2())) be the minimal number of g-balls required to cover F in the L2 Theorem 2.14.2 of Van der Vaart Wellner (2023) states that eq. (A.2.2) holds with
(1+ (8)
2n( (n(
8² 8'
where
(x) sup Q
V + log ((,, (2()))
with the supremum taken over all probability measures Q. Lemma A.3 in Appendix A.7 gives
J,) S Lo
which implies that 8 1 (0()))22-1/2 is decreasing, the condition is satisfied.
Condition (A.2.3): According to Lemma A.4 in Appendix A.7 there are sequences L,, = Ollog=,7), v, = O(en 88' 1226") such that there is neural network fo F(L,; vn) with
sup #EM
fo(2) -o(=)fo)]]=---)000 O(e.).
Together with Assumption 5.4, this implies
-
folli,(Ps) S b sup (fo(z) fo(=)) S
ZEM
as required.
Condition (A.2.4): Using L, Ollog= 2) Vn =O(sn" -d' 125"" and our bound on (8) from Lemma A.3, we get
(8) S 8 log 222 3/2
Now observe that
@n(En)
-
En x -2
log - (E
S En
log 1)
22
vn
2--
2d
En 2.-
log² + En
(og 1112
$ n 1/2ggn)))) + 1/2
where the last step follows from our definition en and the fact that log(e, S log n. In particular, En satisfies on (ex) S Vnef, which concludes the proof of the theorem.
15

---

Adjustment for Confounding using Pre-Trained Representation
A.7. Auxiliary results
Lemma A.3. Let F(L, be set of neural networks with SuP)EF(C.)) llf < co. For all ô > 0 sufficiently small, holds
sup [ Q
6 V V1+ log ((, S /olog(1/))T///o((Ta,
Proof. Denote by VC(F) the Vapnik-Chervonenkis dimension of the set F. By Theorem 2.6.7 in Van der Vaart Wellner (2023), it holds
sup (,,FF, & log(1/e)VC(F),
for > sufficiently small. By Theorem 7 of Bartlett et al. (2019), we have
VC(F(L, & L2v² log Lvv).
For small e, this gives
sup /
S (lo((//))Lv/log(L))
Integrating the right-hand side gives the desired result.
Lemma A.4. Suppose fo satisfies Assumption 5.2 for given constraint set and (s dM). Define (s*,d*)
arg (a.)) s/p. Then for any 0 sufficiently small, there is neural network architecture F(L, with = Olog -1),, ( /28") such that there is fo FFL,v) with
(fo(2)) fo(2) = 0(s). =EM
sup
Proof. The proof proceeds three steps. We first approximate the embedding component vi by neural network v. then the HCM component of by neural network f. Finally, we concatenate the networks to approximate the composition fo f o + by fo f.
Approximation of the embedding component. Recall that w: .M Rd is s--smooth mapping. Write +(2) (1h1(e (z),..., (=)) and note that each wh M - T.. is also sy-smooth. Since M is smooth dyc-dimensional manifold, it has Minkowski dimension dyy Then Theorem of Kohler et al. (2023) (setting M E -1/26 in their notation) implies that there is a neural network b/y E F(L.g.u) ux) with Tap = Ollog s and up O(e-dM/22) such that
sup (2) w)y(2) = O(s).
sEM
3a): M + Rd. By construction, the parallelized network Parallelize the networks by into single network := (.)
has Lg layers, width d Uy Ovv)), and satisfies
sup (1x(2) (22))l 0()). #EM
Approximation of the HCM component. Let a E (0, oo) arbitrary. By Theorem 3(a) of Kohler Langer (2021) (setting Mig. -1/2p)" in their notation), there a neural network f e F(L,, with L, = O(loge and Vy OE /2s²) such that
sup 15(x) [(x) O(e).
16

---

Adjustment for Confounding using Pre-Trainee Representation
Combined approximation. Now concatenate the networks & and to obtain the network So Jo E F(Ly + T. max ). Observe that L + f = Ologg= and Uy f = O(e=a" /2="), the network has the right size. remains to show that its approximation error is sufficiently small. Define
= sup [1v(()) (=))l,,
MEM
which is (s)) by the construction of 1
sup [|)(2)) +% =EM
which is ((1) by assumption, and
[f(z) f(x')l
K sup 2,2²
|| - I'1l
which is finite since f is Lipschitz due to min (s,) Lipschitz. By the triangle inequality, have
> 1 and the fact that finite compositions of Lipschitz functions
sup [o(2) fo(2)1 - sup =E.M
[f(((2) - f(((2)))+ + sup [f(((()) f(+(2))]
=EM
REM
< sup
1f(x) - [xx]+
=E(-a,ji"
O()),
claimed.
A.8. Proof of Theorem 5.7
Proof. We validate the conditions of Theorem II.1 1 of Chernozhukov et al. (2017). Our Assumption 5.6 covers all their moment and boundedness conditions on g and m. By Theorem 5.5, we further know that
|| (k) ml2 + (l)(k) gl l2(2)) = (1().
Further, Theorem 5.5 yields
|| [ (k) m (1)) x [lg(k) gl =
=0p
n 20E x (())))))((e
max
max
(s (s')))ePullddddd
204p zotle
= =0p
max ((s)e (s".p") (n.
max
We have to show that the term on the right is of order op(n-1/2), 1/2). Observe that
2 p/s + - /s
2s +
2s' + pl 2
2
4+p/s p /
1
(2+p/s)22p//ss1)
2
pp'
4+p/s > >2 p/
+
2ss
4> ppl
ss!
Thus, our condition
1
min
min
(sp))EP p
(s (M pl 4'
implies
(k) - -mllca(P2) x (k) (044)-9l(s4te - = opp(1-11
as required.
17

---

Adjustment for Confounding using Pre-Trained Representation
B. Additional Related Literature & Visualizations
B.1. Empirical Evidence of Low Intrinsic Dimensions
Using different intrinsic dimension (ID) estimators such as the maximum likelihood estimator (MLE; Levina & Bickel, 2004) on popular image datasets such as ImageNet (Deng et al., 2009), several works find clear empirical evidence for low ID of both the image data and related latent features obtained from pre-trained NNs (Gong al., 2019; Ansuini et al., 2019; Pope et al., 2021). The existence of the phenomenon of low intrinsic dimensions was also verified in the medical imaging (Konz& & Mazurowski, 2024) and text-domain (Aghajanyan et al., 2020). All of the mentioned research finds striking inverse relation between intrinsic dimensions and ostate-of-the-ar)) model performance, which nicely matches the previously introduced theory about ID>related convergence rates.
B.2. Hierarchical Composition Model (HCM) Visualization
This section provides an illustration of the Hierarchical Composition Model (HCM) that was formally introduced in Definition 5.1. As the name suggests, every HCM is a composition of HCMs of lower level. In Figure 8 we give an illustration
of a particular HCM of level 2 and constraint set P21. which we abbreviated by HCM (2, P21). Following the notation of Definition 5.1, the (2, P21) corresponds the function Id R defined by (xx 1 (n.")(()) (x), (z)()) (x)).
Each (x) for j E .p]} corresponds to HCM of level 1, which itself are compositions of HCMs of level O. Each of the latter corresponds to a feature in the data. The constraint set of each HCM corresponds to the collection of pairs of the degree of smoothness and number of inputs of each HCM it is composed of. For example, assuming that n is a 5-smooth function, then the constraint set of the (2M(2, P2i) function f is P21 -パー both regression and classification. In the latter, the conditional probability would need to satisfy the HCM condition, and any non-linear link function for classification would just correspond to a simple function the final layer of the HCM.
1 Pij U p). The HCM framework fits
1 P
A
a
3d
11
a
J
Zd
21
HCM(O, Pat)
(O, Pos)
HCM(L,Pt) PH)
HCM(1,
HCM(1,
(2M(2,
Figure 8. Visualization of a HCM: The illustration depicts a HCM of level and constraint set Pz. The HCM of level 2 is a composition of HCMs of level 1, i.e. hr, hey..., hy. The latter are itself compositions of HCMs of level O, each corresponding a feature the data.
C. Experimental Details and Computing Environment
We conduct several simulation studies to investigate the performance of different Average Treatment Effect (ATE) estimators of a binary treatment on some outcome in the presence of a confounding induced by non-tabular data. In the experiments, the confounding is induced by the labels, i.e., the pneumonia status or the review as well as more complex functions of the pre-trained features. Nuisance function estimation is based on the pre-trained representations that are obtained from passing the non-tabular data through the pre-trained neural models and extracting the last hidden layer features.
C.1. Data and Pre-trained Models
IMDb For the text data, we utilize the IMDb Movie Reviews dataset from Lhoest et al. (2021) consisting of 50,000 movie reviews labeled for sentiment analysis. For each review, we extract the [CLS] token, a fi8-dimensiooal vector per
18

---

Adjustment for Confounding using Pre-Trainct Representation
review entry, of the pre-trained Transformer-basee model BERT (Devlin et al., 2019). To process the text, we BERT's s subword tokenizer lat--bbbseuuusesee and truncate sequences maximum length of 128 tokens. We use padding if necessary. After preprocessing and extraction of pre-trained representations, we sub-sampled 1,000 and 4,000 pre-trained representations for the two confounding setups to make the simulation study tractable.
X-Ray For the image data simulation, we use the dataset from Kermany et al. (2018) that originally contains 5,863 chest X-ray images of children that were obtained from routine clinical care in the Guangzhou Women and Children's Medical Center, Guangzhou. We preprocess the data such that each patient appears only once the dataset. This reduces the effective sample size to 3,769 chest X-rays. Each image is labeled according to whether the lung disease pneumonia is present or not. The latent features are obtained by passing the images through a pre-trained convolutional neural network and extracting the 1024-dimensional last hidden layer features of the model. For this purpose, we use pre-trained Densenet-12] 121 model from the TorchXRay Vision library (Cohen et al., 2022), Specifically, we use the model called is a Densenet-121 model with resolution 224 x 224 that was pre-trained on all chest X-ray datasets considered in Cohen et al. (2020). We chose this model for the extraction of pre-trained representation in our experiments, based on its superior performance in benchmark studies conducted in prior work (Cohen et al., 2020). Note that the dataset from the Guangzhou Women and Children's 's Medical Center that use, was not used during the training of the model. This is important from theoretical and practical viewpoint, as the confounding simulation via labels might otherwise be too easy to adjust for given that the model could have memorized the input data. However, using this kind of data we rule out this possibility.
-rees222alll which
C.2. Confounding
As introduced in the main text, we simulate confounding both on the true labels of the non-tabular data as well encodings from trained autoencoder. While this induces a different degree of complexity for the confounding. the simulated confounding is somewhat similar in both settings. We first discuss the simpler setting of Label Confounding. In all of the experiments, the true average treatment effect was chosen to be two.
Label Confounding Label Confounding was induced by simulating treatment and outcome both dependent on the binary label. In the case of the label being one (so in case of pneumonia or in case of a positive review), the probability of treatment is 0.7 compared to 0.3 when the label is zero. The chosen probabilities guaranteed a sufficient amount of overlap between the two groups. The outcome Y is simulated based on a linear model including a binary treatment indicator multiplied by the true treatment effect (chosen to be 2), as well as a linear term for the label. Gaussian noise is added to obtain the final simulated outcome. The linear term for the label has a negative coefficient in order to induce a negative bias to the average treatment setup compared to a randomized setting. Given that the confounding simulation is only based on the labels, the study was in fact randomized with respect to any other source of confounding.
Complex Confounding To simulate Complex Confounding with structural smoothness and sparsity, we first train an autoencoder (AE) with 5-dimensional latent space on the pre-trained representations, both in the case of the text and image representations. These AE-encodings are then used to simulate confounding similarly as in the previous experiment. The only difference is that we now sample the coefficients for the 5-dimensional AE-encodings. For the propensity score, these are sampled from a normal distribution, while the sampled coefficients for outcome regression are restricted to be negative, to ensure a sufficiently larger confounding effect, that biases naive estimation. We choose a 5-dimensional latent space to allow for sufficiently good recovery of the original pre-trained representations.
C.3. ATE Estimators
We estimate the ATE using multiple methods across 5 simulation iterations. In each of these, we estimate a Naive estimator that simply regresses the outcome treatment while not adjusting for confounding. The Oracle estimator uses linear regression of outcome on both treatment and the true label that was used to induce confounding. The S-Learner estimates the outcome regression function s(4, 2) E|y T = 1, Z =] by fitting a single model ((,, to all data, treating the treatment indicator as a feature. The average treatment effect estimate of the S-Learnrr is then given by
ATES = Σ
0(1, 20(1,)))-00 21.
i=1
19

---

Adjustment for Confounding using Pre-Trainee Representation
In contrast, the Double Machine Learning (DML) estimators estimates both the outcome regression function and the propensity score obtain its double robustness property. In our experiments, DML estimators use the partialling-out approach for ATE estimation, which is further discussed in the next paragraph.
In the Label Confounding experiments, both the S-Learner and DML estimators are used in combination with linear and random forest-based nuisance estimators. DML (Linear) uses standard linear regression for the estimation of the outcome regression function, and logistic regression with L2-penalty for the estimation of the propensity score. Both nuisance function estimators are ILT-invariant. DML (Lasso) uses L1 L1-penalized linear and logistic regression with cross-validated penalty parameter selection for the outcome regression and propensity score estimation, respectively. The S-Learner (Linear) and S-Learner (Lasso) use unpenalized and L1-pennlized linear regression for the outcome regression, respectively. The random forest-based nuisance estimation (both for DML and S-Learner) is based on the standard random forest implementation from scikit-leam. The number of estimated trees is varied in certain experiments to improve numerical stability.
In the Complex Confounding experiments, we also use neural network- seteork-based nuisance estimators for DML and the S-Learner. For this purpose, we employed neural networks with a depth of 100 and a width of 50 while using ReLUU activation and Adam for optimization. While DML (NN) and S-Learner (NN) use neural networks for the outcome regression, logistic regression is employed in DML (NN) for propensity score estimation to enhance numerical stability. Generally, DML is used with sample splitting and with two folds for vross-valiiation. For the S-Leamer and DML the Python packages CausalMI. (Chen et al., 2020) and DoubleMI. (Bach et al., 2022) are used, respectively.
Partially Linear Model and Orthogonal Scores In our experiments, we simulated two different types of confounding. In both cases we use non-tabular data to adjust for this confounding, given that the confounding inducing information is contained in this data source, but not available otherwise. However, as this information is non-linearly embedded in the non-tabular data, the model that we aim to estimate follows the structure of a so-called partially linear model (PLM). Given a binary treatment variable T, the PLM is a special case of the more general confounding setup that we consider in the theoretical discussion of this paper. Specifically, the PLM considers the case where the outcome regression function in (2) decomposes
W) [T, W] = 0.T + +gW).
(6)
The structure of the propensity score in (3) remains the same. The parameter 6g in the PLM corresponds to the target parameter considered in (1), namely the ATE. In their theoretical investigation, Chernozhukov et al. (2018) discuss ATE estimation both in the partially linear model and in the more general setup, which they refer to interactive model. Given that we consider the more general case in Section 5, the orthogonalized score stated in this section matches that of Chernozhukov et al. (2018) for the ATE in the interactive model. In case of the PL.,, Chernozhukov et al. (2018) consider two other orthogonalized scores, one of which is the so-called partialling- -ot score function, which dates back to Robinson (1988). The partialling-oot score corresponds to an unscaled version of the ATE score in the (binary) interactive model in case the outcome regression decomposes as in (6). The scaling is based on certain estimated weights. Therefore. score functions as the partialling-out score are sometimes referred to as unweighted scores (Young & Shah, 2024). While the theoretical result in Theorem 5.7 could also be obtained for DML with partialling-ou score under similar assumptions, the key requirement being again (5), the approach may not be asymptotically efficient given that it does not use the efficient influence function. However, the potential loss of asymptotic efficiency is often outweighed by increased robustness in finite-sample estimation when using unweighted scores, which has contributed to the popularity of approaches such as the partialling-on out method in practice (van der Vaart (1998, $25.9), Chernozhukov et al. (2018, $2.2.4), Young & Shah (2024)). Accordingly, we also adapted the partialling-out approach in the DML-based ATE estimation in our experiments.
C.4. Intrinsic Dimensions of Pre-trained Representations
In Section 6.2 we also provide empirical evidence that validates the hypothesis of low intrinsic dimensions of pre-trained representations. For this, we use different pre-trained models from the from the TorchXRay Vision library (Cohen et al., 2022). All of these are trained on chest X-rays and use a Densenet-121 (Huang et al., 2017) architecture. Given the same architecture of the models, the dimension of the last layer hidden features is 1024 for all models. The different names of the pre-trained models on the x-axis in Figure 6 indicate the dataset they were trained on. We use the 3,769 chest X-rays from the X-rays dataset described above and pass these through each pre-trained model to extract the last layer features of each model, which we call the pre- pre-trained trained representations of the data. Subsequently, we use standard intrinsic dimension estimators such as the Maximum Likelihood Estimator (MLE) (Levina & Bickel, 2004), the Expected Simplex Skewness (ESS) estimator (Johnsson et al., 2015), and the local Principal Component Analysis (IPCA) estimator (Fukunaga & Olsen,
20

---

Adjustment for Confounding using Pre-Trainee Representation
1971), with choice of number of neighbors set to 5, 25 and 50, respectively. While the intrinsic dimension estimates vary by the pre-trained model and the intrinsic dimension estimator used, the results indicated that the intrinsic dimension of the pre-trained representations is much smaller than the dimension of the ambient space (1024).
C.5. Double Machine Learning with Convolutional Neural Networks Nuisance Estimators
In Section 6.3, compare DML with pre-trained neural networks against DML without pre-trate pre-trained neural networks. This experiment investigates the benefits of pre-training for nuisance estimation in the context of DML-based ATE estimation. Experiments are conducted on the X-Ray dataset, and confounding is simulated based on Label Confounding. DML (Pre- trained) uses the same pre-trained Densenet-12 121 from the TorchXRay Vision library (Cohen et al., 2022) that was previously used as pre-trained neural feature extractors the other image-based experiments. Building on this pre-trained feature extractor, DML (Pre-trained) then uses linear models on the pre-trained features for the nuisance function estimation. In contrast, DML without a pre-trained feature extractor uses standard Convolutional Neural Networks (CNNs) to estimate the nuisance functions directly on top of the images. The experiment of Figure uses five-layer CNN with 3x3 convolutions, batch normalization, ReLU activation, and max pooling, followed by model head consisting of fully connected layers with dropout. Training uses Adam optimization with early stopping. When the network is utilized for propensity score estimation, the outputs are converted to probabilities via a sigmoid activation. Both DML with and without pre-trained feature extractors use the "partialling-or approach in combination with sample-splittng for doubly-robust ATE estimation.
C.6. Computational Environment
All computations were performed on a user PC with Intel(R) Core(TM) i7-8665U CPU @ 1.90GHz, 8 cores, and 16 GB RAM. Run times of each experiment do not exceed one hour. The code to reproduce the results of the experiments can be found at https:: //githubttttbb com/ ri rerrrickkkms ckmer schtlte/Pretrainn Prineedd Causal maasal-Adjuuss Adjuss.
D. Further Experiments
This section provides additional results from experiments that extend those discussed in the main body of the paper.
D.1. Comparison of ATE Estimators
The results depicted in Figure and Figure 10 complement Figure and Figure that discussed in Section 6.
Label Confounding (X-Ray) The results for the Label Confounding simulation based on the X-Ray dataset over 5 simulations are depicted in Figure 9. As before, the naive estimator shows strong negative bias. Similarly, the S-Learner (for all three types of nuisance estimators) and for DML using random forest or lasso exhibit negative bias and too narrow confidence intervals. In contrast, DML using linear nuisance estimator (without sparsity- -nnuuigg penalty) yields less biased estimates with good coverage due to its properly adapted confidence intervals.
Eslimated ATE
2.0
1.8
1.4
1.2
Naive
Oracle
S-Leamer S-Leamer S-Leamer DML (Linear)
DML (Lasso)
(RF)
(Lasso)
(Linear)
(RF)
Figure 9. Label Confounding (X-Ray): Comparison of ATE estimators the X- Ray dataset. DML & S-Learner use pre-trained reprexentations and three types of nuisance estimators: lineur models without L Li-penalization (Linear), linear models with Li- penalization (Lasso), as well random forest (RF), Point estimates and 95% Cly depicted.
21

---

Adjustment for Confounding using Pre-Traines TrriTrained Representation
Complex Confounding (IMDb) A similar pattern can be observed for the Complex Confounding setting on the IMDb data depicted in Figure 10. The naive estimator and both of the random forest-based ATE estimators exhibit strong bias. In contrast, both neural network-bass based estimators show very little bias. This provides further evidence that neural networks can adapt to the low intrinsic dimension of the data. However, unlike the DML estimator, the S S-Learner still produces overly narrow confidence intervals and thus has poor coverage. As in the example discussed the main body of the text, the DML. (NN) estimator is the only one that yields unbiased estimates and valid inference.
Estimated ATE
True ATE
Naive
S-Leamer (NN)
S-Leamer (RF)
DML (NN)
DML (RF)
Figure Complex Confounding (IMDb); Comparison of ATE estimators on the IMDb dataset. DML S-Learner use pre-trained representations and either neural network (NN) random forest (RF) based musance estimators. Point estimates and 95% Cls are depicted.
D.2. DML with and without Pre-Training
The experiment on DML with and without pre-trained representations explored the benefits of pre-training for DML and was discussed in Section 6.3. We extend this line of experiment by considering different sample sizes, as well as neural network architectures for the non-pre-trained model. While the DML (CNN) estimator in Figure 7 uses five-layer CNNs for the nuisance function estimation, the DML (CNN) estimator in Figure 11 uses slightly simpler model architecture of two-layer CNNs with ReLUU activation and max- max-pooling, followed by fully connected layers as model head. The slightly less complex model architecture requires fewer neural network parameters to be trained, which might be beneficial in the context of the X-Ray dataset, considering the comparably small sample size available for model training.
Overall, Figure 11 confirms the previous finding that DML with pre-trained representations performs much better than DML. without pre-training. While the former yields unbiased ATE estimates, the ATE estimates of the latter show a strong negative bias. As the two plots show, this result is independent of the model architecture and sample size used.
2.75
2:50
Edtimated ATE
2.0
TTeATE
Estimated ATE
2.00
1.6
1.75
Nalve
Oracie
DML (Pre-trained)
DML (CNN)
Nalve
Oracle
DML (Pre-trsined)
DML (CNN)
Figure DML with/out Pre- Pre-Training: Comparison of DML using pre-trained representations "DML (Pre-trained)" and DML without pre-training "DML (CNN)" for ATE estimation using 500 (Left) and all 3769 (Right) images from the X-Ray dataset, Point estimates and 95% Cls are depicted.
22

---

Adjustment for Confounding using Pre- Pre-Trained Representation
D.3. Asymptotic Normality of DML ATE Estimation
This line of experiments explores the asymptotic normality of the DML estimator, For this purpose, we extend the ATE estimation experiments of Figure with Label Confounding and Figure with Complex Confounding that are based on the X- X-Ray dataset. While the two figures depict the estimates of different ATEs over 5 simulation iterations, we repeat both experiments with 200 iterations and collect the ATE estimates. We standardize each estimate and plot the corresponding empirical distribution. The results are shown in Figure 12. The left plot depicts the empirical distribution of the 200 standardized point estimates of the Oracle and Naive estimator, as well as DML with linear nuisance estimator in the Label Confounding experiment. The right plot displays the empirical distribution of the 200 standardized point estimates of the Naive, S-Learner, and DML estimator in the Label Confounding experiment. The latter two use neural network- nuisance estimation. While the distributions of the Naive and S-Learner estimators show a strong bias, the distribution of the DML approach matches the theoretical standard normal distribution in both experiments.
A10,
AÃO,
Natvo
1.2
Naiva
Orade
S-Leamer
DML
DML
0.8
0.8
0.1
0.2
0.0
0.0
-15.0
-12.5
-10.0
-7.5
-5.0
-2.5
0.0
2.5
-40
-30
-20
-10
10
(ATÉ ATE)J
(ATE -TTEII
Figure 12. Asymptotic Normality of DML.: Compartson of the empirical distributions standardized point estimates of different ATE estimators on the X-ray dataset. Left: Distribution of 200 standardized point estimates of the Naive, Oracle, and DML with linear nuisance estimation from the Label Confounding experiment. Right: Distribution of 200 standardized point estimates of Naive, S-Learne with NN-based musance estimation (S Learner) and DML with NN-based nuisance estimation from the Complex Confounding experiment.
D.4. Effects of the Hierarchical Composition Model (HCM) Structure on Estimation
In this experiment, we investigate the effect of the Hierarchical Composition Model (HCM) structure in the context of ATE estimation. The HCM was formally introduced in Definition 5.1 and later used Theorem 5.5 to derive convergence rates of neural network-bart the worst-case pair that appeared in the constraint set of the HCM. The core benefit of the HCM in our context is, that constitutes very flexible class of functions, while at the same time, it enables obtain fast convergence rates in case the target function factors favorably according to the hierarchical structure of the HCM. The latter could be fulfilled in case each composition in the hierarchical structure only depends on few prior compositions. This structural sparsity would improve the worst-case pair in the constraint set and thereby could allow for obtaining fast convergence rates.
estimation. This result showed that the convergence rate of neural networks determined by
Now, in the DML ATE estimation, one might be interested in what happens when the nuisance functions do not factor according to the HCM structure, such that sufficiently fast convergence rates for the nuisance estimation, required the ATE estimation, can be achieved. We simulate such scenario by extending the previously introduced setup of Complex Confounding. Instead of simulating confounding based on low-dimensional encodings from AE trained on the pre-trained representations from the X-Ray experiments (as done in some of the previous experiments). do this directly based on the pre-trained representations. For this purpose, we define the nuisance functions (outcome regression and propensity score) to depend on the product of all features in each pre-trained representation. Hence, the nuisance functions depend on all 1024 features, thereby mimicking the curse of dimensionality scenario discussed in Section 4. Further, the nuisance functions are constructed such that bias is introduced in the ATE estimation. In the following, we estimate the same set of ATE estimators that were previously used the Complex Confounding experiments.
23

---

Adjustment for Confounding using Pre-Trained Representation
The results depicted in Figure 13. All estimators show substantial bias (even the DML approach), given that of the estimators is able properly adapt this complex type of confounding structure. The results are validation of the fact that no (nuisance) estimator can escape the curse of dimensionality without utilizing certain beneficial structural assumptions.
25
Estimated ATE
15
Q
10
Tae ATE.
00000
Naive
S-Leamer (NN)
S-Leamer (RF)
DML
(NN)
(RF)
Figure 13. Comparison of ATE estimators: DML & S-Learner use pre-trained representations and either neural network (NN) or random forest (RF) based nuisance estimators. Point estimates and 95% are depicted.
24


Reviews:
[{"review_id": "Kf5Y77IUjw", "reviewer": "Reviewer_vZqT", "summary": "This paper investigates the application of Double Machine Learning (DML) for estimating the Average Treatment Effect (ATE) in non-tabular data contexts, such as text and images. The authors highlight the limitations of traditional causal inference methods in handling non-tabular data and propose leveraging pre-trained representations from deep neural networks (DNNs) to adjust for confounding variables.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Additional experiments** can be found at https://anonymous.4open.science/r/icml2025-6599/add_exp_r1.pdf\n\n- - - -\n\n### **Impact of ILTs on ATE**\n\n> How does ILT invariance specifically impact DML estimation?\n\nThe ILT invariance of pre-trained representation does not only have theoretical but also crucial practical consequences in DML estimation. More specifically, Random Forest (RF) and the Lasso are commonly used for nuisance function estimation in DML applications with tabular data. However, as we show both theoretically and empirically, additivity and sparsity cannot reasonably be assumed given the ILT invariance.  Further extending our previous results on the IMDb and X-Ray datasets, we show in Fig.13 that both the ILT non-invariant nuisance estimators RF (building on additivity) and Lasso (building on sparsity) yield biased ATE estimation when using pre-trained representations, while DML with ILT invariant “Linear” (NN with linear and classification model head) estimators yields unbiased results in both empirical studies. \n\nWe thank the reviewer for this question and hope our explanation clarifies it. We have discussed this aspect in the paper on pages 7-8, but are happy to further highlight this point in a revised version of the manuscript.\n\n- - - -\n\n### **With/out Pre-Training**\n\n> Why is there no experiment comparing DML with and without pre-trained representations?\n\nWe thank the reviewer for bringing up this important aspect. To validate the benefits of pre-training in our context, we have now conducted experiments where we fit DML with pre-trained representations and compare it with DML without pre-training. As models have to be trained from scratch, we expect pre-training to have less bias than training from scratch.  We demonstrate this on the pneumonia (X-ray) data set, where we used 500 and all 3769 X-rays for ATE estimation and fitted CNNs as nuisance function estimates (both for the propensity score and outcome regression). The results are depicted in Fig.14 and demonstrate that DML using pre-trained representations yields unbiased estimates with good coverage while DML with from-scratch training of CNNs does not. Note that these results become even more pronounced in case of smaller sample sizes (<200) –- a setup that is frequently encountered in clinical practice. \n\nWe thank the reviewer again for this very helpful comment. We think that the new results further strengthen our paper. \n\n- - - -\n\n### **HCM**\n\n> Are there scenarios where HCM does not hold? If real-world data does not follow an HCM structure, how does this affect ATE estimation?\n\nWe thank the reviewer for raising this question. Indeed, the HCM structure is a structural assumption that does not necessarily need to hold in all real-world data applications. However, the success of DNNs in many real-world applications and the fact that DNNs precisely mimic such HCM structures suggest that it may be a reasonable assumption in many settings. That being said, our asymptotic normality results in Thm. 5.7 do not necessarily depend on the HCM assumption. In fact, we used it to relax the smoothness assumptions that otherwise would be required. Given a sufficient amount of smoothness of the target function and low ID of the representations, Thm. 5.7 can achieve the same root-n consistency and asymptotic normality. However, there are cases where ATE estimation yields biased results if the HCM structure does not hold. To illustrate this, we conducted an additional experiment, where confounding is based on the product of pre-trained representations (hence HCM no longer holds, which is crucially different from our previous complex conf. experiments). The results in Fig. 15 demonstrate that none of the estimators, not even the DML with NN, can yield unbiased estimates in this non-HCM setup. \n\nWe appreciate the reviewer’s thoughtful comment and hope our response and additional experiments clarified the question.\n\n- - - -\n \n### **Other Comments**\n\n> [The paper] could better position itself in the literature\n\nWe thank the reviewer for the comment. We will revise our related literature section accordingly.\n\n- - - -\n\nWe sincerely appreciate the reviewer's constructive feedback and hope that the additional experimental findings and clarifying explanations address the points raised in the review. Should the reviewer find the response satisfactory, we would appreciate reconsidering the initial score. Otherwise, we remain fully committed to addressing any remaining concerns during the second author response phase.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes, I check the main proofs.", "experimental_designs_or_analyses": "The experimental setup is well-structured, with a clear dataset and a relevant benchmark.", "relation_to_broader_scientific_literature": "The work relates to prior research on DML for causal inference, particularly in tabular data settings (Chernozhukov et al., 2017; 2018).\nIt extends research on representation learning for causal inference (Veitch et al., 2019; 2020), but could better position itself in this literature. Connections to theoretical work on ILTs and non-identifiability (Dai et al., 2022) are relevant but lack empirical grounding.", "essential_references_not_discussed": "No, the related works are thoroughly summarized and appropriately referenced.", "other_strengths_and_weaknesses": "**Strengths**:\n\n-Novel application of DML to non-tabular data.\n\n-Theoretical contributions on ILTs and HCM are interesting and relevant.\n\n**Weaknesses**: \n\n-The manuscript lacks empirical validation of the impact of ILTs on ATE estimation.\n\n-While the benefits of HCM are theoretically motivated, they are not directly tested or empirically verified.\n\n-Critical ablation studies, such as comparing DML with raw data versus pre-trained features, are missing.\n\n-Although ILTs are extensively discussed, their connection to the experimental results is not clearly established or well-articulated.", "other_comments_or_suggestions": "-Consider adding empirical tests on ILTs to evaluate their impact on DML estimation.\n\n-Clarify how HCM enhances DML performance beyond theoretical claims, providing empirical evidence if possible.", "questions_for_authors": "1. How does ILT invariance specifically impact DML estimation?\n\n2. Why is there no experiment comparing DML with and without pre-trained representations?\n\n3. Are there scenarios where HCM does not hold?bIf real-world data does not follow an HCM structure, how does this affect ATE estimation?", "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": "M0wJjT7pRA", "reviewer": "Reviewer_gvFH", "summary": "The paper revisits the problem of estimating Average Treatment Effects (ATE) in observational studies under an assumption of ignorability where confounding factors are available as images or text (non-tabular data). The authors develop a theoretical argument that describes under what conditions the use of pre-trained neural networks to extract relevant features (representations) from non-tabular data gives valid adjustment within the Double Machine Learning (DML) framework. One key theoretical contribution is to show that the intrinsic dimension of the latent representation Z is invariant under linear transformations which combined with properties of the target function f(z) leads to fast convergence of the ATE.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### **Theoretical Results**\n\n> Why frame the results so closely to ATE estimation? Is there not an opportunity here to provide guarantees for any functional of pre-trained representations?\n\nIndeed, the derived convergence rates could be used for convergence guarantees of any functional of the pre-trained representations and thereby establish asymptotic results about other causal estimands other than the ATE. Given the popularity of the ATE (both in theory and practice), we chose to focus on the ATE estimand in the context of DML. However, as correctly pointed out by the reviewer, our convergence rates can be used more generally, e.g., to provide asymptotic inference results for estimands such as the average treatment effect of the treated (ATT) and the conditional ATE (CATE). For the latter, however, stronger assumptions would be required, e.g., $P$-OMS instead of $P$-valid pre-trained representations. \n\nBased on the reviewer’s excellent remark, we will add a discussion to Sec. 5 and elaborate on this aspect.\n\n> I would say that the biggest challenge for the correctness of the procedure is to guarantee Def. 3.1. It is highly non-trivial that the representation will accurately preserve the information content in the original data. (I guess this is a limitation of all DML [...])\n\nWe thank the reviewer for raising this relevant point. Indeed, validating the assumption of Def. 3.1 in practice is not trivial — as with any other causal assumption. However, this limitation is not specific to DML per se, but is instead inherent to any method aiming to estimate the ATE based on pre-trained representation for adjustment, given that Def. 3.1 (i) is necessary for the identification of the ATE in this context. We will make this clear in a revised version of the paper and thank the reviewer for bringing this up. \n\n> Sec. 5.2, in which f is parameterized by a feed forward NN, is not illustrated in the experiments (as far as I know). The authors instead use linear functions and random forest regressors. Is there a reason for this disconnect between theory and experiments? \n\nThis seems to be a misunderstanding. In all experiments, we use DML with neural networks as regressors. The “Label confounding” setup in Figs. 4 & 7 is simple enough that a single layer (“linear”) suffices; in the “Complex confounding” setup in Figs. 5 & 8, we use a 100-layer neural network. The RF regressor is included to highlight that methods not invariant to ILT transformations often fail when used on pre-trained representations. \n\nHence, our experiments are well in line with our theory. We will state this more clearly in the revised version. \n\n----\n\n### **HCM**\n\n> Classification problems typically require a non-linear transformation or a link function, do classification functions satisfy the HCM condition?\n\nYes, this is correct. More precisely, in the classification context, the conditional probability would satisfy the HCM condition. A non-linearity or link is not a problem in this context, since it is just a simple function in the final layer of the HCM. We thank the reviewer for this question and will clarify this in the revised version of the paper.\n\n> The HCM definition is a little bit difficult to visualize. Could more details or examples be given to better convey the intuition.\n\nWe thank the reviewer for this remark. To better illustrate the concept of the HCM in our paper, we will add a visualization of the HCM similar to the one from the original publication (Kohler and Langer, Annals of Statistics, 2021) and add further explanation to it.\n\n- - - -\n\nWe appreciate the reviewer's thoughtful remarks and important points related to the assumptions in our paper. We hope the additional clarifications in our response address the points raised in the review. Should the reviewer find the response satisfactory, we would appreciate reconsidering the initial score. Otherwise, we remain fully committed to addressing any remaining concerns during the second author response phase.", "claims_and_evidence": "All claims are very well supported both theoretically and empirically.", "methods_and_evaluation_criteria": "Yes, all good.", "theoretical_claims": "Did not check proofs in detail but follows established results in the literature.", "experimental_designs_or_analyses": "Yes, experimental design is sound.", "relation_to_broader_scientific_literature": "Yes, related work is discussed appropriately.", "essential_references_not_discussed": "No significant reference appears to be missing.", "other_strengths_and_weaknesses": "The paper is very well executed: it introduces existing convergence guarantees for function approximation, highlights the challenges well, and provides a compelling argument for the validity for using pre-trained features for adjustment. \n\nThe HCM definition is a little bit difficult to visualize. Could more details or examples be given to better convey the intuition.\n\nIt seems to me that a lot of the heavy lifting is done in the function approximation theorems of Secs. 4 and 5.1 and 5.2. Once those are established ATE estimation follows more or less straightforwardly. Why frame the results so closely to ATE estimation? Is there not an opportunity here to provide guarantees for any functional of pre-trained representations?\n\nWith my causality hat on, I would say that the biggest challenge for the correctness of the procedure is to guarantee Def. 3.1. It is highly non-trivial that the representation will accurately preserve the information content in the original data. (I guess this is a limitation of all DML work so might not need addressing here.)", "other_comments_or_suggestions": "See above.", "questions_for_authors": "Sec. 5.2, in which f is parameterized by a feed forward NN, is not illustrated in the experiments (as far as I know). The authors instead use linear functions and random forest regressors. Is there a reason for this disconnect between theory and experiments? If no non-linearities can be incorporated into the target function parameterization then that is a potential limitation of the theory.\n\nClassification problems typically require a non-linear transformation or a link function, do classification functions satisfy the HCM condition?", "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": "625UVzHsx8", "reviewer": "Reviewer_UvN8", "summary": "This paper explores how non-tabular data, such as images and text, can be incorporated into average treatment effect (ATE) estimation to account for confounding factors. The authors propose using latent features from pre-trained neural networks for adjustment. They formalize conditions under which these features enable valid ATE estimation, particularly in the double machine learning framework. The paper highlights challenges related to high-dimensional representations and non-identifiability but argues that neural networks can overcome these issues by adapting to intrinsic sparsity and dimensional structures, enabling fast convergence rates in treatment effect estimation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Additional experiments** can be found at https://anonymous.4open.science/r/icml2025-6599/add_exp_r3.pdf\n\n- - - -\n\n### **Theory & Practical Relevance**\n\n> I found this paper to be a primarily theoretical paper with minimal experimental support.\n\nThis is correct. Our paper is a theoretical contribution to the fast-growing literature that aims to incorporate non-tabular data and pre-trained models in causal inference procedures. While many papers do this empirically (as described in Sec. 2), we establish a novel set of theoretical conditions allowing for valid statistical inference of ATE estimation in this context. Our experiments mainly serve to illustrate these concepts. \n\nNonetheless, we followed the reviewer’s suggestion and added several additional experiments (see answer below). \n\n> [NNs] can achieve fast convergence rates by adapting to intrinsic notions of sparsity and dimension of the learning problem.\" [...] I couldn't find experimental support for these claims.\n\nWe provide empirical support for these claims in our experiments on complex confounding shown in Fig. 5 & 8. The confounding via latent features precisely mimics the idea of low intrinsic dimension and HCM structure of the target function. In contrast to the other depicted ATE estimators, NN-based nuisance estimation denoted by “DML (NN)” shows that NNs can adapt to the intrinsic notions of sparsity and dimension of the learning problem, thereby yielding unbiased ATE estimation. We thank the reviewer for the remark and will emphasize this in the revised version of the paper.\n\n> I find the formulation of the problem unpractical.[...] practically modeling [imaging data] does not really gain you any improved understanding [...] images as their surrogate seems to be a detour and over-complication. \n\nWe politely disagree with this viewpoint. Incorporating non-tabular data in ATE estimation to adjust for confounding does strongly impact scientific understanding and resembles what is done in practical application, as also described in our motivating examples in Sec. 1. In scenarios, where the confounder is available only embedded in non-tabular data (e.g. kidney stone or tumor size can only be measured via medical imaging), incorporating such modality in the ATE estimation is crucial to obtain unbiased estimates and draw valid scientific conclusions, as we show both theoretically and empirically in our paper.\n\n- - - -\n\n### **Empirical Evaluation & Additional Experiments**\n\n> Experimental setups are over simplified. Having [...] a [5-dim. AE latent rep.] for complex confounding have very limited generalizability to real life problems given the confounding effects in imaging applications are high-dimensional in nature.\n\nWe thank the reviewer for raising this point. However, all of the previously mentioned confounding “in nature”, e.g. kidney stone or tumor size, is in fact low dimensional yet embedded in the medical image in high dimensions, making it necessary to extract this information. Our experiments using both text and X-ray data precisely show that pre-trained models can be used to extract the latent confounding (low-dim.) information from the high-dim. non-tabular data modalities and achieve valid inference.\n\n> The study could explore way more different simulation scenarios, e.g., by varying simulation parameters.\n\nBased on the reviewer's suggestion, we conducted several additional experiments, in which we varied simulation parameters including the treatment assignment probabilities and the size of the latent dim. of the confounder.  The new results (Fig. 9-12) are in line with our previous results and reinforce our main theoretical findings. In particular, results are insensitive to changes in assignment prob. and latent dim. We thank the reviewer for this remark and think these results further strengthen our paper. Following the suggestion of reviewer vZqT we also conducted several other experiments, e.g. comparing DML with and without pre-training in Fig. 14.\n\n> Compared to other estimators, the proposed one does not seem to be better in Fig. 4&5.\n\nThis might be a misunderstanding. The estimators suggested by us (DML Linear and the DML NN) are the only estimators in Figs. 4 & 5 that yield unbiased ATE estimation with good coverage (overlap of confidence intervals with the true ATE, i.e., the red line) while the other estimators do not. Note that the unbiased “Oracle” estimator in Fig. 4 is an infeasible model in practice and serves as a gold standard comparison. \n\n- - - -\n\nWe thank the reviewer for the suggestions related to additional experiments. We hope the additional experiments in our response address the points raised in the review. Should the reviewer find the response satisfactory, we would appreciate reconsidering the initial score. Otherwise, we remain fully committed to addressing any remaining concerns during the second author response phase.", "claims_and_evidence": "I found this paper to be a primarily theoretical paper with minimal experimental support. For example, the last few arguments in the abstract are \"Common structural assumptions for obtaining fast convergence rates with additive or sparse linear models are shown to be unrealistic for latent features. We argue, however, that neural networks are largely insensitive to these issues. In particular, we show that neural networks can achieve fast convergence rates by adapting to intrinsic notions of sparsity and dimension of the learning problem.\" but I couldn't find experimental  support for these claims. \n\nMoreover, just like the paper, I find the formulation of the problem unpractical. Indeed, imaging data could be regarded as confounders but practically modeling this does not really gain you any improved understanding. All examples given in the paper have a simplistic form; i.e., the severity of disease, the size of the stone, or the extent of a fracture is a univariate confounder. I'm all for modeling these clear and interpretable confounding effects, but using images as their surrogate seems to be a detour and over-complication. In the end, we don't know whether image latent representations (extracted by a pre-trained network) actually contain those information or not.", "methods_and_evaluation_criteria": "Methods:\n\nThe formulation and theoretical derivation look correct to my eye, but I have to admit that I'm not a statistician so might not be able to fully appreciate the methodological and theoretical sophistication. I made my overall assessment assuming great merit in these formulations.\n\nEvaluation:\n\nThe shortcoming of the study is in the evaluation. The core analysis is two plots (Fig. 4&5) showing the estimated ATE by the proposed estimator and some baseline estimators. These two experiments are toy examples in nature and have extremely simple setup. Even confined to this experimental settings, the study could explore way more different simulation scenarios, e.g., by varying simulation parameters.", "theoretical_claims": "As mentioned above, The formulation and theoretical derivation look correct to my eye but I don't have the expertise to fully assess its correctness.", "experimental_designs_or_analyses": "* I feel the experimental setups are over simplified. Having a simple 0.7/0.3 ratio for simulating the basic label confounding and the 5-dimensional autoencoder latent representation for complex confounding have very limited generalizability to real life problems given the confounding effects in imaging applications are high-dimensional in nature.\n\n* Compared to other estimators, the proposed one does not seem to be better in Fig. 4&5.", "relation_to_broader_scientific_literature": "Yes, the paper presents the scenario that they want to dive into, which is a specific case in the broader confounder analysis literature.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "NA", "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:

Meta-review: This paper presents a theoretically motivated approach to ATE estimation using latent features extracted from pre-trained neural networks, with an emphasis on addressing non-tabular confounding in settings such as image-based data. 

The reviewers generally appreciate the depth and rigor of the theoretical contributions. One reviewer highlights the novelty in bridging DML and high-dimensional feature representations, while another notes the originality in adapting convergence rate analysis to neural estimators under nonstandard assumptions.

However, concerns remain, especially regarding the experimental component. Several reviewers point out that the empirical evaluation is limited to toy datasets and lacks critical ablation studies, such as comparisons between DML using raw data vs. latent features. Others question the practical relevance of HCM and ILTs in real-world data, noting that these assumptions, though theoretically sound, are difficult to validate empirically. Moreover, the absence of meaningful empirical support for claims about the robustness of neural network-based ATE estimation weakens the practical implications. Overall, I recommend accepting if there is room.
Decision: Accept (poster)