Title: IMPLICIT CAUSAL MODELS FOR GENOME-WIDE ASSOCIATION STUDIES

Abstract: Progress in probabilistic generative models has accelerated, developing richer models with neural architectures, implicit densities, and with scalable algorithms for their Bayesian inference. However, there has been limited progress in models that capture causal relationships, for example, how individual genetic factors cause major human diseases. In this work, we focus on two challenges in particular: How do we build richer causal models, which can capture highly nonlinear relationships and interactions between multiple causes? How do we adjust for latent confounders, which are variables influencing both cause and effect and which prevent learning of causal relationships? To address these challenges, we synthesize ideas from causality and modern probabilistic modeling. For the first, we describe implicit causal models, a class of causal models that leverages neural architectures with an implicit density. For the second, we describe an implicit causal model that adjusts for confounders by sharing strength across examples. In experiments, we scale Bayesian inference on up to a billion genetic measurements. We achieve state of the art accuracy for identifying causal factors: we significantly outperform existing genetics methods by an absolute difference of 15-45.3%.

1. INTRODUCTION:

Probabilistic models provide a language for specifying rich and flexible generative processes (Pearl, 1988; Murphy, 2012) . Recent advances expand this language with neural architectures, implicit densities, and with scalable algorithms for their Bayesian inference (Rezende et al., 2014; Tran et al., 2017) . However, there has been limited progress in models that capture high-dimensional causal relationships (Pearl, 2000; Spirtes et al., 1993; Imbens & Rubin, 2015) . Unlike models which learn statistical relationships, causal models let us manipulate the generative process and make counterfactual statements, that is, what would have happened if the distributions changed.

As the running example in this work, consider genome-wide association studies (GWAS) (Yu et al., 2005; Price et al., 2006; Kang et al., 2010) . The goal of GWAS is to understand how genetic factors, i.e., single nucleotide polymorphisms (SNPs), cause traits to appear in individuals. Understanding this causation both lets us predict whether an individual has a genetic predisposition to a disease and also understand how to cure the disease by targeting the individual SNPs that cause it.

With this example in mind, we focus on two challenges to combining modern probabilistic models and causality. The first is to develop richer, more expressive causal models. Probabilistic causal models represent variables as deterministic functions of noise and other variables, and existing work usually focuses on additive noise models (Hoyer et al., 2009) such as linear mixed models (Kang et al., 2010) . These models apply simple nonlinearities such as polynomials, hand-engineered low order interactions between inputs, and assume additive interaction with, e.g., Gaussian noise. In GWAS, strong evidence suggests that susceptibility to common diseases is influenced by epistasis (the interaction between multiple genes) (Culverhouse et al., 2002; McKinney et al., 2006) . We would like to capture and discover such interactions. This requires models with nonlinear, learnable interactions among the inputs and the noise.

The second challenge is how to address latent population-based confounders. In GWAS, both latent population structure, i.e., subgroups in the population with ancestry differences, and relatedness among sample individuals produce spurious correlations among SNPs to the trait of interest. Existing methods correct for this correlation in two stages (Yu et al., 2005; Price et al., 2006; Kang et al., 2010) : first, estimate the confounder given data; then, run standard causal inferences given the estimated confounder. These methods are effective in some settings, but they are difficult to understand as principled causal models, and they cannot easily accommodate complex latent structure.

To address these challenges, we synthesize ideas from causality and modern probabilistic modeling. For the first challenge, we develop implicit causal models, a class of causal models that leverages neural architectures with an implicit density. With GWAS, implicit causal models generalize previous methods to capture important nonlinearities, such as gene-gene and gene-population interaction. Building on this, for the second challenge, we describe an implicit causal model that adjusts for population-confounders by sharing strength across examples (genes). We derive conditions that prove the model consistently estimates the causal relationship. This theoretically justifies existing methods and generalizes them to more complex latent variable models of the confounder.

In experiments, we scale Bayesian inference on implicit causal models on up to a billion genetic measurements. Validating these results are not possible for observational data (Pearl, 2000) , so we first perform an extensive simulation study of 11 configurations of 100,000 SNPs and 940 to 5,000 individuals. We achieve state of the art accuracy for identifying causal factors: we significantly outperform existing genetics methods by an absolute difference of 15-45.3%. In a real-world GWAS, we also show our model discovers real causal relationships-identifying similar SNPs as previous state of the art-while being more principled as a causal model.

1.1. RELATED WORK:

There has been growing work on richer causal models. Louizos et al. (2017) develop variational auto-encoders for causality and address local confounders via proxy variables. Our work is complementary: we develop implicit models for causality and address global confounders by sharing strength across examples. In other work, Mooij et al. (2010) propose a Gaussian process over causal mechanisms, and Zhang & Hyvärinen (2009) study post-nonlinear models, which apply a nonlinearity after adding noise. These models typically focus on the task of causal discovery, and they assume fixed nonlinearities (post-nonlinear models) or impose strong smoothness assumptions (Gaussian processes) which we relax using neural networks. In the potential outcomes literature, much recent work has considered decision trees and neural networks (e.g., Hill (2011); Wager & Athey (2015) ; Johansson et al. (2016) ). These methods tackle a related but different problem of balancing covariates across treatments.

Causality with population-confounders has primarily been studied for genome-wide association studies (GWAS). A popular approach is to first, estimate the confounder using the top principal components of the genotype matrix of individuals-by-SNPs; then, linearly regress the trait of interest onto the genetic factors and these components (Price et al., 2006; Astle & Balding, 2009) . Another approach is to first, estimate the confounder via a "kinship matrix" on the genotypes; then, fit a linear mixed model of the trait given genetic factors, and where the covariance of random effects is the kinship matrix (Yu et al., 2005; Kang et al., 2010) . Other work adjusts for the confounder via admixture models and factor analysis (Song et al., 2015; Hao et al., 2016) . This paper builds on all these methods, providing a theoretical understanding about when causal inferences can succeed while adjusting for latent confounders. We also develop a new causal model with nonlinear, learnable gene-gene and gene-population interactions; and we describe a Bayesian inference algorithm that justifies the two-stage estimation.

The problem of epistasis, that is, nonadditive interactions between multiple genes, dates back to classical work on epigenetics by Bateson and R.A. Fisher (Fisher, 1918) . Primary methods for capturing epistasis include adding interactions within a linear model, permutation tests, exhaustive search, and multifactor dimensionality reduction (Cordell, 2009) . These methods require hand-engineering over all possible interactions, which grows exponentially in the number of genetic factors. Neural networks have been applied to address epistasis for epigenomic data, such as to predict sequence specificities of protein bindings given DNA sequences (Alipanahi et al., 2015) . These methods use discriminative neural networks (unlike neural networks within a generative model), and they focus on prediction rather than causality.

2. IMPLICIT CAUSAL MODELS:

We describe the framework of probabilistic causal models. We then describe implicit causal models, an extension of implicit models for encoding complex, nonlinear causal relationships. 

2.1. PROBABILISTIC CAUSAL MODELS:

Probabilistic causal models (Pearl, 2000) , or structural equation models, represent variables as deterministic functions of noise and other variables. As illustration, consider the causal diagram in Figure 1 . It represents a causal model where there is a global variable

β = f β ( β ), β ∼ s(•),

and for each data point n = 1, . . . , N ,

EQUATION

The noise are background variables, representing unknown external quantities which are jointly independent. Each variable β, x, y is a function of other variables and its background variable.

We are interested in estimating the causal mechanism f y . It lets us calculate quantities such as the causal effect p(y | do(X = x), β), the probability of an outcome y given that we force X to a specific value x and under fixed global structure β. This quantity differs from the conditional p(y | x, β). The conditional takes the model and filters to the subpopulation where X = x; in general, the processes which set X to that value may also have influenced Y . Thus the conditional is not the same as if we had manipulated X directly (Pearl, 2000) .

Under the causal graph of Figure 1 , the adjustment formula says that p(y | do(x), β) = p(y | x, β). This means we can estimate f y from observational data {(x n , y n )}, assuming we observe the global structure β. For example, an additive noise model (Hoyer et al., 2009) posits

y n = f (x n , β | θ) + n , ∼ s(•),

where f (•) might be a linear function of the concatenated inputs, f (•) = [x n , β] θ, or it might use spline functions for nonlinearities. If s(•) is standard normal, the induced density for y is normal with unit variance. Placing a prior over parameters p(θ), Bayesian inference yields

EQUATION

The right hand side is a joint density whose individual components can be calculated. We can use standard algorithms, such as variational inference or MCMC (Murphy, 2012) .

A limitation in additive noise models that they typically apply simple nonlinearities such as polynomials, hand-engineered low-order interactions between inputs, and assume additive interaction with, e.g., Gaussian noise. Next we describe how to build richer causal models which relax these restrictions. (An additional problem is that we typically don't observe β; we address this in § 3.)

Implicit models capture an unknown distribution by hypothesizing about its generative process (Diggle & Gratton, 1984; Tran et al., 2017) . For a distribution p(x) of observations x, recent advances define a function g that takes in noise ∼ s(•) and outputs x given parameters θ,

x = g( | θ), ∼ s(•). (3) z n x nm y n M N w m z n φ θ x nm y n M N Figure 2: (left)

Causal graph for GWAS. The population structure of SNPs for each individual (z n ) confounds inference of how each SNP (x nm ) causes a trait of interest (y n ). (right) Implicit causal model for GWAS (described in § 3.2). Its structure is the same as the causal graph but also places priors over parameters φ and θ and with a latent variable w m per SNP.

Unlike models which assume additive noise, setting g to be a neural network enables multilayer, nonlinear interactions. Implicit models also separate randomness from the transformation; this imitates the structural invariance of causal models (Equation 1).

To enforce causality, we define an implicit causal model as a probabilistic causal model where the functions g form structural equations, that is, causal relations among variables. Implicit causal models extend implicit models in the same way that causal networks extend Bayesian networks (Pearl & Verma, 1991) and path analysis extends regression analysis (Wright, 1921) . They are nonparametric structural equation models where the functional forms are themselves learned.

A natural question is the representational capacity of implicit causal models. In general, they are universal approximators: we can use a fully connected network with a sufficiently large number of hidden units to approximate each causal mechanism. We describe this formally.

Theorem (Universal Approximation Theorem). Let the tuple (E, V, F, s(E)) denote a probabilistic causal model, where E represent the set of background variables with probability s(E), V the set of endogenous variables, and F the causal mechanisms. Assume each causal mechanism is a continuous function on the m-dimensional unit cube f ∈ C([0, 1] m ). Let σ be a nonconstant, bounded, and monotonically-increasing continuous function.

For each causal mechanism f and any error δ > 0, there exist parameters θ = {v, w, b}, for real constants v i , b i ∈ R and real vectors w i ∈ R m for i = 1, . . . , H and fixed H, such that the following function approximates f :

g(x | θ) = H i=1 v i σ w T i x + b i , |g(x | θ) -f (x)| < δ for all x ∈ [0, 1] m .

The implicit model defined by the collection of functions g and same noise distributions universally approximates the true causal model.

(This directly follows from the approximator theorem of, e.g., Cybenko (1989) .) A key aspect is that implicit causal models are not only universal approximators, but that we can use fast algorithms for their Bayesian inference (to calculate Equation 2). In particular, variational methods both scale to massive data and provide accurate posterior approximations ( § 4). This lets us obtain good performance in practice with finite-sized neural networks; § 5 describes such experiments.

3. IMPLICIT CAUSAL MODELS WITH LATENT CONFOUNDERS:

We described implicit causal models, a rich class of models that can capture arbitrary causal relations. For simplicity, we assumed that the global structure is observed; this enables standard inference methods. We now consider the typical setting when it is unobserved.

3.1. CAUSAL INFERENCE WITH A LATENT CONFOUNDER:

Consider the running example of genome-wide association studies (GWAS) (Figure 2 ). There are N data points (individuals). Each data point consists of an input vector of length M (measured SNPs), x n = [x n1 , . . . , x nM ] and a scalar outcome y n (the trait of interest). Typically, the number of measured SNPs M ranges from 100,000 to 1 million and the number of individuals N ranges from 500 to 10,000.

We are interested in how changes to each SNP X m cause changes to the trait Y . Formally, this is the causal effect p(y | do(x m ), x -m ), which is the probability of an outcome y given that we force SNP X m = x m and consider fixed remaining SNPs x -m . Standard inference methods are confounded by the unobserved population structure of SNPs for each individual, as well as the individual's cryptic relatedness to other samples in the data set. This confounder is represented as a latent variable z n , which influences x nm and y n for each data index n; see Figure 2 . Because we do not observe the z n 's, the causal effect p(y | do(x m ), x -m ) is unidentifiable (Spirtes et al., 1993) .

Building on previous GWAS methods ( Price et al., 2006; Yu et al., 2005; Astle & Balding, 2009) , we build a model that jointly captures z n 's and the mechanisms for X m → Y . Consider the implicit causal model where for each data point n = 1, . . . , N and for each SNP m = 1, . . . , M ,

EQUATION

The function g z (•) for the confounder is fixed. Each function g xm (• | w m ) per SNP depends on the confounder and has parameters w m . The function g y (• | θ) for the trait depends on the confounder and all SNPs, and it has parameters θ. We place priors over the parameters p(w m ) and p(θ). 2only requires a model from X → Y , and the rest of the graph is "ignorable" (Imbens & Rubin, 2015) .

To estimate the mechanism f y , we calculate the posterior of the outcome parameters θ,

EQUATION

Note how this accounts for the unobserved confounders: it assumes that p(z | x, y) accurately reflects the latent structure. In doing so, we perform inferences for p(θ | x, y, z), averaged over posterior samples from p(z | x, y).

In general, causal inference with latent confounders can be dangerous: it uses the data twice (once to estimate the confounder; another to estimate the mechanism), and thus it may bias our estimates of each arrow X m → Y . Why is this justified? We answer this below. Proposition 1. Assume the causal graph of Figure 2 (left) is correct and that the true distribution resides in some configuration of the parameters of the causal model (Figure 2 (right)). Then the posterior p(θ | x, y) provides a consistent estimator of the causal mechanism f y .

(See Appendix A for the proof.) Proposition 1 rigorizes previous methods in the framework of probabilistic causal models. The intuition is that as more SNPs arrive ("M → ∞, N fixed"), the posterior concentrates at the true confounders z n , and thus we can estimate the causal mechanism given each data point's confounder z n . As more data points arrive ("N → ∞, M fixed"), we can estimate the causal mechanism given any confounder z n as there are infinity of them.

Connecting to Two-Stage Estimation. Existing GWAS methods adjust for latent population structure using two stages (Astle & Balding, 2009) : first, estimate the confounders z 1:N ; second, infer the outcome parameters θ given the data set and the estimate of the confounders. To incorporate uncertainty, a Bayesian version would not use a point estimate of z 1:N but the full posterior p(z 1:N | x, y); then it would infer θ given posterior samples of z 1:N . Following Equation 5, this is the same as joint posterior inference. Thus the two stage approach is justified as a Bayesian approximation that uses a point estimate of the posterior.

3.2. IMPLICIT CAUSAL MODEL WITH A LATENT CONFOUNDER:

Above, we outlined how to specify an implicit causal model for GWAS. We now specify in detail the functions and priors for the confounders z n , the SNPs x nm , and the traits y n (Equation 4). Figure 2 (right) visualizes the model we describe below. Appendix B provides an example implementation in the Edward probabilistic programming language (Tran et al., 2016) .

Generative Process of Confounders z n . We use standard normal noise and set the confounder function g z (•) to the identity. This implies the distribution of confounders p(z n ) is standard normal.

Their dimension z n ∈ R K is a hyperparameter. The dimension K should be set to the highest value such that the latent space most closely approximates the true population structure but smaller than the total number of SNPs to avoid overfitting.

Generative Process of SNPs x nm . Designing nonlinear processes that return matrices is an ongoing research direction (e.g., Lawrence (2005); Lloyd et al. ( 2012)). To design one for GWAS (the SNP matrix), we build on an implicit modeling formulation of factor analysis; it has been successful in GWAS applications (Price et al., 2006; Song et al., 2015) . Let each SNP be encoded as a 0, 1, or 2 to denote the three possible genotypes. This is unphased data, where 0 indicates two major alleles; 1 indicates one major and one minor allele; and 2 indicates two minor alleles. Set

logit π nm = z n w m , x nm = I[ 1 > π nm ] + I[ 2 > π nm ],

1 , 2 ∼ Uniform(0, 1). This defines a Binomial(2, π nm ) distribution on x nm . Analogous to generalized linear models, the Binomial's logit probability is linear with respect to z n . We then sum up two Bernoulli trials: they are represented as indicator functions of whether a uniform sample is greater than the probability. (The uniform noises are newly drawn for each index n and m.)

Assuming a standard normal prior on the variables w m , this generative process is equivalent to logistic factor analysis. The variables w m act as "principal components," embedding the M -many SNPs within a subspace of lower dimension K.

Logistic factor analysis makes strong assumptions: linear dependence on the confounder and that one parameter per dimension has sufficient representational capacity. We relax these assumptions using a neural network over concatenated inputs,

logit π nm = NN([z n , w m ] | φ).

Similar to the above, the variables w m serve as principal components. The neural network takes an input of dimension 2K and outputs a scalar real value; its weights and biases φ are shared across SNPs m and individuals n. This enables learning of nonlinear interactions between z n and w m , preserves the model's conditional independence assumptions, and avoids the complexity of a neural net that outputs the full N × M matrix. We place a standard normal prior over φ.

Generative Process of Traits y n . To specify the traits, we build on an implicit modeling formulation of linear regression. It is the mainstay tool in GWAS applications (Price et al., 2006; Song et al., 2015) . Formally, for real-valued y ∈ R, we model each observed trait as

y n = [x n,1:M , z n ] θ + n ,

n ∼ Normal(0, 1), This process assumes linear dependence on SNPs, no gene-gene and gene-population interaction, and additive noise. We generalize this model using a neural network over the same inputs,

y n = NN([x n,1:M , z n , ] | θ),

n ∼ Normal(0, 1). The neural net takes an input of dimension M +K +1 and outputs a scalar real value; for categorical outcomes, the output is discretized over equally spaced cutpoints. We also place a group Lasso prior on weights connecting a SNP to a hidden layer. This encourages sparse inputs: we suspect few SNPs affect the trait (Yuan & Lin, 2006) . We use standard normal for other weights and biases.

4. LIKELIHOOD-FREE VARIATIONAL INFERENCE:

We described a rich causal model for how SNPs cause traits and that can adjust for latent populationconfounders. Given GWAS data, we aim to infer the posterior of outcome parameters θ (Equation 5). Calculating this posterior reduces to calculating the joint posterior of confounders z n , SNP parameters w m and φ, and trait parameters θ,

p(z 1:N , w 1:M , φ, θ | x, y) ∝ p(φ)p(θ) N n=1 p(z n )p(y n | x n,1:M , z n , θ) M m=1 p(w m )p(x nm | z n , w m , φ) .

This means we can use typical inference algorithms on the joint posterior. We then collapse variables to obtain the marginal posterior of θ. (For Monte Carlo methods, we drop the auxiliary samples; for variational methods, it is given if the variational family follows the posterior's factorization.)

One difficulty is that with implicit models, evaluating the density is intractable: it requires integrating over a nonlinear function with respect to a high-dimensional noise (Equation 3). Thus we require likelihood-free methods, which assume that one can only sample from the model's likelihood (Marin et al., 2012; Tran et al., 2017) . Here we apply likelihood-free variational inference (LFVI), which we scale to billions of genetic measurements (Tran et al., 2017) .

As with all variational methods, LFVI posits a family of distributions over the latent variables and then optimizes to find the member closest to the posterior. For the variational family, we specify normal distributions with diagonal covariance for the SNP components w m and confounder z n , q(w m ) = Normal(w m ; µ wm , σ wm I), q(z n ) = Normal(z n ; µ zn , σ zn I).

We specify a point mass for the variational family on both neural network parameters φ and θ. (This is equivalent to point estimation in a variational EM setting.)

For LFVI to scale to massive GWAS data, we use stochastic optimization by subsampling SNPs (Gopalan et al., 2016) . At a high level, the algorithm proceeds in two stages. In the first stage, LFVI cycles through the following steps:

1. Sample SNP location m and collect the observations at that location from all individuals. 2. Use the observations and current estimate of the confounders z 1:N to update the m th SNP component w m and SNP neural network parameters φ. 3. Use the observations and current estimate of SNP components w 1:M to update the confounders z 1:N .

This first stage infers the posterior distribution of confounders z n and SNPs parameters w m and φ. Each step's computation is independent of the number of SNPs, allowing us to scale to millions of genetic factors. In experiments, the algorithm converges while scanning over the full set of SNPs only once or twice.

In the second stage, we infer the posterior of outcome parameters θ given the inferred confounders from the first stage. Appendix C describes the algorithm in more detail; it expands on the LFVI implementation in Edward (Tran et al., 2016) .

5. EMPIRICAL STUDY:

We described implicit causal models, how to adjust for latent population-based confounders, and how to perform scalable variational inference. In general, validating causal inferences on observational data is not possible (Pearl, 2000) . Therefore to validate our work, we perform an extensive simulation study on 100,000 SNPs, 940 to 5,000 individuals, and across 100 replications of 11 settings. The study indicates that our model is significantly more robust to spurious associations, with a state-of-the-art gain of 15-45.3% in accuracy. We also apply our model to a real-world GWAS of Northern Finland Birth Cohorts; our model indeed captures real causal relationships-identifying similar SNPs as previous state of the art.

We compare against three methods that are currently state of the art: PCA with linear regression (Price et al., 2006) ("PCA"); a linear mixed model (with the EMMAX software) (Kang et al., 2010) ("LMM"); and logistic factor analysis with inverse regression (Song et al., 2015) ("GCAT"). In all experiments, we use Adam with a initial step-size of 0.005, initialize neural network parameters uniformly with He variance scaling (He et al., 2015) , and specify the neural networks for traits and SNPs as fully connected with two hidden layers, ReLU activation, and batch normalization (hidden layer sizes described below). For the trait model's neural network, we found that including latent variables as input to the final output layer improves information flow in the network.

5.1. SIMULATION STUDY: ROBUSTNESS TO SPURIOUS ASSOCIATIONS:

We analyze 11 simulation configurations, where each configuration uses 100,000 SNPs and 940 to 5,000 individuals. We simulate 100 GWAS data sets per configuration for a grand total of 4,400 fitted models (4 methods of comparison (Pritchard et al., 2000) ; and four variations of a configuration where population structure is determined by a latent spatial position of individuals. Only 10 of the 100,000 SNPs are set to be causal. Appendix D provides more detail.

Table 1 displays the precision for predicting causal factors across methods. When failing to account for population structure, "spurious associations" occur between genetic markers and the trait of interest, despite the fact that there is no biological connection. Precision is the fraction of the number of true positives over the number of true and false positives. This measures a method's robustness to spurious associations: higher precision means fewer false positives and thus more robustness. 1Table 1 shows that our method achieves state of the art across all configurations. Our method especially dominates in difficult tasks with sparse (small a), spatial (Spatial), and/or mixed membership structure (PSD): there is over a 15% margin in difference to the second best in general, and up to a 45.3% margin on the Spatial (a = 0.01) configuration. For simpler configurations, such as a mixture model (HapMap), our method has comparable performance.

5.2. NORTHERN FINLAND BIRTH COHORT DATA:

We analyze a real-world GWAS of Northern Finland Birth Cohorts (Sabatti et al., 2009) , which measure several metabolic traits and height and which contain 324,160 SNPs and 5,027 individuals. We separately fitted 10 implicit causal models, each of which models the effect of SNPs on one of ten traits. To specify the implicit causal models, we set the latent dimension of confounders to be 6 (following Song et al. (2015) ). We use 512 units in both hidden layers of the SNP neural network and use 32 and 256 units for the trait neural network's first and second hidden layers respectively. Appendix E provides more detail.

Table 2 compares the number of identified causal SNPs across methods, with an additional "uncorrected" baseline, which does not adjust for any latent population structure. Each method is performed with a subsequent correction as measured by the genomic control inflation factor (Sabatti et al., 2009) . Our models identify similar causal SNPs as previous methods. Interestingly, our model tends to agree with Song et al. (2015) , identifying a total of 15 significant loci (Song et al. (2015) identified 16; others identified 11-14 loci). This makes sense intuitively, as Song et al. (2015) uses logistic factor analysis which, compared to all methods, most resembles our model. 

6. DISCUSSION:

We described implicit causal models, a rich class of models that can capture high-dimensional, nonlinear causal relationships. With genome-wide association studies, implicit causal models generalize previous successful methods to capture important nonlinearities, such as gene-gene and gene-population interaction. In addition, we described an implicit causal model that adjusts for confounders by sharing strength across examples. Our model achieves state-of-the-art accuracy, significantly outperforming existing genetics methods by 15-45.3%.

There are several limitations to learning true causal associations. For example, alleles at different loci typically exhibit linkage disequilibrium, which is a local non-random association influenced by factors such as the rate of recombination, mutation, and genetic drift. The implicit causal model might be extended with variables shared across subsets of SNPs to model the recombination process.

Another limitation involves the data, where granularity of sequenced loci may lose signal or attribute causation to a region involving multiple SNPs. Better technology, and accounting for mishaps in the sequencing process in the model, can help.

While we focused on GWAS applications in this paper, we also believe implicit causal models have significant potential in other sciences: for example, to design new dynamical theories in high energy physics; and to accurately model structural equations of discrete choices in economics. We're excited about applications to these new domains, leveraging modern probabilistic modeling and causality to drive new scientific understanding.

A CONSISTENCY:

Consider the simplest setting in § 2, where the causal graph is as shown with a global confounder with finite dimension and where we observe the data set {(x n , y n )}. Assume the specified model family over the causal graph includes the true data generating process.

First consider an atomic intevention do(X = x) and let β * be the true structural value that generated our observations. The probability of a new outcome given the intervention and global structure is

p * (y | do(X = x), β * )) = p * (y | x, β * ).

This follows from the backdoor criterion on the empty set. By Bernstein von-Mises, the posterior for our model p(β | x, y) concentrates at β * . Thus, similarly, our posterior for θ given β * concentrates to the true functional mechanism f y . This implies we have a consistent estimate of

p * (y | do(•), β * ).

This simple proof rigorizes the ideas behind learning and fixing population structure, a common heuristic in GWAS methods. Moreover, it lets us understand how to extend them to more complex latent variable models of the confounder and also provide uncertainty estimates of the latent structure become important as we apply these methods to finite data in practice.

B IMPLICIT CAUSAL MODEL IN EDWARD:

We provide an example of an implicit causal model written in the Edward language below. It writes neural net weights and biases as model parameters. It does not include priors on the weights or biases; we add these as penalties to the objective function during training. The log-likelihood per-individual and per-SNP is

log p(x nm , y n | w m , z n , θ, φ) = log p(y n | x n,1:M , z n , θ) + log p(x nm | w m , z n , φ).

There are local priors p(w m ), p(z n ) and global priors p(θ), p(φ). The posterior factorizes as

p(z 1:N , w 1:M , φ, θ | x, y) = p(φ | x)p(θ | x, y) M m=1 p(w m | x 1:N,M , φ) N n=1 p(z n | x n,1:M , y n , θ, φ, w 1:M ).

Let the variational family follow the posterior's factorization above. For notational convenience, we drop the data dependence in q

(φ) = q(φ | x), q(θ) = q(θ | x, y) q(w m | φ) = q(w m | x 1:N,m , φ), q(z n | θ, φ, w 1:M ) = q(z n | x n,1:M , y n , θ, φ, w 1:M ).

We write the evidence lower bound and decompose the model's log joint density and the variational density,

L = M m=1 N n=1 E q(φ)q(wm | φ)q(θ)q(zn | θ,φ,w 1:M ) log p(x nm | w m , z n , φ) + N n=1 E q(φ)q(w 1:M | φ)q(θ)q(zn | θ,φ,w 1:M ) log p(y n | x n,1:M , z n , θ) + M m=1 E q(φ)q(wm | φ) log p(w m ) -log q(w m | φ) + N n=1 E q(φ)q(w 1:M | φ)q(θ)q(zn | θ,φ,w 1:M ) log p(z n ) -log q(z n | θ, φ, w 1:M ) + E q(φ) [log p(φ) -log q(φ)] + E q(θ) [log p(θ) -log q(θ)].

Assume that the variational family for z n and w m are independent of other variables, q(z n ) and q(w m ). Also assume delta point masses for q(θ) = I[θ-θ ] parameterized by θ and q(φ) = I[φ-φ ] parameterized by φ . This simplifies the objective, reducing to

L = M m=1 N n=1 E q(wm)q(zn) log p(x nm | w m , z n , φ ) + N n=1 E q(zn) log p(y n | x n,1:M , z n , θ ) + M m=1 E q(wm) log p(w m ) -log q(w m ) + N n=1 E q(zn) log p(z n ) -log q(z n ) + log p(φ ) + log p(θ ).

Each expectation can be unbiasedly estimated with Monte Carlo. For gradient-based optimization, we use reparameterization gradients (Rezende et al., 2014) . We describe them next.

C.2 FIRST STAGE: LEARNING THE CONFOUNDER:

We provide details for the gradients. Let λ wm parameterize q(w m ; λ wm ) and λ zn parameterize q(z n ; λ zn ). We are interested in training the parameters λ wm , λ zn , θ , φ . Subsample a SNP location m ∈ {1, . . . , M }. Draw a sample z n ∼ q(z n ; λ zn ) for n = 1, . . . , N and w m ∼ q(w m ; λ wm ) for m = 1, . . . , M , where the samples are reparameterizable (see Rezende et al. ( 2014) for details).

The gradient with respect to parameters λ zn is unbiasedly estimated by

∇ λz n L ≈ ∇ λz n log p(x nm | w m , z n , φ ) + log p(y n | x n,1:M , z n , θ ) + log p(z n ) -log q(z n ) .

This gradient scales linearly with the number of SNPs M . This is undesirable as the number of SNPs ranges from the hundreds of thousands to millions. We prevent this scaling by observing that for large M , the information in x n,1:M will influence the posterior far more than the single scalar y n . In math, p(z n | x n,1:M , y n , θ, φ, w 1:M ) ≈ p(z n | x n,1:M , θ, φ, w 1:M ). This is a tacit assumption in all GWAS methods that adjust for the confounder (Yu et al., 2005; Price et al., 2006; Astle & Balding, 2009; Kang et al., 2010) .

The gradient with respect to parameters λ zn simplifies to

∇ λz n L ≈ ∇ λz n log p(x nm | w m , z n , φ ) + log p(z n ) -log q(z n ) ,

which scales to massive numbers of SNPs.

The gradients with respect to parameters λ wm and φ are unbiasedly estimated by

∇ λw m L ≈ N n=1 ∇ λw m log p(x nm | w m , z n , φ ) + ∇ λw m log p(w m ) -∇ λw m log q(w m ), ∇ φ L ≈ N n=1 ∇ φ log p(x nm | w m , z n , φ ).

Note how none of this depends on the trait y or trait parameters θ. We can thus perform inference to first approximate the posterior p(z 1:N , w 1:M , φ | x, y). In a second stage, we can then perform inference to approximate the posterior p(θ | z 1:N , w 1:M , x, y). The computational savings is significant not only within task but across tasks: when modelling many traits of interest (for example, § 5.2), inference over the SNP confounders only needs to be done once and can be re-used. We perform stochastic gradient ascent using these gradients to maximize the variational objective.

C.3 SECOND STAGE: LEARNING THE TRAIT Above we described the first stage of an algorithm which performs stochastic gradient ascent to optimize parameters so that q(z n )q(w 1:M )q(θ ) ≈ p(z 1:N , w 1:M , φ | x, y). Given these parameters, we are interested in training θ . Dropping constants with respect to θ in the objective, we have

L ∝ N n=1

E q(zn) log p(y n | x n,1:M , z n , θ ) + log p(θ ).

We maximize this objective using stochastic gradients with a single sample z n ∼ q(z n ),

∇ θ L ≈ N n=1 ∇ θ log p(y n | x n,1:M , z n , θ ) + ∇ θ log p(θ )

This corresponds to Monte Carlo EM. Its primary computation per-iteration is the backward pass of the trait's neural network. Unlike in the first stage, we do not subsample SNPs as the likelihood depends on all SNPs.

C.4 SECOND STAGE: HANDLING LIKELIHOOD-FREE TRAITS:

In general, the density of y n is intractable: we exploit its tractable density if y n is discrete (it induces a categorical likelihood); otherwise for real-valued traits, we perform likelihood-free inference with respect to y n . Following LFVI (Tran et al., 2017) , define q(y) to be the empirical distribution over observed data {y n }. Then subtract it as a constant to the objective, so

L ∝ N n=1

E q(zn) log p(y n | x n,1:M , z n , θ ) -log q(y n ) + log p(θ ).

We approximate this log-ratio with a ratio estimator, r(y n , x n,1:M , z n , θ ; λ r ). It is a function of all inputs in the log-ratio and is parameterized by λ r .

We train the ratio estimator by minimizing a loss function with respect to its parameters,

E p(yn | x n,1:M ,zn,θ ) [-log σ(r(y n , x n,1:M , z n , θ ; λ r ))]+E q(yn) [-log(1-σ(r(y n , x n,1:M , z n , θ ; λ r )))].

The global minima of this objective with respect to the ratio estimator is the desired log-ratio, r * (y n , x n,1:M , z n , θ ) = log p(y n | x n,1:M , z n , θ ) -log q(y n ).

Unfortunately, the ratio estimator has inputs of many dimensions. In particular, it has the problematic property of scaling with the number of SNPs, which can be on the order of hundreds of thousands.

We can efficiently parameterize the ratio estimator by studying two extreme cases with respect to computational efficiency and statistical efficiency. In one extreme, suppose y n has a tractable Gaussian density with mean given by the outcome model's neural network and unit variance (that is, the neural net is parameterized to apply a location-shift on the noise input, y n = NN(•) + n ). Up to additive and multiplicative constants, the optimal log-ratio is

r * (y n , x n,1:M , z n , θ ) ∝ (y n -NN([x n,1:M , z n ] | θ )) 2 .

This implies the ratio estimator must relearn the neural network's forward pass in order to estimate the optimal log-ratio. This is computationally redundant and can lead to unstable training. On another extreme, suppose we parameterize r as

r(y n , x n,1:M , z n , θ ; λ r ) = r(y n , NN([x n,1:M , z n , ] | θ ); λ r ).

This dramatically reduces r's input dimensions, from hundreds of thousands to just two. However, while computationally efficient, this is a poor statistical approximation: there is only a single dimension to preserve information about x n,1:M , z n , n , and θ relevant to y n ; this dimension is lossy for even Gaussian densities.

As a middleground, we use the neural net's first hidden layer as input into the ratio estimator, r(y n , x n,1:M , z n , θ ; λ r ) = r(y n , h n ; λ r ).

This reduces the ratio estimator's inputs to be the trait y n and first hidden layer's units h n . This hidden layer has much fewer dimensions than the raw inputs, such as 32 units (making it computationally efficient). Moreover, under the data processing inequality (Cover & Thomas, 1991) , h n preserves more information relevant to y n than subsequent layers of the neural network (making it statistically efficient). For all experiments, we parameterized r with two fully connected hidden layers with equal number of hidden units.

The gradient with respect to parameters θ is estimated by

∇ θ L ≈ N n=1 ∇ θ r(y n , x n,1:M , z n , θ ; λ r ) + log p(θ ).

This substitutes in the ratio estimator as a proxy to the intractable likelihood.

We minimize the auxiliary loss function in order to train the ratio estimator. Sample y n ∼ p(y n | x n,1:M , z n , θ ) and subsample a data point y n ∼ q(y n ). The gradient is estimated by

∇ λr • ≈ ∇ λr -log σ(r(y n , h n ; λ r )) -log(1 -σ(r(y n , h n ; λ r ))) .

This corresponds to maximum likelihood, balanced with an adversarial objective to estimate the likelihood, and is relatively fast. We perform stochastic gradient ascent, alternating between these two sets of gradients.

D SIMULATION STUDY:

We provide more detail to § 5.1. Implicit causal models can not only represent many causal structures but, more importantly, learn them from data. To demonstrate this, we simulate data from a comprehensive collection of popular models in GWAS and analyze how well the fitted model can capture them. These configurations exactly follow Hao et al. (2016) with same hyperparameters, which we describe below.

For each of the 11 simulation configurations, we generate 100 independent data sets. Each data set consists of a M × N matrix of genotypes X and vector of N traits y. Each individual n has M SNPs and one trait.

D.1 GENOTYPE MATRIX:

To simulate the M × N matrix of genotypes X, we draw x mn ∼ Binomial(2, π mn ) for m = 1, . . . , M SNPs and n = 1, . . . , N individuals. The probabilities π mn can be encoded under a real-valued M × N matrix of allele frequencies F where π mn = [logit(F )] mn .

Many models in GWAS can be described under the factorization F = ΓS, where Γ is a M ×K matrix and S is a K ×N matrix for a fixed rank K ≤ N . This includes principal components analysis (Price et al., 2006) , the Balding-Nichols model (Balding & Nichols, 1995) , and the Pritchard-Stephens-Donnelly model (Pritchard et al., 2000) . The M × K matrix Γ describes how structure manifests in the allele frequencies across SNPs. The K × N matrix S encapsulates the genetic population structure across individuals.

We describe how we form Γ and S for each of the 11 simulation configurations.

Balding-Nichols Model (BN) + HapMap. The BN model describes individuals according to a discrete mixture of ancestral subpopulations (Balding & Nichols, 1995) . The HapMap data set was collected from three discrete populations (Gibbs et al., 2003) 1000 Genomes Project (TGP). TGP is a project that comprehensively catalogs human genetic variation by producing complete genome sequences of well over 1000 individuals of diverse ancestries (Consortium et al., 2010) . To form Γ, we sample Γ mk ∼ 0.9 Uniform(0, 1/2) for k = 1, 2 and set Γ m3 = 0.05. To form S, we compute the first two principal components of the TGP genotype matrix after mean centering each SNP. We then transform each principal component to be between (0, 1) and set the first two rows of S to be the transformed principal components. The third row of S is set to 1 as an intercept. The simulated data has M = 100, 000 SNPs and N = 1500 individuals (the total number of individuals in the TGP data set).

Human Genome Diversity Project (HGDP). HGDP is an international project that has genotyped a large collection of DNA samples from individuals distributed around the world, aiming to assess worldwide genetic diversity at the genomic level (Rosenberg et al., 2002) . We followed the same scheme as for TGP above. The simulated data has M = 100, 000 SNPs and N = 940 individuals (the total number of individuals in the HGDP data set).

Pritchard-Stephens-Donnelly (PSD) + HGDP. The PSD model describes individuals according to an admixture of ancestral subpopulations (Pritchard et al., 2000) . The rows of Γ are drawn from three i.i.d. draws from the Balding-Nichols model:

Γ mk ∼ BN(p m , F m ),

where k ∈ {1, 2, 3}. The pairs (p m , F m ) are computed by randomly selecting a SNP in the HGDP data set, calculating its observed allele frequency, and estimating its F ST value using the estimator of Weir & Cockerham (1984) . The estimator requires each individual to be assigned to a subpopulation, which are made according to the K = 5 subpopulations from the analysis in Rosenberg et al. (2002) . The columns of S are sampled (s 1n , s 2n , s 3n ) ∼ Dirichlet(α = (a, a, a)) for n = 1, . . . , N . We apply four PSD configurations with hyperparameter settings of a = 0.01, 0.1, 0.5, 1. Varying a from 1 to 0 varies the level of sparsity as individuals are placed from uniformly to corners of the simplex. The simulated data has M = 100, 000 SNPs and N = 5000 individuals.

Spatial. In this setting, we simulate genotypes such that the population structure relates to the spatial position of individuals. The matrix Γ is populated by sampling Γ mk ∼ 0.9 Uniform(0, 1/2) for k = 1, 2 and setting Γ m3 = 0.05. The first two rows of S correspond to coordinates for each individual on the unit square and are set to be independent and identically distributed samples from Beta(a, a), while the third row of S is set to 1 as an intercept. We apply four spatial configurations with hyperparameter settings of a = 0.01, 0.1, 0.5, 1. As with the Dirichlet distribution in the PSD model, varying a from 1 to 0 varies the level of sparsity as individuals are placed from uniformly to corners of the unit square. The simulated data has M = 100, 000 SNPs and N = 5000 individuals.

D.2 TRAITS OF INTEREST:

To simulate traits y, we simulate from a linear model: for each individual n's SNPs {x mn },

y n = M m=1 β m x mn + λ n + n , n ∼ Normal(0, σ 2 n ).

Each trait is real-valued and is determined by a linear combination of SNPs, a per-individual offset, and heteroskedastic noise. Below we describe how we set {β m }, {λ n }, and {σ n }.

Without loss of generality, we set the first 10 SNPs to be true alternative SNPs (β m = 0), where β m ∼ Normal(0, 0.5) for m = 1, 2, . . . , 10; β m = 0 for m > 10. In order to simulate λ j and j so that they are also influenced by the latent variables z 1 , . . . , z n , we performed the following:

1. Run K-means clustering on the columns of S with K = 3 using Euclidean distance. This assigns each individual j to one of three partitions S 1 , S 2 , S 3 where S k ⊂ 1, 2, . . . , n. 2. Set λ j = k for all j ∈ S k for each k = 1, 2, 3.

3.. Draw τ 2:

1 , τ2 2 , τ 2 3 ∼ InverseGamma(3, 1). Set σ j = τ k for all j ∈ S k .

Following Song et al. (2015) ,

This strategy simulates non-genetic effects and random variation that manifest among K discrete groups over a more continuous population genetic structure defined by S. This is meant to emulate the fact that environment (specifically lifestyle) may partition among individuals in a manner distinct from, but highly related to population structure.

We apply this strategy for each of the 11 configurations where each involved up to M = 100, 000 SNPs and N = 5000 individuals. For each configuration, we simulated 100 independent data sets, thus requiring a total of 1100 fitted models per method of comparison.

D.3 OTHER DETAILS:

For testing, we obtain one p-value for each SNP, which is the probability that the SNP's effect on the trait of interest is zero. To do so, we apply a likelihood ratio test statistic for each SNP m, T (x, y) = 2 max where θ m corresponds to weights connecting x m to hidden units in the trait neural networks' first layer. When θ m = 0 in the model, following Wilks' theorem, the null distribution of this test statistic is χ 2 H where H is the number of first layer hidden units, i.e., the test statistic converges to this distribution as the number of individuals N → ∞. This statistic and null distribution is a simple extension of Song et al. (2015) , which applies a linear model. (See also their proof for the null distribution.) 2

In summary, the testing procedure works as follows:

1. Formulate and estimate a model of population structure that provides well-behaved estimates of the posterior p(z | x, y). 2. First, perform a regression of the trait given SNPs and estimated posterior p(z | x, y). Second, for each SNP, perform a regression of the trait given SNPs and estimated posterior p(z | x, y), but with θ m fixed at 0.

3. Calculate a p-value for each SNP. Use the test statistic above, which follows a χ 2 H null distribution for large sample sizes.

:

We also measured recall. All true-positives were found across all methods (like a real-world experiment, only few (10) SNPs are causal). Rarely did a method deviate and if so, it missed at most one.

y = trait_neural_network(z, x)

The test only holds approximately for two reasons. First, it assumes asymptotically large sample size (as all frequentist tests). Second, given a multi-layered neural network, there exist permutations of the weights and biases beyond the first layer weights which set the effect of xm on y to 0.

Figures and Tables:

Figure 1: Probabilistic causal model. (left) Variable x causes y coupled with a shared variable β. (right) A more explicit diagram where variables (denoted with a square) are a deterministic function of other variables and noise (denoted with a triangle).

Figure 2 (right) visualizes the model. It is a model over the full causal graph (Figure 2 (left)) and differs from the unconfounded case: Equation2only requires a model from X → Y , and the rest of the graph is "ignorable"(Imbens & Rubin, 2015).

.tile(tf.reshape(z, [N, 1, K]), [1, M, 1]) 11 w_tile = tf.tile(tf.reshape(w, [1, M, K]), [N, 1, 1]) 12 h = tf.concat([z_tile, w_tile], 2) 13 h = tf.layers.dense(h, 512, activation=tf.nn.relu) 14 h = tf.layers.dense(h, 512, activation=tf.nn.relu) 15 h = tf.layers.dense(h, 1, activation=None) .layers.dense(h, 32, activation=tf.nn.relu) 22 h = tf.layers.dense(h, 256, activation=tf.nn.relu) 23 h = tf.concat([z, h], 1) # include connection to z for output layer 24 h = tf.layers.dense(h, 1, activation=None) Binomial(total_count=2.0, logits=logits) 31 32 # Note: Sometimes it's preferable to have neural net parameterize a 33 # distribution for tractable density such as for binary-valued traits. 34 # To do this, specify, e.g., 'y = Bernoulli(logits=trait_neural_network(z, x))' C LIKELIHOOD-FREE VARIATIONAL INFERENCE FOR GWAS C.1 VARIATIONAL OBJECTIVE

log p(y | x; θ) -max θ-m log p(y | x; θ -m , θ m = 0)

Precision accuracy over an extensive set of configurations and methods; we average over 100 simulations for a grand total of 4,400 fitted models. The setting a in PSD and Spatial determines the amount of sparsity in the latent population structure: lower a means higher sparsity. ICM is significantly more robust to spurious associations, outperforming other methods by up to 45.3%.

Table content: <table><tr><td>Trait</td><td colspan="4">ICM PCA (Price+06) LMM (Kang+10) GCAT (Song+10)</td></tr><tr><td>HapMap</td><td>99.2</td><td>34.8</td><td>30.7</td><td>99.2</td></tr><tr><td>TGP</td><td>85.6</td><td>2.7</td><td>43.3</td><td>70.3</td></tr><tr><td>HGDP</td><td>91.8</td><td>6.8</td><td>40.2</td><td>72.3</td></tr><tr><td>PSD (a = 1)</td><td>97.0</td><td>80.4</td><td>92.3</td><td>95.3</td></tr><tr><td>PSD (a = 0.5)</td><td>94.3</td><td>79.5</td><td>90.1</td><td>93.6</td></tr><tr><td>PSD (a = 0.1)</td><td>92.2</td><td>38.1</td><td>38.6</td><td>90.4</td></tr><tr><td>PSD (a = 0.01)</td><td>92.7</td><td>24.2</td><td>35.1</td><td>90.7</td></tr><tr><td>Spatial (a = 1)</td><td>90.9</td><td>56.4</td><td>60.0</td><td>75.2</td></tr><tr><td>Spatial (a = 0.5)</td><td>86.2</td><td>50.5</td><td>46.6</td><td>72.5</td></tr><tr><td>Spatial (a = 0.1)</td><td>80.9</td><td>2.4</td><td>26.6</td><td>35.6</td></tr><tr><td colspan="2">Spatial (a = 0.01) 75.5</td><td>1.8</td><td>15.3</td><td>30.2</td></tr></table>

Number of significant loci at genome-wide significance (p < 7.2 × 10 -8 ) for each of the ten traits from NFBC data. The counts for GCAT are obtained from Song et al. (2015, Table 1); counts for LMM, PCA, and uncorrected are obtained from Kang et al. (2010, Table 2). The implicit causal model (ICM) captures causal relationships comparable to previously work.

Table content: <table><tr><td>Trait</td><td colspan="5">ICM GCAT LMM PCA Uncorrected</td></tr><tr><td>Body mass index</td><td>0</td><td>0</td><td>0</td><td>0</td><td>0</td></tr><tr><td>C-reactive protein</td><td>2</td><td>2</td><td>2</td><td>2</td><td>2</td></tr><tr><td colspan="2">Diastolic blood pressure 0</td><td>0</td><td>0</td><td>0</td><td>0</td></tr><tr><td>Glucose levels</td><td>3</td><td>3</td><td>2</td><td>2</td><td>2</td></tr><tr><td>HDL cholesterol levels</td><td>4</td><td>4</td><td>4</td><td>2</td><td>4</td></tr><tr><td>Height</td><td>1</td><td>1</td><td>0</td><td>0</td><td>0</td></tr><tr><td>Insulin levels</td><td>0</td><td>0</td><td>0</td><td>0</td><td>0</td></tr><tr><td>LDL cholesterol levels</td><td>3</td><td>4</td><td>3</td><td>3</td><td>3</td></tr><tr><td>Systolic blood pressure</td><td>0</td><td>0</td><td>0</td><td>0</td><td>0</td></tr><tr><td>Triglyceride levels</td><td>2</td><td>2</td><td>3</td><td>2</td><td>2</td></tr></table>

, which allows us to populate each row m of Γ with three i.i.d. draws from the Balding-Nichols model:Γ mk ∼ BN(p m , F m ),where k ∈ {1, 2, 3}. Each Γ mk is interpreted to be the allele frequency for subpopulation k at SNP m. The pairs (p m , F m ) are computed by randomly selecting a SNP in the HapMap data set, calculating its observed allele frequency, and estimating its F ST value using the estimator ofWeir & Cockerham (1984). The columns of S are populated with indicator vectors such that each individual is assigned to one of the three subpopulations. The subpopulation assignments are drawn independently with probabilities 60/210, 60/210, and 90/210, which reflect the subpopulation proportions in the HapMap data set. The simulated data has M = 100, 000 SNPs and N = 5000 individuals.

Table content: <table/>

Peer review (iclr2018):
 The paper presents a non-linear generative model for GWAS that models population structure. Non-linearities are modeled using neural networks as non-linear function approximators and inference is performed using likelihood-free variational inference. The paper is overall well-written and makes new and non-trivial contributions to model inference and the application. Stated contributions are that the model captures causal relationships, models highly non-linear interactions between causes and accounts for confounders. However, not all claims are well-supported by the data provided in the paper. Especially, the aspect of causality does not seem to be considered in the application beyond a simple dependence test between SNPs and phenotypes. The paper also suffers from unconvincing experimental validation: - The evaluation metric for simulations based on precision is not meaningful without reporting the recall at the same time. - The details on how significance in each experiment has been determined are not sufficient. From the description in D.3 the p-value a p-value threshold of 0.0025 has been applied. Has this threshold been used for all methods? The description in D.3 seems to describe a posterior probability of the weight being zero, instead of a Frequentist p-value, which would be the probability of estimating a parameter at least as large on a data set that had been generated with a 0-weight. - Genomic control is applied in the real world experiment but not on the simulations. Genomic control changes the acceptance threshold of each method in a different way. Both precision and recall depend on this acceptance threshold. Genomic control is a heuristic that adjusts for being too anti-conservative, but also for being too conservative, making it hard to judge the performance of each method on its own. Consequently, the paper should provide additional detail on the results and should contrast the performance of the method without the use of genomic control. minor: The authors claim to model nonlinear, learnable gene-gene and gene-population interactions. While neural networks may approximate highly non-linear functions, it still seems as if the confounders are modeled largely as linear. This is indicated by the fact that the authors report performance gains from adding the confounders as input to the final layer. The two step approach to confounder correction is compared to PCA and LMMs, which are stated to first estimate confounders and then use them for testing. For LMMs this is not really true, as LMMs treat the confounder as a latent variable throughout and only estimate the induced covariance.