Original: Deep Learning Methods for Proximal Inference via Maximum Moment Restriction
Revised: Deep Learning Methods for Proximal Inference via Maximum Moment Restriction

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AfﬁliationAddress email
Revised: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: The No Unmeasured Confounding Assumption is widely used to identify causaleffects in observational studies.
Revised: The No Unmeasured Confounding Assumption is widely used to identify causaleffects in observational studies.

Original: Recent work on proximal inference has provided alternative identiﬁcation results that succeed even in the presence of unobserved confounders, provided that one has measured a sufﬁciently rich set of proxy variables ,satisfying speciﬁc structural conditions.
Revised: Recent work on proximal inference has provided alternative identiﬁcation results that succeed even in the presence of unobserved confounders, provided that one has measured a sufﬁciently rich set of proxy variables ,satisfying speciﬁc structural conditions.

Original: However, proximal inference requiressolving an ill-posed integral equation.
Revised: However, proximal inference requiressolving an ill-posed integral equation.

Original: Previous approaches have used a varietyof machine learning techniques to estimate a solution to this integral equation,commonly referred to as the bridge function .
Revised: Previous approaches have used a varietyof machine learning techniques to estimate a solution to this integral equation,commonly referred to as the bridge function .

Original: However, prior work has often beenlimited by relying on pre-speciﬁed kernel functions, which are not data adaptive andstruggle to scale to large datasets.
Revised: However, prior work has often beenlimited by relying on pre-speciﬁed kernel functions, which are not data adaptive andstruggle to scale to large datasets.

Original: In this work, we introduce a ﬂexible and scalablemethod based on a deep neural network to estimate causal effects in the presenceof unmeasured confounding using proximal inference.
Revised: In this work, we introduce a ﬂexible and scalablemethod based on a deep neural network to estimate causal effects in the presenceof unmeasured confounding using proximal inference.

Original: Our method achieves stateof the art performance on two well-established proximal inference benchmarks.
Revised: Our method achieves stateof the art performance on two well-established proximal inference benchmarks.

Original: Finally, we provide theoretical consistency guarantees for our method.
Revised: Finally, we provide theoretical consistency guarantees for our method.

Original:  Introduction
Revised:  Introduction

Original: Causal inference is concerned with estimating the effect of a treatment A on an outcome Y fromeither observational data or the results of a randomized experiment.
Revised: Causal inference is concerned with estimating the effect of a treatment A on an outcome Y fromeither observational data or the results of a randomized experiment.

Original: An estimand of primaryimportance is the average causal effect (ACE), which is the expected difference in Y caused bychanging the treatment from value a to a 0 for each unit in the study population, and is deﬁned asa contrast between the expected value of the potential outcomes at the two levels of the treatment:
Revised: An estimand of primaryimportance is the average causal effect (ACE), which is the expected difference in Y caused bychanging the treatment from value a to a 0 for each unit in the study population, and is deﬁned asa contrast between the expected value of the potential outcomes at the two levels of the treatment:

Original: E [ Y a 0 ]  E
Revised: E [ Y a 0 ]  E

Original: [ Y a ] .
Revised: [ Y a ] .

Original: However, in observational settings, the ACE is rarely equal to the observeddifference in conditional means, E [ Y | A
Revised: However, in observational settings, the ACE is rarely equal to the observeddifference in conditional means, E [ Y | A

Original: = a 0 ]  E
Revised: = a 0 ]  E

Original: [ Y |
Revised: [ Y |

Original: A = a ] due to confounding.
Revised: A = a ] due to confounding.

Original: In an attemptto eliminate the inﬂuence of confounding, investigators measure putative confounders X andsubsequently make adjustments for X in their analyses.
Revised: In an attemptto eliminate the inﬂuence of confounding, investigators measure putative confounders X andsubsequently make adjustments for X in their analyses.

Original: Given X , common approaches, such as standardization and inverse probability weighting (Hernánand Robins [1]), obtain valid estimates of the ACE given that the following assumptions hold: i)
Revised: Given X , common approaches, such as standardization and inverse probability weighting (Hernánand Robins [1]), obtain valid estimates of the ACE given that the following assumptions hold: i)

Original: Positivity: Pr
Revised: Positivity: Pr

Original: [ A = a | X = x ] > 0 for all x in the population, ii) Consistency: Y a = Y for allindividuals with A = a , iii) No unmeasured confounding which results in conditional exchangeability:
Revised: [ A = a | X = x ] > 0 for all x in the population, ii) Consistency: Y a = Y for allindividuals with A = a , iii) No unmeasured confounding which results in conditional exchangeability:

Original: Y a ??
Revised: Y a ??

Original: A | X ., and iv) No model misspeciﬁcation.
Revised: A | X ., and iv) No model misspeciﬁcation.

Original: While positivity can be empirically veriﬁed, the remaining 3 assumptions are not.
Revised: The assumptions are typically unveriﬁable for continuous data.

Original: While modelmispeciﬁcation is likely in all real-world scenarios, ﬂexible models and doubly robust estimators havebeen developed to mitigate the effect of this assumption[2].
Revised: While model mispeciﬁcation is likelyin all real-world scenarios, ﬂexible models and doubly robust estimators have been developed to mitigate the effect of this assumption[2].

Original: Therefore, the assumption of conditionalexchangeability, or equivalently, the No Unmeasured Confounding Assumption (NUCA), is thedeﬁning characteristic of this broad set of approaches to causal effect estimation (Hernán and Robins[3]).
Revised: Therefore, the assumption of conditional exchangeability, or equivalently, the No Unmeasured Confounding Assumption (NUCA), is the deﬁning characteristic of this broad set of approaches to causal effect estimation (Hernán and Robins [3]).

Original: However, in many settings, it is unrealistic to assume that we are able to measure a sufﬁcient setof confounders for A and Y such that conditional exchangeability holds.
Revised: However, in manysettings, it is unrealistic to assume that we are able to measure a sufﬁcient set of confounders for Aand Y such that conditional exchangeability holds

Original: Proximal inference is a recently introduced framework that allows for the identiﬁcation of causal effects even in the presence of unmeasured confounders [4, 5].
Revised: Proximal inference is a recently introduced framework that allows for the identiﬁcation of causal effects even in the presence of unmeasured confounders [4, 5].

Original: Proximal inference requires categorizingthe measured covariates into threemutually exclusive (but not exhaustive) groups: treatment-inducingproxy variables Z , outcome-inducing proxy variables W , and “backdoor” variables X that affectboth A and Y (i.e. typical confounders).
Revised: Proximal inference requires categorizingthe measured covariates into three groups: treatment-inducing proxy variables Z , outcome-inducingproxy variables W , and “backdoor” variables X that affect both A and Y (i.e. typical confounders).

Original: See Figure 1 for an example of a directed acyclic graph(DAG) that admits identiﬁcation under the assumptions of proximal inference.
Revised: See Figure 1 for an example of a directed acyclic graph (DAG) that admits identiﬁcation under theassumptions of proximal inference.

Original: The proxy sets W and Z must contain sufﬁcient information about the remaining unobserved confounders U , a conditionthat can be formalized by completeness assumptions.
Revised: The proxy sets W and Z must contain sufﬁcient informationabout the remaining unobserved confounders U , a condition that can be formalized by completenessassumptions.

Original: Under these and several other conditions, onecan estimate average potential outcomes from data even in the presence of unmeasured confounding.
Revised: Under these and several other conditions, one can estimate average potential outcomesfrom data even in the presence of unmeasured confounding.

Original: Proximal inference methods have potential applications in medical settings, where a natural questionis considering the effect of a treatment on an outcome in the presence of unmeasured confounding.
Revised: Proximal inference has potential applications in medical settings, where a natural question is the effect of a treatment on an outcome in thepresence of unmeasured confounding.

Original: Before applying proximal inference to real world problems, more validation is required before theycan be used safely to inform medical decision-making.
Revised: Before applying proximal inference to real world problems,more validation is required before they can be used safely to inform medical decision-making.

Original: Existing methods for proximal inference can be divided into two categories: two-stage regressionprocedures and methods that impose a maximum moment restriction (MMR).
Revised: Existing methods for proximal inference can be divided into two categories: two-stage regressionprocedures and methods that impose a maximum moment restriction (MMR).

Original: In two-stage regressionprocedures, the ﬁrst stage aims to predict outcome-inducing proxy variables W as a function of A ,
Revised: In two-stage regressionprocedures, the ﬁrst stage aims to predict outcome-inducing proxy variables W as a function of A ,

Original: X , and Z .
Revised: X , and Z .

Original: Then, the second stage regression estimates outcomes Y as a function of the predictedˆ W and the treatment A , and measured confounders X .
Revised: Then, the second stage regression estimates outcomes Y as a function of the predictedˆ W and the treatment A , and measured confounders X .

Original: Tchetgen Tchetgen et al.
Revised: Tchetgen Tchetgen et al.

Original: [5] introduced theﬁrst estimation technique for proximal inference which was a two-stage procedure that used a modelbased on ordinary least squares regression.
Revised: [5] introduced theﬁrst estimation technique for proximal inference which was a two-stage procedure that used a modelbased on ordinary least squares regression.

Original: Mastouri et al.
Revised: Mastouri et al.

Original: [6] extended this framework by replacingsimple linear regression with kernel ridge regression.
Revised: [6] extended this framework by replacingsimple linear regression with kernel ridge regression.

Original: Xu et al.
Revised: Xu et al.

Original: [7] increased feature ﬂexibility furtherby incorporating neural networks as feature maps instead of kernels.
Revised: [7] increased feature ﬂexibility furtherby incorporating neural networks as feature maps instead of kernels.

Original: In contrast, MMR methods are single-stage procedures to estimate average potential outcomes.
Revised: In contrast, MMR methods are single-stage procedures to estimate average potential outcomes.

Original: Muandet et al.
Revised: Muandet et al.

Original: [8] introduced MMR for reproducing kernel Hilbert spaces (RKHS).
Revised: [8] introduced MMR for reproducing kernel Hilbert spaces (RKHS).

Original: MMR critically relieson the optimization of a V-statistic or U-statistic for learning a function needed to calculate the ACE.
Revised: MMR critically relieson the optimization of a V-statistic or U-statistic for learning a function needed to calculate the ACE.

Original: Zhang et al.
Revised: Zhang et al.

Original: [9] used an MMR method to obtain point identiﬁcation of the ACE in the instrumental variable (IV) setting and incorporated neural networks into their method by training with the V-statistic asa loss function and optimized using stochastic gradient descent.
Revised: [9] used an MMR method to obtain point identiﬁcation of the ACE in the instrumentalvariable (IV) setting and incorporated neural networks into their method trained with the V-statistic asa loss function and optimized using stochastic gradient descent.

Original: Mastouri et al.
Revised: Mastouri et al.

Original: [6] demonstrated thatthe MMR framework with kernel functions can be used for proximal inference as well as IV regression.
Revised: [6] demonstrated thatthe MMR framework with kernel functions can be used for proximal inference as well as IV regression.

Original: In this work, we introduce a new method, Neural Maximum Moment Restriction (NMMR) which isa ﬂexible neural network approach that is trained to minimize a loss function derived from either a
Revised: In this work, we introduce a new method, Neural Maximum Moment Restriction (NMMR) which isa ﬂexible neural network approach that is trained to minimize a loss function derived from either a

Original: U-statistic or V-statistic to satisfy MMR in the proximal setting.
Revised: U-statistic or V-statistic to satisfy MMR in the proximal setting.

Original: The method introduced in this workmakes several novel contributions to the proximal inference literature:
Revised: The method introduced in this workmakes several novel contributions to the proximal inference literature:

Original: • We introduce a new, single stage method based on neural networks for estimating potential
Revised: • We introduce a new, single stage method based on neural networks for estimating potential

Original: outcomes and the ACE in the presence of unmeasured confounding.
Revised: outcomes and the ACE in the presence of unmeasured confounding.

Original: • We provide new theoretical consistency guarantees for our method.
Revised: • We provide new theoretical consistency guarantees for our method.

Original: • We demonstrate state-of-the-art (SOTA) performance on two well-established proximalinference benchmark tasks.
Revised: • We demonstrate state-of-the-art (SOTA) performance on two well-established proximal inference benchmark tasks.

Original: • We show for the ﬁrst time how to incorporate domain-speciﬁc inductive biases using aconvolutional model on a proximal inference task that uses images.
Revised: • We show for the ﬁrst time how to incorporate domain-speciﬁc inductive biases using aconvolutional model on a proximal inference task that uses images.•

Original: • We provide the ﬁrst unbiased estimate of the MMR risk function using the U-statistic rather
Revised: We provide the ﬁrst unbiased estimate of the MMR risk function using the U-statistic rather

Original: than V-statistic in the proximal setting.
Revised: than V-statistic in the proximal setting.

Original:  Background: Proximal Inference
Revised:  Background: Proximal Inference

Original: Throughout the rest of this manuscript we will use capital letters (e.g. A ) to denote random variablestaking values in measurable spaces, (e.g. A ; typically subsets of R ). Lower case lettersdenoterealizations of these random variables (e.g. A = a ).
Revised: In this manuscript, capital, caligraphic, and lower case letters (e.g. A , A , and a ) denote randomvariables and their corresponding domains and realizations, respectively.

Original: Estimates of random variables will be indicatedusing “hat” notation, e.g. ˆ Y would be an estimate for the random variable Y .
Revised: Estimates of randomvariables and functions will be indicated by hats, e.g. ˆ Y and ˆ h are estimates of Y and h , respectively.

Original: Our goal is to estimate the average potential outcome a speciﬁed level ofan treatment, i.e. E [ Y a ]isthe average potential outcome under the treatment a .
Revised: Our goal is to estimate, for each level of treatment, a , the expected potential outcome E [ Y a ] .

Original: Without loss of generality, we refer to A as atreatment, though it could refer to any (possibly continuous) intervention.
Revised: Withoutloss of generality, we refer to A as a treatment, though it could refer to any (possibly continuous)intervention.

Original: Proximal inference allowsfor the estimation of mean potential outcomes for Y a in the presence of unobserved confounders Uprovided that we have a sufﬁciently rich set of proxies Z, W along with observed confounders Xthat obey certain structural assumptions.
Revised: Proximal inference allows us to do this, in the presence of unobserved confounders, U , provided that we have a sufﬁciently rich set of proxies, ( W, Z ) , that obey certain structuralassumptions. We may also include observed confounders, X .

Original: The assumptions needed for identiﬁcation are given below:
Revised: We also require the following:

Original: Assumption 1.
Revised: Assumption 1.

Original:  A, U, W, X, Y, Z satisfy the conditional independences Y ??
Revised: Given ( A, U, W, X, Y, Z ) , Y ??

Original: Z | A, U, X and W ??
Revised: Z | A, U, X and W ??

Original: 90 ( A, Z ) | U, X .
Revised: ( A, Z ) | U, X .

Original: Note: Assumption 1 can be equivalently stated as the requirement that these random variables mustbe generated by one of 22 valid causal directed acyclic graphs (DAGs) for proximal inference [5].

Original: Figure 1 provides an example of a DAG that satisﬁes these assumptions.
Revised: Figure 1 provides an example of a DAG that satisﬁes these assumptions.

Original: We can also formalizethenotion of “rich enough” proxies with two assumptions about the completeness of U and Z :
Revised: The following completenessconditions formalize the notion that the proxies are “sufﬁciently rich”:

Original: Assumption 2.
Revised: Assumption 2.

Original: For any square integrable function g , E [ g ( U ) |
Revised: For all f 2 L 2 and all a 2 A , x 2 X , E [ f ( U ) |

Original: A = a, X = x, Z = z ] = 0 , [Equation] if and only if 
Revised: A = a, X = x, Z = z ] = 0 for all 90 z 2 Z if and only if f ( U )

Original: g

Original: ( u ) =  almost surely.
Revised: = 0 almost surely.

Original: Assumption 3.
Revised: Assumption 3.

Original: For any square integrable function g , E [ g ( Z ) |
Revised: For all f 2 L 2 and all a 2 A , x 2 X , E [ f ( Z ) |

Original: A = a, W = w, X = x ] = 0 , 8 ( a, w, x ) 2A ⇥ W⇥ X if and only if g ( z ) = 0 almost surely.
Revised: A = a, W = w, X = x ] = 0 for all 92 w 2 W if and only if f ( Z ) = 0 almost surely.

Original: We will use two other assumptions at various points in the paper.
Revised: We will use two other assumptions at various points in the paper.

Original: The ﬁrst guarantees the uniquenessof the bridge function, while the second ensures the risk function does not have false zeroes.
Revised: The ﬁrst guarantees the uniquenessof the bridge function, while the second ensures the risk function does not have false zeroes.

Original: Assumption 4.
Revised: Assumption 4.

Original: E [ f ( A, W, X ) | A, X, Z ] = 

Revised: E [ f ( A, W, X ) | A, X, Z ] =

Original: P A,X,Z -almost
Revised: P A,X,Z -almost

Original: surely if and only if f ( A, W, X ) = 

Revised: surely if and only if f ( A, W, X ) =

Original: P A,W,X -almost surely.
Revised: P A,W,X -almost surely.

Original: Assumption 5.
Revised: Assumption 5.

Original: k : ( A ⇥ X ⇥ Z ) 2 !
Revised: k : ( A ⇥ X ⇥ Z ) 2 !

Original: R is continuous, bounded, and Integrally Strictly Positive 104 Deﬁnite (ISPD), so that
Revised: R is continuous, bounded, and Integrally Strictly Positive 98 Deﬁnite (ISPD), so that

Original: [Equation] 0 ) f ( ⇠ 0 )
Revised: R f ( ⇠ ) k ( ⇠ , ⇠ 0 ) f ( ⇠ 0 )

Original: d ⇠ d ⇠ 0 > 

Revised: d ⇠ d ⇠ 0 >

Original: if and only if f 6 = 0 P A,Z,X -almost surely.
Revised: if and only if f 6 = 0 P A,Z,X -almost surely.

Original: Assumptions 1-3 together with several regularity assumptions (see assumptions (v)-(vii) in [4]) ensurethat there exists a function h such that:
Revised: Assumptions 1-3 together with several regularity assumptions (see assumptions (v)-(vii) in [4]) ensurethat there exists a function h such that:

Original: [Equation]
Revised: [Equation]

Original: Equation 1 is a Fredholm integral equation of the ﬁrst kind and in general is difﬁcult to solve for h . The function h  is often referred to as the “bridge function” and Theorem 1 of Miao et al.
Revised: Equation 1 is a Fredholm integral equation of the ﬁrst kind; its solution, h , is often called the “bridgefunction.” Theorem 1 of Miao et al.

Original: [4] alsoshowed that if one has a solution h , then average potential outcomes can be computed as:
Revised: [4] shows that the expected potential outcomes are given by:

Original: [Equation]
Revised: [Equation]

Original: Thus, under the proximal inference framework an unbiased estimate of the average potential outcome
Revised: We can obtain unbiased estimates of the expected potential outcomes, E

Original: E [ Y a ] is obtained by  ﬁrst estimating the bridge function h , by some ˆ h , and then ﬁxing a and takingthe empirical average of ˆ h over the remaining variables using a held-out dataset with Mdata points
Revised: [ Y a ] , by splitting thesample, using the ﬁrst part of the data to estimate the bridge function h , by some ˆ h , and using thesecond part of the data to compute the empirical mean of ˆ h with a ﬁxed to the value of interest,ˆ E [ Y a ] = 1 M

Original: [Equation]
Revised: P Mi =1 ˆ h ( a, w i , x i ) .

Original: Once these mean potential outcomes havebeen estimated, then contrasts of interest like the ACE can be calculated in a straight-forward manner.
Revised: From these, we can obtain other quantities of interest, like the ACE. In what follows, all norms will be L 2 with respect to the relevant probability measure, unlessotherwise noted.

Revised: If necessary, we will explicitly denote the norm of f with respect to a probabilitymeasure P U,V,... by k f k P U,V,... or k f k 2 , P U,V,... .

Original:  Related Work
Revised:  Related Work

Original: Kuroki and Pearl [10] ﬁrst established identiﬁcation of a causal effect in the setting of unobservedconfounders by leveraging noisy proxy variables and estimating the distribution of U using external datasets that allowed for estimation of a distribution p ( w | u ) .
Revised: Kuroki and Pearl [10] ﬁrst established identiﬁcation of a causal effect in the setting of unobservedconfounders by leveraging noisy proxy variables, W , to “recover” the distribution of U , potentiallyusing external datasets to estimate p ( w | u ) .

Original: These results were extended to allow foridentiﬁcation without recovery of U in Miao et al.[4] and Tchetgen Tchetgen et al.
Revised: Tchetgen-Tchetgen and colleagues extended these resultsto allow for identiﬁcation without recovery of U in Miao et al.

Original: In particular,
Revised: [4] and Tchetgen Tchetgen et al.[5] also providing a 2-Stage Least Squares (2SLS) method to identify and estimate causal effectsunder the assumption that the bridge function is linear.

Original: Tchetgen Tchetgen et al.
Revised: Cui et al.

Original: [5] set out the 22 valid causal DAGs that satisfy Assumption 1 and allowfor identiﬁcation provided one has access to asufﬁciently rich set of proxy variables. These authorsalso provide a 2-Stage Least Squares (2SLS) method to identify and estimate causal effects underthe assumption that the bridge function has a linear form.
Revised: [11] introduced a bridge function for Inverse Probability Weighting (IPW), which enabled IPW and Doubly Robust (DR) proximalestimators, for which they presented inﬂuence functions under binary treatment. This was extended in

Original: Ghassami et al.
Revised: Ghassami et al.

Original: [11] established a doublyrobust method for proximal inference using inﬂuence functions.
Revised: [12] and further explored by Kallus et al. [13], who provided alternative identiﬁcationassumptions as well as results for general treatments. The latter two works also consider the use ofadversarial methods for estimation, which were previously utilized by Lewis and Syrgkanis [14] and Bennett, et al.

Original: Earlier work applying machine learning techniques to the task of identiﬁcation in the  setting ofproximal inference included Deaner [12], which relied on a two-stage penalized sieve distanceminimization.
Revised: Bennett et al. [15] in the Instrumental Variable (IV) setting and Dikkala et al. [16] inconditional moment models. Other early investigators include Deaner [17] who developed machine learning techniques forproximal inference introducing a method based on a two-stage penalized sieve distance minimization.

Original: Several later works similarly employed two-stage regressions with increasinglyﬂexible basis functions to estimate potential outcomes.
Revised: Several later works similarly employed two-stage regressions with increasingly ﬂexible basis functionsto estimate potential outcomes.

Original: Mastouri et al.
Revised: Mastouri et al.

Original: [6] developed a two-stagekernel ridge regression (Kernel Proxy Variables “KPV”) to estimate the bridge function h . Whilea kernel basis has more ﬂexibility than the linear basis of Tchetgen Tchetgen et al.
Revised: [6] developed a two-stage kernel ridge regression(Kernel Proxy Variables “KPV”) to estimate the bridge function h , allowing more ﬂexibility thanthe linear basis of Tchetgen Tchetgen et al.

Original: [5], Xu et al.[7] introduced an adaptive basis derived from neural networks.
Revised: Xu et al. [7] further improved upon this with anadaptive basis derived from neural networks.

Original: Their two stage regression method, 
Revised: Their two stage regression method, Deep Feature Proxy

Original: Deep Feature Proxy Variables (DFPV), established the previous SOTA performance on the proximalbenchmark tasks that we consider in our work.
Revised: Variables (DFPV), established the previous SOTA performance on the proximal benchmark tasks thatwe consider in our work.

Revised: Singh and colleagues also considered two stage kernel models in the IVsetting [18] and RKHS techniques for proximal inference [19].

Original: An alternative approach based on maximum moment restriction (MMR) uses single-stage estimatorsof the bridge function.
Revised: An alternative approach based on maximum moment restriction (MMR) uses single-stage estimatorsof the bridge function.

Original: MMR-based methods were established in Muandet et al.
Revised: MMR-based methods were established in Muandet et al.

Original: [8] as a way toenforce conditional moment restrictions [13].
Revised: [8] as a way toenforce conditional moment restrictions [20].

Original: Zhang et al.
Revised: Zhang et al.

Original: [9] introduced the MMR framework totheinstrumental variable (IV) setting. The IV setting assumes a speciﬁc DAG where unobservedconfounders can becircumvented by use of an instrument that is not confounded by U . IV DAGs canbe considered a subset of proximal DAGs without outcome-inducing proxies W  [5].
Revised: [9] introduced the MMR framework tothe IV setting, which can be considered a subset of proximal inference without outcome-inducingproxies W and with additional exclusion restrictions [5].

Original: There are nowseveral machine learning methods that can be applied in the IV setting [14–18]
Revised: There are now several machine learningmethods that can be applied in the IV setting [21, 14, 18, 16, 22]

Original: Of note, Zhang et al.
Revised: Of note, Zhang et al.

Original: [9] introduced MMR-IV which is related to  Lewis and Syrgkanis [15] and Dikkala et al.
Revised: [9] introduced MMR-IV which is related to work by Lewis and Syrgkanis[14] and Dikkala et al.

Original: MMR-IV involves optimizing a family of risk functions based on U- or V-statistics[19].
Revised: MMR-IV involves optimizing a family of risk functions based on U or V-statistics [23].

Original: However, Zhang et al.
Revised: However, Zhang et al.

Original: [9] only consider the IV setting, which signiﬁcantly differs from theproximal setting because of the absence of outcome-inducing proxy variables. Additionally, Zhanget al. [9] only optimize neural networks by a loss that corresponds to the V-statistic.
Revised: [9] only consider IV, rather than proximal, inference andonly optimize neural networks by a loss that corresponds to the V-statistic.

Original: The V-statisticprovides a biased estimate [19] of its corresponding risk function, such as R ( h ) in Equation 3.
Revised: The V-statistic provides abiased estimate [23] of its corresponding risk function, such as R ( h ) in Equation 3.

Original: Finally, Mastouri et al.
Revised: Finally, Mastouri et al.

Original: [6] introduced an MMR-based method for proximal inference called Proximal
Revised: [6] introduced an MMR-based method for proximal inference called Proximal

Original: Maximum Moment Restriction (PMMR).
Revised: Maximum Moment Restriction (PMMR).

Original: PMMR extends the MMR framework to the proximalsetting through the use of kernel functions and also optimizes Equation 3 via a V-statistic.
Revised: PMMR extends the MMR framework to the proximalsetting through the use of kernel functions and also optimizes Equation 3 via a V-statistic.

Original: For acomparison of our model to PMMR and MMR-IV, see Table 1.
Revised: For acomparison of our model to PMMR and MMR-IV, see Table 1.

Original:  Our Method: Neural Maximum Moment Restriction (NMMR)
Revised:  Our Method: Neural Maximum Moment Restriction (NMMR)

Original: In this work we propose Neural Maximum Moment Restriction (NMMR) as a method to estimateaverage potential outcomes
Revised: In this work we propose Neural Maximum Moment Restriction (NMMR): a method to estimateexpected potential outcomes

Original: E [ Y a ] in the presence of unmeasured confounding.
Revised: E [ Y a ] in the presence of unmeasured confounding.

Original: We consider thehypothesis class of h belonging to the family of deep neural networks given their well-establishedﬂexibility, ability to scale to large data sets, and for the ease with which domain-speciﬁc inductivebiases can be included directly into the model (e.g. convolutions for images).
Revised: We use deepneural networks due to their ﬂexibility, scalability, and adaptability to diverse data types (e.g. usingconvolutions for images).

Original: Using a connectionbetween maximum moment restrictions [8]  and U- and V-statistics, we show how a single-stageneural network procedure can be used to provide estimates of the bridge function.
Revised: By rewriting maximum moment restrictions [8] as U and V-statistics [6, 9],we show how a single-stage neural network procedure can be used to estimate the bridge function.

Original: Following Muandet et al.
Revised: Following Muandet et al.

Original: [8] and Zhang et al.
Revised: [8] and Zhang et al.

Original: [9], wecan rewrite the integral equation (1)as a conditional moment restriction : E [ Y  h
Revised: [9], we rewrite the integral equation (1) as a conditionalmoment restriction : E [ Y  h

Original: ( A, W, X ) | A,
Revised: ( A, W, X ) | A,

Original: X, Z ] = 0 .
Revised: X, Z ] = 0 .

Original: Then,E [( Y  161 h ( A, W, X )) g ( A, X, Z ) | A, X, Z ] = 0 for any measurablefunction g on A ⇥ X ⇥ Z  ,as well.
Revised: Then, for any measurable g : A ⇥ X ⇥ Z ! R ,

Original: By taking the expectation over A, X, Z we obtain E [( Y  h
Revised: E [( Y  h ( A, W, X )) g ( A, X, Z ) | A, X, Z ] = 0 , and, thus, E [( Y  h

Original: ( A, W, X ))
Revised: ( A, W, X ))

Original: g ( A, X, Z )]
Revised: g ( A, X, Z )]

Original: = eliminatingthe need to take conditional expectations entirely.
Revised: = 0 ,so we can use unconditional expectations instead of conditional ones.

Original: The intuition for this result is that any solution forh in the original conditional moment restriction will cause E [ Y  h

Original: ( A, W, X ) | A, Z, X ] to be 0 and,thus, the product with g ( A, Z, W ) will also be 0 , as well as the unconditional expectation, regardlessof which function g one chooses.

Original: This creates an inﬁnite number of moment restrictions and servesas the basis of the MMR framework of Muandet et al.
Revised: This is the basis of the MMRframework of Muandet et al.

Original: This suggests a minimax strategy ofsearching for an estimateof h such that risk R ( h ) for the worst-case value of g is minimized:
Revised: Since this yields an inﬁnite number of moment restrictions weemploy a minimax strategy to estimate h by minimizing the risk R ( h ) for the worst-case value of g :

Original: [Equation]
Revised: [Equation]

Original: Following Zhang et al.
Revised: Following Zhang et al.

Original: [9]’s work in the IV setting, Mastouri et al.
Revised: [9]’s work in the IV setting, Mastouri et al.

Original: [6] (Lemma 2) showed that, if gis an element of an RKHS, R ( h ) can be rewritten in the form R k ( h )
Revised: [6] (Lemma 2) showed that, if gis an element of an RKHS, R ( h ) can be rewritten in the form R k ( h )

Original: =
Revised: =

Original: E [( Y  h
Revised: E [( Y  h

Original: ( A, W, X ))
Revised: ( A, W, X ))

Original: ( Y 0 h ( A 0 , W 0 , X 0 ))
Revised: ( Y 0 h ( A 0 , W 0 , X 0 ))

Original: k
Revised: k

Original: (( A, X, Z ) , ( A 0 , X 0 , Z 0 ))
Revised: (( A, X, Z ) , ( A 0 , X 0 , Z 0 ))

Original: ] where ( A 0 , W 0 , X 0 , Y 0 , Z 0 ) are independent copies ofthe random variables ( A, W, X, Y, Z ) and k :
Revised: ] where ( A 0 , W 0 , X 0 , Y 0 , Z 0 ) are independent copies ofthe random variables ( A, W, X, Y, Z ) and k :

Original: ( A ⇥ Z ⇥ X ) 2 !
Revised: ( A ⇥ Z ⇥ X ) 2 !

Original: R is a continuous, bounded,and Integrally Strictly Positive Deﬁnite (ISPD) kernel.
Revised: R is a continuous, bounded,and Integrally Strictly Positive Deﬁnite (ISPD) kernel.

Original: Then, if h satisﬁes R k ( h ) = 0 , E [ Y h ( A, W, X ) | A, X, Z ] = 

Revised: Then, if h satisﬁes R k ( h ) = 0 , E [ Y h ( A, W, X ) | A, X, Z ] =

Original: P A,X,Z -almost surely.
Revised: P A,X,Z -almost surely.

Original: Thus, if we can ﬁnd a neural network
Revised: Thus, if we can ﬁnd a neural network

Original: h thatsatisﬁes R k ( h ) = 0 , we will have obtained a P A,X,Z -almost sure solution to Equation 1 and cancompute any desired average potential outcome with Equation 2.
Revised: h thatsatisﬁes R k ( h ) = 0 , we will have obtained a P A,X,Z -almost sure solution to Equation 1 and cancompute any expected potential outcome by using Equation 2.

Original: We can minimize R k ( h ) for data D = { ( a i , w i , x i , y i , z i ) } Ni =1 by minimizing the expectation in
Revised: The empirical risk ˆ R k,n given data D = { ( a i , w i , x i , y i , z i ) } Ni =1 can be written as either a U or

Original: Mastouri et al. [6] (Lemma 2), which can be approximated by either a V-statistic [19] [Equation] or with a U-statistic [19]:
Revised: V-statistic, respectively [23]:

Original: [Equation]
Revised: [Equation]

Original: where h i = h ( a i , w i , x i ) and k ij = k
Revised: where h i = h ( a i , w i , x i ) and k ij = k

Original: (( a i , z i , x i ) , ( a j , z j , x j )) .
Revised: (( a i , z i , x i ) , ( a j , z j , x j )) .

Original: ˆ R U ( h ) is the minimum variance unbiased estimator of R k ( h ) [19], while ˆ R V ( h ) is a biased estimator of R k ( h ) .
Revised: ˆ R k,U,n ( h ) is the minimum varianceunbiased estimator of R k ( h ) [23], while ˆ R k,V,n ( h ) is a biased estimator of R k ( h ) .

Revised: In order to preventoverﬁtting, we add an additional penalty to our risk function ⇤ :

Revised: H ⇥ ⇥ H !

Revised: R + , which is a function 174 of h as well as, possibly, its parameters, ✓ h (e.g. network weights).

Original: In practice, a penalized version of ˆ R U ( h ) has a slight amount of bias, but in practice is much less biased than even an unpenalized ˆ R V ( h ).
Revised: Speciﬁcally, we take ⇤ to bean L 2 penalty on the weights so ⇤ [ h, ✓ h ] = P i ✓ 2 h,i . We then denote the penalized risk functionsby ˆ R k,U,  ,n ( h )

Revised: =

Revised: ˆ R k,U,n ( h ) +  ⇤ [ h, ✓ h ] and ˆ R k,V,  ,n ( h ) =

Revised: ˆ R k,V,n ( h ) +  ⇤

Revised: [ h, ✓ h ] , respectively.

Revised: In practice, ˆ R k,U,  ,n ( h ) is slightly biased, but, in simulations, is much less biased than even theunpenalized ˆ R k,V,n ( h ) .

Original: Previous work either did not consider the U-statistic [6], or did not utilize the U-statistic [9].
Revised: Previous work either did not consider the U-statistic [6], or did not utilizethe U-statistic [9].

Original: In our work, we introduce two variants of our method, NMMR-U and NMMR-V, where the former is optimized with a U-statistic and the latter a V-statistic.
Revised: In our work, we introduce two variants of our method, NMMR-U and NMMR-V,where the former is optimized with a U-statistic and the latter a V-statistic.

Original: We train the neural networks in both variants with a regularized loss function:
Revised: We train the neuralnetworks in both variants with the regularized loss function:

Original: [Equation]
Revised: [Equation]

Original: where ( Y  h ( A, W, X )) is a vector of residuals from the neural network’s predictions and K is akernel matrix with entries k ij .
Revised: where ( Y  h ( A, W, X )) is a vector of residuals from the neural network’s predictions and K is akernel matrix with entries k ij .

Original: Throughout, we choose k to be the RBF kernel (see Appendix B).
Revised: We choose k to be an RBF kernel (see Appendix B).

Original: If 
Revised: If L represents a

Original: L is representing a V-statistic,then we include  main diagonal elements of K . If L is representing a U-statistic,then we set the main diagonal to be 0.
Revised: V-statistic, we include the main diagonal elements of K , while if L represents a U-statistic, we setthe main diagonal to be 0.

Original: Once we’ve obtained an optimal neural network ˆ h , we can compute an estimate of the averagepotential outcome with data from a held-out dataset with M data pointsD W , X = { ( w i , x i ) } Mi =1 ,ˆ E [ Y a ] = 1 M P
Revised: Once we’ve obtained an optimal neural network ˆ h , we can compute anestimate of the expected potential outcome with data from a held-out dataset with M data points

Original: [Equation]
Revised: [Equation]

Original: i )

Original: In contrast to PMMR [6], which uses kernels as feature maps for proxy and treatment variables,
Revised: In contrast to PMMR [6], which uses kernels as feature maps for proxy and treatment variables,

Original: NMMR uses adaptive feature maps from neural networks.
Revised: NMMR uses adaptive feature maps from neural networks.

Original: NMMR is similar to MMR-IV [9], but
Revised: NMMR is similar to MMR-IV [9], but

Original: MMR-IV is restricted to the instrumental variable (IV) setting rather than the proximal inference setting.
Revised: MMR-IV is restricted to the instrumental variable (IV) setting rather than the proximal inference setting.

Original: Table 1 places NMMR in context with existing methods for proximal inference and IV regression. 
Revised: Table 1 places NMMR in context with existing methods for proximal inference and IV regression. [Table] [Table] Kernels

Original: [Table]
Revised: [Table]

Original: Consistency of NMMR
Revised: statistic NMMR-V (ours)

Original: In this section we provide a probabilistic bound on the distance of the estimated bridge function,ˆ h k,  ,n
Revised: Proximal [Table] V- [Table]

Revised: NMMR

Original: [Table]
Revised: [Table]

Original: from the true bridge function
Revised: Networks

Original: [Table]
Revised: [Table]

Original:  h ⇤ , in terms of the Radamacher complexity R n ( F ) of a class of functions F derived from elements of the hypothesis space H and the ﬁxed kernel, k .
Revised: In this section we provide a probabilistic bound on the distance of the estimated bridge function,ˆ h k,  ,n , from the true bridge function, h ⇤ , in terms of the Radamacher complexity R n ( F ) of a classof functions F derived from elements of the hypothesis space H and the ﬁxed kernel, k .

Original: Notethat ˆ R k,  ,n ( h ) =
Revised: Note thatˆ R k,U,  ,n ( h ) =

Original: ˆ R U ( h ) +  k h k 22 (see Section 4).
Revised: ˆ R k,U,n ( h ) +  ⇤ [ h, ✓ h ] (see Section 4).

Original: We use this bound to demonstrate that, undermild conditions, ˆ h k,  ,n converges in probability to h ⇤ , and that, under an additional completenessassumption, h ⇤ is unique P A,W,X -almost surely.
Revised: We use this bound to demonstrate that, undermild conditions, ˆ h k,  ,n converges in probability to h ⇤ , and that, under an additional completenessassumption, h ⇤ is unique P A,W,X -almost surely.

Original: This provides a consistent estimate of E [ Y a ] .
Revised: This provides a consistent estimate of E [ Y a ] .

Original: Theorem 1.
Revised: Theorem 1.

Original: Let ˜ h
Revised: Let ˜ h

Original: k minimize R k
Revised: k minimize R k

Original: ( h ) and ˆ h k,  ,n minimize ˆ R k,  ,n ( h ) for h 2 H and let h ⇤ :
Revised: ( h ) and ˆ h k,U,  ,n minimize ˆ R k,U,  ,n ( h ) for h 2 H , k :

Original:  A ⇥
Revised: ( A ⇥

Revised: X ⇥ Z ) 2 !

Revised: [  M k , M k ] , ⇤ :

Revised: H ⇥ ⇥ h !

Revised: [  M  , M  ] , and let h ⇤ :

Original:  W ⇥ X !
Revised: A ⇥ W ⇥ X !

Original: R satisfyE
Revised: R satisfy

Original:  [ Y  h
Revised: E [ Y  h

Original: ⇤ ( A, W, X ) | A, X, Z ] = 

Revised: ⇤ [Equation] A, W, X ) | A, X, Z ] =

Original: P A,X,Z -almost surely, where
Revised: P A,X,Z -almost surely, where

Original: [Equation]
Revised: [Equation]

Original: Also let,
Revised: Also let,

Original: [Equation]
Revised: [Equation]

Original: Then, d 2 k ( h ⇤ , h ) =
Revised: Then, d 2 k ( h ⇤ , h ) =

Original: R k ( h ) and, with probability at least 1   ,
Revised: R k ( h ) and, with probability at least 1   ,

Original: [Equation]
Revised: [Equation]

Original: ⇣ h ⇤ , ˜ h
Revised: ˆ

Original: k ⌘ +  M 2 + 8 M E A,X,Z  R n  F 0 A,X,Z  + R n  F 0 A,X,Z  + 16 M 2 M k ✓ n log  ◆ + 10 (2 log 2) 12 M 2 M k n  d 2 k ⇣
Revised: [Equation]

Original: h ⇤ , ˜ h k ⌘ +  M 2 + 8 M ( R n  1 ( F 0 ) + R n ( F 0 ))
Revised: k,U,  ,n

Original: [Equation]
Revised: [Equation]

Revised: M

Original: [Equation]
Revised: [Equation]

Original: 2 M k ✓ n log
Revised: Further, if Assumption 5 holds, so k is ISPD, then d k is a metric on L 2 AXZ and, if the right hand side 205 of the inequality goes to zero as n goes to inﬁnity,

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: E h

Revised: ˆ h k

Original: [Equation]
Revised: [Equation]

Original: 201 Further, if Assumption 5 holds, so k is ISPD, then d k is a metric and, if the right hand side goes tozero as n goes to inﬁnity,
Revised:

Original:  [Equation]
Revised: A, [Equation]

Revised: , Z i

Revised: P !

Revised: E [ h ⇤ |

Revised: A, X, Z ]in d k .

Revised: Also,

Revised: [Equation]

Revised:

Revised: [Equation]

Revised: i

Revised: P

Revised: [Equation]

Revised: E [ h ⇤ | A, X, Z ] in L 2 (P A , X , Z )  norm .

Revised: [Equation]

Revised:  a

Revised: [Equation]

Revised: w

Revised: [Equation]

Revised: x

Revised: [Equation]

Revised: z = h  a 0 , w 0 , x k

Revised:  a 0 , x 0 , z , ( a, x, z )

Revised: [Equation]

Revised: k

Original:  Corollary 8 provides a similar result for V-statistic estimators of R ( h ) , meaning we can choose to useeither U or V-Statistics and have similar guarantees.
Revised: [Equation] a [Equation]  Corollary 8 provides a similar result for V-statistic estimators of R ( h ) , meaning we can choose to useeither U or V-Statistics and have similar guarantees.

Original:  Theorem 2.
Revised: In Theorem 1 if the quadratic form converges ata particular rate, say n  12 , E h ˆ h  A, X, Z i

Revised: P !

Revised: E [ h ⇤ |

Revised: A, X, Z ] , under the metric induced by the kernel,d k , at half the rate, in this case n  14 .

Revised: This is similar to Kallus et al.

Revised: [13]’s ﬁndings in the unstabilizedcase.

Revised: Theorem 2.

Original: Under Assumption 4, h ⇤ is the unique solution to the integral equation P A,W,X -almostsurely.
Revised: Under Assumption 4, h ⇤ is the unique solution to the integral equation P A,W,X -almostsurely.

Original: Further, if E
Revised: Further, if E

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: See Appendix A for proofs of Theorems 1 and 2 and Corollary 8.
Revised: See Appendix A for proofs of Theorems 1 and 2.

Original: Taken together, these results tell usthat, as long as our optimization algorithm is successful in estimating ˆ h k,  ,n , it will asymptoticallyapproach the true bridge function, h ⇤ .
Revised: Taken together, these results tell us that, as long asour optimization algorithm is successful in estimating ˆ h k,  ,n , it will asymptotically approach the truebridge function, h ⇤ .

Original: In order for this to occur, the right hand side of the inequalitiesin Theorem 1 must go to zero, which requires not only that the Rademacher terms vanish, but alsothat ˜ h
Revised: In order for this to occur, the right hand side of the inequalities in Theorem 1must go to zero, which requires not only that the Rademacher terms vanish, but also that ˜ h

Original: k must approach h ⇤ arbitrarily closely as n increases.
Revised: k mustapproach h ⇤ arbitrarily closely as n increases.

Original: In practice, this means increasingthecomplexity of the  neural network, but doing so slowly enough the Rademacher complexity terms stilldecrease with sample size.
Revised: In practice, this means increasing the complexity ofthe neural network, but doing so slowly enough the Rademacher complexity terms still decreasewith sample size.

Original: Following Xu et al.
Revised: Following Xu et al.

Original: [7], we note that recent results from Neyshabur et al.[20] suggest that the Rademacher complexity of a ﬁxed network scales like n  12 (similar to manyother popular hypothesis classes) and that, although we cannot compute the scaling of the Rademacherterms directly due to the presence of the kernel function, we expect that they will decline with samplesize and that, as the neural network becomes more complex, their scaling will more closely resembleterms derived from a pure neural network.
Revised: [7], we note that recent results from Neyshabur et al. [24]suggest that the Rademacher complexity of a ﬁxed network scales like n  12 (similar to many otherpopular hypothesis classes) and that, although we cannot compute the scaling of the Rademacherterms directly due to the presence of the kernel function, we expect that they will decline with samplesize and that, as the neural network becomes more complex, their scaling will more closely resembleterms derived from a pure neural network.

Original: Finally, we require that the regularization parameterdecrease as sample size grows, which will, again, depend on the balance between increasing samplesize, which tends to decrease the need for regularization, and increasing complexity, which tendsto increase its importance.
Revised: Finally, we require that the regularization parameterdecrease as sample size grows, which will, again, depend on the balance between increasing samplesize, which tends to decrease the need for regularization, and increasing complexity, which tends toincrease its importance.

Original: Thus, while it is difﬁcult to determine an optimal rate of convergence,by choosing an appropriate growth rate for the network complexity, we expect the aforementionedterms to vanish as n increases to inﬁnity, and, with them, the entire right hand side, making ˆ h k,  ,n aconsistent estimator of h ⇤ .
Revised: Thus, by choosing an appropriate growth rate for the network complexity,we expect the aforementioned terms to vanish as n increases to inﬁnity, and, with them, the entireright hand side, making ˆ h k,  ,n a consistent (likely p n ) estimator of h ⇤ .

Original: Experiments
Revised: Experiments

Original: 6.1 Overview of Baseline Models
Revised: 6.1 Overview of Baseline Models

Original: We compare the performance of NMMR-U and NMMR-V to that of several previous approaches,which we describe brieﬂy here.
Revised: We compare the performance of NMMR-U and NMMR-V to that of several previous approaches,which we describe brieﬂy here.

Original: The baselines can be divided into two categories: structural andnaive.
Revised: The baselines can be divided into two categories: structural and naive.

Original: The structural approachesall leveragesome kind of causal information about the datagenerating process.
Revised: Structural approaches leverage causal information about the data generating process.

Original: The structural methods includeKernel Proxy Variables (KPV) [6], Proximal
Revised: They include

Original:  Maximum Moment Restriction (PMMR) [6], DeepFeature Proxy Variables (DFPV) [21], Causal
Revised: Kernel Proxy Variables (KPV) [6], Proximal Maximum Moment Restriction (PMMR) [6], Deep

Original:  Effect Variational Autoencoder[22] (CEVAE) and the two-stage least squares model (2SLS) from Miao et al.
Revised: Feature Proxy Variables (DFPV) [25], Causal Effect Variational Autoencoder[26] (CEVAE), and thetwo-stage least squares model (2SLS) from Miao et al.

Original: For a review of KPV, PMMR, and DFPV, see Section 3.
Revised: For a review of KPV, PMMR, and DFPV,see Section 3.

Original: CEVAE is an autoencoderapproach derived by Xu et al.
Revised: CEVAE is an autoencoder approach derived by Xu et al.

Original: [7] from Louizos et al.
Revised: [7] from Louizos et al.

Original: 2SLS is a two-stage least squares methodwhich assumes that the bridge function h is linear [5].
Revised: [26].2SLS is a two-stage least squares method which assumes that the bridge function h is linear [5].

Original: The naive approaches serve as baselines and do not use causal information and instead attempt todirectly regress
Revised: The naive approaches serve as baselines and do not use causal information, instead directly regressing

Original: A, Z, W on the outcome Y .
Revised: A and W on the outcome Y .

Original: These methods include a naive neural network (Naivenet), ordinary least squares regression (LS), and ordinary least squares with quadratic features(LS-QF).
Revised: These methods include a naive neural network (Naive net), ordinaryleast squares regression (LS), and ordinary least squares with quadratic features (LS-QF).

Original: Naive Net is a neural network with the same architecture  as NMMR (described furtherin Appendix B) that is trained to predict Y directly from A, Z, W by minimizing observational MSE,1 n
Revised: Naive Net is a neural network that has undergone the same architecture search as NMMR (described furtherin Appendix B) that is trained to predict Y directly from A and W by minimizing observational MSE, 1 n

Original: P n i =1{+ni =
+} ( y  ˆ y ) 2 .
Revised: P {+ni =
+} ( y  ˆ y ) 2 .

Original: Least Squares (LS) is the standard linear regression model that predicts Y usinga linear combination of A, Z, W .
Revised: Least Squares (LS) is the standard linear regression model that predicts Y using a linear combination of A and W .

Original: Least Squares with Quadratic Features (LS-QF) is the same as LS but with additional quadratic terms A 2 , W 2 , Z 2 , AW, AZ, WZ .
Revised: Least Squares with Quadratic Features (LS-QF) is thesame as LS but with additional quadratic terms A 2 , W 2 , AW .

Original: We evaluate NMMR-U, NMMR-V and baseline methods on two synthetic benchmark tasks from Xuet al.
Revised: We evaluate NMMR-U, NMMR-V and baseline methods on two synthetic benchmark tasks from Xuet al.

Original: The ﬁrst is a simulation of how ticket price set by an airline affects the number of ticket soldin the presence of a latent confounder of demand for travel (the Demand experiment).
Revised: The ﬁrst is a simulation of how ticket prices affect the number of tickets sold in the presenceof a latent confounder: demand for travel (the Demand experiment).

Original: The second isan experiment where the goal is to recover a property of an image that is inﬂuenced by an unobservedconfounder (the dSprite experiment).
Revised: The second is an experimentwhere the goal is to recover a property of an image that is inﬂuenced by an unobserved confounder(the dSprite experiment).

Original: The Demand experiment is a low-dimensional estimationproblem, whereas dSprite is high-dimensional as A and W are 64x64=4096-dimensional.
Revised: The Demand experiment is a low-dimensional estimation problem, whereasdSprite is high-dimensional as A and W are 64x64=4096-dimensional.

Original: dSprite alsooffers the opportunity to leverage image-speciﬁc models that are rarely used in the causal inferenceliterature.
Revised: dSprite leverages imagespeciﬁc models, which are rarely used in the causal inference literature.

Original: We note here that both tasks take the measured confounders X to be the null set.
Revised: Neither task uses X .

Original: For the
Revised: For the

Original: Demand experiment we evaluate all the methods mentioned above, whereas for the dSprite experiment,2SLS, LS, and LS-QFwere omitted because of their lack of scalability to high-dimensional settings.
Revised: Demand experiment we evaluate all the methods mentioned above, whereas for the dSprite experiment,we omit 2SLS, LS, and LS-QF because of their lack of scalability to high-dimensional settings.

Original: The experiments are conducted in PyTorch 1.9.0 (Python 3.9.7), using an A100 40GB or TitanX12GB GPU and CUDA version 11.2.
Revised: Experiments were conducted in PyTorch 1.9.0 (Python 3.9.7), using an A100 40GB or TitanX 12GB GPU and CUDA version 11.2.

Original: The experiments can be run in minutes for simpler models(LS, LS-QF, 2SLS) and within several hours for the larger experiments and more complex models(DFPV, NMMR).
Revised: They can be run in minutes for simpler models (LS, LS-QF, 2SLS)and in several hours for the larger experiments and more complex models (DFPV, NMMR).

Original: Implementations of PMMR, KPV, CEVAE, and DFPV and the code to reproduceour experiments can be accessed on GitHub.
Revised: Thecode to reproduce our experiments can be accessed on GitHub.

Original: 6.2 Demand Experiment
Revised: 6.2 Demand Experiment

Original: Hartford et al.
Revised: Hartford et al.

Original: [14] introduced a data generating process for studying instrumental variable regression,and Xu et al.
Revised: [21] introduced a data generating process for studying instrumental variable regression,and Xu et al.

Original: [7] adapted it to the proximal setting.
Revised: [7] adapted it to the proximal setting.

Original: The goal is to estimate the effect of airlineticket price A on sales Y , where these are confounded by demand U (e.g. seasonal ﬂuctuations).
Revised: The goal is to estimate the effect of airlineticket price A on sales Y , which is confounded by demand U (e.g. seasonal ﬂuctuations).

Original: For a treatment-inducing proxy, we have information on the cost of fuel Z = ( Z 1 , Z 2 ) ,  andfor anoutcome-inducing proxy, we have information on the number of views on the ticket reservationwebsite W  (Figure S1).
Revised: Weuse the cost of fuel, Z = ( Z 1 , Z 2 ) , as a treatment-inducing proxy and number of views at a ticketreservation website, W , as an outcome-inducing proxy (Figure S1).

Original: Additional simulation details and the structural equations underlying thecausal DAG can be found in Appendix C.1.
Revised: Additional simulation detailsand the structural equations underlying the causal DAG can be found in Appendix C.1.

Original: Each method was trained on simulated datasets with sample sizes of 1000, 5000, 10,000, and50,000.
Revised: Each method was trained on simulated datasets with sample sizes of 1000, 5000, 10,000, and50,000.

Original: To assess the performance of each method, we evaluated a at 10 equally-spaced intervalsbetween 10 and 30.
Revised: To assess the performance of each method, we evaluated a at 10 equally-spaced intervalsbetween 10 and 30.

Original: We comparedthe difference between each method’s estimated average potentialoutcome ˆ E [ Y a ] against the true average potential outcome E
Revised: We compared each method’s estimated potential outcomes, ˆ E [ Y a ] , againstestimates of the truth, E

Original: [ Y a ] ,which was obtained by a Monte Carlo simulation (10,000 replicates) from the data generating process with A = a ﬁxed.
Revised: [ Y a ] , obtained from Monte Carlo simulations (10,000 replicates) of thedata generating process for each a .

Original: Theevaluation metric is the causal mean squared error (c-MSE) across the 10 evaluation points of a :
Revised: The evaluation metric is the causal mean squared error (cMSE) across the 10 evaluation points of a :

Original: P 10 i =1 ( E [ Y a i ]  ˆ E
Revised: P 10 i =1 ( E [ Y a i ]  ˆ E

Original: [ Y a i ]) 2 .
Revised: [ Y a i ]) 2 .

Original: For MMR-based methods, their predictions are computed on a heldoutdataset, D W with 1,000 draws from W where ˆ E[ Y a i ] =
Revised: For MMR-based methods,predictions are computed using a heldout dataset, D W with 1,000 draws from W so ˆ E

Original: | D W | [Equation]
Revised: [ Y a i ] =| D W |

Original:  ˆ h ( a i , w j ) , i.e. a sampleaverage of the estimated bridge function over W .
Revised: P | D W | j ˆ h ( a i , w j ) , i.e. a sample average of the estimated bridge function over W .

Original: We performed 20 replicates for each method onCausal MSE (c-MSE) of NMMR and baseline methods in the Demand experiment. Each method was replicated 20 times and evaluated on the same 10 test values of E [ Y a ] each replicate. Each individual box plot represents 20 values of c-MSE.
Revised: Weperformed 20 replicates for each method on each sample size, where a single replicate yields onec-MSE value.

Original: See Table S4 for the statistics of each boxplots each sample size, where a single method replicate yields one c-MSE value.

Original: Figure 2 summarizes thec-MSE distribution for each method across the four sample sizes.
Revised: Figure 2 summarizes the c-MSE distribution for each method across the four samplesizes.

Original: NMMR-U has the lowest c-MSEacross all sample sizes, with NMMR-Vas a close second.
Revised: NMMR-U has the lowest c-MSE across all sample sizes, with NMMR-V a close second.

Original: Importantly, Least Squares consistentlyoutperforms all other methods besides NMMR.

Original: It should also be noted that DFPV encountersdifﬁculties with the larger sample sizes of 10,000 and 50,000, potentially due to convergence issueswith its feature maps.
Revised: DFPV encounters difﬁculties with the larger sample sizes of 10,000 and 50,000, potentially due toconvergence issues with its feature maps.

Original: Similarly, PMMR and KPV could not be scaled to n = 50 , 000 .
Revised: Similarly, PMMR and KPV could not scale to n = 50 , 000 .

Revised: Causal MSE (c-MSE) of NMMR and baseline methods in the Demand experiment.

Revised: Each method was replicated 20 times and evaluated on the same 10 test values of E [ Y a ] each replicate.

Revised: Each individual box plot represents 20 values of c-MSE.

Revised: See Table S4 for the statistics of each boxplots

Original: For a more in-depth view of the potential outcome curve that each method estimates, we provide thereplicate-wise potential outcome prediction curves for each of the 4 sample sizes in Figures S3-S6.
Revised: For a more in-depth view of the potential outcome curve estimated by each method, we providereplicate-wise potential outcome prediction curves for each of the 4 sample sizes in Figures S3-S6.

Original: We can see that Least Squares estimates relatively unbiased prediction curves due to the nature ofthe data generating process and further beneﬁts from having very low variance.
Revised: Least Squares estimates relatively unbiased prediction curves due to the nature of the data generatingprocess and has very low variance.

Original: LS-QF managesto match some of the curvature, although its c-MSE distribution (not shown) is not better than LS.
Revised: LS-QF matches some of the curvature, although its c-MSEdistribution (not shown) is not better than LS.

Original: Kernel-based methods, KPV and PMMR, are highly biased.
Revised: Kernel-based methods, KPV and PMMR, are highlybiased.

Original: DFPV is less biased, though it stillsuffers from a lack of ﬂexibility.
Revised: DFPV is less biased, but still suffers from a lack of ﬂexibility.

Original: Both NMMR variants demonstrate the beneﬁt of added ﬂexibilityand additionally has relatively lower variance, which results in a lower c-MSE.
Revised: Both NMMR variantsdemonstrate the beneﬁt of added ﬂexibility and have lower variance, resulting in a lower c-MSE.

Original: Finally, we also varied themagnitude of the variance on the Gaussian noise terms in the structuralequations for Z and W  in order to examine how each method performs with varying quality proxiesfor U (see Appendix E).
Revised: Finally, we also varied the variance of the Gaussian noise terms in the structural equations for Z and W , in order to examine how each method performs with varying quality proxies for U (see Appendix E).

Original: In Figure S9, we can see that NMMR-V is more robust to noised proxiesthan NMMR-U. This could derive from the fact that the U-statistic is an unbiased, but higher variance,estimator than the V-statistic, so when the proxies are less reliable, the estimate of the risk function R k ( h ) iscorrespondingly less stable.
Revised: Figure S10, shows that NMMR-V is more robust to proxy noise than NMMR-U. This could bebecause U-statistics yield unbiased, but higher variance, estimators than V-statistics, so, when proxiesare less reliable, the estimated risk function R k ( h ) is less stable.

Original: We can also see that the kernel-based methods (KPV and
Revised: Kernel-based methods (KPV and

Original: PMMR) rank increasingly well with noisier proxies, which is likely related to the fact that they areless data-adaptive methods.
Revised: PMMR) perform increasingly well with noisier proxies, which is likely related to the fact that theyare less data-adaptive.

Original: Figures S10 through S17 showthe replication-wise prediction curvesacross all 72 noise levels, with one grid plot per method.
Revised: Figures S11 through S18 show replication-wise prediction curves across all 72noise levels, with one grid plot per method.

Original: 6.3 dSprite Experiment
Revised: 6.3 dSprite Experiment

Revised: [Figure]

Original: The second benchmark uses the dSprite dataset from Matthey et al.
Revised: The second benchmark uses the dSprite dataset from Matthey et al.

Original: [23], which was also initiallyadapted for instrumental variable regression in Xu et al.
Revised: [27], which was initially adaptedto instrumental variable regression in Xu et al.

Original: [21], followed by repurposing for proximalinference in Xu et al.
Revised: [25], and repurposed for proximal inference in Xu et al.[7].

Original: This image dataset consists of 2D shapes procedurally generated from6 independent parameters: color, shape, scale, rotation, posX, and posY. All possible combinationsof these parameters are present exactly once, generating 737,280 total images.
Revised: This image dataset consists of 2D shapes procedurally generated from 6 independent parameters:color, shape, scale, rotation, posX, and posY. All possible combinations of these parameters arepresent exactly once, generating 737,280 total images.

Original: In this experiment,we ﬁx shape = heart, color = white, resulting in 245,760 images, each of which contains 64x64=4096pixels.
Revised: In this experiment, we ﬁx shape = heart,color = white, resulting in 245,760 images, each of which contains 64x64=4096 pixels.

Original: The causal DAG for this problem is shown in Figure S7.
Revised: The causal DAG for this problem is shown in Figure S7.

Original: The structural equations and detaileddata generating mechanism underlying the causal DAG can be found in Appendix C.6.
Revised: The structural equations and detailed data generatingmechanism underlying the causal DAG can be found in Appendix C.6.

Original: In the DAG, Fig ( · ) represents the act of retrieving the image from the dSprite dataset with thegiven arguments.
Revised: In the DAG, Fig ( · ) represents the act of retrieving the image from the dSprite dataset with the given arguments.

Original: In this experiment, A and W are vectors representing noised images of a heart , where the heart has a ( scale ) orientation ( rotation ) horizontal position ( posX ) position ( posY ) an image
Revised: A and W are vectors representing noised images of a heart shape, where the heart has a size( scale ), orientation ( rotation ), horizontal position ( posX ), and vertical position ( posY ).

Original:  A and W  see Figure S
Revised: For an exemplar image A and W , see Figure S8.

Revised: The benchmark computes E

Original: [Equation]
Revised: [ Y a ] =

Original: [Figure]
Revised: 110 || vec ( a ) T B || 22  5000where

Original: where B is a 4096 ⇥ 10 matrix of U (0 , 1) weights  Xu et al
Revised: B is a 4096 ⇥ 10 matrix of U (0 , 1) weights from Xu et al.

Original: The observed [Equation] [Figure]
Revised: The observed outcome is computedas Y =

Original: [Equation]
Revised: 110 || vec ( A ) T B

Original: [Figure] [Equation] [Figure]{+|| 22  
+}
Revised: {+|| 22  
+}

Original: [Equation]
Revised: [Equation]

Original: So U dictates the vertical position of the shape in A , as well as the value of Y , making U aconfounder of A, Y .
Revised: + ✏ , ✏ ⇠ N (0 , 0 . 5) .

Original: U is a discrete uniform random variable with E [ (31 ⇥ U 
Revised: U is a discrete uniform randomvariable with E [ (31 ⇥ U

Original: [Figure]
Revised: [Equation]

Original: ) 2 85 .] = 1  .
Revised: [Equation] ] = 1 that dictates the vertical position of the shape in A , as well as thevalue of Y , making U a confounder of A, Y .

Original: Wehypothesized that a convolutional neural network would be exceptionally strong at recovering thisinformation about U from the images A and W .
Revised: We hypothesized that a convolutional neural networkwould be exceptionally strong at recovering this information about U from the images A and W .

Original: Similar to the Demand experiment, we trained each method on simulated datasets with sizes 1,000,5,000, and 7,500, followed by an evaluation on the same test set as Xu et al.
Revised: Similar to the Demand experiment, we trained each method on simulated datasets with sizes1,000, 5,000, and 7,500, followed by an evaluation on the same test set as Xu et al.

Original: This test set contains588 images A that span the range of scale, rotation, posX and posY values (seeAppendix C.9)
Revised: Thistest set contains 588 images A that span the range of scale, rotation, posX and posY values (see

Original: andthe 588 corresponding values of E [ Y a ] .
Revised: Appendix C.9) and the 588 corresponding values of E [ Y a ] .

Original: The evaluation metric is again c-MSE:P
Revised: The evaluation metric is again c-MSE:

Original: [Figure]

Original: i

Original: [Equation]
Revised: [Equation]

Original: 1 ( E [ Y a i ] ˆ E

Original: [-[ Y a i ]) 
-]
Revised: [-[ Y a i ]) 
-]

Original: and we performed 20 replicates for each method on each sample size.
Revised: ; we performed 20 replicates for each method on each sample size.

Original: Figure 3 showsthat NMMR-U or NMMR-V is consistently lowest in c-MSE, with NMMR-V showing substantialimprovement with increasing sample size.
Revised: Figure 3 shows that NMMR-U or NMMR-V is consistently lowest in c-MSE, with NMMR-Vshowing substantial improvement with increasing sample size.

Original: Due to the high dimensionality of the images A and W , wecould not evaluate Least Squares, LS-QF or 2SLS on this experiment.
Revised: Due to the high dimensionality ofthe images A and W , we could not evaluate Least Squares, LS-QF or 2SLS on this experiment.

Original: We can also see that KPV and PMMR do not improve much with increasing sample size.
Revised: KPVand PMMR do not improve much with increasing sample size.

Original: Finally, we also highlight that the Naive net,which uses the same underlying convolutional neural network architecture as NMMR but is trainedusing observational MSE, performs second-to-worst, with a much larger c-MSE than NMMR-U or NMMR-V. This reinforces the need to use causal knowledge in scenarios where it is available.
Revised: The Naive net, which uses the sameunderlying convolutional neural network architecture as NMMR but is trained using observational MSE, performs second-to-worst, with a much larger c-MSE than NMMR-U or NMMR-V. Thisreinforces the need to use causal knowledge in scenarios where it is available.

Original:  Conclusion
Revised:  Conclusion

Original: In this work we have presented a novel method to estimate potential outcomes in the presence ofunmeasured confounding using deep neural networks.
Revised: In this work we have presented a novel method to estimate potential outcomes in the presence ofunmeasured confounding using deep neural networks.

Original: Though our method is promising, there are several limitations.
Revised: Though our method is promising, it has severallimitations.

Original: For very high dimensional data, calculating the kernel matrix K in the loss functioncan be computationally intensive (see Appendix D).
Revised: For very high dimensional data, calculating the kernel matrix K in the loss functioncan be computationally intensive (see Appendix D).

Original: Additionally, mapping real world scenarios to DAGs that satisfy Assumption 1 is non-trivial and technically unveriﬁable (e.g. we cannot be trulysure that W has no impact on A ), though unveriﬁable assumptions are inherent to causal inference.
Revised: Additionally, mapping real world scenariosto DAGs that satisfy Assumption 1 is non-trivial and technically unveriﬁable (e.g. we cannot be trulysure that W has no impact on A ), though unveriﬁable assumptions are inherent to causal inference.

Original: In summary, we provide a new single stage estimator and show how it can be trained on a U-statistic based loss in addition  existing approaches based on V-statstics.
Revised: In summary, we provide a new single stage estimator and show how it can be trained on a U-statisticbased loss in addition to existing approaches based on V-statstics.

Original: We further prove theoreticconvergence properties of our method.
Revised: We further prove theoreticconvergence properties of our method.

Original: On established proximal inference benchmarks, our methodachieves state of the art performance in estimating causal quantities.
Revised: On established proximal inference benchmarks, our methodachieves state of the art performance in estimating causal quantities.

Original: Finally, since our approachis a single-stage neural network, it potentially unlocks new domains for causal inference where deeplearning has had success, such as imaging.
Revised: Finally, since our approachis a single-stage neural network, it potentially unlocks new domains for causal inference where deeplearning has had success, such as imaging.

Original: Checklist
Revised: Checklist

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope?
Revised: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope?

Original: Yes ] The claims made in the abstract and introduction arereﬂected in Sections 4, 5, and [-
 (b)-] Did you describe the limitations of your work?
Revised: [Yes] The claims made in the abstract and introduction arereﬂected in Sections 4, 5, and [-
 (b)-]6(b) Did you describe the limitations of your work?

Original: [Yes] See Section 
Revised: [Yes] See Section 7(c)

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: Did you discuss any potential negative societal impacts of your work?

Original: [Yes] See
Revised: [Yes] See

Original: Section 1 where we discuss potential medical impacts.
Revised: Section 1 where we discuss potential medical impacts.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes] See Section 5and Appendix A
Revised: [Yes] See Section 5and Appendix A

Original: (b) Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

Original: [Yes] See Appendix A3.
Revised: [Yes] See Appendix A3.

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] Available herehttps://github.com/neurips-2022-11442/anonymous-NMMR
Revised: [Yes] Available here https://github.com/neurips-2022-11442/anonymous-NMMR

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?

Original: [Yes] See Appendix B and C as well as our code
Revised: [Yes] See Appendix B and C as well as our code

Original: (c) Did you report error bars (e.g., with respect to the random seed after running exper- iments multiple times)?
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes] We reported results across 20 random seeds in all
Revised: [Yes] We reported results across 20 random seeds in all

Original: our
Revised: our

Original: Figures and Tables.
Revised: Figures and Tables.

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes] See details in Section64.
Revised: [Yes] See details in Section

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] We described the
Revised: [Yes] We described the

Original: fork of assets from Xu et al.
Revised: fork of assets from Xu et al.

Original: [7] in Section6(b)
Revised: [7] in Section

Original:  Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

Original: [Yes] See 6, MIT license(c) Did you include any new assets either in the supplemental material or as a URL?
Revised: [Yes] See 6, MIT license(c) Did you include any new assets either in the supplemental material or as a URL?

Original: [Yes]Available here https://github.com/neurips-2022-11442/anonymous-NMMR
Revised: [Yes]Available here https://github.com/neurips-2022-11442/anonymous-NMMR

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Revised: [Table]

Original: using/curating?
Revised: using/curating?

Original: [N/A]
Revised: [N/A]

Original: Synthetic data
Revised: Synthetic data

Original:  (e) Did you discuss whether the data you are using/curating contains personally identiﬁable 
Revised: [Table] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable [Table]

Original: information or offensive content?
Revised: information or offensive content?

Original: [N/A]
Revised: [N/A]

Original: Synthetic data
Revised: Synthetic data

Revised: [Table]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
Revised: [Table] ) Did you include the full text of instructions given to participants and screenshots, if [Table] applicable?

Original: [N/A]
Revised: [N/A]

Original:  Did you describe any potential participant risks, with links to Institutional Review 
Revised: [Table] (b) Did you describe any potential participant risks, with links to Institutional Review [Table]

Original: Board 
Revised: Board (

Original:  approvals, if applicable
Revised: ) approvals, if applicable?

Original: [N
Revised: [N/A

Original:  you include the estimated hourly wage paid to participants and the total amount spent on participant compensation
Revised: [Table] (c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?

Original: N/A
Revised: [N/A]

Revised: [Table]