summaries: This paper considers the problem of minimizing a non-smooth matrix function  $f(X)$ over the spectrahedron $S_n$, i.e., positive semidefinite matrices with trace equal to 1. Such optimization problems often arise as the convex relaxation of various low-rank matrix recovery problems. In general solving these problems with projected subgradient descent requires a projection onto the spectraherdon at every iteration, which is expensive because of the need to compute a full SVD at every step.  Under the assumption that the non-smooth objective function $f(X)$ can be written as the pointwise maximum of a family of smooth functions,  the authors reformulate the original optimization problem as a saddle point problem which can be solved via the projected extragradient method. The key contribution is that under a “generalized strict complementarity” condition and when close to the ground truth, the projection onto $S_n$ can be replaced by a low-rank projection. Therefore, the overall algorithm can be significantly more efficient, because only a low-rank SVD is needed.

summaries: This paper develops an extragradient method for low-rank and nonsmooth matrix optimization problems. Under certain conditions, the proposed method converges to an optimal solution with rate O(1/t). Numerical experiments are provided to support the theoretical results.

summaries: In this paper, motivated by a new efficient method in solving smooth low-rank matrix optimization problems, the authors propose to extend the method to nonsmooth problems. However, direct generalization is not possible, as is shown for a specific failure case. Thus, the authors propose to introduce some auxiliary parameters and transform the original nonsmooth problem to a smooth saddle-point problem. The authors then extend to the method of this new problem and derive a new algorithm that is suitable for the saddle-point problem. Theoretically, the authors proved the convergence of their proposed method. And empirical results verify the correctness of their proposed method.

summaries: This paper studies the problem of minimizing a convex, non-smooth objective function of matrices with low-rank constraints. For the smooth counterpart of the same problem, it has been previously shown that SVDs computed during the iterations of algorithms such as projected sub-gradient descent can be replaced by low-rank SVDs, thus improving the computational complexity. This paper studies the conditions under which such a replacement can be used for non-smooth objective functions.   The main result of the paper consists of describing the class of non-smooth functions for which low-rank SVDs can be used. The paper proves theoretically and empirically that for non-smooth functions that have a saddle point structure (can be written as a sum of a smooth function and a point-wise maximum of affine functions), and when the initialization is within a certain radius w.r.t the ground truth, low0rank SVDs lead to the same solution as full SVDs and a O(1/t) convergence can be achieved.   Further, the paper also shows empirically that using low-rank SVDs on non-smooth objectives leads to the same solution as using full SVDs.

summaries: The paper considers a class of nonsmooth and low-rank matrix optimization problems. Major contributions include: (1) Discuss strict complementarity for nonsmooth problems and challenges of low-rank projected subgradient steps. (2) Under certain assumptions, a special type of nonsmooth problems can be reformulated as maximum of smooth ones and thereby a minimax form, and develop a projected extragradient method with low-rank projections. (3) Use the proposed method to solve several classical low-rank approximation problems in applications.

abstractText: Low-rank and nonsmooth matrix optimization problems capture many fundamental tasks in statistics and machine learning. While significant progress has been made in recent years in developing efficient methods for smooth low-rank optimization problems that avoid maintaining high-rank matrices and computing expensive high-rank SVDs, advances for nonsmooth problems have been slow paced. In this paper we consider standard convex relaxations for such problems. Mainly, we prove that under a natural generalized strict complementarity condition and under the relatively mild assumption that the nonsmooth objective can be written as a maximum of smooth functions, the extragradient method, when initialized with a “warm-start” point, converges to an optimal solution with rate O(1/t) while requiring only two low-rank SVDs per iteration. We give a precise trade-off between the rank of the SVDs required and the radius of the ball in which we need to initialize the method. We support our theoretical results with empirical experiments on several nonsmooth low-rank matrix recovery tasks, demonstrating that using simple initializations, the extragradient method produces exactly the same iterates when full-rank SVDs are replaced with SVDs of rank that matches the rank of the (low-rank) ground-truth matrix to be recovered.

title: Community detection and stochastic block models: recent developments

venue: The Journal of Machine Learning Research,

authors:
Praneeth Netrapalli
UN Niranjan
Sujay Sanghavi
Animashree Anandkumar
Prateek Jain

title: Non-convex robust pca

venue: Advances in Neural Information Processing Systems,

title: First-order methods in optimization

authors:
Amir Beck
Marc Teboulle

title: A fast iterative shrinkage-thresholding algorithm for linear inverse problems

venue: SIAM journal on imaging sciences,

title: Nonlinear programming

venue: Athena Scientific,

authors:
Afonso S. Bandeira
Nicolas Boumal
Amit Singer

title: Tightness of the maximum likelihood semidefinite relaxation for angular synchronization

venue: Mathematical Programming,

title: Convex optimization: Algorithms and complexity

venue: Foundations and Trends in Machine Learning,

authors:
Vasileios Charisopoulos
Yudong Chen
Damek Davis
Mateo Díaz
Lijun Ding
Dmitriy Drusvyatskiy

title: Low-rank matrix recovery with composite optimization: good conditioning and rapid convergence

venue: Foundations of Computational Mathematics,

authors:
PXinyang Yi
Dohyung Park
Yudong Chen
Constantine Caramanis

title: Fast algorithms for robust pca via gradient descent

authors:
Yudong Chen
Huan Xu
Constantine Caramanis
Sujay Sanghavi

title: Robust matrix completion and corrupted columns

venue: In Proceedings of the 28th International Conference on Machine Learning

title: Optimization and Nonsmooth Analysis

authors:
Alexandre d’Aspremont
Laurent El Ghaoui
Michael I Jordan
Gert RG Lanckriet

title: A direct formulation for sparse pca using semidefinite programming

venue: SIAM review,

authors:
Lijun Ding
Jicong Fan
Madeleine Udell

title: kfw: A frank-wolfe style algorithm with stronger subproblem oracles, 2020

authors:
Dmitriy Drusvyatskiy
Adrian S Lewis

title: Error bounds, quadratic growth, and linear convergence of proximal methods

venue: Mathematics of Operations Research,

authors:
Francisco Facchinei
Jong-Shi Pang

title: Finite-Dimensional Variational Inequalities and Complementarity Problems, volume II

title: Linear convergence of frank-wolfe for rank-one matrix recovery without strong convexity

venue: arXiv preprint arXiv:1912.01467,

title: On the convergence of projected-gradient methods with low-rank projections for smooth convex minimization over trace-norm balls and related problems

venue: SIAM Journal on Optimization,

title: On the convergence of stochastic gradient descent with low-rank projections for convex low-rank matrix problems

venue: Conference on Learning Theory, COLT,

authors:
Dan Garber
Atara Kaplan

title: Fast stochastic algorithms for low-rank and nonsmooth matrix problems

venue: The 22nd International Conference on Artificial Intelligence and Statistics,

title: On the efficient implementation of the matrix exponentiated gradient algorithm for low-rank matrix optimization

venue: arXiv preprint arXiv:2012.10469,

authors:
Hui Zou
Trevor Hastie
Robert Tibshirani

title: Sparse principal component analysis

venue: Journal of Computational and Graphical Statistics,

authors:
Olga Klopp
Karim Lounici
Alexandre B Tsybakov

title: Robust matrix completion

venue: Probability Theory and Related Fields,

title: The extragradient method for finding saddle points and other problems

authors:
Srinadh Bhojanapalli
Anastasios Kyrillidis
Sujay Sanghavi

title: Dropping convexity for faster semi-definite optimization

venue: In 29th Annual Conference on Learning Theory,

authors:
Xiao Li
Zhihui Zhu
Anthony Man-Cho So
Rene Vidal

title: Nonconvex robust low-rank matrix recovery

authors:
Francesco Locatello
Alp Yurtsevert
Olivier Fercoq
Volkan Cevhert

title: Stochastic frankwolfe for composite convex minimization

authors:
Canyi Lu
Jinhui Tang
Shuicheng Yan
Zhouchen Lin

title: Generalized nonconvex nonsmooth low-rank minimization

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Emmanuel J. Candès
Xiaodong Li
Yi Ma
John Wright

title: Robust principal component analysis

venue: Journal of the ACM,

authors:
Adel Javanmard
Andrea Montanari
Federico Ricci-Tersenghi

title: Phase transitions in semidefinite relaxations

venue: Proceedings of the National Academy of Sciences,

title: Prox-method with rate of convergence o(1/t) for variational inequalities with lipschitz continuous monotone operators and smooth convex-concave saddle point problems

venue: SIAM J. on Optimization,

authors:
Gergely Odor
Yen-Huan Li
Alp Yurtsever
Ya-Ping Hsieh
Quoc Tran-Dinh
Marwa El Halabi
Volkan Cevher

title: Frank-wolfe works for non-lipschitz continuous gradient objectives: scalable poisson phase retrieval

venue: In 2016 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP),

authors:
Stephen Boyd
Neal Parikh
Eric Chu
Borja Peleato
Jonathan Eckstein

title: Distributed optimization and statistical learning via the alternating direction method of multipliers

authors:
John Wright
Arvind Ganesh
Shankar Rao
Yigang Peng
Yi Ma

title: Robust principal component analysis: Exact recovery of corrupted low-rank matrices via convex optimization

venue: Advances in Neural Information Processing Systems

authors:
Julien Mairal
Francis Bach
Jean Ponce
Guillermo Sapiro

title: Online learning for matrix factorization and sparse coding

venue: Journal of Machine Learning Research,

authors:
Emile Richard
Pierre-Andr‘e Savalle
Nicolas Vayatis

title: Estimation of simultaneously sparse and low rank matrices

venue: Proceedings of the 29th International Conference on Machine Learning,

authors:
Kiran K Thekumparampil
Prateek Jain
Praneeth Netrapalli
Sewoong Oh

title: Projection efficient subgradient method and optimal nonsmooth frank-wolfe method

authors:
Nicolas Boumal
Vladislav Voroninski
Afonso S. Bandeira

title: Deterministic guarantees for burer-monteiro factorizations of smooth semidefinite programs

venue: Communications on Pure and Applied Mathematics,

authors:
Cun Mu
Yuqian Zhang
John Wright
Donald Goldfarb

title: Scalable robust matrix recovery: Frank-wolfe meets proximal methods

venue: SIAM Journal on Scientific Computing,

authors:
Lijun Ding
Yingjie Fei
Qiantong Xu
Chengrun Yang

title: Spectral frank-wolfe algorithm: Strict complementarity and linear convergence

authors:
Quanming Yao
James T Kwok
Taifeng Wang
Tie-Yan Liu

title: Large-scale low-rank matrix learning with nonconvex regularizers

venue: IEEE transactions on pattern analysis and machine intelligence,

authors:
Yiqiao Zhong
Nicolas Boumal

title: Near-optimal bounds for phase synchronization

authors:
Zirui Zhou
Anthony Man-Cho So

title: A unified approach to error bounds for structured convex optimization problems

heading: 1 Introduction

text:
Low-rank and nonsmooth matrix optimization problems have many important applications in statistics, machine learning, and related fields, such as sparse PCA [21, 34], robust PCA [28, 33, 2, 9, 38], phase synchronization [41, 6, 29], community detection and stochastic block models [1]1, low-rank and sparse covariance matrix recovery [35], robust matrix completion [22, 10], and more. For many of these problems, convex relaxations, in which one replaces the nonconvex low-rank constraint with a trace-norm constraint, have been demonstrated in numerous papers to be highly effective both in theory (under suitable assumptions) and empirically (see references above). These convex relaxations can be formulated as the following general nonsmooth optimization problem:
min X∈Sn g(X), (1)
where g : Sn → R is convex but nonsmooth, and Sn = {X ∈ Sn | Tr(X) = 1, X 0} is the spectrahedron in Sn, Sn being the space of n× n real symmetric matrices. Problem (1), despite being convex, is notoriously difficult to solve in large scale. The simplest and most general approach applicable to it is the projected subgradient method [3, 7], which requires on
1in [41, 6, 29] and [1] the authors consider SDPs with linear objective function and affine constraints of the form A(X) = b. By incorporating the linear constraints into the objective function via a `2 penalty term of the form λ‖A(X)− b‖2, λ > 0, we obtain a nonsmooth objective function.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
each iteration to compute a Euclidean projection onto the spectrahedron Sn, which in worst case amounts to O(n3) runtime per iteration. In many applications g(X) follows a composite model, i.e., g(X) = h(X) + w(X), where h(·) is convex and smooth and w(·) is convex and nonsmooth but admits a simple structure (e.g., nonsmooth regularizer). For such composite objectives, without the spectrahedron constraint, proximal methods such as FISTA [4] or splitting methods such as ADMM [32] are often very effective. However, with the spectrahderon constraint, all such methods require on each iteration to apply a subprocedure (e.g., computing the proximal mapping) which in worst case amounts to at least O(n3) runtime. A third type of off-the-shelf methods include those which are based on the conditional gradient method and adapted to nonsmooth problems, see for instance [31, 19, 36, 26]. The advantage of such methods is that no expensive high-rank SVD computations are needed. Instead, only a single leading eigenvector computation (i.e., a rank-one SVD) per iteration is required. However, whenever the number of iterations is not small, these methods still require to store high-rank matrices in memory, even when the optimal solution is low-rank. Thus, to conclude, standard first-order methods for Problem (1) require in worst case Ω(n3) runtime per iteration and/or to store high-rank matrices in memory.
In the recent works [17, 18] it was established that for smooth objective functions, the high-rank SVD computations required for Euclidean projections onto the spectrahedron in standard gradient methods, can be replaced with low-rank SVDs in the close proximity of a low-rank optimal solution. This is significant since the runtime to compute a rank-r SVD of a given n×n matrix using efficient iterative methods typically scales with rn2 (and further improves when the matrix is sparse), instead of n3 for a full-rank SVD. These results depend on the existence of eigen-gaps in the gradient of the optimal solution, which we refer to as a generalized strict complementarity condition. These results also hinge on a unique property of the Euclidean projection onto the spectrahedron. The projection onto the spectrahedron of a matrix X ∈ Sn, which admits an eigen-decomposition X = ∑n i=1 λiviv > i , is given by
ΠSn [X] = n∑ i=1 max{0, λi − λ}viv>i , (2)
where λ ∈ R is the unique scalar satisfying ∑n i=1 max{0, λi−λ} = 1. This operation thus truncates all eigenvalues that are smaller than λ, while leaving the eigenvectors unchanged, thereby returning a matrix with rank equal to the number of eigenvalues greater than λ. Importantly, when the projection of X onto Sn is of rank r, only the first r components in the eigen-decomposition of X are required to compute it in the first place, and thus, only a rank-r SVD of X is required. In other words and simplifying, [17, 18] show that under strict complementary, at the proximity of an optimal solution of rank r, the exact Euclidean projection equals the rank-r truncated projection given by:
Π̂rSn [X] := ΠSn [ r∑ i=1 λiviv > i ] . (3)
Extending the results of [17, 18] to the nonsmooth setting is difficult since the smoothness assumption is critical to the analysis. Moreover, while [17, 18] rely on certain eigen-gaps in the gradients at optimal points, for nonsmooth problems, since the subdifferential set is often not a singleton, it is not likely that a similar eigen-gap property holds for all subgradients of an optimal solution.
In this paper we show that under the mild assumption that Problem (1) can be formulated as a smooth convex-concave saddle-point problem, i.e., the nonsmooth term can be written as a maximum over (possibly infinite number of) smooth convex functions, we can obtain results in the spirit of [17, 18]. Concretely, we show that if a generalized strict complementarity (GSC) assumption holds for a low-rank optimal solution (see Assumption 1 in the sequel), the extragradient method for smooth convex-concave saddle-point problems [23, 30] (see Algorithm 1 below), when initialized in the proximity of the optimal solution, converges with its original convergence rate of O(1/t), while requiring only two low-rank SVDs per iteration2. It is important to recall that while the extragradient method requires two SVDs per iteration, it has the benefit of a fast O(1/t) convergence rate, while simpler saddle-point methods such as mirror-descent-based only achieve a O(1/ √ t) rate [7].
Our contributions can be summarized as follows: 2note that the extradgradient method computes two projected-gradient steps on each iteration, and thus two SVDs are needed per iteration.
• We prove that even under (standard) strict complementarity, the projected subgradient method, when initialized with a “warm-start", may produce iterates with rank higher than that of the optimal solution. This phenomena further motivates our saddle-point approach. See Lemma 5.
• We suggest a generalized strict complementarity (GSC) condition for saddle-point problems and prove that when g(·) — the objective function in Problem (1), admits a highly popular saddle-point structure (one which captures all applications we mentioned in this paper), GSC w.r.t. an optimal solution to Problem (1) implies GSC (with the same parameters) w.r.t. a corresponding optimal solution of the equivalent saddle-point problem (the other direction always holds). See Section 3.
• Main result: we prove that for a smooth convex-concave saddle-point problem and an optimal solution which satisfies GSC, the extragradient method, when initialized with a “warm-start", converges with its original rate of O(1/t) while requiring only two low-rank SVDs per iteration. Moreover, we prove GSC facilitates a precise and powerful tradeoff: increasing the rank of SVD computations (beyond the rank of the optimal solution) can significantly increase the radius of the ball in which the method needs to be initialized. See Theorem 1.
• We present extensive numerical evidence that demonstrate both the plausibility of the GSC assumption in various tasks, and more importantly, demonstrate that indeed the extragradient method with simple initialization converges correctly (i.e., produces exactly the same sequences of iterates) when the rank of the SVDs used to compute the (truncated) projections matches the rank of the (low-rank) ground-truth matrix to be recovered, instead of naively using full-rank SVDs (as suggested by (2)). See Section 5.

heading: 1.1 Additional related work

text:
Since, as in the works [17, 18] mentioned before which deal with smooth objectives, strict complementarity plays a key role in our analysis, we refer the interested reader to the recent works [16, 39, 13, 20] which also exploit this property for efficient smooth and convex optimization over the spectrahedron. Strict complementarity has also played an instrumental role in two recent and very influential works which used it to prove linear convergence rates for proximal gradient methods [42, 14].
Besides convex relaxations such as Problem (1), considerable advances have been made in the past several yeas in developing efficient nonconvex methods with global convergence guarantees for low-rank matrix problems. In [37] the authors consider semidefinite programs and prove that under a smooth manifold assumption on the constraints, such methods converge to the optimal global solution. In [24] the authors prove global convergence of factorized nonconvex gradient descent from a ”warm-start” initialization point for non-linear smooth minimization on the positive semidefinite cone. Very recently, [8] has established, under statistical conditions, fast convergence results from “warm-start” initialization of nonconvex first-order methods, when applied to nonsmooth nonconvex matrix recovery problems which are based on the explicit factorization of the low-rank matrix. A result of similar flavor concerning nonsmooth and nonconvex formulation of robust recovery of lowrank matrices from random linear measurements was presented in [25]. Finally, several recent works have considered nonconvex low-rank regularizers which result in nonconvex nonsmooth optimization problems, but guarantee convergence only to a stationary point [27, 40].

heading: 2 Strict complementarity for nonsmooth optimization and difficulty of applying low-rank projected subgradient steps

text:
Our analysis of the nonsmooth Problem (1) naturally depends on certain subgradients of an optimal solution which, in many aspects, behave like the gradients of smooth functions. The existence of such a subgradient is guaranteed from the first-order optimality condition for constrained convex minimization problems:
Lemma 1 (first-order optimality condition, see [3]). Let g : Sn → R be a convex function. Then X∗ ∈ Sn minimizes g over Sn if and only if there exists a subgradient G∗ ∈ ∂g(X∗) such that 〈X−X∗,G∗〉 ≥ 0 for all X ∈ Sn.
For some G∗ ∈ ∂g(X∗) which satisfies the first-order optimality condition for an optimal solution X∗, if the multiplicity of the smallest eigenvalue equals r∗ = rank(X∗), then it can be shown that the optimal solution satisfies a strict complementarity assumption. The equivalence between a standard strict complementarity assumption on some low-rank optimal solution of a smooth optimization problem over the spectrahedron and an eigen-gap in the gradient of the optimal solution was established in [39]. We generalize this equivalence to also include nonsmooth problems. The proof follows similar arguments and is given in Appendix A.1. Definition 1 (strict complementarity). An optimal solution X∗ ∈ Sn of rank r∗ for Problem (1) satisfies the strict complementarity assumption with parameter δ > 0, if there exists an optimal solution of the dual problem3 (Z∗, s∗) ∈ Sn ×R such that rank(Z∗) = n− r∗, and λn−r∗(Z∗) ≥ δ. Lemma 2. Let X∗ ∈ Sn be a rank-r∗ optimal solution to Problem (1). X∗ satisfies the (standard) strict complementarity assumption with parameter δ > 0 if and only if there exists a subgradient G∗ ∈ ∂g(X∗) such that 〈X−X∗,G∗〉 ≥ 0 for all X ∈ Sn and λn−r∗(G∗)− λn(G∗) ≥ δ.
Throughout this paper we assume a weaker and more general assumption than strict complementarity, namely generalized strict complementarity (GSC), which we present now. Assumption 1 (generalized strict complementarity). We say an optimal solution X∗ to Problem (1) satisfies the generalized strict complementarity assumption with parameters r, δ, if there exists a subgradient G∗ ∈ ∂g(X∗) such that 〈X−X∗,G∗〉 ≥ 0 for all X ∈ Sn and λn−r(G∗)−λn(G∗) ≥ δ.
In [17] the author presents several characteristic properties of the gradient of the optimal solution in optimization problems over the spectrahedron. Using the existence of subgradients which satisfy the condition in Lemma 1, we can extend these properties also to the nonsmooth setting. The following lemma shows that GSC with parameters (r, δ) for some δ > 0 (Assumption 1) is a sufficient condition for the optimal solution to be of rank at most r. The proof follows immediately from the proof of the analogous Lemma 7 in [17], by replacing the gradient of the optimal solution with a subgradient for which the first-order optimality condition holds. Lemma 3. Let X∗ be an optimal solution to Problem (1) and write its eigen-decomposition as X∗ = ∑r∗ i=1 λiviv T i . Then, any subgradient G
∗ ∈ ∂g(X∗) which satisfies 〈X−X∗,G∗〉 ≥ 0 for all X ∈ Sn, admits an eigen-decomposition such that the set of vectors {vi}r ∗
i=1 is a set of leading eigenvectors of (−G∗) which corresponds to the eigenvalue λ1(−G∗) = −λn(G∗). Furthermore, there exists at least one such subgradient.
One motivation for assuming (standard) strict complementarity (Assumption 1 with parameters r = rank(X∗) and δ > 0) is that it guarantees a certain notion of robustness of the problem to small perturbations in the parameters. It is well known (see for instance [3]) that a projected subgradient step from X∗ with respect to a subgradient G∗ ∈ ∂g(X∗) for which the first-order optimality condition holds, returns the optimal solution X∗ itself. This implies that rank (ΠSn [X
∗ − ηG∗]) = rank(X∗) (here η is the step-size). Without (standard) strict complementarity however, a small change in the parameters could result in a higher rank matrix. This is captured in the following lemma which is analogous to Lemma 3 in [18], where again the proof is straightforward from the proof in [18] by replacing the gradient of the optimal solution with a subgradient for which the first-order optimality condition holds. Lemma 4. Let X∗ be an optimal solution of rank r∗ to Problem (1). Let G∗ ∈ ∂g(X∗) be a subgradient at X∗ such that 〈X −X∗,G∗〉 ≥ 0 for all X ∈ Sn. Then, λn−r∗(G∗) = λn(G∗) if and only if for any arbitrarily small ζ > 0 it holds that rank ( Π(1+ζ)Sn [X ∗ − ηG∗] ) > r∗, where η > 0, (1 + ζ)Sn = {(1 + ζ)X |X ∈ Sn}, and Π(1+ζ)Sn [·] denotes the Euclidean projection onto the set (1 + ζ)Sn.

heading: 2.1 The challenge of applying low-rank projected subgradient steps

text:
We now demonstrate the difficulty of replacing the full-rank SVD computations required in projected subgradient steps over the spectrahedron, with their low-rank SVD counterparts when attempting to solve Problem (1). We prove that a projected subgradient step from a point arbitrarily close to a
3Denote q(Z, s) = minX∈Sn{g(X) + s(1− Tr(X))− 〈Z,X〉}. The dual problem of Problem (1) can be written as: max{Z 0, s∈R}{q(Z, s) | (Z, s) ∈ dom(q)}.
low-rank optimal solution — even one that satisfies strict complementarity (Definition 1), may result in a higher rank matrix. The problem on which we demonstrate this phenomena is a well known convex formulation of the sparse PCA problem [12]. The proof is given in Appendix A.2. Lemma 5 (failure of low-rank subgradient descent on sparse PCA). Consider the problem minX∈Sn{g(X) := − 〈 zz> + z⊥z > ⊥,X 〉 + 12k‖X‖1}, where z = (1/ √ k, . . . , 1/ √ k, 0, . . . , 0)> is
supported on the first k entries, z⊥ = (0, . . . , 0, 1/ √ n− k, . . . , 1/ √ n− k)> is supported on the last n−k entries, and k ≤ n/4. Then, zz> is a rank-one optimal solution for which strict complementarity holds. However, for any η < 23 and any v ∈ R
n such that ‖v‖ = 1, support(v) ⊆ support(z), and 〈z,v〉2 = 1 − 12‖vv > − zz>‖2F ≥ 1 − 12k2 , it holds that rank ( ΠSn [vv > − ηGvv> ] ) > 1, where Gvv> = −zz> − z⊥z>⊥ + 12k sign(vv >) ∈ ∂g(vv>).
Note that the subgradient of the `1-norm which we choose for the projected subgradient step simply corresponds to the sign function, which is arguably the most natural choice.

heading: 3 From nonsmooth to saddle-point problems

text:
To circumvent the difficulty demonstrated in Lemma 5 in incorporating low-rank SVDs into standard subgradient methods for solving Problem (1), we propose tackling the nonsmooth problem with saddle-point methods.
We assume the nonsmooth Problem (1) can be written as a maximum of smooth functions, i.e., g(X) = maxy∈K f(X,y), where K ⊂ Y is some compact and convex subset of the finite linear space over the reals Y onto which it is efficient to compute Euclidean projections. We assume f(·,y) is convex for all y ∈ K and f(X, ·) is concave for all X ∈ Sn. That is, we rewrite Problem (1) as the following equivalent saddle-point problem:
min X∈Sn max y∈K f(X,y). (4)
Finding an optimal solution to problem (4) is equivalent to finding a saddle-point (X∗,y∗) ∈ Sn ×K such that for all X ∈ Sn and y ∈ K, f(X∗,y) ≤ f(X∗,y∗) ≤ f(X,y∗). We make a standard assumption that f(·, ·) is smooth with respect to all components. That is, we assume there exist βX , βy, βXy, βyX ≥ 0 such that for any X, X̃ ∈ Sn and y, ỹ ∈ K the following four inequalities hold: ‖∇Xf(X,y)−∇Xf(X̃,y)‖F ≤ βX‖X− X̃‖F , ‖∇yf(X,y)− ∇yf(X, ỹ)‖2 ≤ βy‖y − ỹ‖2, ‖∇Xf(X,y) − ∇Xf(X, ỹ)‖F ≤ βXy‖y − ỹ‖2, ‖∇yf(X,y) − ∇yf(X̃,y)‖2 ≤ βyX‖X− X̃‖F , where∇Xf = ∂f∂X and ∇yf = ∂f ∂y .
We denote by β the full Lipschitz parameter of the gradient, that is for any X, X̃ ∈ Sn and y, ỹ ∈ K,
‖(∇Xf(X,y),−∇yf(X,y))− (∇Xf(X̃, ỹ),−∇yf(X̃, ỹ))‖ ≤ β‖(X,Y)− (X̃, ỹ)‖, where ‖ · ‖ denotes the Euclidean norm over the product space Sn × Y. The relationship between the full Lipschitz parameter β and its components βX , βy, βXy, βyX can be written as β = √ 2 max {√
β2X + β 2 yX , √ β2y + β 2 Xy } . See proof in Appendix B.
The following lemma highlights a connection between the gradient of a saddle-point of (4) and subgradients of an optimal solution to (1) for which the first order optimality condition holds. One of the connections we will be interested in, is that GSC for Problem (1) implies GSC (with the same parameters) for Problem (4). However, to prove this specific connection we require an additional structural assumption on the objective function g(·). We note that this assumption holds for all applications mentioned in this paper. Assumption 2. g(X) is of the form g(X) = h(X) + maxy∈K y>(A(X)−b), where h(·) is smooth and convex, and A is a linear map. Lemma 6. If (X∗,y∗) is a saddle-point of Problem (4) then X∗ is an optimal solution to Problem (1), ∇Xf(X∗,y∗) ∈ ∂g(X∗), and for all X ∈ Sn it holds that 〈X − X∗,∇Xf(X∗,y∗)〉 ≥ 0. Conversely, under Assumption 2, if X∗ is an optimal solution to Problem (1), and G∗ ∈ ∂g(X∗) which satisfies 〈X −X∗,G∗〉 ≥ 0 for all X ∈ Sn, then there exists y∗ ∈ arg maxy∈K f(X∗,y) such that (X∗,y∗) is a saddle-point of Problem (4), and ∇Xf(X∗,y∗) = G∗.
The proof is given in Appendix C.1. The connection between the gradient of an optimal solution to the saddle-point problem and a subgradient of a corresponding optimal solution in the equivalent nonsmooth problem established in Lemma 6, naturally leads to the formulation of the following generalized strict complementarity assumption for saddle-point problems.
Assumption 3 (generalized strict complementarity for saddle-points). We say a saddle-point (X∗,y∗) ∈ Sn×K of Problem (4) with rank(X∗) = r∗ satisfies the generalized strict complementarity assumption with parameters r ≥ r∗, δ > 0, if λn−r(∇Xf(X∗,y∗))− λn(∇Xf(X∗,y∗)) ≥ δ. Remark 1. Note that under Assumption 2, due to Lemma 6, GSC with parameters r, δ for some optimal solution X∗ to Problem (1) implies GSC with parameters r, δ to a corresponding saddle-point (X∗,y∗) of Problem (4). Nevertheless, Assumption 2 is not necessary for proving our convergence results for Problem (4), which are directly stated in terms of Assumption 3.

heading: 4 Projected extragradient method with low-rank projections

text:
In this section we formally state and prove our main result: the projected extragradient method for the saddle-point Problem (4), when initialized in the proximity of a saddle-point which satisfies GSC (Assumption 3), converges with its original O(1/t) rate while requiring only two low-rank SVD computations per iteration.
Algorithm 1 Projected extragradient method for saddle-point problems (see also [23, 30]) Input: sequence of step-sizes {ηt}t≥1 Initialization: (X1,y1) ∈ Sn ×K for t = 1, 2, ... do Zt+1 = ΠSn [Xt − ηt∇Xf(Xt,yt)] wt+1 = ΠK[yt + ηt∇yf(Xt,yt)] Xt+1 = ΠSn [Xt − ηt∇Xf(Zt+1,wt+1)] yt+1 = ΠK[yt + ηt∇yf(Zt+1,wt+1)]
end for
Theorem 1 (main theorem). Fix an optimal solution (X∗,y∗) ∈ Sn ×K to Problem (4). Let r̃ denote the multiplicity of λn(∇Xf(X∗,y∗)) and for any r ≥ r̃ define δ(r) = λn−r(∇Xf(X∗,y∗)− λn(∇Xf(X∗,y∗). Let {(Xt,yt)}t≥1 and {(Zt,wt)}t≥2 be the sequences of iterates generated by Algorithm 1 with a fixed step-size η = min { 1
2 √ β2X+β 2 yX , 1 2 √ β2y+β 2 Xy , 1βX+βXy , 1 βy+βyX
} . Assume
the initialization (X1,y1) satisfies ‖(X1,y1)− (X∗,y∗)‖ ≤ R0(r), where
R0(r) := η
(1 + √ 2) ( 1 + (2 + √ 2)ηmax{βX , βXy} ) max{√r̃δ(r − r̃ + 1) 2 , δ(r) (1 + 1/ √ r̃) } .
Then, for all t ≥ 1, the projections ΠSn [Xt − η∇Xf(Xt,yt)] and ΠSn [Xt − η∇Xf(Zt+1,wt+1)] can be replaced with their rank-r truncated counterparts (see (3)) without changing the sequences {(Xt,yt)}t≥1 and {(Zt,wt)}t≥2, and for any T ≥ 0 it holds that
1
T T∑ t=1 max y∈K f(Zt+1,y)− 1 T T∑ t=1 min X∈Sn f(X,wt+1)
≤ D2 max
{√ β2X + β 2 yX , √ β2y + β 2 Xy, 1 2 (βX + βXy), 1 2 (βy + βyX) } T ,
where D := sup(X,y),(Z,w)∈Sn×K ‖(X,y)− (Z,w)‖. Remark 2. Note that Theorem 1 implies that if standard strict complementarity holds for Problem (4), that is Assumption 3 holds with r = r∗ = rank(X∗) and some δ > 0, then only rank-r∗ SVDs are required so that Algorithm 1 converges with the guaranteed convergence rate of O(1/t), when initialized with a “warm-start”. Furthermore, by using SVDs of rank r > r∗, with moderately higher values of r, we can increase the radius of the ball in which Algorithm 1 needs to be initialized quite significantly.
The complete proof of Theorem 1 is given in Appendix D. Below we give a short sketch of the proof.
Proof sketch for Theorem 1. The convergence rate of Algorithm 1 is well known in the literature (see Appendix D.1 for completeness). Thus, we focus on the novel part of the low-rank projections. For simplicity consider the case r = r̃. From (2) it can be deduced that for any P ∈ Sn it holds that rank(ΠSn [P]) ≤ r̃ if and only if the condition ∑r̃ i=1 λi(P) ≥ 1 + r̃λr̃+1(P) holds. Denote P∗ = X∗ − η∇Xf(X∗,y∗). From Lemma 6 and Lemma 3 it follows that ∀i ≤ rank(X∗) : λi(P∗) = λi(X∗)− ηλn(∇Xf(X∗,y∗)); ∀i > rank(X∗) : λi(P∗) = −ηλn−i+1(∇Xf(X∗,y∗)).
Using this and the fact that λn−i+1(∇fX(X∗,y∗)) = λn(∇fX(X∗,y∗)) for all i ≤ r̃ we have, r̃∑ i=1 λi(P ∗) = rank(X∗)∑ i=1 λi(X ∗)− η r̃∑ i=1 λn(∇Xf(X∗,y∗)) = Tr(X∗)− ηr̃λn(∇Xf(X∗,y∗)) = 1− ηr̃λn−r̃(∇Xf(X∗,y∗)) + ηr̃δ(r̃) = 1 + r̃λr̃+1(P∗) + ηr̃δ(r̃). (5) Thus, P∗ = X∗ − η∇Xf(X∗,y∗) not only satisfies the condition that ensures its projection onto Sn is of rank at most r̃, but it satisfies it with a positive slack of ηr̃δ(r̃). Also, for any (X,y) sufficiently close to (X∗,y∗), using the smoothness of f , we have that P = X− η∇Xf(X,y) is close to P∗. Thus, by applying perturbation bounds for the eigenvalues of symmetric matrices to Eq. (5) and using the fact that the positive slack ηr̃δ(r̃) in the RHS of (5) allows to absorb sufficiently small errors, we can establish that for such points (X,y) it holds that rank(ΠSn [P]) ≤ r̃. A similar argument shows that the second primal step in Algorithm 1: Xt+1 = ΠSn [Xt − η∇Xf(Zt+1,wt+1)], also results in matrices of rank at most r̃ when Zt,wt are also in the proximity of X∗,y∗. Finally, we prove a complementary argument, that when initialized in the proximity of (X∗,y∗), the iterates of Algorithm 1 stay in the proximity of (X∗,y∗) throughout all iterations.
Remark 3. Note that when applying Algorithm 1 towards minimizing a nonsmooth objective of the form g(X) = maxy∈K f(X,y) (as discussed in Section 3), a guarantee of the form 1 T ∑T t=1 maxy∈K f(Zt+1,y) − 1 T ∑T t=1 minX∈Sn f(X,wt+1) ≤ , for some > 0, which is what we get from Theorem 1, implies in particular that mint∈[T ] g(Xt)− g(X∗) ≤ , where X∗ is a minimizer of g(·) over Sn. Remark 4. A downside of considering the saddle-point formulation (4) when attempting to solve Problem (1) that arises from Theorem 1, is that not only do we need a “warm-start” initialization for the original primal matrix variable X, in the saddle-point formulation we need a “warm-start” for the primal-dual pair (X,y). Nevertheless, as we demonstrate extensively in Section 5, it seems that very simple initialization schemes work very well in practice.

heading: 4.1 Efficiently-computable certificates for correctness of low-rank projections

text: Since Theorem 1 only applies in some neighborhood of an optimal solution, it is of interest to have a procedure for verifying if the rank-r truncated projection of a given point indeed equals the exact Euclidean projection. In particular, from a practical point of view, it does not matter whether the conditions of Theorem 1 hold. In practice, as long as the truncated projection (see (3)) equals the exact projection (see (2)), we are guaranteed that Algorithm 1 converges correctly with rate O(1/t), without needing to verify any other condition. Luckily, the expression in (2) which characterizes the structure of the Euclidean projection onto the spectrahedron, yields exactly such a verification procedure. As already noted in [17], for any X ∈ Sn, we have Π̂rSn [X] = ΠSn [X] if an only if the condition ∑r i=1 λi(X) ≥ 1 + r · λr+1(X) holds. Note that verifying this condition simply requires increasing the rank of the SVD computation by one, i.e., computing a rank-(r+ 1) SVD of the matrix to project rather than a rank-r SVD.

heading: 5 Empirical evidence

text:
The goal of this section is to bring empirical evidence in support of our theoretical approach. We consider various tasks that take the form of minimizing a composite objective, i.e., the sum of a
smooth convex function and a nonsmooth convex function, where the nonsmoothness comes from either an `1-norm or `2-norm regularizer / penalty term, over a τ -scaled spectrahedron. In all cases the nonsmooth objective can be written as a saddle-point with function f(X,y) which is linear in y and in particular satisfies Assumption 2.
The tasks considered include 1. sparse PCA, 2. robust PCA, 3. low-rank and sparse recovery, 4. phase synchronization, and 5. linearly-constrained low-rank estimation, under variety of parameters. Due to lack of space many of the results are deferred to Appendix F.
For all tasks considered we generate random instances, and examine the sequences of iterates generated by Algorithm 1 {(Xt,yt)}t≥1, {(Zt,wt)}t≥2, when initialized with simple initialization procedures. Out of both sequences generated, we choose our candidate for the optimal solution to be the iterate for which the dual-gap, which is a certificate for optimality, is smallest. See Appendix E.
In all tasks considered the goal is to recover a ground-truth low-rank matrix M0 ∈ Sn from some noisy observation of it M = M0 + N, where N ∈ Sn is a noise matrix. We measure the signal-to-noise ratio (SNR) as ‖M0‖2F / ‖N‖2F . In all experiments we measure the relative
initialization error by ∥∥∥Tr(M0)τ X1 −M0∥∥∥2
F
/ ‖M0‖2F , and similarly we measure the relative re-
covery error by ∥∥∥Tr(M0)τ X∗ −M0∥∥∥2
F
/ ‖M0‖2F . Note that in some of the experiments we take
τ < Tr(M0) to prevent the method from overfitting the noise. In addition, we measure the (standard) strict complementarity parameter which corresponds to the eigen-gap gap(∇Xf(X∗,y∗)) := λn−r(∇Xf(X∗,y∗))− λn(∇Xf(X∗,y∗)), r = rank(M0). In all experiments we use SVDs of rank r = rank(M0) to compute the projections in Algorithm 1 according to the truncated projection given in (3). To certify the correctness of these low-rank projections (that is, that they equal the exact Euclidean projection) we confirm that the inequality∑r i=1 λi(Pj) ≥ τ + r · λr+1(Pj) always holds for P1 = Xt − η∇Xf(Xt,Yt) and P2 = Xt − η∇Xf(Zt+1,Wt+1) (see also Section 4.1). Indeed, we can now already state our main observation from the experiments:
In all tasks considered and for all random instances generated, throughout all iterations of Algorithm 1, when initialized with a simple “warm-start” strategy and when computing only rank-r truncated projections, r = rank(M0), the truncated projections of P1 = Xt − η∇Xf(Xt,Yt) and P2 = Xt − η∇Xf(Zt+1,Wt+1) equal their exact full-rank counterparts. That is, Algorithm 1, using only rank-r SVDs, computed exactly the same sequences of iterates it would have computed if using full-rank SVDs.
Aside from the above observation, in the sequel we demonstrate that all models considered indeed satisfy that: 1. the returned solution, denoted (X∗,y∗), is of the same rank as the ground-truth matrix and satisfies the strict complementarity condition with non-negligible parameter (measured by the eigengap λn−r(∇Xf(X∗,y∗))− λn(∇Xf(X∗,y∗))), 2. the recovery error of the returned solution indeed improves significantly over the error of the initialization point.
Sparse PCA: We consider the sparse PCA problem in a well known convex formulation taken from [12] and its equivalent saddle-point formulation:
min Tr(X)=1,
X 0
〈X,−M〉+ λ‖X‖1 = min Tr(X)=1,
X 0
max ‖Y‖∞≤1
{〈X,−M〉+ λ〈X,Y〉},
where M = zz> + c2 (N + N >) is a noisy observation of a rank-one matrix zz>, with z being a sparse unit vector. Each entry zi is chosen to be 0 with probability 0.9 and U{1, . . . , 10} with probability 0.1, and then we normalize z to be of unit norm. We test the results obtained when adding different magnitudes of Gaussian or uniform noise. We set the signal-to-noise ratio (SNR) to be a constant. Thus, we set the noise level to c = 2SNR·‖N+N>‖F for our choice of SNR. We initialize the X variable with the rank-one approximation of M. That is, we take X1 = u1u>1 , where u1 is the top eigenvector of M. For the Y variable we initialize it with Y1 = sign(X1) which is a subgradient of ‖X1‖1. We set the step-size to η = 1/(2λ) and we set the number of iterations to T = 1000 and for any set of parameters we average the measurements over 10 i.i.d. runs.
Robust PCA: We consider the robust PCA problem [28] in the following formulation: min
Tr(X)=τ, X 0
‖X−M‖1 = min Tr(X)=τ,
X 0
max ‖Y‖∞≤1
〈X−M,Y〉,
where M = rZ0Z>0 + 1 2 (N+N >) is a sparsely-corrupted observation of some rank-r matrix Z0Z>0 . We choose Z0 ∈ Rn×r to be a random unit Frobenius norm matrix. For N ∈ Rn×n, we choose each entry to be 0 with probability 1 − 1/ √ n and otherwise 1 or −1 with equal probability. We initialize the X variable with the projection X1 = Π{Tr(X)=τ, X 0}[M], and the Y variable with Y1 = sign(X1 −M). We vary the rank of Z0, which in turn determines the step-size η. We set the trace bound to τ = 0.95 · Tr(rZ0Z>0 ), and the number of iterations to T = 30, 000. For every set of parameters we average the measurements over 10 i.i.d. runs.
Linearly-constrained low-rank estimation: Consider the following penalized formulation: min
Tr(X)=1, X 0
〈X,−M〉+ λ‖A(X)− b‖2 = min Tr(X)=1,
X 0
max ‖y‖2≤1
〈X,−M〉+ λ〈A(X)− b,y〉,
where M = z0z>0 + c 2 (N + N >) is a noisy observation of some rank-one matrix z0z>0 such that ‖z0‖2 = 1 and the noise matrix is chosen N ∼ N (0, In). We take A(X) = (〈A1,X〉, . . . , 〈Am,X〉)> with matrices A1, . . . ,Am ∈ Sn of the form Ai = viv>i such that vi ∼ N (0, 1). We take b ∈ Rm such that bi = 〈Ai, z0z>0 〉. We initialize the X variable with the rank-one approximation of M. That is, we take X1 = u1u>1 , where u1 is the top eigenvector of M. The y variable is initialized with y1 = (A(X1) − b)/‖A(X1) − b‖2. We set the number of constraints to m = n, the penalty parameter to λ = 2, and the step-size to η = 1/(2λ). We set the number of iterations in each experiment to T = 2000 and for each value of n we average the measurements over 10 i.i.d. runs.
Low-rank and sparse matrix recovery: We consider the problem of recovering a simultaneously low-rank and sparse covariance matrix [35], which can be written as the following saddle-point optimization problem:
min Tr(X)=1,
X 0
1 2 ‖X−M‖2F + λ‖X‖1 = min
Tr(X)=τ, X 0
max ‖Y‖∞≤1
1 2 ‖X−M‖2F + λ〈X,Y〉,
where M = Z0Z0> + c2 (N + N >) is a noisy observation of some low-rank and sparse covariance matrix Z0Z0>. We choose Z0 ∈ Rn×r to be a sparse matrix where each entry Z0i,j is chosen to be
0 with probability 0.9 and U{1, . . . , 10} with probability 0.1, and then we normalize Z0 to be of unit Frobenius norm. We choose N ∼ N (0.5, In). We test the model with rank(Z0Z>0 ) = 5, 10. We set the signal-to-noise ratio (SNR) to be a constant and set the noise level to c = 2‖Z0Z > 0 ‖F
SNR·‖N+N>‖F for our choice of SNR. We initialize the X variable with the rank-r approximation of M. That is, we take X1 = Urdiag ( Π∆τ,r [diag(−Λr)] ) U>r , where UrΛrU > r is the rank-r eigen-decomposition of M
and ∆τ,r = {z ∈ Rr | z ≥ 0, ∑r i=1 zi = τ} is the simplex of radius τ in Rr. For the Y variable we initialize it with Y1 = sign(X1) which is a subgradient of ‖X1‖1. We set the step-size to η = 1, τ = 0.7 · Tr(Z0Z>0 ), and the number of iterations in each experiment to T = 2000. For each value of r and n we average the measurements over 10 i.i.d. runs.

heading: 6 Discussion

text: This work expands upon a line of research that aims to harness the ability of convex relaxations to produce low-rank and high-quality solutions to important low-rank matrix optimization problems, while insisting on methods that, at least locally, store and manipulate only low-rank matrices. Focusing on the challenging case of nonsmooth objective functions and following our evidence for the difficulties of obtaining such a result for subgradient methods (Lemma 5), we consider tackling nonsmooth objectives via saddle-point formulations. We prove that indeed under a generalized strict complementarity condition, a state-of-the-art method for convex-concave saddle-point problems converges locally while storing and manipulating only low-rank matrices. Extensive experiments over several tasks demonstrate that our conceptual approach of utilizing low-rank projections for more efficient optimization is not only of theoretical merit, but indeed seems to work well in practice.

heading: Acknowledgements

text: This research was supported by the ISRAEL SCIENCE FOUNDATION (grant No. 1108/18).

title: Low-Rank Extragradient Method for Nonsmooth and Low-Rank Matrix Optimization Problems


summaries: This paper proposed a CATE estimation approach when the treatment variable $T$ is a graph. Under a separable assumption (A3) between confounders $X$ and the treatment $T$ in the outcome structural equation ($Y$), Robinson's decomposition of CATE for binary $T$ can be applied. The performance of the authors' proposal was evaluated in both synthetic experiments and real data analysis.

summaries: The paper proposed a CATE estimation method for graph-structured treatment by Robinson Decomposition. It is an extension of Robinson Decomposition, and the experiment shows the effectiveness of GIN on two small world datasets. The writing is good, but some contents are missing, so that the paper is not that easy to read. My detailed comments and concerns are described in the following.

summaries: This paper proposes graph intervention network (GIN) to estimate conditional average causal effects when treatments are graph-structured. GIN builds on an extension of traditional Robinson decomposition to arbitrary treatment types. In particular, GIN factorizes the outcome Y into an inner product of two terms $g(X)^\top h(T)$. During training, GIN first estimates a mean outcome model $m(X)$ and then estimate $g,h$ using the estimated $m(X)$ function. The method is evaluated on two datasets with synthetic outcome functions.

summaries: The paper considers the problem of estimating conditional average treatment effect when the treatment is a graph. The paper generalizes the Robinson decomposition, adapts the R-learner to this graph-treatment setting and introduces a two-stage algorithm. The performance of the proposed algorithm is compared to a few benchmarks through numerical experiments.

summaries: This paper studies the estimation of conditional average treatment effect (CATE) with multiple treatments that are graph structured. The proposed method extends the R-learner (Nie and Wager, 2020), which is designed for binary treatment, by learning a representation that reduces the dimension of the treatments. The representation is parametrized by a graph neural network that takes the graph structure into account. Unlike the standard R-learner that learns the nuisance components $E[Y|X]$ and $E[T|X]$ separately from the CATE, the proposed method only estimates $E[Y|X]$ separately while learns the representation $h(T)$, the propensity features $E[h(T)|X]$, and CATE jointly via gradient-based methods.

abstractText: We address the estimation of conditional average treatment effects (CATEs) for structured treatments (e.g., graphs, images, texts). Given a weak condition on the effect, we propose the generalized Robinson decomposition, which (i) isolates the causal estimand (reducing regularization bias), (ii) allows one to plug in arbitrary models for learning, and (iii) possesses a quasi-oracle convergence guarantee under mild assumptions. In experiments with small-world and molecular graphs we demonstrate that our approach outperforms prior work in CATE estimation.

authors:
A. Alaa
M. van der Schaar

title: Limits of estimating heterogeneous treatment effects: Guidelines for practical algorithm design

venue: In Dy,

authors:
A.M. Alaa
M. van der Schaar

title: Bayesian inference of individualized treatment effects using multi-task gaussian processes

venue: Advances in Neural Information Processing Systems,

authors:
S. Athey
G. Imbens

title: Recursive partitioning for heterogeneous causal effects

authors:
S. Athey
S. Wager

title: Estimating treatment effects with causal forests: An application

venue: arXiv preprint arXiv:1902.07409,

authors:
S. Athey
J. Tibshirani
S. Wager

title: Generalized random forests

venue: The Annals of Statistics,

authors:
I. Bica
J. Jordon
M. van der Schaar

title: Estimating the effects of continuous-valued interventions using generative adversarial networks

venue: Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems

authors:
A. Caron
I. Manolopoulou
G. Baio

title: Estimating individual treatment effects using non-parametric regression models: a review

venue: arXiv preprint arXiv:2009.06472,

authors:
V. Chernozhukov
D. Chetverikov
M. Demirer
E. Duflo
C. Hansen
W. Newey
J. Robins

title: Double/debiased machine learning for treatment and structural parameters

venue: The Econometrics Journal, 21(1):C1–C68,

authors:
A. Curth
M. van der Schaar

title: Nonparametric estimation of heterogeneous treatment effects: From theory to learning algorithms

venue: The 24th International Conference on Artificial Intelligence and Statistics,

authors:
M. Fey
J.E. Lenssen

title: Fast graph representation learning with PyTorch Geometric

venue: In ICLR Workshop on Representation Learning on Graphs and Manifolds,

authors:
S. Flaxman
S. Mishra
A. Gandy
H.J.T. Unwin
T.A. Mellan
H. Coupland
C. Whittaker
H. Zhu
T. Berah
Eaton
J. W

title: Estimating the effects of non-pharmaceutical interventions on covid-19 in europe

authors:
Health Statistics
N. C

title: 2007–2008 national health and nutrition examination survey (nhanes). US Department of Health and Human Services, Centers for Disease Control and Prevention

authors:
A. Gretton
K. Fukumizu
C.H. Teo
L. Song
B. Schölkopf
A.J. Smola

title: A kernel statistical test of independence

venue: Advances in Neural Information Processing Systems 20, Proceedings of the Twenty-First Annual Conference on Neural Information Processing Systems,

authors:
P.R. Hahn
J.S. Murray
Carvalho
C. M

title: Bayesian regression tree models for causal inference: Regularization, confounding, and heterogeneous effects (with discussion)

venue: Bayesian Analysis,

authors:
S. Harada
H. Kashima

title: Graphite: Estimating individual effects of graph-structured treatments

venue: arXiv preprint arXiv:2009.14061,

authors:
J. Hartford
G. Lewis
K. Leyton-Brown
M. Taddy

title: Deep IV: A flexible approach for counterfactual prediction

venue: Proceedings of the 34th International Conference on Machine Learning,

authors:
T. Hatt
S. Feuerriegel

title: Estimating average treatment effects via orthogonal regularization

venue: arXiv preprint arXiv:2101.08490,

title: Bayesian nonparametric modeling for causal inference

authors:
G. Imbens
D. Rubin

title: Causal Inference for Statistics, Social, and Biomedical Sciences

title: Nonparametric estimation of average treatment effects under exogeneity: A review

venue: The Review of Economics and Statistics,

authors:
A. Jesson
S. Mindermann
U. Shalit
Y. Gal

title: Identifying causal-effect inference failure with uncertainty-aware models

authors:
A. Jesson
S. Mindermann
Y. Gal
U. Shalit

title: Quantifying ignorance in individual-level causal-effect estimates under hidden confounding

venue: In International Conference on Machine Learning. PMLR,

authors:
F. Johansson
U. Shalit
D. Sontag

title: Learning representations for counterfactual inference

venue: In International conference on machine learning,

title: DeepMatch: Balancing deep covariate representations for causal inference using adversarial training

venue: Proceedings of the 37th International Conference on Machine Learning,

title: Optimal doubly robust estimation of heterogeneous causal effects

venue: arXiv preprint arXiv:2004.14497,

authors:
E.H. Kennedy
Z. Ma
M.D. McHugh
D.S. Small

title: Nonparametric methods for doubly robust estimation of continuous treatment effects

venue: Journal of the Royal Statistical Society. Series B, Statistical Methodology,

authors:
D.P. Kingma
J. Ba

title: Adam: A method for stochastic optimization

venue: 3rd International Conference on Learning Representations,

authors:
S.R. Künzel
J.S. Sekhon
P.J. Bickel
B. Yu

title: Metalearners for estimating heterogeneous treatment effects using machine learning

title: Propensity score weighting for causal inference with multiple treatments

venue: Annals of Applied Statistics,

authors:
M.J. Lopez
R. Gutman

title: Estimation of causal effects with multiple treatments: a review and new ideas

venue: Statistical Science,

authors:
C. Louizos
U. Shalit
J.M. Mooij
D. Sontag
R. Zemel
M. Welling

title: Causal effect inference with deep latent-variable models

authors:
D. Lu
C. Tao
J. Chen
F. Li
F. Guo
L. Carin

title: Reconsidering generative objectives for counterfactual reasoning

authors:
K.W. Ma
J.P. Lewis
W.B. Kleijn

title: The HSIC bottleneck: Deep learning without back-propagation

venue: In The Thirty-Fourth AAAI Conference on Artificial Intelligence,

authors:
A. Mastouri
Y. Zhu
L. Gultchin
A. Korba
R. Silva
M.J. Kusner
A. Gretton
K. Muandet

title: Proximal causal learning with kernels: Two-stage estimation and moment restriction

authors:
S. Mendelson
J. Neeman

title: Regularization in kernel learning

authors:
W. Miao
Z. Geng
E.T. Tchetgen

title: Identifying causal effects with proxy variables of an unmeasured confounder

venue: arXiv preprint arXiv:1609.08816,

authors:
C. Morris
M. Ritzert
M. Fey
W.L. Hamilton
J.E. Lenssen
G. Rattan
M. Grohe

title: Weisfeiler and leman go neural: Higher-order graph neural networks

venue: Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
R. Nabi
T. McNutt
I. Shpitser

title: Semiparametric causal sufficient dimension reduction of high dimensional treatments

venue: arXiv preprint arXiv:1710.06727,

authors:
B. Neal
Huang
C.-W
S. Raghupathi

title: Realcause: Realistic causal inference benchmarking

venue: arXiv preprint arXiv:2011.15007,

authors:
L. Nie
M. Ye
qiang liu
D. Nicolae

title: VCN}et and functional targeted regularization for learning causal effects of continuous treatments

venue: In International Conference on Learning Representations,

authors:
X. Nie
S. Wager

title: Quasi-oracle estimation of heterogeneous treatment

venue: effects. Biometrika,

authors:
A. Paszke
S. Gross
F. Massa
A. Lerer
J. Bradbury
G. Chanan
T. Killeen
Z. Lin
N. Gimelshein
L. Antiga
A. Desmaison
A. Köpf
E. Yang
Z. DeVito
M. Raison
A. Tejani
S. Chilamkurthy
B. Steiner
L. Fang
J. Bai
S. Chintala

title: Pytorch: An imperative style, high-performance deep learning library

venue: Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems

title: Causality: models, reasoning, and inference

title: Causal inference for spatial treatments

venue: arXiv preprint arXiv:2011.00373,

authors:
R. Ramakrishnan
P.O. Dral
M. Rupp
O.A. von Lilienfeld

title: Quantum chemistry structures and properties of 134 kilo molecules

venue: Scientific Data,

title: Root-n-consistent semiparametric regression

authors:
P.R. Rosenbaum
D.B. Rubin

title: The central role of the propensity score in observational studies for causal effects

authors:
L. Ruddigkeit
R. van Deursen
L.C. Blum
Reymond
J.-L

title: Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17

venue: Journal of Chemical Information and Modeling,

authors:
M.S. Schlichtkrull
T.N. Kipf
P. Bloem
R. van den Berg
I. Titov
M. Welling

title: Modeling relational data with graph convolutional networks

venue: The Semantic Web - 15th International Conference,

authors:
P. Schwab
L. Linhardt
W. Karlen

title: Perfect match: A simple method for learning representations for counterfactual inference with neural networks

venue: arXiv preprint arXiv:1810.00656,

authors:
P. Schwab
L. Linhardt
S. Bauer
J.M. Buhmann
W. Karlen

title: Learning Counterfactual Representations for Estimating Individual Dose-Response Curves

venue: In AAAI Conference on Artificial Intelligence,

authors:
U. Shalit
F.D. Johansson
D. Sontag

title: Estimating individual treatment effect: generalization bounds and algorithms

venue: In International Conference on Machine Learning,

authors:
C. Shi
D. Blei
V. Veitch

title: Adapting neural networks for the estimation of treatment effects

authors:
C. Shi
V. Veitch
D. Blei

title: Invariant representation learning for treatment effect estimation

venue: arXiv preprint arXiv:2011.12379,

title: Observational-interventional priors for dose-response learning

authors:
R. Singh
M. Sahani
A. Gretton

title: Kernel instrumental variable regression

authors:
E.J.T. Tchetgen
A. Ying
Y. Cui
X. Shi
W. Miao

title: An introduction to proximal causal learning

venue: arXiv preprint arXiv:2009.10982,

authors:
S. Wager
S. Athey

title: Estimation and inference of heterogeneous treatment effects using random forests

venue: Journal of the American Statistical Association,

authors:
J.N. Weinstein
E.A. Collisson
G.B. Mills
K.R.M. Shaw
B.A. Ozenberger
K. Ellrott
I. Shmulevich
C. Sander
J.M. Stuart

title: The cancer genome atlas pan-cancer analysis project

venue: Nature genetics,

authors:
L. Xu
Y. Chen
S. Srinivasan
N. de Freitas
A. Doucet
A. Gretton

title: Learning deep features in instrumental variable regression

authors:
L. Yao
S. Li
Y. Li
M. Huai
J. Gao
A. Zhang

title: Representation learning for treatment effect estimation from observational data

authors:
D.W. Zhang
G.J. Burghouts
C.G.M. Snoek

title: Set prediction without imposing structure as conditional density estimation

authors:
Z. Zhou
S. Athey
S. Wager

title: Offline multi-action policy learning: Generalization and optimization

venue: arXiv preprint arXiv:1810.04778,

authors:
Y. Zhu
D.L. Coffman
D. Ghosh

title: A boosting algorithm for estimating generalized propensity scores with continuous treatments

venue: Journal of Causal Inference,

heading: 1 Introduction

text:
Estimating feature-level causal effects, so-called conditional average treatment effects (CATEs), from observational data is a fundamental problem across many domains. Examples include understanding the effects of non-pharmaceutical interventions on the transmission of COVID-19 in a specific region [12], how school meal programs impact child health [13], and the effects of chemotherapy drugs on cancer patients [52]. Supervised learning methods face two challenges in such settings: (i) missing interventions, the fact that we only observe one treatment for each individual means models must extrapolate to new treatments without access to ground truth, and (ii) confounding factors that affect both treatment assignment and the outcome means that extrapolation from observation to intervention requires assumptions. Many approaches have been proposed to overcome these issues [1, 2, 3, 4, 5, 6, 7, 9, 10, 15, 18, 19, 21, 22, 23, 25, 27, 29, 33, 39, 41, 42, 45, 52, 56, 57, 60, 64, 67].
In many cases, treatments are naturally structured. For instance, a drug is commonly represented by its molecular structure (graph), the nutritional content of a meal as a food label (text), and geographic regions affected by a new policy as a map (image). Taking this structure into account can provide several advantages: (i) higher data-efficiency, (ii) capability to work with many treatments, and (iii) generalizing to unseen treatments during test time. However, the vast majority of prior work operates on either binary or continuous scalar treatments (structured treatments are rarely considered, a notable exception to this trend is Harada & Kashima [16] which we describe in Section 2).
To estimate CATEs with structured interventions, our contributions include:
• Generalized Robinson decomposition (GRD): A generalization of the Robinson decomposition [47] to treatments that can be vectorized as a continuous embedding. This GRD reveals a learnable ∗Correspondence to jean.kaddour.20@ucl.ac.uk
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
pseudo-outcome target that isolates the causal component of the observed signal by partialling out confounding associations. Further, it allows one to learn the nuisance and target functions using any supervised learning method, thus extending recent work on plug-in estimators [42, 29].
• Quasi-oracle convergence guarantee: A result that shows that given access to estimators of certain nuisance functions, as long as the estimates converge at an O(n−1/4) rate, the target estimator for the CATE achieves the same error bounds as an oracle who has ground-truth knowledge of both nuisance components, the propensity features, and conditional mean outcome.
• Structured Intervention Networks (SIN): A practical algorithm using GRD, representation learning, and alternating gradient descent. Our PyTorch [43] implementation is online.2
• Evaluation metrics designed for structured treatments. Since previous evaluation protocols of CATE estimators have mostly focused on binary or scalar-continuous treatment settings, we believe that our proposed evaluation metrics can be useful for comparing future work.
• Experimental results with graph treatments in which SIN outperforms previous approaches.

heading: 2 Related Work

text: Closest to our work is GraphITE [16], a method that learns representations of graph interventions for CATE estimation. They propose to minimize prediction loss plus a regularization term that aims to control for confounding based on the Hilbert-Schmidt Independence Criterion (HSIC) [14]. This technique suffers from two drawbacks: (i) the HSIC requires multiplication of kernel matrices and scales quadratically in the batch size; (ii) selecting the HSIC kernel hyper-parameter is not straightforward, as ground-truth CATEs are never observed, and empirical loss does not bound CATE estimation error [1]. We discuss other related work not on structured treatments in Appendix A.

heading: 3 Preliminaries

heading: 3.1 Conditional Average Treatment Effects (CATEs)

text:
Imagine a dataset where each example (xi, ti, yi) ∈ D represents a hospital patient’s medical history record xi, prescribed drug treatment ti, and health outcome yi, as illustrated in Figure 1 (Center). Further, we wish to understand how changing the treatment changes a patient’s health outcome. The CATE, τ ( t′, ti,xi ) , describes the expected change in outcome for individuals with history xi, when treatment ti is replaced by t′, depicted in Figure 1 (Right). In real-world scenarios, we only observe one outcome for each patient at one treatment level. Further, the patient’s pre-treatment health conditions xi influence both the doctor’s treatment prescription and outcome, thereby confounding the effect of the treatment on the outcome.
Formally, we have the dataset D= { (xi, ti, yi) }n i=1
sampled from a joint distribution p (X,T, Y ), where Y = f (X,T) + ε, as depicted in Figure 1 (Left). We define the causal effect of fixing
2https://github.com/JeanKaddour/SIN
treatment variable T ∈ T to a value t on outcome variable Y ∈ R using the do-operator [44] as E [ Y | do(T = t) ] . Crucially, this estimate differs from the conditional expectation E [ Y | T = t ] in that it describes the effect of an external entity intervening on T by fixing it to a value t (removing the edge X → T). We further condition on pre-treatment covariates X to define the conditional causal estimand E [ Y | X = x, do(T = t) ] . The conditional average treatment effect (CATE) is the difference between expected outcomes at different treatment values t, t′ for given covariates x,
τ(t′, t,x) , E [ Y | X = x, do(T = t′) ]︸ ︷︷ ︸ =:µt′ (x) −E [ Y | X = x, do(T = t) ]︸ ︷︷ ︸ =:µt(x) , (1)
where µt (x) is defined as the expected outcome for a covariate vector x under treatment t.
Because we do not observe both treatments t, t′ for a single covariate x, we need to make assumptions that allow us to identify the CATE from observational data.
Assumption 1. (Unconfoundedness) There are no confounders of the effect between T and Y beyond X. Therefore, Pr ( Y ≤ y | x, do(t) ) = Pr ( Y ≤ y | x, t ) , for all (x, t, y).
Assumption 2. (Overlap) It holds that 0 < p ( t | x ) < 1, for all (x, t).
Assumption 2 means that all sub-populations have some probability of receiving any value of treatment (otherwise, some τ(t′, t,x) may be undefined or impossible to estimate.) These assumptions allow us to estimate the causal quantity τ(t′, t,x) through statistical estimands:
τ ( t′, t,x ) = µt′ (x)− µt (x) = E [ Y | X = x,T = t′ ] − E [ Y | X = x,T = t ] . (2)
While one can model µt(x) with regression models, such approaches suffer from bias [9, 26, 29] due to two factors: (i) associations between X and T, due to confounding, makes it hard to identify the distinct contributions of X and T on Y , and (ii) regularization for predictive performance can harm effect estimation. Mitigating these biases relies on exposing and removing nuisance components. This transforms the optimization into a (regularized) regression problem that isolates the causal effect.

heading: 3.2 Robinson Decomposition

text:
One way to formulate such nuisance components is via the Robinson decomposition [47]. Originally a reformulation of the CATE for binary treatments, it was used by the R-learner [42] to construct a plug-in estimator. The R-learner exploits the decomposition by partialling out the confounding of X on T and Y . It also isolates the CATE, thereby removing regularization bias.
Let the treatment variable be T ∈ {0, 1} and the outcome model p ( y | x, t ) parameterized as
Y = f(X, T ) + ε ≡ µ0(X) + T × τb(X) + ε, (3)
where we define error term ε such that E [ ε | x, t ] = E [ ε | x ] = 0, and τb (x) , τ (1, 0,x).
Define the propensity score [48] e (x) , p ( T = 1 | x ) and the conditional mean outcome as
m (x) , E [ Y | x ] = µ0 (x) + e (x) τb (x) . (4)
From model (3) and the previous definitions, it follows that Y −m (X) = ( T − e (X) ) τb (X) + ε, (5)
allowing us to define the estimator
τ̂b (·) = arg min τb
{ 1
n n∑ i=1 ( ỹi − t̃i × τb (xi) )2 + Λ ( τb (·) )} , (6)
where ỹi , yi−m̂ (xi) and t̃i , ti−ê (xi) are pseudo-data points defined through estimated nuisance functions m̂(·), ê(·), which can be learned separately with any supervised learning algorithm.

heading: 4 The Generalized Robinson Decomposition

text: Our goal is to estimate the CATE τ(t′, t,x) for structured interventions t′, t (e.g., graphs, images, text) while accounting for the confounding of X on T and Y . Inspired by the Robinson decomposition, which has enabled flexible CATE estimation for binary treatments [6, 9, 33, 42], we propose the Generalized Robinson Decomposition from which we extract a pseudo-outcome that targets the causal effect. We demonstrate the usefulness of this decomposition from both a theoretical view (quasi-oracle convergence rate in Section 4.2) and practical view (Structured Intervention Networks in Section 5). For details on its motivation and derivation, we refer the reader to Appendix B.

heading: 4.1 Generalizing the Robinson Decomposition

text:
To generalize the Robinson decomposition to structured treatments, we introduce two concepts: (a) we assume that the causal effect is a product effect: the outcome function f∗ (X,T) can be written as an inner product of two separate functionals, one over the covariates and one over the treatment, and (b) propensity features, which partial out the effects from the covariates on the treatment features. Similar techniques have been previously shown to add to the robustness of estimation [9, 42]. Assumption 3. (Product effect) We consider the following partial parameterization of p(y | x, t),
Y = g (X) > h (T) + ε, (7)
where g : X → Rd, h : T → Rd and E[ε | x, t] = E [ ε | x ] = 0, for all (x, t) ∈ X × T .
This assumption is mild, as we can formally justify its universality. The following asserts that provided we allow the dimensionality of g and h to grow, we may approximate any arbitrary bounded continuous functions in C (X × T ) where X × T is compact. Proposition 1. (Universality of product effect) LetHX×T be a Reproducing Kernel Hilbert Space (RKHS) on the set X × T with universal kernel k. For any δ > 0, and any f ∈ HX×T , there is a d ∈ N such that there exist two d-dimensional vector fields g : X → Rd and h : T → Rd, where ‖f − g>h‖L2(PX×T ) ≤ δ. (Proof in Appendix C)
This assumption allows us to simplify the expression of the CATE for treatments t′, t, given x,
τ ( t′, t,x ) = g (x) > ( h ( t′ ) − h (t) ) . (8)
Define propensity features eh (x) , E [ h (T) | x ] and m (x) , E [ Y | x ] = g (x) > eh (x) .
Following the same steps as in Section 3.2, the Generalized Robinson Decomposition for eq. (7) is Y −m (X) = g (X)> ( h (T)− eh (X) ) + ε. (9)
Given nuisance estimates m̂(·), êh(·), we can use this decomposition to derive an optimization problem for h(·), g(·) (note êh(·) implicitly depends on h(·), we address this dependence in Section 5).
ĝ (·) , ĥ (·) , arg min g,h  1n n∑ i=1 ( Yi − m̂ (Xi)− g (Xi)> ( h (Ti)− êh (Xi) ))2 + Λ ( g (·) ) (10)

heading: 4.2 Quasi-oracle error bound of Generalized Robinson Decomposition

text:
We establish the main theoretical result of our paper: a quasi-oracle convergence guarantee for the Generalized Robinson Decomposition under a finite-basis representation of the outcome function. This result is analogous to the R-learner for binary CATEs [42]: when the true e (·) ,m (·) are unknown, and we only have access to the estimators ê (·) , m̂ (·), then as long as the estimates converge at n−1/4 rate, the estimator τ̂b (·) achieves the same error bounds as an oracle who has ground-truth knowledge of these two nuisance components.
More formally, provided the nuisance estimators m̂(·) and êh(·) converge at an O ( n−1/4 ) rate, our
CATE estimator will converge at an Õ(n− 1
2(1+p) ) rate for arbitrarily small p > 0, recovering the parametric convergence rate for when the true m(·) and eh(·) are provided as oracle quantities.
Our analysis assumes that the outcome E [ Y | X = x,T = t ] can be written as a linear combination of fixed basis functions. By Proposition 1, as long as we have enough basis functions, this representation is flexible enough to capture the true outcome function. Assumption 4. Let α(X) ∈ Rdα , β(T) ∈ Rdβ be fixed, known orthonormal basis features on X ∈ Rdx , T ∈ Rdt , respectively. The true outcome function f∗(x, t) = E[Y | X = x,T = t] can be written as f∗(x, t) = α>(x)Θ∗β(t) for some (unknown) matrix of coefficients Θ∗.
Note that by setting g = α>Θ∗ and h = β, we recover eq. (7). Additionally, we will need overlap in the basis features α(X ),β(T ). Assumption 5 (Overlap in features). The marginal distribution of features Pα(X )×β(T ) is positive, i.e. supp[Pα(X )×β(T )] = α(X )× β(T ).
Assumption 5 is typically weaker than requiring overlap in X and T, i.e., when dα, dβ dx, dt. With further technical assumptions specified in Appendix F, we establish the following theorem. Theorem 2. Let Θ∗ denote the representer of the true outcome function. Suppose Assumptions 5, 6, and 4 hold. Moreover, suppose that the propensity estimate êh is uniformly consistent,
sup x∈X ‖êh(x)− eh(x)‖ →p 0 (11)
and the L2 errors converge at rate E [{ m̂(X)−m∗(X) }2] ,E [ ‖êh (X)− eh (X) ‖2 ] = O(a2n) (12)
for some sequence an → 0, where (an) is such that an = O(n−κ) with κ > 14 . Further, we define the regret as the excess risk
R ( Θ̂n ) , L ( Θ̂n ) − L (Θ∗) , L (Θ) , E [{( Y −m∗ (X) ) −α (X) Θ ( β (T)− eh (X) )}2] .
(13)
Suppose that we obtain Θ̂n via a penalized basis function regression variant of the Generalized Robinson Decomposition, with a properly chosen penalty Λn ( ‖Θ̂n‖2 ) (specified in the proof). Then,
Θ̂n satisfies the regret bound: R ( Θ̂n ) = Õ(r2n) with rn = n − 1 2(1+p) for arbitrarily small p > 0.

heading: 5 Structured Intervention Networks

text: We introduce Structured Intervention Networks (SIN), a two-stage training algorithm for neural networks, which enables flexibility in learning complex causal relationships, and scalability to large data-sets. This implementation of GRD strikes a balance between theory and practice: while we assumed fixed basis-functions in Section 4.2, in practice, we often need to learn the feature maps from data. We leave the convergence analysis of this representation learning setting for future work.

heading: 5.1 Training Algorithm

text:
We propose to simultaneously learn feature maps ĝ (X) , ĥ (T) using alternating gradient descent, so that they can adapt to each other. A remaining challenge is that learning êh(X) is now entangled with learning ĥ (T). While the R-learner is based on the idea of cross-fitting, where at each data point i we pick estimates of the nuisances that do not use that data point, we introduce a pragmatic representation learning approach for (ĝ, ĥ) that does not use cross-fitting3.
3We could in principle use cross-fitting for êh, although the loop between fitting ĥ alternating with êh would break the overall independence between êhi (X) and data point i. While it is possible that cross-fitting for ê
h is still beneficial in this case, for simplicity and for computational savings, we did not implement it.
a SIN Training. Input: Stage 1 data D1 := {(xi, yi)}mi=1, Stage 2 data D2 := {(xi, ti, yi)}ni=1 Step sizes λθ, λη, λψ, λφ. Number of update steps K. Mini-batch sizes B1, B2.
1: Initialize parameters: θ,η,ψ,φ 2: while not converged do . Stage 1 3: Sample mini-batch {(xb, yb)} mB1 b=1 4: Evaluate Jm (θ) 5: Update θ ← θ − λθ∇̂θJ(θ) 6: end while 7: while not converged do . Stage 2 8: Sample mini-batch {(xb, tb, yb)} nB2 b=1
9: Evaluate Jg,h (ψ,φ) , Jeh (η) 10: for k = 1 to K do 11: Update φ← φ− λφ∇̂φJg,h (ψ,φ) 12: Updateψ ← ψ−λψ∇̂ψJg,h (ψ,φ) 13: end for 14: Update η ← η − λη∇̂ηJeh (η) 15: end while
b Pseudocode in a PyTorch-like style. # Initialize submodels and optimizers m, e, g, h = MLP(...), MLP(...), MLP(...),
GNN(...) m_opt, e_opt, g_opt, h_opt = Adam(m.params(),
m_lr), Adam(e.params(), e_lr), ...
# Stage 1 for batch in train_loader:
X, Y = batch.X, batch.Y m_opt.zero_grad() F.mse_loss(m(X), Y).backward() m_opt.step()
# Stage 2 for batch in train_loader:
X, T, Y = batch.X, batch.T, batch.Y for _ in range(num_update_steps):
g_opt.zero_grad() h_opt.zero_grad() F.mse_loss((g(X)*(h(T) - e(X))).sum
(-1), (Y-m(X))).backward() g_opt.step() h_opt.step()
e_opt.zero_grad() F.mse_loss(e(X), h(T)).backward() e_opt.step()
Figure 2: The two-stage algorithm for training SIN.
We learn surrogate models for the mean outcome and propensity features m̂θ(X) and êhη(X) with parameters θ ∈ Rdθ ,η ∈ Rdη , as well as feature maps for covariates and treatments ĝψ(X), ĥφ(T), parameterized by ψ ∈ Rdψ ,φ ∈ Rdφ . We denote regularizers by Λ (·). Figure 2 summarizes the algorithm. As the mean outcome model m̂θ (X) does not depend on the other components, we learn it separately in Stage 1. In Stage 2, we alternate between learning ψ,φ,η.
Stage 1: Learn parameters θ of the mean outcome model m̂θ (X) based on the objective
Jm (θ) = 1
m m∑ i=1 ( yi − m̂θ (xi) )2 + Λ (θ) , (14)
which relies only on covariates and outcome data D1 := { (xi, yi) }m i=1 .
Stage 2: Learn parameters ψ,φ for the covariates and treatments feature maps ĝψ (X) , ĥφ (T), as well as parameters η for the propensity features êhη (X).
Jg,h (φ,ψ) = 1
n n∑ i=1
( yi − { m̂θ (xi) + ĝψ (xi) > ( ĥφ (ti)− êhη (xi) )})2 + Λ (ψ) + Λ (φ) .
(15)
This loss hinges on êhη (X), which needs to be learned by
Jeh (η) = n∑ i=1 ∥∥∥ĥφ (ti)− êhη (xi)∥∥∥2 2 + Λ (η) , (16)
note again the dependence on ĥφ (T). While it may be tempting to learn ψ,φ and η jointly, they have fundamentally different objectives (êhη (X) is defined as an estimate of the expectation E [ h (T) | x ] ). Therefore, we employ an alternating optimization procedure, where we take k ∈ {1, . . . ,K} optimization steps for ψ,φ towards Jg,h (ψ,φ) and one step for learning η. We observe that setting K > 1, i.e. updating ψ,φ more frequently than η, stabilizes the training process.

heading: 5.2 Advantages of SIN

text:
We conclude by describing the beneficial properties of SIN, particularly in finite-sample regimes:
1. Targeted regularization: Regularizing ĝψ (X) , ĥφ (T) in eq. (15) after partialing out confounding is a type of targeted regularization of the isolated causal effect. In contrast, outcome estimation methods can suffer from regularization-induced confounding, e.g., regularizing the effect estimate away from zero in the service of trying to improve predictive performance [29].
2. Propensity features: Learning propensity features can help us to (i) partial out parts of X that cause the treatment but not the outcome, and (ii) dispose unnecessary components of T.
3. Data-efficiency: In contrast to methods that split the data into disjoint models for each treatment group (known as T-learners for binary treatments [8, 10]), sharing causal effect parameters between all covariates regardless of their assigned treatment increases data-efficiency.
4. Partial data: In settings without access to both the treatment assignment and the outcome but only access to one of them, one can leverage that data to improve the (nuisance) estimator further, e.g., when a patient’s recovery is observed one year after a drug was administered [33].

heading: 6 Experiments

text: Here we evaluate how CATE estimation with our proposed model SIN compares with prior methods.

heading: 6.1 Experimental Setup

text:
Datasets. To be able to compute CATE estimation error w.r.t. a ground truth, we design two causal models: a simpler synthetic model with small-world graph treatments and a more complex model with real-world molecular graph treatments and gene expression covariates. The Small-World (SW) simulation contains 1,000 uniformly sampled covariates and 200 randomly generated Watts–Strogatz small-world graphs [61] as treatments. The Cancer Genomic Atlas (TCGA) simulation uses 9,659 gene expression measurements of cancer patients for covariates [62] and 10,000 sampled molecules from the QM9 dataset [46] as treatments. Appendix D details the data-generating schemes.
Baselines. We compare our method to (1) Zero, a sanity-check baseline that consistently predicts zero treatment effect and equals the mean squared treatment effect (poorly regularized models may perform worse than that due to confounding), (2) CAT, a categorical treatment variable model using one-hot encoded treatment indicator vectors, (3) GNN, a model that first encodes treatments with a GNN and then concatenates treatment and individual features for regression, (4) GraphITE [16], a CATE estimation method designed for graph treatments (more details in Section 2). GNN and CAT reflect the performance of standard regression models. The contrast between these two provides insight into whether the additional graph structure of the treatment improves CATE estimation. To deal with unseen treatments during CATE evaluation, we map such to the most similar ones seen during training based on their Euclidean distance in the embedding space of the GNN baseline.
Graph models. For small-world networks, we use k-dimensional GNNs [38], as to distinguish graphs they take higher-order structures into account. To model molecular graphs, we use Relational Graph Convolutional Networks [50], where the nodes are atoms and each edge type corresponds to a specific bond type. We use the implementations of PyTorch Geometric [11].
Evaluation metrics. We extend the expected Precision in Estimation of Heterogeneous Effect (PEHE) commonly used in binary treatment settings [19] to arbitrary pairs of treatments (t, t′) as follows. We denote the Unweighted PEHE (UPEHE) and the Weighted PEHE (WPEHE) as
UPEHE(WPEHE) , ∫ X ( τ̂ ( t′, t,x ) − τ ( t′, t,x ))2 p ( t | x ) p ( t′ | x ) p (x) dx, (17)
where the weighted version gives less importance to treatment pairs that are less likely; to account for the fact that such pairs will have higher estimation errors. In fact, as the reliability of estimated effects decreases by how likely they are in the observational study, we evaluate all methods on U/WPEHE truncated to the top K treatments, which we call U/WPEHE@K. To compute this, for each x, we rank all treatments by their propensity p ( t | x ) (given by the causal model) in descending order. We
take the top K treatments and compute the U/WPEHE for all ( K 2 ) treatment pairs.
In-sample vs. out-sample. A common benchmark for causal inference methods is the in-sample task, which we include here for completeness: estimating CATEs for covariate values x found in the training set. This task is still non-trivial, as the outcome of only one treatment is observed during training 4. In contrast, and arguably of more relevance to decision making, the goal of the out-sample task is to estimate CATEs for completely unseen covariate realizations x′.
Hyper-parameter tuning. To ensure a fair comparison, we perform hyper-parameter optimization with random search for all models on held-out data and select the best hyper-parameters over 10 runs.
Propensity. We define the propensity (or treatment selection bias) as p ( T | x ) =
softmax ( κW>X ) , where W ∈ R|T |×d,∀i, j : Wij ∼ U [0, 1] is a random matrix (sampled then fixed for each run). Recall |T | is the number of available treatments and let d be the dimensionality of the covariates. Here the bias strength κ is a temperature parameter that determines the flatness of the propensity (the lower the flatter, i.e., κ=0 corresponds to the uniform distribution).

heading: 6.2 Comparison of Performances on different K Treatments

text:
Figure 3 shows the UPEHE@K of all methods for K ∈ {2, . . . , 10}. We also report the WPEHE@6 of all methods in Table 1. Unless stated otherwise, we report results for bias strengths κ = 10 and κ = 0.1 in the SW and TCGA datasets, respectively across 10 random trials.
The results indicate that the relative performance of each method, for both the in-sample and outsample estimation tasks, is consistent. Further, they suggest that, overall, the performance of SIN is best due to a better isolation of the causal effect from the observed data compared to other methods. The performance difference between CAT and GNN across all results indicate that accounting for graph information significantly improves the estimates. We observe from the SW experiments that GraphITE [16] performs slightly better than GNN, while it is nearly the same as GNN on TCGA.
Surprisingly, the results of the TCGA experiments with low bias strength κ = 0.1 expose that all models but SIN fail to isolate causal effects better than the Zero baseline. These results confirm that confounding effects of X on Y combined with moderate causal effects can cause severe regularization bias for black-box regression models, while SIN partials these out from the outcome by m̂θ (X). We include additional results on convergence and larger values of K in Appendix E.1.
4The original motivation comes from Fisherian designs where the only source of randomness is on the treatment assignment [20]. Our motivation is simpler: rule out the extra variability from different covariates, highlighting the difference between methods due to different loss functions and less due to smoothing abilities.

heading: 6.3 Comparison of Robustness to different Bias Strengths κ

text: A strong selection bias (i.e. large κ) in the observed data makes CATE estimation more difficult, as it becomes unlikely to see certain treatments t ∈ T for particular covariates x. Here, we assess each model’s robustness to varying levels of selection bias, determined by κ, across 5 random seeds. In Figure 4, we see that SIN outperforms the baselines across the entire range of considered biases. Interestingly, SIN performs competitively even in a case with no selection bias (κ=0, which corresponds to a randomized experiment). Importantly, all performances seem to either stagnate (SW) or to increase (TCGA) with increasing biases. Notably, the poor performance of CAT suddenly improves on datasets with high bias. We believe this is because, in high bias regimes, we see fewer distinct treatments overall, which allows the CAT model to approach the performance of GNN.

heading: 7 Limitations, Future Work and Potential Negative Societal Impacts

text:
Limitations and future work. Firstly, in some real-life domains, Assumption 1 (Unconfoundedness) can be too strong, as there may exist hidden confounders. There are two common strategies to deal with them: utilizing instrumental variables [17, 58, 63] or proxy variables [35, 37, 59]. Developing new approaches for structured interventions in such settings is a promising future direction. Secondly, SIN is based on neural networks; however, neural network initialization can impact final estimates. To obtain consistency guarantees, GRD can be combined with kernel methods [35, 58].
Potential negative societal impacts. Because causal inference methods make recommendations about interventions to apply in real-world settings, misapplying them can have a negative real-world impact. It is crucial to thoroughly test these methods on realistic simulations and alter aspects of them to understand how violations of assumptions impact estimation. We have aimed to provide a comprehensive evaluation of structured treatment methods by showing how estimation degrades as less likely treatments are considered (Figure 3) and as treatment bias increases (Figure 4).

heading: 8 Conclusion

text: The main contributions of this paper are two-fold: (i) the generalized Robinson decomposition that yields a pseudo-outcome targeting the causal effect while possessing a quasi-oracle convergence guarantee under mild assumptions, and (ii) Structured Intervention Networks, a practical algorithm using representation learning that outperforms prior approaches in experiments with graph treatments.

heading: Acknowledgements

text: We thank Antonin Schrab, David Watson, Jakob Zeitler, Limor Gultchin, Marc Deisenroth and Shonosuke Harada for useful discussions and constructive feedback on the paper. JK and YZ acknowledge support by the Engineering and Physical Sciences Research Council with grant number EP/S021566/1. This work was partially supported by an ONR grant number N62909-19-1-2096 to RS. We thank the Alan Turing Institute for the provision of Azure cloud computing resources.

title: Causal Effect Inference for Structured Treatments