Model setup

The LMM log likelihood for phenotype data \(\mathbf{\color{darkgreen}{y}}\) (dimension \(n \times 1\)) given fixed effects \(\mathbf{X}\):

\[ \log \mathcal{L} = -\frac{1}{2} \left[ n \log(2\pi) + \log\bigl|\sigma_e^2 \mathbf{I} + \sigma_g^2 \mathbf{\color{darkblue}{K}}\bigr| + \\ \Bigl(\mathbf{\color{darkgreen}{y}} - \mathbf{X}\mathbf{\beta}\Bigr)^\top \Bigl(\sigma_e^2 \mathbf{I} + \sigma_g^2 \mathbf{\color{darkblue}{K}}\Bigr)^{-1} \Bigl(\mathbf{\color{darkgreen}{y}} - \mathbf{X}\mathbf{\beta}\Bigr) \right] \]

• \(\mathbf{\color{darkblue}{K}}\): Genetic similarity matrix (\(n \times n\))
• \(\sigma_e^2\): Residual variance
• \(\sigma_g^2\): Genetic variance, \(\mathbf{\beta}\): Fixed-effect weights (\(d \times 1\))

Log likelihood after factorizing \(\mathbf{\color{darkblue}{K}}\)

Let \(\delta = \sigma_e^2 / \sigma_g^2\) and decompose \(\mathbf{\color{darkblue}{K}} = \mathbf{U} \mathbf{S} \mathbf{U}^\top\)
\[ \log \mathcal{L} = -\frac{1}{2} \left[ n \log(2\pi\sigma_g^2) + \log\bigl|\mathbf{S} + \delta \mathbf{I}\bigr| + \\ \frac{1}{\sigma_g^2} \Bigl(\mathbf{U}^\top\mathbf{\color{darkgreen}{y}} - \mathbf{U}^\top\mathbf{X} \mathbf{\beta}\Bigr)^\top \Bigl(\mathbf{S} + \delta \mathbf{I}\Bigr)^{-1} \Bigl(\mathbf{U}^\top\mathbf{\color{darkgreen}{y}} - \mathbf{U}^\top\mathbf{X}\mathbf{\beta}\Bigr) \right] \]

Simplifications:

1. \(|\sigma_g^2 (\mathbf{S} + \delta \mathbf{I})| = \sigma_g^{2n} |\mathbf{S} + \delta \mathbf{I}|\)
2. \((\mathbf{\color{darkblue}{K}} + \delta \mathbf{I})^{-1} = \mathbf{U} (\mathbf{S} + \delta \mathbf{I})^{-1} \mathbf{U}^\top\)

Simplified log likelihood

Rotate data using \(\mathbf{U}^\top\) and let \(\tilde{\mathbf{\color{darkgreen}{y}}} = \mathbf{U}^\top\mathbf{\color{darkgreen}{y}}\), \(\tilde{\mathbf{X}} = \mathbf{U}^\top\mathbf{X}\):

\[ \log \mathcal{L} = -\frac{1}{2} \left[ n \log(2\pi\sigma_g^2) + \sum_{i=1}^n \log(s_i + \delta) + \\ \frac{1}{\sigma_g^2} \sum_{i=1}^n \frac{\Bigl(\tilde{\color{darkgreen}{y}}_i - \tilde{\mathbf{X}}_{i:} \mathbf{\beta})^2}{s_i + \delta} \right] \]

where \(s_i\) are eigenvalues of \(\mathbf{\color{darkblue}{K}}\).

This reduces to a sum of univariate normal terms:

\[ \prod_{i=1}^n \mathcal{N}\left(\tilde{\color{darkgreen}{y}}_i \,\Big|\, \tilde{\mathbf{X}}_{i:} \boldsymbol{\beta}, \; \sigma_g^2 (s_i + \delta)\right) \]

Optimization Process

1. Solve for \(\mathbf{\beta}(\delta)\)
   Differentiate LL w.r.t. \(\mathbf{\beta}\) and set to zero: \[ \mathbf{\beta}(\delta) = \left(\tilde{\mathbf{X}}^\top (\mathbf{S} + \delta \mathbf{I})^{-1} \tilde{\mathbf{X}}\right)^{-1} \tilde{\mathbf{X}}^\top (\mathbf{S} + \delta \mathbf{I})^{-1} \tilde{\mathbf{\color{darkgreen}{y}}} \]

2. Solve for \(\sigma_g^2(\delta)\)
   Substitute \(\mathbf{\beta}(\delta)\) into LL and differentiate w.r.t. \(\sigma_g^2\): \[ \sigma_g^2(\delta) = \frac{1}{n} \sum_{i=1}^n \frac{\Bigl(\tilde{\color{darkgreen}{y}}_i - \tilde{\mathbf{X}}_{i} \mathbf{\beta}(\delta)\Bigr)^2}{s_i + \delta} \]

3. Optimize \(\delta\)
   Plug \(\mathbf{\beta}(\delta)\) and \(\sigma_g^2(\delta)\) into (2) and use Brent’s method for 1D optimization over \(\delta\).

FaST-LMM TLDR
The “Fa” in FaST-LMM refers to the factorization \(\mathbf{\color{darkblue}{K}} = \mathbf{U}\mathbf{S}\mathbf{U}^\top\), which diagonalizes the covariance matrix and enables computationally efficient \(O(n)\) updates.

However, expensive eigen decomposition \(O(n^3)\) must occur up front.

Fast-GWA with GCTA

Jiang et al., Nature Genetics, 2019

Fast-GWA with GCTA

Jiang et al., Nature Genetics, 2019

Reduce computational complexity with sparse matrices

using SparseArrays
K = simulate_GRM(N, 10, 0.02) # third argument is baseline-relatedness
Ktrunc = copy(K)
Ktrunc[Ktrunc .< 0.05] .= 0.0 # truncate values
spK = SparseMatrixCSC(Ktrunc)
@time z1 = K * x; # 0.015 seconds
@time z2 = spK * x; # 0.0002 seconds

cor(z1, z2) # 0.999

Using a dense GRM matrix is 75x slower than a sparse matrix multiplication!

What is a sparse matrix?

struct SparseMatrix
    m::Int32                   # Number of rows
    n::Int32                   # Number of columns
    colptr::Vector{Int32}      # Column j is in colptr[j]:(colptr[j+1]-1)
    rowval::Vector{Int32}      # Row indices of stored values
    nzval::Vector{Float64}     # Stored values, typically nonzeros
end;

What are the tradeoffs of using a sparse GRM in place of a dense GRM?

Review

Compute burden in LMMs can be reduced by:

Case-control phenotypes

So far - we’ve discussed efficient methods for running LMMs at biobank-scale. How well will these work for case-control phenotypes?

Case-control phenotypes

Figure 1, Zhou et al., Nature Genetics, 2019

Saddle point approximation for case-control phenotypes

1. Logistic Mixed Model

The null model is specified as: \[ \text{logit}\Bigl(\mathbb{E}[\color{darkgreen}{Y_i}]\Bigr) = \mathbf{X}_i^\top \alpha + \mathbf{g}_i^\top \mathbf{u} \]

2. Variance Component Estimation

The variance ratio \(\hat{r}\) is estimated as: \[ \hat{r} = \frac{\sigma_g^2}{\sigma_e^2} \] where \(\sigma_g^2\) is the genetic variance and \(\sigma_e^2\) is the residual variance (fixed at 1 for logistic models).

3. Predicted Means and Weights

• Predicted means under the null model: \[ \hat{\mu}_i = \text{logit}^{-1}\left( \mathbf{X}_i^T \hat{\boldsymbol{\alpha}} + \mathbf{g}_i^T \hat{\mathbf{u}} \right) \]
• Weight matrix \(\hat{\mathbf{W}}\) (diagonal entries): \[ \hat{W}_{ii} = \hat{\mu}_i (1 - \hat{\mu}_i) \]

4. Variance-Adjusted Test Statistic

The adjusted test statistic for Step 2 is: \[ T_{\text{adj}} = \frac{\tilde{\mathbf{G}}^T (\mathbf{\color{darkgreen}{Y}} - \hat{\boldsymbol{\mu}})}{\sqrt{\hat{r} \cdot \tilde{\mathbf{G}}^T \hat{\mathbf{W}} \tilde{\mathbf{G}}}} \]

Computing pvalues

• Under the null, the distribution of the test statistic is assumed to converge in distribution to a gaussian
• However, if the case-control imbalance is particularly large, usual asymptotics don’t quite kick in, making this a poor approximation

Simulation

function gwas_sim(N = 50_000, P = 10_000)
    X = rand(Normal(0, 1), N, 10)
    α = rand(Normal(0, 1), 10)
    μ0 = -7.0
    η = μ0 .+ X * α
    μ = logistic.(η)
    Y = rand.(Bernoulli.(μ)) # 3% cases
    W = Diagonal(μ .* (1 .- μ))

    stats = zeros(P)
    pvals = zeros(P)

    @inbounds Threads.@threads :static for i in 1:P
        p = rand(Beta(3, 12))
        g = rand(Binomial(2, p), N)
        stats[i] = g' * (Y - μ) / sqrt(g' * W * g)
        pvals[i] = 2.0 * (1.0 - cdf(Normal(0, 1), abs(stats[i])))
    end

    return stats, pvals
end;

Random.seed!(3);
stats, pvals = gwas_sim();

Pvalues are not distributed uniformly under the null

Substantial deviation in qq plot

How can we generate properly calibrated test statistics?

To apply fastSPA to \(T_{adj}\), first derive its cumulant generating function (CGF). Given \(\widehat{\boldsymbol{b}}\), \(T_{adj}\) is modeled as a weighted sum of independent Bernoulli random variables. The approximated CGF is:

\[ K\left(t;\ \widehat{\mu}, c\right) = \sum_{i=1}^{N} \log\left(1 - \widehat{\mu}_{i} + \widehat{\mu}_{i} e^{c t \widehat{G}_{i}}\right) - c t \sum_{i=1}^{N} \widehat{G}_{i}^{\prime} \widehat{\mu}_{i} \]

where \(c = \text{Var}^{*}(\mathbf{T})^{-1/2}\).

FastSPA, Dey et al., AJHG, 2017

Derivatives and Probability Calculation

Let \(K^{\prime}(t)\) and \(K^{\prime\prime}(t)\) denote the first and second derivatives of \(K\) with respect to \(t\). To compute the probability \(\Pr(T_{adj} < q)\) for an observed test statistic \(q\), use:

\[ \Pr\left(T_{adj} < q\right) \approx F(q) = \Phi\left\{w + \frac{1}{w} \log\left(\frac{v}{w}\right)\right\} \]

where: \[ w = \text{sign}\left(\widehat{\zeta}\right) \left[2\left\{\widehat{\zeta} q - K\left(\widehat{\zeta}\right)\right\}\right]^{1/2}, \quad v = \widehat{\zeta} \left\{K^{\prime\prime}\left(\widehat{\zeta}\right)\right\}^{1/2} \]

and \(\widehat{\zeta} = \widehat{\zeta}(q)\) solves the equation \(K^{\prime}(\widehat{\zeta}) = q\).

\[ \Pr(T_{adj} < q) \approx \Phi\left(\frac{q - \mu}{\sigma}\right) \]

Review

Further optimizations

Mbatchou et al., Nature Genetics, 2021

REGENIE benchmarks

Mbatchou et al., Table 1

Optimization for multiple phenotypes

• No need to process each phenotype completely independently
• For example, rather than projecting out covariates for each phenotype one at a time, this can be done for a matrix storing all phenotypes:

\(\mathbb{Y} = \left[ Y_1, Y_2, \cdots, Y_P \right]\)

Summary