Training Flexible Models of Genetic Variant Effects from Functional Annotations using Accelerated Linear Algebra
Authors: Alan Nawzad Amin
Andres Potapczynski
Andrew Gordon Wilson
Venue: ICML 2025 poster (2025)
Keywords: GWAS
Transformers
Iterative Methods
Machine Learning
Numerical Linear Algebra
Primary area: probabilistic_methods->bayesian_models_and_methods
TL;DR: We improve disease prediction on GWAS data by developing a methodology that allows us to fit flexible neural networks.

Abstract: To understand how genetic variants in human genomes manifest in phenotypes - traits like height or diseases like asthma - geneticists have sequenced and measured hundreds of thousands of individuals. Geneticists use this data to build models that predict how a genetic variant impacts phenotype given genomic features of the variant, like DNA accessibility or the presence of nearby DNA-bound proteins. As more data and features become available, one might expect predictive models to improve. Unfortunately, training these models is bottlenecked by the need to solve expensive linear algebra problems because variants in the genome are correlated with nearby variants, requiring inversion of large matrices. Previous methods have therefore been restricted to fitting small models, and fitting simplified summary statistics, rather than the full likelihood of the statistical model. In this paper, we leverage modern fast linear algebra techniques to develop DeepWAS (Deep genome Wide Association Studies), a method to train large and flexible neural network predictive models to optimize likelihood. Surprisingly, we find that larger models only improve performance when using our full likelihood approach; when trained by fitting traditional summary statistics, larger models perform no better than small ones. We find larger models trained on more features make better predictions, potentially improving disease predictions and therapeutic target identification.

Training Flexible Models of Genetic Variant Effects from Functional Annotations using Accelerated Linear Algebra
Alan N. Amin 1
Andres Potapezynski 11 Andrew Gordon Wilson
Abstract
(GWAS) use data from up to hundreds of thousands of indi- viduals to build models that correlate the presence of genetic variants with phenotypes such as disease or height (Yang ct al., 2010; Visscher et al., 2017: Halldorsson et al., 2021). However there are orders of magnitude more variants than measurements, making GWAS underpowered to predict phe- notype determine the effects of all but the most impactful variants.
To understand how genetic variants in human genomes manifest in phenotypes - traits like height or diseases like asthma - geneticists have sequenced and measured hundreds of thousands of individuals. Geneticists use this data to build models that predict how a genetic variant impacts phenotype given genomic features of the variant, like DÑA accessibility or the presence of nearby DNA-bound proteins. As more data and features become available, one might expect predictive models to improve. Unfortunately, training these models is bottlenecked by the need to solve ex- pensive linear algebra problems because variants in the genome are correlated with nearby vari- ants, requiring inversion of large matrices. Pre- vious methods have therefore been restricted to fitting small models, and fitting simplified sum- mary statistics, rather than the full likelihood of the statistical model. In this paper, we leverage modern fast linear algebra techniques develop DeepWAS (Deep genome Wide Association Stud- ies), a method to train large and flexible neural network predictive models to optimize likelihood. Surprisingly, we find that larger models only im- prove performance when using our full likelihood approach; when trained by fitting traditional sum- mary statistics, larger models perform no better than small ones. We find larger models trained on more features make better predictions, potentially improving disease predictions and therapeutic tar- get identification.
To increase prediction accuracy and uncover the molecu- lar causes of disease, geneticists have leveraged the fact that complex phenotypes are extremely polygenic that is, they are affected by a huge number of variants spread throughout the genome (Manolio et al., 2009; Boyle et al., 2017). Geneticists look for features that distinguish variants that do and do not effect a phenotype on a set of chromo- somes and use these features to build "functionally informed" priors to analyze variants on other chromosomes (Gusev et al., 2014; Finucane et al., 2015; Kichaev et al., 2019). To build these priors they use functional genomic features (EN- CODE Project Consortium, 2012: Lizio al., 2015), such as measurements of DNA "accessibilityy" or binding of tran- scription factor proteins near the variant; and comparative genomics features (Cooper et al., 2005; Pollard et al., 2010), such as whether the variant is in a region of the genome that is conserved across primates. As more accurate mea- surements of genomic features are made and more individu- als have their genomes sequenced, in principle, geneticists should be able to build more accurate functionally informed priors with more flexible models that learn from more fea- tures.
In practice, however, significant computational challenges have prevented the development of large models. Function- ally informed priors are typically phrased as priors on the effect of each variant in hierarchical Bayesian model of the genetic and phenotypic data (Loh et al., 2015; Zheng et al., 2024). Ideally, we could fit the prior using empirical Bayes approach to maximize the marginal likelihood (Ni et al., 2018). Unfortunately this is numerically challenging due to linkage disequilibrium (LD) - the presence of variants in the genome can be strongly correlated, and accounting for this correlation in the marginal likelihood involves inverting and calculating the log determinant of large matrix known
1. Introduction
To predict the risk of genetic disease and understand its molecular causes, Genome Wide Association Studies
'quual contribution 'New York University. Correspondence to: Alan Amin <alanamin nyu.edu>.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Predicting variant effects from functional annotations
the LD matrix. To avoid inverting this matrix, state- of-the-art methods 1) fit simple parametric models of the
>100 features:
Effect on phenotype
Binding Consenvation
relation between functional annotations and phenotype, and 2) sacrifice statistical efficiency by fitting summary statistics or an approximation of the marginal likelihood (Finucane et al., 2015; Li et al., 2024).
Predictions
Chromosome
similar challenge of having to invert a large matrix per- form empirical Bayesian inference was addressed in works on Gaussian process regression with two strategies (Gardner et al., 2018 8). First, using iterative algorithm, inversion of an M x M matrix could be reduced from (MM)) to K) where << M is the number of iterations. Second, by rearranging the problem so that the large matrix well-conditioned the number of steps K could be reduced by orders of magnitude.
Hidden effect
M, x M, LD matrix
V ggLikeihhodd :
Observed association
Naive:
DeepWAS:
OM O(M))
Predicted effect fa
Here we introduce method to train large models that predict variant effects from functional annotations Deep genome Wide Association Studies (DeepWAS) (Fig. 1). We outline our contributions:
Figure I. DeepWAS enables training large models to predict the effect of variants from genomic features by leveraging fast linear algebra. Top: We want to train model, fe. to predict the effect of variant in our genome from a large set of curated genomic features in a window around the variant. Bottom: We train fo to maximize the likelihood of observed associations be- tween variants and traits. We efficiently compute the likelihood by applying accelerated linear algebra on the correlation matrix of variants in sliding window. See section for full details.
We propose method to train large neural networks phenotype association data with millions of variants by leveraging banded approximation to the LD matrix and using the approximating slices as mini-batches.
We train models that maximize likelihood rather than fit summary statistics or approximations. To do so, we rearrange likelihood to make it amenable to acceleration from iterative linear algebra algorithms, allowing us to efficiently perform challenging linear algebra operations at each training step.
In Sec. 2.1, we introduce hierarchical Bayesian model of heritability with functionally informed prior on the effects of individual variants. We describe how to fit such a model using public data. Then, in Sec. 2.2 we describe LD score regression, the state-of-the-art method to such a model in practice.
We curate a large set of genomic features to train our models.
We train large functionally informed priors on large public phenotype association data. We see that, when fitting summary statistics (LD score regression), larger models fit the data worse than small models. However, when maximizing likelihood with WeepWAS, larger models fit the data better than small models.
2.1. Functionally informed priors to predict variant effect
learn the heritibility of a trait, suppose that we have measured the genotypes of N (> 10*) subjects measured the presence absence of variants at M (= 10 10*) positions on both chromosomes to get genotype matrix X E {0, 20 MxN, and the presence of the trait to get a phenotype vector y E NR.. We can assume is centered to have mean 0 and variance 1 and X is a centered X with all rows mean and variance
have
We show that larger models trained more features achieve better data fit, suggesting that even larger models may yield further improvements.
Our code for training DeepWAS models is avail- able at https: /tt:://gghttbb../ com hoomm/llantn AlaaNaazaaAAmnn DeepWAS.
Measured traits that we are interested in, such as height, smoking status schizophrenia, are polygenic they are in- fluenced by many variants scattered throughout the genome rather than a small number of positions (Manolio et al., 2009; Boyle et al., 2017). This is captured by the infinitesi- mal model in which each variant has a small effect drawn from prior (Barton et al., 2016; Trippe et al., 2021).
2. Background
In this section we explain how to train models that describe
phenotypic traits using variants in the genome as features.
2

---

Predicting variant effects from functional annotations
A popular infinitesimal model is the linear model' y 8 + with iid noise ~ (O, T² where the effect size at position m, Pm. is independently drawn from prior normal distribution 9m ~ (O, fm)) Therefore we can describe the marginal distribution of y as
2.2. LD score regression (LDSR)
One way to solve problem 2 is avoid the need to invert large matrix by only fitting summary statistics. From Eqn 2 we can note that variant m expected to have large association if it is correlated to other variants expected to have large effects?:
y -~ ~~~ (0, + +o21), where F = dfa (l)
Our first goal is estimating the effect size 3. The challenge is that there are many more variables than observations, N << M, so it is challenging to get enough statistical power to predict the values of many B. Our second goal is to identify the features that characterize variants m with large effect sizes 8m-
E|N B
N Σ Sm R Remme + 0².
(4)
172
The simplest model of heritability gives each variant the same expected heritability, fm = =f. in which case we can write Eqn. 4 as E[N N [E R + 0². The term in the brackets, known as the LD score, measures how many other variants m is correlated with and can be precomputed before fitting f-
We can make progress towards these goals with good prior f. which will increase our power to determine 8.m. and pre- dict which variants are expected to have large magnitude since E333 = fm- To build such a prior, we can take advan- tage of large datasets of genomic features Cm (elaborated in Sec. 5), to predict fm with a model with parameters 0, fo)). Naively, we may train fo and o² by maximizing the marginal likelihood of Eqn. 1.
By fitting a line to the magnitudes of the association statis- tics N pm and precomputed LD scores, one can recover N f as the slope and o² the intercept. This method, known as LD score regression, gives accurate predictions of how much of a trait is explained by genetics. f. and how much is caused by noise or the environment, o² (Bulik-Sullivan et al., 2015). Finucane et al. (2015) extended this approach to fit a linear fm that depends on d genomic features - in this case one performs a lilliiidimensonall linear regres- sion with d precomputed variables. Unfortunately, LD score regression loses statistical efficiency by not making use of correlations between 3 (Ni et al., 2018).
Public Statistics However, to protect the privacy of study participants, we are not given the precise value of and X; rather we are given another set of public statistics. In particular, we are given the empirical correlation matrix known as the "Linkage Disequilibrium (L.D) matrix" R R= x and the empirical associations 3 = Xy. We can then write Eqn 1 in terms of public statistics with x3-- N²?
3. Previous work
N (0, o2
(2)
Training functionally informed priors Training a large neural network as a functionally informed prior by directly optimizing the likelihood of the data has, up until now, been computationally prohibitive due to the cost of linear algebra operations on the LD matrix. Previous methods have used a number of strategies to restrict the flexibility of their prior or looked at other approximate or derived objectives in order to do inference. First, most GWAS methods pick their prior with only a handful of parameters (usually I or 2) and fit it by grid search or other bespoke methods that struggle to scale (Yang et al., 2010; Loh et al., 2015; Speed et al., 2017; Spence et al., 2022). Second, Finucane et al. (2015) fit a linear prior by performing LD score regression in Eqn. 4. Third, Lu et al. (2016) and Fabiha et al. (2024) fit a small model by teaching it to classify the small number (*x 2000) of available high confidence positive and negative causal variants. Fourth, Li al. (2024) considered fitting simple generalized linear model fe by approximating Eqn. 3 using an approximation of R- All of these methods run on CPU and use parallelism to compute the gradient of the likelihood
The second term in the variance, 0% R, comes from spu- rious correlations with the noise e; if the presence of two variants m and m' are correlated (mim,'' large) then the
associations 3. and 9m will have similar correlations with the noise c. The first term in the variance comes from the effect variants have on the trait. Specifically, the m, m' entry of RFR is Σn Rem, fks which is large if there variants k correlated to both and m' - large Rm, and Rm A - which are expected to have large effect fi.
Now, in principle, we could build a prior by maximizing the likelihood of Eqn. using gradient descent:
3 3T (RFoR+ -1 3 log RFoR + Ro
1
+ (3) where is a constant value. The challenge is the need to 1) calculate all M non-zero entries in Fg. which becomes challenging for large models, and 2) invert and calculate the log determinant of the huge M x M matrix.
- 2
We ignore fixed effects such as age, sex, and population strati- fication in this model, assuming they have been projected out of y and X (Loh et al., 2018).
IDD score regression can also be derived in infinitesimal mod- els more general than Bayesian linear models with normal prior (Bulik-Sullivan et al., 2015).
3

---

Predicting variant effects from functional annotations
across the entire genome for each update.
place of the true R. DeepWAS uses an iterative algorithm to perform linear algebra operations the exact matrix; we rearrange our loss so it is amenable to acceleration from iterative algorithms.
Unfortunately these methods are unsuitable for training large flexible prior as they 1) require prior values for all variants in the genome for a single gradient update, making it challenging to train a large model or 2) lose statistical power by fitting summary statistics or approximations to the likelihood. In contrast, our method DeepWAS 1) updates the model using its predictions in minibatches, and accelerates linear algebra operations in each mini-batch with GPUs, 2) directly optimizes the likelihood of data from millions of variants.
Fast linear algebra for fitting large Bayesian models Fitting Gaussian processes similarly involves inverting a large matrix known as the Gram matrix. While one can avoid inverting the matrix with variational inference, state of the art methods invert the Gram matrix with an iterative algorithm with Nyström preconditioner (Gardner et al., 201 8). We build bespoke preconditioner leveraging the structure of LD matrices to quickly invert LD matrices with iterative algorithms; in our setting, our preconditioner per- forms much better than a general purpose Nyström precon- ditioner (Frangella et al., 2021).
In related work. Huang et al. (2024) fit a graph neural net- work of variants to predict 3 directly; they use their model to increase power find more associated variants. However such model does not distinguish between variants with large effects 3 and variants they are associated with.
4. Efficient training of the likelihood
Downstream uses of functionally informed priors A number of works have built methods use functionally informed priors to increase the power of downstream analy- ses. Huang et al. (2024) and Kichaev et al. (2019) demon- strated that models that can predict the effect of variants can improve the power of GWAS. Weissbrod et al. (2020) demonstrated such models can also identify causal variants and Li al. (2020) used such variants to identify causal genes. The DeepWAS prior can in principle fit into these same pipelines.
Our goal of directly optimizing the likelihood in Eqn 3 comes with two challenges. First, in contrast to previous methods, we train a neural network model for so with mil- lions of parameters, so computing sa.m for every variant m is prohibitively expensive. Second, we must perform expen- sive linear algebra operations like inverting and calculating the log determinant of As RF + o R. at every step".
First in Sec. 4.1, we utilize a banded approximation of R which allows us to amortize the training of A by treating each slice as a mini-batch. Then in Sec. 4.2 we reduce expensive linear algebra operations on Ag to operations on well-conditioned matrix, which can quickly be performed using iterative algorithms.
Flexible models of heritability In addition more flexi- ble models predicting variant effects from functional anno- tations, we can improve fits to association data with models that are more flexible than mixed linear models. Zhang et al. (2021) consider different, non-normal, priors, and Loh et al. (2015) consider mixture of normal priors on the effect sizes. There have also been a number of nonlinear models for predicting y from X (Conard et al., 2023). For simplicity, DeepWAS considers the popular normal prior with linear model and leaves more flexible models to future work.
4.1. Using submatrices for mini-batching
A key challenge in calculating Eqn 3 is computing fe.m for every m in the genome. To address this, we first break the genome up into 2700 windows of size one million and assume the associations in each window are generated independently. This can be justified by the fact that It is approximately block diagonal, so there is little correlation between further than one million positions apart (Berisa & Pickrell, 2016; Salchi Nowbandegani et al., 2023). Eqn. then becomes
Fast linear algebra with large genotype matrices Large linear algebra problems appear throughout population ge- netics. As such, number of other works have looked at approximately inverting large matrices in genetics. Loh et al. (2015) used a conjugate gradient algorithm to invert the ma- trix of correlations of variants between study individuals the empirical kinship matrix X Loh et al. (2018) noted that their algorithm converges much faster after removing the top eigenvalues of the kinship matrix, improving its con- dition number. Berisa & Pickrell (2016) approximated R with a block diagonal matrix, Shi et al. (2016) approximated R with a low rank matrix, and Salehi Nowbandegani et al. (2023) approximated the inverse of the R with an extremely sparse matrix; these works use these approximations in
Σ (450 100 + 1og/44"1
(5)
Next, note that A!" takes the following form
1 ( Ro + o R
'Here we focus on fitting 0 ignoring a. Typically can be estimated accurately using other simpler methods (Bulik-Sullivan et al., 2015)
4

---

Predicting variant effects from functional annotations
where B()); represents the rectangular submatrix of R whose rows variants in window i and R(),6) is similar. Calculating (0 still requires calculating fo,m for every variant m. To reduce this calculation, we then use the well- established fact that variants that are distant in the genome should have little correlation, and so we can use a banded approximation of R (Bulik-Sullivan et al., 2015); in particu- lar, we assume that Rgr when the positions of the k-th and r-th variants are more than one million apart. Thus,
where L(p Bay (1) R Rayco and W(s)
B R R00000 Re), and |A00 as
log |4 = log 1 R(ty.e
W Fo /
+ log |1 +
NN
Note the matrices Blo) L(j) and (()) do not depend on O and can therefore be calculated before training begins.
Ag R0),
(1) Boo (0) + o R00000
Given the previous simplifications, now the only challenging linear algebra computation to perform at every step is to invert and calculate the log determinant of the matrix
where (i)+ is the set of all variants within 10 positions of a variant in window (i) and F6)) is the (i)* x (i)+ sub- matrix of Fe. Eqn 5 then allows us to optimize 0, through
// = -
(0) W WO (1)
stochastic gradient descent, by sampling windows (i) and only calculating fe.m for the roughly 10a variants in (i)+.
N
Although the dimensions of BB are the size of (i)+ - pos- sibly 3 times the size of (i) - it is strictly positive definite and, since x x Fg = N x M-1 < 1, it is typically well- conditioned. This last fact means that despite its increased size, we can actually invert it substantially faster than A using iterative algorithms.
Connection to LD score regression Due the large size of our windows, we expect both approximations above to be accurate. In contrast, if we focus on the extreme case of a window size of 1 the objective, Eqn. 5 becomes
N
Σ NN Emn S' R mom' + 0²
+ +log(N ( Σ sm Raan + +²) ml
4.2.2. ITERATIVE ALGORITTMS
To invert and calculate the log determinant of A or B, we could perform Cholesky decomposition. Unfortunately, this is ( where M, is the number of rows of the matrix; this is too computationally expensive to perform at each step and becomes more so if we wish to include more variants in our study in the future.
which tries to fit N Emm fm Ramm + o² to N BR. This is exactly the idea of LD score regression (Eqn. 4). Therefore LDSR can be thought of as our objective in the extreme case assuming every /3 was generated independently.
4.2. Fast linear algebra with iterative algorithms
Luckily, for well-conditioned matrices, we can achieve much better computational complexity by using iterative methods like stochastic Lanczos quadrature (SLQ) (Golub & Loan, 2018; Saad, 2011) for log |B,²| and conjugate gradients (CG) (Nocedal & Wright, 2006; Golub & Loan, 2018; Saad, 2003) for solves (BB )-1 Both methods rely on performing multiplications against B8²) at each iteration and each iteration typically exponentially improves the qual- ity of the approximation, typically until machine precision. Thus, the computational cost of both methods is manage- able O(M? K) where K is the number of iterations and O (M?) is an upper bound on the cost of doing a matrix- multiply with B,a) The number of iterations required to converge below an error threshold of these iterative methods is directly linked to the eigenspectrum of the matrix (No- cedal & Wright, 2006; Saad, 2011; Hogben, 2013) - since B is well conditioned, we expect it to need few iterations K and therefore fewer computational resources.
To calculate the likelihood, we must now invert and calcu- late the log determinant of A5") However, this matrix is large, and often singular, making calculations expensive and unstable. We address these problems by rewriting the loss in terms of well-conditioned matrix and then efficiently calculating its inverse using iterative algorithms.
4.2.1. 1. REFORMULATING THE LOSS
Previous works like Salehi Nowbandegani et al. (2023) or Hormozdiari et al. (2014), deal with the singularity issues by adding regularization to R()),0) R())() eT for some small c, which trades off numerical stability for bias. In con- trast, we can work with the exact R by using the Woodbury identity and the matrix determinant lemma to write the loss in terms of the pseudo-inverse R/ (0)(0) (see Appendix C.1): we re- re-teite )4
(1())-1==
1 R()000)
Rt
To implement these methods, we use CoLA (Potapczynski et al., 2023), a numerical linear algebra library that is com- patible with diverse deep learning frameworks and that pro- vides backpropagation capabilities for SLQ and CG. COLA
Nº
T+
1
W
W
F
L()
0
N
5

---

Predicting variant effects from functional annotations
CPU
GPU
15
Time (s)
1001
Effect enrichment
10
10
1
Conservation
0.1
(phastCons470)
Iter, Aa
Nys. Aa
Chol. Aa
Iter. Be
Iter. Aa
Nys. Aa
Chol. Aa
Iter. Ba
Accessibility (DNase-se q)
Histone mod. (H3K9me1)
Figure WeS efficiently computes the loss and its gradi- ent. We measure the time it takes to compute our loss Eqn. 3 as well as its gradients with respect to We this for mini-batches of real UKBB data and display the mean runtime as harplots, Chol stands for Cholesky decomposition. Iter stands for the iterative algorithms SLQQ and CG. Nys stands for the itera- tive algorithms with Nyström preconditioning. We set a relative tolerance of 10 for CG and use 100 samples of SLQ yielding average relative error on the gradient of 2.5%. For GPU we used an NVIDIA 1000SXM4-80GG and for CPU Intel(R) Xeon(R) Platinum 8268 CPU @ 2.90GHzz.
binding (TRIM24)
Position on chromosome
Figure 3. DeepWAS flexibly predicts the effect of variant on a phenotype from variety of genomic features. We plot the pre- dicted effect enrichment / fo.mm/ /B..ff.. m for hypothetical variants at all positions in 5000-Iong non non-coding window of chromosome 6, with the expectation over variants in UKBB. We compare with selected genomic features fed into fe. We predictions from DeepWAS transformer- model trained on height in UKBB.
computes the gradients of B5 and log [B=1] by using the following identities
features than previous works in Sec. 5.1. We describe large flexible neural network architectures for fe in Sec. 5.2. The result is flexible function that predicts the effect of vari- ant from a variety of genomic information (Fig. 3).
V = B, B
Vo log B7' trace Bj BB)
Ew B3 Vo Bou
where both quantities require backpropagating through Be only and where we use the Hutchinson trace estimator. Ad- ditionally, CoLA allows us to leverage GPU acceleration for our numerical techniques which significantly reduces the runtime.
5.1. Features
Previous methods have built fo using functional genomics features such as DNA accessibility, proximity to functional elements, and presence genomics features such as conservation scores (Finucane et al., 2015; Li et al., 2024). Many of these features defined as annotations each position in the genome; to get single value, annotations were averaged in window before being passed to fo-
coding region and comparative
In Fig. 2 we compare the cost of computing the likelihood Eqn. 3 by using iterative algorithms on A or ( or by performing Cholesky decomposition on A (performing Cholesky on B²a)) would take longer due to its increased number of rows). We also compare with Nyström precon- ditioning, a popular method for solving large systems in machine learning that effectively makes A!" better condi- tioned (Gardner et al., 2018). We see that, despite being larger, we can invert B faster than A/ by leveraging its being well conditioned using iterative algorithms, even when using Nyström preconditioning. We also further accel- erate our calculations by performing them on GPUs rather than CPUs, which had previously been used to train func- tionally informed priors.
We expand this set in two ways. First we consider a signif- icantly expanded set of functional genomics annotations binding and accessibility annotations from ENCODE (EN- CODE Project Consortium, 2012), enhancer annotations from FANTOM (Lizio et al., 2015) and comparative genomics annotations conservation scores such Phy- loP (Pollard et al., 2010), and predictions of effects of mu- tations in coding regions such as those from ESM2 (Liu et al., 2020). Details these data are in App. D.2 and D.3. Second, instead of considering an average of the values of annotations in window around the variant, we pass the model the exact values of the annotations at all positions in the window.
5. Predicting variant effects from genomic features
Gazal et al. (2017) used LDSR to demonstrate that the recent history of a variant in humans can also be predictive of its ef- fect size. To account for this, we also included the frequency of each variant m, freq.: its "minor allele" frequency,
Now we have a method for accurately and efficiently train- ing a large model fe. Here we specify how we build fo that include many more functional and comparative genomies
6

---

Predicting variant effects from functional annotations
min ((reg freqm J: and its LD score Em Re features.
LDSR
DeepWAS
log fa error (1)
1.5
5.2. Architecture
1
For all genomics annotations but coding mutation effect predictions, consider a window around each variant of size w. We pass these functional annotations Cuunc.m E Rdaaue X along with predictions of the effects of muta- tions if the mutation is in a coding region and genomic architecture information, Cpred,m Rdpred, to a neural net- work funnnommm m.) In our case, drune = 165, and dpred 9; we also choose a window size of w 256. We transformer-CNN hybrid architecture adapted from network used to predict functional annotations from se- quence, Enformer (Avsec et al., 2021)*; this architecture uses a mix of convolutional and attention layers. We can change the dimension of the internal representation to alter the number of parameters of the model.
0.5
0
Constant
Linear
Transformer
Figure 4. DeepWAS using transformer model best recovers the true f. The bars represent the RMSE difference between the learnt fo and the ground truth f in the log space evaluated over set of validation functional annotations.
ropean ancestry (Weissbrod et al., 2020) (see App. D.1). We generate synthetic 3 as follows:
Bm ~ N (0, Jm (Crune,m; Cpred, ~ ~ (R3, N R)
Speed al. (2017) suggested that setting f constant in model makes the implicit assumption that rare variants have larger effects. They remove this assumption with more general model fm = (freqm - freq...)) where a is a fit parameter. In our case, we consider
where f represents that function that we are trying to learn. We consider So as randomly initialized transformer- based neural network model (52 million parameters). See App. A.5 for details.
We tried fitting this data with models based on Eqn. 6. We used simple models NN, = constant and NN, = gen- eralized linear model (see App. A.2) and a more flexible NNg with the transformer architecture with LD score regres- sion (LDSR) and DeepWAS. In Figure we show DeepWAS with a large model can closely recover the true variant effect distribution fo achieves low error in predictions fo. Furthermore, this model better predictions than restricted constant and linear fa. We also see that our method makes more accurate predictions than models trained with LD score regression. In Sec. B.1 we also consider inferring a more challenging ground truth f.
fa.m (freqm (1 )
Cpred.m) (6)
where NNg is neural network or any other model. In many of our experiments, typically converged to value between 0.6 and 0.7 regardless of its initialization; to sim- plify our model comparisons below, we fixed it at 0.7 for all experiments.
6. Empirical Results
We now apply DeepWAS to train flexible neural network models in order to better explain which variants are associ- ated with phenotypic traits. See A for experimental details.
6.2. Predicting phenotype on UK Biobank
6.1. Deep WAS learns true f from semi-synthetic data
Now we attempt to better explain phenotypes y of three traits from UK Biobank calculated in Loh et al. (2015) (see App. D.I) - body mass index, height, and asthma. Unfor- tunately, we only have access the public 3 rather than the private y. Using just public data out some chromosomes and evaluate our models on the following task:
We first wish to check if our gradient descent procedure can recover the true f from noisy genetic data. Unfortunately we do not have access to real f real data. We therefor turn to semi-synthetic simulations, which are staple of statistical genetics literature (ex. Candès et al. (2016) O'Connor et al. (2019)), used to validate methods when true effects 3 are unavailable.
however hold
train model 9 phenotypes y and genotypes X of held-in chromosomes,
We use real LD patterns R from public data from the UK Biobank (UKBB) study of over 300,000 individuals of Eu-
then predict phenotypes y using only genotypes of held-out chromosomes.
"We do not use pretrained weights as the original Enformer architecture regressed from sequence to predict functional annota- tions, rather than our task which is to take in functional annotations to predict effect on phenotype.
We report how much the marginal likelihood (Eqn. 1) increases using the model on held-out chromosomes
7

---

Predicting variant effects from functional annotations
% likelihood increase
BMI
Height
Asthma
per person (T)
11.7
0.694
2.77
10.9
0.296
0.319
9.39
0.226
2.24 2.28
0.148
2.19
FZZZ
Constant
Linear
Transformer
Constant
Linear
Transformer
Constant
Linear
Transformer
Figure 5. More flexible models trained with DeepWAS better explain phenotype using held-out chromosomes. We try predict the phenotype y of individuals in UK Biobank using only variants held out chromosomes 6, 7, as in Eqn. 1. We report how much the marginal likelihoods inereases when phenotype (F' 0). When training with DeepWAS, larger models better fit the data. While larger models perform worse when training with LDSR. Legend is identical to Fig. 4.
different priors F, compared model which assumes there is genetic effect the
compared to a "null" model that assumes there is no ge- netic effect on y (f O). We then divide the marginal log likelihood by N and exponentiate to get quantity "per person".
likellhood Increase
BMI
Height
Asthma
per person (T)
We show in App. C.2 that this is equivalent to training our model using public association data 3 on held-in chromo- somes and reporting the marginal likelihood (Eqn. 3) of association data 3 on held-out chromosomes. In our experi- ments, we hold out chromosomes 6, 7, and 8 and train all other autosomal chromosomes.
%
1w
params features Full model
400 1 params features Full model
4(1²*
1 params feetures Full model
4
Figure 6. Ablations show that larger models with more features better explain genetic associations. We ablate the number of vari- ants in an LD window ((1) ). the feature set (w and the presence of conservation features) and the model size of our transformer based model. We report statistics in Fig. 5, with the same y-axes. We compare them to the full model (dark blue) reporting the same value in Fig. 5. For exact numbers App. B.2.
Using flexible models with Wee results in better explanations of phenotype data. We evaluate the ability of various architectures trained with LDSR and DeepWAS to predict phentotype using held-out chromosomes. LDSR and DeepWAS used similar computational resources in training.
Fig. shows that a model with constant/ f better explains the data than null model, f 0 this represents how much better our prediction of phenotype is simply from including information about genotype X in our model. We next see that including functional annotations in linear prior increases prediction quality this represents how much better our prediction is from including any informa- tion about genomic features. The varying improvements across the phenotypes correlate with their reported "SNP- heritabilities" (Hou et al., 2019) more than 50% of the be explained
Larger features and larger models improve prediction. We finally set out to determine the importance of the number of variants we trained on, the feature set, and model size for performance. We trained models with fewer variants per window we reduced the size of (i)+ relative to (i) from windows of 1'000'000 000 to 100'000 (see Sec. 4.1) -a reduced feature set we looked rather than the conservation features from phastCons and PhyloP (see Sec. 5.1) and with a reduced size we reduce the number of parameters from 52 million to 3.9 million. We chose to ablate the conservation features as these have been seen to be particularly important for effect size prediction in previous work (Finucane et al.. 2015).
window of w 16 256 around each variant, and removed
variation of height seen in the population by variants in our study, while that number is roughly 30% for BMI and 10-16% for asthma.
We next increase the model size to a transformer. Surpris- ingly, when training using LDSR, the quality of prediction does not significantly increase, and the case of BMI, even decreases. The more flexible architecture potentially over- fits the data due to the loss of statistical efficiency when performing LDSR. When accounting for correlations in the 3 with Deep WAS however, the larger transformer substan- tially outperform all other models.
Fig. 6 shows that ablating the number of variants in each window, the model size, or feature set often harms model performance. This shows that saving compute by consid- ering smaller (i)+ severely harms performance. This also suggests that our model benefits from its flexibility and features to better predict the effects of variants.
8

---

Predicting variant effects from functional annotations
7. Discussion
Acknowledgments
By efficiently inverting LD matrices, DeepWAS allows us to train large models that better predict the effects of variants on phenotype and to learn their functional causes. Our re- sults demonstrate that larger models make better predictions than the simple models used in practice, and that increasing the model size and using more features improves predictive power. Our work suggests that even larger models trained on more data may result in further improvements in the future.
We thank Benjamin Strober and David Knowles for help- ful advice on dealing with the idiosyncrasies of genetic data and literature. We thank Shenglong Wang at NYU IT High Performance Computing for help optimizing the data loader to work efficiently on high performance com- pute resources. This work was supported in part by NSF CAREER IIS-2145492, NSF CDS&E- MSS 2134216, NSF HDR-2118310, BigHat Biosciences, Capital One, and an Amazon Research Award.
Confounding and other priors Future work may address some of the limitations of WAS. First, LD score re- gression can deal with some confounding from population stratification (Bulik- Sullivv-S et al., 2015), while such effects were only accounted for in DeepWAS by using population- covariate-adjusted BOLT-LMM public statistics (Loh et al., 2015). Future work may investigate the effectiveness of this method. Next, DeepWAS considers simple linear model with a single component normal prior. More flexible models with multi-component priors or accounting for non-linearity would be a natural next step (Loh al., 2015; Zhang et al., 2021).
Impact Statement
This paper presents work whose goal is to leverage large data to better predict the genetic causes of human traits. It could be used to better predict disease in the future. On the other hand, genetics data is not collected uniformly across the population; therefore building models using existing large data may only improve prediction for subsets of the population, exacerbating health outcome inequalities.
References
Avsec, Z., Agarwal, V., Visentin, D., Ledsam, J. R., Grabska- Barwinska, A., Taylor, K. R., Assael, Y., Jumper, J., Kohli, P., and Kelley, D. R. Effective gene expression predic- tion from sequence by integrating long-range interactions. bioRxiv, pp.
Larger models on more data DeepWAS could be ex- panded to include other sources of information, such as em- beddings of genes that are near variants (Weeks al., 2020; Theodoris et al., 2023) and use pre-trained models (Avsec et al., 2021; Dalla-Torre et al., 2023). In this work, we only looked at single European population; for more robust estimates DeepWAS could be trained to learn across popu- lations (Spence et al., 2022). Finally, rather than training on a single phenotype at a time, Wee could learn from multiple phenotypes (Morgante et al., 2023) boosting power to learn common patterns.
2021.
Barton, N. H., Etheridge, A. M., and Véber, A. The infinites- imal model. bioRxiv, February 2016.
Berisa, T. and Pickrell, K. Approximately indepen-
dent linkage disequilibrium blocks in human populations. Bioinformatics, 32(2):283-2855 January 2016.
Boyle, E. A., Li, Y. I., and Pritchard, J. K. An expanded view of complex traits: From polygenic to omnigenic. Cell, 169(7):117--1188)) June 2017.
Boosting signal in rare or low-heritability disease We saw in Fig. 5 that the flexibility of larger model made the largest impact in asthma, where the "SNP-heritbbility" was low since signal-to-noise was lower, an improved prior made a larger impact. DeepWAS could potentially impact other settings with low signal-to-noins
Bulik- SulikvSullivan, B. K., Loh, P.-R., Finucane, H. K., Ripke, S., Yang. J., Schizophrenia Working Group of the Psychi- atric Genomics Consortium, Patterson, N., Daly, M. J., Price, A. L., and Neale, B. M. LD score regression distin- guishes confounding from polygenicity in genome-wide association studies. Nat. Genet., 47(3): March 2015.
such as the interpre- tation of rare variants (Lee et al., 2014), or, using pre-trained WeS models, the analysis of small cohorts (Huang et al., 2024) or rare disease (Samocha et al., 2014).
As well, inference in low signal-to-noise settings in genetics typically require bespoke tools and analyses. Using an em- pirical Bayesian framework to account for uncertainty. Deep- WAS can both leverage functional insights from common variants and extract signal from rare variants themselves, providing a unified framework across the allele frequency spectrum.
Candès, E., Fan, Y., Janson, L., and Lv, J. Panning for gold: 'model-x' knockoffs for high dimensional controlled vari- able selection. Journal of the Royal Statistical Society: Series B (Statistical Methodology). 80, October 2016.
Conard, A. M., DenAdel, A., and Crawford, L. A spec- trum of explainable and interpretable machine learning
9

---

Predicting variant effects from functional annotations
approaches for genomic studies. Wiley Interdiscip. Rev. Comput. Stat., May 2023.
Golub, G. H. and Loan, C. F. V. Matrix Computations. The Johns Hopkins University Press, 2018. Fourth Edition.
Cooper, G. M., Stone, E. A., Asimenos, G., NISC Compar- ative Sequencing Program, Green, E. D., Batzoglou, S., and Sidow, A. Distribution and intensity of constraint in mammalian genomic sequence. Genome Res., 15(7): 901-913, July 2005.
Gusev, A., Lee, S. H., Trynka, G., Finucane, H., Vilhjälmsson, B. J., Xu, H., Zang, C., Ripke, S., Bulik- Sullivan, B., Stahl, E., Schizophrenia Working Group of the Psychiatric Genomics Consortium, SWE-SCZ Con- sortium, Kähler, A. K., Hultman, C. M., Purcell, S. M., McCarroll, S. A., Daly, M., Pasamiuc, B., Sullivan, P. F., Neale, B. M., Wray, N. R., Raychaudhuri, S., Price, A. L., Schizophrenia Working Group of the Psychiatric Genomics Consortium, and SWE-SCZ Consortium. Par- titioning heritability of regulatory and cell-type-specifie variants across 11 common diseases. Am. J. Hum. Genet., 95(5):535-55 552,, November 2014.
Dalla-Torre, H., Gonzalez, L., Revilla, M., Carranza, N. L.. Grzywaczewski, A. H., Oteri, F., Dallago, C., Trop. E., Sirelkhatim, H., Richard, G., Skwark, M., Beguir, K., Lopez, M., and Pierrot, T. The nucleotide transformer: Building and evaluating robust foundation models for hu- man genomics. bioRxiv, pp. 2023.01.11.523677 January 2023.
ENCODE Project Consortium. An integrated encyclopedia of DNA elements in the human genome. Nature, 489 (7414):57-74, September 2012.
Halldorsson, B. V., Eggertsson, H. P., Moore, K. H. S.. Hauswedell, H., Eiriksson, O., Ulfarsson, M. O., Pals- son, G., Hardarson, M. T., Oddsson, A., Jensson, B. O., Kristmundsdottir, S., Sigurpalsdottir, B. D., Stefansson, O. A., Beyter, D., Holley, G., Tragante, V., Gylfason, A., Olason, P. I., Zink, F., Asgeirsdottir,, M., Sverrisson, S. T., Sigurdsson, B., Gudjonsson, A., Sigurdsson, G. T., Halldorsson, G. H., Sveinbjornsson, G., Norland, K., Styrkarsdottir, U., Magmusdottir, D. N., Snorradottir, S., Kristinsson, K., Sobech, E., Thorleifsson, G., Jons- son, F., Melsted, P., Jonsdottir, I., Rafnar, T., Holm, H., Stefansson, H., Saemundsdottir, J., Gudbjartsson, D. F. Magnusson, O. T., Masson, G., Thorsteinsdottir, U., Hel- gason, A., Jonsson, H., Sulem, P., and Stefansson, K. The sequences of 150,1 genomes in the UK biobank. bioRxiv, pp. 2021. 11. 16.46824, November 2021.
Fabiha, T., Evergreen, I., Kundu, S., Pampari, A., Abramov, S., Boytsov, A., Strouse, K., Dura, K., Fang, W., Kerner, G., Butts, J., Ali, T., Gschwind, A., Mualim, K. S., Moore, J. E.,, Weng. 7., Ulirsch, J., Ji, H. E., Vierstra, J., Reddy, T. E., Montgomery, S. B., Engreitz, J., Kundaje, A., Tewhey, R., Price, A., and Dey, K. A consensus variant- to-function score to functionally prioritize variants for disease. bioRxiv, November 2024.
Finucane, H. K., Sulliva,, B., Gusev, A., Trynka, G., Reshef, Y., Loh, P.-R., Anttila, V., Xu, H., Zang, C., Farh, K., Ripke, S., Day, F. R., ReproGen Consor- tium, Schizophrenia Working Group of the Psychiatric Genomics Consortium, RACI Consortium, Purcell, S., Stahl, E., Lindstrom, S., Perry, J. R. B., Okada, Y., Ray- chaudhuri, S., Daly, M. J., Patterson, N., Neale, B. M., and Price, L.. Partitioning heritability by functional an- notation using genome-wide wide association summary statis- tics. Nat. Genet., 47 I): 1228-1222112 1235, November 2015.
Hogben, L. Handbook of Linear Algebra. Chapman and Hall/CRC, 2013.
Hormozdiari, F., Kostem, E., Kang, E. Y., Pasaniue, B., and Eskin, E. Identifying causal variants at loci with multiple signals of association. Genetics, 198(2):497- 508, October 2014.
Frangella, Z., Tropp, J. A., and Udell, M. Randomized Nyström Preconditioning. arXiv 2110.02820v2, 2021.
Hou, K., Burch, K. S., Majumdar, A., Shi, H., Mancuso, N., Wu, Y., Sankararaman, S., and Pasaniuc, B. Accurate estimation of SNP herrhbrityyylity from biobank-scall data irrespective of genetic architecture. Nat. Genet., 51()): 1244-1251, August 2019,
Gardner, J. R., Pleiss, G., Bindel, D., Weinberger, K. Q., and Wilson, A. G. GPyTorch: blackbox matrix-matrix gaussian process inference with GPU acceleration. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, NIPS' 18, pp- 7587-7597, Red Hook, NY, USA, December 2018. Cur- Associates Inc.
Huang, K., Zeng, T., Koc, S., Pettet, A., Zhou, J., Jain, M., Sun, D., Ruiz, C., Ren, H., Howe, L., Richard- son, T. G., Cortes, A., Aiello, K., Branson, K., Pfen- ning, A., Engreitz, J. M., Zhang, M. J., and Leskovec, J. Small-cohort GWAS discovery with Al over mas- sive functional genomics knowledge graph. Genetic and Genomic Medicine, (medxxxi:20224112.3 211112200 December 2024.
Gazal, S., Finucane, H. K., Furlotte, N. A., Loh, P.-R., Palamara, P. F., Liu, X., Schoech, A., Bulik-Sullivan, B., Neale, B. M., Gusev, A., and Price, A. L. Linkage disequilibruum-deeeeeen architecture of human complex traits shows action of negative selection. Nat. Genet., 49 140:114211427, October 2017.
10

---

Predicting variant effects from functional annotations
Hubisz, M. J., Pollard, K. S., and Siepel, A. PHAST and RPHAST: phylogenetic analysis with space/time models. Brief. Bioinform., 12(1):41-51, January 2011.
E. A., Hoen, P. A. C., Tatum, Z.., Thompson, M., Toy- oda, TT,, Wright, D. W., Daub, C. O., Itoh, M., Carninci, P., Hayashizaki, Y., Forrest, A. R. R., Kawaji, H., and FANTOM consortium. Gateways to the FANTOMS pro- moter level mammalian expression atlas. Genome Biol., 16(1):22, January 2015.
Kichaev, G., Bhatia, G., Loh, P.-R., Gazal, S., Burch, K., Freund, M. K., Schoech, A., Pasaniue, B., and Price, A. L. Leveraging polygenic functional enrichment to improve GWAS power, Am. J. Hum. Genet., 104(1);65- 75, January 2019.
Loh, P.-R., Tucker, G., B. J., Finucane, H. K., Salem, R. M., Chasman, D. I., Ridker, P. M., Neale, B. M., Berger, B., Patterson, N., and Price, A. L. Efficient bayesian mixed-modes analysis increases association power in large cohorts. Nat. Genet., 47(3):284 290, March 2015.
B. K., Vilhjälmsson,
Lee, S., Abecasis, G. R., Boehnke, M., and Lin, X. Rare- variant association analysis: study designs and statistical tests. Am. J. Hum. Genet., 95(1):5-23, 20-,2,,,J,, July 2014.
Li, H., Kamath, T., Mazumder, R., Lin, X., and O' O'Connor, L. Improved heritability partitioning and enrichment analyses using summary statistics with graphREML. Genetic and Genomic Medicine, 21t.::2202..1114424111666666
Loh, P.-R., Kichaev, G., Gazal, S., Schoech, A. P., and Price. A. L. Mixedd-mdee association for biobank-scale datasets. Nut. Genet., 50(7):906-908, July 2018.
Lu, Q. Yao, Hu, Y., and Zhao, H. GenoWAP: GWAS sig- nal prioritization through integrated analysis of genomic functional annotation. Bioinformatics, 32(4):542-548, February 2016
November 2024.
Li, X., Li, Z., Zhou, H., Gaynor, S. M., Liu, Y., Chen, H., Sun, R., Dey, R., Arett, D. K., Aslibekyan, S., Ballan- tyne, C. M., Bielak, L. F., Blangero, J., Boerwinkle, E., Bowden, D. W., Broome, J. G., Conomos, M. P., Correa, A., Cupples, L. A., Curran, J. E., Freedman, B. I., Guo, X., Hindy, G., Irvin, M. R., Kardia, S. I., R., Kathire- san, S., Khan, A. T., Kooperberg, C. L., Laurie, C. C., Liu, X. S., Mahaney, M. C., Manichaikul, A. W., Mar- tin, I.. W., Mathias, R. A., McGarvey, S. T., Mitchell, B. D.. Montasser, M. E., Moore, J. E., Morrison, A. C., O' O'Connell, J. R., Palmer, N. D., Pampana, A., Peralta, J. M., Peyser, P. A., Psaty, B. M., Redline, S., Rice, K. M., Rich, S. S., Smith, J. A., Tiwari, H. K., Tsai, M. Y., Vasan, R. S., Wang. F. F., Weeks, D. E., Weng, Z., Wilson, G., Yanek, L. R., NHLBI Maans for Precision Medicine (TOPMed) Consortium, TOPMed Lipids Work- ing Group, Neale, B. M., Sunyaev, S. R., Abecasis, G. R., Rotter, J. I., Willer, C. J., Peloso, G. M., Natarajan, P., and Lin, X. Dynamic incorporation of multiple in silico functional annotations empowers rare variant association analysis of large whole whole-genome genome sequencing studies at scale. Nat. Genet., 52(9):969 9983,999833888 September 2020.
Manolio, T. A., Collins, F. S., Cox, N. J., Goldstein, D. B., Hindorff, L. A., Hunter, D. J., McCarthy, M. I., Ramos, E. M., Cardon, L. R., Chakravarti, A., Cho, J. H., Guttmacher, A. E.,, Kong, A., Kruglyak, L.,, Mardis, E., Rotimi, C. N., Slatkin, M., Valle, D., Whittemore, A. S., Boehnke, M., Clark, A. G., Eichler, E. E., Gibson, G., Haines, J. L., Mackay, F. C., McCarroll S. A., and Viss- cher, P. M. Finding the missing heritability of complex diseases. Nature, (7765))) October 2009.
Morgante, F., Carbonetto, P., Wang, G., Zou, Y., Sarkar, A., and Stephens, M. A flexible empirical bayes approach to multivariate multiple regression, and its improved ac- curacy in predicting multi-tissue gene expression from genotypes. PLoS Genet., 19())xe0010539, July 2023.
Ni, G., Moser, G., Schizophrenia Working Group of the Psychiatric Genomies Consortium, Wray, N. R., and Lee, S. H. Estimation of genetic correlation via linkage dise- quilibrium score regression and genomic restricted maxi- mum likelihood. Am. J. Hum. Genet., 102(6): June 2018
Liu, X., Li, C., Mou, C., Dong, Y., and Tu, Y. dbNSFP v4: comprehensive database of transcript-specifie functional predictions and annotations for human nonsynonymous and splice-site SNVs. Genome Med., 12(1):103, Decem- ber 2020.
Nocedal, J. and Wright, S. J. Numerical Optimization. Springer, 2006. Second Edition.
O'Connor, L. J., Schoech, A. P., Hormozdiari, F., Gazal, S., Patterson, N., and Price, A. L. Extreme polygenicity of complex traits is explained by negative selection. Am. J. Hum. Genet., 105(3):456 476, September 2019.
Lizio, M., Harshbarger, J., Shimoji, H., Severin, J., Ka- sukawa, T., Sahin, S., Abugessaisa, I., Fukuda, S., Hori, F., Ishikawa-Kato, S., Mungall, C. J., Arner, E.,, Baillie, J. K., Bertin, N., Bono, H., de Hoon, M., Diehl, A. D., Di- mont, E., Freeman, T. C., Fujieda, K., Hide, W., Kaliyape- rumal, R., Katayama, T., Lassmann. T., Meehan, T. F., Nishikata, K., Ono, H., Rehli, M., Sandelin, A., Schultes,
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury. J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Köpf, A., Yang, E.,, Vitooo, Z., Rai- son, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang,

---

Predicting variant effects from functional annotations
L.,, Bai, J., and Chintala, S. PyTorch: An imperative style, gerh-merformacce deep learning library. arXiv [cs.LG), December 2019.
Trippe, B., Finucane, H., and Broderick, T. For high- dimensional hierarchical models, consider exchangeabil- ity of effects across covariates instead of across datasets. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P. S., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 13471- 13484. Curran Associates, Inc., 2021.
Pollard, K. S., Hubisz, M. J., Rosenbloom, K. R., and Sie- pel, A. Detection of nonneutral substitution rates on mammalian phylogenies. Genome Res., 20(1):110-121, January 2010.
Visscher, P. M., Wray, N. R.,, Zhang, Q., Sklar, P., McCarthy, M. I., Brown, M. A., and Yang. J. 10 years of GWAS discovery: Biology, function, and translation. Am. J. Hum. Genet., 101(1):5-22, July 2017.
Potapczynski, A., Finzi, M., Pleiss, G., and Wilson, A. G. CoLA: Exploiting Compositional Structure for Automatic and Efficient Numerical Linear Algebra. Advances in Neural Information Processing Systems (NeurIPS), 2023.
Weeks, E. M., Ulirsch, J. C.,, Cheng, N. Y., Trippe, B. L., Fine, R. S., Miao, J., Patwardhan, T. A., Kanai, M.. Nasser, J., Fulco, C. P., Tashman, K. C., Aguet, F,, Li, T., Ordovas-Montanes, J., Smillie, C. S., Biton, M., Shalek, A. K., Ananthakrishnan, A. N., Xavier, R. J., Regev, A., Gupta, R. M., Lage, K., Ardlie, K. G., Hirschhorn, J. N., Lander, E. S., Engreitz, J. M., and Finucane, H. K. Lever- aging polygenic enrichments of gene features to predict genes underlying complex traits and diseases. bioRxiv, September 2020.
Saad, Iterative Methods for Sparse Linear Systems. SIAM, 2003.
Saad, Y. Numerical methods for large eigenvalue problems. SIAM, 2011.
Salehi Nowbandegani, P., Wohns, A. W., Ballard, L., Lan- der, E. S., Bloemendal, A., Neale, B. M., and O' O'Connor, L. J. Extremely sparse models of linkage disequilibrium in ancestrally diverse association studies. Nat. Genet., 55 (9):1494-5502, September 2023.
Weissbrod, O.,, Hormozdiari, F., Benner, C., Cui, R., Ulirsch, J., Gazal, S., Schoech, A. P., van de Geijn, B., Reshef, Y., Márquez-Luna, C., O'Connor, L.,, Pirinen, M., Fin- ucane, H. K., and Price, A. L. Functionally informed fine- map-map and polygenic localization of complex trait heritability. Nat. Genet., 52(12) 1363, December 2020.
Samocha, K. E., Robinson, E. B., Sanders, S. J., Stevens, C.,, Sabo, A., McGrath, L. M., Kosmicki, J. A., Rehnström, K., Mallick, S., Kirby, A., Wall, D. P., MacArthur, D. G., Gabriel, S. B., DePristo, M., Purcell, S. M., Palotie, A., Boerwinkle, E.,, Buxbaum, J. D., Cook. Jr, E. H., Gibbs, R. A., Schellenberg, G. D.,, Sutcliffe, J. S., Devlin, B., Roeder, K., Neale, B. M., and Daly, M. J. A framework for the interpretation of de novo mutation in human dis- case. Nat. Genet., (6(9))944-950 September 2014.
Yang. J., Benyamin, B., McEvoy, B. P., Gordon, S., Henders, A. K., Nyholt, D. R., Madden, P. A., Heath, A. C., Martin, N. G., Montgomery, G. W., Goddard, M. E., and Visscher, P. M. Common SNPs explain a large proportion of the heritability for human height. Nat. Genet., 42(7))5655556, July 2010.
Shi, H., Kichaev, G., and Pasaniuc, B. Contrasting the genetic architecture of 30 complex traits from summary association data. Am. J. Hum. Genet., 99(1):139-153, July 2016.
Zhang, Q., Privé, F., Vilhjálmsson, B., and Speed, D. Improved genetic prediction of complex traits from individual-level data summary statisties. Nat. Com- mun., :2(1):4112, July 2021.
Speed, D., Cai, N., UCLEB Consortium, Johnson, M. R., Nejentsev, S., and Balding, D. Reevaluation of SNP heritability in complex human traits. Nat. Genet., 49(7): 986-992 July 2017.
Zheng, Z., Liu, S., Sidorenko, J., Wang, Y., Lin, T., Yengo, L.. Turley, P., Ani, A., Wang, R., Nolte, M., Snieder, H., LifeLines Cohort Study, Yang, J., Wray, N. R., God- dard. M. E., Visscher, P. M., and Zeng, J. Leveraging functional genomic annotations and genome coverage to improve polygenic prediction of complex traits within and between ancestries. Nat. Gener,, 5665):667-777, May 2024.
Spence, J. P.,
N., Assimes, T. L., and Pritchard, J. K. A flexible modeling and inference frame- work for estimating variant effect sizes from GWAS sum- mary statistics. bioRxiv, April 2022.
Theodoris, C. V., Xiao, L., Chopra,
A., Chaffin, M. D., Al Sayed, Z. R., Hill, M. C., Mantinco, H., Brydon, E. M., Zeng, Z., Liu, X. S., and Ellinor, P. T. Transfer learn- ing enables predictions in network biology. Nature, 618 (7965):616-63 May 2023.
12

---

Predicting variant effects from functional annotations
A. Experimental details
A.1. Models
We obtained code for enformer from license. We set the internal dimension 1536 and the number of transformer layers to Our "smaller model" in the ablations reduced the internal dimension 384.
under the MIT
We normalize features to have mean 0 and variance 1 across the genome before passing them any model.
A.2. Generalized linear model
As baseline we consider generalized linear model as suggested in Li al. (2024) using averages of each functional annotation in the window in Finucane al. (2015) (although we only consider single window size):
fo.m (freqm (1 freqm)) ) exp Σ wa Σ
Clunc.m.ddec
Σ Wd'' Cpred,m.dd +
where (wa)d. (w/a )a, and c are learnable parameters.
A.3. Training
We trained our models with A optimizer with default hyperparameters, 100 warmup steps with linear schedule. We use a learning rate of 0.0001 for transformer models and 0.001 for linear and constant models. We train all models for 10 epochs. We train models on single A100 GPUS on an academic cluster; transformer- based models were trained for 15 to 20 hours.
For all models we use calculated using BOLT-LMM from Loh al. (2015).
12 data loader workers in PyTorch (Paszke et al., 2019) Dataloader object. To most efficiently use 10 capacity, each loader loads 100 batches linearly from the same chromosome before switching to different location. To average across this bias, we accumulate gradient steps across 12 steps.
For all models,
For the reduced LD window ablation, for each original (1) window of size 3'000'000 we set (i)+ as a random window of up to 6' 6'000 contiguous variants we chose 6'000 as it is near the average size of (i) in our standard approach. Then we set (i) as all the variants further than 100'000 from the ends of (i)+.
A.4. LD score regression (LDSR)
Finucane et al. (2015) suggested performing the linear LD score regression with a square loss 1) dividing by the (rough) standard deviation of chi-squared variable and 2) downweighting variants LDD with many other variants: calling I 1 2º2 and h3 fe, (ballpark estimate made before training), we minimize (we also multiply numerator and denominator by N)
Σ ( +)² (Nigh/M
M R So + N3² -²
2
A.5. Simulation
Here we describe how we chose realistic f for semi-synthetic simulation. Recall,
y ~~N N(0, N
Therefore
Var(yr.) 0² Xx2
N
Assuming presence of variant Xm,s is independent of Fms we have
M
1 = Var(y:s) = 00
Em Em IX2 m Em Fmal
= σ²
Em FmJ.
N
N
13

---

Predicting variant effects from functional annotations
Thus, in our simulated data, ideally would ensure that
N M (1 (1-
EmFm
In our case, we choose highly heritable disease with 0² 1/2 so half of the variance of is from the noise and the other half genetic. Using real values N 407527 and M 11904424 for our data, we E. fm calculating Fm.m. and defining fm 2M
_N.
by initializing
2M
N Sm-
We defined fm cxp(10 x NNo(Crunc,)) mpred.m) where NN₂ is randomly initialized Enformer model.
B. Additional results
B.1. Biologically realistic ground truth for semi-synthetic experiment
In Sec. we considered a semi-synthetic setting where the ground truth f was a transformer-based model. In this case, when model class is well-specified, recovered the true f. Here, we consider biologically-motivaae ground truth f which designed to be challenging model class is misspecified.
interested in measuring if there is still benefit to using more flexible model if our
We consider function which adds an effect enrichment whenever the sum of functional annotations in window is above a threshold. Let C = [Ctune, Cprea] be the concatenation for Ciune and Cpred in the same order they're 're presented App. D with Cpred broadcasted to the same window size as Cruno. Thus, the ground truth f takes the shape
log f(C)
Σ Val
Σ
Cau > ea - -log M
(E(02,77,12)
wE(118....1483/
where tu = 0.7, 12 = 0.5, 17 = 1.37, 012 = 0.5 and co 0, C2 0, -2 and 612 100 This a challenging f for transformer-based model because the signal is sparse (depending only 4 features), discontinuous, multi-modal and it only depends 32 positions at the center of the window of size 256. We chose the values of ea that feature d 0 is active about & 50% of the time across the genome, d 2 about x 25%, d 7 always active and d 12 about zs 50%. Then, we set the values of ea so that f is multi-modaa with values ranging from 0.072 up to 0.40.
In Fig. we see, even when misspecified, using a more flexible model class leads to a better fit the true function.
LDSR
DeepWAS
log fe error (L)
1
0.5
Constant
Linear
Transformer
Figure 7. DeepWAS using transformer model best recovers the true even in challenging setting. The bars represent the RMSE difference between the learnt fe and the ground truth in the log space evaluated over a set of validation functional annotations.
B.2. A Ablation results in table
We present the numerical results Fig. in Tab.
14

---

Predicting variant effects from functional annotations
Table I. Model Ablation Results: Percentage Likelihond Increase per Person
Model Configuration
BMI
Height
Asthma
Reduced LD window (i)
2.38
10.32
0.358
Reduced Feature window w
2.68
10.43
0.465
Reduced Parameters (reduce hidden dimension)
2.64
10.69
0.492
Reduced Features (remove 15 conservation features)
3.01
10.42
0.633
Full Model
3.02
11.70
0.694
C. Derivations
C.1. Loss derivation
In its more general form, the Woodbury matrix identity (also known as the Aherman-Morrison Woodbury formula) states that (A + UCV)-1 A- A ACCUUCC1 + VA -VV---VVV-VV- and the matrix determinant lemma (also known as the Weinstein-Aronszziy identity) states that + UCV |C]| U
Morriron-Wooduury
Recall that A = (0)+ F ( + R
Applying Woodbury to (A!))- and defining L = Ro)+..0f R
and W BB0 (0) R R/o,00 B(.,60+ we get
)A
R
R RO) ((T())) + R .(0) R (())1()))))
Rot Rott (..0)
N
(1)00
N
N
(1)(0)))))
oN D (,(
N
R
(0,(0)
N
1
(
Po
(1)
xx
F
(1) (>() (0
N
1
R
(1())
1
I
Fo
(1)
)
(1)
Fg 2 1.
N
N
N
Moreover, for | we have
|A!0 = 1 R (1)00)++11) ((F()) -1+
W
1 ((),
|
(
1
1
Fg
(0
w
F0
N
10 Ro,00 1BB
where | 1+ stands for the -ecuud-determmnants. As we discussed before, ()) should be strictly positive-definite but, due to numerical precision, R()),0) could contain eigenvalues close to zero, or negative. This is why we use R (( and why we would only consider de positive eigenvalues for the computation of (o% Ro0
C.2. Held-out chromosome prediction using public data
We essentially show that, since Eqn. derived from Eqn. 1, their likelihoods can only differ by a constant.
The negative log likelihood of y using only held-in chromosomes Xin (Eqn. minus the log likelihood of the null (F 0)
is
1 y (X -
y v'
log Fm + 0²1 log |o²1|
(7)
2
while the same equation for 3 (Eqn. 2) is, defining pscudo-determinants | |
1 2
3 (Rn Fin Rin + o Rim)t
3 R + log Rin Fu Rin + o Riul+ log 1o Rmn J
(8)
15

---

Predicting variant effects from functional annotations
By Woodbury, the first two terms Eqn. are equal to
0l (1 + -o-- 2x Xin) -inn!
(o F: + R13)-13
2
9 (o Fim - Rin Rin Run) Bin
2
31 N Fin + Ran) 20
3 R
20 N
which is equal to the first terms of Eqn. 8. Meanwhile by similar application of the matrix determinant lemma, the last two terms Eqn. are equal the last two terms of Eqn. 8. The same is true replacing X with only held out chromosomes Xout-
D. Data Collection
D.1. UKBB public statisties
We downloaded UK biobank LD matrices computed in Weissbrod et al. (2020) from the Amazon web services S3 container 1U/UKBB LD/. These matrices can have small negative eigenvalues, which we removed prior training. We also, is standard, do not include regions with long-range LD in chromosomes 6, and 11.
We downloaded UK biobank association statistics computed using BOLT-LMM (Loh et al., 2015)
from the UKBB 409K folder in https:// /ttpss://coscle..../ cloud. google. wwt.le..oossttoo wttp//ooo/gggeeeeeeeee www.///ttrageeboowseeeeee. beteggroup-pablc--eqqeee
These association statistics also contained frequencies of each variant. Any variants that have LD information but that are missing associations are discarded; all variants with association information also had LDD information.
UKBB coordinates are in GrCh37 but many of our features below the GrCh38 build. We used rsid's and pyliftOver (https: /guthub...... com/ woonsanntiin coattint/ppllli hpyplllflover/// to map to GrCh38. For the handful of variants we could not map we gave them the location of nearby variant.
D.2. Coding variant annotations
We downloaded the predictions of the effects of variants in coding regions from various models from https: /www. dbnsfp. org/. We used six predictions labeled ESM1 b score, GERP-+ RS, SIFT score, PROVEAN score, FATHMM.score, EVE_score For non-coding variants or variants missing a prediction, we set C== 0.
D.3. Functional and conservation annotation data
Conservation We downloaded big files of our phylogenetic correlation annotations from http: /hggdonloaad soe uese. weweeoluleeeeeee wtl.denaathhhhhhg (Pollard et al., 2010; Hubisz et al., 2011). We used 15 PhyloP and phastCons scores made from various alignments: phylop470wag phyloP20way, phyloP1 ooPlwwy, phyloP Tway, phyloP4way, phastCons470way, phastCons 00way, phastCons30way, phastCons20way, phastCons! 1way,, phastCons7way PoastCons 4way.
phy1oP447way, phyloP100way,
FANTOM We downloaded ICAGE FANTOM annotations of human tissues from https: http://fatttomogo gsc. riken. hptp///aaaubb///gg/eee http://ttm//ma...teee tissue.hCAGE/
hCAGE/ (Lizio et al., 2015). This gave us roughly 400 annotations: we picked a random 20 tissues from this set and collected forward and backward CAGE annotations for each tissue, giving us total of 40 features. The tissues were 1ymph node, adult, donorl; heart, adult, diseased post-infarction. donorl; skeletal muscle, adult, pooll; occipital lobe, adult, donorl; parietal cortex, adult, donor10258 thymus, adult, pooll; thyroid, adult, pooll; pons, adult, pooll; parotid gland, adult; Fingernail (including nail plate, eponychium and hyponychium), donor2; thalamus, adult, donor10258; caudate nucleus, adult, donor10252; parietal lobe, adult, donor10252; cerebrospinal fluid, donor2; kidney, fetal, pooll; eye - muscle
16

---

Predicting variant effects from functional annotations
inffriir rectus, donorl; nucleus accumbens, adult, pooll; parietal lobe adult, donor10196; cerebral meninges, adult; threat, adult.
We downloaded big files of functional genomics annotations httt:s://ww.ennooppojjjoo org/search/ (ENCODE Project Consortium, 2012). We did not use annotations with warnings errors, or that were non-compliant. We used assays with titles TF ChIP-seq, Histone ChIP-seq, eCLIP, total RNA-seq, polyA plus RNA-seq, polyA minus RNA-sed, small RNA-seq, micrORNA-seq, ChTA-PET, WGBS, DNase-seq, ATAC-seq, PRO-cap, PRO-seq, Bru-seq, BruChase-seq; RAMPAGE, PAS-seq and those that had available bigWig files for GrCh38, We got over 100 eCLIP annotations of RNA binding; since each of these annotations are sparse, we summed them together create a single _ll_ecLIP annotation. For TF ChIP-seq experiments that targeted transcription factor, we only used assays from the 24 targets that had measurements from two more labs.
ENCODE
Each of these experiments had multiple data annotations. We used the Fold_change www.ovverooooooooo corr_contoo for random replicate if it was available, otherwise we used a randomly chosen annotation. In total we had 91 annotations from ENCODE; the full list with bioproject ids is as follows: al all.eCLIP, polyA.minus Ranyy..inuus.RN--eee RNA-seq (ENCSRO00CQE), TF TF_ChIP - of MTA3 (ENCSR391KgC), Histone-ChIP-see seq of H3K4me3 (ENCSR393ZCI), TF ChIP-seq of CREB1 (ENCSR897JAS), TF ChIP-seq of MCM3 (ENCSR990A2C), TF ChIP-seq of FOLR2Aphosphes5 (ENCSR000BTW), TF_ChIP-seq of NFIB (ENCSR702BYX) TF TF_ChIP-seq
of SUZ12 (ENCSR757EMK) IF-ChIP-seq of CAMTA2 (ENCSR336GFK), TF-ChIP-seq of NFRKB (ENCSR145BHD) Histone Chetsoe of НЗк3 6me3 (ENCSR524VDR) TF ChIP-seq of SIN3A (ENCSR468LJO), IF-ChIP-seq of PкNOX1 (ENCSR233FAG) TF-ChIP- of
NR3C1 (ENCSR516IVX), TF TF_ChIP-seq of HLTF (#NCSRO90JNM) Histone Mistoone_hhIP-eee of H3к27me3 (BNCSR660TLD), IF-ChIP-seq of NONO (ENCSR476B0A), TF-ChIP-seq of CBX8 (ENCSR61 6MOB), TF ChIP-seq of MCM7 (ENCSRO68VRA), TF ChIP-seq of JUN (ENCSRO48CVK), Histone -hetoeeehhII-eegg of Нзк2зас (ENCSR592JNN), TF-ChIP-See seq of Cebpa (ENCSR827IOM) small RNA- (ENCSRO00CRR), TF TE_ChIP-seq of 2FX (ENCSRO27UFT), TF_ChIP-seq of JUND (ENCSR000RSA), TF-ChIP-seq of MLX (#NC5R873LYH), TF_ChIP-seq of HDGF (ENCSR200CUA), TF TF-ChIP-seq of GATAD2A (ENCSR1600YK), PAS-seq (ENCSR233JPT), WGBS (ENCSR999CXD), RAMPAGE (ENCSR413FKS), TF_ChIP-sed of HDAC1 (ENCSR711VWL), TF-ChIP-seq of DPF2 (ENCSR715CCR), Histone.ChIP-sen of H2AFZ (BNCSR859FGM), TE_ChIP-seq of сТвΩ1 (ENCSR636EYA), TF_ChIP-seq of SMARCA5 (ENCSR895HSJ), TF-ChIP-seg of MNT (ENCSR460XGV), TF-ChIP-seq of BCOR (ENCSR808AK2) TF ChIP-seq of GTF2F1 (FNCSR557377), Histone ChIP-seq of н3K56ac (ENCSRO36NSK) Histome.ChIP-sey of H3К9aс (ENCSR192IRQ), TF-ChIP-seq of JUNB (ENCSR431 3RL)), Histone of H4K20me1 (ENCSR258TUP), TF-ChIP-seq of TRIM24 (ENCSR957LDM), TF_ChIP-seq of PLRG1 (BNCSR019KFC), TF-ChIP-seq of FOXK2
(ENCSR171FUX), Histone ChIP-seq of H3K9me3 (ENCSR444YIP), TF ChIIh-eee of GATAD2B Sestone_ChIP-Be of H2BK15ac (ENCSR739B2R), TF-ChIP-Sed of KHSRP (ENCSR686EYO), TF-ChIP-seq of LARP7 (ENCSR2B8MO2), PRO-cap (ENCSR10OLIJ), TF_ChIP-seq of SMARCA4 (ENCSR5870Q1), TF_ChIP-seq of PHR (#NCSR650AWK) ChIA-PET of POLR2A (ENCSR217IFN) Histone.ChiP seq of Н4K5ac (ENCSR035B2I), TF-ChIP seq of CSDE1 (ENCSR6260JO), Bru-seq (ENCSR849FAX), TF TF_ChIP-seq of DMAP1 (ENCSR670YPO), microRNA-seq (ENCSR934NMC), ATAC-seq (ENCSR2652XX), TF-ChIP-seg of ARNT
(ENCSR389BLX),
(ENCSRO29IBC), ChIA-PET of CTCF (ENCSRO30KAB), TF TE.ChIP-seq of CERPB (ENCSROOOBUR), TF_ChIP FF-ChIP-seq of REXANK (BNCSR823ADL), TF-ChIP-sed of NBN (ENCER210ZYL), TF-ChIP-sed of MLLTI (ENCSR675LRO), TF ChIP-seq of HDAC2 (ENCSR3300E0), TF ChIP-seq of SP1 (ENCSR906PEI), TF TF-ChIP-sed of RBBP5 (ENCSR330EXS), TF-ChIP-sed of MTA2 (ENC5R5512D2), TF-ChIP-seq of PBX3 (ENCSRO0OBTN), Histone.ChIP-sed of
H3K4me2 (#NCSR251DKX), TF PP.hhI--sed of POLR2A (ENCSR388Q7F), IF.ChIP-seq of CBFA2T3 (ENCSR697YLJ), TF FT.ChI--seg of MAX (ENCSROOOBSH), TF PT.ChIP-seq of YBX3 (ENCSR567JEU) TF-ChIP-seq of RAD21 (ENCSROOOBUC), (ENCSR213JMO), DNase-seq (ENCSR366NBE), TELChIP-seq of CTCF (ENCSR817HTJ) Histone ChIP-seq of H3K4me1 (ENCSR874NDX), TF ChIP-seq of TARDBP (ENCSR801 (1SW))
of H3K7 9me1
17

---

Predicting variant effects from functional annotations
TF-ChIP-See seq of FOXP1 (ENCSR369YUK) Tee of IKžF1 (ENCSR278JQG) PRO-seq (ENCSR988B2M) TF ChIP-seq of ZBTBl (ENCSR309ELI), TF ChIP-seq of NCO33 (ENCSR5730UP) , Histone.ChIP-seq of н2вк20ас (ENCSR462XRE) , TF TF_ChIP-sed of SUPTSH (ENCSR894CGX), Histone.ChIP-seq of изк27ас (ENCSRO10OVI), , Bruchase-seq (ENCSR672HUM) TF_ChIP-seq of EP300 (ENCSR686BQM)
18


Reviews:
[{"review_id": "Uw5dMoGCtr", "reviewer": "Reviewer_j54Y", "summary": "The paper introduces WASP, a method for training large-scale neural network models to predict the effects of genetic variants from functional annotations. The paper begins by introducing linear models for variant effects, the need for good priors, and how they can be fit using anonymous summary statistics. This introduces the challenge of LD matrix $R$ inversion. The paper then introduces a batching approach and a preconditioner for computing the likelihood, overcoming the coputational challenge.\n\nAn Enformer model for $f_{\\theta}$ is then trained on large scale functional genomic data from ENCODE, FANTOM, PhyloP and ESM-2. The resulting WASP method is validated using UKBB data against altnerative approaches, demonstrating improved performance.\n\n## update after rebuttal\nThe authors provided additional experiments and clarficiations. My score is unchanged at 4 as this is a good paper that should be accepted.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your support and thoughtful review! Below we discuss new experiments that address the sensitivity to the LD window distance cutoff, larger model sizes and clarify the purpose of Section 4.3.\n\n**On distance sensitivity and assumptions.** We make two distance-based approximations: (1) we break chromosomes into windows of size 1,000,000 and treat observed associations in each window as independent of other associations, and (2) there is no LD for variants further than 1,000,000 positions away from each other.\n\nAssumption (1) is mostly of no concern as WASP would still be a consistent statistical estimator, that is, given enough data the estimator will converge to the correct optimum. Intuitively, this is the case as we train our model using parts of the evidence and treating all the other variables as missing (a proof can be found on the LD score regression paper). Moreover, we do not expect our method to be very sensitive to this parameter, and shrinking the window size should interpolate between the behavior of WASP and LDSR since LDSR is WASP with a window size of 1 (section 4.1).\n\nAccording to prior work we took a very conservative approach to assumption (2) by having no LD variants further than 1M positions away. However, to probe your question we trained a model that incorporates a stronger assumption where variants that are more than 100,000 positions from each window have no LD with variants inside the window. We noted that there is no significant drop in performance on the test set: the likelihood only dropped by 14, while training with LDSR instead of WASP drops the likelihood by 96, and using a linear rather than an Enformer architecture drops the likelihood by 51. Finally, we do find this analysis to be interesting and we’ll expand it on the paper with other window changes.\n\n**On Enformer size.** We picked the size of a network that we could fit with reasonable compute. In principle, training a large-enough model for long enough should result in overfitting. Based on your question and our computational resources, we trained a larger Enformer model on the BMI data (we increased the number of attention layers from 2 to 8) and noted slight improvements on the test likelihood (increased by 11) suggesting that we are not yet overfitting. Indeed it appears that there is more potential in leveraging the huge data collected in large genetics efforts with even larger models. Measuring the scaling laws of these models could therefore be an interesting direction for future work.\n\n**On purpose and contributions in Section 4.3.** Indeed Section 4.3 is a review of iterative methods used for optimizing systems of large matrices. In particular, it introduces the identities for the gradients of the matrix inverse and log determinant; these formulas have been previously introduced in [1,2,3]. Although the construction of the preconditioner in section 4.4 is ours. We’ll clarify our contributions and prior work in section 4.3.\n\n[1] Gardner et al. 2018. GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration.\n\n[2] Saad 2003. Iterative Methods for Sparse Linear Systems.\n\n[3] Saad 2011. Numerical Methods for Large Eigenvalue Problems.", "claims_and_evidence": "The idea for batching turns out to be quite straightforward: chunk the genome into blocks, by assuming block-diagonal LD structure. Then again assume locality, with distant variants not correlating. \n\nBoth assumptions are presented as-is with relevant citations of prior work. I still wonder whether there is any sensitivity to the distances used here?", "methods_and_evaluation_criteria": "Benchmarking approach is adequate, with realistic UKBB data used. Empirical evidence for the used assumptions would be great still.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "- How was the size of the Enformer network optimized? Is this a critical parameter - the experiment shows that making it smaller reduces performance, but what happens in the other direction - is there a risk of overfitting, and if so, how would we detect it.", "relation_to_broader_scientific_literature": "The paper introduces related work and makes a clear case for what limitations WASP seeks to overcome.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "- The paper is very well written and introduces all required concepts to clearly motivate the need for WASP. The method itself is elegant, breaking down a numerically intractable problem into chuks that are amenable to minibatch neural network training approaches.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "- It is unclear to me if there is a contribution in 4.3. Is my understanding correct that this in essence a rationale of what prior work to leverage (without introducing new concepts)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ELx5A3IZyh", "reviewer": "Reviewer_fAUH", "summary": "This paper introduces WASP, a method leveraging accelerated linear algebra to train flexible neural network models for predicting genetic variant effects from functional annotations. By employing banded LD matrix approximations and a structured preconditioner, WASP efficiently handles large-scale genomic data while avoiding costly matrix inversions. Experimental results demonstrate superior performance over LD score regression in predicting phenotypic associations and recovering causal variants. The approach enables training larger models with richer genomic features, advancing functionally informed priors for polygenic trait analysis.\n\n## update after rebuttal.\n      The authors have provided additional results and explanations. I have raised the score accordingly.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed and insightful review. We ran your suggested experiments which have significantly strengthened our paper. \n\n**On semi-synthetic simulations.** You raised a fair point that having a randomly initialized Enformer as a ground-truth model might benefit WASP. Thus, we re-ran all the semi-synthetic simulations but now set a biologically-informed model of inheritance as the ground-truth function. That is, we considered a threshold-based model that whenever the sum of a track in a window around a variant is above a threshold, this track adds a multiplicative increase to the effect size: $$f_{\\theta, m}=\\sum_{d}\\mathbb{1}\\left(\\left(\\sum_{w}C_{\\mathrm{track}, m, d, w} \\right)> \\mathrm{thresh}_d\\right)\\times\\mathrm{enrich}_d.$$\nWe chose the parameters such that $f$ is sparse (mostly only 4 tracks pass the value $\\mathrm{thresh}_d$) and multi-modal (with randomly selected multipliers $\\mathrm{enrich}_d$). Indeed, given the flexibility of the Enformer NN, WASP approximates the correct $f$ much better than other alternatives used in practice and that we compare to as seen in this [figure](https://drive.google.com/file/d/1mf-xaMQnK3znZLDeFI1nIeTT5FDkF4rL/view?usp=sharing). This experiment better highlights the impact of having a flexible neural network model.\n\n**On computational complexity.** WASP is roughly 36% more computationally expensive than LD score regression due to the likelihood computation. However, the discrepancy between WASP and LD score regression is not that large as the cost of the likelihood is not as expensive as the forward and backward pass of the NN that both LD and WASP have to do. Moreover, increasing the model size would decrease the gap. We will discuss this runtime consideration in the paper. Also, to clarify, the goal of section 4 and, in particular, Figure 2, is to discuss alternatives techniques to most efficiently compute the likelihood for WASP.\n\n**On the effect of window size**. We make two distance-based approximations: (1) we break chromosomes into windows of size 1,000,000 and treat observed associations in each window as independent of other associations, and (2) there is no LD for variants further than 1,000,000 positions away from each other.\n\nAssumption (1) is mostly of no concern as WASP would still be a consistent statistical estimator, that is, given enough data the estimator will converge to the correct optimum. Intuitively, this is the case as we train our model using parts of the evidence and treating all the other variables as missing (a proof can be found on the LD score regression paper). Moreover, we do not expect our method to be very sensitive to this parameter, and shrinking the window size should interpolate between the behavior of WASP and LDSR since LDSR is WASP with a window size of 1 (section 4.1).\n\nAccording to prior work we took a very conservative approach to assumption (2) by having no LD variants further than 1M positions away. However, to probe your question we trained a model that incorporates a stronger assumption where variants that are more than 100,000 positions from each window have no LD with variants inside the window. We noted that there is no significant drop in performance on the test set: the likelihood only dropped by 14, while training with LDSR instead of WASP drops the likelihood by 96, and using a linear rather than an Enformer architecture drops the likelihood by 51. Finally, we do find this analysis to be interesting and we’ll expand it on the paper with other window changes.\n\n**WASP on training process.** WASP is involved in the forward and backward pass of the loss computation. Given a mini-batch of $\\beta$ and tracks (where the mini-batching could arguably be considered part of WASP), WASP computes the log determinant and quadratic terms of the loss using fast linear algebra methods. For the backward pass WASP uses a stochastic estimator of the loss. Below we’ve added pseudocode to illustrate these two steps.\n\n// Pseudocode for WASP fwd\n\n1: $\\log(f)$ = NN(geno, anno)  // compute NN output for mini-batch\n\n2: $A = R F R + \\sigma^2 R + \\epsilon I$ // form matrix section 4.1\t\n\n3: $ P $   // compute preconditioner from section 4.4\n\n4: $\\log(|A|) = SLQ(A, P)$ // compute see section 4.3\n\n5: $A^{-1} \\beta = CG(A, P, \\beta)$ // compute see section 4.3\n\n6: $\\ell = log(|A|) + \\beta^T A^{-1} \\beta$  // compute log-likelihood\n\n7: return $\\ell$\n\n// Pseudocode for WASP bwd \n\n1: $u_i \\sim N(0, I)$   // sample M probes\n\n2: $ \\frac{1}{M} \\sum_i \\text{VJP}(u_i, A, A^{-1} u_i ) = \\frac{1}{M} \\sum_i u_i^T \\nabla A (A^{-1} u_i) \\approx \\nabla \\log(A) $ // section 4.3\n\n3: $ \\text{VJP}(-\\beta^{T}A^{-1}, A, A^{-1}\\beta ) = - (A^{-1} \\beta)^T \\nabla A (A^{-1}\\beta) = \\nabla \\beta^{T} A^{-1} \\beta$ // section 4.3\n\n4: return  $\\frac{1}{M} \\sum_i u_i^T \\nabla A (A^{-1} u_i) - (A^{-1} \\beta)^T \\nabla A (A^{-1}\\beta) $ // gradient of $\\ell$", "claims_and_evidence": "The authors use semi-synthetic data to validate WASP, but the data is generated by an Enformer-based model, which biases the evaluation in favor of neural networks and does not prove real-world generations. A stronger validation would involve testing on data generated by different models or real causal variants.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "I have reviewed the formulas and theoretical claims presented in the paper. Based on my examination, I did not find any major flaws or inconsistencies in the mathematical derivations.", "experimental_designs_or_analyses": "The paper lacks a comprehensive model comparison that focuses on both computational complexity and performance. In Figure 2, the authors compare the calculation complexity on WASP and other methods (no preconditioner, Nystrom and Cholesky). However, in Table 1, the authors compare the performance on WASP and LDSR. To provide a more balanced comparison, the authors may consider adding a computational complexity evaluation on LDSR. This would allow them to demonstrate that WASP outperforms LDSR not only in terms of performance but also in computational efficiency, presenting a more holistic view of its advantages across all relevant aspects.", "relation_to_broader_scientific_literature": "By adapting iterative algorithms and structured preconditioners from Gaussian process literature, WASP overcomes the bottleneck of LD matrix inversion to accelerate linear algebra operations.", "essential_references_not_discussed": "I did not identify any major omissions in the reference.", "other_strengths_and_weaknesses": "This paper effectively integrates neural network training with modern fast linear algebra techniques, providing a computationally efficient approach to GWAS data analysis. This combination is a promising direction for accelerating model training and improving scalability in genomic studies. \n\nHowever, the experimental results and analysis can be more diversified and detailed. One potential limitation is the lack of discussion on the impact of sliding window algorithm and the window size on the final result. The method assumes that only variants inside the window are considered, but this choice could slightly influence model accuracy and interpretation. A deeper analysis of how window size - especially in regions spanning different chromosomes or chromosomal structures - affects performance would strengthen the paper’s claim.", "other_comments_or_suggestions": "The writing of the paper still needs to be improved. For example, there is a lack of punctuation: “however there are orders of…” should be “however, there are orders of…”; “Unfortunately this is numerically” should be “Unfortunately, this is numerically”. Typo errors: “we the use well established” should be “we then use well established”.\n\nThe tables in the paper are labeled as “Table 1” and “Table 2”, but in the text they are referred as “Table 6.3”, which might lead to confusion.", "questions_for_authors": "How exactly is WASP incorporated into the model training process?\nThe paper provides mathematical formulas showing how WASP works. However, it does not clearly specify when, where and how WASP is applied in the training pipeline. Is it used during forward propagation, back propagation, loss computation, or another stage? A detailed explanation or a pipeline diagram would greatly improve clarity.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hwLxeTcyA7", "reviewer": "Reviewer_zKNT", "summary": "The paper addresses the challenge of predicting how genetic variants affect phenotypes from large datasets. LD score regression, make simplifying assumptions to avoid computationally expensive linear algebra across genomic metrics. Their method leverages preconditioned linear algebra and GPU acceleration for improved results.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review! In our paper, we train the first large scale models on recently released huge genomics and functional datasets. Below we address your points on additional baselines and extensions of our model.\n\n**On additional baselines.** Note there are no methods that have been able to train models at the scale we do. Even in the original LD score regression paper, they considered a simple model that could be fit with a closed form formula. Therefore we devised a gradient descent method to train a large model using LD score regression. As we show, this unfortunately does not work well, necessitating exact likelihood inference and therefore our development of WASP.\n\nIndeed there are variational inference methods for fitting these models among our citations. However, these models were at best used to fit a handful of hyperparameters, sometimes using grid search; there is no straightforward way to adapt them to scalably fit a large scale model using gradient descent. In particular, this is because one must update the posteriors over the effect sizes over the whole genome before taking a step to update the hyperparameters.\n\n**On whole-genome datasets.** Our mini-batching on the SNP level, similar to SGD, is used to accelerate training. However, our resulting model applies to the whole genome. Note we showed that LDSR makes an extreme version of the WASP mini-batching assumption (Sec. 4.1) and is still used for interpreting the whole genome! \n\nWhat is the connection between models that predict disease from a whole genome and WASP, which is trained on subsets of chromosomes? We begin with a whole-genome generative process that explains the traits of each individual $y$, (Eqn. 1) and we multiply by $X$ to get a generative model of associations on chromosomes $\\hat\\beta$ (Eqn. 2). The likelihood of these two models is only a constant away from each other, so doing well on one is equivalent to doing well on the other. In particular, we explain below how to interpret our metrics in Tables 1 and 2 as our ability to do whole-genome prediction using the held-out chromosomes.\n\n**On more seeds**. We didn’t provide multiple runs as fitting each model is computationally expensive. Here, we ran two more seeds for Enformer trained on BMI via LDSR and WASP. We saw a standard deviation in log likelihood of 14 for LDSR and 8 for WASP, which is much smaller than the differences in the performance of different methods: 96. Nevertheless, we will report error bars for all reported numbers in future drafts.\n\n**On non Gaussian priors**. Although the Gaussianity assumption is by far the most popular in the population genetic literature, in our citations we include references to methods with sparsity-inducing non-Gaussian priors. Applying modern linear algebra techniques like WASP could help us train large models with these priors in future work!\n\n**On metric interpretation**. We take, for each variant $m$ in held-out chromosomes 21 and 22, its observed association with BMI, height, and asthma $\\hat\\beta_m$; then we try to explain this data using our Gaussian model in Eqn 2. The metric is the difference likelihood of these associations under a trained prior $f$ vs a model that assumes there’s no genetic effect on these traits ($f_m=0$ for all $m$). It can therefore be interpreted as how well our trained models explain the associations on held-out chromosomes.\n\nAnother interpretation is the likelihood of how well we could predict traits in the UKBiobank cohort $y$ if we were to only use held-out chromosomes 21 and 22 for our predictions. This is because the models in Eqns 1 and 2 are equivalent. In practice, to do whole-genome estimation, we would train a different model holding out each chromosome and combine their predictions on each held-out chromosome; the likelihood of such a “whole-genome” model would be the sums of likelihoods like those in tables 1 and 2.\n\nThis is why the numbers we show are not in the regular BMI range. We will elaborate on this explanation in the paper.", "claims_and_evidence": "Yes. This paper is out of my area of expertise, so I may not be the best judge.", "methods_and_evaluation_criteria": "Performance is evaluated using semi-synthetic data (controlled experiments) and real-world GWAS data (UK Biobank).\n\nThe comparison is mainly with LDSR, there could be other baselines potentially, like vae? Is this part of the literature?", "theoretical_claims": "The maths makes sense.", "experimental_designs_or_analyses": "The authors perform experiments on semi-synthetic data and real data. They are well motivated, however variances are missing. Some of the results can be very close, running multiple seeds might allow for better comparision", "relation_to_broader_scientific_literature": "The contributions of this paper fit into the broader landscape of genome-wide association studies (GWAS) and Bayesian modeling.", "essential_references_not_discussed": "na", "other_strengths_and_weaknesses": "S1. The authors aim to tackle a realistic problem statement and offer solution in O(n^2) as compare to slower O(n^3).\n\nS2. Instead of using precomputed LD scores, WASP directly optimizes the marginal likelihood, which theoretically leads to better parameter estimates.\n\nW1. They use chromosome-level mini-batching, but it's unclear how it performs on whole-genome datasets.\n\nW2. The paper is missing uncertainty quantification.", "other_comments_or_suggestions": "- The different metrics such as BMI, Athama, etc do not have an introduction in the paper. Some context might be more helpful to better understand the significance of the numbers. To a layperson (aka me, I thought BMI is supposed to be between 18-25 but the numbers are much higher here)", "questions_for_authors": "1. The paper is motivated for genomic studies, however they mainly focuses on chromosome-level analysis rather than whole-genome scalability. Is this a good psuedotask?\n\n2. Are non gaussian priors explored?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "o2MMYhKxXs", "reviewer": "Reviewer_crZx", "summary": "This paper proposed WASP for phenotype prediction based on the genomic data, which utilizes an iterative algorithm and an approximation inverse to optimize the goal efficiently.  Experimental results on GWAS show its superiority over LD Score Regression.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. In our paper, we train the first large scale models on recently released huge genomics and functional datasets, enabled by modern linear algebra techniques. Below we address your points on additional baselines and downstream uses of our model.\n\n**On weak baselines**. LD score regression is the de facto framework that is currently used for modeling genetic variability due to its tractable computational demands. Its simplest version – where f is assumed to be constant --  was indeed published ~10 years ago as we cited, but there has been continued research on it ([example](https://doi.org/10.1101/2025.03.07.25323578)). \n\nIndeed, no prior work has been able to train at the scale that we have! Even in the original LD score regression paper, they considered a simple model that could be fit with a closed form formula. Therefore, we devised a gradient descent method to train a large model using LD score regression. As we show, this unfortunately does not work well, necessitating exact likelihood inference and therefore our development of WASP. \n\nThe fact that there is no better alternative method to LD score regression highlights the importance and impact of our contribution.\n\nNext, we clarify that the methods in lines 179 and 212 are not competitors to WASP. The methods in line 179 take pre-trained functionally-informed priors and use them to improve genetics tasks – they therefore complement our development of WASP, which is a method for pre-training functionally-informed priors.\nMoreover, we build methods for fast linear algebra operations on large LD matrices. The methods on line 212 include other methods for speeding up linear algebra problems in genetics; genetics is a broad field and these techniques do not solve the problems involved in functionally-informed priors. In particular, they approximate $R$ without necessarily describing how to invert the matrix in the WASP likelihood, $A_\\phi^{(i)}$. As an example, we cite Berisa, T. and Pickrell, J. K 2016 who discovered the block-diagonal approximation of $R$ to learn about biophysical processes in the cell. We adapted their discovery to devise our mini-batching strategy for WASP.\n\n**On one dataset.** UKBiobank is the only dataset that satisfies two considerations: 1) the summary association and LD data is public while other databases such as China Kadoorie Biobank are not, 2) UKBiobank is a huge dataset of hundreds of thousands of individuals that we expect to benefit from a larger model. Put differently, UKBiobank is a foundational dataset similar to OpenWebText. In contrast to vision models, where there are a myriad of datasets (CIFAR, SVHN, ImageNet, etc), the genetics datasets are naturally much larger and most publications focus on fitting to the largest amount of data available.\n\n**On informed priors and disease prediction accuracy.** The improved posterior that WASP provides can be used for a myriad of downstream tasks such as disease prediction (simply take a MAP estimate of the effect of each variant) but also for causal variant identification, or to interpret the importance of the diverse functional annotations that we used (Phylo, Encode, Fantom). Indeed the goal of WASP is to solve the machine learning problem of fitting the data in Eqn2 as accurately as possible and, in our previous work section, we highlight some of the many efforts that use the functionally informed priors (that result from solving the prediction problem) for downstream applications that we just mentioned.\n\n**On lack of framework and theoretical convergence.** Succinctly, WASP proposes a loss function Eqn 2 and a computationally efficient way to train the loss in section 4. The approximation that we used for mini-batching has been well studied as cited [Berisa & Pickrell, 2016 & Salehi Nowbandegani et al. 2023) and all the convergence of our linear algebra techniques is well-understood (Saad 2011 or Hobgen, 2013). Additionally, WASP is a consistent estimator as explained on the first point [here](https://openreview.net/forum?id=oOtdWiLb1e&noteId=Uw5dMoGCtr). Thus WASP is a principled method for computing the likelihood in Eqn 2. If your theoretical convergence concern is about our use of neural networks then that is an empirical question that we address with the performance on held-out data.", "claims_and_evidence": "Not very clear\n1. The introduction notes that researchers in phenotype prediction typically emphasize incorporating functionally informed priors. However, my understanding of this paper is that its primary focus is on improving the optimization efficiency of e.q. (2) and (3). I did not observe any explicit methods for integrating or modeling such priors.\n2. The paper mentioned that LD Score Regression is the SOTA method, which was published around 10 years ago. In the related work, there are some improved approximation methods in line 212 and some methods that utilize these priors in line 179. Did they improve the performance for phenotype prediction?", "methods_and_evaluation_criteria": "For baselines, the only use of LD Score Regression is too weak. The paper mentioned that a bunch of approximation approaches, they should be considered in verifying the efficiency of the WASP.", "theoretical_claims": "There are no theoretical guarantees.", "experimental_designs_or_analyses": "The baselines are too weak.", "relation_to_broader_scientific_literature": "The", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Weakness: \nThere is no clear framework of the proposed WASP. \nNo theoretical guarantee for the convergence.\nOnly one dataset and one baseline.", "other_comments_or_suggestions": "No other comments. The paper would benefit from clearer organization.", "questions_for_authors": "For this paper, I am still not clear how to involve the informed priors to improve the accuracy in prediction?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_flag": "Flag this paper for an ethics review."}}]

Author rebuttal:

Meta-review: The majority of the reviewers assigned a weakly positive score to this paper. They all seem to agree that the paper is clearly motivated and well-written, that a relevant problem is studied and that the solution provided in the form of a combination of neural networks with fast linear algebra techniques is interesting and technically sound. There were also several more negative comments, such as the some experimental design choices (effect of sliding window, use of chromosome-level mini-batching etc.), missing convergence guarantees and missing uncertainty quantification. Some of these concerns, however, could be addressed in the rebuttal (at least to some extent). In summary, I think that although this paper is still in the "borderline region", to me the positive aspects regarding relevance and technical soundness slightly dominate to points of criticism.
Decision: Accept (poster)