summaries: This paper aims to improve the density distribution of generative radiance field models by generating relitable reflectance fields and using different lightings to shade them during training. The core idea is that artifacts will arise when the generated density distribution ("shape") is unnatural and is getting shaded by different lightings, which should be detected and resolved by a discriminator. Experiments verify that this shading regularization can yield naturally looking shapes induced by the generated radiance fields.

summaries: This paper proposes a GAN capable of generating a relightable radiance field when trained on an unlabeled dataset of face images. The relightable radiance field generator is conditioned on the 3D location and a latent code and outputs a volume density field (à la NeRF) and an "albedo" field (I strongly advise against the use of this word; see below) that allows explicit control of the light location. The training is conducted in the GAN style, where the training signal comes from only a discriminator loss. The intuition is that the generator is better off generating a meaningful 3D volume that, when rendered from a random camera view, generates a photorealistic face image that falls into the training data distribution, than generating some intricately-designed volume that gives faces of different identities when viewed from different angles.  The authors demonstrate that the model is able to learn 3D shapes in an unsupervised fashion from just 2D images. The work is mainly compared against pi-GAN, which is shown to suffer from what the authors call "color-shape ambiguities." For efficient rendering, the authors also propose an auxiliary network that predicts the surface location given the latent code and a viewing direction, such that the network can sample around the predicted surface to avoid expensive sampling of unoccupied space.

summaries: This paper proposes a new method to train generative radiance fields that can achieve more accurate 3D shape reconstruction. The key idea is that instead of directly synthesizing the images from different viewpoints through volume ray tracing, which may suffer from color-shape ambiguity, it synthesizes albedo and normal and then renders the image with different lighting conditions. This new method requires the synthesized albedo and normal to be able to render realistic images under various lighting conditions, providing an extra regularization for shape reconstruction. To further improve the rendering speed, it uses an extra 2D CNN to predict a depth map from latent code and camera poses so that it only needs to sample near the surface in volume ray tracing. Experiments show that the proposed method greatly improves shape reconstruction with similar image synthesis quality compared to state-of-the-arts.

summaries: The paper proposes a novel generative model for neural radiance field, where the key idea is to synthetically add shading effects based on random lighting configurations. Shading will make images unrealistic when the geometry is corrupted. Therefore, the system gets trained to generate accurate shapes in its implicit radiance field representation. The system makes convincing results both qualitatively and quantitatively.

abstractText: The advancement of generative radiance fields has pushed the boundary of 3Daware image synthesis. Motivated by the observation that a 3D object should look realistic from multiple viewpoints, these methods introduce a multi-view constraint as regularization to learn valid 3D radiance fields from 2D images. Despite the progress, they often fall short of capturing accurate 3D shapes due to the shapecolor ambiguity, limiting their applicability in downstream tasks. In this work, we address this ambiguity by proposing a novel shading-guided generative implicit model that is able to learn a starkly improved shape representation. Our key insight is that an accurate 3D shape should also yield a realistic rendering under different lighting conditions. This multi-lighting constraint is realized by modeling illumination explicitly and performing shading with various lighting conditions. Gradients are derived by feeding the synthesized images to a discriminator. To compensate for the additional computational burden of calculating surface normals, we further devise an efficient volume rendering strategy via surface tracking, reducing the training and inference time by 24% and 48%, respectively. Our experiments on multiple datasets show that the proposed approach achieves photorealistic 3D-aware image synthesis while capturing accurate underlying 3D shapes. We demonstrate improved performance of our approach on 3D shape reconstruction against existing methods, and show its applicability on image relighting. Our code will be released at https://github.com/XingangPan/ShadeGAN.

authors:
T. Karras
S. Laine
T. Aila

title: A style-based generator architecture for generative adversarial networks

authors:
T. Karras
S. Laine
M. Aittala
J. Hellsten
J. Lehtinen
T. Aila

title: Analyzing and improving the image quality of stylegan

authors:
A. Brock
J. Donahue
K. Simonyan

title: Large scale gan training for high fidelity natural image synthesis

authors:
E.R. Chan
M. Monteiro
P. Kellnhofer
J. Wu
G. Wetzstein

title: pi-gan: Periodic implicit generative adversarial networks for 3d-aware image synthesis

authors:
K. Schwarz
Y. Liao
M. Niemeyer
A. Geiger

title: Graf: Generative radiance fields for 3d-aware image synthesis

venue: NeurIPS, 2020.

authors:
B. Mildenhall
P.P. Srinivasan
M. Tancik
J.T. Barron
R. Ramamoorthi
R. Ng

title: Nerf: Representing scenes as neural radiance fields for view synthesis

title: Photometric method for determining surface orientation from multiple images

venue: Optical engineering, vol. 19, no. 1, p. 191139, 1980.

authors:
L. Liu
J. Gu
K. Zaw Lin
T.-S. Chua
C. Theobalt

title: Neural sparse voxel fields

authors:
A. Yu
R. Li
M. Tancik
H. Li
R. Ng
A. Kanazawa

title: PlenOctrees for real-time rendering of neural radiance fields

authors:
C. Reiser
S. Peng
Y. Liao
A. Geiger

title: Kilonerf: Speeding up neural radiance fields with thousands of tiny mlps

venue: arXiv preprint arXiv:2103.13744, 2021.

authors:
S.J. Garbin
M. Kowalski
M. Johnson
J. Shotton
J. Valentin

title: Fastnerf: High-fidelity neural rendering at 200fps

venue: arXiv preprint arXiv:2103.10380, 2021.

authors:
T. Neff
P. Stadlbauer
M. Parger
A. Kurz
C.R.A. Chaitanya
A. Kaplanyan
M. Steinberger

title: Donerf: Towards real-time rendering of neural radiance fields using depth oracle networks

authors:
P. Paysan
R. Knothe
B. Amberg
S. Romdhani
T. Vetter

title: A 3d face model for pose and illumination invariant face recognition

venue: 2009 Sixth IEEE International Conference on Advanced Video and Signal Based Surveillance, Ieee, 2009.

authors:
M. Boss
R. Braun
V. Jampani
J.T. Barron
C. Liu
H. Lensch

title: Nerd: Neural reflectance decomposition from image collections

venue: arXiv preprint arXiv:2012.03918, 2020.

authors:
S. Bi
Z. Xu
P. Srinivasan
B. Mildenhall
K. Sunkavalli
M. Hašan
Y. Hold-Geoffroy
D. Kriegman
R. Ramamoorthi

title: Neural reflectance fields for appearance acquisition

venue: arXiv preprint arXiv:2008.03824, 2020.

authors:
P.P. Srinivasan
B. Deng
X. Zhang
M. Tancik
B. Mildenhall
J.T. Barron

title: Nerv: Neural reflectance and visibility fields for relighting and view synthesis

authors:
D.B. Lindell
J.N. Martel
G. Wetzstein

title: Autoint: Automatic integration for fast neural volume rendering

venue: arXiv preprint arXiv:2012.01714, 2020.

authors:
I. Goodfellow
J. Pouget-Abadie
M. Mirza
B. Xu
D. Warde-Farley
S. Ozair
A. Courville
Y. Bengio

title: Generative adversarial nets

authors:
J. Wu
C. Zhang
T. Xue
B. Freeman
J. Tenenbaum

title: Learning a probabilistic latent space of object shapes via 3d generative-adversarial modeling

authors:
H.A. Alhaija
S.K. Mustikovela
A. Geiger
C. Rother

title: Geometric image synthesis

venue: ACCV, Springer, 2018.

authors:
X. Chen
D. Cohen-Or
B. Chen
N.J. Mitra

title: Towards a neural graphics pipeline for controllable image generation

venue: Computer Graphics Forum, vol. 40, no. 2, 2021.

authors:
J.-Y. Zhu
Z. Zhang
C. Zhang
J. Wu
A. Torralba
J. Tenenbaum
B. Freeman

title: Visual object networks: Image generation with disentangled 3d representations

venue: NeurIPS, 2018.

authors:
T. Nguyen-Phuoc
C. Li
L. Theis
C. Richardt
Y.-L. Yang

title: Hologan: Unsupervised learning of 3d representations from natural images

authors:
T. Nguyen-Phuoc
C. Richardt
L. Mai
Y.-L. Yang
N. Mitra

title: Blockgan: Learning 3d object-aware scene representations from unlabelled images

venue: NeurIPS, Nov 2020.

authors:
Y. Liao
K. Schwarz
L. Mescheder
A. Geiger

title: Towards unsupervised learning of generative models for 3d controllable image synthesis

authors:
P. Henzler
N.J. Mitra
T. Ritschel

title: Escaping plato’s cave: 3d shape from adversarial rendering

authors:
A. Szabó
G. Meishvili
P. Favaro

title: Unsupervised generative 3d shape learning from natural images

venue: arXiv preprint arXiv:1910.00287, 2019.

authors:
M. Niemeyer
A. Geiger

title: Giraffe: Representing scenes as compositional generative neural feature fields

authors:
A. Tewari
M. Elgharib
G. Bharaj
F. Bernard
H.-P. Seidel
P. Pérez
M. Zollhofer
C. Theobalt

title: Stylerig: Rigging stylegan for 3d control over portrait images

authors:
T. Yenamandra
A. Tewari
F. Bernard
H.-P. Seidel
M. Elgharib
D. Cremers
C. Theobalt

title: i3dmm: Deep implicit 3d morphable model of human heads

authors:
S. Tulsiani
N. Kulkarni
A. Gupta

title: Implicit mesh reconstruction from unannotated image collections

venue: arXiv preprint arXiv:2007.08504, 2020.

authors:
S. Goel
A. Kanazawa
J. Malik

title: Shape and viewpoint without keypoints

authors:
L. Tran
X. Liu

title: Nonlinear 3d face morphable model

authors:
A. Kanazawa
S. Tulsiani
A.A. Efros
J. Malik

title: Learning category-specific mesh reconstruction from image collections

authors:
A. Kanazawa
M.J. Black
D.W. Jacobs
J. Malik

title: End-to-end recovery of human shape and pose

authors:
S. Sanyal
T. Bolkart
H. Feng
M.J. Black

title: Learning to regress 3d face shape and expression from an image without 3d supervision

authors:
J. Shang
T. Shen
S. Li
L. Zhou
M. Zhen
T. Fang
L. Quan

title: Self-supervised monocular 3d face reconstruction by occlusion-aware multi-view geometry consistency

authors:
M. Sahasrabudhe
Z. Shu
E. Bartrum
R. Alp Guler
D. Samaras
I. Kokkinos

title: Lifting autoencoders: Unsupervised learning of a fully-disentangled 3d morphable model using deep non-rigid structure from motion

venue: ICCV Workshops, 2019.

authors:
S. Wu
C. Rupprecht
A. Vedaldi

title: Unsupervised learning of probably symmetric deformable 3d objects from images in the wild

authors:
X. Li
S. Liu
K. Kim
S. De Mello
V. Jampani
M.-H. Yang
J. Kautz

title: Self-supervised single-view 3d reconstruction via semantic consistency

authors:
X. Pan
B. Dai
Z. Liu
C.C. Loy
P. Luo

title: Do 2d gans know 3d shape? unsupervised 3d shape reconstruction from 2d image gans

authors:
L. Mescheder
A. Geiger
S. Nowozin

title: Which training methods for gans do actually converge

authors:
Z. Liu
P. Luo
X. Wang
X. Tang

title: Deep learning face attributes in the wild

authors:
W. Zhang
J. Sun
X. Tang

title: Cat head detection-how to effectively exploit shape and texture features

authors:
V. Sitzmann
J. Martel
A. Bergman
D. Lindell
G. Wetzstein

title: Implicit neural representations with periodic activation functions

authors:
M. Heusel
H. Ramsauer
T. Unterthiner
B. Nessler
S. Hochreiter

title: Gans trained by a two time-scale update rule converge to a local nash equilibrium

venue: NeurIPS, 2017.

title: Illumination for computer generated pictures

venue: Communications of the ACM, vol. 18, no. 6, pp. 311–317, 1975.

authors:
E. Perez
F. Strub
H. De Vries
V. Dumoulin
A. Courville

title: Film: Visual reasoning with a general conditioning layer

authors:
R. Liu
J. Lehman
P. Molino
F.P. Such
E. Frank
A. Sergeev
J. Yosinski

title: An intriguing failing of convolutional neural networks and the coordconv solution

authors:
K. He
X. Zhang
S. Ren
J. Sun

authors:
Y. Wu
K. He

title: Group normalization

authors:
S. Ioffe
C. Szegedy

title: Batch normalization: Accelerating deep network training by reducing internal covariate shift

authors:
D.P. Kingma
J. Ba

title: Adam: A method for stochastic optimization

venue: arXiv preprint arXiv:1412.6980, 2014.

authors:
Y. Choi
Y. Uh
J. Yoo
J.-W. Ha

title: Stargan v2: Diverse image synthesis for multiple domains

heading: 1 Introduction

text:
Advanced deep generative models, e.g., StyleGAN [1, 2] and BigGAN [3], have achieved great successes in natural image synthesis. While producing impressive results, these 2D representationbased models cannot synthesize novel views of an instance in a 3D-consistent manner. They also fall short of representing an explicit 3D object shape. To overcome such limitations, researchers have proposed new deep generative models that represent 3D scenes as neural radiance fields [4, 5]. Such 3D-aware generative models allow explicit control of viewpoint while preserving 3D consistency during image synthesis. Perhaps a more fascinating merit is that they have shown the great potential of learning 3D shapes in an unsupervised manner from just a collection of unconstrained 2D images. If we could train a 3D-aware generative model that learns accurate 3D object shapes, it would broaden various downstream applications such as 3D shape reconstruction and image relighting.
Existing attempts for 3D-aware image synthesis [4, 5] tend to learn coarse 3D shapes that are inaccurate and noisy, as shown in Fig.1 (a). We found that such inaccuracy arises from an inevitable ambiguity inherent in the training strategy adopted by these methods. In particular, a form of
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
regularization, which we refer to as "multi-view constraint", is used to enforce the 3D representation to look realistic from different viewpoints. The constraint is commonly implemented by first projecting the generator’s outputs (e.g., radiance fields [6]) to randomly sampled viewpoints, and then feeding them to a discriminator as fake images for training. While such a constraint enables these models to synthesize images in a 3D-aware manner, it suffers from the shape-color ambiguity, i.e., small variations of shape could lead to similar RGB images that look equally plausible to the discriminator, as the color of many objects is locally smooth. Consequently, inaccurate shapes are concealed under this constraint.
In this work, we propose a novel shading-guided generative implicit model (ShadeGAN) to address the aforementioned ambiguity. In particular, ShadeGAN learns more accurate 3D shapes by explicitly modeling shading, i.e., the interaction of illumination and shape. We believe that an accurate 3D shape should look realistic not only from different viewpoints, but also under different lighting conditions, i.e., satisfying the "multi-lighting constraint". This idea shares similar intuition with photometric stereo [7], which shows that accurate surface normal could be recovered from images taken under different lighting conditions. Note that the multi-lighting constraint is feasible as real-world images used for training are often taken under various lighting conditions. To fulfill this constraint, ShadeGAN takes a relightable color field as the intermediate representation, which approximates the albedo but does not necessarily satisfy viewpoint independence. The color field is shaded under a randomly sampled lighting condition during rendering. Since image appearance via such a shading process is strongly dependent on surface normals, inaccurate 3D shape representations will be much more clearly revealed than in earlier shading-agnostic generative models. Hence, by satisfying the multi-lighting constraint, ShadeGAN is encouraged to infer more accurate 3D shapes as shown in Fig.1 (b).
The above shading process requires the calculation of the normal direction via back-propagation through the generator, and such calculation needs to be repeated dozens of times for a pixel in volume rendering [4, 5], introducing additional computational overhead. Existing efficient volume rendering techniques [8, 9, 10, 11, 12] mainly target static scenes, and could not be directly applied to generative models due to their dynamic nature. Therefore, to improve the rendering speed of ShadeGAN, we formulate an efficient surface tracking network to estimate the rendered object surface conditioned on the latent code. This enables us to save rendering computations by just querying points near the predicted surface, leading to 24% and 48% reduction of training and inference time without affecting the quality of rendered images.
Comprehensive experiments are conducted across multiple datasets to verify the effectiveness of ShadeGAN. The results show that our approach is capable of synthesizing photorealistic images while capturing more accurate underlying 3D shapes than previous generative methods. The learned
distribution of 3D shapes enables various downstream tasks like 3D shape reconstruction, where our approach significantly outperforms other baselines on the BFM dataset [13]. Besides, modeling the shading process enables explicit control over lighting conditions, achieving image relighting effect. Our contributions can be summarized as follows: 1) We address the shape-color ambiguity in existing 3D-aware image synthesis methods with a shading-guided generative model that satisfies the proposed multi-lighting constraint. In this way, ShadeGAN is able to learn more accurate 3D shapes for better image synthesis. 2) We devise an efficient rendering technique via surface tracking, which significantly saves training and inference time for volume rendering-based generative models. 3) We show that ShadeGAN learns to disentangle shading and color that well approximates the albedo, achieving natural relighting effects in image synthesis.

heading: 2 Related Work

text:
Neural volume rendering. Starting from the seminal work of neural radiance fields (NeRF) [6], neural volume rendering has gained much popularity in representing 3D scenes and synthesizing novel views. By integrating coordinate-based neural networks with volume rendering, NeRF performs high-fidelity view synthesis in a 3D consistent manner. Several attempts have been proposed to extend or improve NeRF. For instance, [14, 15, 16] further model illumination, and learn to disentangle reflectance with shading given well-aligned multi-view and multi-lighting images. Besides, many studies accelerate the rendering of static scenes from the perspective of spatial sparsity [8, 9], architectural design [10, 11], or efficient rendering [17, 12]. However, it is not trivial to apply these illumination and acceleration techniques to volume rendering-based generative models [5, 4], as they typically learn from unposed and unpaired images, and represent dynamic scenes that change with respect to the input latent codes.
In this work, we take the first attempt to model illumination in volume rendering-based generative models, which serves as a regularization for accurate 3D shape learning. We further devise an efficient rendering technique for our approach, which shares similar insight with [12], but does not rely on ground truth depth for training and it is not limited to a small viewpoint range.
Generative 3D-aware image synthesis. Generative adversarial networks (GANs) [18] are capable of generating photorealistic images of high-resolution, but lack explicit control over camera viewpoint. In order to enable them to synthesis images in a 3D-aware manner, many recent approaches investigate how 3D representations could be incorporated into GANs [19, 20, 21, 22, 23, 24, 25, 26, 27, 5, 4, 28, 29, 30]. While some works directly learn from 3D data [19, 20, 21, 22, 30], in this work we focus on approaches that only have access to unconstrained 2D images, which is a more practical setting. Several attempts [23, 24, 25] adopt 3D voxel features with learned neural rendering. These methods produce realistic 3D-aware synthesis, but the 3D voxels are not interpretable, i.e., they cannot be transferred to 3D shapes. By leveraging differentiable renderer, [26] and [27] learn interpretable 3D voxels and meshes respectively, but [26] suffers from limited visual quality due to low voxel resolution while the learned 3D shapes of [27] exhibit noticeable distortions. The success of NeRF has motivated researchers to use radiance fields as the intermediate 3D representation in GANs [5, 4, 28]. While achieving impressive 3D-aware image synthesis with multi-view consistency, the extracted 3D shapes of these approaches are often imprecise and noisy. Our main goal in this work is to address the inaccurate shape by explicitly modeling illumination in the rendering process. This innovation helps achieve better 3D-aware image synthesis with broader applications.
Unsupervised 3D shape learning from 2D images. Our work is also related to unsupervised approaches that learn 3D object shapes from unconstrained, monocular view 2D images. While several approaches use external 3D shape templates or 2D key-points as weak supervisions to facilitate learning [31, 32, 33, 34, 35, 36, 37], in this work we consider the harder setting where only 2D images are available. To tackle this problem, most approaches adopt an “analysis-by-synthesis” paradigm [38, 39, 40]. Specifically, they design photo-geometric autoencoders to infer the 3D shape and viewpoint of each image with a reconstruction loss. While succeed in learning the 3D shapes for some object categories, these approaches typically rely on certain regularization to prevent trivial solutions, like the commonly used symmetry assumption on object shapes [39, 40, 31, 32]. Such assumption tends to produce symmetric results that may overlook the asymmetric aspects of objects. Recently, GAN2Shape [41] shows that it is possible to recover 3D shapes for images generated by 2D GANs. This method, however, requires inefficient instance-specific training, and recovers depth maps instead of full 3D representations.
The proposed 3D-aware generative model also serves as a powerful approach for unsupervised 3D shape learning. Compared with aforementioned autoencoder-based methods, our GAN-based approach avoids the need to infer the viewpoint of each image, and does not rely on strong regularizations. In experiments, we demonstrate superior performance over recent state-of-the-art approaches Unsup3d [39] and GAN2Shape [41].

heading: 3 Methodology

text: We consider the problem of 3D-aware image synthesis by learning from a collection of unconstrained and unlabeled 2D images. We argue that modeling shading, i.e., the interaction of illumination and shape, in a generative implicit model enables unsupervised learning of more accurate 3D object shapes. In the following, we first provide some preliminaries on neural radiance fields (NeRF) [6], and then introduce our shading-guided generative implicit model.

heading: 3.1 Preliminaries on Neural Radiance Fields

text:
As a deep implicit model, NeRF [6] uses an MLP network to represent a 3D scene as a radiance field. The MLP fθ : (x,d)→ (σ, c) takes a 3D coordinate x ∈ R3 and a viewing direction d ∈ S2 as inputs, and outputs a volume density σ ∈ R+ and a color c ∈ R3. To render an image under a given camera pose, each pixel color C of the image is obtained via volume rendering along its corresponding camera ray r(t) = o+ td with near and far bounds tn and tf as below:
C(r) = ∫ tf tn T (t)σ(r(t))c(r(t),d)dt, where T (t) = exp(− ∫ t tn σ(r(s))ds). (1)
In practice, this volume rendering is implemented with a discretized form using stratified and hierarchical sampling. As this rendering process is differentiable, NeRF could be directly optimized via posed images of a static scene. After training, NeRF allows the rendering of images under new camera poses, achieving high-quality novel view synthesis.

heading: 3.2 Shading-Guided Generative Implicit Model

text:
In this work, we are interested in developing a generative implicit model that explicitly models the shading process for 3D-aware image synthesis. To achieve this, we make two extensions to the MLP network in NeRF. First, similar to most deep generative models, it is further conditioned on a latent code z sampled from a prior distribution N (0, I)d. Second, instead of directly outputting the color c, it outputs a relightable pre-cosine color term a ∈ R3, which is conceptually similar to albedo in the way that it could be shaded under a given lighting condition. While albedo is viewpoint-independent, in this work we do not strictly enforce such independence for a in order to account for dataset bias. Thus, our generator gθ : (x,d, z)→ (σ,a) takes a coordinate x, a viewing direction d, and a latent
code z as inputs, and outputs a volume density σ and a pre-cosine color a. Note that here σ is independent of d, while the dependence of a on d is optional. To obtain the color C of a camera ray r(t) = o+ td with near and far bounds tn and tf , we calculate the final pre-cosine colorA via:
A(r, z) = ∫ tf tn T (t, z)σ(r(t), z)a(r(t),d, z)dt, where T (t, z) = exp(− ∫ t tn σ(r(s), z)ds).
(2)
We also calculate the normal direction n with: n(r, z) = n̂(r, z)/‖n̂(r, z)‖2, where n̂(r, z) = − ∫ tf tn T (t, z)σ(r(t), z)∇r(t)σ(r(t), z)dt,
(3)
where ∇r(t)σ(r(t), z) is the derivative of volume density σ with respect to its input coordinate, which naturally captures the local normal direction, and could be calculated via back-propagation. Then the final color C is obtained via Lambertian shading as:
C(r, z) = A(r, z)(ka + kdmax(0, l · n(r, z))), (4)
where l ∈ S2 is the lighting direction, ka and kd are the ambient and diffuse coefficients. We provide more discussions on this shading formulation at the end of this subsection.
Camera and Lighting Sampling. Eq.(2 - 4) describe the process of rendering a pixel color given a camera ray r(t) and a lighting condition µ = (l, ka, kd). Generating a full image Ig ∈ R3×H×W requires one to sample a camera pose ξ and a lighting condition µ in addition to the latent code z, i.e., Ig = Gθ(z, ξ,µ). In our setting, the camera pose ξ could be described by pitch and yaw angles, and is sampled from a prior Gaussian or uniform distribution pξ, as also done in previous works [4, 5]. Sampling the camera pose randomly during training would motivate the learned 3D scene to look realistic from different viewpoints. While this multi-view constraint is beneficial for learning a valid 3D representation, it is often insufficient to infer the accurate 3D object shape. Thus, in our approach, we further introduce a multi-lighting constraint by also randomly sampling a lighting condition µ from a prior distribution pµ. In practice, pµ could be estimated from the dataset using existing approaches like [39]. We also show in our experiments that a simple and manually tuned prior distribution could also produce reasonable results. As the shading process is sensitive to the normal direction due to the diffuse term kdmax(0, l · n(r, z)) in Eq.(4), this multi-lighting constraint would regularize the model to learn more accurate 3D shapes that produce natural shading, as shown in Fig.1 (b).
Training. Our generative model follows the paradigm of GANs [18], where the generator is trained together with a discriminator D with parameters φ in an adversarial manner. During training, the generator generates fake images Ig = Gθ(z, ξ,µ) by sampling the latent code z, camera pose ξ and lighting condition µ from their corresponding prior distributions pz , pξ, and pµ. Let I denotes real images sampled from the data distribution pI . We train our model with a non-saturating GAN loss with R1 regularization [42]:
L(θ, φ) = Ez∼pz,ξ∼pξ,µ∼pµ [ f ( Dφ(Gθ(z, ξ,µ)) )] + EI∼pD [ f(−Dφ(I)) + λ‖∇Dφ(I)‖2 ] ,
(5) where f(u) = − log(1 + exp(−u)), and λ controls the strength of regularization. More implementation details are provided in the supplementary material.
Discussion. Note that in Eq.(2 - 4), we perform shading after A and n are obtained via volume rendering. An alternative way is to perform shading at each local spatial point as c(r(t),d, z) = a(r(t),d, z)(ka + kdmax(0, l · n(r(t), z))), where n(r(t), z) = −∇r(t)σ(r(t), z)/‖∇r(t)σ(r(t), z)‖2 is the local normal. Then we could perform volume rendering using c(r(t), z) to get the final pixel color. In practice, we observe that this formulation obtains suboptimal results. An intuitive reason is that in this formulation, the normal direction is normalized at each local point, neglecting the magnitude of∇r(t)σ(r(t), z), which tends to be larger near the object surfaces. We provide more analysis in experiments and the supplementary material.
The Lambertian shading we used is an approximation to the real illumination scenario. While serving as a good regularization for improving the learned 3D shape, it could possibly introduce an additional gap between the distribution of generated images and that of real images. To compensate
for such risk, we could optionally let the predicted a be conditioned on the lighting condition, i.e., a = a(r(t),d,µ, z). Thus, in cases where the lighting condition deviates from the real data distribution, the generator could learn to adjust the value of a and reduce the aforementioned gap. We show the benefit of this design in the experiments.

heading: 3.3 Efficient Volume Rendering via Surface Tracking

text:
Similar to NeRF, we implement volume rendering with a discretized integral, which typically requires to sample dozens of points along a camera ray, as shown in Fig. 3 (a). In our approach, we also need to perform back-propagation across the generator in Eq.(3) to get the normal direction for each point, which introduces additional computational cost. To achieve more efficient volume rendering, a natural idea is to exploit spatial sparsity. Usually, the weight T (t, z)σ(r(t), z) in volume rendering would concentrate on the object surface position during training. Thus, if we know the rough surface position before rendering, we could sample points near the surface to save computation. While for a static scene it is possible to store such spatial sparsity in a sparse voxel grid [8, 9], this technique cannot be directly applied to our generative model, as the 3D scene keeps changing with respect to the input latent code.
To achieve more efficient volume rendering in our generative implicit model, we further propose a surface tracking network S that learns to mimic the surface position conditioned on the latent code. In particular, the volume rendering naturally allows the depth estimation of the object surface via:
ts(r, z) = ∫ tf tn T (t, z)σ(r(t), z)t dt, (6)
where T (t, z) is defined the same way as in Eq.(2). Thus, given a camera pose ξ and a latent code z, we could render the full depth map ts(z, ξ). As shown in Fig. 3 (b), we mimic ts(z, ξ) with the surface tracking network Sψ , which is a light-weighted convolutional neural network that takes z, ξ as inputs and outputs a depth map. The depth mimic loss is:
L(ψ) = Ez∼pz,ξ∼pξ [‖Sψ(z, ξ)− ts(z, ξ)‖1 + Prec(Sψ(z, ξ), ts(d(z, ξ))] , (7) where Prec is the perceptual loss that motivates Sψ to better capture edges of the surface.
During training, Sψ is optimized jointly with the generator and the discriminator. Thus, each time after we sample a latent code z and a camera pose ξ, we can get an initial guess of the depth map as Sψ(z, ξ). Then for a pixel with predicted depth s, we could perform volume rendering in Eq.(2,3,6) with near bound tn = s−∆i/2 and far bound tf = s+ ∆i/2, where ∆i is the interval for volume rendering that decreases as the training iteration i grows. Specifically, we start with a large interval ∆max and decrease to ∆min with an exponential schedule. As ∆i decreases, the number of points used for rendering m also decreases accordingly. Note that the computational cost of our efficient surface tracking network is marginal compared to the generator, as the former only needs a single forward pass to render an image while the latter will be queried for H ×W ×m times. Thus, the reduction of m would significantly accelerate the training and inference speed for ShadeGAN.

heading: 4 Experiments

text:
In this section, we evaluate the proposed ShadeGAN on 3D-aware image synthesis. We also show that ShadeGAN learns much more accurate 3D shapes than previous methods, and in the meantime allows explicit control over lighting conditions. The datasets used include CelebA [43], BFM [13], and Cats [44], all of which contain only unconstrained 2D RGB images.
Implementation. In terms of model architectures, we adopt a SIREN-based MLP [45] as the generator and a convolutional neural network as the discriminator following [4]. For the prior distribution of lighting conditions, we use Unsup3d [39] to estimate the lighting conditions of real data and subsequently fit a multivariate Gaussian distribution of µ = (l, ka, kd) as the prior. A hand-crafted prior distribution is also included in the ablation study. In quantitative study, we let the pre-cosine color a be conditioned on the lighting condition µ as well as the viewing direction d unless otherwise stated. In qualitative study, we observe that removing view conditioning achieves slightly better 3D shapes for CelebA and BFM datasets. Thus, we show results without view conditioning for these two datasets in the main paper, and put those with view conditioning in Fig. 4 of the supplementary material. Other implementation details are also provided in the supplementary.
Figure 5: Generated face images and their 3D meshes.
Normal AlbedoMeshImage
(a) ShadeGAN
(b) Local normal
(c) Manual prior
Figure 6: Qualitative ablation. See the main text for discussions.
Comparison with baselines. We compare ShadeGAN with two state-of-the-art generative implicit models, namely GRAF [5] and pi-GAN [4]. Specifically, Fig. 4 includes both synthesized images as well as their corresponding 3D meshes, which are obtained by performing marching cubes on the volume density σ. While GRAF and pi-GAN could synthesize images with controllable poses, their learned 3D shapes are inaccurate and noisy. In contrast, our approach not only synthesizes photorealistic 3D-consistent images, but also learns much more accurate 3D shapes and surface normals, indicating the effectiveness of the proposed multi-lighting constraint as a regularization. More synthesized images and their corresponding shapes are included in Fig.5. Besides more accurate 3D shapes, ShadeGAN can also learn the albedo and diffuse shading components inherently. As shown in Fig. 4, although not perfect, ShadeGAN has managed to disentangle shading and albedo with satisfying quality, as such disentanglement is a natural solution to the multi-lighting constraint.
The quality of learned 3D shapes is quantitatively evaluated on the BFM dataset. Specifically, we use each of the generative implicit models to generate 50k images and their corresponding depth maps. Image-depth pairs from each model are used as training data to train an additional convolutional neural network (CNN) that learns to predict the depth map of an input image. We then test each trained CNN on the BFM test set and compare its predictions to the ground-truth depth maps as a measurement of the quality of learned 3D shapes. Following [39], we report the scale-invariant depth error (SIDE) and mean angle deviation (MAD) metrics. The results are included in Tab. 1, where ShadeGAN significantly outperforms GRAF and pi-GAN. Besides, ShadeGAN also outperforms other advanced unsupervised 3D shape learning approaches including Unsup3d [39] and GAN2Shape [41], demonstrating its large potential in unsupervised 3D shapes learning. In terms of image quality, Tab. 1 includes the FID [46] scores of images synthesized by different models, where the FID score of ShadeGAN is slightly inferior to pi-GAN in BFM and CelebA. Intuitively, this is caused by the gap between our approximated shading (i.e. Lambertian shading) and the real illumination, which can be potentially avoided by adopting more realistic shading models and improving the lighting prior.
In Tab. 2, we also show the quantitative results of different models on CelebA and Cats. To evaluate the learned shape, we use each generative implicit model to generate 2k front-view images and their corresponding depth maps. While these datasets do not have ground truth depth, we report MAD obtained by testing pretrained Unsup3d models [39] on these generated image-depth pairs as
Table 3: Ablation study on the BFM dataset.
No. Method FID ↓ SIDE ↓ MAD ↓ (1) ShadeGAN 17.7 0.607 14.52 (2) local shading 30.1 0.754 18.18 (3) w/o light 19.2 0.618 14.53 (4) w/o view 18.6 0.622 14.88 (5) manual prior 20.2 0.643 15.38 (6) +efficient 18.2 0.673 14.72
Table 4: Training and inference time cost on CelebA. The efficient volume rendering significantly improves training and inference speed.
Method Train (h) Inference (s) FID ShadeGAN 92.3 0.343 16.4 +efficient 70.2 0.179 16.2 pi-GAN 56.8 0.204 15.7 +efficient 46.9 0.114 15.9
A lb
ed o
(a) w/o light condition (b) with light condition
S h
ad in
g Im
ag e
(c) with specular
Figure 8: Illumination-aware image synthesis. ShadeGAN allows explicit control over the lighting. The pre-cosine color (albedo) is independent of lighting in (a) and is conditioned on lighting in (b). We show results of adding a specular term in (c).
a reference. As we can observe, results on CelebA and Cat are consistent with those on the BFM dataset.
Ablation studies. We further study the effects of several design choices in ShadeGAN. First, we perform local points-specific shading as mentioned in the discussion of Sec. 3.2. As Tab. 3 No.(2) and Fig. 6 (b) show, the results of such a local shading strategy are notably worse than the original one, which indicates that taking the magnitude of∇xσ into account is beneficial. Besides, the results of Tab. 3 No.(3) and No.(4) imply that removing a’s dependence on the lighting µ or the viewpoint d could lead to a slight performance drop. The results of using a simple manually tuned lighting prior are provided in Tab. 3 No.(5) and Fig. 6 (c), which are only moderately worse than the results of using a data-driven prior, and the generated shapes are still significantly better than the ones produced by existing approaches.
To verify the effectiveness of the proposed efficient volume rendering technique, we include its effects on image quality and training/inference time in Tab. 3 No.(6) and Tab. 4. It is observed that the efficient volume rendering has marginal effects on the performance, but significantly reduces the training and inference time by 24% and 48% for ShadeGAN. Moreover, in Fig. 7 we visualize the depth maps predicted by our surface tracking network and those obtained via volume rendering. It is shown that under varying identities and camera poses, the surface tracking network could consistently predict depth values that are quite close to the real surface positions, so that we can sample points near the predicted surface for rendering without sacrificing image quality.
Illumination-aware image synthesis. As ShadeGAN models the shading process, it by design allows explicit control over the lighting condition. We provide such illumination-aware image synthesis results in Fig.8, where ShadeGAN generates promising images under different lighting directions. We also show that in cases where the predicted a is conditioned on the lighting condition µ, a would slightly change w.r.t. the lighting condition, e.g., it would be brighter in areas having a overly dim shading in order to make the final image more natural. Besides, we could optionally add a specular term ksmax(0,h · n)p in Eq. 4 (i.e., Blinn-Phong shading [47], where h is the bisector of the angle between the viewpoint and the lighting direction) to create specular highlight effects, as shown in Fig.8 (c).
GAN inversion. ShadeGAN could also be used to reconstruct a given target image by performing GAN inversion. As shown in Fig. 9 such inversion allows us to obtain several factors of the image, including the 3D shape, surface normal, approximated albedo, and shading. Besides, we can further perform view synthesis and relighting by changing the viewpoint and lighting condition. The implementation of GAN inversion is provided in the supplementary material.
Discussions. As the Lambertian shading we used is an approximation to the real illumination, the albedo learned by ShadeGAN is not perfectly disentangled. Our approach does not consider the spatially-varying material properties of objects as well. In the future, we intend to incorporate more sophisticated shading models to learn better disentangled generative reflectance fields.

heading: 5 Conclusion

text:
In this work, we present ShadeGAN, a new generative implicit model for shape-accurate 3D-aware image synthesis. We have shown that the multi-lighting constraint, achieved in ShadeGAN by explicit illumination modeling, significantly helps learning accurate 3D shapes from 2D images. ShadeGAN also allows us to control the lighting condition during image synthesis, achieving natural image relighting effects. To reduce the computational cost, we have further devised a light-weighted surface tracking network, which enables an efficient volume rendering technique for generative implicit models, achieving significant acceleration on both training and inference speed. A generative model with shape-accurate 3D representation could broaden its applications in vision and graphics, and our work has taken a solid step towards this goal.
Acknowledgment. We would like to thank Eric R. Chan for sharing the codebase of pi-GAN. This study is supported under the ERC Consolidator Grant 4DRepLy (770784). This study is also supported under the RIE2020 Industry Alignment Fund – Industry Collaboration Projects (IAF-ICP) Funding Initiative, as well as cash and in-kind contribution from the industry partner(s).

title: A Shading-Guided Generative Implicit Model for Shape-Accurate 3D-Aware Image Synthesis


summaries: The paper concerns the uncertain quantification for instrumental variable (IV) regression that it is an approach for estimating causal effect from confounded observational data and, proposes a scalable quasi-Bayesian procedure. The procedure was built upon the kernelized IV models (Muandet et al (2020)). The contribution is forming a quasi-posterior by employing a Gaussian process prior and constructing a kernel conditional expectation estimator using a randomized prior trick. Theoretical properties and simulation studies are provided.

summaries: The paper studies a quasi-Bayesian approach to instrumental variable (IV) regression, based on using an empirical estimate of the ‘likelihood’ to obtain a Gibbs posterior. Such an approach has been studied in earlier literature due to the difficulty of incorporating moment conditions in IV via a likelihood. The paper pursues a Gaussian process approach and proves some theory, namely that the resulting posterior puts most of its mass on functions (almost) satisfying the moment constraint. The paper provides a computation approach based on the randomized prior trick and illustrates the method on various simulated datasets.

summaries: The paper proposes a scalable quasi-Bayesian inference for instrumental variable regression. Specifically, the paper builds a quasi-posterior on the recent quadratic kernel IV loss function. For scalability, the paper extends the randomized prior trick based on the recent dual formulation of the quadratic loss. The approach has a theoretical guarantee that incorrect solutions are excluded while valid kept. Moreover, the convergence analysis for the randomized prior trick with the stochastic gradient descent-ascent is present. Experiments validate the approach.

summaries: The authors introduce a scalable quasi-Bayesian method for nonparametric instrumental variable (IV) regression that provides uncertainty quantification for the IV models. They base their approach on kernelized IV models and define a quasi-posterior as the Radon-Nikodym derivative of a generalized method of moments (GMM) objective function with respect to the standard Gaussian process prior. They go on to show that the quasi-posterior exhibits satisfies certain posterior contraction criteria that lead to asymptotically sound results, and prove a conservative posterior contraction rate under a set of typical assumptions. They then show how to achieve scalable inference for the quasi-posterior using a previously proposed “randomized prior trick” for Gaussian process regression, and provide a proof for the convergence of the inference algorithm. Finally, the authors showcase the performance of their method in a series of simulated experiments of varying complexity, by comparing it against other approaches that provide uncertainty quantification (Bayesian, bootstrap).

abstractText: Recent years have witnessed an upsurge of interest in employing flexible machine learning models for instrumental variable (IV) regression, but the development of uncertainty quantification methodology is still lacking. In this work we present a scalable quasi-Bayesian procedure for IV regression, building upon the recently developed kernelized IV models. Contrary to Bayesian modeling for IV, our approach does not require additional assumptions on the data generating process, and leads to a scalable approximate inference algorithm with time cost comparable to the corresponding point estimation methods. Our algorithm can be further extended to work with neural network models. We analyze the theoretical properties of the proposed quasi-posterior, and demonstrate through empirical evaluation the competitive performance of our method.

title: Pischke, Mostly harmless econometrics: An empiricist’s companion

venue: Princeton university press,

title: An introduction to instrumental variables for epidemiologists

venue: International Journal of Epidemiology, vol. 29, no. 4, pp. 722–729, 2000.

authors:
J. Cuzick
R. Edwards
N. Segnan

title: Adjusting for non-compliance and contamination in randomized clinical trials

venue: Statistics in Medicine, vol. 16, no. 9, pp. 1017–1029, 1997. 10

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

title: Deep IV: A flexible approach for counterfactual prediction

venue: International Conference on Machine Learning. PMLR, 2017, pp. 1414–1423.

authors:
G. Lewis
V. Syrgkanis

title: Adversarial generalized method of moments

venue: arXiv preprint arXiv:1803.07164, 2018.

authors:
A. Bennett
N. Kallus
T. Schnabel

title: Deep generalized method of moments for instrumental variable analysis

venue: arXiv preprint arXiv:1905.12495, 2019.

authors:
R. Singh
M. Sahani
A. Gretton

title: Kernel Instrumental Variable Regression

authors:
A. Bennett
N. Kallus

title: The Variational Method of Moments

authors:
J.H. Stock
J.H. Wright
M. Yogo

title: A survey of weak instruments and weak identification in generalized method of moments

venue: Journal of Business & Economic Statistics, vol. 20, no. 4, pp. 518–529, 2002.

title: Nonparametric instrumental variables estimation

venue: American Economic Review, vol. 103, no. 3, pp. 550–56, 2013.

authors:
F. Kleibergen
H.K. Van Dijk

title: Bayesian simultaneous equations analysis using reduced rank structures

venue: Econometric theory, pp. 701–743, 1998.

authors:
F. Kleibergen
E. Zivot

title: Bayesian and classical approaches to instrumental variable regression

venue: Journal of Econometrics, vol. 114, no. 1, pp. 29–72, 2003.

authors:
T.G. Conley
C.B. Hansen
R.E. McCulloch
P.E. Rossi

title: A semi-parametric Bayesian approach to the instrumental variable problem

venue: Journal of Econometrics, vol. 144, no. 1, pp. 276–305, 2008.

authors:
H.F. Lopes
N.G. Polson

title: Bayesian instrumental variables: priors and likelihoods

venue: Econometric Reviews, vol. 33, no. 1-4, pp. 100–121, 2014.

authors:
M. Wiesenfarth
C.M. Hisgen
T. Kneib
C. Cadarso-Suarez

title: Bayesian nonparametric instrumental variables regression based on penalized splines and dirichlet process mixtures

venue: Journal of Business & Economic Statistics, vol. 32, no. 3, pp. 468–482, 2014.

authors:
V. Chernozhukov
H. Hong

title: An MCMC approach to classical estimation

venue: Journal of Econometrics, vol. 115, no. 2, pp. 293–346, Aug. 2003.

authors:
Y. Liao
W. Jiang

title: Posterior consistency of nonparametric conditional moment restricted models

venue: The Annals of Statistics, vol. 39, no. 6, pp. 3003–3031, Dec. 2011.

title: Quasi-Bayesian analysis of nonparametric instrumental variables models

venue: The Annals of Statistics, vol. 41, no. 5, pp. 2359–2390, Oct. 2013.

authors:
J.-P. Florens
A. Simoni

title: Nonparametric estimation of an instrumental regression: A quasi- Bayesian approach based on regularized posterior

venue: Journal of Econometrics, vol. 170, no. 2, pp. 458–475, Oct. 2012.

authors:
K. Muandet
A. Mehrjou
S.K. Lee
A. Raj

title: Dual Instrumental Variable Regression

authors:
L. Liao
Y.-L. Chen
Z. Yang
B. Dai
Z. Wang
M. Kolar

title: Provably efficient neural estimation of structural equation model: An adversarial approach

venue: arXiv preprint arXiv:2007.01290, 2020.

authors:
I. Osband
J. Aslanides
A. Cassirer

title: Randomized prior functions for deep reinforcement learning

venue: Proceedings of the 32nd International Conference on Neural Information Processing Systems, 2018, pp. 8626–8638. 11

authors:
W.K. Newey
J.L. Powell

title: Instrumental variable estimation of nonparametric models

authors:
K. Fukumizu
F.R. Bach
A. Gretton

title: Statistical Consistency of Kernel Canonical Correlation Analysis

venue: Journal of Machine Learning Research, p. 23, 2004.

authors:
B. Dai
N. He
Y. Pan
B. Boots
L. Song

title: Learning from Conditional Distributions via Dual Embeddings

authors:
N. Dikkala
G. Lewis
L. Mackey
V. Syrgkanis

title: Minimax estimation of conditional moment models

venue: Advances in Neural Information Processing Systems, H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, Eds., vol. 33. Curran Associates, Inc., 2020, pp. 12 248–12 262.

title: Inverse problems: a Bayesian perspective

venue: Acta numerica, vol. 19, pp. 451– 559, 2010.

authors:
M. Kanagawa
P. Hennig
D. Sejdinovic
B.K. Sriperumbudur

title: Gaussian processes and kernel methods: A review on connections and equivalences

venue: arXiv preprint arXiv:1807.02582, 2018.

title: From -entropy to KL-entropy: Analysis of minimum information complexity density estimation

venue: The Annals of Statistics, vol. 34, no. 5, pp. 2180–2210, 2006.

authors:
A.W. van der Vaart
J.H. van Zanten

title: Rates of contraction of posterior distributions based on Gaussian process priors

venue: The Annals of Statistics, vol. 36, no. 3, Jun. 2008.

title: Information Rates of Nonparametric Gaussian Process Regression

venue: Journal of Machine Learning Research, vol. 12, pp. 2095–2119, 2011.

authors:
I. Steinwart
C. Scovel

title: Mercer’s Theorem on General Domains: On the Interaction between Measures, Kernels, and RKHSs

venue: Constructive Approximation, vol. 35, no. 3, pp. 363– 417, Jun. 2012.

authors:
L.H. Dicker
D.P. Foster
D. Hsu

title: Kernel ridge vs. principal component regression: Minimax bounds and the qualification of regularization operators

venue: Electronic Journal of Statistics, vol. 11, no. 1, Jan. 2017.

authors:
M. Moreira
J.R. Porter
G.A. Suarez

title: Bootstrap and higher-order expansion validity when instruments may be weak

title: Finite sample evidence of IV estimators under weak instruments

venue: Journal of Applied Econometrics, vol. 22, no. 3, pp. 677–694, Apr. 2007.

authors:
R. Davidson
J.G. Mackinnon

title: Wild Bootstrap Tests for IV Regression

venue: Journal of Business & Economic Statistics, vol. 28, no. 1, pp. 128–144, 2010, publisher: [American Statistical Association, Taylor & Francis, Ltd.].

authors:
T. Pearce
F. Leibfried
A. Brintrup

title: Uncertainty in neural networks: Approximately Bayesian ensembling

venue: Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, ser. Proceedings of Machine Learning Research, S. Chiappa and R. Calandra, Eds., vol. 108. PMLR, 26–28 Aug 2020, pp. 234–244.

authors:
B. He
B. Lakshminarayanan
Y.W. Teh

title: Bayesian deep ensembles via the neural tangent kernel

authors:
A. Jacot
F. Gabriel
C. Hongler

title: Neural Tangent Kernel: Convergence and Generalization in Neural Networks

venue: Advances in Neural Information Processing Systems, S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, Eds., vol. 31. Curran Associates, Inc., 2018.

title: Bayesian learning for neural networks

venue: Springer Science & Business Media,

authors:
J. Lee
Y. Bahri
R. Novak
S.S. Schoenholz
J. Pennington
J. Sohl-Dickstein

title: Deep neural networks as Gaussian processes

venue: arXiv preprint arXiv:1711.00165, 2017.

authors:
A.G. d. G. Matthews
M. Rowland
J. Hron
R.E. Turner
Z. Ghahramani

title: Gaussian process behaviour in wide deep neural networks

venue: arXiv preprint arXiv:1804.11271, 2018.

authors:
W. Hu
Z. Li
D. Yu

title: Simple and effective regularization methods for training on noisily labeled data with generalization guarantee

venue: International Conference on Learning Representations, 2020.

authors:
T. Hu
W. Wang
C. Lin
G. Cheng

title: Regularization matters: A nonparametric perspective on overparametrized neural network

venue: Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, ser. Proceedings of Machine Learning Research, A. Banerjee and K. Fukumizu, Eds., vol. 130. PMLR, 13–15 Apr 2021, pp. 829–837.

authors:
J. Lee
L. Xiao
S. Schoenholz
Y. Bahri
R. Novak
J. Sohl-Dickstein
J. Pennington

title: Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent

venue: Advances in Neural Information Processing Systems, H. Wallach, H. Larochelle, A. Beygelzimer, F. d. Alché-Buc, E. Fox, and R. Garnett, Eds., vol. 32. Curran Associates, Inc., 2019.

authors:
A. Rahimi
B. Recht

title: Random features for large-scale kernel machines.

venue: in NIPS,

title: Bayesian Method of Moments (BMOM) Analysis of Mean and Regression Models

title: Limited information likelihood and Bayesian analysis

venue: Journal of Econometrics, p. 19, 2002.

authors:
T. Jaakkola
M. Meila
T. Jebara

title: Maximum entropy discrimination

venue: Proceedings of the 12th International Conference on Neural Information Processing Systems, 1999, pp. 470– 476.

authors:
M. Dudík
S.J. Phillips
R.E. Schapire

title: Maximum entropy density estimation with generalized regularization and an application to species distribution modeling

venue: Journal of Machine Learning Research, 2007.

authors:
J. Zhu
E.P. Xing

title: Maximum entropy discrimination Markov networks.

venue: Journal of Machine Learning Research,

title: Gaussian processes and bayesian moment estimation

venue: Journal of Business & Economic Statistics, vol. 39, no. 2, pp. 482–492, 2021.

authors:
R. Zhang
M. Imaizumi
B. Schölkopf
K. Muandet

title: Maximum Moment Restriction for Instrumental Variable Regression

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

venue: arXiv preprint arXiv:2105.04544, 2021.

authors:
A. Puli
R. Ranganath

title: General control functions for causal effect estimation from instrumental variables

venue: Advances in Neural Information Processing Systems, vol. 33, p. 8440, 2020.

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

title: Learning Deep Features in Instrumental Variable Regression

authors:
S. Chib
M. Shin
A. Simoni

title: Bayesian estimation and comparison of conditional moment models

authors:
X. Chen
T.M. Christensen

title: Optimal sup-norm rates and uniform inference on nonlinear functionals of nonparametric IV regression

venue: Quantitative Economics, vol. 9, no. 1, pp. 39–84, 2018.

authors:
B.T. Knapik
A.W. van der Vaart
J.H. van Zanten

title: Bayesian inverse problems with Gaussian priors

venue: The Annals of Statistics, vol. 39, no. 5, pp. 2626–2657, Oct. 2011, arXiv: 1103.2692.

authors:
A. Van der Vaart
H. Van Zanten

title: Bayesian inference with rescaled Gaussian process priors

venue: Electronic Journal of Statistics, vol. 1, pp. 433–448, 2007.

authors:
C. Ai
X. Chen

title: Efficient estimation of models with conditional moment restrictions containing unknown functions

authors:
S. Ghosal
J.K. Ghosh
A.W. Van Der Vaart

title: Convergence rates of posterior distributions

venue: Annals of Statistics, pp. 500–531, 2000.

authors:
A. Caponnetto
E. De Vito

title: Optimal rates for the regularized least-squares algorithm

venue: Foundations of Computational Mathematics, vol. 7, no. 3, pp. 331–368, 2007.

title: Nonparametric Bayesian inference with kernel mean embedding

venue: Modern Methodology and Applications in Spatial-Temporal Modeling. Springer, 2015, pp. 1–24.

authors:
E. De Vito
L. Rosasco
A. Caponnetto
U. De Giovannini
F. Odone
P. Bartlett

title: Learning from examples as an inverse problem.

title: High-dimensional probability: An introduction with applications in data science

heading: 1 Introduction

text:
Instrumental variable (IV) regression is a standard approach for estimating causal effect from confounded observational data. In the presence of confounding, any regression method estimating E(y | x) cannot recover the causal relation f† between the outcome y and the treatment x, since the residual u = y − f†(x) is correlated with x due to the unobserved confounders. IV regression enables identification of the causal effect through the introduction of instruments, variables z that are known to influence y only through x.
IV regression is widely used in economics [1], epidemiology [2] and clinical research [3], but modeling nonlinear effect in IV regression can be challenging. Recent years have seen great development in adopting modern machine learning models for IV regression [4–8]. However, there is still a lack of uncertainty quantification measures for these flexible IV models. Uncertainty quantification is especially important for IV analysis, since unlike in standard supervised learning scenarios, we do not have (unconfounded) validation data, from which we could deduce the error pattern of the estimated model. Moreover, the instrument of choice may be weak, meaning that it only provides limited information for x; in such cases point estimators suffer from high variance [9]. This problem is exacerbated in the nonparametric setting, where IV estimation is usually an ill-posed inverse problem, where instruments provide vanishing information for the higher-order nonlinear effects [10].
The IV setting brings unique challenges for uncertainty quantification. For example, while it is natural to consider a Bayesian approach, specification of the likelihood requires knowledge of the entire data generating process, which is typically unavailable in IV regression. Consequently, most, if not all, existing work on Bayesian IV [11–15] assumes the following data generating process:
∗ZW and YZ contribute equally. ‡JZ is the corresponding author. An extended version is available at https://arxiv.org/abs/2106.08750v2.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
x = g(z) + u1, y = f(x) + u2, where (u1,u2) are correlated and independent of z. These methods then conduct posterior inference on g, f as well as the generative model for the unobserved confounders (u1,u2). However, the additive error in x is an unnecessary assumption for most point estimation procedures, and is difficult to check in high dimensions. Furthermore, the need to model the generating process of (u1,u2) introduces an extra risk of model misspecification, and Bayesian inference on the generative model is computationally expensive, especially on complex high-dimensional datasets. None of these issues present if only point estimation is needed.
For the above reasons, it is appealing to turn to an alternative quasi-Bayesian approach [16–18]. Quasi-Bayesian analysis views IV estimation as a generalized method-of-moments (GMM) procedure, and defines the quasi-posterior as a Gibbs distribution constructed from a chosen prior and violation of the moment constraints. It does not require full knowledge of the data generating process, and thus does not suffer from the aforementioned drawbacks of Bayesian inference. However, computation of the quasi-posterior is non-trivial, as its density contains a conditional expectation term E(y − f(x) | z), which itself needs to be estimated from data. This makes both approximate inference and theoretical analysis difficult. So far, quasi-Bayesian analysis for nonparametric IV is only developed on classical models such as wavelet basis [17, 18] or Nadaraya-Watson smoothing [19], while adoption of flexible machine learning models such as kernel machines or neural networks remains an open challenge. Moreover, numerical study has been limited in previous work, so little is known about the empirical performance of the quasi-Bayesian approach.
In this work, we present a novel quasi-Bayesian procedure for IV regression, building upon the recent development in kernelized IV models [7, 20]. We employ a Gaussian process (GP) prior and construct a quasi-likelihood using a kernel conditional expectation estimator. We establish theoretical properties of the resultant quasi-posterior, proving its consistency and showing that it may quantify instrument strength. Furthermore, inspired from the minimax formulation of IV estimation [5, 8, 20, 21], we design a principled approximate inference algorithm using random feature expansion and a novel adaptation of the “randomized prior trick” [22]. The algorithm has the form of stochastic gradient descent-ascent and is thus scalable and easy to implement. It can also be adapted to work with flexible neural network models, in which case its behavior can be formally justified by analyzing the neural networks in the kernel regime. Empirical evaluation shows that the proposed method produces informative uncertainty estimates, scales to high-dimensional and complex nonlinear problems, and is especially advantageous when the instrument is weak.
The rest of the paper is organized as follows: In Section 2 we set up the problem. We then derive the quasi-posterior and analyze its theoretical properties in Section 3, and present the approximate inference algorithm in Section 4. Section 5 reviews related work, and Section 6 presents numerical studies. Finally, we discuss conclusion and future work in Section 7.

heading: 2 Notations and Setup

text:
Notations We use boldface (x,y, z) to represent random variables on the space X × Y × Z , regular font (x, y, z) to denote deterministic values. [n] denotes the set {1, 2, · · · , n}. {(xi, yi, zi) : i ∈ [n]} indicates the training data. We use the notations X := (x1, . . . , xn) ∈ Xn, f(X) := (f(x1), . . . , f(xn)); likewise for Y,Z. For finite-dimensional vectors θ, θ′ ∈ Rm, we use ‖θ‖2, 〈θ, θ′〉2 to denote the Euclidean norm and inner product, respectively. For any operator A : H1 → H2 between Hilbert spaces H1 and H2, we denote its adjoint by A∗ : H2 → H1. When H1 = H2 and λ ∈ R, we use the notation Aλ := A+ λI for simplicity.
IV regression Denote the treatment and response variables as x,y, the instrument as z, and the true structural function of interest as f†. Consider the data generating process y = f†(x) + u, where the unobserved u satisfies E(u | z) = 0, but may correlate with x. Then f† satisfies
E(y − f†(x) | z) = 0 a.s. [P (dz)], (1)
where P denotes the data distribution. This conditional moment restriction (CMR) formulation is the standard definition in literature [e.g., 23, 18], and is used in the recent work on machine learning models for IV. It connects to GMM as (1) can be viewed as a continuum of generalized moment constraints. Note that (1) does not place any structural constraint on the conditional distribution p(x | z), such as additive noise; hence, it does not require full knowledge of the data generating
process. Also, as discussed in, e.g., Hartford et al. [4], the setup can also be extended to incorporate observed confounders v, by including v in both x and z.
Kernelized IV and a dual view Let H, I be suitably chosen function spaces on X ,Z , respectively. The generalized moment constraint (1) motivates the use of the following objective
min f∈H
L(f) := d2n(Ênf − b̂) + λ̄Ω(f), (2)
where Ên : H → I is an empirical approximation to the conditional expectation operator E : f 7→ E(f(x) | z = ·), b̂ is an estimator of E(y | z = ·), {dn} is a sequence of suitable (semi-)norm on I, and Ω : H → R is a regularization term. When H, I are reproducing kernel Hilbert spaces (RKHS) with corresponding kernels kx, kz , it is natural to set Ω(f) = 12‖f‖ 2 H, and define b̂ as the result of kernel ridge regression on y with respect to z. In this case Ên can be defined with the empirical kernel conditional expectation operator: let Czx = E(k(x, ·) ⊗ k(z, ·)), Czz = E(k(z, ·) ⊗ k(z, ·)). Assuming E maps all f ∈ H to Ef ∈ I, we have [24] CzzEf = Czxf, ∀f ∈ H. This motivates the use of Ên := Ĉ−1zz,ν̄Ĉzx where Ĉzz := 1 n ∑n i=1 k(zi, ·) ⊗ k(zi, ·), Ĉzx is defined similarly, and ν̄ is a regularization hyperparameter.2 The choice of dn is flexible; the dual IV formulation uses d2n(g) := 1 2n ∑n j=1 g(zj) 2 + ν̄2‖g‖ 2 I = 1 2 〈g, Ĉzz,ν̄g〉I . Introducing the evaluation operator Sz : I → Rn, Szg := (g(z1), . . . , g(zn)), we have b̂ = Ĉ−1zz,ν̄ S∗zY n , and
L(f) = 1 2 ∥∥∥∥Ĉ−1/2zz,ν̄ (Ĉzxf − S∗zYn )∥∥∥∥2 I + λ̄ 2 ‖f‖2H (3)
= max g∈I
1
n n∑ j=1 ( (f(xj)− yj)g(zj)− g(zj) 2 2 ) − ν̄ 2 ‖g‖2I + λ̄ 2 ‖f‖2H, (4)
where (4) holds because of the Fenchel duality 12u 2 = supv∈R ( uv − 12v 2 )
and the equality E(u(x,y)g(z)) = E(E(u(x,y) | z)g(z)); see [25, 8, 21, 26]. The dual formulation (4) circumvents the need to compute Ên directly, an operator between the typically infinite-dimensional spaces H and I, and leads to a scalable estimation procedure based on stochastic gradient descent-ascent (SGDA). It can also be generalized to work with deep neural networks (DNNs) instead of kernel machines, by replacing the RKHS regularizer with a suitable regularizer for DNNs, although theoretical analysis for the resulted algorithm requires separate effort [8, 26]. As we shall see, (4) will also enable the construction of a scalable approximate inference algorithm.

heading: 3 Quasi-Bayesian Analysis of Dual IV

text:
Introduce the notations Sx : H → Rn, Sxf := (f(x1), . . . , f(xn)), so that Ĉzx = 1nS ∗ zSx, Ĉzz = 1 nS ∗ zSz . Define λ := nλ̄, ν = nν̄. Now we can re-express (3) in an equivalent form as:
L̄(f) := n λ L(f) = 1 2 (f(X)− Y )>(λ−1L)(f(X)− Y ) + 1 2 ‖f‖2H, (5)
where L := 1nSzĈ −1 zz,ν̄S ∗ z is a linear map from Rn to Rn, and thus can be identified with an n × n matrix. Since the first term above is equivalent to the log density of the multivariate normal distribution N (Y | f(X), λL−1), we can view (5) as the objective of a kernel ridge regression problem, which has a data-dependent noise covariance λL−1.3 The connection between kernel ridge regression and Gaussian process regression [28] thus motivates the use of the quasi-posterior
2With an abuse of notation, k refers to the reproducing kernel of the corresponding RKHS (kx forH or kz for I) whenever the denotation is clear.
3Here we assume the invertibility of L for brevity. Alternatively, observe that (5) defines a linear inverse problem with the finite-dimensional observation operator f 7→ √ Lf(X) and noise variance λ−1I , and we can follow [27, Chapter 6] to derive the same quasi-posterior.
Π(df | D(n)), defined through the following Radon-Nikodym derivative w.r.t. the standard GP prior Π = GP(0, kx):
dΠ(· | D(n)) dΠ
(f) ∝ exp ( − 1
2 (f(X)−Y )>(λ−1L)(f(X)−Y )
) = exp ( −n λ d2n(Ênf − b̂) ) . (6)
Note that contrary to standard Bayesian modeling, we do not assume Y ∼ N (f(X), λL−1) is part of the true data generating process, nor does the theoretical analysis below rely on it. Instead, the quasi-posterior (6) should be interpreted as a Gibbs distribution which trades off between the properly scaled evidence λ−1nd2n(Ênf − b̂), which characterizes the estimated violation of the GMM constraint (1), and our prior belief Π(df). This trade-off is most clear from the well-known variational characterization of the Gibbs distribution [29],
Π(· | D(n)) = arg min Ψ Ef∼Ψ[λ−1nd2n(Ênf − b̂)] + KL(Ψ ‖Π). (7)
Nonetheless, the fictitious data generating process f ∼ GP(0, k), Y ∼ N (f(X), λL−1) is useful for deriving the quasi-posterior, since its conditional distribution pfic(df | Y ) coincides with (6) [27, Chapter 6]. In the probability space of this fictitious data generating process, for any finite set of test inputs x∗, we have
pfic(f(x∗), Y ) ∼ N ( 0, [ K∗∗ K∗x Kx∗ Kxx + λL −1 ]) ,
where K(·) denote the corresponding Gram matrices with subscript ∗ denoting x∗ and x denoting X (so, e.g., K∗x := k(x∗, X)). Thus by the Gaussian conditioning formula, we have
Π(f(x∗) | D(n)) = pfic(f(x∗) | Y ) = N (m,S), (8) where m := K∗x(λI + LKxx)−1LY, (9)
S := K∗∗ −K∗xL(λI +KxxL)−1Kx∗, (10) L = Kzz(Kzz + νI) −1. (11)
In the above Kzz := k(Z,Z) denotes the Gram matrix, and (11) follows from the definition L = n−1Sz(n −1S∗zSz + ν̄I) −1S∗z , the Woodbury identity, and the observation that SzS ∗ z = Kzz .
Theoretical Analysis For the quasi-posterior Π(· | D(n)) to be a useful measure of uncertainty, it needs to satisfy the following informal criteria: as n→∞, we expect
(C1) Π(· | D(n)) will exclude incorrect solutions; (C2) In cases of non-identification, Π(· | D(n)) will not exclude any valid solution in the model.
In the following, we formalize these criteria, and demonstrate that with appropriately chosen hyperparameters, the quasi-posterior satisfies both criteria. We will work with the following assumptions: Assumption 3.1. The restriction of the conditional expectation operator f 7→ E(f(x) | z = ·) on H, denoted as E, has its image contained in I. E : H → I is bounded. Assumption 3.1 is intuitive, and is also slightly more general than some previous work [20, 26] that require the conditional expectation operator to be bounded on the entire hypothesis space, which typically corresponds to the “sample space” of the GP prior in our setting, and is much larger than H (see Appendix A). Nonetheless, we note that Hypothesis 4 in Singh et al. [7] may be more general, although they impose extra smoothness assumptions on E. Assumption 3.2. The true structural function f†(x) is such that, there exists a sequence {τm : m ∈ N} satisfying τm → 0, and
mτ2m ≥ inf f†m∈H:‖f†−f†m‖≤τm ‖f†m‖2H − log Π({f : ‖f‖ ≤ τm}), (12)
where ‖ · ‖ denotes the sup norm. Assumption 3.2 requires that f† can be well approximated in H; this is more general than previous work [e.g., 20, 7] that require f† ∈ H, but is typical in the GP literature [30, 31]. The sequence {τm} is determined by the complexity of H and its ability for approximating f†, and is usually the optimal posterior contraction rate for Gaussian process regression on the unconfounded dataset {(xi, f†(xi) + i) : i ∈ [m]} [31].
Assumption 3.3. X and Z are Polish spaces. The kernels kx, kz are continuous, supx∈X k(x, x) + supz∈Z k(z, z) ≤ κ2, and Mercer’s representations [32] of kx and kz exist. The random variable y − f†(x) is sub-exponential. Assumption 3.3 imposes technical conditions frequently assumed in literature [e.g., 33]. The requirements on the kernels can be satisfied by e.g. continuous kernels on compact subsets of Rd. Given the assumptions above, we characterize (C1) by showing that as n→∞, the posterior places vanishing mass on the region of functions that violates the GMM constraints (1). Concretely, Theorem 3.1 (Proof in Appendix B). There exist a constantM > 0 depending on the data distribution and the kernels of choice, and a sufficiently slowly growing sequence {γn} → ∞ (e.g., we can always have γn ≤ log log n) such that when taking λ̄ = n−1/2, ν̄ = min{τ2d√ne, n −1/2γn}, it has
ED(n)Π({f : E2(f(x)− y | z) > Mτ2d√ne} | D (n))→ 0. (13)
Remark 3.1. The above result should primarily be interpreted as a consistency result, although it also provides a posterior contraction rate [34] in the order of τd√ne in the semi-norm f 7→ ‖Ef‖L2(P (dz)), where {τm} is defined in Assumption 3.2. As an example, suppose the regularity of f† is similar to the Matérn-1/2 RKHS; then for suitable kernels we have τd√ne = O(n−1/8), see Remark A.3.
The rate is suboptimal, and can be immediately improved if we impose further regularity assumption on kz: if we assume the critical radius of the local Rademacher complexity [35] of the unit-norm ball I1 is δn τn n− b 2(b+1) , a rate of O(n− b−1 2(b+1) ) for the semi-norm can be established following similar arguments. This is still worse than [26], which provides the rate O(max{τn, δn}) for what corresponds to our posterior mean estimator. Nonetheless, their result is also generally suboptimal, because the choice of δn τn is actually suboptimal when the IV regression problem is ill-posed. To date, minimax optimal rates have only been established under additional assumptions on the relations between the conditional expectation operator and the model for f . Our extended version establishes such results, for both this semi-norm and more intuitive norms such as L2(P (dx)).
Remark 3.2. The use of an increasing λ = nλ̄ is a technical artifact due to our minimal assumptions on kz . It is common to impose extra regularization in nonparametric IV, due to the need to estimate E from data [19], and, in the quasi-Bayesian setting, also due to the additional technical challenges [18]. However, this is not necessary for our method if additional assumptions are imposed, as we show in the extended version. In practice, our hyperparameter selection procedure does not produce λ with significant growth.
The following proposition characterizes (C2): Proposition 3.1 (Proof in Appendix B). The scaled log quasi-likelihood estimate `n(f) = d2n(Ênf − b̂) satisfies
Π ({ ∀f : ∀δ > 0, lim
n→∞ PD(n)( ∣∣`n(f)− Ez∼P (z)(E(f(x)− y | z)2)∣∣ > δ) = 0}) = 1. In words, for Π-almost every f , the log quasi-likelihood estimate scaled by n−1 converges to Ez∼P (z)(E(f(x) − y | z)2), which characterizes the violation of (1). While the restriction to the probability-1 subset can be concerning in the nonparametric setting [34], the proof only depends on f satisfying similar approximability conditions to f†; see (45). Remark 3.3. To understand why the proposition characterizes (C2), suppose there are multiple f that satisfy (1). Then all of them will eventually have significantly higher log likelihood than functions violating (1): the difference is Θ(λ−1n). If these functions have a similar level of regularity in the sense that their concentration functions, which is the right-hand side of (12), have the same asymptotics as → 0, Borell’s inequality [e.g., 34, Proposition 11.17] implies that the posterior mass of small ‖ · ‖-norm balls around them will also have the same asymptotics.
Finally, we provide the following (over-)simplified example, which compares the behavior of the quasi-posterior and bootstrap in the context of (C2): Example 3.1. Suppose z is completely non-informative so that E(f(x) | z) ≡ Ef(x) for all f ; and suppose the estimated conditional expectation Ên is sufficiently accurate, so that we replace it
with E.4 In this case bootstrap on (3) will always return the point estimator f ≡ 0 due to the nonzero regularization on ‖f‖H, while the quasi-posterior behaves like the prior, correctly reflecting the complete lack of evidence in data.
While this example is oversimplified, and in practice the estimation error of E plays an important role, it is known that bootstrap uncertainty estimates can be problematic given weak IVs [36–38]. We also observe similar failures for bootstrap in the experiments (see Section 6.1).

heading: 4 Scalable Approximate Inference via a Randomized Prior Trick

text:
We now turn to approximate inference with parametric models such as random feature expansion or wide NNs. Scalable inference for the IV quasi-posterior appears difficult, since for any f , computing the quasi-likelihood involves computing Ênf , which in turn requires either inverting an n × n Gram matrix, judging from (5) and (11), or solving an optimization problem specific to f from (4). Nonetheless, we show that it is possible, by extending the “randomized prior” trick for Gaussian process regression [22] to work with (quasi-)likelihoods with an optimization formulation as in (4).
Our algorithm works with random feature models. A random feature model for kz approximates kz(z, z ′) ≈ k̃z,m(z, z′) := 1mφz,m(z) >φz,m(z
′), where φz,m takes value in Rm. Then the map ϕ 7→ 1√
m ϕ>φz,m(·) =: g(·;ϕ) parameterizes an approximate RKHS H̃; and for all c > 0, the
random function g(·;ϕ), where ϕ ∼ N (0, cI), is distributed as GP(0, ck̃z,m). The notations related to kx are similar and thus omitted. Now we can state the objective function: Proposition 4.1 (Proof in Appendix C.1). Let φ0 ∼ N (0, λν−1I), θ0 ∼ N (0, I), ỹi ∼ N (yi, λ). Then the optima θ∗ of
min θ∈Rm max φ∈Rm n∑ i=1 ( (f(xi; θ)− ỹi)g(zi;φ)− g(zi;φ) 2 2 ) − ν 2 ‖φ− φ0‖22 + λ 2 ‖θ − θ0‖22 (14)
parameterizes a random function which follows the quasi-posterior distribution (6), where the kernels are replaced by the random feature approximations.
Given the above proposition, we can sample from the random feature-approximated quasi-posterior by solving (14) with stochastic gradient descent-ascent; the approximation errors will be analyzed in the following. The objective (14) is closely related to (4); as we show in Appendix C.1.1, it is equivalent to
min f∈H̃ max g∈Ĩ n∑ i=1 ( (f(xi)− ỹi)g(zi)− g(zi) 2 2 ) − ν 2 ‖g − g0‖2Ĩ + λ 2 ‖f − f0‖2H̃, (15)
which differs from (4) only in the regularizers: instead of regularizing the norm of f and g, (15) encourages the functions to stay close to randomly sampled anchors [39]. Alternatively, we can view (15) as perturbing the point estimator (4), so that it has a covariance matching that of the quasi-posterior. A similar relation is also observed in the original randomized prior trick, which transforms GP regression to the optimization problem minf∈H̃ ∑n i=1(f(xi)− ỹi)2 + λ‖f − f0‖22. In both cases, the resultant algorithm for approximate inference has the same time complexity as ensemble training for point estimation.
While the algorithm can be directly applied to neural network models as in [22], we follow [40] and modify the objective, to account for the difference between the neural tangent kernel (NTK) [41] of a wide neural network architecture and the NNGP kernel of the corresponding infinite-width Bayesian neural network [42–44]. Concretely, we modify (14) as
min θ max φ n∑ i=1 ( (f̃θ(xi)− ỹi)g̃φ(zi)− g̃φ(zi) 2 2 ) − ν 2 ‖φ− φ0‖22 + λ 2 ‖θ − θ0‖22, (16)
where g̃φ(z) := g(z;φ)− g(z;φ0) + g̃0(z), g̃0(z) := √ λ
ν
〈 φ̄0, ∂g
∂φ
∣∣ φ=φ0 (z) 〉 ,
4This setting could be realistic in the sample splitting setup considered in [7]. Note that as long as we have finite samples for the estimation of f (the “second stage”), we still need to have λ̄ > 0.
and φ0 denotes the initial value of φ, and φ̄0 ∼ N (0, I) is a set of randomly initialized NN parameters independent of φ0; and f̃θ is defined similarly.
We only give a formal justification for the modification, under the assumption5 that the NNs remain in the kernel regime throughout training, so that g(z;φ)−g(z;φ0) = 〈φ−φ0, ∂g(z)∂φ |φ0〉2 [47]. Thus for the purpose of analyzing g(·;φ) − g(·;φ0), we can view g as a random feature model with the parameterization φ 7→ 〈φ, ∂g(z)∂φ |φ0〉2. Thus by the argument in Appendix C.1.1, we can show that the weight regularizer ‖φ − φ0‖2 is equivalent to ‖g(·;φ) − g(·;φ0)‖Ĩ = ‖g̃φ − g̃0‖Ĩ , where Ĩ is determined by the NTK kg,ntk(z, z′) := 〈∂g(z)∂φ |φ0 , ∂g(z′) ∂φ |φ0〉2. Similar arguments can be made for f̃θ and f̃0. Consequently, (16) is equivalent to an instance of (15) with H̃, Ĩ defined by the NTKs. Implementation details for the algorithm, including hyperparameter selection, are discussed in Appendix D.
Convergence analysis We now complete the analysis of the inference algorithm, by showing that for any fixed set of test points x∗, SGDA can approximate the marginal distribution Π(f(x∗) | D(n)) arbitrarily well given a sufficient computational budget. This implies that the approximate posterior is good enough for prediction purposes.
We place several mild assumptions on the random feature model, listed in Appendix C.2; they are satisfied by common approximations such as the random Fourier features [48]. The SGDA algorithm is described in detail in Appendix C.4. Under this setup, we have Proposition 4.2 (Proof in Appendix C.5). Fix D(n) and λ, ν > 0. Then there exist a sequence of choices of m and SGDA step-size schemes, such that for any l ∈ N, we have
supx∗∈X l max(‖µ̂m − µm‖2, ‖Ŝm − Sm‖F ) p→ 0.

heading: In the above, µ̂m, Ŝm denote the mean and covariance of the approximate marginal posterior for f(x∗), µ, S correspond to the true posterior, ‖·‖F denotes the Frobenius norm, and the convergence in probability is defined with respect to the sampling of random feature basis.

heading: 5 Related Work

text:
Quasi-Bayesian analysis for GMM estimation problems was first developed in [49, 50, 16], which provided theoretical analysis for parametric models. The use of the quasi-posterior is motivated from the maximum entropy principle, based on which similar ideas have been developed in the machine learning literature [51–53]. Alternatively, it can also be obtained as a non-informative limit of certain Bayesian posteriors [54]. [17, 18] first analyzed the use of nonparametric models for quasi-Bayesian IV; no numerical study was presented. Closer to our work is [19] which constructed a quasi-posterior using Nadaraya-Watson smoothing as Ên. While similar in form to our method, the smoothing approach relies on stronger assumptions, especially for high-dimensional data; see [7, Appendix A.2.1] which compared the corresponding point estimators. Their approach is also harder to scale to large datasets.
Our quasi-Bayesian procedure builds upon the kernelized IV models [7, 20] and the dual formulation of IV regression [6, 20, 21, 26]. The formulation (3)-(4) is from [20, 26]. [7] proposes a similar method. Although it was motivated differently as a kernelized two-stage least squares (2SLS) estimator, its objective is asymptotically equivalent to (3); however, the slight difference in regularization prevents the use of estimation procedures similar to (4), as we discuss in Appendix C.1.4. An alternative kernel-based estimator is proposed in [55, 56]; in particular, [55] also includes a Gaussian process construction, but for the different purpose of computing a leave-one-out validation statistics. Other recent work on ML models for IV include [4, 57, 58]. It remains interesting future work to develop scalable quasi-Bayesian procedures for these methods, although the mean estimator derived from our quasi-posterior implementation has competitive performance.
[59] studies semi-parametric IV estimation using an exponentially tilted empirical likelihood (ETEL) estimator. ETEL also connect to Bayesian methods. However, similar to the quasi-Bayesian
5The same linearization assumption has been employed in [40, 45]. We refer readers to [46, 21] for analysis of the linearization error in various settings similar to ours.
approach, the need to estimate the conditional moment restrictions complicates the understanding of its frequentist behavior, as well as the design of scalable inference algorithms; in their case, it is especially unclear if the empirical likelihood can be computed in sublinear time. Other approaches for uncertainty quantification include Bayesian inference and bootstrap. We have discussed previous works on Bayesian IV and their limitations in Section 1. Bootstrap is typically justified in the asymptotic regime, which does not cover many scenarios where uncertainty is most needed; this is different from the quasi-Bayesian approach which can always be justified through (7). Moreover, standard bootstrap inference on 2SLS is known to be unreliable when the instrument strength is weak [36–38]; while remedies exist (e.g., [38]), they heavily rely on the linearity and additive noise assumptions, and are thus difficult to generalize to the nonlinear setting we are interested in. As the kernelized IV methods generalize 2SLS, we expect similar issues to exist in our setting.

heading: 6 Experiments

text: Code to reproduce the experiments is available at https://github.com/meta-inf/qbdiv.

heading: 6.1 1D Simulation

text:
We first experiment on a variety of 1D synthetic datasets, constructed by modifying the setup in [5] to incorporate a nonlinear first stage, in a way similar to [7, 60]:
z := sigmoid(w), x := sigmoid ( αw + (1− α)u′√ α2 + (1− α)2 ) , yi ∼ N (f0(2x− 1) + 2u, 0.1),
where (u, u′) are normal random variables with unit variance and a correlation of 0.5, w ∼ N (0, 1) is independent of (u, u′), α is a parameter controlling the instrument strength, and f0 is constructed from the sine, step, abs or a linear function. We choose N ∈ {200, 1000} and α ∈ {0.05, 0.5}. Our baselines include BayesIV [15], a state-of-the-art Bayesian model based on B-splines and Dirichlet process mixture; we also include bootstrap on 2SLS with ridge regularization, either applied directly to the input features (Linear), on their polynomial expansion (Poly), or on the same kernelized models (KIV) as ours.6 Hyperparameter for the kernelized IV methods are selected by cross validation based on the observable first-stage and second-stage losses as in previous work [7, 20]; see Appendix D.1. For kernels we choose the RBF and Matérn kernels, although results for Matérn kernels are deferred to appendix for brevity. See Appendix D.3 for the detailed setup.
Normalized MSE and coverage rate of 95% credible intervals (CI) on the sine datasets are plotted in Figure 1. We report results on 20 independently generated datasets. As we can see, quasi-Bayesian inference provides the most reliable uncertainty estimates, and is especially advantageous in the weak IV setting (α = 0.05). While its CI can be conservative, we note that it is still informative, and properly reflects the sample size and instrument strength; see Appendix D.4 for visualizations. These results connect to previous work showing (in a different setting) that the radius of credible set produced by a GP posterior can have the correct order of magnitude [61].
We report the average run time in Table 1. We can see that in addition to the improved flexibility, our method is also more scalable compared with the standard Bayesian baseline, as it avoids the costly modeling of the joint residual distribution.
6We do not compare with [19] since their source code is unavailable. Note that their smoothing-based method does not scale to large datasets, and there is no numerical study on the credible interval in [19].
Full results on all datasets, visualizations and additional experiments are deferred to Appendix D.4. As a summary, (i) on the abs and linear datasets where all modeling assumptions hold, the results are qualitatively similar to the sin dataset. Moreover, the over-smoothed RBF kernel appears to have similar coverage comparing with the optimal kernel, and follow a similar contraction rate, as the previous work on GP regression [62] suggests. (ii) On the step dataset which violates Assumption 3.2, the quasi-posterior still provides more coverage than the baselines. (iii) Uncertainty estimates produced by the approximate inference algorithm are similar to that from the exact quasi-posterior.

heading: 6.2 Airline Demand

text:
We now turn to the more challenging demand simulation first proposed by [4]. The dataset simulates a scenario where we need to predict the demand of airline tickets y, as a function of the price x, and two observed confounders: customer type s, and time of year t. The data generating process is
x := (z + 3)ψ(t) + 25 + u′, y := f0(x, t, s) + u, f0(x, t, s) := 100 + (10 + x) · s · ψ(t)− 2x where (u, u′) are standard normal variables with correlation ρ, z ∼ N (0, 1) is independent of (u, u′), and ψ is a nonlinear function whose shape is given in Figure 2. The variable s either varies across {0, . . . , 6} (the lower-dimensional setting), or is observed as an MNIST image representing the corresponding digit; the latter case represents the real-world scenario where only highdimensional surrogates of the true confounder is observed. We only report results for ρ = 0.5, noting that results using other choices of ρ have been similar. We use n ∈ {1000, 10000} for the lower-dimensional setting, and use n = 50000 for the image-based setting.
We compare our method with bootstrap on the same model, BayesIV, and bootstrap on linear or polynomial models. Performance of other point estimators on this dataset has been reported in [7, 20, 58], compared with which our method is generally competitive. We implement the dual IV model using both an RBF kernel and DNN models. See Appendix D.5 for details.
We report the test MSE and CI coverage for N = 1000 in Table 2, and visualize all approximate quasi-posteriors for the NN models in Figure 2. As we can see, when implemented with DNNs, our method produces uncertainty estimates with excellent coverage, which also correctly reflects the information available in the dataset: the CI is wider when N is smaller, or in the highdimensional experiment where estimation is harder. Bootstrap has a noticeably worse performance when N = 1000. Still, it performs well in the (arguably less interesting) large-sample setting, with a CI coverage similar to our method; see Appendix D.6. This is because on this dataset, the total instrument strength is stronger due to the presence of observed confounders, and the NN model is a good fit. Consequently, the asymptotic behavior of bootstrap can be observed when N is large.
Both methods have poorer performances when we switch to the RBF kernel, although the quasiposterior is still more reliable. We conjecture that both H and I are misspecified in this case, and
thus the credible intervals can only reflect the uncertainty within the prior model. These results show that NN models can be advantageous, which our inference algorithm supports.
The other baselines perform poorly due to their inflexibility; in particular, note that BayesIV uses additive models for both stages (e.g., f(x, t, s) = f1(x) + f2(t) + f3(s)) which do not approximate this data generating process well. Full results and visualizations are deferred to Appendix D.6.

heading: 7 Conclusion

text:
In this work we propose a scalable quasi-Bayesian procedure for IV regression. We analyze the theoretical properties of the proposed quasi-posterior, and derive a scalable algorithm for approximate inference. Empirical evaluations show that the proposed method scales to large and highdimensional datasets, and can be particularly advantageous when the instrument strength is weak.
Beyond IV regression, formulations like (1) also appear in various other problems in causal inference and statistics, as discussed in, e.g., [21]; our method can be readily applied to these problems. Future work includes extension to more general conditional moment restriction problems with nonlinear constraints [63, 17].

heading: Acknowledgement

text:
We thank the anonymous reviewers for their helpful feedback and references. Z.W., Y.Z. and J.Z. were supported by NSFC Projects (Nos. 61620106010, 62061136001, 61621136008, 62076147, U19B2034, U19A2081, U1811461), Beijing NSF Project (No. JQ19016), Beijing Academy of Artificial Intelligence (BAAI), Tsinghua-Huawei Joint Research Program, a grant from Tsinghua Institute for Guo Qiang, Tiangong Institute for Intelligent Computing, and the NVIDIA NVAIL Program with GPU/DGX Acceleration.
Broader Impact
IV analysis is widely used in social science and clinical research, where it can be highly undesirable to have a systematic error or increased prediction variance on subpopulations which are underrepresented in the training dataset, or on which the instrument is locally weak. Uncertainty quantification is an important first step in detecting such issues, and the quasi-Bayesian approach we take can be preferable since it has a clear interpretation in a non-asymptotic setting, and when the instrument strength is weak. However, epistemic uncertainty estimates are only as good as the model allows, and it should not create a false sense of complete security. Caution must be taken when employing complex models in scenarios with potential fairness implications.

title: Scalable Quasi-Bayesian Inference for Instrumental Variable Regression