# Convergence analysis of stochastic gradient descent with adaptive
preconditioning for non-convex and convex functions

Dmitrii A. Pasechnyuk1,2,3,4, Alexander Gasnikov2,3,4, Martin Takáč1
1Mohamed bin Zayed University of Artificial Intelligence, UAE,
2Moscow Institute of Physics and Technology, Russia
3Kharkevich Institute for Information Transmission Problems RAS, Russia
4Ivannikov Institute for System Programming RAS, Russia
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Preconditioning is a crucial operation in gradient-based numerical
optimisation. It helps decrease the local condition number of a function by
appropriately transforming its gradient. For a convex function, where the
gradient can be computed exactly, the optimal linear transformation
corresponds to the inverse of the Hessian operator, while the optimal convex
transformation is the convex conjugate of the function. Different conditions
result in variations of these dependencies. Practical algorithms often employ
low-rank or stochastic approximations of the inverse Hessian matrix for
preconditioning. However, theoretical guarantees for these algorithms
typically lack a justification for the defining property of preconditioning.
This paper presents a simple theoretical framework that demonstrates, given a
smooth function and an available unbiased stochastic approximation of its
gradient, that it is possible to refine the dependency of the convergence rate
on the Lipschitz constant of the gradient.

##  1 Introduction

Although a significant portion of the history of preconditioning was dedicated
to solving systems of linear equations and quadratic programming, we direct
the reader to [1, 2] and the references therein instead of discussing it here.
In general optimisation which considers problems of the following form:

$\min_{x\in Q}f(x),$ (1)

preconditioning was one of the first methods developed to generalise the
gradient descent method. It has multiple interpretations: change of variables,
modification of the function’s model to minimize at the current step, and
aligning the step direction closer to the Newton one [3]. The latter
interpretation is the primary inspiration for this paper. Let us introduce a
simple, yet general, update formula for a preconditioned method, without a
detailed description of the context:

$x_{t+1}\leftarrow x_{t}-\alpha_{t}k_{t}(g_{t}).$ (2)

Unless stated otherwise, $P_{t}$ is defined by $k_{t}(g)=P_{t}^{-1}g$, i.e.,
linear preconditioning is considered, and the objective function $f$ is
assumed to be convex. Preconditioning is closely related to Newton-type
methods, which are preconditioned methods with $P_{t}=\nabla^{2}f(x_{t})$. If
$g_{t}=\nabla f(x_{t})$, the (damped) Newton method converges (asymptotically)
as fast as the optimal gradient method [4] with a proper choice of step sizes
$\alpha_{t}$, and super-linearly in the vicinity of the solution [5]. If
$g_{t}=\nabla f(x_{t})+\xi_{t}$, where $\xi_{t}$ is a random variable with a
known probability density function $p$, the convex preconditioned method with
$k_{t}(g)=(\nabla^{2}f(x^{*}))^{-1}(J(p))^{-1}\nabla(\log p)(g)$, where $J(p)$
is the Fisher information matrix for $p$, is asymptotically optimal among
convex preconditioned methods [6]. Despite the important role of the Fisher
information matrix in stochastic preconditioned methods [7], it is beyond the
scope of this paper. It is evident that methods involving $(\nabla^{2}f)^{-1}$
in the preconditioning operation are worth studying.

Inverse Hessian computation is a time- and memory-consuming procedure in
large-scale optimisation problems. Historically, the first response to this
challenge was the development of variable metric [8] and, especially, quasi-
Newton [9] methods, which are preconditioned methods with $P_{t}$ being an
approximation of $\nabla^{2}f(x_{t})$. Variable metric methods also include
adaptive methods [10], which are more studied in the stochastic (and online)
optimisation setting and allow obtaining local guarantees [11]. However, these
approaches only serve as inspiration and are beyond the scope of this paper.
Quasi-Newton methods require the Hessian approximation to be asymptotically
exact, which allows demonstrating that key properties of the Newton method
remain valid under some conditions [12]. If $f$ is strongly self-concordant,
one can obtain explicit convergence rates of Broyden’s family of quasi-Newton
methods by taking into account the discrepancy between the Hessian and its
approximation [13]. Some recent results [14] are obtained by considering the
inexact cubic regularized Newton method [15]. At the same time, variable
metric methods, which allow the use of arbitrary (in general, stochastically
biased) Hessian approximations, are understudied. Despite this, such
approximations (for instance, Jacobi’s preconditioning [16] and Hutchinson’s
trace estimator [17]) are used to design efficient optimisation methods [18,
19]. The latter methods inspire the preconditioning operation proposed in this
paper with the aim to demonstrate the theoretical framework allowing a wider
range of Hessian approximations and applicable to any variable metric method.

##  2 Theory

Let us consider optimisation problem (1) for some Hilbert space
$(E,\langle\cdot,\cdot\rangle)$ with induced norm $\|\cdot\|$ and dual space
$E^{*}$ with dual norm $\|\cdot\|^{*}$, closed convex non-empty $Q\subseteq E$
and $f\in C^{2}(Q,\mathbb{R})$ with $M$-Lipschitz continuous w.r.t.
$T$-Mahalanobis norm Hessian $\nabla^{2}f\in\mathcal{L}(E,E^{*})$, i.e. it
holds that

$f(y)\leqslant f(x)+\langle\nabla
f(x),y-x\rangle+\frac{1}{2}\|y-x\|_{\nabla^{2}f(x)}+\frac{M}{6}\|y-x\|_{T(x)}^{3},\quad\forall
x,y\in Q$ (3)

for some $T:E\to\mathcal{L}(E,E^{*})$ such that $T(x)$ is positive-definite
$\forall x\in Q$. $\|x\|_{A}^{2}:=\langle Ax,x\rangle$ for any $x\in E$ and
$A\in\mathcal{L}(E,E^{*})$. $\mathcal{L}(E,E^{*})$ consists of bounded linear
operators from $E$ to $E^{*}$. We assume that there exists second-order
minimum $x^{*}\in Q$ of $f$, i.e. such that $\inf_{x\in Q}f(x)=f(x^{*})$ and
$\nabla^{2}f(x^{*})$ is positive-definite. $f$ is equipped with unbiased
stochastic gradient oracle $g$ with bounded variance and stochastic Hessian
approximation $H$, defined as follows.

Let $(\Xi,\mathcal{F}_{\xi},P_{\xi})$ and $(N,\mathcal{F}_{\eta},P_{\eta})$ be
probability triples with independent $P_{\xi}$, $P_{\eta}$. $g:Q\times\Xi\to
E^{*}$ is unbiased stochastic gradient oracle [10] (Definition 3.4.1), i.e.
$\forall x\in Q$ it holds that

$\mathbb{E}_{P_{\xi}}[g(x,\xi)]=\nabla f(x).$ (4)

Hereinafter $\xi$ is a sample drawn from $P_{\xi}$. $g$ has bounded variance,
i.e. $\forall x\in Q$ it holds that

$\mathbb{E}_{P_{\xi}}[\|g(x,\xi)-\nabla
f(x)\|^{*}]^{2}\leqslant\sigma_{g}^{2}$ (5)

for some $\sigma_{g}>0$. In turn, $H:Q\times N\to\mathcal{L}(E,E^{*})$ is
stochastic Hessian approximation, i.e. $\forall x\in Q$ it holds that

$\mathbb{E}_{P_{\eta}}[\|H(x,\eta)-\nabla^{2}f(x)\|_{op}^{2}]\leqslant\sigma_{H}^{2}$
(6)

for some $\sigma_{H}>0$. Hereinafter $\eta$ is a sample drawn from $P_{\eta}$.
$\|A\|_{op}:=\sup_{\|x\|=1}\|Ax\|^{*}$ for any $A\in\mathcal{L}(E,E^{*})$.
Besides, we require $H(x,\eta)$ to be positive-definite for every $x,\eta$ in
the entire domain of $H$.

Let us apply Linear Averaged Preconditioning (LAP) to Stochastic Gradient
Descent (SGD). The proposed algorithm is listed in Algorithm 1. Note that if
$E=\mathbb{R}^{n}$ for some $n\in\mathbb{N}$, one can choose $P_{0}$ to be
equal to the $n\times n$ identity matrix. Also, note that the gradient step
can be equivalently rewritten as follows:

$x_{t+1}\leftarrow\operatorname{proj}_{Q}\left\\{x_{t}-\alpha_{t}P_{t}^{-1}g_{t}\right\\},$
(7)

where $\operatorname{proj}_{Q}$ is a projection onto $Q$, which is well-
defined due to the Hilbert projection theorem. It is worth noticing that the
averaging in Algorithm 1 is an exponential moving average with a variable
smoothing constant $\beta_{t}$. This is the only averaging to consider if
greedy-optimal preconditioning is to be found. Indeed, the decrease of the
objective function is upper bounded by the Taylor expansion, which is
dependent only on $P_{t}$ and $\nabla^{2}f(x_{t})$, so no previous
preconditionings have an impact on this upper bound. Further, there is no
$P_{t+1}^{\prime}$ that ensures a better decrease of the objective function
than any $P_{t}+\beta(H_{t}-P_{t})$, $\beta\in[0,1]$ (preconditioning operator
lying on a segment between $P_{t}$ and $H_{t}$ in $\mathcal{L}(E,E^{*})$)
because of a geometrical reason: there exists $P_{t+1}$ on this segment, such
that

$\|P_{t+1}-P_{t}\|_{op}\leqslant\|P_{t+1}^{\prime}-P_{t}\|_{op}\ \mbox{and}\
\|P_{t+1}-H_{t}\|_{op}\leqslant\|P_{t+1}^{\prime}-H_{t}\|_{op}$

with at least one inequality being strict (the orthogonal projection of
$P_{t+1}^{\prime}$ onto the segment satisfies this due to the Hilbert
projection theorem). At the same time, the decrease of the objective function
depends only on the norms of operators $P_{t+1}-P_{t}$ and $P_{t+1}-H_{t}$
along certain directions, which both are upper bounded by their operator norms
$\|\cdot\|_{op}$. Note that $P_{t}$ is positive-definite for all $t=0,...,T$
by construction.

Algorithm 1 SGD with LAP

initial point $x_{0}\in Q$

initial positive-definite preconditioning $P_{-1}\in\mathcal{L}(E,E^{*})$

iterations number $T>0$

step sizes $\\{\alpha_{t}>0:t=0,...,T-1\\}$

averaging rates $\\{\beta_{t}\in[0,1]:t=0,...,T-1\\}$

draw sample $t\in\\{0,...,T-1\\}$ from $P_{T}$, defined by
$P_{T}(t):=\alpha_{t}/\sum_{k=0}^{T-1}\alpha_{k}$, $x_{t}$

for $t=0,...,T-1$ do

draw samples $\xi_{t}$, $\eta_{t}$ from $P_{\xi}$, $P_{\eta}$, correspondingly

evaluate oracles $g_{t}\leftarrow g(x_{t},\xi_{t})$, $H_{t}\leftarrow
H(x_{t},\eta_{t})$

average preconditioning $P_{t}\leftarrow P_{t-1}+\beta_{t}(H_{t}-P_{t-1})$

make gradient step $x_{t+1}\leftarrow\arg\min_{x\in
Q}\left\\{\alpha_{t}\langle
g_{t},x-x_{t}\rangle+\frac{1}{2}\|x-x_{t}\|^{2}_{P_{t}}\right\\}$

end for

Firstly, let us consider problems of optimising non-convex functions given the
biased, i.e. not satisfying condition (4), stochastic gradient oracle. Plain
stochastic gradient descent applied to functions with $L$-Lipschitz continuous
gradient requires $\alpha_{t}=O(1/L)$ and ensures
$O(L(f(x_{0})-f(x^{*}))/T)+O(\sigma_{g}^{2})$ convergence rate for squared
norm of the gradient, where $g\leqslant O(f)$ means that there exists
$c\in\mathbb{R}_{+}$ such that $g(t)\leqslant cf(t)$ for all $t\in\mathbb{N}$.
Intuitively, preconditioning must let choose $\alpha_{t}$ closer to $1$.
Theorem 1 below gives explicit formula for $\alpha_{t}$, such that
$1/\alpha_{t}-1$ tends to value, proportional to $\sigma_{H}$, $\sigma_{T}$ or
$\sigma_{g}$, conditions allowing arbitrary choice of $\beta_{t}\in(0,1)$, and
non-asymptotic convergence rate of Algorithm 1 for expectation of $(\|\nabla
f(x_{t})\|_{P_{t}}^{*})^{2}$.

######  Theorem 1 (Convergence rate of SGD with LAP, non-convex function).

If $\beta_{0}=1$, $\beta_{t}\in[0,1]$, and
$\lambda_{\min}(P_{t})\geqslant\mu_{t}$ for all $t=1,...,T-1$, where sequence
$\\{\mu_{t}>0:t=1,...,T-1\\}$ is chosen satisfying the following condition:

$\frac{1}{\mu_{t}}>\frac{1}{\mu_{t-1}}+\frac{M\alpha_{t-1}\|g_{t-1}\|_{P_{t-1}}^{*}(1+\sigma_{T}^{1/2})}{\Delta_{t-1}},$

where $\mu_{0}>0$, $\Delta_{0}=\sigma_{H}/\mu_{0}$, and sequence
$\\{\Delta_{t}>0:t=0,...,T-1\\}$ is decreasing, and where step sizes are
chosen satisfying the following condition:

$\alpha_{t}\leqslant\frac{1}{1+\left(1+\sqrt{M\|g_{t}\|_{P_{t}}^{*}}\right)\Delta_{t}+\sqrt{M\|g_{t}\|_{P_{t}}^{*}}(\sigma_{T}+2)},$

it holds that

$\displaystyle\mathbb{E}_{P_{T}P_{\xi}P_{\eta}}\left[\|\nabla
f(x_{t})\|_{P_{t}}^{*}\right]^{2}$
$\displaystyle\leqslant\frac{2(f(x_{0})-f(x^{*}))}{T}\left(1+\frac{\sigma_{H}}{\mu_{0}}+\sqrt{MB}\left(\sigma_{T}+\frac{\sigma_{H}}{\mu_{0}}+2\right)\right)+\frac{3\sigma_{g}^{2}}{\mathbb{E}_{P_{T}}[\mu_{t^{\prime}}]},$

where $B:=\max_{k=0,...,T-1}\|g_{k}\|^{*}_{P_{k}}>0$.

Secondly, let us consider convex optimisation problems with unbiased
stochastic gradient oracle. Convergence rate of plain stochastic gradient
method is guaranteed to be $O(R^{2}L/T)+O(\sigma R/\sqrt{T})$ in this case
[20], which is achieved by choice $\alpha_{t}=O(1/\sqrt{T})$, in addition to
$O(1/L)$. Theorem 2 gives conditions on $\beta_{t}$ and $\alpha_{t}$ ensuring
similar convergence rate for SGD with LAP, which are more strict than in non-
convex case, and provide non-asymptotic convergence rate of Algorithm 1 for
expectation of $f(x_{t})-f(x^{*})$.

######  Theorem 2 (Convergence rate of SGD with LAP, convex function).

If $\beta_{0}=1$, $\beta_{t}\in[0,1]$, and
$\lambda_{\min}(P_{t})\geqslant\mu>0$ for all $t=1,...,T-1$, and sequence
$\\{\Delta_{t}>0:t=0,...,T-1\\}$ is decreasing, and where step sizes and
averaging rates are chosen satisfying the following condition:

$\alpha_{t}=\min\left\\{\frac{1}{\frac{1}{c_{1}}(1+\Delta_{t})+(\|g_{t}\|_{P_{t}}^{*})^{3/4}(1+(\Delta_{t}+\sigma_{T})^{3/4})},\sqrt{\frac{2c_{3}\sigma_{H}+c_{3}M(1+\sigma_{T}^{1/2})+c_{2}+1}{T\left(\frac{\sigma_{g}^{2}}{\mu(1-c_{1})}+\frac{M}{2}\right)}}\right\\},$
$\frac{1}{1-\beta_{t}}\geqslant\max\left\\{1+\frac{\alpha_{t-1}M\|g_{t-1}\|_{P_{t-1}}^{*}(1+\sigma_{T}^{1/2})+2\sigma_{H}}{\|H_{t-1}-P_{t-1}\|_{op}},(t+1)^{2}\right\\},$

where $c_{1}$, $c_{2}$, and $c_{3}$ denote some small positive constants, it
holds that

$\displaystyle\mathbb{E}_{P_{\eta}P_{\xi}P_{T}}[f(x_{t})-f(x^{*})]\leqslant
O(1)\sqrt{\frac{(\sigma_{g}^{2}/\mu+MB^{3/2})(\sigma_{H}+M(1+\sigma_{T}^{1/2})+1)}{T}}$
$\displaystyle\quad\quad+O(1)\left(1+\sigma_{H}+(1+(\sigma_{H}+\sigma_{T})^{3/4})B^{3/4}\right)\frac{R^{2}(\sigma_{H}+M(1+\sigma_{T}^{1/2})+1)}{T},$

where $R:=\|x_{0}-x^{*}\|_{H_{0}}>0$ and $O(1)$ denotes some small positive
constant, hiding $c_{1}$, $c_{2}$, and $c_{3}$.

Summarising Theorems 1 and 2, it worth noting that Lipschitz constant of the
gradient $L$, included in convergence rates of plain SGD, was replaced
everywhere by the norm of the gradient $B$, which is decreasing while
approaching to $x^{*}$, together with relative Lipschitz constant of the
Hessian $M$. This allows to address problems of optimising function which
gradient is not Lipschitz-continuous, i.e. growth of gradient is not limited.
Besides, dependency of convergence rates on $B$ improves upon that of $L$ for
plain SGD. Omitting terms proportional to $\sigma_{H}$, $\sigma_{T}$ or
$\sigma_{g}$, for non-convex functions, decreasing term of convergence rate is
$O((1+\sqrt{MB})(f(x_{0})-f(x^{*}))/T)$, and for convex functions, convergence
rate is $O(MB^{3/4}(1/\sqrt{T}+R^{2}/T))$.

##  3 Practical example

Let us consider the following problem:

$\min_{x\in\mathbb{R}^{n}}\langle Ax,x\rangle-\langle b,x\rangle,$ (8)

where $A\in\mathbb{R}_{+}^{n\times n}$ is positive definite with uniform
random non-diagonal elements drawn from $(0,0.001)$ and uniform random
diagonal elements drawn from $(0,0.006)$, and $b\in\mathbb{R}_{+}^{n}$ with
uniform random elements drawn from $(0,1)$.

Figure 1 shows convergence curves of SGD and SGD with LAP, Algorithm 1, with
vertical axis measuring function value $f(x_{t})$ (left, $f(x^{*})=0$) or norm
of the gradient $\|\nabla^{2}f(x_{t})\|_{2}$ (right) in logarithmic scale, and
horizontal axis measuring iterations number $t$. One can see that
preconditioning with properly chosen $\beta_{t}$, $t=0,...,T-1$ improves upon
plain stochastic gradient descent, even if $\sigma_{H}$ is significantly big
(bigger than $\lambda_{\max}(A)$, in our experiment). Note, that the case of
biased stochastic gradient oracle is considered and therefore both algorithms
do not converge to $x^{*}$, but only to its vicinity.

(a) $\beta_{t}\in(0,1),t=0,...,T-1$, $\alpha_{0}=0.8$

(b) $\beta_{t}\in\left(1-\frac{1}{t^{2}},1\right),t=0,...,T-1$, $\alpha_{0}=4$

Figure 1: Convergence curves of SGD with and without LAP applied to (8) with
$T=5000$, $n=10$, $\sigma_{g}=0.004$ (biased), $\sigma_{H}=1$,
$\mu_{t}>0.008$, $\beta_{t}\approx\arg\min_{\beta}f(x_{t+1})$ using Brent’s
method on a corresponding interval, and $\alpha_{t}=\alpha_{0}/\sqrt{t}$ with
$\alpha_{0}=10$ for SGD and corresponding $\alpha_{0}$ for SGD with LAP.
$\alpha_{0}$ are chosen to be optimal using grid search.

##  4 Discussion

This paper continues the development of theoretical framework for analysing
preconditioning in stochastic gradient methods, initiated by authors in [21].

Theorem 1 improves upon previous analysis of stochastic gradient descent with
preconditioning in terms of:

  1. 1.

requirements on $\beta_{t}$, $t=1,...,T-1$, which were completely compensated
by requirements on margins from zero $\mu_{t}$, $t=1,...,T-1$;

  2. 2.

conditions on Hessian and its approximation, which are now presented in
“additive” notation through $\|\nabla^{2}f(x_{t})-P_{t}\|_{op,P_{t}}$ instead
of previously used “multiplicative” notation through bounding
$\|\nabla^{2}f(x_{t})P_{t}^{-1}\|_{op}$, which narrowed the class of
acceptable $\nabla^{2}f(x_{t})$.

Newly initiated direction of work is analysis of stochastic gradient descent
with preconditioning for convex functions in Theorem 2. The proposed analysis
improves upon [22] by ensuring the absence of diverging terms. The dependency
on $B$ and requirements on $\beta_{t}$, $t=1,...,T-1$ in the Theorem 2 are to
be improved in future works.

## References

  * Benzi [2002] Michele Benzi.  Preconditioning techniques for large linear systems: a survey.  _Journal of computational Physics_ , 182(2):418–477, 2002.
  * Qu et al. [2022] Zhaonan Qu, Wenzhi Gao, Oliver Hinder, Yinyu Ye, and Zhengyuan Zhou.  Optimal diagonal preconditioning: Theory and practice.  _arXiv preprint arXiv:2209.00809_ , 2022.
  * Boyd and Vandenberghe [2004] Stephen P Boyd and Lieven Vandenberghe.  _Convex optimization_.  Cambridge university press, 2004.
  * Hanzely et al. [2022] Slavomír Hanzely, Dmitry Kamzolov, Dmitry A Pasechnyuk, Alexander Gasnikov, Peter Richtárik, and Martin Takác.  A damped newton method achieves global O(1/k²) and local quadratic convergence rate.  _Advances in Neural Information Processing Systems_ , 35:25320–25334, 2022.
  * Polyak [1987] Boris T Polyak.  _Introduction to optimization_.  New York, Optimization Software, 1987.
  * Polyak and Tsypkin [1980] Boris T Polyak and Yakov Z Tsypkin.  Optimal pseudogradient adaptation algorithms.  _Avtomatika i Telemekhanika_ , 8:74–84, 1980.
  * Li [2018] Xi-Lin Li.  Preconditioner on matrix lie group for sgd.  _arXiv preprint arXiv:1809.10232_ , 2018.
  * Davidon [1991] William C Davidon.  Variable metric method for minimization.  _SIAM Journal on optimization_ , 1(1):1–17, 1991\.
  * Dennis and Moré [1977] John E Dennis, Jr and Jorge J Moré.  Quasi-newton methods, motivation and theory.  _SIAM review_ , 19(1):46–89, 1977.
  * Duchi [2018] John C Duchi.  Introductory lectures on stochastic optimization.  _The mathematics of data_ , 25:99–186, 2018.
  * Lu et al. [2022] Zhou Lu, Wenhan Xia, Sanjeev Arora, and Elad Hazan.  Adaptive gradient methods with local guarantees.  _arXiv preprint arXiv:2203.01400_ , 2022.
  * Powell [1976] Michael JD Powell.  Some global convergence properties of a variable metric algorithm for minimization without exact line searches.  _Nonlinear programming_ , 9(1):53–72, 1976.
  * Rodomanov and Nesterov [2021] Anton Rodomanov and Yurii Nesterov.  Greedy quasi-newton methods with explicit superlinear convergence.  _SIAM Journal on Optimization_ , 31(1):785–811, 2021.
  * Kamzolov et al. [2023] Dmitry Kamzolov, Klea Ziu, Artem Agafonov, and Martin Takáč.  Accelerated adaptive cubic regularized quasi-newton methods.  _arXiv preprint arXiv:2302.04987_ , 2023.
  * Ghadimi et al. [2017] Saeed Ghadimi, Han Liu, and Tong Zhang.  Second-order methods with cubic regularization under inexact information.  _arXiv preprint arXiv:1710.05782_ , 2017.
  * Jacobi [1845] Carl GJ Jacobi.  Ueber eine neue auflösungsart der bei der methode der kleinsten quadrate vorkommenden lineären gleichungen.  _Astronomische Nachrichten_ , 22(20):297–306, 1845.
  * Hutchinson [1989] Michael F Hutchinson.  A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines.  _Communications in Statistics-Simulation and Computation_ , 18(3):1059–1076, 1989.
  * Jahani et al. [2021] Majid Jahani, Sergey Rusakov, Zheng Shi, Peter Richtárik, Michael W Mahoney, and Martin Takáč.  Doubly adaptive scaled algorithm for machine learning using second-order information.  _arXiv preprint arXiv:2109.05198_ , 2021.
  * Sadiev et al. [2022] Abdurakhmon Sadiev, Aleksandr Beznosikov, Abdulla Jasem Almansoori, Dmitry Kamzolov, Rachael Tappenden, and Martin Takáč.  Stochastic gradient methods with preconditioned updates.  _arXiv preprint arXiv:2206.00285_ , 2022.
  * Lan [2012] Guanghui Lan.  An optimal method for stochastic composite optimization.  _Mathematical Programming_ , 133(1-2):365–397, 2012.
  * Pasechnyuk et al. [2022] Dmitry A Pasechnyuk, Alexander Gasnikov, and Martin Takáč.  Effects of momentum scaling for sgd.  _arXiv preprint arXiv:2210.11869_ , 2022.
  * Scheinberg and Tang [2016] Katya Scheinberg and Xiaocheng Tang.  Practical inexact proximal quasi-newton method with global complexity analysis.  _Mathematical Programming_ , 160:495–529, 2016.
  * Bertsekas [2009] Dimitri Bertsekas.  _Convex optimization theory_ , volume 1.  Athena Scientific, 2009.

##  Appendix A Proofs

######  Lemma 1 (Non-expansiveness of $\operatorname{proj}_{Q}$ [23]).

For any closed convex non-empty $Q\subseteq E$, $x\in Q$, and $y\in E$, it
holds that

$\|\operatorname{proj}_{Q}y-x\|\leqslant\|y-x\|.$ (9)

######  Lemma 2 (Direction of $\operatorname{proj}_{Q}$).

For any closed convex non-empty $Q\subseteq E$, $x\in Q$, $g\in E^{*}$, and
positive-definite $P\in\mathcal{L}(E,E^{*})$, it holds that

$\langle g,\operatorname{proj}_{Q}(x-P^{-1}g)-(x-P^{-1}g)\rangle\geqslant 0.$
(10)

###### Proof.

It follows from that, in the same context and any $h\in E$, it holds that

$\langle h,\operatorname{proj}_{Q}(x-h)-(x-h)\rangle\geqslant 0,$

by applying it to $h=P^{-1}g$, due to the equivalence of $\|\cdot\|$ and
$\|\cdot\|_{P}^{*}$ together with self-adjointness of $P^{-1}$ (implied by its
positive-definiteness). ∎

######  Lemma 3 (Length of $\operatorname{proj}_{Q}$).

For any closed convex non-empty $Q\subseteq E$, $x\in Q$, $g,d\in E^{*}$, and
positive-definite $P\in\mathcal{L}(E,E^{*})$, it holds that

$|\langle
d,(x-P^{-1}g)-\operatorname{proj}_{Q}(x-P^{-1}g)\rangle|\leqslant\|d\|_{P}^{*}\|g\|_{P}^{*}.$
(11)

###### Proof.

It follows from extended Cauchy–Schwarz inequality. ∎

######  Lemma 4 (Correction to coefficient of cubic Taylor’s expansion term).

For any $g\in E^{*}$, and positive-definite $P,T\in\mathcal{L}(E,E^{*})$, it
holds that

$\|P^{-1}g\|^{3}_{T}\leqslant\frac{3}{2}(\|g\|_{P}^{*})^{3}(1+|\|T-P\|_{op,P}|^{3/2}).$
(12)

###### Proof.

Firstly,

$\displaystyle\langle P^{-1}g,TP^{-1}g\rangle$ $\displaystyle=\langle
P^{-1}g,g\rangle+\langle P^{-1}g,(T-P)P^{-1}g\rangle$
$\displaystyle\leqslant(\|g\|_{P}^{*})^{2}+\|T-P\|_{op,P}(\|g\|_{P}^{*})^{2}$
$\displaystyle\leqslant(\|g\|_{P}^{*})^{2}(1+|\|T-P\|_{op,P}|)$

The fact that $(1+x)^{3/2}\leqslant\frac{3}{2}(1+x^{3/2})$ for any
$x\in\mathbb{R}_{+}$ implies lemma’s statement. ∎

######  Lemma 5 (Harmonic mean upper bound).

For any $I\in\mathbb{N}$ and sequence
$\\{a_{i}\in\mathbb{R}_{+}:i=1,...,I\\}$, it holds that

$\frac{I}{\sum_{i=1}^{I}\frac{1}{a_{i}}}\leqslant\frac{\sum_{i=1}^{I}a_{i}}{I}.$
(13)

######  Lemma 6 (Abel’s inequality).

For any $I\in\mathbb{N}$ and sequences $\\{a_{i}\in\mathbb{R}:i=1,...,I\\}$,
$\\{b_{i}\in\mathbb{R}:i=1,...,I\\}$, it holds that

$\displaystyle\sum_{i=1}^{I}a_{i}b_{i}=\left(\sum_{i=1}^{I}b_{i}\right)a_{I}-\sum_{j=1}^{I-1}\left(\sum_{i=1}^{j}b_{i}\right)(a_{j+1}-a_{j}),$
(14)
$\displaystyle\sum_{i=1}^{I}a_{i}b_{i}=\left(\sum_{i=1}^{I}b_{i}\right)a_{0}-\sum_{j=1}^{I-1}\left(\sum_{i=1}^{j}b_{i}\right)(a_{j}-a_{j+1}).$

######  Lemma 7 (Abel’s test).

For any $I\in\mathbb{N}$ and sequences
$\\{a_{i}\in\mathbb{R}:i=1,...,+\infty\\}$,
$\\{b_{i}\in\mathbb{R}_{+}:i=1,...,+\infty\\}$, such that $b_{i+1}\leqslant
b_{i}$ for any $i=1,...,+\infty$ and series $\sum_{i=1}^{+\infty}a_{i}$
converges, it holds that there exists $c\in\mathbb{R}_{+}$, such that

$\sum_{i=1}^{\infty}a_{i}b_{i}=c.$ (15)

######  Proposition 1.

For any $a,b\in E^{*}$ and positive-definite $P\in\mathcal{L}(E,E^{*})$, it
holds that

$-\langle
a,P^{-1}b\rangle=-\frac{1}{2}\|a\|_{P}^{*}-\frac{1}{2}\|b\|_{P}^{*}+\frac{1}{2}\|a-b\|_{P}^{*}.$
(16)

######  Proposition 2.

For any $a,x\in\mathbb{R}_{+}$, it holds that
$ax+\frac{a^{2}}{2}\leqslant\frac{x^{2}}{4}+\frac{3a^{2}}{2}$ (17).

######  Proposition 3.

For any $A,B\in\mathcal{L}(E,E^{*})$ and positive-definite
$P\in\mathcal{L}(E,E^{*})$, such that $\mu\leqslant\lambda_{\min}(P)$, it
holds that

$\displaystyle\mu\frac{\langle Ax,x\rangle}{\langle
Px,x\rangle}\leqslant\frac{\langle Ax,x\rangle}{\langle
x,x\rangle}\leqslant\|A\|_{op},$ (18) $\displaystyle\langle
Bx,x\rangle\leqslant\|B\|_{op}\langle
x,x\rangle\leqslant\frac{\|B\|_{op}}{\mu}\langle Px,x\rangle.$

###### Proof of Theorem 1 via diminishing step sizes.

Let us introduce $\Delta_{t}>0$, such that

$\|\nabla^{2}f(x_{t})-P_{t}\|_{op,P_{t}}\leqslant\Delta_{t}.$ (19)

$\|A\|_{op,P}:=\sup_{\|x\|_{P}=1}\|Ax\|$ for any positive-definite
$P\in\mathcal{L}(E,E^{*})$ and $A\in\mathcal{L}(E,E^{*})$ (not necessarily
positive-definite). Let us also introduce auxiliary point

$x_{t+1}^{\prime}:=x_{t}-P_{t}^{-1}g_{t}.$ (20) $\displaystyle f(x_{t+1})$
$\displaystyle\leqslant f(x_{t})+\langle\nabla
f(x_{t}),x_{t+1}-x_{t}\rangle+\frac{1}{2}\|x_{t+1}-x_{t}\|_{\nabla^{2}f(x_{t})}^{2}+\frac{M}{6}\|x_{t+1}-x_{t}\|_{T(x_{t})}$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:proj-
exp}}}{{\leqslant}}f(x_{t})+\langle\nabla
f(x_{t}),x_{t+1}-x_{t}\rangle+\frac{\alpha_{t}^{2}}{2}\|P_{t}^{-1}g_{t}\|_{\nabla^{2}f(x_{t})}^{2}+\frac{\alpha_{t}^{3}M}{6}\|P_{t}^{-1}g_{t}\|_{T(x_{t})}$
$\displaystyle=f(x_{t})+\langle\nabla
f(x_{t}),x_{t+1}-x_{t}\rangle+\frac{\alpha_{t}^{2}}{2}(\|g_{t}\|_{P_{t}}^{*})^{2}+$
$\displaystyle\quad+\frac{\alpha_{t}^{2}}{2}\langle(\nabla^{2}f(x_{t})-P_{t})P_{t}^{-1}g_{t},P_{t}^{-1}g_{t}\rangle+\frac{\alpha_{t}^{3}M}{6}\|P_{t}^{-1}g_{t}\|_{T(x_{t})}$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:delta}}}{{\leqslant}}f(x_{t})+\langle\nabla
f(x_{t}),x_{t+1}-x_{t}\rangle+\frac{\alpha_{t}^{2}(1+\Delta_{t})}{2}(\|g_{t}\|_{P_{t}}^{*})^{2}+\frac{\alpha_{t}^{3}M}{6}\|P_{t}^{-1}g_{t}\|_{T(x_{t})}^{3}$
$\displaystyle=f(x_{t})+\langle\nabla
f(x_{t}),x_{t+1}^{\prime}-x_{t}\rangle+\frac{\alpha_{t}^{2}(1+\Delta_{t})}{2}(\|g_{t}\|_{P_{t}}^{*})^{2}+\frac{\alpha_{t}^{3}M}{6}\|P_{t}^{-1}g_{t}\|_{T(x_{t})}^{3}+$
$\displaystyle\quad+\langle\nabla
f(x_{t})-g_{t},x_{t+1}-x_{t+1}^{\prime}\rangle-\langle
g_{t},x_{t+1}^{\prime}-x_{t+1}\rangle$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:proj-dir},\eqref{eq:aux-
point}}}{{\leqslant}}f(x_{t})-\alpha_{t}\langle\nabla
f(x_{t}),P_{t}^{-1}g_{t}\rangle+\frac{\alpha_{t}^{2}(1+\Delta_{t})}{2}(\|g_{t}\|_{P_{t}}^{*})^{2}+\frac{\alpha_{t}^{3}M}{6}\|P_{t}^{-1}g_{t}\|_{T(x_{t})}^{3}+$
$\displaystyle\quad+\langle\nabla
f(x_{t})-g_{t},x_{t+1}-x_{t+1}^{\prime}\rangle$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:brackets}}}{{=}}f(x_{t})-\frac{\alpha_{t}}{2}(\|\nabla
f(x_{t})\|_{P_{t}}^{*})^{2}+\left(\frac{\alpha_{t}^{2}(1+\Delta_{t})}{2}-\frac{\alpha_{t}}{2}\right)(\|g_{t}\|_{P_{t}}^{*})^{2}+\frac{\alpha_{t}^{3}M}{6}\|P_{t}^{-1}g_{t}\|_{T(x_{t})}^{3}+$
$\displaystyle\quad+\langle\nabla
f(x_{t})-g_{t},x_{t+1}-x_{t+1}^{\prime}\rangle+\frac{\alpha_{t}}{2}(\|\nabla
f(x_{t})-g_{t}\|_{P_{t}}^{*})^{2}$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:proj-
len}}}{{\leqslant}}f(x_{t})-\frac{\alpha_{t}}{2}(\|\nabla
f(x_{t})\|_{P_{t}}^{*})^{2}+\left(\frac{\alpha_{t}^{2}(1+\Delta_{t})}{2}-\frac{\alpha_{t}}{2}\right)(\|g_{t}\|_{P_{t}}^{*})^{2}+\frac{\alpha_{t}^{3}M}{6}\|P_{t}^{-1}g_{t}\|_{T(x_{t})}^{3}+$
$\displaystyle\quad+\alpha_{t}\|\nabla
f(x_{t})-g_{t}\|_{P_{t}}^{*}\|g_{t}\|_{P_{t}}^{*}+\frac{\alpha_{t}}{2}(\|\nabla
f(x_{t})-g_{t}\|_{P_{t}}^{*})^{2}$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:upb-
var}}}{{\leqslant}}f(x_{t})-\frac{\alpha_{t}}{2}(\|\nabla
f(x_{t})\|_{P_{t}}^{*})^{2}+\left(\frac{\alpha_{t}^{2}(1+\Delta_{t})}{2}-\frac{\alpha_{t}}{4}\right)(\|g_{t}\|_{P_{t}}^{*})^{2}+\frac{\alpha_{t}^{3}M}{6}\|P_{t}^{-1}g_{t}\|_{T(x_{t})}^{3}+$
$\displaystyle\quad+\frac{3\alpha_{t}}{2}(\|\nabla
f(x_{t})-g_{t}\|_{P_{t}}^{*})^{2}$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:upb-
cubic}}}{{\leqslant}}f(x_{t})-\frac{\alpha_{t}}{2}(\|\nabla
f(x_{t})\|_{P_{t}}^{*})^{2}+\frac{\alpha_{t}}{4}Q_{t}(\alpha_{t})(\|g_{t}\|_{P_{t}}^{*})^{2}+\frac{3\alpha_{t}}{2}(\|\nabla
f(x_{t})-g_{t}\|_{P_{t}}^{*})^{2},$ (21)

where $Q_{t}$ is defined as follows:

$Q_{t}(\alpha):=M\|g_{t}\|_{P_{t}}^{*}(1+(\Delta_{t}+\sigma_{T})^{3/2})\cdot\alpha^{2}+2(1+\Delta_{t})\cdot\alpha-1.$
(22)

We choose $\alpha_{t}$ small enough to make $Q_{t}(\alpha_{t})\leqslant 0$.
Since coefficients of $Q_{t}(\alpha)$ of $\alpha^{2}$ and $\alpha$ are
positive, it is sufficient to choose $\alpha_{t}$ less than the larger of two
its roots. Using the fact that $\sqrt{x+y}\leqslant\sqrt{x}+\sqrt{y}$ for any
$x,y\in\mathbb{R}_{+}$, we find that it sufficient to choose $\alpha_{t}$ such
that

$\alpha_{t}\leqslant\frac{1}{1+\Delta_{t}+(M\|g_{t}\|_{P_{t}}^{*})^{1/2}(1+(\Delta_{t}+\sigma_{T})^{3/4})}.$

To make dependency of $\alpha_{t}$ on $\Delta_{t}$ simpler, we apply the fact
that $x^{3/4}\leqslant\frac{3}{4}x+\frac{1}{4}$ for any $x\in\mathbb{R}_{+}$,
to find that it sufficient to choose $\alpha_{t}$ such that

$\alpha_{t}\leqslant\frac{1}{1+\left(1+\frac{3}{4}(M\|g_{t}\|_{P_{t}}^{*})^{1/2}\right)\Delta_{t}+(M\|g_{t}\|_{P_{t}}^{*})^{1/2}\left(\frac{3}{4}\sigma_{T}+\frac{5}{4}\right)}.$
(23)

(21) and (23) imply the following “descent lemma” inequality:

$f(x_{t+1})\leqslant f(x_{t})-\frac{\alpha_{t}}{2}(\|\nabla
f(x_{t})\|_{P_{t}}^{*})^{2}+\frac{3\alpha_{t}}{2}(\|\nabla
f(x_{t})-g_{t}\|_{P_{t}}^{*})^{2}.$ (24)

Rearranging, summing up for $t=0,...,T-1$, and taking into account that
$f(x_{t+1})\geqslant f(x^{*})$,

$\displaystyle\sum_{t=0}^{T-1}\alpha_{t}(\|\nabla f(x_{t})\|_{P_{t}}^{*})^{2}$
$\displaystyle\leqslant
2(f(x_{0})-f(x^{*}))+3\sum_{t=0}^{T-1}\alpha_{t}(\|\nabla
f(x_{t})-g_{t}\|_{P_{t}}^{*})^{2}.$

Taking $\mathbb{E}_{P_{\xi}}$ and $\mathbb{E}_{P_{T}}$, and dividing by
$\sum_{k=0}^{T-1}\alpha_{t}$,

$\mathbb{E}_{P_{T}P_{\xi}}[\|\nabla
f(x_{t})\|_{P_{t}}^{*}]^{2}\leqslant\frac{2(f(x_{0})-f(x^{*}))}{\sum_{k=0}^{T-1}\alpha_{k}}+3\mathbb{E}_{P_{T}P_{\xi}}[\|\nabla
f(x_{t^{\prime}})-v_{t^{\prime}}\|_{P_{t^{\prime}}}^{*}]^{2}.$ (25)

Let us introduce the following denotation:

$\Sigma_{t}:=\sum_{k=0}^{t-1}(\|v_{k}\|_{P_{k}}^{*})^{1/2}.$ (26)
$\displaystyle\frac{2T}{\sum_{k=0}^{T-1}\alpha_{k}}$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:upb-
harmonic}}}{{\leqslant}}\frac{2\sum_{k=0}^{T-1}\frac{1}{\alpha_{k}}}{T}\stackrel{{\scriptstyle\eqref{eq:alpha_sgd}}}{{=}}2+\frac{2\sum_{k=0}^{T-1}(1+\frac{3}{4}(M\|v_{k}\|_{P_{k}}^{*})^{1/2})\Delta_{t}}{T}+M^{1/2}\left(\frac{3\sigma_{T}}{2}+\frac{5}{2}\right)\frac{\Sigma_{T}}{T}$
$\displaystyle=2+\frac{2\sum_{k=0}^{T-1}\Delta_{k}}{T}+M^{1/2}\frac{3}{2}\frac{\sum_{k=0}^{T-1}(\|v_{k}\|_{P_{k}}^{*})^{1/2}\Delta_{t}}{T}+M^{1/2}\left(\frac{3\sigma_{T}}{2}+\frac{5}{2}\right)\frac{\Sigma_{T}}{T},$
(27) $\displaystyle\mathbb{E}_{P_{\eta}}$
$\displaystyle[\|\nabla^{2}f(x_{t})-P_{t}\|_{op}]=\mathbb{E}_{P_{\eta}}[\|\nabla^{2}f(x_{t})-H_{t}+(1-\beta_{t})(H_{t}-P_{t-1})\|_{op}]$
$\displaystyle\qquad\qquad\qquad=\mathbb{E}_{P_{\eta}}[\|\beta_{t}(\nabla^{2}f(x_{t})-H_{t})+(1-\beta_{t})(\nabla^{2}f(x_{t})-P_{t-1})\|_{op}]$
$\displaystyle\qquad\qquad\qquad\stackrel{{\scriptstyle\eqref{eq:hess-
var},\eqref{eq:hess-
cont}}}{{\leqslant}}\beta_{t}\sigma_{H}+(1-\beta_{t})M\|x_{t}-x_{t-1}\|_{T(x_{t-1})}+(1-\beta_{t})\|\nabla^{2}f(x_{t-1})-P_{t-1}\|_{op}$
$\displaystyle\qquad\qquad\qquad\stackrel{{\scriptstyle\eqref{eq:proj-
exp}}}{{\leqslant}}\beta_{t}\sigma_{H}+(1-\beta_{t})M\alpha_{t-1}\|P_{t-1}^{-1}g_{y-1}\|_{T(x_{t-1})}+(1-\beta_{t})\|\nabla^{2}f(x_{t-1})-P_{t-1}\|_{op}$
$\displaystyle\qquad\qquad\qquad\stackrel{{\scriptstyle\eqref{eq:upb-
cubic}}}{{\leqslant}}\beta_{t}\sigma_{H}+(1-\beta_{t})M\alpha_{t-1}\|g_{t-1}\|^{*}_{P_{t-1}}(1+\sigma_{T}^{1/2})+(1-\beta_{t})\|\nabla^{2}f(x_{t-1})-P_{t-1}\|_{op}.$
$\displaystyle\mathbb{E}_{P_{\eta}}[\|\nabla^{2}f(x_{t})-P_{t}\|_{op,P_{t}}]$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:op-b}}}{{\leqslant}}\mu_{t}\beta_{t}\sigma_{H}+\mu_{t}(1-\beta_{t})M\alpha_{t-1}\|g_{t-1}\|_{P_{t-1}}^{*}(1+\sigma_{T}^{1/2})$
$\displaystyle\quad+\frac{\mu_{t}}{\mu_{t-1}}(1-\beta_{t})\|\nabla^{2}f(x_{t-1})-P_{t-1}\|_{op,P_{t-1}}.$

(27), together with (26), implies

$\frac{2T}{\sum_{k=0}^{T-1}\alpha_{k}}\leqslant
2+\frac{2\sum_{k=0}^{T-1}\Delta_{k}}{T}-M^{1/2}\frac{3}{2}\frac{\sum_{k=1}^{T-1}\Sigma_{k}(\Delta_{k-1}-\Delta_{k})}{T}+M^{1/2}\left(\frac{3(\sigma_{T}+\Delta_{0})}{2}+\frac{5}{2}\right)\frac{\Sigma_{T}}{T}.$
(28)

Let us ensure that sequence $\Delta_{t}$ is monotone decreasing. It holds if
$\beta_{t}$ satisfies the following condition:

$\displaystyle\beta_{t}\geqslant
1-\frac{\Delta_{t-1}/\mu_{t}-\sigma_{H}}{\Delta_{t-1}/\mu_{t-1}+M\alpha_{t-1}\|g_{t-1}\|_{P_{t-1}}^{*}(1+\sigma_{T}^{1/2})-\sigma_{H}}.$

Note, that this condition reduces to $\beta_{t}\geqslant 0$, if

$\displaystyle\frac{1}{\mu_{t}}>\frac{1}{\mu_{t-1}}+\frac{M\alpha_{t-1}\|g_{t-1}\|_{P_{t-1}}^{*}(1+\sigma_{T}^{1/2})}{\Delta_{t-1}}.$

Taking $\mathbb{E}_{P_{\eta}}$ in (28), and setting $\beta_{0}=1$,

$\displaystyle\frac{2T}{\sum_{k=0}^{T-1}\alpha_{k}}$ $\displaystyle\leqslant
2+\frac{2\sum_{k=0}^{T-1}\Delta_{k}}{T}+M^{1/2}\frac{3(\sigma_{T}+\Delta_{0})+5}{2}\frac{\Sigma_{T}}{T}$
$\displaystyle\leqslant
2(1+\Delta_{0})+(MB)^{1/2}\frac{3(\sigma_{T}+\Delta_{0})+5}{2}$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:op-b}}}{{\leqslant}}2\left(1+\frac{\sigma_{H}}{\mu_{0}}+(MB)^{1/2}\left(\sigma_{T}+\frac{\sigma_{H}}{\mu_{0}}+2\right)\right).$

∎

###### Proof of Theorem 2 à la [20].

Let us introduce $\Delta_{t}>0$, such that

$\|\nabla^{2}f(x_{t})-P_{t}\|_{op}\leqslant\Delta_{t}.$ (29)

$\displaystyle\alpha_{t}f(x_{t+1})$
$\displaystyle\leqslant\alpha_{t}f(x_{t})+\alpha_{t}\langle\nabla
f(x_{t}),x_{t+1}-x_{t}\rangle+\frac{\alpha_{t}}{2}\|x_{t+1}-x_{t}\|_{\nabla^{2}f(x_{t})}^{2}+\frac{\alpha_{t}M}{6}\|x_{t+1}-x_{t}\|_{T(x_{t})}^{3}$
$\displaystyle\leqslant\alpha_{t}f(x_{t})+\alpha_{t}\langle
g_{t},x_{t+1}-x_{t}\rangle+\frac{\alpha_{t}}{2}\|x_{t+1}-x_{t}\|_{\nabla^{2}f(x_{t})}^{2}+\frac{\alpha_{t}M}{6}\|x_{t+1}-x_{t}\|_{T(x_{t})}^{3}$
$\displaystyle\quad+\alpha_{t}\langle\nabla
f(x_{t})-g_{t},x_{t+1}-x_{t}\rangle$
$\displaystyle\leqslant\alpha_{t}f(x_{t})+\alpha_{t}\langle
g_{t},x_{t+1}-x_{t}\rangle+\frac{1}{2}\|x_{t+1}-x_{t}\|_{P_{t}}^{2}$
$\displaystyle\quad+\frac{\alpha_{t}-1}{2}\|x_{t+1}-x_{t}\|_{\nabla^{2}f(x_{t})}^{2}+\frac{\Delta_{t}}{2}\|x_{t+1}-x_{t}\|_{P_{t}}^{2}$
$\displaystyle\quad+\frac{\alpha_{t}M}{6}\|x_{t+1}-x_{t}\|_{T(x_{t})}^{3}+\alpha_{t}\langle\nabla
f(x_{t})-g_{t},x_{t+1}-x_{t}\rangle$
$\displaystyle\leqslant\alpha_{t}f(x_{t})+\alpha_{t}\langle
g_{t},x-x_{t}\rangle+\frac{1}{2}\|x_{t}-x\|_{P_{t}}^{2}-\frac{1}{2}\|x_{t+1}-x\|_{P_{t}}^{2}+\frac{\alpha_{t}M}{6}\|x_{t+1}-x_{t}\|_{T(x_{t})}^{3}$
$\displaystyle\quad+\alpha_{t}\langle\nabla
f(x_{t})-g_{t},x_{t+1}-x_{t}\rangle-\frac{1-\alpha_{t}(1+\Delta_{t})}{2}\|x_{t+1}-x_{t}\|_{P_{t}}^{2}$
$\displaystyle\leqslant\alpha_{t}f(x_{t})+\alpha_{t}\langle
g_{t},x-x_{t}\rangle+\frac{1}{2}\|x_{t}-x\|_{P_{t}}^{2}-\frac{1}{2}\|x_{t+1}-x\|_{P_{t}}^{2}$
$\displaystyle\quad+\frac{\alpha_{t}M}{6}\|x_{t+1}-x_{t}\|_{T(x_{t})}^{3}+\frac{\alpha_{t}^{2}\|\nabla
f(x_{t})-g_{t}\|^{2}}{2\mu(1-\alpha_{t}(1+\Delta_{t}))}$
$\displaystyle\leqslant\alpha_{t}f(x_{t})+\alpha_{t}\langle\nabla
f(x_{t}),x-x_{t}\rangle+\frac{1}{2}\|x_{t}-x\|_{P_{t}}^{2}-\frac{1}{2}\|x_{t+1}-x\|_{P_{t}}^{2}$
$\displaystyle\quad+\frac{\alpha_{t}M}{6}\|x_{t+1}-x_{t}\|_{T(x_{t})}^{3}+\alpha_{t}\langle
g_{t}-\nabla f(x_{t}),x-x_{t}\rangle+\frac{\alpha_{t}^{2}\|\nabla
f(x_{t})-g_{t}\|^{2}}{2\mu(1-\alpha_{t}(1+\Delta_{t}))}$
$\displaystyle\leqslant\alpha_{t}f(x)+\frac{1}{2}\|x_{t}-x\|_{P_{t}}^{2}-\frac{1}{2}\|x_{t+1}-x\|_{P_{t}}^{2}$
$\displaystyle\quad+\frac{\alpha_{t}M}{6}\|x_{t+1}-x_{t}\|_{T(x_{t})}^{3}+\alpha_{t}\langle
g_{t}-\nabla f(x_{t}),x-x_{t}\rangle+\frac{\alpha_{t}^{2}\|\nabla
f(x_{t})-g_{t}\|^{2}}{2\mu(1-\alpha_{t}(1+\Delta_{t}))}$

Let us denote the latter two terms by the following introduced value:

$\displaystyle E_{t}:=\alpha_{t}\langle g_{t}-\nabla
f(x_{t}),x-x_{t}\rangle+\frac{\alpha_{t}^{2}\|\nabla
f(x_{t})-g_{t}\|^{2}}{2\mu(1-\alpha_{t}(1+\Delta_{t}))}.$
$\displaystyle\alpha_{t}(f(x_{t+1})-f(x^{*}))$
$\displaystyle\leqslant\frac{1}{2}\|x_{t}-x^{*}\|_{P_{t}}^{2}-\frac{1}{2}\|x_{t+1}-x^{*}\|_{P_{t+1}}^{2}+\frac{\alpha_{t}M}{6}\|x_{t+1}-x_{t}\|_{T(x_{t})}^{3}$
$\displaystyle\quad+\frac{1}{2}\langle(P_{t}-P_{t+1})(x_{t+1}-x^{*}),x_{t+1}-x^{*}\rangle+E_{t}$

Summing up for $t=0,...,T-1$ and taking $\mathbb{E}_{P_{T}}$,

$\displaystyle\left(\sum_{t=0}^{T-1}\alpha_{t}\right)\mathbb{E}_{P_{T}}$
$\displaystyle[f(\widetilde{x}_{T})-f(x^{*})]\leqslant\frac{1}{2}\|x_{0}-x^{*}\|_{P_{0}}^{2}-\frac{1}{2}\|x_{T}-x^{*}\|_{P_{T}}^{2}+\sum_{t=0}^{T-1}E_{t}$
$\displaystyle\quad+\frac{M}{6}\sum_{t=0}^{T-1}\alpha_{t}^{4}\|x_{t+1}-x_{t}\|_{T(x_{t})}^{3}+\frac{R^{2}}{2}\sum_{t=0}^{T-1}\|\beta_{t}(P_{t}-H_{t})\|_{op}$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:upb-
cubic},\eqref{eq:alpha_conv}}}{{\leqslant}}\frac{R^{2}}{2}+\sum_{t=0}^{T-1}\alpha_{t}\langle
g_{t}-\nabla
f(x_{t}),x-x_{t}\rangle+\frac{1}{2\mu(1-c_{1})}\sum_{t=0}^{T-1}\alpha_{t}^{2}\|\nabla
f(x_{t})-g_{t}\|^{2}$
$\displaystyle\quad+\frac{M}{4}\sum_{t=0}^{T-1}\alpha_{t}^{4}(\|g_{t}\|_{P_{t}}^{*})^{3}(1+(\Delta_{t}+\sigma_{T})^{3/2})+\frac{R^{2}}{2}\sum_{t=0}^{T-1}\beta_{t}\|H_{t}-P_{t}\|_{op}$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:alpha_conv}}}{{\leqslant}}\frac{R^{2}}{2}+\sum_{t=0}^{T-1}\alpha_{t}\langle
g_{t}-\nabla
f(x_{t}),x-x_{t}\rangle+\frac{1}{2\mu(1-c_{1})}\sum_{t=0}^{T-1}\alpha_{t}^{2}\|\nabla
f(x_{t})-g_{t}\|^{2}$
$\displaystyle\quad+\frac{M}{4}\sum_{t=0}^{T-1}\alpha_{t}^{2}(\|g_{t}\|_{P_{t}}^{*})^{3-2\delta}+\frac{R^{2}}{2}\sum_{t=0}^{T-1}\beta_{t}\|H_{t}-P_{t}\|_{op},$
(30)

where it is required that $\alpha_{t}$ satisfies the following inequality for
some $c_{1}\in(0,1)$ and $\delta\in(0,3/2)$:

$\displaystyle\alpha_{t}$
$\displaystyle\leqslant\frac{1}{\gamma_{t}}:=\frac{1}{\frac{1}{c_{1}}(1+\Delta_{t})+(\|g_{t}\|_{P_{t}}^{*})^{\delta}(1+(\Delta_{t}+\sigma_{T})^{3/4})}$
$\displaystyle\leqslant\min\left\\{\frac{c_{1}}{1+\Delta_{t}},\frac{1}{(\|g_{t}\|_{P_{t}}^{*})^{\delta}(1+(\Delta_{t}+\sigma_{T})^{3/4})}\right\\}.$
(31)

Substituting $x=x^{*}$ into (30), taking $\mathbb{E}_{P_{\xi}}$, and using
(4),

$\left(\sum_{t=0}^{T-1}\alpha_{t}\right)\mathbb{E}_{P_{\xi}P_{T}}[f(x_{t})-f(x^{*})]\leqslant\frac{R^{2}}{2}+\left(\frac{\sigma_{g}^{2}}{2\mu(1-c_{1})}+\frac{MB^{3-2\delta}}{4}\right)\sum_{t=0}^{T-1}\alpha_{t}^{2}+\frac{R^{2}}{2}\sum_{k=0}^{T-1}\beta_{k}\|H_{k}-P_{k}\|_{op},$
(32)

where $B=\max_{k=0,...,T-1}\|g_{t}\|_{P_{t}}^{*}$. For any $t>0$,

$\displaystyle\mathbb{E}_{P_{\eta}}[\|H_{t}-P_{t}\|_{op}]$
$\displaystyle=\mathbb{E}_{P_{\eta}}[\|H_{t}-P_{t-1}-\beta_{t}(H_{t}-P_{t-1})\|_{op}]=(1-\beta_{t})\mathbb{E}_{P_{\eta}}[\|H_{t}-P_{t-1}\|_{op}]$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:hess-
var}}}{{\leqslant}}(1-\beta_{t})(\|H_{t-1}-P_{t-1}\|_{op}+\|\nabla^{2}f(x_{t})-\nabla^{2}f(x_{t-1})\|_{op}+2\sigma_{H})$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:hess-
cont}}}{{\leqslant}}(1-\beta_{t})(\|H_{t-1}-P_{t-1}\|_{op}+M\|x_{t}-x_{t-1}\|_{T(x_{t-1})}+2\sigma_{H})$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:proj-
exp}}}{{\leqslant}}(1-\beta_{t})(\|H_{t-1}-P_{t-1}\|_{op}+M\|\alpha_{t-1}P_{t-1}^{-1}g_{t-1}\|_{T(x_{t-1})}+2\sigma_{H})$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:upb-
cubic}}}{{\leqslant}}(1-\beta_{t})(\|H_{t-1}-P_{t-1}\|_{op}+\alpha_{t-1}M\|g_{t-1}\|_{P_{t-1}}^{*}(1+\sigma_{T}^{1/2})+2\sigma_{H})$
(33)

Taking $\mathbb{E}_{P_{\eta}}$ in (32), substituting (33) into (32), and
taking $\mathbb{E}_{P_{\eta}}$ again,

$\displaystyle\left(\sum_{k=0}^{T-1}\alpha_{k}\right)\mathbb{E}_{P_{\eta}P_{\xi}P_{T}}$
$\displaystyle[f(x_{t})-f(x^{*})]\leqslant\frac{R^{2}}{2}+\left(\frac{\sigma_{g}^{2}}{2\mu(1-c_{1})}+\frac{MB^{3-2\delta}}{4}\right)\sum_{k=0}^{T-1}\alpha_{k}^{2}+\sigma_{H}R^{2}\sum_{k=1}^{T-1}\beta_{k}(1-\beta_{k})$
$\displaystyle\quad+\frac{MR^{2}(1+\sigma_{T}^{1/2})}{2}\sum_{k=1}^{T-1}\beta_{k}(1-\beta_{k})\alpha_{k-1}\|g_{k-1}\|_{P_{k-1}}^{*}$
$\displaystyle\quad+\frac{R^{2}}{2}\sum_{k=1}^{T-1}\beta_{k}(1-\beta_{k})\mathbb{E}_{P_{\eta}}[\|H_{k-1}-P_{k-1}\|_{op}]+(1-\beta_{0})\|H_{0}-P_{-1}\|_{op}$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:alpha_conv}}}{{\leqslant}}\frac{R^{2}}{2}+(1-\beta_{0})\|H_{0}-P_{-1}\|_{op}+\left(\frac{\sigma_{g}^{2}}{2\mu(1-c_{1})}+\frac{MB^{3-2\delta}}{4}\right)\sum_{k=0}^{T-1}\alpha_{k}^{2}$
$\displaystyle\quad+R^{2}\left(\sigma_{H}+\frac{M(1+\sigma_{T}^{1/2})}{2}\right)\sum_{k=1}^{T-1}\beta_{k}(1-\beta_{k})$
$\displaystyle\quad+\frac{R^{2}}{2}\sum_{k=1}^{T-1}\beta_{k}(1-\beta_{k})\mathbb{E}_{P_{\eta}}[\|H_{k-1}-P_{k-1}\|_{op}]$
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:beta-conv},\eqref{eq:abel-
test}}}{{\leqslant}}R^{2}\left(c_{3}\sigma_{H}+\frac{c_{3}M(1+\sigma_{T}^{1/2})+c_{2}+1}{2}\right)+(1-\beta_{0})\|H_{0}-P_{-1}\|_{op}$
$\displaystyle\quad+\left(\frac{\sigma_{g}^{2}}{2\mu(1-c_{1})}+\frac{MB^{3-2\delta}}{4}\right)\sum_{k=0}^{T-1}\alpha_{k}^{2}$
(34)

where it is required that $\beta_{t}$ satisfies the following condition:

$\frac{1}{1-\beta_{t}}\geqslant\max\left\\{1+\frac{\alpha_{t-1}M\|g_{t-1}\|_{P_{t-1}}^{*}(1+\sigma_{T}^{1/2})+2\sigma_{H}}{\|H_{t-1}-P_{t-1}\|_{op}},(t+1)^{2}\right\\}.$
(35)

Setting $\beta_{0}=1$ and $\delta=\frac{3}{4}$, (34) implies

$\displaystyle\mathbb{E}_{P_{\eta}P_{\xi}P_{T}}[f(x_{t})-f(x^{*})]$
$\displaystyle\quad\leqslant\frac{R^{2}}{\alpha\sum_{k=0}^{T-1}\frac{1}{\gamma_{k}}}\left(c_{3}\sigma_{H}+\frac{c_{3}M(1+\sigma_{T}^{1/2})+c_{2}+1}{2}\right)+\alpha\frac{\sum_{k=0}^{T-1}\frac{1}{\gamma_{k}^{2}}}{\sum_{k=0}^{T-1}\frac{1}{\gamma_{k}}}\left(\frac{\sigma_{g}^{2}}{2\mu(1-c_{1})}+\frac{MB^{3-2\delta}}{4}\right)$
$\displaystyle\quad\stackrel{{\scriptstyle\eqref{eq:alpha_conv}}}{{\leqslant}}\frac{R^{2}}{\alpha\sum_{k=0}^{T-1}\frac{1}{\gamma_{k}}}\left(c_{3}\sigma_{H}+\frac{c_{3}M(1+\sigma_{T}^{1/2})+c_{2}+1}{2}\right)+\alpha\frac{T}{\sum_{k=0}^{T-1}\frac{1}{\gamma_{k}}}\left(\frac{\sigma_{g}^{2}}{2\mu(1-c_{1})}+\frac{MB^{3-2\delta}}{4}\right)$
$\displaystyle\quad\stackrel{{\scriptstyle\eqref{eq:alpha-
cvx}}}{{\leqslant}}O(1)\left(\frac{R^{2}(\sigma_{H}+M(1+\sigma_{T}^{1/2})+1){\displaystyle\max_{k=0,...,T-1}\gamma_{k}}}{T}+\sqrt{\frac{(\sigma_{g}^{2}/\mu+MB^{3-2\delta})(\sigma_{H}+M(1+\sigma_{T}^{1/2})+1)}{T}}\right)$
$\displaystyle\quad\leqslant
O(1)\left(1+\sigma_{H}+(1+(\sigma_{H}+\sigma_{T})^{3/4})B^{3/4}\right)\frac{R^{2}(\sigma_{H}+M(1+\sigma_{T}^{1/2})+1)}{T}$
$\displaystyle\quad\quad+O(1)\sqrt{\frac{(\sigma_{g}^{2}/\mu+MB^{3/2})(\sigma_{H}+M(1+\sigma_{T}^{1/2})+1)}{T}},$

where

$\alpha_{t}=\min\left\\{\frac{1}{\gamma_{t}},\sqrt{\frac{2c_{3}\sigma_{H}+c_{3}M(1+\sigma_{T}^{1/2})+c_{2}+1}{T\left(\frac{\sigma_{g}^{2}}{\mu(1-c_{1})}+\frac{M}{2}\right)}}\right\\}.$
(36)

∎


# High Dimensional Causal Inference
with Variational Backdoor Adjustment

Daniel Israel
CS Department
UCLA
<EMAIL_ADDRESS>
&Aditya Grover∗
CS Department
UCLA
<EMAIL_ADDRESS>
&Guy Van den Broeck∗
CS Department
UCLA
<EMAIL_ADDRESS>
&∗ Equal Contribution

###### Abstract

Backdoor adjustment is a technique in causal inference for estimating
interventional quantities from purely observational data. For example, in
medical settings, backdoor adjustment can be used to control for confounding
and estimate the effectiveness of a treatment. However, high dimensional
treatments and confounders pose a series of potential pitfalls: tractability,
identifiability, optimization. In this work, we take a generative modeling
approach to backdoor adjustment for high dimensional treatments and
confounders. We cast backdoor adjustment as an optimization problem in
variational inference without reliance on proxy variables and hidden
confounders. Empirically, our method is able to estimate interventional
likelihood in a variety of high dimensional settings, including semi-synthetic
X-ray medical data. To the best of our knowledge, this is the first
application of backdoor adjustment in which all the relevant variables are
high dimensional.

##  1 Introduction

Understanding causal relationships is central to many scientific disciplines
such as healthcare and economics. In these settings, professionals need to
assess the importance of interventions (e.g. treatments, policy changes) on
societal outcomes (e.g. patient health, economic well-being). However,
determining interventional effects is challenging due to the presence of
confounders, such as age, gender, and income of participants. One approach is
to directly collect interventional data using randomized control trials
(RCTs), but RCTs must be designed carefully (Schulz & Grimes, 2002) and are
costly with respect to time and money (Sørensen et al., 2006). As an
alternative, causal inference can be used to estimate interventional
quantities from observational data alone (Pearl, 1995). Given modeling
assumptions encoded in a DAG, Pearl’s backdoor adjustment (Pearl, 2009b)
estimates interventional likelihood by re-weighting the observational
likelihood of the outcomes to block the influence of confounding variables.

This paper makes progress towards the application of backdoor adjustments for
high-dimensional datasets. In many real-world scenarios, we frequently
encounter treatments, outcomes, and confounders as high dimensional objects.
For example, it is often the case that treatments are text embeddings of
recommended procedures, outcomes are images such as X-ray and MRI screenings,
and confounders are high dimensional genetic and environmental factors. High
dimensional backdoor adjustment suffers from two major challenges:
intractability in integrating out high-dimensional confounders and
expressivity in learning non-linear dependencies from the observational
distributions.

Prior work in this space can partly address these challenges. Inverse
propensity weighting methods (Austin, 2011) can be used for high dimensional
confounders, but do not generalize well to high dimensional treatments.
Another common approach is using variational autoencoders (VAEs) (Kingma &
Welling, 2013) as models of proxy confounding, where CEVAE is a notable
example (Louizos et al., 2017). These models are overly expressive, in a
sense, because they model unobserved confounders that are unidentifiable
(Rissanen & Marttinen, 2021).

We propose Variational Backdoor Adjustment (VBA), a novel variational approach
for backdoor adjustment under direct confounding. Our method uses variational
inference to optimize a lower bound of the interventional likelihood and
compute backdoor adjustment in high dimensions. We encounter and overcome the
following key challenges:

  1. 1.

For high dimensional confounders and treatments, backdoor adjustment is either
intractable to compute exactly or impractical to estimate with sampling due to
high variance. We apply variational inference to efficiently obtain a more
accurate approximate estimate.

  2. 2.

Variational inference is typically applied in a latent variable setting in
which the confounder in question is unobserved. Because the latent variable is
unidentifiable, applying backdoor adjustment leads to an inconsistent
estimate. We apply variational inference such that the latent space is
restricted to the observed confounder.

  3. 3.

However, in the new regime we cannot jointly optimize over all parameters, as
is typical in VAEs. We introduce an optimization method that respects a latent
space constrained by observation.

In our framework, we define three distributional components: encoder, decoder,
and prior. Note that although we utilize standard VAE terminology to invoke
similar components, our method is not to be confused with VAE due to different
assumptions. These components must first be optimized separately to ensure
identifiability, such that backdoor adjustment is correctly computed over the
observed distributions. Once the components are optimized separately, the
encoder can be further optimized to obtain better interventional density
estimates.

This framework proves empirically effective at computing backdoor adjustment
in a variety of synthetic and semi-synthetic settings. We construct a high
dimensional linear Gaussian training set to empirically verify interventional
density estimation with VBA. In this setting, we show that our optimization
technique provides significant improvement over more naive estimation
strategies. We then construct image datasets to test backdoor adjustment in
more nonlinear settings. As a proof of concept, we demonstrate a potential
application of VBA on X-ray medical data. We once again show in this setting
the importance of the optimization technique used in VBA.

$X$$Y$$Z$ (a) Direct Confounding

$X$$Y$$Z$ (b) Direct Confounding $do(X)$

$X$$Y$$U$$Z$ (c) Proxy Model Confounding

Figure 1: A visualization of the various discussed DAGs. Figure (a) and (c)
show different approaches to modeling confounding. In this work, we opt for
(a) due to identifiability. Figure (b) shows the $do$ operation graphically.

##  2 Related Work

The existing body of literature pertaining to treatment effects is extensive
(Shalit et al., 2017; Shi et al., 2019; Chen et al., 2019; Tesei et al.,
2023). While these methods relate to individual treatment effects, backdoor
adjustment is applied at the population level. Recent work uses neural mean
embeddings (Xu & Gretton, 2022) for backdoor adjustment. While related, it is
not applied to high dimensional outcomes. We shall also address related works
that primarily fall into two approaches.

#### Propensity Score Methods

In the literature, a common way to handle high dimensional confounding is by
using propensity score methods (Austin, 2011). The key idea is that if a high
dimensional confounder $z$ is sufficient for backdoor adjustment, then a
single dimensional propensity score $g(z)=p(X=1\mid z)$ will suffice for
backdoor adjustment (Rosenbaum & Rubin, 1983). Note that $X$ in this case is
assumed to be a binary treatment. To our knowledge, there is no work applying
propensity score methods to high dimensional treatments. While there may be no
reason in principle that propensity score methods cannot be applied to such
scenarios, they do not address the fundamental inference problem posed when
$X$ is high dimensional; namely, that the conditional density $p(y\mid
x,g(z))$ will have high variance due to a small effective sample size for each
treatment. Shifting the burden of adjustment onto $g$ will not solve this
issue for high dimensional treatments. Hence, we argue for the use of
variational inference.

#### Proxy Model VAE and Identifiability

Past work applies variational autoencoders to causal inference, but with
different modeling assumptions. Methods such as CEVAE (Louizos et al., 2017),
VSR (Zou et al., 2020), and TEDVAE (Zhang et al., 2021) operate under the
assumption that the true confounder is unobserved and only a proxy of the
latent variable is observed, as depicted in Figure 1(c). Recent work by
Rissanen & Marttinen (2021) demonstrates that these techniques do not yield
consistent estimates of causal effect because in general they cannot, in
practice or even in theory, model unobserved latent variables correctly.
Vanilla VAEs are unidentifiable (Locatello et al., 2019), which means that
there exists infinite number of transformations on the latent variable that
would emit the same marginal distribution. For this reason, performing
backdoor adjustment over an unidentifiable latent variable can lead to
completely inaccurate results. Recent works such as iVAE (Khemakhem et al.,
2020) show that a factorized prior conditioned on additional observations give
identifiability up to a class of transformations, and Intact-VAE (Wu &
Fukumizu, 2021) exploits these results. In practice, the assumptions needed
for identifiability cannot be verified for unobserved variables in real-world
data. Rather than modeling a proxy of an unobserved confounder, we perform
causal inference over observed variables only, as seen in Figure 1(a). While
the proxy model approach must argue that unobserved variables are
identifiable, our approach is identifiable by definition because all variables
are observed.

##  3 Approach

###  3.1 Preliminaries

We are given variables $X,Y,Z$ as depicted in Figure 1(a). We shall refer to
$X,Y,Z$ respectively as treatment, outcome, and confounder. Suppose we are
given probability density functions $p(z)=p(Z=z)$ and $p(y\mid x,z)=p(Y=y\mid
X=x,Z=z)$. We can apply Pearl’s backdoor adjustment (Pearl, 2009b) to obtain
the interventional density

$p(y\mid do(x))=\sum_{z}p(z)p(y\mid x,z).$ (1)

In practice, we will not know the prior $p(z)$ and outcome likelihood $p(y\mid
x,z)$ but will assume access to an observational dataset
$D={(x_{i},y_{i},z_{i})_{i=1}^{n}}$ consisting of $n$ triplets for $(x,y,z)$.
The $do$ operator represents an intervention on the value of $X$, which severs
the influence of $Z$ on $X$. The result is the distribution induced by
modifying the DAG as seen in Figure 1(b). Observe that the right-hand-side of
the equation does not contain the $do$ operator, thus allowing interventional
quantities to be obtained from observational data alone. Note that backdoor
adjustment applies more broadly than to the DAG in Figure 1. In general,
backdoor adjustment can be applied if the variables in question satisfy the
backdoor criterion (Pearl, 2009a), which slightly broadens the scope of this
work, as many DAGs can be equivalent to Figure 1(a) with respect to backdoor
adjustment.

This formula alone is satisfactory in low dimensional settings when the number
of confounders is limited. In high dimensional settings however, it is known
that exact inference is intractable and the sum will be exponential in the
dimension of $Z$ (Wang & Kwiatkowska, 2022). Naturally, for higher dimensions
one may learn a conditional distribution $p(y\mid x,z)$ and generative model
of the confounder $p(z)$. The interventional density $p(y\mid do(x))$ would
then be approximated by drawing samples from the confounder $z\sim p(z)$ and
computing an expectation over $p(y\mid x,z)$. However, it is well-known that a
naive Monte Carlo estimate will have high variance in large dimensions because
a sampled high dimensional $Z$ will almost never give high probability to
$p(y\mid x,z)$ for a chosen $Y$ and $X$. Thus, even with a perfectly learned
$p(z)$ and $p(y\mid x,z)$, inference of $p(y\mid do(x))$ using naive sampling
is impractical.

###  3.2 Variational Backdoor Adjustment

Variational inference is a commonly used technique for generative modeling
tasks (Kingma & Welling, 2013; Rezende & Mohamed, 2015). In variational
inference, the key idea is to approximate the posterior $p(z\mid x)$ with a
simpler distribution, thus allowing marginal likelihood to be approximated
with an evidence lower bound (ELBO). With variational inference, we can
estimate a lower bound on the interventional density given in Equation 1.
Given some auxiliary “simple” distribution $q$, we have

$\displaystyle\log p(y\mid do(x))$ $\displaystyle=\log\sum_{z}p(z)p(y\mid
x,z)$ (2) $\displaystyle=\log\mathbb{E}_{q(z\mid x,y)}\left(\frac{p(z)p(y\mid
x,z)}{q(z\mid x,y)}\right)$ (3) $\displaystyle\geq\mathbb{E}_{q(z\mid
x,y)}\left(\log p(z)+\log p(y\mid x,z)-\log q(z\mid x,y)\right).$ (4)

The inequality in Equation 4 is given by Jensen’s inequality. In theory it is
possible to condition $q$ on exclusively $x$ or $y$, but $q$ will be learned
so it is best to parameterize it in the most expressive way. Intuitively, the
purpose of $q(z\mid x,y)$ is as an encoder, giving high probability samples of
$Z$, which will help decrease variance in high dimensions. Contrast this with
the aforementioned Monte Carlo estimate in which we sampled from $p(z)$,
leading to high variance. The penalty incurred by the encoder will be its KL
divergence with the true prior distribution $p(z)$. It is well known in
probabilistic inference that sampling is unbiased but has high variance, while
variational inference will have some bias in exchange for much lower variance
(Lange et al., 2022). Note that while our method utilizes variational
inference, it is not a standard VAE because in Equation 4 the interpretation
of $z$ is not as a latent variable. Instead, we have observed data on $Z$.

###  3.3 Optimization Method

Because $z$ is not a latent variable, we cannot optimize the lower bound in
the same manner as a VAE. To illustrate this, we shall first introduce our two
step optimization method to perform Variational Backdoor Adjustment (VBA).
After introducing our method and giving notation, it will become more clear
why we cannot apply standard VAE “joint” optimization in our setting.

#### Separate Training Phase

Let $p(z)$ and $p(y\mid x,z)$ be parameterized with a model of the prior
$p_{\theta}(z)$ and decoder model $p_{\gamma}(y\mid x,z)$. We optimize these
models with maximum likelihood objectives. Recall that the observational data
is given by $D={(x_{i},y_{i},z_{i})_{i=1}^{n}}$. Each loss respectively will
be computed over the dataset as $\mathbb{E}_{D}(\mathcal{L}(*))$.

$\mathcal{L}^{\text{MLE}}_{\theta}(z)=-\log p_{\theta}(z)$ (5)
$\mathcal{L}^{\text{MLE}}_{\gamma}(x,y,z)=-\log p_{\gamma}(y\mid x,z)$ (6)

Let $q_{\phi}(z\mid x,y)$ be an encoder model. It can also be optimized with
maximum likelihood training, although it need not be limited in this way.

$\mathcal{L}^{\text{MLE}}_{\phi}(x,y,z)=-\log q_{\phi}(z\mid x,y)$ (7)

We can optimize these components separately using a gradient-based optimizer
and obtain a lower bound on the interventional likelihood by plugging in the
components as suggested by Equation 4. The model optimized by such loss
functions is

$f_{\phi,\theta,\gamma}(x,y)=\mathbb{E}_{q_{\phi}(z\mid x,y)}\left(\log
p_{\theta}(z)+\log p_{\gamma}(y\mid x,z)-\log q_{\phi}(z\mid x,y)\right)$ (8)

This step of VBA will be referred to as separate training.

#### Finetuning Phase

Separate optimization will not lead to the tightest lower bound on
interventional likelihood. We can further optimize the encoder to obtain more
accurate interventional density estimation. Let $\hat{\theta}$ and
$\hat{\gamma}$ be parameters optimizing their own respective MLE objectives in
separate training. We give the following finetuning objective, which can be
optimized with respect to $x,y\sim D$

$\mathcal{L}^{\text{ELBO}}_{\phi}(x,y)=-\sum\nolimits_{z^{\prime}_{1},...,z^{\prime}_{n}\sim
q_{\phi}(z\mid
x,y)}\left(\mathcal{L}^{\text{MLE}}_{\hat{\theta}}(z^{\prime}_{j})+\mathcal{L}^{\text{MLE}}_{\hat{\gamma}}(x,y,z^{\prime}_{j})-\mathcal{L}^{\text{MLE}}_{\phi}(x,y,z^{\prime}_{j})\right)$
(9)

Following separate training, we now optimize $\mathcal{L}^{\text{ELBO}}$ with
gradient descent. We refer to this step as finetuning the encoder. This loss
will optimize the expectation seen in Equation 4. Observe that the loss
function differs from the typical ELBO objective seen in VAEs because the
parameters $\theta,\gamma$ of the “decoder” $p_{\gamma}(y\mid x,z)$ and
“prior” $p_{\theta}(z)$ are held fixed. We show experimentally that the
additional finetuning phase improves finite sample performance compared to MLE
separate training alone. The mechanism for this improvement is that the
encoder is able to increase likelihood of the decoder $p(y\mid x,z)$ or prior
$p(z)$ by generating less or more likely $z^{\prime}$, as determined by the
ELBO loss.

#### Pitfall of Joint Optimization

Instead of fixing $\theta$ and $\gamma$ to their maximum likelihood estimates
as shown in Equation 9, it may seem tempting to optimize all of the parameters
“jointly” as follows

$\mathcal{L}^{\text{ELBO}}_{\phi,\theta,\gamma}(x,y)=-\sum\nolimits_{z^{\prime}_{1},...,z^{\prime}_{n}\sim
q_{\phi}(z\mid
x,y)}\left(\mathcal{L}^{\text{MLE}}_{\theta}(z^{\prime}_{j})+\mathcal{L}^{\text{MLE}}_{\gamma}(x,y,z^{\prime}_{j})-\mathcal{L}^{\text{MLE}}_{\phi}(x,y,z^{\prime}_{j})\right)$
(10)

As previously explained, for VAEs the interpretation of $z$ in Equation 4 is
as a latent variable. However, in our case, $z$ represents observed data,
therefore backdoor adjustment necessitates that those quantities must be
maximum likelihood estimates. Equation 10 is not guaranteed to converge to
$\log p(y\mid do(x))$.

###### Proposition 1

There exists a Structural Causal Model (SCM) over $X,Y,Z$ such that for the
optimal parameters
$\phi^{*},\theta^{*},\gamma^{*}=\operatorname*{arg\,min}\mathcal{L}^{\text{ELBO}}_{\phi,\theta,\gamma}$,
we have $f_{\phi^{*},\theta^{*},\gamma^{*}}(x,y)=\log p(y\mid x)$.

Proof of this proposition is in Appendix B. Intuitively, “joint” optimization
does not work because the distribution over $Z$ will become untethered from
the data, which means we are no longer performing backdoor adjustment. In the
loss of Equation 10, because $\theta$ and $\gamma$ are not fixed, we can
minimize the loss by making the KL divergence between $q_{\phi}(z\mid x,y)$
and $p_{\theta}(z)$ arbitrarily small while $p_{\gamma}(y\mid x,z)$ loses
dependence on $Z$. This is very similar to a problem in VAEs called posterior
collapse (Lucas et al., 2019), but in our situation it makes estimating
backdoor adjustment impossible.

#### Implementation Details

In this framework, each distribution can be parameterized as needed depending
on the setting. In general, $p_{\theta}(z)$ will be a generative model to
estimate the density of a high dimensional confounder. Depending on the
circumstance, $p_{\gamma}(y\mid x,z)$ and $q_{\phi}(z\mid x,y)$ can be
conditional generative models or vanilla neural networks. In most cases
vanilla neural networks are sufficient, but in some cases $Z$ or $Y$ may
contain significant exogenous stochasticity that will be better captured with
conditional generative models.

##  4 Experiments

Evaluation is one of the most challenging aspects to causal inference because
the real-world does not grant access to ground truth causal effect. This is
because we do not have a complete understanding of the data generating
process, so if some outcome is observed, we do not have access to the
counterfactual outcome given different circumstances. Some refer to this issue
as the fundamental problem of causal inference (Holland, 1986).

To gain access to ground truth, one must make assumptions about the data
generating process at the expense of realism. In this work, we attempt a
balanced approach by evaluating at different points along this tradeoff.
First, we evaluate our method on high dimensional linear Gaussian data. In
this setting, ground truth causal effect can be solved for analytically. For a
more realistic high dimensional setting that incorporates image data, we
design a toy example involving MNIST images. Unfortunately, this setting does
not have ground truth causal effect, but we can demonstrate qualitative
improvements when using backdoor adjustment over naively modelling the
observational distribution $p(y\mid x)$. Finally, we apply variational
backdoor adjustment to a synthetic dataset containing real X-ray images and
show the benefit of finetuning over separate training. In our empirical image
experiments, we can only make qualitative observations about whether we
estimate the true interventional density. Our main claim to the accuracy of
VBA relies on experiments using linear Gaussian data. The purpose of these
experiments can be summarized as follows:

  * •

Linear Gaussian: shows that VBA is more accurate than sampling, and it
consistently converges to the correct estimate.

  * •

MNIST: gives a high dimensional image setting and compares VBA, which
estimates the interventional distribution, to naively modelling the
observational distribution.

  * •

X-ray: presents a medical scenario and demonstrates the benefit of the
“finetuning” mechanism over “separate training” alone.

The relevant code for reproducability can be found on Github.

###  4.1 Linear Gaussian

To evaluate the VBA density estimation against ground truth, we consider the
following setting. Suppose the relationships between variables are linear and
the base distributions are Gaussian. The structural equations for $X,Y,Z$ will
take the following form

$Z:=\mathcal{N}(0,1)\hskip 28.45274ptX:=c_{1}Z+\mathcal{N}(0,\sigma_{1})\hskip
28.45274ptY:=c_{2}X+c_{3}Z+\mathcal{N}(0,\sigma_{2})$ (11)

The linear Gaussian setting allows for fully analytic solutions to ground
truth interventional density. The ground truth for interventional distribution
is given by

$y\mid
do(x)\sim\mathcal{N}\left(c_{2}x,\sqrt{c_{3}^{2}+\sigma_{2}^{2}}\right)$ (12)

and we can compare it with the observational distribution

$y\mid
x\sim\mathcal{N}\left(\left(\frac{c_{1}c_{3}}{(c_{1}+\sigma_{1}^{2})}+c_{2}\right)x,\sqrt{\frac{c_{3}^{2}\sigma_{1}^{2}}{c_{1}^{2}+\sigma_{1}^{2}}+\sigma_{2}^{2}}\right)$
(13)

(a) Sampling vs Variational Inference
Log Error Comparison

(b) Finetuning vs Separate Training
Log Error Comparison

Figure 2: Figures (a) and (b) contrast variational inference with naive sampling at inference time. Variational inference proves beneficial, especially as dimensionality increases. Finetuning yields a better interventional density estimate than separate training. We average results over 10 random generations of data. Table 1: Linear Gaussian Results. We report values in nats and report the mean and standard error over 5 random seeds. The top row is mean absolute error from the ground truth for $\mathbb{E}_{x,y\sim D}\log p(y|do(x))$. Rows 2-4 below are decomposed from the ELBO. Row 5 gives the likelihood of the encoder. Finetuning the encoder $q$ gives better results as opposed to generating maximum likelihood $z$. Lower likelihood samples of $z$ results in a tighter ELBO. | In Distribution |  | Out of Distribution |
---|---|---|---|---
| Separate Training | Finetuned | Separate Training | Finetuned
Ground Truth MAE ($\downarrow$) | 5.58 (0.003) | 1.89 (0.001) | 11.74 (0.005) | 7.51 (0.001)
$\mathbb{E}_{x,y\sim D}\mathbb{E}_{q(z\mid x,y)}\log p(z)$ | -21.17 (0.003) | -20.83 (0.000) | -129.41 (0.001) | -110.82 (0.000)
$\mathbb{E}_{x,y\sim D}\mathbb{E}_{q(z\mid x,y)}\log p(y\mid x,z)$ | -21.80 (0.001) | -22.14 (0.002) | -26.17 (0.003) | -34.13 (0.006)
$\mathbb{E}_{x,y\sim D}\mathbb{E}_{q(z\mid x,y)}\log q(z\mid x,y)$ | 17.20 (0.002) | 6.27 (0.001) | 17.20 (0.002) | 6.28 (0.001)
$\mathbb{E}_{D}[\log q(z\mid x,y)]$ | 14.97 | 4.17 | 3.28 | -25.02

Constants $c_{1},c_{2},c_{3}$ can be thought of as determining the strength of
the causal relationships. It therefore follows that the mean of Equation 12 is
only dependent on $c_{2}$ because $c_{2}$ determines the causal relationship
between $X$ and $Y$. The mean of Equation 13 also contains $c_{1}$ and
$c_{3}$, which are responsible for confounding. We give a concrete example of
the difference between observational and interventional distribution in Figure
7 of Appendix A. So as not to draw conclusions based on cherry-picked
parameters, for each experiment we randomly sample and fix
$c_{1},c_{2},c_{3}$. We sample these constants in a manner that emphasizes the
difference between interventional and observational distribution. Linear
Gaussian hyperparameters are sampled as

$c_{1}\sim\mathcal{U}(0,5)\hskip 14.22636ptc_{2}\sim\mathcal{U}(0,3)\hskip
14.22636ptc_{3}\sim\mathcal{U}(-5,-10)\hskip
14.22636pt\sigma_{1}\sim\mathcal{U}(0,1)\hskip
14.22636pt\sigma_{2}\sim\mathcal{U}(0,3).$ (14)

We obtain a high dimensional linear Gaussian dataset by repeatedly sampling a
set of parameters for each dimension and treating each dimension
independently. To perform backdoor adjustment on linear Gaussian data, we
model the confounder with a VAE and both encoder and decoder with Gaussians
parameterized by MLPs.

Figure 2(a) shows the difference between sampling and variational inference.
For 10 unique generations of linear Gaussian data, we train the respective
components using their own MLE objective (also known as separate training). At
inference time, we compare two approaches for obtaining $\log p(y\mid do(x))$:
sampling from the prior $z\sim p(z)$ and computing an average over $\log
p(y\mid x,z)$, or sampling from the encoder $z\sim q(z\mid x,y)$ and utilizing
variational inference given by Equation 4. Variational inference outperforms
naive sampling as the dimensionality increases.

We use the same analysis to compare optimization methods. Figures 8(b) and
2(b) show a substantive increase in performance when finetuning the encoder,
and the performance is robust in higher dimensions. To better understand why
finetuning gives better estimates, we focus on a single run at a fixed
dimension. We examine 15 dimensional linear Gaussian data generated with the
process previously described, and Table 1 demonstrates that utilizing the
finetuning objective for the encoder achieves log-likelihood estimates closer
to the ground truth. We generate data out-of-distribution by intervening on
$Z$ and setting it to a $\mathcal{U}(-7,7)$ variable. This in turn generates
out-of-distribution $X$ and $Y$ by construction. These results analyze the
components given by the formula for variational backdoor adjustment (Equation
4) and show the reason for increased performance: optimizing the encoder gives
it the flexibility to give lower likelihood to samples without dramatically
sacrificing likelihood of the decoder and prior. This optimization results in
a tighter bound on Equation 4, and therefore a more accurate interventional
likelihood.

###  4.2 MNIST

To demonstrate that our method can work on images, we construct a synthetic
dataset of binary MNIST images that simulate confounding. The setup is as
follows: sample a random image $Z$, sample an image $X$ such that $X$ is the
consecutive number after $Z$, and concatenate them together to form image $Y$.
The causal structure is depicted in Figure 3. In the training data, $Y$ will
only contain concatenated images with successive numbers due to the
confounding of image $Z$. However, by training with backdoor adjustment, the
model can learn that $Y$ can consist of any two numbers given an intervention
on $X$.

Figure 3: Data generating process for MNIST dataset. On the left is the
observational distribution in which $Y$ can only have successive numbers. On
the right, interventions on $X$ can produce images $Y$ with two numbers not
seen on the left.

(a) Observational $p(y\mid x)$

(b) Interventional $p(y\mid do(x))$

(c) Difference ((a) - (b))

Figure 4: We generate interventional data for MNIST as seen on the right of
Figure 3. Figures (a) and (b) show the respective log-likelihood, where the
axes delineate how $Y$ is generated. For example, $Z=3$ and $do(X=9)$ will
produce an image $Y$ reassembling a “39”. The key result is (c), where we see
that in general, interventional likelihood will give higher likelihoods. The
blue diagonal in (c) shows that images with consecutive numbers have higher
likelihood in the observational distribution.

For observational density $p(y\mid x)$, we train a conditional VAE with 50
latents and 2 hidden layers, each with 200 hidden units. For the
interventional density $p(y\mid do(x))$, we apply VBA, parameterizing the
encoder as a vanilla neural net, and the decoder and prior as VAEs with 50
latents and 2 layers of 200 hidden units. Because the data is binary, we
sample in training using the Gumbel-Softmax trick (Jang et al., 2016), which
allows gradient descent through discrete samples. The difference between
observational and interventional density is born out in Figure 4. It shows
that using variational backdoor adjustment, images with non-consecutive
numbers in $Y$ will be more likely under the interventional density than under
a naively estimated observational density.

###  4.3 Causal X-ray

Setup We introduce a new high dimensional dataset containing X-ray images to
emulate a potential real-world scenario. The data models the following
situation: a patient receives an initial X-ray scan ($Z$), a doctor provides
annotations indicating location for treatment ($X$), and the patient obtains a
new X-ray scan following treatment ($Y$). The treatment is administered with a
certain accuracy with respect to location based on the patient X-ray
($Z\rightarrow X$), and when administered, the treatment has a certain level
of efficacy ($X\rightarrow Y$). The goal of backdoor adjustment is to isolate
the effectiveness of the treatment and ignore the accuracy of how it is
administered.

To generate the data, we utilize chest X-ray scans from MedMNIST (Yang et al.,
2021), and apply a slight augmentation to add synthetic “tumors”. A tumor is a
circle composited with the image, and it occurs on the left or right lung with
equal probability. The treatments are annotations that are synthetically
created through shrinking MNIST zero digits to approximately the size of a
tumor (see Figure 5 for visual). It is placed either on the left or right
(either accurately or inaccurately depending on whether the tumor is placed on
the left or right). In this setup, we assume that the doctor administers the
treatment accurately 60% of the time. The outcome will be an X-ray image where
if the treatment is effective and administered accurately, the patient will
recover with 50% probability, resulting in an X-ray image without the tumor.
If the treatment is administered incorrectly, the tumor will always remain in
the patient.

Figure 5: Data generating process for Causal X-ray dataset. The small circle
in patient $Z$ is the tumor. On the right, intervening on the treatment will
make recovery rate independent of whether the annotation is correctly placed,
i.e. the accuracy of the doctor.

(a)

(b)

Figure 6: Left: Performance on Causal X-ray measured in bits/dim (lower is
better). Estimates do not vary with treatment accuracy, as expected in the
interventional density. Right: Samples drawn from separately trained encoder
(top row) and finetuned encoder (bottom row). The images from the finetuned
encoder are more smoother and more realistic, as they do not contain “double”
tumor artifacts.

Results We parameterize VBA for this task in the following manner. We model
the confounder images $Z$ with an auto-regressive flow model called FFJORD
(Grathwohl et al., 2018). We opt for this model because it gives exact
likelihood estimates in continuous settings and is amongst the state-of-the-
art models for density estimation. Because in this setting $Z$ and $Y$ are
closely related, parameterizing the encoder $p(z\mid x,y)$ and $p(y\mid x,z)$
with vanilla deep neural networks using a Gaussian predictive distribution
with diagonal covariance is sufficient. To evaluate the efficacy of backdoor
adjustment, we generate new interventional data in which we change the
accuracy of the doctors treatments. In the training data, treatment occurs at
the location of the tumor 60% of the time by construction. We can vary this
rate when generating test data, and because treatment $X$ as no dependence on
patient $Z$ in the interventional distribution, the likelihood should not
significantly change.

Table 2: Causal X-ray Results. We report in nats, with mean and standard error
averaged over the test Causal X-ray data for 5 runs. The rows given are
decomposed from the terms in the ELBO. Finetuning the encoder $q$ outperforms
$z$ sampled from MLE estimates of $q$. In this case, higher likelihood samples
of $z$ result in a tighter ELBO.
| Separate Training | Finetuned
---|---|---
$\mathbb{E}_{x,y\sim D}\mathbb{E}_{q(z\mid x,y)}\log p(z)$ | 979.67 (1.67) | 2012.51 (0.10)
$\mathbb{E}_{x,y\sim D}\mathbb{E}_{q(z\mid x,y)}\log p(y\mid x,z)$ | 1705.91 (0.23) | 1837.63 (0.08)
$\mathbb{E}_{x,y\sim D}\mathbb{E}_{-q(z\mid x,y)}\log q(z\mid x,y)$ | -1923.18 (0.08) | -2589.35 (0.08)

Figure 6(a) shows little dependence between treatment accuracy and reported
bits per dimension of VBA. It also shows that VBA with finetuning performs
strictly better than separate training. We investigate this difference in
performance by visualizing samples from each encoder. For a given $X,Y$, the
finetuned encoder produces smoother, more realistic images $Z$. Better looking
images along with greater likelihoods are indications of a better generative
model and a more accurate estimate of $p(y\mid do(x))$. Intuitively, this is
because a separately trained encoder must fit to the data, but the finetuned
encoder is optimized to better fit with the prior $p(z)$ and decoder $p(y\mid
x,z)$ to obtain a better lower bound. We confirm these intuitions in Table 2,
which explicitly shows these trade-offs. Finetuning makes encoder samples $z$
more likely, but in turn also improves prior and decoder likelihood.
Interestingly, the opposite effect is observed in Table 1, where lower
likelihood samples $z$ improved the ELBO. We restate the importance and power
of optimizing interventional density in this manner, since the interventional
density estimate can be improved in multiple ways.

##  5 Conclusion

We explore the various challenges that surface when trying to compute backdoor
adjustment in high dimensions. We show that using variational inference, it is
possible to compute the backdoor formula in the presence of high-dimensional
treatments and confounders. While variational inference is typically used in a
latent variable setting, we show with a novel optimization framework that
variational inference can be used as a powerful tool to compute identifiable
quantities in causal inference.

## References

  * Atkinson et al. (2022) Eric Atkinson, Charles Yuan, Guillaume Baudart, Louis Mandel, and Michael Carbin.  Semi-symbolic inference for efficient streaming probabilistic programming, 2022.
  * Austin (2011) Peter C Austin.  An introduction to propensity score methods for reducing the effects of confounding in observational studies.  _Multivariate behavioral research_ , 46(3):399–424, 2011.
  * Chen et al. (2019) Peipei Chen, Wei Dong, Xudong Lu, Uzay Kaymak, Kunlun He, and Zhengxing Huang.  Deep representation learning for individualized treatment effect estimation using electronic health records.  _Journal of biomedical informatics_ , 100:103303, 2019.
  * Fink (1997) Daniel Fink.  A compendium of conjugate priors.  _See http://www. people. cornell. edu/pages/df36/CONJINTRnew% 20TEX. pdf_ , 46, 1997.
  * Grathwohl et al. (2018) Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud.  Ffjord: Free-form continuous dynamics for scalable reversible generative models.  _arXiv preprint arXiv:1810.01367_ , 2018.
  * Holland (1986) Paul W Holland.  Statistics and causal inference.  _Journal of the American statistical Association_ , 81(396):945–960, 1986.
  * Jang et al. (2016) Eric Jang, Shixiang Gu, and Ben Poole.  Categorical reparameterization with gumbel-softmax.  _arXiv preprint arXiv:1611.01144_ , 2016.
  * Khemakhem et al. (2020) Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen.  Variational autoencoders and nonlinear ica: A unifying framework.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 2207–2217. PMLR, 2020.
  * Kingma & Welling (2013) Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Lange et al. (2022) Richard D. Lange, Ari Benjamin, Ralf M. Haefner, and Xaq Pitkow.  Interpolating between sampling and variational inference with infinite stochastic mixtures, 2022.
  * Locatello et al. (2019) Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem.  Challenging common assumptions in the unsupervised learning of disentangled representations.  In _international conference on machine learning_ , pp. 4114–4124. PMLR, 2019.
  * Louizos et al. (2017) Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard Zemel, and Max Welling.  Causal effect inference with deep latent-variable models.  _Advances in neural information processing systems_ , 30, 2017.
  * Lucas et al. (2019) James Lucas, George Tucker, Roger Grosse, and Mohammad Norouzi.  Understanding posterior collapse in generative latent variable models.  2019\.
  * Pearl (1995) Judea Pearl.  Causal diagrams for empirical research.  _Biometrika_ , 82(4):669–688, 1995.
  * Pearl (2009a) Judea Pearl.  Causal inference in statistics: An overview.  2009a.
  * Pearl (2009b) Judea Pearl.  _Causality_.  Cambridge university press, 2009b.
  * Rezende & Mohamed (2015) Danilo Rezende and Shakir Mohamed.  Variational inference with normalizing flows.  In _International conference on machine learning_ , pp. 1530–1538. PMLR, 2015.
  * Rissanen & Marttinen (2021) Severi Rissanen and Pekka Marttinen.  A critical look at the consistency of causal estimation with deep latent variable models.  _Advances in Neural Information Processing Systems_ , 34:4207–4217, 2021.
  * Rosenbaum & Rubin (1983) Paul R Rosenbaum and Donald B Rubin.  The central role of the propensity score in observational studies for causal effects.  _Biometrika_ , 70(1):41–55, 1983.
  * Schulz & Grimes (2002) Kenneth F Schulz and David A Grimes.  Blinding in randomised trials: hiding who got what.  _The Lancet_ , 359(9307):696–700, 2002.
  * Shalit et al. (2017) Uri Shalit, Fredrik D Johansson, and David Sontag.  Estimating individual treatment effect: generalization bounds and algorithms.  In _International Conference on Machine Learning_ , pp. 3076–3085. PMLR, 2017.
  * Shi et al. (2019) Claudia Shi, David Blei, and Victor Veitch.  Adapting neural networks for the estimation of treatment effects.  _Advances in neural information processing systems_ , 32, 2019.
  * Sørensen et al. (2006) Henrik Toft Sørensen, Timothy L Lash, and Kenneth J Rothman.  Beyond randomized controlled trials: a critical comparison of trials with nonrandomized studies.  _Hepatology_ , 44(5):1075–1082, 2006.
  * Tesei et al. (2023) Gino Tesei, Stefanos Giampanis, Jingpu Shi, and Beau Norgeot.  Learning end-to-end patient representations through self-supervised covariate balancing for causal treatment effect estimation.  _Journal of Biomedical Informatics_ , 140:104339, 2023.
  * Wang & Kwiatkowska (2022) Benjie Wang and Marta Kwiatkowska.  Symbolic causal inference via operations on probabilistic circuits.  In _NeurIPS 2022 Workshop on Neuro Causal and Symbolic AI (nCSI)_ , 2022.
  * Wu & Fukumizu (2021) Pengzhou Wu and Kenji Fukumizu.  Intact-vae: Estimating treatment effects under unobserved confounding.  _arXiv preprint arXiv:2101.06662_ , 2021.
  * Xu & Gretton (2022) Liyuan Xu and Arthur Gretton.  A neural mean embedding approach for back-door and front-door adjustment.  _arXiv preprint arXiv:2210.06610_ , 2022.
  * Yang et al. (2021) Jiancheng Yang, Rui Shi, Donglai Wei, Zequan Liu, Lin Zhao, Bilian Ke, Hanspeter Pfister, and Bingbing Ni.  Medmnist v2: A large-scale lightweight benchmark for 2d and 3d biomedical image classification.  _CoRR_ , abs/2110.14795, 2021.  URL https://arxiv.org/abs/2110.14795.
  * Zhang et al. (2021) Weijia Zhang, Lin Liu, and Jiuyong Li.  Treatment effect estimation with disentangled latent factors.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 35, pp. 10923–10930, 2021.
  * Zou et al. (2020) Hao Zou, Peng Cui, Bo Li, Zheyan Shen, Jianxin Ma, Hongxia Yang, and Yue He.  Counterfactual prediction for bundle treatment.  _Advances in Neural Information Processing Systems_ , 33:19705–19715, 2020.

##  Appendix A Supplemental Linear Gaussian Figures

First we show here a useful visual to gain intuition about the linear Gaussian
setting. We visualize the distributions given by Equation 11 using concrete
values: $c_{1}=3,c_{2}=2,c_{3}=-6,\sigma_{1}=0.5,\sigma_{2}=1$. We observe the
potential difference between observational and interventional densities. In
7(b) and 7(c), we depict two distributions that are consistent with the
observed joint distribution depicted in 7(a). Without the assumptions made by
backdoor adjustment, we cannot estimate the interventional density seen in
7(a).

(a) Joint Density
$p(x,y)$

(b) Observational Density
$p(y\mid x)$

(c) Interventional Density $p(y\mid do(x))$

Figure 7: Linear Gaussian Visualization

Below are figures that give a different view on Figures 2(a) and 2(b). Instead
of plotting the error we give the log likelihood plotted with the ground truth
interventional likelihood. As given by the ELBO, our estimates will be a lower
bound in expectation of the ground truth.

(a) Sampling vs Variational Inference
Log Likelihood Comparison

(b) Finetuning vs Separate Training
Log Likelihood Comparison

Figure 8: Linear Gaussian Results

##  Appendix B Proof of Proposition 1

Consider the objective given in Equation 10.

$\mathcal{L}^{\text{ELBO}}_{\phi,\theta,\gamma}(x,y)=-\sum\nolimits_{z^{\prime}_{1},...,z^{\prime}_{n}\sim
q_{\phi}(z\mid
x,y)}\left(\mathcal{L}^{\text{MLE}}_{\theta}(z^{\prime}_{j})+\mathcal{L}^{\text{MLE}}_{\gamma}(x,y,z^{\prime}_{j})-\mathcal{L}^{\text{MLE}}_{\phi}(x,y,z^{\prime}_{j})\right)$
(15)

There exists parameters $\phi,\theta,\gamma$ such that

$q_{\phi}(z\mid x,y)=p_{\theta}(z)$ (16) $p_{\gamma}(y\mid x,z)=p(y\mid x)$
(17)

Given these parameters, we get the following ELBO.

$\displaystyle\mathbb{E}_{q_{\phi}(z\mid x,y)}\left(\log p_{\theta}(z)+\log
p_{\gamma}(y\mid x,z)-\log q_{\phi}(z\mid x,y)\right)$ (18)
$\displaystyle=\mathbb{E}_{q_{\phi}(z\mid x,y)}\left(\log p(y\mid x)\right)$
(19) $\displaystyle=\log p(y\mid x)$ (20)

Thus, when $E_{p(x,y)}\log p(y\mid x)>E_{p(x,y)}\log p(y\mid do(x))$, the
model $f_{\phi,\theta,\gamma}(x,y)$ will converge to $p(y\mid x)$. We can show
by construction that this can occur for some SCM. Let $\mathcal{B}$ denote a
Bernoulli distribution. Consider a simple example in which variables $X,Y,Z$
induce distributions

$\displaystyle p(z)=\mathcal{B}(1/2)$ (21) $\displaystyle p(x\mid
z)=\begin{cases}\mathcal{B}(3/4)&\text{if }z,\\\ \mathcal{B}(1/2)&\text{if
}\neg z\end{cases}$ (22) $\displaystyle p(y\mid
x,z)=\begin{cases}\mathcal{B}(1/2)&\text{if }x\vee z,\\\
\mathcal{B}(3/4)&\text{if }\neg x\wedge\neg z\end{cases}$ (23)

It then follows from this example that

$\displaystyle E_{p(x,y)}\log p(y\mid x)$ (24)
$\displaystyle=(1/8)\log(1/3)+(1/4)\log(2/3)+(5/16)\log(1/2)+(5/16)\log(1/2)$
(25)
$\displaystyle>(1/8)\log(3/8)+(1/4)\log(5/8)+(5/16)\log(1/2)+(5/16)\log(1/2)$
(26) $\displaystyle=E_{p(x,y)}\log p(y\mid do(x))$ (27)

##  Appendix C Pitfall of Fully Joint Training

As state in Proposition 1 we cannot simply optimize the parameters of the
prior $\theta$ and decoder $\gamma$, and that they must remain fixed in the
objective $\mathcal{L}^{\text{ELBO}}$ given in Equation 9. If these parameters
are freely optimized in a fully joint manner, the backdoor adjustment will not
be computed correctly because $Z$ will not be constrained by observation.
Below, we have a plot similar to Figure 8(b).

Figure 9: Pitfall of Fully Joint Training

This figure empirically shows that naively optimizing Equation 9 with fully
joint training, that is optimizing all the parameters, will result in
estimation of the wrong query. It can be seen that fully joint training is
estimating $p(y\mid x)$ rather than $p(y\mid do(x))$

##  Appendix D Proof of Linear Gaussian Analytical Solutions

We prove the solutions found for the interventional and observational
distributions seen in Equation 12 and Equation 13.

###  D.1 Interventional

We aim to show that

$y\mid
do(x)\sim\mathcal{N}\left(c_{2}x,\sqrt{c_{3}^{2}+\sigma_{2}^{2}}\right)$

First observe that the setup seen in Equation 11 entails
$z\sim\mathcal{N}(0,1)$ and $y\mid
x,z\sim\mathcal{N}(c_{2}x+c_{3}z,\sigma_{2})$. We can thus compute

$\displaystyle p(y\mid do(x))$ $\displaystyle=\sum_{z}p(z)p(y\mid x,z)$ (28)
$\displaystyle=\sum_{z}\left(\frac{1}{\sqrt{2\pi}}e^{\frac{-z^{2}}{2}}\right)\left(\frac{1}{\sigma_{2}\sqrt{2\pi}}e^{\frac{-(y-c_{2}x-c_{3}z)^{2}}{2\sigma_{2}^{2}}}\right)$
(29)
$\displaystyle=\frac{1}{\sqrt{(c_{3}^{2}+\sigma_{2}^{2})2\pi}}e^{\frac{-(y-c_{2}x)^{2}}{2(c_{3}^{2}+\sigma_{2}^{2})}}$
(30)

The same proof can be found in Rissanen & Marttinen (2021) Appendix A Equation
12.

###  D.2 Observational

To compute $p(y\mid x)$ in this setting, we must express
$Y:=c_{2}X+c_{3}Z+\mathcal{N}(0,\sigma_{2})$ in terms of only $X$. To do this,
we must express $Z$ in terms of $X$. It is known that this swap can be
performed using conjugate priors (Fink, 1997; Atkinson et al., 2022). We find

$Z\mid
X=x\sim\mathcal{N}\left(\frac{c_{1}x}{c_{1}^{2}+\sigma_{1}^{2}},\sqrt{\frac{\sigma_{1}^{2}}{c_{1}^{2}+\sigma_{1}^{2}}}\right)$

Through simple addition and scaling properties of Gaussians, we combine the
above with the definition of $Y$ to obtain

$y\mid
x\sim\mathcal{N}\left(\left(\frac{c_{1}c_{3}}{(c_{1}+\sigma_{1}^{2})}+c_{2}\right)x,\sqrt{\frac{c_{3}^{2}\sigma_{1}^{2}}{c_{1}^{2}+\sigma_{1}^{2}}+\sigma_{2}^{2}}\right)$

##  Appendix E Image Samples

In our generative models, we can draw samples to qualitatively ensure that the
correct distribution is being learned.

###  E.1 MNIST Samples

Figure 10: Examples of $X$ from MNIST dataset. Figure 11: Samples from
$p(y\mid x)$, with $X$ given by Figure 10. Here, we see that the numbers
generated are successive because that is how $Y$ is observed. Figure 12:
Samples from $p(y\mid do(x))$, with $X$ given by Figure 10. Here, we are able
to sample $Y$ such that digits are not successive, which indicates the
difference between the observational and interventional distributions.

###  E.2 Causal X-Ray Samples

Figure 13: Samples from FFJORD

Figure 14: More Causal X-ray samples from the encoder. Top 3 rows are sampled
from an encoder that is trained separately. Bottom 3 rows are sampled from a
finetuned encoder. More samples further reinforce the conclusion from Figure
6(b)

.