# Learning Invariant Representations
with Missing Data

Mark Goldstein Jörn-Henrik Jacobsen1 Olina Chau1 Adriel Saporta1 &Aahlad Puli2
Rajesh Ranganath2 Andrew C. Miller1

1Apple 2New York University Work done while at Apple. Correspondence to
<EMAIL_ADDRESS>and<EMAIL_ADDRESS>
###### Abstract

Spurious correlations allow flexible models to predict well during training
but poorly on related test populations. Recent work has shown that models that
satisfy particular independencies involving correlation-inducing nuisance
variables have guarantees on their test performance. Enforcing such
independencies requires nuisances to be observed during training. However,
nuisances, such as demographics or image background labels, are often missing.
Enforcing independence on just the observed data does not imply independence
on the entire population. Here we derive mmd estimators used for invariance
objectives under missing nuisances. On simulations and clinical data,
optimizing through these estimates achieves test performance similar to using
estimators that make use of the full data.

##  1 Introduction

Spurious correlations allow models that predict well on training data to have
worse than chance performance on related populations at test time [9, 27, 35,
22, 13, 32]. For example, diabetes is associated with high body mass index
(bmi) in the United States. However, in India and Taiwan, diabetes also
frequently co-occurs with low and average bmi [37]. Due to their shifting
relationship with the label, nuisance variables (e.g., bmi) can cause models
to exploit correlations in training data, leading them to generalize poorly on
test sets of interest.

Invariant prediction methods are designed to improve performance on a range of
test distributions when training data exhibits spurious correlations [26, 3].
We focus on methods that enforce independencies between the model and nuisance
given some assumed causal structure [22, 35, 27]. These methods require the
nuisance to be specified explicitly and observed. However, in large health
datasets, nuisances are often missing. For example, not all people who report
diabetes status report other correlated conditions (e.g., hypertension,
depression) or demographics (e.g., gender).

To improve generalization on a range of test distributions, it is necessary to
handle missingness appropriately. However, extending invariant prediction
methods to handle missing data is not straightforward. This difficulty stems
from the invariant method’s optimization objective, which usually includes a
measure of dependence — e.g., Maximum Mean Discrepancy (mmd) or Mutual
Information (mi)— that requires a sample from the fully-observed data to
estimate consistently.

We propose mmd estimators for measuring nuisance-model dependence under
missingness. First, we show that enforcing independence on only the nuisance-
observed data does not imply independence on the full data, and vice versa.
Next, we derive three estimators, including one that is doubly-robust: it is
consistent when either the nuisance or missingness can be consistently modeled
[4]. Using simulations, a semi-simulation using textured mnist and clinical
data from mimic, we show that the estimators perform close to ground-truth
estimation with no missingness and that they improve test accuracy relative to
mmd computation using only the data with nuisances observed.

##  2 Notation and background

#### Notation.

Let $X$ denote features. Let $Y$ be a label such as disease status. Let $Z$
denote the nuisance, e.g., another disease correlated with $Y$, demographics,
or image backgrounds. Denote the nuisance missingness indicator as $\Delta$.
Instead of $(X,Y,Z)$, we observe $(X,Y,\Delta,\tilde{Z}=\Delta Z)$, where
$\tilde{Z}=Z$ when $\Delta=1$ and $Z$ is unobserved otherwise. We write
functions as $f_{X}=f(X)$ to avoid excess parentheses. Let $h_{X}=h(X)$ denote
a model to predict $Y$. Let $Y\sim\mathcal{B}(p)$ denote
$Y\sim\text{Bernoulli}(p)$.

#### Assumptions and scope.

The estimators require ignorability,
$Z\rotatebox[origin={c}]{90.0}{$\models$}\Delta|X,Y$ [15]. Some distributions
that satisfy this are shown in Figure 1. We require positivity $0<\epsilon\leq
P(\Delta=1|X,Y)$ to observe $Z$ appropriately. While we focus on the graph in
Figure 1(a) with binary $Z$, the presented method can extend to continuous $Z$
and to other graphs (e.g., Figure 1(b)).

(a) Anti-Causal

(b) Causal

Figure 1: Generative processes we consider in this work. The $Y\rightarrow Z$
edge is dashed to emphasize that the $Z|Y$ may change at test time. $\Delta$
determines missingness of $Z$ and in general depends on both $X$ and $Y$.

#### Modeling under spurious correlations.

Nuisance-based prediction arises in training data when $Z$ is predictive of
$Y$ and associated with $X$, causing models to use information about $Z$ in
$X$ to predict $Y$. This is a problem when the test distribution is expected
to have a different $(Y,Z)$ relationship from the training distribution. For
example, define a family of distributions, indexed by $D$, that varies along
$P(Z~{}|~{}Y)$:

$\displaystyle\mathcal{F}=\\{P_{D}(X,Y,Z)=P(Y)P_{D}(Z|Y)P(X|Y,Z)\\}.$

When $P_{test}(Z|Y)\neq P_{train}(Z|Y)$, in general $P_{train}(Y~{}|~{}X)\neq
P_{test}(Y|X)$ and a model built on $P_{train}$ can generalize poorly
(Appendix C). When it is possible to anticipate and observe nuisances during
training, enforcing certain independence constraints [22, 35, 27] helps
guarantee performance regardless of the nuisance-label relationship. For
example, for this choice of $\mathcal{F}$, maximum likelihood estimation for
$Y|h_{X}$ while enforcing the constraint
$h_{X}\rotatebox[origin={c}]{90.0}{$\models$}Z~{}|~{}Y=y$ for $y\in\\{0,1\\}$
implies equal performance on all members, and better than chance on all
members (Appendix D).

#### Missingness.

When $Z$ is subject to missingness, two parts of the data distribution help
estimate functionals like $\mathop{\mathbb{E}}[Z]$: the missingness process
$G_{X}\triangleq\mathbb{E}[\Delta~{}|~{}X,Y]$ and conditional expectation
$m_{X}\triangleq\mathbb{E}[Z~{}|~{}X,Y]$. We review estimators that use either
$G_{X}$ [16, 5, 29] or $m_{X}$ [30, 33] in Appendix H. The doubly-robust (dr)
estimator [28, 4, 18] combines both by noting the equality:

$\displaystyle\mathop{\mathbb{E}}[Z]=\mathop{\mathbb{E}}\Big{[}\frac{\Delta\tilde{Z}}{G_{X}}-\frac{\Delta-
G_{X}}{G_{X}}m_{X}\Big{]}\,.$ (1)

Replacing $G_{X}$ or $m_{X}$ with models, Monte Carlo estimates of the right
side of Equation 1 are consistent for $\mathop{\mathbb{E}}[Z]$ when, for all
$X$, either $G_{X}$ or $m_{X}$ are consistently estimated (Appendix I).

##  3 Invariant representations with missing data

In the anti-causal setting, [35] enforce
$h_{X}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y$ for $y\in\\{0,1\\}$ by
minimizing the mmd:

$\displaystyle\max_{h}\,\quad\log
p(y|h_{X})-\lambda\cdot\sum_{y\in\\{0,1\\}}\lx@glossaries@gls@link{acronym}{mmd}{\leavevmode\textsc{mmd}}\left(p(h_{X}|Z=1,Y=y)\,,\,p(h_{X}|Z=0,Y=y)\right)\,.$
(2)

First, we demonstrate what can go wrong when enforcing this mmd penalty only
on samples where $Z$ is observed. We then derive estimators of the full-data
mmd under missingness.

###  3.1 Failures of restricting to observed data

Restricting computation to data with non-missing $Z$ enforces
$h_{X}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y,\Delta=1$ instead of
$h_{X}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y$. We show that these
conditions do not imply each other.

######  Proposition 1.

There exist distributions on $(X,Y,\Delta,Z)$ such that

$\displaystyle\exists h_{X}^{\star}\quad\text{s.t.}\quad
h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y,\quad\text{but}\quad
h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y=y,\Delta=1$

and there exist distributions on $(X,Y,\Delta,Z)$ such that

$\displaystyle\exists h_{X}^{\star}\quad\text{s.t.}\quad
h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y,\Delta=1\quad\text{but}\quad
h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y=y$

The proof is in Appendix G. This implies that (1) optimizing the observed-only
mmd can discard a solution to the full-data mmd, and (2) using the observed-
only data may lead one to believe a model is invariant when it is not. This
means one must enforce independence on the full data.

###  3.2 Mmd estimation under missingness

We present estimators of the full-data unconditional mmd. One can compute the
mmd conditional on $Y=y$ simply by restricting samples to those with $Y=y$.
For a kernel $k$ and $k_{XX^{\prime}}\triangleq k(h_{X},h_{X^{\prime}})$,

$\displaystyle\lx@glossaries@gls@link{acronym}{mmd}{\leavevmode\textsc{mmd}}\big{(}p(h_{X}|Z=1),p(h_{X}|Z=0)\big{)}=\mathop{\mathbb{E}}_{\begin{subarray}{c}X|Z=1\\\
X^{\prime}|Z^{\prime}=1\end{subarray}}k_{XX^{\prime}}+\mathop{\mathbb{E}}_{\begin{subarray}{c}X|Z=0\\\
X^{\prime}|Z^{\prime}=0\end{subarray}}k_{XX^{\prime}}-2\mathop{\mathbb{E}}_{\begin{subarray}{c}X|Z=1\\\
X^{\prime}|Z^{\prime}=0\end{subarray}}k_{XX^{\prime}}\,.$ (3)

Estimation is challenging due to missingness in the conditioning set. For
$b\in\\{0,1\\}$, let $N(b,b^{\prime})=P(Z=b)P(Z^{\prime}=b^{\prime})$ and let
$Z_{1}\triangleq Z$ and $Z_{0}\triangleq 1-Z$. The dependence on $Z$ can be
re-written:

$\displaystyle\mathop{\mathbb{E}}_{\begin{subarray}{c}X|Z=b\\\
X^{\prime}|Z^{\prime}=b^{\prime}\end{subarray}}k_{XX^{\prime}}=\frac{1}{N(b,b^{\prime})}\mathop{\mathbb{E}}\Big{[}k_{XX^{\prime}}\cdot
Z_{b}\cdot Z_{b^{\prime}}\Big{]}\,.$ (4)

Under no missingness, each expectation could be estimated with Monte Carlo. We
now develop three mmd estimators. We derive simpler $G_{X}$-based and
$m_{X}$-based estimators in Equations 10 and 11 (Appendix J). Here we combine
them. Let $m_{X1}\triangleq m_{X},\phantom{.}m_{X0}\triangleq 1-m_{X}$, and
$G_{XX^{\prime}}\triangleq G_{X}G_{X^{\prime}}$.

######  Proposition 2.

(dr estimator). Assume positivity, ignorability, and $\forall X$,
$G_{X}=\mathop{\mathbb{E}}[\Delta|X,Y]$ or $m_{X}=\mathop{\mathbb{E}}[Z|X,Y]$.
Then,

$\displaystyle\mathop{\mathbb{E}}_{\begin{subarray}{c}X|Z=b\\\
X^{\prime}|Z^{\prime}=b^{\prime}\end{subarray}}\Big{[}k_{XX^{\prime}}\Big{]}$
$\displaystyle=\frac{1}{N(b,b^{\prime})}\mathop{\mathbb{E}}\Big{[}\Big{(}\frac{\Delta\Delta^{\prime}\tilde{Z}_{b}\tilde{Z^{\prime}}_{b^{\prime}}}{G_{XX^{\prime}}}-\frac{\Delta\Delta^{\prime}-G_{XX^{\prime}}}{G_{XX^{\prime}}}\cdot
m_{Xb}\cdot m_{X^{\prime}b^{\prime}}\Big{)}k_{XX^{\prime}}\Big{]}\,.$ (5)

The proof is in Appendix J. We can use any of Equations 10, 11 and 5 to
estimate the terms in eq. 3. Each of Equations 10, 11 and 5 is a ratio of two
expectations: the normalization constant $N(b,b^{\prime})$ depends on
$\mathop{\mathbb{E}}[Z]$ and must itself be estimated under missingness (e.g.,
with Equation 1). The ratio of consistent estimates of these quantities is
consistent by Weak Law of Large Numbers and Slutsky’s theorem. We discuss
estimation in practice, trade-offs among the three estimators, and variance in
Appendix E. We review recent related work in Appendix A.

##  4 Experiments

We compare accuracy and mmd minimization using different estimators: none (mle
only, no mmd), full (mle and mmd using data with $Z$ fully-observed), obs (mle
and observed-only mmd), dr (mle and dr estimator, called dr + when using true
$G_{X}$), ip (mle and re-weighted estimator, called ip + when using true
$G_{X}$), and reg (mle and regression estimator).

We first compare these algorithms in a simulation study. We then use textured
mnist to show the utility of the proposed estimators on high-dimensional data.
In quantitative tables, we show mean $\pm$ standard deviation over three
seeds. We then predict hospital length of stay in the mimic dataset, and
compare performance when demographic nuisances are subject to missingness. For
the $Y|X$ predictive loss, we use negative Bernoulli log likelihood with logit
equal to $h_{X}$.

#### Comparing mmds.

In all tables, the training set mmd for each method is computed using the
ground-truth full-data mmd estimation method (see eq. 4). This is also what
the full model optimizes. True $Z$’s are available in both simulated and real
data as missingness is simulated. However, each model trains and validates the
$\log p(Y|X)+\lx@glossaries@gls@link{acronym}{mmd}{\leavevmode\textsc{mmd}}$
loss using its own estimation method.

###  4.1 Experiment 1: Simulation.

We set up strong $(Y,Z)$ correlation. With $\overline{Y}=1-Y$, the training
and validation sets are drawn:

$\displaystyle\begin{split}Y\sim\mathcal{B}(0.5),\quad
Z\sim\mathcal{B}(.9Y+.1\overline{Y}),\quad
X\sim[\mathcal{N}(Y-Z,\sigma_{X}^{2}),\mathcal{N}(Y+Z,\sigma_{X}^{2})]\,.\end{split}$
(6)

The test set has the opposite relationship
$Z\sim\mathcal{B}(.1Y+.9\overline{Y})$. Here $h_{X}^{\star}=(X_{1}+X_{2})/2$
predicts $Y$ with smallest MSE among representations satisfying independence.
We construct $\Delta$ to show the failure of computing mmd on the observed-
only subset. For this, we use $\hat{Z}\triangleq-(X_{1}-X_{2})/2$, which is
correlated with $Z$. We draw $\Delta$ conditional on $h_{X}^{\star}$ and
$\hat{Z}$ (both are functions of $X$):

$\displaystyle
Q=\mathbbm{1}\left[h_{X}^{\star}>0.6\right]\cdot\mathbbm{1}\left[\hat{Z}<0.6\right],\quad\Delta\sim\mathcal{B}(Q+0.2\overline{Q})\,.$

This example construction leads to
$h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y$ but
$h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y,\Delta=1$. For
$h,G_{X}$ and $m_{X}$ we use small feed-forward neural networks (Section L.1).

#### Results.

In Table 1, the dr estimators achieve indistinguishable performance to the
full-data mmd, both in mmd and accuracy, and better than none and obs. We
include more results in Appendix F.

Table 1:  Simulation. $\lambda=1$. None has highest mmd and lowest test accuracy. Obs improves over this. The dr and reg methods are able to bring the mmd close to $0.0$ and attain best test accuracy. | none | obs | full | dr |  dr + | reg
---|---|---|---|---|---|---
tr mmd | $0.21\pm 0.04$ | $0.05\pm 0.04$ | $0.00\pm 0.01$ | $0.00\pm 0.01$ | $0.00\pm 0.01$ | $0.01\pm 0.00$
tr acc | $0.89\pm 0.00$ | $0.87\pm 0.00$ | $0.86\pm 0.01$ | $0.85\pm 0.01$ | $0.84\pm 0.02$ | $0.86\pm 0.00$
te acc | $0.67\pm 0.02$ | $0.77\pm 0.02$ | $0.80\pm 0.01$ | $0.81\pm 0.02$ | $0.81\pm 0.01$ | $0.79\pm 0.02$

###  4.2 Experiment 2: Textured mnist.

Following [11]111 We adapt this repository (linked) to construct textured
mnist and will make our code available., we correlate mnist digits $0$ and $1$
with two textures from the Brodatz dataset (Figure 2). The missingness is
based on the average pixel intensity of $X$ and its class.

Figure 2: Textured mnist with digits 0,1 on two textures from the Brodatz
dataset.

For $h,G_{X}$ and $m_{X}$ we use small convolutional networks. We include more
details in Section L.2.

#### Results.

In Table 2, none and obs perform poorly on test. In contrast, the dr
estimators — including the one with a learned $G_{X},m_{X}$ — achieve close to
full’s performance.

Table 2:  mnist $\lambda=1$. dr and reg estimators achieve close to full performance as measured by full mmd $=0$ and high test accuracy. None and obs perform poorly on test. Obs is notably high variance. | none | obs | full | dr |  dr + | reg
---|---|---|---|---|---|---
tr mmd | $2.05\pm 0.18$ | $0.02\pm 0.04$ | $0.00\pm 0.01$ | $0.00\pm 0.01$ | $0.00\pm 0.01$ | $0.00\pm 0.01$
tr acc | $0.90\pm 0.01$ | $0.74\pm 0.03$ | $0.76\pm 0.01$ | $0.77\pm 0.0$ | $0.76\pm 0.01$ | $0.76\pm 0.01$
te acc | $0.13\pm 0.01$ | $0.63\pm 0.17$ | $0.74\pm 0.01$ | $0.72\pm 0.04$ | $0.73\pm 0.01$ | $0.73\pm 0.01$

###  4.3 Experiment 3: Predicting length of stay in the icu

We predict length of stay in the intensive care unit (icu) in mimic [17]222The
mimic critical-care database is available on Physionet [10]. using
demographics and first day labs/vitals among patients that stay at least one
day. The prediction task is whether the stay is more than $2.5$ days. To
demonstrate that spurious correlations cause issues at deployment, we choose
$Z=1$ to indicate the patient is recorded as white. While race may be
correlated with health outcomes (e.g., due to unobserved socioeconomic factors
[24]), it may not always be appropriate for a model to use this information
[6]. The test set represents a new population with different outcome-
demographic structure: we split the data so that the training/validation set
has mostly samples with $Y\neq Z$ while the test set has mostly samples with
$Y=Z$. We set non-male patients to have $Z$ observed with probability $0.2$.
We include more details in Appendix M.

Table 3:  mimic $\lambda=1$. reg estimator matches full’s performance and improves upon obs while dr does not, due to high objective variance during training (not shown in table). | none | obs | full | dr | reg
---|---|---|---|---|---
tr mmd | $0.017\pm 0.02$ | $0.002\pm 0.01$ | $0.00\pm 0.00$ | $0.009\pm 0.01$ | $0.00\pm 0.00$
tr acc | $0.71\pm 0.02$ | $0.68\pm 0.01$ | $0.70\pm 0.01$ | $0.70\pm 0.01$ | $0.71\pm 0.00$
te acc | $0.64\pm 0.00$ | $0.64\pm 0.00$ | $0.66\pm 0.00$ | $0.62\pm 0.00$ | $0.66\pm 0.01$

#### Results.

In this real data setting with strong $(Y,Z)$ correlation, the full-data mmd
estimator reported in the table for all methods may have high variance. We
focus on the attained accuracies. The reg estimator matches the ground-truth
full estimator and performs better than obs and dr. This is not unexpected,
since it is possible for the reg estimator to be lower variance than dr when
the true $G_{X}$ is small or $G_{X}$ is not modeled well [8], especially under
strong $(Y,Z)$ correlation (Appendix E).

##  5 Conclusion

We present estimators for the mmd that extend recent invariant prediction
methods to missing data. Unlike prior estimators that only leverage data with
nuisances observed, or consider worst-case estimation (see related work in
Appendix A), the presented estimators of the full data objective are
consistent when either auxiliary model can be learned. As we show in
proposition 1, estimation of the full data objective is necessary to preserve
the theoretical properties of invariant prediction methods. In the
experiments, the dr and reg estimators are able to match full-data mmd
performance and improve test accuracy relative to the obs estimator. In
practice, we recommend exploring the two simpler proposed estimators (reg and
ip) in addition to the dr estimator and selecting the model based on the
validation metric.

Moving forward, one limitation is that the full-data estimator — used as
ground-truth mmd evaluation for the experiments — may itself have high
variance on small datasets with strong nuisance-label correlation. Variance
reduction is an important avenue both for optimizing and evaluating with the
mmd using smaller batch sizes (in our experiments, batch sizes $1500$ for
mnist and $4000$ for mimic are large). Beyond variance reduction, it is a
promising direction to apply the methodology in this work to the mutual
information objective in [27], which sidesteps the choice of kernel and may be
better suited for continuous and high dimensional nuisances.

## Acknowledgments and Disclosure of Funding

The authors thank Scotty Fleming, Joe Futoma, Leon Gatys, Sean Jewell, Tayor
Killian, and Guillermo Sapiro for feedback and discussions. This work was in
part supported by NIH/NHLBI Award R01HL148248, NSF Award 1922658 NRT-HDR:
FUTURE Foundations, Translation, and Responsibility for Data Science, and NSF
Award 1815633 SHF.

## References

  * Ahmed et al. [2020] F. Ahmed, Y. Bengio, H. van Seijen, and A. Courville.  Systematic generalisation with group invariant predictions.  In _International Conference on Learning Representations_ , 2020.
  * Anonymous [2022] Anonymous.  Invariant learning with partial group labels.  In _Submitted to The Tenth International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=sWbXSWzHPa.  under review.
  * Arjovsky et al. [2019] M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz.  Invariant risk minimization.  _arXiv preprint arXiv:1907.02893_ , 2019.
  * Bang and Robins [2005] H. Bang and J. M. Robins.  Doubly robust estimation in missing data and causal inference models.  _Biometrics_ , 61(4):962–973, 2005.
  * Binder [1983] D. A. Binder.  On the variances of asymptotically normal estimators from complex surveys.  _International Statistical Review/Revue Internationale de Statistique_ , pages 279–292, 1983.
  * Chen et al. [2018] I. Chen, F. D. Johansson, and D. Sontag.  Why is my classifier discriminatory?  _arXiv preprint arXiv:1805.12002_ , 2018.
  * Creager et al. [2021] E. Creager, J.-H. Jacobsen, and R. Zemel.  Environment inference for invariant learning.  In _International Conference on Machine Learning_ , pages 2189–2200. PMLR, 2021.
  * Davidian [2005] M. Davidian.  Double robustness in estimation of causal treatment effects, 2005.
  * Geirhos et al. [2020] R. Geirhos, J.-H. Jacobsen, C. Michaelis, R. Zemel, W. Brendel, M. Bethge, and F. A. Wichmann.  Shortcut learning in deep neural networks.  _Nature Machine Intelligence_ , 2(11):665–673, 2020.
  * Goldberger et al. [2000] A. L. Goldberger, L. A. Amaral, L. Glass, J. M. Hausdorff, P. C. Ivanov, R. G. Mark, J. E. Mietus, G. B. Moody, C.-K. Peng, and H. E. Stanley.  Physiobank, physiotoolkit, and physionet: components of a new research resource for complex physiologic signals.  _circulation_ , 101(23):e215–e220, 2000.
  * Goodfellow et al. [2013] I. J. Goodfellow, D. Warde-Farley, P. Lamblin, V. Dumoulin, M. Mirza, R. Pascanu, J. Bergstra, F. Bastien, and Y. Bengio.  Pylearn2: a machine learning research library.  _arXiv preprint arXiv:1308.4214_ , 2013.  URL http://arxiv.org/abs/1308.4214.
  * Gretton et al. [2012] A. Gretton, K. M. Borgwardt, M. J. Rasch, B. Schölkopf, and A. Smola.  A kernel two-sample test.  _The Journal of Machine Learning Research_ , 13(1):723–773, 2012.
  * Gulrajani and Lopez-Paz [2020] I. Gulrajani and D. Lopez-Paz.  In search of lost domain generalization.  _arXiv preprint arXiv:2007.01434_ , 2020.
  * Hébert-Johnson et al. [2018] U. Hébert-Johnson, M. Kim, O. Reingold, and G. Rothblum.  Multicalibration: Calibration for the (computationally-identifiable) masses.  In _International Conference on Machine Learning_ , pages 1939–1948. PMLR, 2018.
  * Hernan and Robins [2021] M. A. Hernan and J. M. Robins.  _Causal Inference: What If_.  Chapman & Hall/CRC, 1st edition, 2021.
  * Horvitz and Thompson [1952] D. G. Horvitz and D. J. Thompson.  A generalization of sampling without replacement from a finite universe.  _Journal of the American statistical Association_ , 47(260):663–685, 1952.
  * Johnson et al. [2016] A. E. Johnson, T. J. Pollard, L. Shen, H. L. Li-wei, M. Feng, M. Ghassemi, B. Moody, P. Szolovits, L. A. Celi, and R. G. Mark.  Mimic-iii, a freely accessible critical care database.  _Scientific data_ , 3:160035, 2016.
  * Kang and Schafer [2007] J. D. Kang and J. L. Schafer.  Demystifying double robustness: A comparison of alternative strategies for estimating a population mean from incomplete data.  _Statistical science_ , pages 523–539, 2007.
  * Lahoti et al. [2020] P. Lahoti, A. Beutel, J. Chen, K. Lee, F. Prost, N. Thain, X. Wang, and E. H. Chi.  Fairness without demographics through adversarially reweighted learning.  _arXiv preprint arXiv:2006.13114_ , 2020.
  * Little and Rubin [2019] R. J. Little and D. B. Rubin.  _Statistical analysis with missing data_ , volume 793.  John Wiley & Sons, 2019.
  * Liu et al. [2021] E. Z. Liu, B. Haghgoo, A. S. Chen, A. Raghunathan, P. W. Koh, S. Sagawa, P. Liang, and C. Finn.  Just train twice: Improving group robustness without training group information.  In _International Conference on Machine Learning_ , pages 6781–6792. PMLR, 2021.
  * Makar et al. [2021] M. Makar, B. Packer, D. Moldovan, D. Blalock, Y. Halpern, and A. D’Amour.  Causally-motivated shortcut removal using auxiliary labels.  _arXiv preprint arXiv:2105.06422_ , 2021.
  * Namkoong and Duchi [2016] H. Namkoong and J. C. Duchi.  Stochastic gradient methods for distributionally robust optimization with f-divergences.  In _NeurIPS_ , volume 29, pages 2208–2216, 2016.
  * Obermeyer et al. [2019] Z. Obermeyer, B. Powers, C. Vogeli, and S. Mullainathan.  Dissecting racial bias in an algorithm used to manage the health of populations.  _Science_ , 366(6464):447–453, 2019.
  * Pearl [2009] J. Pearl.  _Causality_.  Cambridge university press, 2009.
  * Peters et al. [2016] J. Peters, P. Bühlmann, and N. Meinshausen.  Causal inference by using invariant prediction: identification and confidence intervals.  _Journal of the Royal Statistical Society. Series B (Statistical Methodology)_ , pages 947–1012, 2016.
  * Puli et al. [2021] A. Puli, L. H. Zhang, E. K. Oermann, and R. Ranganath.  Predictive modeling in the presence of nuisance-induced spurious correlations.  _arXiv preprint arXiv:2107.00520_ , 2021.
  * Robins and Rotnitzky [2001] J. Robins and A. G. Rotnitzky.  Comment on the bickel and kwon article, ’inference for semiparametric models: Some questions and an answer’.  _Statistica Sinica_ , 11:920–936, 01 2001.
  * Robins et al. [1994] J. M. Robins, A. Rotnitzky, and L. P. Zhao.  Estimation of regression coefficients when some regressors are not always observed.  _Journal of the American statistical Association_ , 89(427):846–866, 1994.
  * Rubin [1976] D. B. Rubin.  Inference and missing data.  _Biometrika_ , 63(3):581–592, 1976.
  * Rubin [2004] D. B. Rubin.  _Multiple imputation for nonresponse in surveys_ , volume 81.  John Wiley & Sons, 2004.
  * Sagawa et al. [2020] S. Sagawa, A. Raghunathan, P. W. Koh, and P. Liang.  An investigation of why overparameterization exacerbates spurious correlations.  In _International Conference on Machine Learning_ , pages 8346–8356. PMLR, 2020.
  * Schafer [1997] J. L. Schafer.  _Analysis of incomplete multivariate data_.  CRC press, 1997.
  * Van der Laan and Robins [2003] M. J. Van der Laan and J. M. Robins.  _Unified methods for censored longitudinal data and causality_.  Springer Science & Business Media, 2003.
  * Veitch et al. [2021] V. Veitch, A. D’Amour, S. Yadlowsky, and J. Eisenstein.  Counterfactual invariance to spurious correlations in text classification.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Wang et al. [2020] S. Wang, W. Guo, H. Narasimhan, A. Cotter, M. Gupta, and M. I. Jordan.  Robust optimization for fairness with noisy protected groups.  _arXiv preprint arXiv:2002.09343_ , 2020.
  * WHO [2004] E. C. WHO.  Appropriate body-mass index for Asian populations and its implications for policy and intervention strategies.  _Lancet (London, England)_ , 363(9403):157–163, 2004.

##  Appendix A Related work

We focus on recent work in fairness and invariant prediction on missing
group/environment labels. Motivated by fairness, [36] study a related problem
of optimizing invariance-inducing objectives when the protected group label
(analogous to our nuisance variables) is noisy. Given bounds on the level of
label noise, this work proposes optimizing an objective based on the
distributionally robust optimization framework [23]. Additionally, if given a
small amount of true labels the authors suggest fitting a model to de-noise
the noisy group labels and re-weight examples in the objective, which is
similar in spirit to our work. In our approach, however, we exploit structural
assumptions about the missingness process to build a doubly-robust estimator
of the mmd penalty used during optimization.

[19] optimize worst-case-over-groups performance without known group labels.
They rely on the assumption that groups are computationally-identifiable (i.e.
that there exists some function on the data that labels their protected group
membership) [14] and use a model to identify groups on which performance is
worst. They pose an adversarial optimization between the group-labeling model
— which searches for groups with poor performance — and the primary predictive
model. Inspired by this work, [7] find worst-case group assignments based on
an empirical risk minimization (erm) model that maximizes invariance penalties
and [1] illustrate that this objective performs well on a wide range of
benchmarks. Relatedly, [21] run usual erm training and then a second iteration
of erm that upweights the loss for datapoints on which the model performs
badly. This identifies groups with bad model performance without explicit
group labels. In both of these works, the groups could be seen either as a
nuisance variable or as a confounder that correlates the label and some
nuisance variable. However, in our setting (and in that of [22, 35, 27]), in
exchange for being willing to make assumptions on the test distribution
family, we do not need to observe samples with poor model performance at
training time (and may not see any) to prevent sudden decreases in performance
on held-out data at test time.

[2] too consider invariant learning with partial/missing group labels. They
maximize over unknown group assignments to get worst-case (over group
assignments) versions of invariance objectives. As opposed to the dr estimator
presented here, in their work the objective still depends on the missingness
distribution: the objective value can take values closer or further from $0$
for the same model but different missingness distributions.

##  Appendix B Measuring dependence with mmds

To enforce independencies it is necessary to measure dependence, i.e., to
measure the distance between a joint distribution and the product of its
marginals. The general Integral Probability Metrics (ipms) class defines a
metric on distributions. A special case, the kernel-based mmd [12], has a
closed form. Let $X_{1}\sim P,X_{2}\sim Q$. Let $X_{j}^{\prime}$ be an
independent sample identically distributed as $X_{j}$. For kernel $k$,

$\displaystyle\lx@glossaries@gls@link{acronym}{mmd}{\leavevmode\textsc{mmd}}(X_{1},X_{2})=\mathop{\mathbb{E}}[k(X_{1},X_{1}^{\prime})]+\mathop{\mathbb{E}}[k(X_{2},X_{2}^{\prime})]-2\mathop{\mathbb{E}}[k(X_{1},X_{2})]\,.$
(7)

mmds can measure dependence between $X,Z$ by computing
$\lx@glossaries@gls@link{acronym}{mmd}{\leavevmode\textsc{mmd}}(p(X,Z),p(X)p(Z))$.

##  Appendix C Out-of-Distribution example

Example. Let $\overline{Y}=1-Y$ and define $P_{train}$ by:

$\displaystyle Y\sim\mathcal{B}(0.5),\quad
Z\sim\mathcal{B}(0.9Y+0.1\overline{Y}),\quad
X\sim[\mathcal{N}(Y,10),\mathcal{N}(Z,1)]$

$X_{1}$ directly contains $Y$ and $X_{2}$ contains information about $Y$
through $Z$. The MLE solution to predicting $Y$ from $X$ on $P_{train}$ is a
linear classifier with a higher weight on $X_{2}$ because of the smaller
variance. However, if for $a\in\\{0,1\\}$ we have
$P_{test}(Z=a~{}|~{}Y=a)=0.9(1-a)+0.1a$, then unlike in $P_{train}$, higher
values correspond to label $Y=1$.

##  Appendix D Invariant predictor

There are at least two distinct usages of the word invariance in the
literature. For one usage, invariance to the nuisance usually implies
independence of model and nuisance, potentially conditional on label. For the
other usage, invariance refers to invariant risk, i.e., the risk is the same
for all test distributions in some family. In some distribution families and
for some independence constraints, these can imply each other. Here, we show
that the particular conditional independence involving the nuisance studied in
this work and originally from [22, 35, 27] implies invariant risk.

Satisfying independence
$h_{X}\rotatebox[origin={c}]{90.0}{$\models$}Z~{}|~{}Y$ means
$P_{train}(h_{X}|Y,Z)=P_{train}(h_{X}|Y)$ and the graphical assumptions on the
anti-causal family mean $P_{train}(h_{X}|Y,Z)=P(h_{X}|Y,Z)$ in any member of
the family. Combined this means $P_{train}(h_{X}|Y,Z)=P(h_{X}|Y)$ in any
member of the family when the model $P_{train}(Y|h_{X})$ satisfies
$h_{X}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y$. Consider test set
performance $\mathop{\mathbb{E}}_{P_{test}(Y,X)}[\log P_{train}(Y|h_{X})]$. By
the assumption on the family, by Bayes, and by satisfying the independence
constraint:

$\displaystyle\mathop{\mathbb{E}}_{P_{test}(Y,X)}[\log P_{train}(Y|h_{X})]$
$\displaystyle=\mathop{\mathbb{E}}_{P_{test}(Y,X)}\Big{[}\log\frac{P_{train}(h_{X}|Y)P(Y)}{P_{train}(h_{X})}\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{P_{test}(Y,X,Z)}\Big{[}\log\frac{P_{train}(h_{X}|Y,Z)P(Y)}{\mathop{\mathbb{E}}_{P(Y)}\Big{[}P_{train}(h_{X}|Y,Z)\Big{]}}\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{P_{test}(Y,h_{X},Z)}\Big{[}\log\frac{P_{train}(h_{X}|Y,Z)P(Y)}{\mathop{\mathbb{E}}_{P(Y)}\Big{[}P_{train}(h_{X}|Y,Z)\Big{]}}\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{P_{test}(Y,h_{X},Z)}\Big{[}\log\frac{P(h_{X}|Y,Z)P(Y)}{\mathop{\mathbb{E}}_{P(Y)}\Big{[}P(h_{X}|Y,Z)\Big{]}}\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{P_{test}(Y,h_{X},Z)}\Big{[}\log\frac{P(h_{X}|Y)P(Y)}{\mathop{\mathbb{E}}_{P(Y)}\Big{[}P(h_{X}|Y)\Big{]}}\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{P_{test}(Y,h_{X})}\Big{[}\log\frac{P(h_{X}|Y)P(Y)}{\mathop{\mathbb{E}}_{P(Y)}\Big{[}P(h_{X}|Y)\Big{]}}\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{P_{test}(h_{X}|Y)P_{test}(Y)}\Big{[}\log\frac{P(h_{X}|Y)P(Y)}{\mathop{\mathbb{E}}_{P(Y)}\Big{[}P(h_{X}|Y)\Big{]}}\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{P(h_{X}|Y)P(Y)}\Big{[}\log\frac{P(h_{X}|Y)P(Y)}{\mathop{\mathbb{E}}_{P(Y)}\Big{[}P(h_{X}|Y)\Big{]}}\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{P(h_{X},Y)}\Big{[}\log\frac{P(h_{X}|Y)P(Y)}{\mathop{\mathbb{E}}_{P(Y)}\Big{[}P(h_{X}|Y)\Big{]}}\Big{]}$

The last quantity does not depend on any specific $P_{D}(Z|Y)$. This means
that performance of the $P_{train}(Y|h_{X})$ model, when the independence is
satisfied, is the same on all $P_{test}$ in $\mathcal{F}$.

##  Appendix E Estimation in practice

###  E.1 Splitting samples

For a given batch, we use $1/4$ of the samples for the normalization term and
$3/4$ for the main term, though this number may be changed. Further, the main
term of any of the three estimators is defined on a pair of independent
samples, i.e. it is a U-statistic. There are two ways to estimate such
expectations. One option is to further break the samples left for the main
term in half into two batches $S_{1}$ and $S_{2}$ and then compute on all
pairs $i\in S_{1},j\in S_{2}$. The alternative, which has slightly higher
sampler efficiency and is the method we use, is to compute on all pairs of
samples and then leave out any diagonal terms $k(X_{i},X_{i})$ from the
average.

###  E.2 Trade-offs among the 3 proposed estimators

For large samples, dr estimates with correct $G_{X}$ and correct $m_{X}$ are
lower variance than the regression with correct $m_{X}$, and lower variance
than re-weighting with correct $G_{X}$. Even when $m_{X}$ is mis-specified but
$G_{X}$ is correct, the dr estimator may still be lower variance than the re-
weighting estimator with correct $G_{X}$ alone. However, the dr estimator with
correct $m_{X}$ but mis-specified $G_{X}$ may be higher variance than the
regression estimator with correct $m_{X}$ [8]. For this reason, when the
missingness model $G_{X}$ is wrong, the regression estimator may out-perform
the dr estimator even in large sample sizes.

The variance of the dr and re-weighting estimators comes from two distinct
places. One is general to missingness: small observation probabilities $G_{X}$
in the denominator. The other reason is general Monte Carlo error: we need
individual samples of $\tilde{Z}$ in the numerator. This is especially a
problem in the spurious correlation setting: $Y$ and $Z$ are possibly strongly
correlated. We need to compute the mmd conditional on $Y=y$ which involves,
for each $Y=y$, expectations using samples where $Z=1$ and where $Z=0$, but we
may have very few samples for one of these $Z$ values. This second source of
variance also applies to estimates of the full-data mmd under no missingness
(eq. 4). We compare the mean and variance of these estimators empirically in
Section E.3.

###  E.3 Empirical investigation of variance

(a) Means

(b) Std. Dev.’s

Figure 3:  Figure 3(a): Mean of 100 MMD estimates at each batch size. Figure
3(b): Standard Deviation of 100 MMD estimates at each batch size

As discussed, when $\mathop{\mathbb{E}}[\Delta|X]$ small, or $(Y,Z)$ highly
correlated, or both, all estimators will be high variance. We train a model on
the experiment 1 simulation using the none method and then study the mean and
variance of dr, dr \+ (to study the effect of using the true $G_{X}$), reg
(since it yielded better performance on mimic) and full (since this method is
used to report the mmds in the tables). In this simulation, we are free to
generate as many large batches of samples as needed. Keeping the model fixed,
for each batch size between $1000$ and $10,000$ incrementing by $250$ we draw
$100$ new batches of that size and estimate the mmd using each method. For
each method, we report the mean (fig. 3(a)) and standard deviation (fig. 3(b))
of these estimates.

Notably, we cannot compute an actual ground-truth for the mmd of this model,
but we could take the mean of the full estimate (no missingness) at the
largest sample size of $10,000$ samples. This is about $0.2$. We see that the
regression estimator stays closer to this number for all sample sizes relative
to the dr methods. Interesting, for standard deviation, we see that the dr
estimator is more well-behaved than the dr \+ estimator that uses the true
$G_{X}$. This has also been observed for learned versus true propensity scores
in treatment effect estimation and usually results from models learning less
extreme probabilities than the true ones, trading some bias. In this case,
there is not a substantial difference in estimated weights or in bias, but
there is a large difference in variance. More investigation is required.

The main take-away from both plots is that the regression method seems more
stable than dr and that $G_{X}$ may be the part of the dr estimator that is
not being learned well. On the other hand the dr estimator may possibly be
safer when it is unknown if it is easier to estimate $G_{X}$ or $m_{X}$. We
recommend using all $3$ of the proposed estimators and comparing validation
objectives.

##  Appendix F Full experiments

Table 4: Simulation. $\lambda=1$. | none | obs | full | dr |  dr + | reg | ip |  ip +
---|---|---|---|---|---|---|---|---
tr mmd | $0.21\pm 0.04$ | $0.05\pm 0.04$ | $0.00\pm 0.01$ | $0.00\pm 0.01$ | $0.00\pm 0.01$ | $0.01\pm 0.00$ | $0.00\pm 0.00$ | $0.00\pm 0.01$
tr acc | $0.89\pm 0.00$ | $0.87\pm 0.00$ | $0.86\pm 0.01$ | $0.85\pm 0.01$ | $0.84\pm 0.02$ | $0.86\pm 0.00$ | $0.84\pm 0.01$ | $0.84\pm 0.01$
te acc | $0.67\pm 0.02$ | $0.77\pm 0.02$ | $0.80\pm 0.01$ | $0.81\pm 0.02$ | $0.81\pm 0.01$ | $0.79\pm 0.02$ | $0.82\pm 0.02$ | $0.81\pm 0.00$
Table 5: Simulation. $\lambda=5$. | none | obs | full | dr |  dr + | reg | ip |  ip +
---|---|---|---|---|---|---|---|---
tr mmd | $0.21\pm 0.04$ | $0.03\pm 0.02$ | $0.00\pm 0.01$ | $0.00\pm 0.00$ | $0.00\pm 0.0$ | $0.00\pm 0.01$ | $0.00\pm 0.0$ | $0.00\pm 0.0$
tr acc | $0.89\pm 0.0$ | $0.85\pm 0.02$ | $0.84\pm 0.01$ | $0.82\pm 0.01$ | $0.78\pm 0.06$ | $0.84\pm 0.00$ | $0.81\pm 0.02$ | $0.81\pm 0.03$
te acc | $0.67\pm 0.02$ | $0.78\pm 0.02$ | $0.83\pm 0.01$ | $0.82\pm 0.02$ | $0.77\pm 0.04$ | $0.82\pm 0.01$ | $0.81\pm 0.02$ | $0.80\pm 0.01$
Table 6:  mnist $\lambda=1$. | none | obs | full | dr |  dr + | reg | ip |  ip +
---|---|---|---|---|---|---|---|---
tr mmd | $2.05\pm 0.18$ | $0.02\pm 0.04$ | $0.00\pm 0.01$ | $0.00\pm 0.01$ | $0.00\pm 0.01$ | $0.00\pm 0.01$ |  $0.07\pm 0.12$. | $0.03\pm 0.06$
tr acc | $0.90\pm 0.01$ | $0.74\pm 0.03$ | $0.76\pm 0.01$ | $0.77\pm 0.00$ | $0.76\pm 0.01$ | $0.76\pm 0.01$ | $0.67\pm 0.16$ | $0.68\pm 0.15$
te acc | $0.13\pm 0.01$ | $0.63\pm 0.17$ | $0.74\pm 0.01$ | $0.72\pm 0.04$ | $0.73\pm 0.01$ | $0.73\pm 0.01$ | $0.64\pm 0.14$ | $0.61\pm 0.11$
Table 7:  mnist $\lambda=5$. | none | obs | full | dr |  dr + | reg | ip |  ip +
---|---|---|---|---|---|---|---|---
tr mmd | $2.05\pm 0.18$ | $0.01\pm 0.02$ | $0.00\pm 0.00$ | $0.00\pm 0.00$ | $0.00\pm 0.01$ | $0.00\pm 0.01$ | $0.01\pm 0.01$ | $0.01\pm 0.02$
tr acc | $0.9\pm 0.01$ | $0.66\pm 0.15$ | $0.75\pm 0.01$ | $0.65\pm 0.14$ | $0.65\pm 0.13$ | $0.75\pm 0.01$ | $0.71\pm 0.08$ | $0.60\pm 0.12$
te acc | $0.13\pm 0.01$ | $0.65\pm 0.15$ | $0.75\pm 0.01$ | $0.73\pm 0.02$ | $0.70\pm 0.09$ | $0.75\pm 0.01$ | $0.55\pm 0.3$ | $0.60\pm 0.12$

##  Appendix G Failures of restricting to observed data

######  Proposition.

There exist distributions on $(X,Y,\Delta,Z)$ such that

$\displaystyle\exists h_{X}^{\star}\quad\text{s.t.}\quad
h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y,\quad\text{but}\quad
h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y=y,\Delta=1$

and there exist distributions on $(X,Y,\Delta,Z)$ such that

$\displaystyle\exists h_{X}^{\star}\quad\text{s.t.}\quad
h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y,\Delta=1\quad\text{but}\quad
h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y=y$

#### First direction.

There exist distributions on $(X,Y,\Delta,Z)$ such that

$\displaystyle\exists h_{X}^{\star}\quad\text{s.t.}\quad
h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y,\quad\text{but}\quad
h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y=y,\Delta=1$

It suffices to illustrate this even when $Z,Y$ are not correlated. Consider

$\displaystyle Y\sim\mathcal{N}(0,1),\quad
Z\sim\mathcal{N}(0,\sigma_{Z}^{2}),\quad\epsilon_{X}\sim\mathcal{N}(0,\sigma_{X}^{2}),\quad
X$ $\displaystyle=[Y-Z+\epsilon_{X},Y+Z]$

For $h_{X}^{\star}=(X_{1}+X_{2})$, we first show
$h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y$. We have

$\displaystyle h_{X}^{\star}|Y\sim\mathcal{N}(2Y,\sigma^{2}_{X})$

and in particular $h_{X}^{\star}=2Y+\epsilon_{X}$. Given $Y=y$, the only
randomness in $h_{X}^{\star}$ is due $\epsilon_{X}$. But $\epsilon_{X}$ is
independent of the joint variable $(Z,Y)$ meaning
$\epsilon_{X}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y$ and therefore
$h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y$.

We now construct $\Delta|(X,Y)$ such that
$h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y=y,\Delta=1$. Let

$\displaystyle\Delta=\text{OR}\Big{(}\mathbbm{1}\left[X_{1}+X_{2}<0\right],\mathbbm{1}\left[X_{2}-Y<0\right]\Big{)}.$

Checking the condition

$\displaystyle
h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y=y,\Delta=1$

(using definition of $h_{X}^{\star}$) is equivalent to checking

$\displaystyle(X_{1}+X_{2})\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y=y,\Delta=1$

(using definition of $\Delta$) is equivalent to checking

$\displaystyle(X_{1}+X_{2})\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y=y,\text{OR}\Big{(}\mathbbm{1}\left[X_{1}+X_{2}<0\right],\mathbbm{1}\left[X_{2}-Y<0\right]\Big{)}=1$

(using defintion of $X_{2}$) is equivalent to checking

$\displaystyle(X_{1}+X_{2})\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y=y,\text{OR}\Big{(}\mathbbm{1}\left[X_{1}+X_{2}<0\right],\mathbbm{1}\left[Z<0\right]\Big{)}=1$

To check that, we need to check if the distribution of
$(X_{1}+X_{2})|Y=y,\Delta=1$ changes when conditioning on different events
involving the random variable $Z$. For example, $\mathbbm{1}\left[Z<0\right]$
and $\mathbbm{1}\left[Z\geq 0\right]$:

  1. 1.

$(X_{1}+X_{2})\quad|\quad
Y=y,\text{OR}\Big{(}\mathbbm{1}\left[X_{1}+X_{2}<0\right],\mathbbm{1}\left[Z<0\right]\Big{)}=1,\mathbbm{1}\left[Z<0\right]=1$

  2. 2.

$(X_{1}+X_{2})\quad|\quad
Y=y,\text{OR}\Big{(}\mathbbm{1}\left[X_{1}+X_{2}<0\right],\mathbbm{1}\left[Z<0\right]\Big{)}=1,\mathbbm{1}\left[Z\geq
0\right]=1$.

We can show these two conditional variables differ in distribution simply by
showing they differ in support. The first conditional variable can be full
support because the event $\mathbbm{1}\left[Z<0\right]$ satisfies one of the
OR conditions leaving the other condition
$\mathbbm{1}\left[X_{1}+X_{2}<0\right]=\mathbbm{1}\left[h_{X}^{\star}<0\right]$
free to take either value. However, the second conditional variable needs
$X_{1}+X_{2}=h_{X}^{\star}<0$ because $\mathbbm{1}\left[Z<0\right]$ is not
satisfied (since we condition on $\mathbbm{1}\left[Z\geq 0\right]=1$) but the
OR has to be $1$. These different supports imply the distributions differ.
That the variables differ on two non-measure zero sets is enough to show
dependence. Then
$(X_{1}+X_{2})\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y=y,\Delta=1$ which
means
$h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y=y,\Delta=1$.

#### Second direction.

There exist distributions on $(X,Y,\Delta,Z)$ such that

$\displaystyle\exists h_{X}^{\star}\quad\text{s.t.}\quad
h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y,\Delta=1\quad\text{but}\quad
h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y=y$

Let the data be drawn as

$\displaystyle Y\sim\mathcal{N}(0,1),\quad Z\sim\mathcal{B}(0.5),\quad X$
$\displaystyle=[Y-Z,Y+Z]$

Let $h_{X}^{\star}=\mathbbm{1}\left[X_{1}\geq 0\right]$. We first show
$h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\not\models$}Z|Y=y$. We have

$\displaystyle h_{X}^{\star}$ $\displaystyle=\mathbbm{1}\left[X_{1}\geq
0\right]$ $\displaystyle=\mathbbm{1}\left[Y-Z\geq 0\right]$

Given $Y=y$, we ask if the random variable $\mathbbm{1}\left[y-Z\geq 0\right]$
is independent of $Z$. To show dependence, we show that the random variable
$\mathbbm{1}\left[y-Z\geq 0\right]$ changes in distribution when $Z$ takes on
its two values:

  1. 1.

$\mathbbm{1}\left[y-Z\geq 0\right]|Y=y,Z=0$

  2. 2.

$\mathbbm{1}\left[y-Z\geq 0\right]|Y=y,Z=1$

Suppose $y\in(0,1)$. When $Z=0$ we have that $\mathbbm{1}\left[y-Z\geq
0\right]=1$ with probability one. When $Z=1$, we have
$\mathbbm{1}\left[y-Z\geq 0\right]=0$ with probability one. Therefore the
variables are dependent.

We now let $\Delta=\mathbbm{1}\left[X_{1}\geq
0\right]=\mathbbm{1}\left[Y-Z\geq 0\right]$ and show
$h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y,\Delta=1$. Note
that $\Delta(X,Y)=h_{X}^{\star}$. We ask whether

$\displaystyle\mathbbm{1}\left[Y-Z\geq
0\right]\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y,\mathbbm{1}\left[Y-Z\geq
0\right]$

The conditioning set fully determines the variable $\mathbbm{1}\left[Y-Z\geq
0\right]$ meaning it is a constant and is therefore independent of $Z$.
Therefore $h_{X}^{\star}\rotatebox[origin={c}]{90.0}{$\models$}Z|Y=y,\Delta=1$
as desired.

##  Appendix H IP and outcome estimators

We review estimation of $\mathop{\mathbb{E}}[Z]$ under missingness. Two pieces
of the data generation process can help, the missingness process $G_{X}$ and
the conditional expectation $m_{X}$ of the missing variable:

$\displaystyle G_{X}\triangleq\mathbb{E}[\Delta~{}|~{}X,Y],\quad\quad
m_{X}\triangleq\mathbb{E}[Z~{}|~{}X,Y]$

Inverse-weighting estimators use $G_{X}$ [16, 5, 29, 34, 15]

$\displaystyle\begin{split}\mathop{\mathbb{E}}[Z]&=\mathop{\mathbb{E}}_{X}\mathop{\mathbb{E}}_{Z|X}[Z]\\\
&=\mathop{\mathbb{E}}_{X}\mathop{\mathbb{E}}_{Z|X}\Big{[}\frac{\mathop{\mathbb{E}}[\Delta|X]}{\mathop{\mathbb{E}}[\Delta|X]}Z\Big{]}\\\
&=\mathop{\mathbb{E}}_{X}\mathop{\mathbb{E}}_{Z|X}\mathop{\mathbb{E}}_{\Delta|X}\Big{[}\frac{\Delta
Z}{\mathop{\mathbb{E}}[\Delta|X]}\Big{]}\\\
&=\mathop{\mathbb{E}}_{XZ\Delta}\Big{[}\frac{\Delta
Z}{\mathop{\mathbb{E}}[\Delta|X]}\Big{]}\\\ &=\mathop{\mathbb{E}}_{X\Delta
Z}\Big{[}\frac{\Delta Z}{G_{X}}\Big{]}\\\ &=\mathop{\mathbb{E}}_{X\Delta
Z}\Big{[}\frac{\Delta\tilde{Z}}{G_{X}}\Big{]}\end{split}$ (8)

This means we can estimate $\mathop{\mathbb{E}}[Z]$ provided that (1)
ignorability and positivity hold and (2) $G_{X}$ is known. $G_{X}$ can be
estimated by regressing $\Delta$ on $X$. Alternatively, standardization
estimators use $m_{X}$ [30, 33, 31, 25, 20, 15]:

$\displaystyle\mathop{\mathbb{E}}[Z]=\mathop{\mathbb{E}}_{X}\Big{[}\mathop{\mathbb{E}}[Z|X]\Big{]}=\mathop{\mathbb{E}}_{X}\Big{[}\mathop{\mathbb{E}}[Z|X,\Delta=1]\Big{]}=\mathop{\mathbb{E}}_{X}[m_{X}]$
(9)

$m_{X}$ can be estimated by regressing $\tilde{Z}$ on $X$ where $\Delta=1$ (by
ignorability).

##  Appendix I DR estimator of mean of Z

The inverse weighting and regression estimators can be combined. Equation 1
defines the dr estimator of $\mathop{\mathbb{E}}[Z]$ by

$\displaystyle\mathop{\mathbb{E}}[Z]=\mathop{\mathbb{E}}\Big{[}\frac{\Delta\tilde{Z}}{G_{X}}-\frac{\Delta-
G_{X}}{G_{X}}m_{X}\Big{]}$

Let us re-write this expectation until we see it equals
$\mathop{\mathbb{E}}[Z]$ when $G$ or $m$ are correct.

$\displaystyle\mathop{\mathbb{E}}\Big{[}\frac{\Delta\tilde{Z}}{G_{X}}-\frac{\Delta-
G_{X}}{G_{X}}m_{X}\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}\Big{[}\frac{\Delta Z}{G_{X}}-\frac{\Delta-
G_{X}}{G_{X}}m_{X}\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}\Big{[}Z+\frac{\Delta
Z}{G_{X}}-Z-\frac{\Delta-G_{X}}{G_{X}}m_{X}\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}\Big{[}Z+\frac{\Delta
Z}{G_{X}}-\frac{G_{X}}{G_{X}}Z-\frac{\Delta-G_{X}}{G_{X}}m_{X}\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}\Big{[}Z+\frac{\Delta-
G_{X}}{G_{X}}Z-\frac{\Delta-G_{X}}{G_{X}}m_{X}\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}\Big{[}Z+\frac{\Delta-
G_{X}}{G_{X}}(Z-m_{X})\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}\Big{[}Z\Big{]}+\mathop{\mathbb{E}}\Big{[}\frac{\Delta-
G_{X}}{G_{X}}(Z-m_{X})\Big{]}$

The first term is what we want, so we just have to check if the second term is
$0$ when either $G$ or $m$ are correct. If $G$ is correct (regardless of $m$)
then:

$\displaystyle\mathop{\mathbb{E}}\Big{[}\frac{\Delta-
G_{X}}{G_{X}}(Z-m_{X})\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}\Bigg{[}\mathop{\mathbb{E}}\Big{[}\frac{\Delta-
G_{X}}{G_{X}}(Z-m_{X})\Big{|}X,Z\Big{]}\Bigg{]}$
$\displaystyle=\mathop{\mathbb{E}}\Bigg{[}\frac{\mathop{\mathbb{E}}[\Delta|X,Z]-G_{X}}{G_{X}}(Z-m_{X})\Bigg{]}$
$\displaystyle=\mathop{\mathbb{E}}\Bigg{[}\frac{\mathop{\mathbb{E}}[\Delta|X]-G_{X}}{G_{X}}(Z-m_{X})\Bigg{]}$
$\displaystyle=\mathop{\mathbb{E}}\Bigg{[}\frac{G_{X}-G_{X}}{G_{X}}(Z-m_{X})\Bigg{]}=0$

When $m$ is correct (regardless of $G$):

$\displaystyle\mathop{\mathbb{E}}\Big{[}\frac{\Delta-
G_{X}}{G_{X}}(Z-m_{X})\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}\Bigg{[}\mathop{\mathbb{E}}\Big{[}\frac{\Delta-
G_{X}}{G_{X}}(Z-m_{X})\Big{|}X,\Delta\Big{]}\Bigg{]}$
$\displaystyle=\mathop{\mathbb{E}}\Bigg{[}\frac{\Delta-
G_{X}}{G_{X}}(\mathop{\mathbb{E}}[Z|X,\Delta]-m_{X})\Big{|}X,\Delta\Bigg{]}$
$\displaystyle=\mathop{\mathbb{E}}\Bigg{[}\frac{\Delta-
G_{X}}{G_{X}}(\mathop{\mathbb{E}}[Z|X,\Delta]-\mathop{\mathbb{E}}[Z|X,\Delta=1])\Big{|}X,\Delta\Bigg{]}$
$\displaystyle=\mathop{\mathbb{E}}\Bigg{[}\frac{\Delta-
G_{X}}{G_{X}}(\mathop{\mathbb{E}}[Z|X]-\mathop{\mathbb{E}}[Z|X])\Big{|}X,\Delta\Bigg{]}=0$

##  Appendix J MMD estimators under missingness

###  J.1 Statement of $G_{x}$ and $m_{x}$ based estimators of the mmd

######  Proposition.

($G_{X}$-based re-weighted estimator) Assume positivity, ignorability, and,
for each $X$, $G_{X}=\mathop{\mathbb{E}}[\Delta|X,Y]$. Then,

$\displaystyle\mathop{\mathbb{E}}_{\begin{subarray}{c}X|Z=b\\\
X^{\prime}|Z^{\prime}=b^{\prime}\end{subarray}}\Big{[}k_{XX^{\prime}}\Big{]}$
$\displaystyle=\frac{1}{N(b,b^{\prime})}\mathop{\mathbb{E}}\Big{[}\frac{\Delta\Delta^{\prime}\tilde{Z}_{b}\tilde{Z^{\prime}}_{b^{\prime}}}{G_{XX^{\prime}}}k_{XX^{\prime}}\Big{]}.$
(10)

######  Proposition.

($m_{X}$-based regression estimator) Assume ignorability, and, for each $X$,
$m_{X}=\mathop{\mathbb{E}}[Z|X,Y]$. Then,

$\displaystyle\mathop{\mathbb{E}}_{\begin{subarray}{c}X|Z=b\\\
X^{\prime}|Z^{\prime}=b^{\prime}\end{subarray}}\Big{[}k_{XX^{\prime}}\Big{]}$
$\displaystyle=\frac{1}{N(b,b^{\prime})}\mathop{\mathbb{E}}\Big{[}m_{Xb}\cdot
m_{X^{\prime}b^{\prime}}\cdot k_{XX^{\prime}}\Big{]}.$ (11)

###  J.2 Deriving the $G_{X}$-based re-weighted estimator

Here we start at the target quantity and derive the estimator. We give the
derivation for $Z=1,Z^{\prime}=1$. The other cases are analogous.

$\displaystyle\mathop{\mathbb{E}}_{P(X|Z=1)P(X^{\prime}|Z^{\prime}=1)}\Big{[}k_{XX^{\prime}}\Big{]}$
$\displaystyle=\int_{X,X^{\prime}}kP(X|Z=1)P(X^{\prime}|Z^{\prime}=1)dXdX^{\prime}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\int_{X,X^{\prime}}kP(Z=1,X)P(Z^{\prime}=1,X^{\prime})dXdX^{\prime}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\int_{X,X^{\prime}}kP(Z=1|X)P(Z^{\prime}=1|X)P(X)P(X^{\prime})dXdX^{\prime}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\int_{X,X^{\prime}}k\mathop{\mathbb{E}}(Z=1|X)\mathop{\mathbb{E}}(Z^{\prime}=1|X)P(X)P(X^{\prime})dXdX^{\prime}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{\begin{subarray}{c}X,Z\\\
X^{\prime},Z^{\prime}\end{subarray}}\Big{[}k\cdot Z\cdot Z^{\prime}\Big{]}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{\begin{subarray}{c}X,Z\\\
X^{\prime},Z^{\prime}\end{subarray}}\Big{[}\frac{\mathop{\mathbb{E}}[\Delta|X]\mathop{\mathbb{E}}[\Delta^{\prime}|X^{\prime}]}{\mathop{\mathbb{E}}[\Delta|X]\mathop{\mathbb{E}}[\Delta^{\prime}|X^{\prime}]}k\cdot
Z\cdot Z^{\prime}\Big{]}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta,Z\\\
X^{\prime},\Delta^{\prime},Z^{\prime}\end{subarray}}\Big{[}\frac{\Delta\Delta^{\prime}}{G_{X}G_{X^{\prime}}}k\cdot
Z\cdot Z^{\prime}\Big{]}$

###  J.3 Deriving the $m_{X}$-based standardization estimator

Here we start at the target quantity and derive the estimator. We give the
derivation for $Z=1,Z^{\prime}=1$. The other cases are analogous.

$\displaystyle\mathop{\mathbb{E}}_{P(X|Z=1)P(X^{\prime}|Z^{\prime}=1)}\Big{[}k_{XX^{\prime}}\Big{]}$
$\displaystyle=\int_{X,X^{\prime}}kP(X|Z=1)P(X^{\prime}|Z^{\prime}=1)dXdX^{\prime}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\int_{X,X^{\prime}}kP(Z=1,X)P(Z^{\prime}=1,X^{\prime})dXdX^{\prime}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\int_{X,X^{\prime}}kP(Z=1|X)P(Z^{\prime}=1|X)P(X)P(X^{\prime})dXdX^{\prime}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\int_{X,X^{\prime}}k\mathop{\mathbb{E}}(Z=1|X)\mathop{\mathbb{E}}(Z^{\prime}=1|X)P(X)P(X^{\prime})dXdX^{\prime}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{X,X^{\prime}}\Big{[}m_{X}\cdot
m_{X^{\prime}}\cdot k\Big{]}$

###  J.4 Deriving the DR estimator

Here we start at the estimator and derive the target quantity. We give the
derivation for $Z=1,Z^{\prime}=1$. The other cases are analogous.

$\displaystyle\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\frac{1}{N(N-1)}\sum_{i\neq
j}\Big{[}\frac{\Delta_{ij}\tilde{Z}_{ij}}{G_{ij}}k_{ij}-\frac{\Delta_{ij}-G_{ij}}{G_{ij}}m_{ij}k_{ij}\Big{]}$
$\displaystyle\approx\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta,Z\\\
X^{\prime},\Delta^{\prime},Z^{\prime}\end{subarray}}\Big{[}\frac{\Delta\Delta^{\prime}\tilde{Z}\tilde{Z}^{\prime}}{G_{X}G_{X^{\prime}}}k-\frac{\Delta\Delta^{\prime}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}m_{X}m_{X^{\prime}}k\Big{]}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta,Z\\\
X^{\prime},\Delta^{\prime},Z^{\prime}\end{subarray}}\Big{[}\frac{\Delta\Delta^{\prime}ZZ^{\prime}}{G_{X}G_{X^{\prime}}}k-\frac{\Delta\Delta^{\prime}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}m_{X}m_{X^{\prime}}k\Big{]}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta,Z\\\
X^{\prime},\Delta^{\prime},Z^{\prime}\end{subarray}}\Big{[}ZZ^{\prime}k+\frac{\Delta\Delta^{\prime}}{G_{X}G_{X^{\prime}}}ZZ^{\prime}k-ZZ^{\prime}k-\frac{\Delta\Delta^{\prime}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}m_{X}m_{X^{\prime}}k\Big{]}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta,Z\\\
X^{\prime},\Delta^{\prime},Z^{\prime}\end{subarray}}\Big{[}ZZ^{\prime}k+\frac{\Delta\Delta^{\prime}}{G_{X}G_{X^{\prime}}}ZZ^{\prime}k-\frac{G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}ZZ^{\prime}k-\frac{\Delta\Delta^{\prime}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}m_{X}m_{X^{\prime}}k\Big{]}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta,Z\\\
X^{\prime},\Delta^{\prime},Z^{\prime}\end{subarray}}\Big{[}ZZ^{\prime}k+\frac{\Delta\Delta^{\prime}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}ZZ^{\prime}k-\frac{\Delta\Delta^{\prime}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}m_{X}m_{X^{\prime}}k\Big{]}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta,Z\\\
X^{\prime},\Delta^{\prime},Z^{\prime}\end{subarray}}\Big{[}ZZ^{\prime}k+\frac{\Delta\Delta^{\prime}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}\Big{(}ZZ^{\prime}-m_{X}m_{X^{\prime}}\Big{)}k\Big{]}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta,Z\\\
X^{\prime},\Delta^{\prime},Z^{\prime}\end{subarray}}\Big{[}ZZ^{\prime}k\Big{]}+\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta,Z\\\
X^{\prime},\Delta^{\prime},Z^{\prime}\end{subarray}}\Big{[}\frac{\Delta\Delta^{\prime}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}\Big{(}ZZ^{\prime}-m_{X}m_{X^{\prime}}\Big{)}k\Big{]}$

Our estimator equals two terms. We first show that the first term equals the
desired quantity, and then show the second term equals 0 when either auxiliary
model is correct.

$\displaystyle\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta,Z\\\
X^{\prime},\Delta^{\prime},Z^{\prime}\end{subarray}}\Big{[}ZZ^{\prime}k\Big{]}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{X,X^{\prime}}\Big{[}k\mathop{\mathbb{E}}[Z,Z^{\prime}|X,X^{\prime}]\Big{]}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\mathop{\mathbb{E}}_{X,X^{\prime}}\Big{[}kP(Z=1,Z^{\prime}=1|X,X^{\prime})\Big{]}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\int_{X,X^{\prime}}kP(Z=1,Z^{\prime}=1|X,X^{\prime})P(X,X^{\prime})dXdX^{\prime}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\int_{X,X^{\prime}}kP(Z=1|X)P(Z^{\prime}=1|X)P(X)P(X^{\prime})dXdX^{\prime}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\int_{X,X^{\prime}}kP(Z=1,X)P(Z^{\prime}=1,X^{\prime})dXdX^{\prime}$
$\displaystyle=\frac{1}{P(Z=1)}\frac{1}{P(Z^{\prime}=1)}\int_{X,X^{\prime}}kP(Z=1,X)P(Z^{\prime}=1,X^{\prime})dXdX^{\prime}$
$\displaystyle=\int_{X,X^{\prime}}kP(X|Z=1)P(X^{\prime}|Z^{\prime}=1)dXdX^{\prime}$
$\displaystyle=\mathop{\mathbb{E}}_{P(X|Z=1)P(X^{\prime}|Z^{\prime}=1)}\Big{[}k\Big{]}$

That’s the expectation we want missing just the $P(Z=1)$ constants, so now we
should show the next term is $0$ when either $m$ or $G$ are correct. When $G$
correct:

$\displaystyle\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta,Z\\\
X^{\prime},\Delta^{\prime},Z^{\prime}\end{subarray}}\Big{[}\frac{\Delta\Delta^{\prime}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}\Big{(}ZZ^{\prime}-m_{X}m_{X^{\prime}}\Big{)}k\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}X,Z\\\
X^{\prime},Z^{\prime}\end{subarray}}\Big{[}\frac{\mathop{\mathbb{E}}[\Delta\Delta^{\prime}|X,X^{\prime},Y,Z^{\prime}]-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}\Big{(}ZZ^{\prime}-m_{X}m_{X^{\prime}}\Big{)}k\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}X,Z\\\
X^{\prime},Z^{\prime}\end{subarray}}\Big{[}\frac{\mathop{\mathbb{E}}[\Delta\Delta^{\prime}|X,X^{\prime}]-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}\Big{(}ZZ^{\prime}-m_{X}m_{X^{\prime}}\Big{)}k\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}X,Z\\\
X^{\prime},Z^{\prime}\end{subarray}}\Big{[}\frac{\mathop{\mathbb{E}}[\Delta|X]\mathop{\mathbb{E}}[\Delta^{\prime}|X^{\prime}]-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}\Big{(}ZZ^{\prime}-m_{X}m_{X^{\prime}}\Big{)}k\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}X,Z\\\
X^{\prime},Z^{\prime}\end{subarray}}\Big{[}\frac{G_{X}G_{X^{\prime}}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}\Big{(}ZZ^{\prime}-m_{X}m_{X^{\prime}}\Big{)}k\Big{]}=0$

Likewise, when $m$ correct:

$\displaystyle\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta,Z\\\
X^{\prime},\Delta^{\prime},Z^{\prime}\end{subarray}}\Big{[}\frac{\Delta\Delta^{\prime}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}\Big{(}ZZ^{\prime}-m_{X}m_{X^{\prime}}\Big{)}k\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta\\\
X^{\prime},\Delta^{\prime}\end{subarray}}\Big{[}\frac{\Delta\Delta^{\prime}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}\Big{(}\mathop{\mathbb{E}}[ZZ^{\prime}|X,X^{\prime},\Delta,\Delta^{\prime}]-m_{X}m_{X^{\prime}}\Big{)}k\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta\\\
X^{\prime},\Delta^{\prime}\end{subarray}}\Big{[}\frac{\Delta\Delta^{\prime}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}\Big{(}\mathop{\mathbb{E}}[Z|X,\Delta]\mathop{\mathbb{E}}[Z^{\prime}|X^{\prime},\Delta^{\prime}]-m_{X}m_{X^{\prime}}\Big{)}k\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta\\\
X^{\prime},\Delta^{\prime}\end{subarray}}\Big{[}\frac{\Delta\Delta^{\prime}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}\Big{(}\mathop{\mathbb{E}}[Z|X,\Delta]\mathop{\mathbb{E}}[Z^{\prime}|X^{\prime},\Delta^{\prime}]-\mathop{\mathbb{E}}[Z|X,\Delta=1]\mathop{\mathbb{E}}[Z^{\prime}|X^{\prime},\Delta^{\prime}=1]\Big{)}k\Big{]}$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}X,\Delta\\\
X^{\prime},\Delta^{\prime}\end{subarray}}\Big{[}\frac{\Delta\Delta^{\prime}-G_{X}G_{X^{\prime}}}{G_{X}G_{X^{\prime}}}\Big{(}\mathop{\mathbb{E}}[Z|X]\mathop{\mathbb{E}}[Z^{\prime}|X^{\prime}]-\mathop{\mathbb{E}}[Z|X]\mathop{\mathbb{E}}[Z^{\prime}|X^{\prime}]\Big{)}k\Big{]}=0$

The proof for the other two terms is analogous but with using
$\overline{Z}=(1-Z)$ instead of $Z$ and $\overline{m}=1-m$ when conditioning
on $Z=0$.

##  Appendix K kernel mmd between joint and product of marginals

#### Continuous nuisances.

In this work we study binary nuisance. We can instead measure the mmd between
joint $p(h_{X},Z)$ and product of marginals $p(h_{X})P(Z)$, which allows for
continuous nuisance.

The above formulation of mmd between $h_{X}|Z=1$ and $h_{X}|Z=0$ relied on
optimizing with respect to $h$ only: $P(Z)$ is constant in the optimization so
the distance between conditionals specifies the distance between the product
of marginals and joint and thus the dependence. However, considering the more
general case of mmd between $P(h_{X},Z)$ and $P(h_{X})P(Z)$ has the advantage
that is not necessary to consider a finite set of conditioning values for $Z$.
That means the mmd can be extended to continuous nuisance $Z$. Let $X::Z$
denote the concatenation of $X$ and $Z$. The more general formulation is:

$\displaystyle\mathbb{E}_{\begin{subarray}{c}(X,Z)\sim P(X,Z)\\\
(X^{\prime},Z^{\prime})\sim P(X,Z)\end{subarray}}$
$\displaystyle\Bigg{[}k\Big{(}X\cat\cat Z,X^{\prime}\cat\cat
Z^{\prime}\Big{)}\Bigg{]}+\mathbb{E}_{\begin{subarray}{c}(X,Z)\sim P(X)P(Z)\\\
(X^{\prime},Z^{\prime})\sim
P(X^{\prime})P(Z^{\prime})\end{subarray}}\Bigg{[}k\Big{(}X\cat\cat
Z,X^{\prime}\cat\cat Z^{\prime}\Big{)}\Bigg{]}$
$\displaystyle-2\mathbb{E}_{\begin{subarray}{c}(X,Z)\sim P(X,Z)\\\
(X^{\prime},Z^{\prime})\sim
P(X^{\prime})P(Z^{\prime})\end{subarray}}\Bigg{[}k\Big{(}X\cat\cat
Z,X^{\prime}\cat\cat Z^{\prime}\Big{)}\Bigg{]}$

This leads to the following estimator:

$\displaystyle\mathbb{E}_{\begin{subarray}{c}P(X,Z)\\\
P(X^{\prime},Z^{\prime})\end{subarray}}\Bigg{[}k\Big{(}X\cat\cat
Z,X^{\prime}\cat\cat Z^{\prime}\Big{)}\Bigg{]}$
$\displaystyle=\mathbb{E}\Bigg{[}\frac{\Delta\Delta^{\prime}k(X\cat\cat
Z,X^{\prime}\cat\cat
Z^{\prime})}{G_{XX^{\prime}}}-\frac{\Delta\Delta^{\prime}-G_{XX^{\prime}}}{G_{XX^{\prime}}}\mathop{\mathbb{E}}[k|X,X^{\prime}]\Bigg{]}$

and

$\displaystyle\mathbb{E}_{\begin{subarray}{c}P(X)P(Z)\\\
P(X^{\prime})P(Z^{\prime})\end{subarray}}$
$\displaystyle\Bigg{[}k\Big{(}X\cat\cat Z,X^{\prime}\cat\cat
Z^{\prime}\Big{)}\Bigg{]}$
$\displaystyle=\mathbb{E}_{\begin{subarray}{c}P(X_{1})P(X_{2},Z_{2})\\\
P(X_{3})P(X_{4},Z_{4})\end{subarray}}\Bigg{[}k\Big{(}X_{1}\cat\cat
Z_{2},X_{3}\cat\cat Z_{4}\Big{)}\Bigg{]}$
$\displaystyle=\mathbb{E}\Bigg{[}\frac{\Delta\Delta^{\prime}k(X_{1}\cat\cat
Z_{2},X_{3}\cat\cat
Z_{4})}{G_{X_{1}X_{3}}}-\frac{\Delta\Delta^{\prime}-G_{X_{1}X_{3}}}{G_{X_{1}X_{3}}}\mathop{\mathbb{E}}[k(X_{1}\cat\cat
Z_{2},X_{3}\cat\cat Z_{4})|X_{1},X_{3}]\Bigg{]}$

and

$\displaystyle\mathbb{E}_{\begin{subarray}{c}P(X,Z)\\\
P(X^{\prime})P(Z^{\prime})\end{subarray}}$
$\displaystyle\Bigg{[}k\Big{(}X\cat\cat Z,X^{\prime}\cat\cat
Z^{\prime}\Big{)}\Bigg{]}$
$\displaystyle=\mathbb{E}_{\begin{subarray}{c}P(X_{1},Z_{1})\\\
P(X_{2})P(X_{3},Z_{3})\end{subarray}}\Bigg{[}k\Big{(}X_{1}\cat\cat
Z_{1},X_{2}\cat\cat Z_{3}\Big{)}\Bigg{]}$
$\displaystyle=\mathbb{E}\Bigg{[}\frac{\Delta\Delta^{\prime}k(X_{1}\cat\cat
Z_{1},X_{2}\cat\cat
Z_{3})}{G_{X_{1}X_{3}}}-\frac{\Delta\Delta^{\prime}-G_{X_{1}X_{3}}}{G_{X_{1}X_{3}}}\mathop{\mathbb{E}}[k(X_{1}\cat\cat
Z_{1},X_{2}\cat\cat Z_{3})|X_{1},X_{3}]\Bigg{]}$

The challenging part of applying this estimator is that now instead of one
function $m_{X}$ we have three functions, each of which estimates the mean of
$k$ under a different sampling distribution. Moreover, these conditional
expectations depend on the current representation $h_{X}$. This means they
must be updated each time $h$ changes.

##  Appendix L Experimental Details

###  L.1 Simulation details

for $h$ we use a small feed-forward neural network

      class SmallNet(nn.Module):
        def __init__(self, D_in,D_hid,D_out):
            super(SmallNet, self).__init__()
            self.fc1 = nn.Linear(D_in,D_hid)
            self.fc2 = nn.Linear(D_hid,D_hid)
            self.fc3 = nn.Linear(D_hid,D_out)
        def forward(self,h):
            h = RELU(self.fc1(h))
            h = RELU(self.fc2(h))
            return self.fc3(h)

We use $\texttt{D\\_hid}=128$. For $G_{X}$ and $m_{X}$ we use a similar model
but that also takes $y$ as input:

       class SmallAuxNet(nn.Module):
        def __init__(self, D_in,D_hid,D_out):
            super(SmallAuxNet, self).__init__()
            self.fc1 = nn.Linear(D_in+1,D_hid)
            self.fc2 = nn.Linear(D_hid,D_out)
        def forward(self,x,y):
            h = torch.cat([x,y.unsqueeze(-1)],dim=-1)
            h = RELU(self.fc1(h))
            return self.fc2(h)

###  L.2 MNIST details

#### Data details.

Following [11]333 We adapt this repository (linked) to construct textured
mnist and will make our code available., we correlate mnist digits $0$ and $1$
with two textures from the Brodatz dataset (Figure 2). This is an example of
invariance to image backgrounds when not all background labels are available.
We follow a similar setup to colored mnist [3]: because $Y|X$ is essentially
deterministic, even strong spurious correlations may be ignored by a model on
mnist. To push $Y|X$ closer to what may be expected in noisier real data, we
flip the label with 25% chance. Letting $X$ only predict $Y$ with $75\%$
chance means a strong spurious correlation with better predictive power can be
used (a color or in this case a background). The missingness is based on the
average pixel intensity of $X$ and its class. Let $\mu_{X}$ be the mean pixel
value of a 28x28 mnist image. We set

$\displaystyle
Q=\mathbbm{1}\left[Y=1\right]\cdot\mathbbm{1}\left[\mu_{X}<0.3\right],\quad\Delta\sim\mathcal{B}(Q+.2\overline{Q}).$

The choice of $Q$ is correlated with $Z$ through whether the image is light or
dark grey. Similar to proposition 1 and experiment 1, this means subsetting on
$\Delta=1$ may not imply independence on the full population and may throw
away solutions that do.

#### Model details.

For $h$ we use a small convolutional neural network

    class MNISTNet(nn.Module):
        def __init__(self,args):
            super(MNISTNet,self).__init__()
            self.args = args
            self.conv1 = nn.Conv2d(1,6,5)
            self.conv2 = nn.Conv2d(6,16,5)
            self.pool = nn.MaxPool2d(2,2)
            self.fc1 = nn.Linear(256,64)
            self.fc2 = nn.Linear(64,64)
            self.fc3 = nn.Linear(64,1)
        def forward(self,x):
            h = self.pool(self.conv1(x).relu())
            h = self.pool(self.conv2(h).relu())
            h = torch.flatten(h,1)
            h = self.fc1(h).relu()
            h = self.fc2(h).relu()
            return self.fc3(h)

For $G_{X}$ and $m_{X}$ we use a similar model but that also take $y$ as
input. We include $y$ in the computation after the convolutions.

        class NumAuxNet(nn.Module):
        def __init__(self,args):
            super(NumAuxNet,self).__init__()
            self.args = args
            self.conv1 = nn.Conv2d(1,6,5)
            self.conv2 = nn.Conv2d(6,16,5)
            self.pool = nn.MaxPool2d(2,2)
            self.fc1 = nn.Linear(256+1,64)
            self.fc2 = nn.Linear(64,64)
            self.fc3 = nn.Linear(64,1)
        def forward(self,x,y):
            h = self.pool(self.conv1(x).relu())
            h = self.pool(self.conv2(h).relu())
            h = torch.flatten(h,1)
            h = torch.cat([h,y.unsqueeze(-1)],dim=-1)
            h = self.fc1(h).relu()
            h = self.fc2(h).relu()
            return self.fc3(h)

##  Appendix M MIMIC Data Details

We use the following SQL script to extract the dataset from the MIMIC-IV
dataset following the instructions on Physionet for querying on Google Big-
Query. This could also be done in a local MIMIC-IV database. We are grateful
to the MIMIC team especially for creating the MIMIC-DERIVED tables which
include aggregated vitals/labs measured during the first day of ICU stay.

    SELECT
    -- ids
    pat.subject_id as subject_id, adm.hadm_id as hadm_id,icu.stay_id as stay_id,
    -- demographics
    CASE WHEN pat.gender="M" THEN 1 ELSE 0 END as is_male,
    CASE WHEN adm.ethnicity="WHITE" THEN 1 ELSE 0 END as is_white,
    icu_detail.admission_age as age,
    -- weight height
    fdw.weight , fdh.height ,
    -- LOS
    icu.los as los_icu_days,
    -- death
    adm.hospital_expire_flag as expire_flag,
    -- vitals labs min max mean over first day
    vitals.*, labs.*, sofa.*, bg.*,
    FROM ‘physionet-data.mimic_core.patients‘ pat
    INNER JOIN
        ‘physionet-data.mimic_core.admissions‘ adm
            on pat.subject_id=adm.subject_id
    INNER JOIN
        ‘physionet-data.mimic_icu.icustays‘ icu
          on adm.subject_id=icu.subject_id
          and
          adm.hadm_id=icu.hadm_id
    INNER JOIN
         ‘physionet-data.mimic_derived.first_day_height‘ fdh
            on
            adm.subject_id = fdh.subject_id  and icu.stay_id = fdh.stay_id
    INNER JOIN
         ‘physionet-data.mimic_derived.first_day_weight‘ fdw
            on
            adm.subject_id = fdw.subject_id  and icu.stay_id = fdw.stay_id
    INNER JOIN
         ‘physionet-data.mimic_derived.icustay_detail‘ icu_detail
            on
            adm.subject_id=icu_detail.subject_id
            and
            adm.hadm_id=icu_detail.hadm_id
            and
            icu.stay_id=icu_detail.stay_id
    INNER JOIN
        ‘physionet-data.mimic_derived.first_day_bg‘ bg
        on
         adm.subject_id=bg.subject_id
            and
            icu.stay_id = bg.stay_id
    INNER JOIN
        ‘physionet-data.mimic_derived.first_day_sofa‘ sofa
        on
            adm.subject_id=sofa.subject_id
            and
            adm.hadm_id=sofa.hadm_id
            and
            icu.stay_id=sofa.stay_id
    INNER JOIN
         ‘physionet-data.mimic_derived.first_day_vitalsign‘ vitals
         on
        adm.subject_id=vitals.subject_id
        and
        icu.stay_id=vitals.stay_id
    INNER JOIN
     ‘physionet-data.mimic_derived.first_day_lab‘ labs
         on
        adm.subject_id=labs.subject_id
        and
        icu.stay_id=labs.stay_id
    where icu_detail.los_icu > 1
        and pat.gender is not null
        and adm.ethnicity is not null
        and adm.ethnicity != "UNABLE TO OBTAIN"
        and adm.ethnicity != "UNKNOWN"


error term. First,

$\displaystyle\frac{1}{n}\big{[}(\mu_{y}-\hat{\mu}_{y})^{\top}(X-\mu_{x})\big{]}_{j}=\frac{1}{n}\sum_{k=1}^{2}\sum_{i\in\mathcal{I}_{k}}\big{(}\mu_{y}(Z_{i})-\hat{\mu}_{y}^{(k)}(Z_{i})\big{)}\big{(}X_{i,j}-\mu_{x,j}(Z_{i})\big{)}$

For each $k=1,2$, conditional on $\mathcal{I}\backslash\mathcal{I}_{k}$, the
terms in the summation
$(\mu_{y}(Z_{i})-\hat{\mu}_{y}^{(k)}(Z_{i}))(X_{i,j}-\mu_{x,j}(Z_{i}))$,
$i\in\mathcal{I}_{k}$ are mutually independent with mean zero, since

$\displaystyle\mathbb{E}\Big{[}\big{(}\mu_{y}(Z_{i})-\hat{\mu}_{y}^{(k)}(Z_{i})\big{)}\big{(}X_{i,j}-\mu_{x,j}(Z_{i})\big{)}\,\Big{|}\,\mathcal{I}\backslash\mathcal{I}_{k}\Big{]}$
$\displaystyle=\mathbb{E}\Big{[}\big{(}\mu_{y}(Z_{i})-\hat{\mu}_{y}^{(k)}(Z_{i})\big{)}\cdot\mathbb{E}\big{[}X_{i,j}-\mu_{x,j}(Z_{i})\,\big{|}\,Z_{i},\mathcal{I}\backslash\mathcal{I}_{k}\big{]}\,\Big{|}\,\mathcal{I}\backslash\mathcal{I}_{k}\Big{]}=0$

by the tower property. Also, the absolute value of each term is bounded below
$4c_{0}$ by the boundedness in Assumption 3.3. We now define the event

$\displaystyle\mathcal{E}_{\textrm{cr}}^{1,k}(\delta)=\Bigg{\\{}\Big{|}\sum_{i\in\mathcal{I}_{k}}\big{(}\mu_{y}(Z_{i})-\hat{\mu}_{y}^{(k)}(Z_{i})\big{)}\big{(}X_{i,j}-\mu_{x,j}(Z_{i})\big{)}\Big{|}\leq\max\Big{\\{}16c_{0}\log(2/\delta)/3,~{}2\hat{\epsilon}_{1}\sqrt{|\mathcal{I}_{k}|\log(2/\delta)}\Big{\\}}\Bigg{\\}},$

where we define

$\displaystyle\hat{\epsilon}_{1}^{2}=4\sup_{k=1,2}\big{\|}\mu_{y}-\hat{\mu}_{y}^{(k)}\big{\|}_{L_{2}(\mathbb{P}_{Z})}^{2}$
$\displaystyle\geq\mathbb{E}\Big{[}\big{(}\mu_{y}(Z_{i})-\hat{\mu}_{y}^{(k)}(Z_{i})\big{)}^{2}\big{(}X_{i,j}-\mu_{x,j}(Z_{i})\big{)}^{2}\,\Big{|}\,\mathcal{I}\backslash\mathcal{I}_{k}\Big{]}.$

The Bernstein’s inequality in Lemma C.1, implies
$\mathbb{P}\big{(}\mathcal{E}_{\textrm{cr}}^{1,k}\,\big{|}\,\mathcal{I}\backslash\mathcal{I}_{k}\big{)}\leq\delta$.
Marginalizing out the conditional probability and taking a union bound over
$k=1,2$, we know that $\mathbb{P}(\mathcal{E}_{\textrm{cr}}^{1}(\delta))\geq
1-\delta$ for any constant $\delta$, where we define
$\mathcal{E}_{\textrm{cr}}^{1}(\delta)=\mathcal{E}_{\textrm{cr}}^{1,1}(\delta/2)\cup\mathcal{E}_{\textrm{cr}}^{1,2}(\delta/2)$.
That is, with probability at least $1-\delta$,

$\displaystyle\Big{|}\frac{1}{n}\big{[}(\mu_{y}-\hat{\mu}_{y})^{\top}(X-\mu_{x})\big{]}_{j}\Big{|}\leq\max\bigg{\\{}\frac{16c_{0}\log(4/\delta)}{3n},~{}\hat{\epsilon}_{1}\sqrt{\frac{2\log(4/\delta)}{n}}\bigg{\\}}.$
(34)

For each fixed $j$, applying exactly the same arguments to

$\displaystyle\frac{1}{n}\big{[}(Y-\mu_{y})^{\top}(\mu_{x}-\hat{\mu}_{x})\big{]}_{j}=\frac{1}{n}\sum_{k=1}^{2}\sum_{i\in\mathcal{I}_{k}}\big{(}Y_{i}-\mu_{y}(Z_{i})\big{)}\big{(}\mu_{x,j}(Z_{i})-\hat{\mu}_{x,j}^{(k)}(Z_{i})\big{)},$

with probability at least $1-\delta$,

$\displaystyle\bigg{|}\frac{1}{n}\big{[}(Y-\mu_{y})^{\top}(\mu_{x}-\hat{\mu}_{x})\big{]}_{j}\bigg{|}\leq\max\bigg{\\{}\frac{16c_{0}\log(4/\delta)}{3n},~{}c_{0}\hat{\epsilon}_{2,j}\sqrt{\frac{2\log(4/\delta)}{n}}\bigg{\\}},$
(35)

where we define
$\hat{\epsilon}_{2,j}=2\sup_{k=1,2}\big{\|}\mu_{x,j}-\hat{\mu}_{x,j}^{(k)}\big{\|}_{L_{2}(\mathbb{P}_{Z})}.$
The third term is

$\displaystyle\frac{1}{n}\big{[}(\mu_{y}-\hat{\mu}_{y})^{\top}(\mu_{x}-\hat{\mu}_{x})\big{]}_{j}=\frac{1}{n}\sum_{k=1}^{2}\sum_{i\in\mathcal{I}_{k}}\big{(}\mu_{y}(Z_{i})-\hat{\mu}_{y}^{(k)}(Z_{i})\big{)}\big{(}\mu_{x,j}(Z_{i})-\hat{\mu}_{x,j}^{(k)}(Z_{i})\big{)}.$

Here for each $k=1,2$, conditional on $\mathcal{I}\backslash\mathcal{I}_{k}$,
each term
$(\mu_{y}(Z_{i})-\hat{\mu}_{y}^{(k)}(Z_{i}))(\mu_{x,j}(Z_{i})-\hat{\mu}_{x,j}^{(k)}(Z_{i}))$,
$i\in\mathcal{I}_{k}$ in the above summation is i.i.d. whose expectation is
bounded as

$\displaystyle\bigg{|}\mathbb{E}\Big{[}\big{(}\mu_{y}(Z_{i})-\hat{\mu}_{y}^{(k)}(Z_{i})\big{)}\big{(}\mu_{x,j}(Z_{i})-\hat{\mu}_{x,j}^{(k)}(Z_{i})\big{)}\,\Big{|}\,\mathcal{I}\backslash\mathcal{I}_{k}\Big{]}\bigg{|}\leq\big{\|}\mu_{y}-\hat{\mu}_{y}^{(k)}\big{\|}_{L_{2}(\mathbb{P}_{Z})}\big{\|}\mu_{x,j}-\hat{\mu}_{x,j}^{(k)}\big{\|}_{L_{2}(\mathbb{P}_{Z})}\leq\hat{\epsilon}_{1}\cdot\hat{\epsilon}_{2,j}/4.$

Meanwhile, their second moments are bounded as

$\displaystyle\mathbb{E}\Big{[}\big{(}\mu_{y}(Z_{i})-\hat{\mu}_{y}^{(k)}(Z_{i})\big{)}^{2}\big{(}\mu_{x,j}(Z_{i})-\hat{\mu}_{x,j}^{(k)}(Z_{i})\big{)}^{2}\,\Big{|}\,\mathcal{I}\backslash\mathcal{I}_{k}\Big{]}\leq
4\big{\|}\mu_{y}-\hat{\mu}_{y}^{(k)}\big{\|}_{L_{2}(\mathbb{P}_{Z})}^{2}\leq\hat{\epsilon}_{1}^{2}$

according to the bounededness in Assumption 3.3. Combining the above two
bounds and invoking the Bernstein’s inequliaty in Lemma C.1, it holds with
probability at least $1-\delta$ that

$\displaystyle\Big{|}\sum_{i\in\mathcal{I}_{k}}\big{(}\mu_{y}(Z_{i})-\hat{\mu}_{y}^{(k)}(Z_{i})\big{)}\big{(}\mu_{x,j}(Z_{i})-\hat{\mu}_{x,j}^{(k)}(Z_{i})\big{)}\Big{|}\leq|\mathcal{I}_{k}|\hat{\epsilon}_{1}\hat{\epsilon}_{2,j}/4+\max\Big{\\{}16c_{0}\log(2/\delta)/3,~{}2\hat{\epsilon}_{1}\sqrt{|\mathcal{I}_{k}|\log(2/\delta)}\Big{\\}}.$

Taking a union bound for $k=1,2$ implies that with probability at least
$1-\delta$,

$\displaystyle\Big{|}\frac{1}{n}\big{[}(\mu_{y}-\hat{\mu}_{y})^{\top}(\mu_{x}-\hat{\mu}_{x})\big{]}_{j}\Big{|}\leq\frac{\hat{\epsilon}_{1}\hat{\epsilon}_{2,j}}{4}+\max\bigg{\\{}\frac{16c_{0}\log(4/\delta)}{3n},~{}\hat{\epsilon}_{1}\sqrt{\frac{2\log(4/\delta)}{n}}\bigg{\\}}.$
(36)

Putting together (34), (35) and (36), we know that with probability at least
$1-\delta/2$,

$\displaystyle\big{\|}(\mu_{y}-\hat{\mu}_{y})^{\top}(X-\mu_{x})+(Y-\mu_{y})^{\top}(\mu_{x}-\hat{\mu}_{x})+(\mu_{y}-\hat{\mu}_{y})^{\top}(\mu_{x}-\hat{\mu}_{x})\big{\|}_{\infty}$
$\displaystyle\leq\frac{\hat{\epsilon}_{1}\hat{\epsilon}_{2,j}}{4}+2\max\bigg{\\{}\frac{16c_{0}\log(24/\delta)}{3n},~{}\hat{\epsilon}_{1}\sqrt{\frac{2\log(24/\delta)}{n}}\bigg{\\}}+\max\bigg{\\{}\frac{16c_{0}\log(24/\delta)}{3n},~{}c_{0}\hat{\epsilon}_{2,j}\sqrt{\frac{2\log(24/\delta)}{n}}\bigg{\\}}.$

Further taking a union bound over the above event and (33) for
$\mathcal{E}_{\textrm{est}}^{n,\delta/2}$, and using the fact that
$\max\\{a,b\\}\leq a+b$ for $a,b\geq 0$, we know it holds with probability at
least $1-\delta$ that

$\displaystyle\big{\|}(\mu_{y}-\hat{\mu}_{y})^{\top}(X-\mu_{x})+(Y-\mu_{y})^{\top}(\mu_{x}-\hat{\mu}_{x})+(\mu_{y}-\hat{\mu}_{y})^{\top}(\mu_{x}-\hat{\mu}_{x})\big{\|}_{\infty}$
$\displaystyle\leq\frac{4c_{n}^{2}(\log(3p_{x}/\delta))^{2}}{\sqrt{n}}+\frac{16c_{0}\log(24/\delta)}{n}+\frac{(2+c_{0})\log(3p_{x}/\delta)\sqrt{2\log(24/\delta)}}{n^{3/4}}.$

Recalling (B.2) and letting

$\displaystyle\bar{c}_{n}=4c_{n}^{2}+16c_{0}/\sqrt{n}+(2+c_{0})/n^{1/4},$

we have $\bar{c}_{n}\to 0$ as $n\to\infty$ and with probability at least
$1-\delta$,

$\displaystyle\frac{1}{n}\big{\|}(\theta^{*})^{\top}X^{\top}X-G\big{\|}_{\infty}\leq\frac{1}{n}\big{\|}(Y-\mu_{y})^{\top}(X-\mu_{x})+(X\theta^{*}-Y)^{\top}X\big{\|}_{\infty}+\frac{\bar{c}_{n}(\log(3p_{x}/\delta))^{2}}{\sqrt{n}}.$

Combining this high-probability bound with (31), we know that

$\displaystyle\mathbb{P}\bigg{(}\|\hat{\theta}-\theta^{*}\|_{2}\leq\frac{3\lambda\sqrt{k}}{2\zeta},\quad\forall~{}\lambda\geq\frac{2}{n}\big{\|}(Y-\mu_{y})^{\top}(X-\mu_{x})+(X\theta^{*}-Y)^{\top}X\big{\|}_{\infty}+\frac{2\bar{c}_{n}(\log(3p_{x}/\delta))^{2}}{\sqrt{n}}\bigg{)}$
$\displaystyle\geq\mathbb{P}\bigg{(}\frac{1}{n}\big{\|}(\theta^{*})^{\top}X^{\top}X-G\big{\|}_{\infty}\leq\frac{1}{n}\big{\|}(Y-\mu_{y})^{\top}(X-\mu_{x})+(X\theta^{*}-Y)^{\top}X\big{\|}_{\infty}+\frac{\bar{c}_{n}(\log(3p_{x}/\delta))^{2}}{\sqrt{n}}\bigg{)}\geq
1-\delta,$

which proves (15). ∎

##  Appendix C Supporting lemmas

######  Lemma C.1 (Bernstein’s inequality).

Suppose $X_{1},\dots,X_{n}$ are independent zero-mean random variables such
that $|X_{i}|\leq M$ almost surely for some constant $M>0$. Then for any
constant $t>0$,

$\displaystyle\mathbb{P}\bigg{(}\Big{|}\sum_{i=1}^{n}X_{i}\Big{|}\geq
t\bigg{)}\leq
2\exp\bigg{(}-\frac{t^{2}}{2\sum_{i=1}^{n}\mathbb{E}[X_{i}^{2}]+2Mt/3}\bigg{)}.$

That is, for any $\delta\in(0,1)$, with probability at least $1-\delta$, it
holds that

$\displaystyle\Big{|}\sum_{i=1}^{n}X_{i}\Big{|}\leq\max\Big{\\{}2\sqrt{{\textstyle\sum_{i=1}^{n}}\mathbb{E}[X_{i}^{2}]\log(2/\delta)},4M\log(2/\delta)/3\Big{\\}}.$

##  Appendix D Deferred simulation results

###  D.1 Low-dimensional simulation

Figure 12 plots the RMSE for estimating $\theta_{1}^{*}$ in settings 1 and 4
with sample size $n=1000$; this is in comparison to Figure 2. Figures 13 to 16
present the RMSE for all entries of $\theta^{*}$ in all settings.

Figure 12: Bias for estimating $\theta_{1}^{*}$ when $n=1000$. Details are
otherwise the same as Figure 2. Figure 13: RMSE for estimating
$\theta^{*}_{1}$. Details are otherwise the same as Figure 2. Figure 14: RMSE
for estimating $\theta^{*}_{2}$. Details are otherwise the same as Figure 2.
Figure 15: RMSE for estimating $\theta^{*}_{3}$. Details are otherwise the
same as Figure 2. Figure 16: RMSE for estimating $\theta^{*}_{4}$. Details are
otherwise the same as Figure 2.

###  D.2 High-dimensional simulation

Figure 17: Boxplot of standard deviations for all $\hat{\theta}_{j}$,
$j\in[p_{x}]$ over $N=1000$ replicates in setting 1 (left) and setting 2
(right). Other details of the plot are the same as Figure 6.