LDAK-KVIK performs fast and powerful mixed-model association analysis of quantitative and binary phenotypes

Notation

Suppose there are n individuals, each genotyped for m SNPs, recorded for q covariates and measured for a phenotype. Let the (n x m) matrix X’ contain the genotypes, let the length-n vector Y’ contain the phenotypes, and let the (n x q) matrix Z contain covariates. We use C to denote the total number of chromosomes, use X and Y to denote, respectively, the genotypes and phenotypes after regressing out the covariates, and use λ to denote the test statistic scaling factor. Without loss of generality, we assume X_j and Y are standardized to have mean zero and variance one.

LDAK-KVIK Step 1

This step produces C LOCO Elastic Net PRS, each taking the form , where is the estimated effect size for SNP j, and the sum is across all SNPs not on Chromosome c. This step also estimates λ.

LDAK-KVIK first performs our novel test for structure (described below); for now, we assume this test finds only weak structure. LDAK-KVIK next computes , estimates of the heritability contributed by each SNP. For this, it considers heritability models of the form where f_j is the MAF of SNP j, h² is the proportion of phenotypic variance explained by all SNPs, while α determines how per-SNP heritability depends on MAF. LDAK estimates α and h² using Randomized Haseman-Elston Regression³³ and Monte Carlo restricted maximum likelihood¹² (REML), then sets to its expected value given these estimates. Note that Monte Carlo REML must calculate terms of the form V^-1Y, KV^-1Y, V^-1r and KV^-1r, where V is an (n x n) variance matrix, K is an (n x n) genomic relatedness matrix (GRM), and r is a length-n vector whose elements are drawn from a standard normal distribution; we calculate these terms using our novel Variational Bayes solver (described below).

LDAK-KVIK then constructs C LOCO Elastic Net PRS, for which it assumes²⁶ where DE(a_j) denotes a double exponential distribution with rate a_j. The parameters p and F determine the relative contributions of the double exponential and normal distributions. LDAK-KVIK finds suitable values for these using cross-validation (by default, LDAK-KVIK uses 90% of samples to construct genome-wide PRS corresponding to ten different pairs of values for p and F, then picks the pair that has lowest mean-squared error measured using the remaining 10% of samples). Note that all PRS are constructed using our novel variational

Bayes solver (described below). Finally, LDAK-KVIK calculates λ. Because we are, for now, assuming there is only weak structure, LDAK-KVIK sets λ=1.

LDAK-KVIK Step 2

LDAK-KVIK tests SNP j for association using least-squares linear regression with the model E [Y −P_c]= X_j β_j, then scales the resulting χ²(1) test statistic by λ.

Binary phenotypes

Many of the operations in LDAK-KVIK are based on a linear model of the form Y _i=∑ X_{i, j} γ_j +e_i. In particular, this model is assumed (either explicitly or implicitly) when estimating h², when constructing PRS and when testing SNPs for association. When the phenotype is quantitative, X_j and Y contain standardized residuals from linearly regressing X_j ‘ and Y’, respectively, on Z, and LDAK-KVIK assumes , where I is an (n x n) identity matrix and is the estimate of h². When the phenotype is binary, LDAK-KVIK first computes the length-n vector μ’, which contains estimates of the probabilities that each individual is a case given the covariates, and constructs the (n x n) diagonal matrix D, with D_i,i∝μ ′ _i(1−μ’_i) and trace(D^-1)=n. LDAK-KVIK then sets X_j to the residual from regressing X_j’ on Z using weighted linear regression with weight matrix D, sets Y =D⁻¹(Y’ −μ’), and assumes . These changes are motivated by the observation that the estimated SNP effect sizes from regressing Y’ on Z and X_j’ using logistic regression are approximately equal to those from regressing D⁻¹(Y ‘ −μ’) on Z and X_j’ using weighted linear regression with weight matrix D (see Supplementary Note 5 for a proof). Note that ensuring trace(D^-1)=n allows us to continue to treat h² as a heritability (i.e., h² continues to represent the proportion of variance of Y explained by all SNPs).

When analyzing quantitative phenotypes, LDAK-KVIK obtains the (unscaled) Step 2 test statistics via a Wald Test (specifically, , the square of the estimated effect size for SNP j divided by its estimated variance). When analyzing binary phenotypes, LDAK-KVIK first obtains U_j via a Wald Test, but if the corresponding p-value (after scaling by λ) is below 0.1, recomputes U_j using our novel SPA solver (described below).

Novel test for structure

LDAK-KVIK picks 512 SNPs semi-randomly from across the genome (specifically, LDAK-KVIK first randomly picks 10,000 SNPs, then retains the 512 SNPs with highest variance). It then computes , the average squared correlation between the r pairs of SNPs on different chromosomes. To test for significance, we use the fact that under the null hypothesis (i.e., if the data are homogeneous and so SNPs on different chromosomes are independent), the squared correlations are beta distributed with parameters 1/2 and (n-2)/2, and so will have expectation 1/(n-1) and variance 2(n-2)/[r(n-1)²(n+1)].

As well as using to test for structure, we consider an estimate of the maximum average inflation of test statistics due to structure, which is based on the following logic. Suppose that we could partition the genetic contribution to a phenotype as G=L_j+ D_j where L_j and D_j are the components local and distal to SNP j, respectively, whose heritability contributions are h²_Lj and h²_Dj, respectively. Then when testing SNP j for association with the phenotype, we could write its expected χ²(1) test statistic from classical linear regression as , where and are the proportions of local and distal genetic variation tagged by SNP j, respectively. Under this partitioning, can be viewed as the expected inflation of S_j due to structure. Finally, if we assume that is a reasonable estimator of , the average value of d² across all SNPs, and recognize that , then it follows that is an upper bound for the average inflation of test statistics due to structure.

LDAK-KVIK determines there is strong structure when and the corresponding p-value is below 0.001 (otherwise, it determines the structure is weak). Our test for structure is very fast (e.g., when analyzing 368k individuals, it takes less than one minute) and has low memory demands (because it is only necessary to store genotypes for 512 SNPs). When LDAK-KVIK determines there is strong structure, it makes the following changes. When analyzing a quantitative phenotype, it switches from the elastic net prior to the infinitesimal prior , then estimates λ using the Grammar-Gamma formula¹¹. When analyzing a binary phenotype, LDAK-KVIK uses the approximate version of fastGWA described below.

Approximate version of fastGWA

The LDAK-KVIK version of fastGWA differs from the original version in three ways. Firstly, whereas fastGWA constructs a GRM using all SNPs, our version uses only 512 SNPs (those used in the test for structure, described above). Secondly, when constructing the sparse GRM, fastGWA recommends truncating values below 0.05, whereas we truncate non-significant values (specifically, those with P>0.1/n, which typically corresponds to values below 0.2). Thirdly, when analyzing a binary phenotype, we assume the null model Y ∼N (0, σ² K + E⁻¹), where K is the sparse GRM and E is a diagonal matrix with E_i,i =μ’_i(1−μ’_i). Note that this represents a simplified version of the quasi-likelihood used by fastGWA, because we set μ’, and therefore also E, based only on the covariates, whereas fastGWA updates μ’ to allow for the contribution of the sparse GRM. Supplementary Figure 18 shows that despite these simplifications, our version of fastGWA performs similarly to the original version, both for quantitative and binary phenotypes.

LDAK-KVIK-GBAT

Previously, we developed LDAK-GBAT, a tool for gene-based association analysis.¹⁹ This tests each gene using the mixed model , where K_S is a local GRM computed using only SNPs within the gene. LDAK-GBAT finds the maximum likelihood estimate of , then obtains a p-value by testing if . Importantly, LDAK-GBAT requires only GWAS summary statistics and a reference panel. Therefore, LDAK-KVIK-GBAT applies LDAK-GBAT using the association results from Step 2, and using 5000 randomly-picked individuals from the data as an in-sample reference panel.

Overview of Variational Bayes

This section explains the general idea of using variational Bayes to construct PRS, whereas the next section provides specific details of our novel variational Bayes solver. Please note that a detailed description of variational Bayes is provided in the supplement of the BOLT-LMM paper.¹²

When constructing PRS, LDAK-KVIK estimates P(γ|Y), the posterior distribution of SNP effect sizes given the data. To do this, we first construct a model likelihood L(Y|γ) by assuming that Y has the multivariate normal distribution , where I is an (n x n) identity matrix. Then we specify π(γ), a prior distribution for SNP effect sizes, and use variational Bayes to approximate P(γ|Y) as the product of independent, single-parameter posterior distributions (one for each SNP):

LDAK-KVIK updates Q (γ) one SNP at a time, by replacing the current Q_j(γ_j) with the distribution that minimizes the difference between the approximate and true log likelihoods (measured by the Kullback-Leibler divergence). Once convergence has been achieved, the corresponding PRS is constructed by setting γ_j=η_j, where η_j is the expectation of Q_j(γ_j).

As explained above, Step 1 of LDAK-KVIK uses variational Bayes to construct Elastic Net PRS, assuming the prior distribution defined in Equation 1. In this application, each Q_j(γ_j) is a mixture of a left truncated normal distribution, a right truncated normal distribution, and a (non-truncated) normal distribution:

If denotes a partial PRS where the effect sizes are based on the current estimate of Q(γ), then Q_j(γ_j) is updated by setting

Novel variational Bayes solver

The variational Bayes solver implemented in BOLT-LMM uses sequential genome-wide scans.¹² This means that on each scan, it updates Q_j(γ_j) once for each SNP in the genome (i.e., it first updates Q₁(γ₁), then Q₂(γ₂), and so on until Q_m(γ_m)). BOLT-LMM performs multiple scans, continuing until LL, the approximate log likelihood, changes by less than a specified tolerance (by default, 0.0005). Typically, BOLT-LMM requires 50-150 scans (convergence is slowest when analyzing highly-heritable traits with very large sample sizes).

By contrast, our novel variational Bayes solver partitions the genome into chunks (by default, each chunk contains 256 SNPs), then uses chunk-based scans (see Supplementary Figure 1 for an illustration). On Scan 1, our solver first updates Q_j(γ_j) for SNPs in Chunk 1, continuing until LL has converged (the default tolerance is n/10⁶). Our solver then updates Q_j(γ_j) for SNPs in Chunk 2, then for SNPs in Chunk 3, continuing until it reaches the final chunk in the genome. On subsequent scans, our solver repeats the same process, except that it only considers chunks that had a sizeable impact on LL (specifically, it only revisits a chunk if on the previous scan, the updates for that chunk caused LL to change by more than n/10⁶). Our solver typically requires no more than ten scans to converge (because at this point, no chunks remain that have a sizeable impact on LL).

Our chunk-based solver has three main advantages over the genome-wide solver. Firstly, it prioritizes SNPs in regions that have a larger influence on LL (which is more efficient than simply treating all SNPs the same). Secondly, it enables on-the-fly reading of the genotypes (while this is, in theory, possible with the genome-wide solver, it would be necessary to read the data 50-150 times). Thirdly, it is more cache-friendly. Supplementary Figure 2 compares our chunk-based solver with a genome-wide version when analyzing either 50k or 100k individuals. We see that both solvers have very similar accuracy (in that both produce models with very similar LL). However, we find that the chunk-based solver is substantially faster, reflecting that it performs fewer scans, and in turn fewer updates of Q_j(γ_j). For example, when analyzing 100k individuals, the genome-wide solver on average requires 56 scans (so updates each Q_j(γ_j) 56 times). By contrast, our chunk-based solver requires on average 4 scans, and on average updates each Q_j(γ_j) 11 times. Supplementary Figure 22 shows that the default convergence criterion suffices, in the sense that results are almost unchanged if the tolerance is made five times smaller.

Computing terms of the form V^-1A and KV^-1A

Our variational Bayes solver can be used to construct Ridge Regression PRS by assuming the prior distribution . In this application, each Q_j(γ_j) has the form , and is updated by setting

However, for this choice of prior, the posterior mean has an explicit form (derived in Supplementary Note 4). In particular, it can be shown that if P is a (genome-wide or LOCO) Ridge Regression PRS, then and where V is a (genome-wide or LOCO) variance matrix. It follows that we can estimate terms of the form V^-1Y and KV^-1Y by using our variational Bayes solver to construct Ridge Regression PRS for Y, then dividing the estimated PRS by or dividing the corresponding residuals by . More generally, we can estimate terms of the form V^-1A and KV^-1A, where A is an arbitrary length-n vector, by performing the same steps but replacing Y with A (i.e., instead of constructing a PRS for Y, we construct a PRS for A). Therefore, our novel variational Bayes solver can not only be used to construct PRS, but also to compute the terms required when performing either Monte Carlo REML or using the Grammar-Gamma formula.¹¹

Overview of SPA

Here we summarize the SPA; for a fuller description, we recommend reading the supplement of the REGENIE paper.¹³ If A is a random variable, then its cumulant-generating function (CGF) is K_A(t)=log(E[exp(tA)]).

The first, second and third derivatives of K_A(t) are where m (t)=E[A^k exp(tA)]. The CGF is additive, such that if A and B are independent random variables, then K_aA+bB(t)=K_A(at)+ K_B(bt). Similar relationships hold for the first, second and third derivatives:

Now suppose A arises as a test statistic from regressing a phenotype on X_j. The most common way to calculate a p-value is to compute U_j=(A−E [A])²/ Var(A), where E[A] and Var(A) are, respectively, the expectation and estimated variance of A under the null hypothesis, then to compare U_j to a χ²(1) distribution (or equivalently, compare its square root to a standard normal distribution). However, this approach assumes the null distribution of A is approximately normal, which can be inappropriate (e.g., if the phenotype is binary and very imbalanced, or when testing rare variants).

The SPA provides an alternative way to compute a p-value. It involves computing where t’ is the solution to K¹(t’)=A. The SPA p-value is then obtained by comparing U’_j to a χ²(1) distribution.

Novel SPA solver

Suppose the test statistic from regressing a phenotype B on X_j takes the form A=a₁B₁ + … + a_nB_n. Our solver starts by calculating K_B(at)=log (∑ exp(atB_i) /n) for 41 predetermined values of a (evenly spaced between - 2 and 2), and for 256 predetermined values of t (ranging from t_min to t_max, as described below). Note that these calculations correspond to assuming that the B_i are independent and identically distributed, and that their true distribution matches the observed distribution. Our solver similarly computes 41 × 256 realizations for each of K¹_B(at), K²_B(at) and K³_B(at). We refer to the predetermined values of a and t as “bin centres” and “knots”, respectively.

When LDAK-KVIK analyzes a quantitative phenotype, the (uncalibrated) Step 2 test statistic for SNP j is , so we set a_i=X_ij and B_i=Y_i-P_ci in the above equations; when LDAK-KVIK analyzes a binary phenotype, the test statistic is , so we set a_i=X_ij and B_i=D_ii(Y_i-P_ci). Note that when calculating the CDF of A and its derivatives, it is necessary to calculate terms of the form K_B(a_it), K¹_B(a_it) and K² _B(a_it), and in general, a_i will not match one of the 41 bin centres. Therefore, we approximate these terms using first, second and third order Taylor Series. Specifically, if a’ denotes the bin centre closest to a _i, then we use the three approximations

Our solver uses the second approximation to estimate K¹_A(t) = a₁ K¹_B(a t) + … + a_nK¹_B(a_nt) for each of the 256 knots. It then identifies t_L and t_R, the knots immediately left and right of the solution to K¹_A(t)=A, and uses the first and last approximations to calculate K_A(t_L), K_A(t_R), K²_A(t) and K²_A(t_R). Lastly, our solver uses linear interpolation to estimate t’, K_A(t’) and K² _A(t’). Specifically, it computes , which estimates the location of the solution, relative to the closest two knots (e.g., if ε=0.5, then t’ is approximately halfway between t_L and t_R), then sets

Our SPA solver now has all terms required to compute U’_j, the χ²(1) test statistic defined in Equation 2. If our solver performed the above calculations naively, it would likely be no faster than existing solvers, because computing each of the CGF and its derivatives still requires n operations (e.g., to compute K_A(t) for a particular knot, it is necessary compute K_B(a_it) for each individual). However, the number of operations can be reduced dramatically by first summarizing X_j with respect to the bin centres. Specifically, if b_j denotes the jth bin centre, and the function I(i,j) indicates whether a_i is closest to b_j, then our solver computes C₁, …, C₉, nine length-41 count vectors whose elements are

Given these nine vectors, we can rewrite the CDF of A and its derivatives as which can be computed using 164, 123 and 82 operations, respectively (numbers that, for large datasets, are substantially smaller than n).

When deciding the predetermined knot values, we initially use 256 quantiles from a Cauchy distribution, scaled such that and . However, if our solver encounter a SNP whose test statistic is either below K¹_A(t_min) or above K¹ _A(t_max), then it resets the knot values so that t and t are five times larger than their original values (this was not necessary for any of the analyses in this paper).

As explained above, our SPA assumes the true distribution of B, the adjusted phenotype, matches its empirical distribution. This has the advantage that our SPA solver can be (validly) applied to any phenotype, whereas existing SPA solvers generally assume the phenotype has a Bernouilli distribution, so can only be applied to binary phenotypes. Furthermore, our SPA solver makes no requirements on the form of X_j(although the predetermined bin centres assume that all elements of X_j are between -2 and 2, if this is not the case, X_j can simply be rescaled). This means that our solver can automatically be applied to dosage data produced by genotype imputation (or in fact non-SNP data).

Supplementary Figure 23 shows that the default values for the numbers of bin centres and knots suffice, in the sense that results are almost unchanged if we increase either of these values (e.g., increase the number of bin centres from 41 to 100, or the number of knots from 256 to 512).