SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Theorem 2 (Informative bounds with our CSM). Let the continuous observational distribution  $\mathbb{P}(Y\mid a)$  be induced by some SCM of class  $\mathfrak{B}(C^2,d),d\geq 2$  , which satisfies Assumption  $\kappa$  .Let  $\mathbb{P}(Y\mid a)$  have a compact support  $\mathcal{V}_a = [l_a,u_a]$  . Then, the ignorance interval for the partial identification of ECOU [ECOT] of class  $\mathfrak{B}(C^2,d)$  has informative bounds, dependent on  $\kappa$  and  $d,$  which are given by  $\underline{Q_{a^{\prime}\to a}(y^{\prime})} = l(\kappa ,d) > l_{a}$  and  $\overline{Q_{a^{\prime}\to a}(y^{\prime})} = u(\kappa ,d) <   u_{a}$
Proof. We will now show that ECOU [ECOT] always have informative bounds for every possible SCM satisfying the assumptions of the theorem. Without the loss of generality, let us consider the lower bound, i.e.,
$$
\underline {{Q _ {a ^ {\prime} \rightarrow a} (y ^ {\prime})}} = \inf  _ {\mathcal {M} \in \mathfrak {B} \left(C ^ {k}, d\right)} Q _ {a ^ {\prime} \rightarrow a} ^ {\mathcal {M}} \left(y ^ {\prime}\right) \quad \text {s . t .} \forall a \in \{0, 1 \}: \mathbb {P} (Y \mid a) = \mathbb {P} ^ {\mathcal {M}} (Y \mid a). \tag {91}
$$
The proof contains three steps. First, we show that the level sets of  $C^2$  functions with compact support consist of the countable number of connected components, each with the finite Hausdorff
measure. Each connected component of almost every level set is thus a  $d - 1$ -dimensional Riemannian manifold. Second, we demonstrate that, under Assumption  $\kappa$ , almost all the level set bundles have a nested structure, are diffeomorphic to each other, have boundaries, and their boundaries always coincide with the boundary of the unit hypercube  $[0,1]^d$ . Third, we arrive at a contradiction in that, to obtain non-informative bounds, we have to fit a  $d$ -ball of arbitrarily small radius to the interior of the counterfactual level set (which is not possible with bounded absolute principal curvature). Therefore, CSM with Assumption  $\kappa$  has informative bounds.
Step 1. The structure of the level sets of  $C^2$  functions with the compact support (or more generally, Lipschitz functions) were studied in [1]. We employ Theorem 2.5 from [1], which is a result of Sard's theorem. Specifically, the level sets of  $C^2$  functions with the compact support consist of the countable number of connected components, each with the finite Hausdorff measure (surface volume), namely,  $\mathcal{H}^{d - 1}(E(y,a)) < \infty$  for  $a\in \{0,1\}$ . In Example 3, we provide an SCM with the level sets with indefinitely many connected components. Furthermore, connected components of almost every level set are  $(d - 1)$ -dimensional Riemannian manifolds of class  $C^2$ . Some points  $y$  have  $(d - 2)$ - and lower dimensional manifolds, but their probability measure is zero (e.g., the level sets of the bounds of the support,  $E(l_{a},a)$  and  $E(u_{a},a)$ ).
Step 2. The Assumption  $\kappa$  assumes the existence of the principal curvatures at every point of the space (except for the boundaries of the unit hypercube  $[0,1]^d$ ), for both the factual and the counterfactual function. Thus, the functions of the SCMs are regular, as otherwise the principal curvatures would not be defined at the critical points (see Eq. (17)). As a result of the fundamental Morse theorem (see Appendix B), almost all the level set bundles are nested and diffeomorphic to each other. Specifically, the level set bundles  $\{E(y,a):y\in (l_a,u_a)\}$  for  $a\in \{0,1\}$ . The latter exclude the level sets of the bounds of the  $\mathcal{V}_a$ , i.e.,  $E(l_{a},a)$  and  $E(u_{a},a)$  for  $a\in \{0,1\}$ , which in turn are laying completely in the boundary of the unit hypercube  $[0,1]^d$ . Another important property is that all the level sets have a boundary (otherwise, due to the diffeomorphism, the critical point would exist in their interior) and this boundary lies at the boundary of the unit hypercube  $[0,1]^d$ .
Step 3. Let us fix the factual level set,  $E(y', a')$ . By the observational distribution constraint, we know that the interior occupies  $\mathbb{F}_{a'}(y')$  of the total volume of the unit hypercube  $[0,1]^d$ , where  $\mathbb{F}_{a'}(\cdot)$  is the CDF of the observational distribution. Also, its absolute principal curvatures are bounded by  $\kappa$ .
Then, let us assume there exists a counterfactual function, which obtains non-informative bounds of ECOU [ECOT]. For this, a counterfactual level set for some  $y$  close to  $l_{a}$  has to contain exactly  $F_{a}(y)$  volume in the interior. At the same time, this level set has to contain as much of the surface volume of the factual level set,  $\mathcal{H}^{d - 1}(E(y',a'))$ , as possible in its interior. This is only possible by "bending" the counterfactual level set along one of the directions, so that the boundaries of the counterfactual level set lay at the boundaries of the unit hypercube,  $[0,1]^d$ . Due to the bound on the maximal absolute principal curvature, we have to be able to fully contain a  $d$ -ball with radius  $\frac{1}{\kappa}$  inside the interior of the counterfactual level set. This  $d$ -ball occupies the volume
$$
\operatorname {v o l} ^ {d} (d \text {- b a l l}) = \frac {\pi^ {n / 2}}{\Gamma \left(\frac {d}{2} + 1\right)} \frac {1}{\kappa^ {d}}, \tag {92}
$$
where  $\Gamma(\cdot)$  is a Gamma function. Therefore, the volume of the interior of the counterfactual level set has to be at the same time arbitrarily close to zero, as  $F_{a}(y) \to 0$  as  $y \to l_{a}$ , and at least of the volume of the  $d$ -ball with radius  $\frac{1}{\kappa}$ . Thus, we arrived at the contradiction, which proves the theorem.
Corollary 4 (Monotonicity wrt. to  $\kappa$ ). By construction, it is guaranteed that, with decreasing  $\kappa$ , the bounds for ECOU [ECOT] could only shrink, i.e., for  $\kappa_{1} > \kappa_{2}: l(\kappa_{1}, d) \leq l(\kappa_{2}, d)$  and  $u(\kappa_{1}, d) \geq u(\kappa_{2}, d)$ . This happens, as we decrease the class of functions in the constraints of the optimization problem from Definition 2.
Corollary 5 (Sufficiency of  $d = 2$ ). In general, upper and lower bounds under assumed CSM ( $\kappa$ ) are different for different  $d_1, d_2 \geq 2$ , e.g.,  $l(\kappa, d_1) \leq l(\kappa, d_2)$  and  $u(\kappa, d_1) \geq u(\kappa, d_2)$  for  $d_1 > d_2$ . This follows from the properties of the high-dimensional spaces, e.g., we one can fit more  $d$ -balls of the fixed radius to the hypercube  $[0,1]^d$  with the increase of  $d$ . Nevertheless, without the loss of generality and in the absence of the information on the dimensionality, we can set  $d = 2$ , as we still cover the entire identifiability spectrum by varying  $\kappa$ .
$$
\underline {{Q _ {a ^ {\prime} \rightarrow a} (y ^ {\prime})}} / \overline {{Q _ {a ^ {\prime} \rightarrow a} (y ^ {\prime})}} = \min  / \max  \left\{B G M _ {+} \left(y ^ {\prime}\right), B G M _ {-} \left(y ^ {\prime}\right), E Q T D _ {l} \left(y ^ {\prime}\right), E Q T D _ {u} \left(y ^ {\prime}\right)\right\} \tag {7}
$$
Proof. We provide an informal proof for  $d = 2$  based on geometric intuition. When  $\kappa = 0$ , all the level sets are straight lines (hyperplanes for  $d > 2$ ). It is possible to characterize all the possible relative arrangements of the factual level set and the bundle of counterfactual level sets, and the corresponding values of ECOU [ECOT]. For a particular arrangement, displayed in Fig. 11, we see that the area of the interior of the factual level set is  $\mathrm{vol}^2 = \mathbb{F}_{a'}(y')$ . At the same time, the factual level set crosses the counterfactual bundle corresponding to the interval  $[l_a, \mathbb{F}_a^{-1}(2\mathbb{F}_{a'}(y'))]$ , thus the counterfactual distribution of [un]treated is simply a  $2\mathbb{F}_{a'}(y')$ -quantile-truncated distribution, i.e.,  $\mathbb{P}(Y \mid a, Y < \mathbb{F}_a^{-1}(2\mathbb{F}_{a'}(y)))$ . Therefore, the bounds of the ignorance interval of ECOU [ECOT] are defined by the min/max of both BGMs and expectations of quantile-truncated distributions (EQTDs).
Figure 11: Factual level set (orange), counterfactual level sets (blue).
E Discussion of extensions and limitations
Sharp bounds. To obtain the sharp bounds under CSM  $(\kappa)$ , one has to exactly solve the constrained variational problem task formulated in the Definition 2 in the space of the SCMs, which satisfy the Assumption  $\kappa$ . This is a very non-trivial task, the distributional constraint and the constraint of level sets of bounded curvature both include the partial differential equation with the Hausdorff integrals. For this task, in general, an explicit solution does not exist in a closed form (unlike, e. g., the solution of the marginal sensitivity model [126]).
Extension of CSM to discrete treatments and continuous covariates. Our CSM  $(\kappa)$  naturally scales to bivariate SCMs with categorical treatments,  $A \in \{0, \dots, K\}$ ; and multivariate Markovian SCMs with additional covariates,  $X \in \mathbb{R}^m$ , which are all predecessors of  $Y$ . These extensions, nevertheless, bring additional challenges from estimating  $\mathbb{P}(Y \mid A, X)$  from the observational (interventional) data. Hence, additional smoothness assumptions are required during modeling.
Extension of CSM to semi-Markovian SCMs. CSM  $(\kappa)$  is limited to the Markovian SCMs. For semi-Markovian SCMs, e. g., a setting of the potential outcomes framework [111], additional assumptions are needed. Specifically, in semi-Markovian SCMs, the latent noise variables could be shared for  $X$  and  $Y$ , and this complicates counterfactual inference. Future work may consider incorporating other sensitivity models to our CSM  $(\kappa)$ , e. g., marginal sensitivity model, MSM  $(\Gamma)$  [126], as described in Fig. 12.
Figure 12: Possible combination of our CSM  $(\kappa)$  with the marginal sensitivity model, MSM  $(\Gamma)$  [126], for the potential outcome framework [111].
F Details on architecture and inference of Augmented Pseudo-Invertible Decoder
Our APID provides an implementation of CSM  $(\kappa)$  for SCMs of class  $\mathfrak{B}(C^2, 2)$ . In the following, we list several important requirements, that the underlying probabilistic model has to satisfy. Then, we explain how deep normalizing flows [107, 125] meet these requirements.
1. A probabilistic model has to explicitly model an arbitrary function  $\hat{f}_Y(a,u_Y),u_Y\in [0,1]^2$  of class  $C^2$ . Normalizing flows with free-form Jacobians can implement arbitrary invertible  $C^2$  transformations from  $\hat{F}_a:\mathbb{R}^2\to \mathbb{R}^2$ . Then, by omitting one of the outputs, we can model functions  $f_{Y}(a,\cdot):[0,1]^{2}\rightarrow \mathcal{V}_{a}\subset \mathbb{R}$ . We discuss normalizing flows with free-form Jacobians in Sec. F.1.
2. A model has to fit the observational (interventional) distribution, given a sample from it,  $\mathcal{D} = \{A_i,Y_i\}_{i = 1}^n\sim \mathbb{P}(A,Y)$ . This is possible for the normalizing flows, as they are maximizing the log-likelihood of the data,  $\hat{\mathbb{P}} (Y = Y_{i}\mid A = A_{i})$ , directly with the gradient-based methods. Importantly, we also employ variational augmentations [22] to evaluate the log-likelihood. We discuss this in detail in Sec. F.2.
3. A probabilistic model should be able to perform the estimation of ECOU [ECOT], through abduction-action-prediction steps. In normalizing flows, this can be achieved with the help of the variational augmentations [22]. Specifically, for evidence,  $y'$  and  $a'$ , our proposed APID can infer arbitrarily many points from the estimated factual level set  $\hat{E}(y', a')$ , which follow the estimated posterior distribution of the latent noise, i.e.,  $\hat{\mathbb{P}}(U_Y \mid y', a')$ . We discuss this in detail in Sec. F.2.
4. All the level sets of the modeled function need to have a bounded curvature. This is possible for normalizing flows as a deep learning model. Namely, the curvature  $\kappa_{1}(u_{Y})$  can be evaluated via automatic differentiation tools and added the loss of the model. We discuss this in detail in Sec. F.3.
F.1 Choice of a normalizing flow with free-form Jacobians
Normalizing flows (NFs) [107, 125] differ in how flexible the modeled transformations in high dimensions are. Some models, like planar NF [107] or Sylvester NF [128], only allow for transformations with low-rank Jacobians. Other models employ masked auto-regressive networks [72] or coupling blocks [33, 34] to construct transformations with lower-triangular (structured) Jacobians.
Recently, several models were proposed for modeling free-form Jacobian transformations, e.g., (i) continuous NFs [23, 45] or (ii) residual NFs [9, 24]. However, this flexibility comes at a computational cost. For (i) continuous NFs, we have to solve a system of ordinary differential equations (ODEs) for every forward and reverse transformation. As noted by [45], numerical ODE solvers only work well for the non-stiff differential equations (=for non-steep transformations). For (ii) residual NFs, the computational complexity stems from the evaluation of the determinant of the Jacobian, which is required for the log-likelihood, and, from the reverse transformation, where we have to employ fixed point iterations.
We experimented with both continuous NFs and residual NFs but we found the latter to be more stable. The drawbacks of residual NFs can be partially fixed under our CSM  $(\kappa)$ . First, the determinant of the Jacobian can be evaluated exactly, as we set  $d = 2$ . Second, the numerical complexity of the reverse transformation<sup>8</sup> can be lowered by adding more residual layers, so that each layer models less steep transformation and the total Lipschitz constant is larger for the whole transformation. Hence, in our work, we resorted to residual normalizing flows [9, 24].
F.2 Variational augmentations and pseudo-invertibility
Variational augmentations were proposed for increasing the expressiveness of normalizing flows [22]. They augment the input to a higher dimension and then employ the invertible transformation of
Note that, in our APID, forward transformation corresponds to the inverse function, i.e.,  $\hat{f}_Y(a,\cdot)$
the flow. Hence, normalizing flows with variational augmentations can be seen as pseudo-invertible probabilistic models [10, 54].
We use variational augmentations to model the inverse function,  $\hat{f}_Y(a,\cdot)$ , i.e., sample points from the level sets. For this, we augment the (estimated) outcome  $Y\in \mathcal{V}$  with  $Y_{\mathrm{aug}}\sim N(g^{a}(Y),\varepsilon^{2})\in \mathbb{R}$  where  $g^{a}(\cdot)$  is a fully-connected neural network with one hidden layer and parameters  $\theta_{a}$ , and  $\varepsilon^2$  is a hyperparameter. As such, our proposed APID models  $\hat{f}_Y(a,\cdot)$  through a two-dimensional transformation  $\hat{F}_{a} = (\hat{f}_{Y_{\mathrm{aug}}}(a,\cdot),\hat{f}_{Y}(a,\cdot)):[0,1]^{2}\to \mathbb{R}\times \mathcal{V}_{a}$ . Variational augmentations facilitate our task in two ways.
(i) They allow evaluating the log-likelihood of the data via
$$
\left. \log \hat {\mathbb {P}} _ {\beta_ {a}, \theta_ {a}} (Y = Y _ {i} \mid a) = \mathbb {E} _ {Y _ {\text {a u g}, i}} \left[ \log \hat {\mathbb {P}} _ {\beta_ {a}} \left(Y _ {\text {a u g}} = Y _ {\text {a u g}, i}, Y = Y _ {i}\right) - \log N \left(Y _ {\text {a u g}, i}; g ^ {a} \left(Y _ {i}\right), \varepsilon^ {2}\right) \right], \right. \tag {93}
$$
where  $\beta_{a}$  are the parameters of the residual normalizing flow, and  $N(\cdot ;g^{a}(Y_{i}),\varepsilon^{2})$  is the density of the normal distribution. Here, we see that, by increasing the  $\varepsilon^2$ , the variance of the sample estimate of the log-likelihood of the data increases. On the other hand, for  $\varepsilon^2\rightarrow 0$ , the transformations of the residual flow become steeper, as we have to transport the point mass to the unit square,  $[0,1]^2$ .
(ii) Variational augmentations enable the abduction-action-prediction to estimate ECOU [ECOT] in a differential fashion. At the abduction step, we infer the sample from the posterior distribution of the latent noise, defined at the estimated factual level set  $\hat{E}(y', a')$ . For that, we variationally augment the evidence  $y'$  with  $b$  samples from  $(y', \{Y_{\mathrm{aug},j}'\}_{j=1}^{b} \sim N(g^{a'}(y'), \varepsilon^2))$ , and, then, transform them to the latent noise space with the factual flow
$$
\left\{\left(\tilde {u} _ {Y ^ {1}, j}, \tilde {u} _ {Y ^ {2}, j}\right) \right\} _ {j = 1} ^ {b} = \hat {F} _ {a ^ {\prime}} ^ {- 1} \left(\left\{\left(y ^ {\prime}, Y _ {\text {a u g}, j} ^ {\prime}\right) \right\} _ {j = 1} ^ {b}\right). \tag {94}
$$
Then, the action step selects the counterfactual flow,  $\hat{F}_a(\cdot)$ , and the prediction step transforms the abducted latent noise with it:
$$
\left\{\left(\hat {Y} _ {\text {a u g}, j}, \hat {Y} _ {j}\right) \right\} _ {j = 1} ^ {b} = \hat {F} _ {a} \left(\left\{\left(\tilde {u} _ {Y ^ {1}, j}, \tilde {u} _ {Y ^ {2}, j}\right) \right\} _ {j = 1} ^ {b}\right). \tag {95}
$$
In the end, ECOU [ECOT] is estimated by averaging  $\hat{Y}_j$ :
$$
\hat {Q} _ {a ^ {\prime} \rightarrow a} \left(y ^ {\prime}\right) = \frac {1}{b} \sum_ {j = 1} ^ {b} \hat {Y} _ {j}. \tag {96}
$$
In our APID, the parameters of the residual normalizing flow,  $\beta_{a}$ , and the parameters of the variational augmentations,  $\theta_{a}$ , are always optimized jointly. We use the reparametrization trick, to back-propagate through sampling of the augmentations.
F.3 Penalizing curvatures of the level sets
Although there exist deep learning models which explicitly bound the curvature of the modeled function (e.g., [124]), none of the works (to the best of our knowledge) did enforce the curvature of the level sets of the modeled function.
In our APID, the curvature  $\kappa_{1}(u_{Y})$  is evaluated using automatic differentiation exactly via Eq. (11) and then incorporated into the loss. As the calculation of the second derivatives is a costly operation, we heuristically evaluate the curvature only at the points of the counterfactual level set, which corresponds to the estimated ECOU [ECOT], namely  $(u_{Y^1}, u_{Y^2}) \in \hat{E}(\hat{Q}_{a' \to a}(y'), a)$ . For this, we again use the variationally augmented sample of size  $b$  for  $\hat{y} = \hat{Q}_{a' \to a}(y')$ :  $(\hat{y}, \{\hat{Y}_{\mathrm{aug},j}\}_{j=1}^{b} \sim N(g^a(\hat{y}), \varepsilon^2))$ . Then, we use the counterfactual flow (such as in Eq. (94))
$$
\left\{(\hat {u} _ {Y ^ {1}, j}, \hat {u} _ {Y ^ {2}, j}) \right\} _ {j = 1} ^ {b} = \hat {F} _ {a} ^ {- 1} \left(\left\{(\hat {y}, \hat {Y} _ {\text {a u g}, j}) \right\} _ {j = 1} ^ {b}\right), \tag {97}
$$
while the curvature has to be only evaluated at  $b$  points given by
$$
\left\{\kappa_ {1} \left(\hat {u} _ {Y ^ {1}, j}, \hat {u} _ {Y ^ {2}, j}\right) \right\} _ {j = 1} ^ {b}. \tag {98}
$$
G Details on training of Augmented Pseudo-Invertible Decoder
G.1 Training objective
Our APID satisfies all the requirements set by CSM  $(\kappa)$ , by combining several losses with different coefficients (see the overview in Fig. 13). Given observational data  $\mathcal{D} = \{A_i, Y_i\}_{i=1}^n$  drawn i.i.d. and a counterfactual query,  $Q_{a' \to a}(y')$ , for the partial identification, APID minimizes the following losses:
(1) Negative log-likelihood loss aims to fit the observational data distribution by minimizing the negative log-likelihood of the data,  $\mathcal{D}_a = \{A_i = a,Y_i\}$ , modeled by APID. To prevent the overfitting, we use noise regularization [110], which adds a normally distributed noise:  $\hat{Y}_i = Y_i + \xi_i;\xi_i\sim N(0,\sigma^2)$ , where  $\sigma^2 >0$  is a hyperparameter. Then, the negative log-likelihood loss for  $a\in \{0,1\}$  is
$$
\mathcal {L} _ {\mathrm {N L L}} \left(\beta_ {a}, \theta_ {a}\right) = - \frac {1}{n} \sum_ {i = 1} ^ {n} \log \hat {\mathbb {P}} _ {\beta_ {a}, \theta_ {a}} (Y = \tilde {Y} _ {i} \mid a), \tag {99}
$$
where the log-likelihood is evaluated according to the Eq. (93).
(2) Wasserstein loss prevents the posterior collapse [26, 132] of the APID, as the sample  $\{Y_i\}$  is not guaranteed to cover the full latent noise space when mapped with the estimated inverse function. We thus sample  $b$  points from the latent noise space,  $\{(U_{Y^1,j}, U_{Y^2,j})\}_{j=1}^b \sim \mathrm{Unif}(0,1)^2$  and map with the forward transformation  $F_a(\cdot)$  for  $a \in \{0,1\}$  via
$$
\left\{\left(\hat {Y} _ {\text {a u g}, a, j}, \hat {Y} _ {a, j}\right) \right\} _ {j = 1} ^ {b} = \hat {F} _ {a} \left(\left\{\left(U _ {Y ^ {1}, j}, U _ {Y ^ {2}, j}\right) \right\} _ {j = 1} ^ {b}\right). \tag {100}
$$
Then, we evaluate the empirical Wasserstein distance
$$
\mathcal {L} _ {\mathbb {W}} \left(\beta_ {a}\right) = \int_ {0} ^ {1} \left| \hat {\mathbb {F}} _ {\hat {Y} _ {a}} ^ {- 1} (q) - \hat {\mathbb {F}} _ {Y _ {a}} ^ {- 1} (q) \right| d q, \tag {101}
$$
where  $\hat{\mathbb{F}}_{\hat{Y}}^{-1}(\cdot)$  and  $\hat{\mathbb{F}}_Y^{-1}(\cdot)$  are empirical quantile functions, based on samples  $\{A_i = a,Y_i\}$  and  $\{\hat{Y}_{a,j}\}$ , respectively.
(3) Counterfactual query loss aims to maximize/minimize ECOU [ECOT]:
$$
\mathcal {L} _ {Q} \left(\beta_ {a ^ {\prime}}, \theta_ {a ^ {\prime}}, \beta_ {a}\right) = \text {S o f t p l u s} (\mp \hat {Q} _ {a ^ {\prime} \rightarrow a} (y ^ {\prime})), \tag {102}
$$
where  $\mp$  changes for maximization/minimization correspondingly,  $\hat{Q}_{a' \to a}(y')$  is evaluated as described in the Eq. (96), and Softplus(x) =  $\log(1 + \exp(x))$ . We use the softplus transformation to scale the counterfactual query logarithmically so that it matches the scale of the negative log-likelihood. We also block the gradients and fit only the counterfactual flow, when maximizing/minimizing ECOU [ECOT], to speed up the training.
(4) Curvature loss penalizes the curvature of the level sets of the modeled counterfactual function  $\hat{f}_Y(a,\cdot)$ :
$$
\mathcal {L} _ {\kappa} \left(\beta_ {a}, \theta_ {a}\right) = \frac {1}{b} \sum_ {j = 1} ^ {b} \kappa_ {1} \left(\hat {u} _ {Y ^ {1}, j}, \hat {u} _ {Y ^ {2}, j}\right), \tag {103}
$$
where  $\{\kappa_1(\hat{u}_{Y^1,j},\hat{u}_{Y^2,j})\}_{j = 1}^b$  are defined in Sec. F.3.
All the losses are summed up to a single training objective
$$
\mathcal {L} \left(\beta_ {a ^ {\prime}}, \theta_ {a ^ {\prime}}, \beta_ {a}, \theta_ {a}\right) = \sum_ {a \in \{0, 1 \}} \left[ \mathcal {L} _ {\mathrm {N L L}} \left(\beta_ {a}, \theta_ {a}\right) + \mathcal {L} _ {\mathbb {W}} \left(\beta_ {a}\right) \right] + \lambda_ {Q} \mathcal {L} _ {Q} \left(\beta_ {a ^ {\prime}}, \theta_ {a ^ {\prime}}, \beta_ {a}\right) + \lambda_ {\kappa} \mathcal {L} _ {\kappa} \left(\beta_ {a}, \theta_ {a}\right). \tag {104}
$$
G.2 Training algorithm and hyperparameters
Training algorithm. The training of APID proceeds in three stages; see the pseudocode in Algorithm 1. At a burn-in stage ( $n_B = 500$  training iterations), we only fit two residual normalizing flows (one for each treatment) with  $\mathcal{L}_{\mathrm{NLL}}$  and  $\mathcal{L}_{\mathbb{W}}$ . Then, we copy the counterfactual flow and variational
Figure 13: Overview of the training and inference of our Augmented Pseudo-Invertible Decoder.
augmentation parameters twice, for the task of maximization/minimization, respectively. Also, we freeze the factual flow parameters. At a query stage ( $n_{Q} = 100$  training iterations), we, additionally to previous losses, enable the counterfactual query loss,  $\lambda_{Q}\mathcal{L}_{Q}$ . During this stage, two counterfactual flows are able to realign and start "bending" their level sets around the frozen factual level set. Ultimately, at the last curvature-query stage ( $n_{\mathcal{C}Q} = 500$  training iterations), all the parts of the loss in Eq. (104) are enabled. After training is over, we report the upper and lower bounds for the counterfactual query as an output of the APID with exponential moving average (EMA) of the model parameters [105]. EMA of the model parameters allows us to reduce the variance during training and is controlled via the smoothness hyperparameter  $\gamma = 0.99$ .