SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Limitations and future work. Two limitations of RBPT are that (i) the robustness of RBPT can lead to a more conservative test, as we have seen in the simulations; moreover, (ii) it requires the estimation of the conditional distribution  $P_{X|Z}$ , which can be challenging. To overcome the second problem, we introduced a variation of RBPT, named RBPT2, in which the Rao-Blackwellized predictor is obtained in a supervised fashion by fitting a second model  $\hat{h}: \mathcal{Z} \to \mathcal{Y}$  that predicts the outputs of the first model  $\hat{g}: \mathcal{X} \times \mathcal{Z} \to \mathcal{Y}$ . However, this solution only works if  $\hat{h}$  is better than  $\hat{g}$  in predicting  $Y$  under  $H_0$ , which ultimately depends on the model class for  $\hat{h}$  and how that model is trained. Future research directions may include (i) theoretically studying the power of RBPT in more detail and (ii) better understanding RBPT2 from a theoretical or methodological point of view, e.g., answering questions on how to choose and train the second model.
This paper is based upon work supported by the National Science Foundation (NSF) under grants no. 1916271, 2027737, 2113373, and 2113364.
<sup>10</sup>We run the experiment for 48 different random seeds and report the average Type-I error rate.
[38] Hao Zhang, Shuigeng Zhou, Jihong Guan, and Jun Huan. Measuring conditional independence by independent residuals for causal discovery. ACM Transactions on Intelligent Systems and Technology (TIST), 10(5):1-19, 2019.
[39] Hao Zhang, Shuigeng Zhou, Kun Zhang, and Jihong Guan. Residual similarity based conditional independence test and its application in causal discovery. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 5942-5949, 2022.
[40] Kun Zhang, Jonas Peters, Dominik Janzing, and Bernhard Scholkopf. Kernel-based conditional independence test and application in causal discovery. arXiv preprint arXiv:1202.3775, 2012.
[41] Lu Zhang and Lucas Janson. Floodgate: inference for model-free variable importance. arXiv preprint arXiv:2007.01283, 2020.
[42] Qinyi Zhang, Sarah Filippi, Seth Flaxman, and Dino Sejdinovic. Feature-to-feature regression for a two-step conditional independence test. In 33rd Conference on Uncertainty in Artificial Intelligence, UAI 2017. Association for Uncertainty in Artificial Intelligence (AUAI), 2017.
A Extra content
A.1 Misspecified inductive biases in modern statistics and machine learning
We present a toy experiment to empirically demonstrate how the training algorithm can prevent us from accurately estimating the Bayes predictor even when the model class is correctly specified, leading to invalid significance tests. We work in the context of a high-dimensional (overparameterized) regression with a training set of 250 observations and  $\geq 300$  covariates. We use the Significance Test of Feature Relevance test $^{11}$  (STFR) [7] to conduct the CI test. The data are generated as
$$
Z \sim N (0, I _ {3 0 0}), X \mid Z \sim N (\beta^ {\top} Z, 1), Y \mid X, Z \sim N (\beta^ {\top} Z, 1)
$$
where the first thirty entries of  $\beta$  are set to 1, and the remaining entries are zero. See that  $X \perp Y \mid Z$  and that  $Y$  is linearly related to  $Z$  and  $(X,Z)$ , and then the class of linear predictors is correctly specified when predicting  $Y$  from  $Z$  or  $(X,Z)$ . To perform the STFR test, we use LASSO  $(\| \cdot \|_1$  penalization term added to empirical squared error) and principal components regression (PCR) to train the linear predictors. Since  $\beta$
Figure 6: Type-I error rate is contingent on the training algorithm and not solely on the model classes. Unlike PCR, the LASSO fit provides the correct inductive bias in high-dimensional regression, controlling Type-I error.
is sparse, the LASSO fit provides the correct inductive bias while PCR leads to misspecification. We set the significance level to  $\alpha = 1\%$  and estimate the Type-I error rate for 100 different training sets. Figure 6 provides the Type-I error rate empirical distribution and illustrates that, despite using the same model class for both fitting methods, the training algorithm induces the wrong model in the PCR case, implying an invalid test most of the time.
A.2 Examples on when STFR fails
Examples A.1 and A.2 show simple situations in which Type-I error control is compromised or the conditional independence test has no power due to model misspecification. As we see in the next examples, Type-I error control is directly related to  $\mathcal{G}_2$  misspecification, while Type-II error minimization directly relates to  $\mathcal{G}_1$  misspecification.
Example A.1 (No Type-I error control). Suppose  $Y = Z + Z^2 + \varepsilon_y$  and  $X = Z^2 + \varepsilon_x$ , where  $\varepsilon_y, \varepsilon_x \sim N(0,1)$  are independent noise variables and  $Z$  has finite variance. Consequently,  $X \perp Y \mid Z$ . Let  $\mathcal{G}_1$  and  $\mathcal{G}_2$  be the classes of linear regressors with no intercept, i.e.,
$$
\mathcal {G} _ {1} = \left\{g _ {1} (x, z) = \beta_ {x} x + \beta_ {z} z: \beta_ {x}, \beta_ {z} \in \mathbb {R} \right\} a n d \mathcal {G} _ {2} = \left\{g _ {2} (z) = \beta_ {z} z: \beta_ {z} \in \mathbb {R} \right\}.
$$
If  $\ell$  denotes the mean squared error, we have that  $\mathbb{E}_P[\ell(g_{2,P}^*(Z), Y)] - \mathbb{E}_P[\ell(g_{1,P}^*(X,Z), Y)] > 0$  because the model class  $\mathcal{G}_2$  is misspecified, that is, it does not contain the Bayes predictor given by the conditional expectation  $\mathbb{E}[Y \mid Z]$ .
Example A.2 (Powerless test). Suppose  $Y = Z + \sin(X) + \varepsilon_y$ , where  $Z, X, \varepsilon_y \stackrel{iid}{\sim} N(0,1)$ . Consequently,  $X \not\perp Y \mid Z$ . Define  $\mathcal{G}_1$  and  $\mathcal{G}_2$  as in Example A.1. If  $\ell$  denotes the mean squared error, we have that  ${}^{12} \mathbb{E}_P[\ell(g_{2,P}^*(Z),Y)] - \mathbb{E}_P[\ell(g_{1,P}^*(X,Z),Y)] = 0$  even though the different Bayes' predictors have a difference in performance. This happens because the model  $g_{1,P}^*$  is not close enough to the Bayes predictor  $\mathbb{E}_P[Y \mid X, Z]$ .
A.3 Generalized Covariance Measure (GCM) test
In the GCM test proposed by Shah and Peters [31], the expected value of the conditional covariance between  $X$  and  $Y$  given  $Z$  is estimated and then tested to determine if it equals zero. To simplify the exposition, we consider  $X$  and  $Y$  univariate and work in a setup similar to the STFR's. If  $(X,Y,Z) \sim P$ , the GCM test relies on the observation that we can always write
$$
X = f _ {1, P} ^ {*} (Z) + \epsilon \text {a n d} Y = f _ {2, P} ^ {*} (Z) + \eta ,
$$
See Section 3.1 for more details
$^{12}$ Because  $\mathbb{E}_P[XY] = \mathbb{E}_P[X\mathbb{E}_P[Y|X]] = \mathbb{E}_P[X\cdot 0] = 0.$
where  $f_{1,P}^{*}(Z) = \mathbb{E}_{P}[X\mid Z]$  and  $f_{2,P}^{*}(Z) = \mathbb{E}_{P}[Y\mid Z]$  while the error terms  $\{\epsilon, \eta\}$  have zero mean when conditioned on  $Z$ . Consequently, we can write  $\mathbb{E}_{P}[\operatorname{Cov}_{P}(X,Y\mid Z)] = \mathbb{E}_{P}[\epsilon \eta]$ . To estimate  $\mathbb{E}_{P}[\operatorname{Cov}_{P}(X,Y\mid Z)]$ , we can first fit two models  $\hat{g}_1^{(m)}: \mathcal{Z} \to \mathcal{X}$  and  $\hat{g}_2^{(m)}: \mathcal{Z} \to \mathcal{Y}$ , that approximate  $f_{1,P}^{*}$  and  $f_{2,P}^{*}$ , using the training set  $\mathcal{D}_{tr}^{(m)}$ , and then compute an empirical version of  $\mathbb{E}_{P}[\epsilon \eta] = \mathbb{E}_{P}[(X - f_{1,P}^{*}(Z))(Y - f_{2,P}^{*}(Z))]$  using  $\hat{g}_1^{(m)}, \hat{g}_2^{(m)}$ , and  $\mathcal{D}_{te}^{(n)}$ .
In the GCM test, we reject  $H_0: X \perp Y \mid Z$  if the statistic  $\Gamma^{(n,m)} \triangleq |\sqrt{n} \bar{T}^{(n,m)} / \hat{\sigma}^{(n,m)}|$  exceeds  $\tau_{\alpha/2} \triangleq \Phi^{-1}(1 - \alpha/2)$ , depending on the test significance level  $\alpha \in (0,1)$ . Here,  $\bar{T}^{(n,m)}$  and  $\hat{\sigma}^{(n,m)}$  are defined as in 3.1 with  $T_i^{(m)} \triangleq (X_i - \hat{g}_1^{(m)}(Z_i))(Y_i - \hat{g}_2^{(m)}(Z_i))$ . If the  $p$ -value is defined as  $p(\mathcal{D}_{te}^{(n)}, \mathcal{D}_{tr}^{(m)}) = 2(1 - \Phi(\Gamma^{(n,m)}))$ , the test  $\varphi_{\alpha}^{\mathrm{GCM}}(\mathcal{D}_{te}^{(n)}, \mathcal{D}_{tr}^{(m)})$  is analogously given by Equation 3.2. Like the STFR, the GCM test depends on the models' classes and implicitly on the training algorithm. If the limiting models  $g_{1,P}^*$  and  $g_{2,P}^*$  are not  $f_{1,P}^*$  and  $f_{2,P}^*$ , then Type-I error control is not guaranteed.
We introduce definitions and assumptions. Assumption A.3 gives a rate of convergence for the models  $\hat{g}_j^{(m)}$  in the mean squared error sense. Definition A.4 gives a definition for the misspecification gaps.
Assumption A.3. There are functions  $g_{1,P}^{*}, g_{2,P}^{*}$ , and a constant  $\gamma > 0$  such that
$$
\mathbb {E} _ {P} \left[ \left(\hat {g} _ {j} ^ {(m)} (Z) - g _ {j, P} ^ {*} (Z)\right) ^ {2} \mid \mathcal {D} _ {t r} ^ {(m)} \right] = \mathcal {O} _ {\mathcal {P}} \left(m ^ {- \gamma}\right), f o r j = 1, 2
$$
Definition A.4. For each  $j\in \{1,2\}$ , define the misspecification gap as  $\delta_{j,P}\triangleq g_{j,P}^{*} - f_{j,P}^{*}$ .
In the next result, we approximate GCM test Type-I error rate and power using the gaps in Definition A.4 and Assumptions A.3, 3.2, and 3.3 applied to this context.
Theorem A.5. Suppose that Assumptions 3.2, 3.3, and A.3 hold. If  $n$  is a function of  $m$  such that  $n \to \infty$  and  $n = o(m^{\gamma})$  as  $m \to \infty$ , then
$$
\mathbb {E} _ {P} \left[ \varphi_ {\alpha} ^ {\mathrm {G C M}} \left(\mathcal {D} _ {t e} ^ {(n)}, \mathcal {D} _ {t r} ^ {(m)}\right) \right] = 1 - \Phi \left(\tau_ {\alpha / 2} - \sqrt {\frac {n}{\sigma_ {P} ^ {2}}} \Omega_ {P} ^ {\mathrm {G C M}}\right) + \Phi \left(- \tau_ {\alpha / 2} - \sqrt {\frac {n}{\sigma_ {P} ^ {2}}} \Omega_ {P} ^ {\mathrm {G C M}}\right) + o (1)
$$
$$
\Omega_ {P} ^ {\mathrm {G C M}} \triangleq \mathbb {E} _ {P} [ \operatorname {C o v} _ {P} (X, Y \mid Z) ] + \mathbb {E} _ {P} [ \delta_ {1, P} (Z) \delta_ {2, P} (Z) ]
$$
From Theorem A.5, it is possible to verify that if  $\delta_{j,P}(Z)$  is zero for at least one  $j\in \{1,2\}$ , i.e., if at least one model converges to the conditional expectation, the GCM test asymptotically controls Type-I error. This can be seen as a double-robustness property of the GCM, which is not present<sup>13</sup> in Shah and Peters [31]. If  $\mathbb{E}_P[\delta_{1,P}(Z)\delta_{2,P}(Z)]\neq 0$ , then  $\mathbb{E}_P[\varphi_\alpha^{\mathrm{GCM}}(\mathcal{D}_{te}^{(n)},\mathcal{D}_{tr}^{(m)})]\to 1$  as  $m\to \infty$  even when  $H_0:X\perp Y\mid Z$ . Under the alternative, if  $\Omega_P^{\mathrm{GCM}}\neq 0$ , Type-II error approaches 0 asymptotically.
A.4 REgression with Subsequent Independence Test (RESIT)
As revisited by Zhang et al. [42], the idea behind RESIT is to first residualize  $Y$  and  $X$  given  $Z$  and then test dependence between the residuals. It is similar to GCM, but requires the error terms and  $Z$  to be independent. When that assumption is reasonable, one advantage of RESIT over GCM is that it has power against a broader set of alternatives. In this section, we use a permutation test [18, Example 15.2.3] to assess the independence of residuals. We analyse RESIT's Type-I error control.
If  $(X,Y,Z)\sim P$  and  $(X,Y)$  can be modeled as an additive noise model (ANM), that is, we can write
$$
X = f _ {1, P} ^ {*} (Z) + \epsilon \text {a n d} Y = f _ {2, P} ^ {*} (Z) + \eta , \tag {A.1}
$$
where  $f_{1,P}^{*}(Z) = \mathbb{E}_{P}[X\mid Z],f_{2,P}^{*}(Z) = \mathbb{E}_{P}[Y\mid Z]$ , and the error terms  $(\epsilon ,\eta)$  are independent of  $Z$  it is possible to show that  $X\perp Y\mid Z\Leftrightarrow \epsilon \perp \eta$ . To facilitate our analysis<sup>14</sup>, we consider first fitting two models  $\hat{g}_1^{(m)}$  and  $\hat{g}_2^{(m)}$  that approximate  $f_{1,P}^{*}$  and  $f_{2,P}^{*}$  using the training set  $\mathcal{D}_{tr}^{(m)}$  and then test the independence of the residuals<sup>15</sup>  $\hat{\epsilon}_i = X_i - \hat{g}_1^{(m)}(Z_i)$  and  $\hat{\eta}_i = Y_i - \hat{g}_2^{(m)}(Z_i)$  using the test set  $\mathcal{D}_{te}^{(n)}$ . Define (i)  $(\hat{\epsilon},\hat{\eta})\triangleq \{(\hat{\epsilon}_i,\hat{\eta}_i)\}_{i = 1}^n$  (test set residuals vertically stacked in matrix form) and (ii)  $(\hat{\epsilon},\hat{\eta})^{(b)}$
This property is clear in our result because we consider data splitting.
In practice, data splitting is not necessary. However, this procedure helps when theoretically analyzing the method.
<sup>15</sup>We omit the residuals superscript to ease notation.
as one of the  $B$  permutations, i.e., consider that we fix  $\hat{\epsilon}$  and permute  $\hat{\eta}$  row-wise. Let  $\Psi$  be a test statistic and  $\Psi((\hat{\epsilon}, \hat{\eta}))$  and  $\Psi((\hat{\epsilon}, \hat{\eta})^{(b)})$  its evaluation on the original residuals and the  $b$ -the permuted set. If the permutation  $p$ -value is given by
$$
p \left(\mathcal {D} _ {t e} ^ {(n)}, \mathcal {D} _ {t r} ^ {(m)}\right) = \frac {1 + \sum_ {b = 1} ^ {B} \mathbb {1} \left[ \Psi \left(\left(\hat {\epsilon} , \hat {\eta}\right) ^ {(b)}\right) \geq \Psi \left(\left(\hat {\epsilon} , \hat {\eta}\right)\right) \right]}{1 + B} \tag {A.2}
$$
a test  $\varphi_{\alpha}^{\mathrm{RESIT}}$  aiming level  $\alpha \in (0,1)$  is given by Equation 3.2.
Similarly to STFR and GCM, we consider  $g_{1,P}^{*}$  and  $g_{2,P}^{*}$  to be the limiting models for  $\hat{g}_1^{(m)}$  and  $\hat{g}_2^{(m)}$ . Different from GCM, both models  $g_{1,P}^{*}$  and  $g_{2,P}^{*}$  are multi-output since  $X$  and  $Y$  are not necessarily univariate.
Now, we introduce some assumptions before we present our result for RESIT. Assumption A.6 gives a rate of convergence for the models  $\hat{g}_j^{(m)}$  in the mean squared error sense.
Assumption A.6. There are models  $g_{1,P}^{*}$ ,  $g_{2,P}^{*}$ , and a constant  $\gamma > 0$  such that
$$
\mathbb {E} _ {P} \left[ \left\| \hat {g} _ {j} ^ {(m)} (Z) - g _ {j, P} ^ {*} (Z) \right\| _ {2} ^ {2} \middle | \mathcal {D} _ {t r} ^ {(m)} \right] = \mathcal {O} _ {\mathcal {P} _ {0}} (m ^ {- \gamma}), j = 1, 2
$$
Assumption A.7 puts more structure on the distributions of the error terms  $(\epsilon, \eta)$  and is a mild assumption.
Assumption A.7. Assume that for all  $P \in \mathcal{P}_0$ , the distribution of  $(\epsilon, \eta)$ ,  $P_{\epsilon, \eta}$ , is absolutely continuous with respect to the Lebesgue measure in  $\mathbb{R}^{d_X} \times \mathbb{R}^{d_Y}$  with  $L$ -Lipschitz density  $p_{\epsilon, \eta}$  for a certain  $L > 0$ . That is, for any  $e_1, e_2 \in \mathbb{R}^{d_X}$  and  $h_1, h_2 \in \mathbb{R}^{d_Y}$ , we have
$$
\left| p _ {\epsilon , \eta} \left(e _ {1}, h _ {1}\right) - p _ {\epsilon , \eta} \left(e _ {2}, h _ {2}\right) \right| \leq L \left\| \left(e _ {1}, h _ {1}\right) - \left(e _ {2}, h _ {2}\right) \right\| _ {2}
$$
We assume that  $L$  does not depend on  $P$ .
Assumption A.8 states that some of the variables we work with are uniformly almost surely bounded over all  $P \in \mathcal{P}_0$ . This assumption is realistic in most practical cases.
Assumption A.8. There is bounded Borel set  $A \in \mathcal{B}(\mathbb{R}^{d_X \times d_Y})$  such that
$$
\inf  _ {P \in \mathcal {P} _ {0}} \mathbb {P} _ {P} \left(\left(X, Y\right) \in A\right) = 1
$$
$$
\inf  _ {P \in \mathcal {P} _ {0}} \inf  _ {g _ {1}, g _ {2}} \mathbb {P} _ {P} \left(\left(g _ {1} (Z), g _ {2} (Z)\right) \in A\right) = 1
$$
Here,  $\inf_{g_1,g_2}$  is taken over the model classes we consider (if the model classes vary with  $m$ , consider the union of model classes). In the following, we present the result for RESIT. For that result, let: (i)  $\epsilon^{*} \triangleq \epsilon - \delta_{1,P}(Z_i)$  and  $\eta^{*} \triangleq \eta - \delta_{2,P}(Z_i)$ , where the misspecification gaps are given as in Definition A.4; (ii)  $d_{\mathrm{TV}}$  represent the total variation (TV) distance between two probability distributions [36]; and (iii) the superscript  $n$ , e.g., in  $P_{\epsilon^{*},\eta^{*}}^{n}$ , represent a product measure.
Theorem A.9. Under Assumptions A.6, A.7, and A.8, if  $H_0: X \perp Y \mid Z$  holds and  $n$  is a function of  $m$  such that  $n \to \infty$  and  $n = o(m^{\frac{\gamma}{2}})$  as  $m \to \infty$ , then
$$
\mathbb {E} _ {P} [ \varphi_ {\alpha} ^ {\mathrm {R E S I T}} (\mathcal {D} _ {t e} ^ {(n)}, \mathcal {D} _ {t r} ^ {(m)}) ] \leq \alpha + \min \{d _ {T V} (P _ {\epsilon^ {*}, \eta^ {*}} ^ {n}, P _ {\epsilon , \eta^ {*}} ^ {n}), d _ {T V} (P _ {\epsilon^ {*}, \eta^ {*}} ^ {n}, P _ {\epsilon^ {*}, \eta} ^ {n}) \} + o (1)
$$
where  $o(1)$  denotes uniform convergence over all  $P \in \mathcal{P}_0$  as  $m \to \infty$ .
From Theorem A.9, we can see that if at least one of the misspecification gaps  $\delta_{1,P}$  or  $\delta_{2,P}$  is null, then  $\mathbb{E}_P[\varphi_\alpha^{\mathrm{RESIT}}(\mathcal{D}_{te}^{(n)},\mathcal{D}_{tr}^{(m)})]\leq \alpha$  under  $H_0:X\perp Y\mid Z$ . This can be seen as a double-robustness property of the RESIT. If none of the misspecification gaps are null, we do not have any guarantees on RESIT's Type-I error control. It can be shown that the proposed upper bound converges to 1.
A.5 RBPT extra derivation
Let  $\mu_{X}$  be a dominating measure of  $Q_{X|Z}^{*}$  and  $P_{X|Z}$  that does not depend on  $Z$ . Let  $q_{X|Z}^{*}$  (resp.  $p_{X|Z}$ ) be  $Q_{X|Z}^{*}$  (resp.  $P_{X|Z}$ ) density with respect to  $\mu_{X}$ .
$$
\begin{array}{l} \Omega_ {P, 1} ^ {\mathrm {R B P T}} = \mathbb {E} _ {P} \left[ \ell \left(\int g _ {P} ^ {*} (x, Z) d Q _ {X | Z} ^ {*} (x), Y\right) - \ell \left(\int g _ {P} ^ {*} (x, Z) d P _ {X | Z} (x), Y\right) \right] \\ \leq L \cdot \mathbb {E} _ {P} \left\| \int g _ {P} ^ {*} (x, Z) d Q _ {X | Z} ^ {*} (x) - \int g _ {P} ^ {*} (x, Z) d P _ {X | Z} (x) \right\| _ {2} \\ \leq L \cdot \mathbb {E} _ {P} \left\| \int g _ {P} ^ {*} (x, Z) d Q _ {X | Z} ^ {*} (x) - \int g _ {P} ^ {*} (x, Z) d P _ {X | Z} (x) \right\| _ {1} \\ = L \cdot \mathbb {E} _ {P} \left\| \int g _ {P} ^ {*} (x, Z) \left(q _ {X | Z} ^ {*} (x | Z) - p _ {X | Z} (x | Z)\right) d \mu_ {X} (x) \right\| _ {1} \\ \leq L \cdot \mathbb {E} _ {P} \int \left\| g _ {P} ^ {*} (x, Z) \right\| _ {1} \left| q _ {X | Z} ^ {*} (x | Z) - p _ {X | Z} (x | Z) \right| d \mu_ {X} (x) \\ \leq M L \cdot \mathbb {E} _ {P} \int \left| q _ {X | Z} ^ {*} (x | Z) - p _ {X | Z} (x | Z) \right| d \mu_ {X} (x) \\ = 2 M L \cdot \mathbb {E} _ {P} \left[ d _ {\mathrm {T V}} \left(Q _ {X | Z} ^ {*}, P _ {X | Z}\right) \right] \\ \end{array}
$$
$$
\mathbb {E} _ {P} \left[ d _ {\mathrm {T V}} \left(Q _ {X \mid Z} ^ {*}, P _ {X \mid Z}\right) \right] \leq \Omega_ {P, 2} ^ {\mathrm {R B P T}} / (2 M L)
$$
$$
\Omega_ {P} ^ {\mathrm {R B P T}} \leq 0
$$
A.6 How to obtain a result for RBPT2?
We informally give some ideas on extending Theorem 4.4. Theorem 4.4 states that
where  $\Omega_P^{\mathrm{RBPT}} = \Omega_{P,1}^{\mathrm{RBPT}} - \Omega_{P,2}^{\mathrm{RBPT}}$  with
$$
\Omega_ {P, 1} ^ {\mathrm {R B P T}} \triangleq \mathbb {E} _ {P} \left[ \ell \left(\int g _ {P} ^ {*} (x, Z) d Q _ {X | Z} ^ {*} (x), Y\right) \right] - \mathbb {E} _ {P} \left[ \ell \left(\int g _ {P} ^ {*} (x, Z) d P _ {X | Z} (x), Y\right) \right]
$$
$$
\Omega_ {P, 2} ^ {\mathrm {R B P T}} \triangleq \mathbb {E} _ {P} \left[ \ell \left(g _ {P} ^ {*} (X, Z), Y\right) \right] - \mathbb {E} _ {P} \left[ \ell \left(\int g _ {P} ^ {*} (x, Z) d P _ {X | Z} (x), Y\right) \right]
$$
If we wanted to adapt that result for RBPT2, we would have to redefine  $\Omega_{P,1}^{\mathrm{RBPT}}$ . The analogue of  $\Omega_{P,1}^{\mathrm{RBPT}}$  for RBPT2 would be
$$
\Omega_ {P, 1} ^ {\mathrm {R B P T 2}} \triangleq \mathbb {E} _ {P} \left[ \ell \left(\tilde {\mathbb {E}} _ {P} \left[ g _ {P} ^ {*} (X, Z) \mid Z \right], Y\right) \right] - \mathbb {E} _ {P} \left[ \ell \left(\mathbb {E} _ {P} \left[ g _ {P} ^ {*} (X, Z) \mid Z \right], Y\right) \right]
$$
where  $\tilde{\mathbb{E}}_P[g_P^*(X,Z) \mid Z = z]$  denotes the limiting regression model to predict  $g_P^*(X,Z)$  given  $Z$ . If we assume the existence of a big unlabeled dataset, deriving a result might be easier as we can avoid the asymptotic details on the convergence of the second regression model by assuming that the limiting Rao-Blackwellization model, for a fixed initial predictor, is known. The only challenge is proving the convergence of  $\tilde{\mathbb{E}}_P[\hat{g}(X,Z) \mid Z]$  to  $\tilde{\mathbb{E}}_P[g_P^*(X,Z) \mid Z]$ .
B Technical proofs
B.1 STFR
Lemma B.1. Assume we are under the conditions of Theorem 3.4. Then:
$$
\left(\hat {\sigma} ^ {(n, m)}\right) ^ {2} - \operatorname {V a r} _ {P} \left[ T _ {1} ^ {(m)} \mid \mathcal {D} _ {t r} ^ {(m)} \right] = o _ {\mathcal {P}} (1) a s m \rightarrow \infty
$$
Proof. First, see that for an arbitrary  $\varepsilon > 0$ , there must be a sequence of probability measures in  $\mathcal{P}$ ,  $(P^{(m)})_{m \in \mathbb{N}}$ , such that
$$
\sup _ {P \in \mathcal {P}} \mathbb {P} _ {P} \left[ \left| (\hat {\sigma} ^ {(n, m)}) ^ {2} - \mathsf {V a r} _ {P} [ T _ {1} ^ {(m)} \mid \mathcal {D} _ {t r} ^ {(m)} ] \right| > \varepsilon \right] \leq \mathbb {P} _ {P (m)} \left[ \left| (\hat {\sigma} ^ {(n, m)}) ^ {2} - \mathsf {V a r} _ {P (m)} [ T _ {1} ^ {(m)} \mid \mathcal {D} _ {t r} ^ {(m)} ] \right| > \varepsilon \right] + \frac {1}{m}
$$
Pick one of such sequences. Then, to prove that  $(\hat{\sigma}^{(n,m)})^2 -\mathsf{Var}_P[T_1^{(m)}\mid \mathcal{D}_{tr}^{(m)}] = o_{\mathcal{P}}(1)$  as  $m\to \infty$  it suffices to show that
$$
\mathbb {P} _ {P ^ {(m)}} \left[ \left| (\hat {\sigma} ^ {(n, m)}) ^ {2} - \mathsf {V a r} _ {P ^ {(m)}} [ T _ {1} ^ {(m)} \mid \mathcal {D} _ {t r} ^ {(m)} ] \right| > \varepsilon \right] \to 0 \mathrm {a s} m \to \infty
$$
Now, expanding  $(\hat{\sigma}^{(n,m)})^2$  we get
$$
\begin{array}{l} (\hat {\sigma} ^ {(n, m)}) ^ {2} = \frac {1}{n} \sum_ {i = 1} ^ {n} (T _ {i} ^ {(m)}) ^ {2} - \left(\frac {1}{n} \sum_ {i = 1} ^ {n} T _ {i} ^ {(m)}\right) ^ {2} \\ = \frac {1}{n} \sum_ {i = 1} ^ {n} \left[ (T _ {i} ^ {(m)}) ^ {2} - \mathbb {E} _ {P ^ {(m)}} [ (T _ {1} ^ {(m)}) ^ {2} \mid \mathcal {D} _ {t r} ^ {(m)} ] \right] + \mathbb {E} _ {P ^ {(m)}} [ (T _ {1} ^ {(m)}) ^ {2} \mid \mathcal {D} _ {t r} ^ {(m)} ] - \left(\frac {1}{n} \sum_ {i = 1} ^ {n} T _ {i} ^ {(m)}\right) ^ {2} \\ = \frac {1}{n} \sum_ {i = 1} ^ {n} \left[ (T _ {i} ^ {(m)}) ^ {2} - \mathbb {E} _ {P ^ {(m)}} [ (T _ {1} ^ {(m)}) ^ {2} \mid \mathcal {D} _ {t r} ^ {(m)} ] \right] - \left(\frac {1}{n} \sum_ {i = 1} ^ {n} T _ {i} ^ {(m)} - \mathbb {E} _ {P ^ {(m)}} [ T _ {1} ^ {(m)} \mid \mathcal {D} _ {t r} ^ {(m)} ]\right) ^ {2} \\ - 2 \left(\frac {1}{n} \sum_ {i = 1} ^ {n} T _ {i} ^ {(m)} \mathbb {E} _ {P ^ {(m)}} [ T _ {1} ^ {(m)} \mid \mathcal {D} _ {t r} ^ {(m)} ] - (\mathbb {E} _ {P ^ {(m)}} [ T _ {1} ^ {(m)} \mid \mathcal {D} _ {t r} ^ {(m)} ]) ^ {2}\right) + \mathsf {V a r} _ {P ^ {(m)}} [ T _ {1} ^ {(m)} \mid \mathcal {D} _ {t r} ^ {(m)} ] \\ = \frac {1}{n} \sum_ {i = 1} ^ {n} \left[ (T _ {i} ^ {(m)}) ^ {2} - \mathbb {E} _ {P ^ {(m)}} [ (T _ {1} ^ {(m)}) ^ {2} \mid \mathcal {D} _ {t r} ^ {(m)} ] \right] - \left(\frac {1}{n} \sum_ {i = 1} ^ {n} T _ {i} ^ {(m)} - \mathbb {E} _ {P ^ {(m)}} [ T _ {1} ^ {(m)} \mid \mathcal {D} _ {t r} ^ {(m)} ]\right) ^ {2} \\ - 2 \mathbb {E} _ {P ^ {(m)}} [ T _ {1} ^ {(m)} \mid \mathcal {D} _ {t r} ^ {(m)} ] \left(\frac {1}{n} \sum_ {i = 1} ^ {n} T _ {i} ^ {(m)} - \mathbb {E} _ {P ^ {(m)}} [ T _ {1} ^ {(m)} \mid \mathcal {D} _ {t r} ^ {(m)} ]\right) + \mathsf {V a r} _ {P ^ {(m)}} [ T _ {1} ^ {(m)} \mid \mathcal {D} _ {t r} ^ {(m)} ] \\ \end{array}
$$