Title: PRSformer: Disease Prediction from Million-Scale Individual Genotypes
Venue: NeurIPS 2025
Accepted: yes
Keywords: Multitask Learning

Transformer Architectures

Neighborhood Attention

Deep learning for Healthcare
Abstract: Predicting disease risk from DNA presents an unprecedented emerging challenge as biobanks approach population scale sizes ($N>10^6$ individuals) with ultra-high-dimensional features ($L>10^5$ genotypes). Current methods, often linear and reliant on summary statistics, fail to capture complex genetic interactions and discard valuable individual-level information. We introduce **PRSformer**, a scalable deep learning architecture designed for end-to-end, multitask disease prediction directly from million-scale individual genotypes.  PRSformer employs neighborhood attention, achieving linear $O(L)$ complexity per layer, making Transformers tractable for genome-scale inputs. Crucially, PRSformer utilizes a stacking of these efficient attention layers, progressively increasing the effective receptive field to model local dependencies (e.g., within linkage disequilibrium blocks) before integrating information across wider genomic regions. This design, tailored for genomics, allows PRSformer to learn complex, potentially non-linear and long-range interactions directly from raw genotypes. We demonstrate PRSformer's effectiveness using a unique large private cohort ($N \approx 5$M) for predicting 18 autoimmune and inflammatory conditions using $L \approx 140$k variants. PRSformer significantly outperforms highly optimized linear models trained on the *same individual-level data* and state-of-the-art summary-statistic-based methods (LDPred2) derived from the *same cohort*, quantifying the benefits of non-linear modeling and multitask learning at scale. Furthermore, experiments reveal that the advantage of non-linearity emerges primarily at large sample sizes ($N > 1$M), and that a multi-ancestry trained model improves generalization, establishing PRSformer as a new framework for deep learning in population-scale genomics.


Peer reviews:
comments: SUMMARY:
The authors introduce PRSformer, an efficient Transformer-based model designed for disease risk prediction. On their autoimmune disease benchmark, PRSformer demonstrates superior performance compared to LDPred2 and a linear baseline. The study also explores the advantages of Transformer non-linearity and its potential for cross-ancestry generalization.

STRENGTHS AND WEAKNESSES:
Strengths

0. The empirical results are promising.

Weaknesses
1. Limited technical novelty: The performance advantage of Transformers over linear models or statistical approaches is well-established across multiple domains, making the core contribution unsurprising.
2. Lack of biological insight: Despite improvements in benchmark metrics, the work does not explore the biological implications of the results. It remains unclear what biological patterns or interactions PRSformer captures that are missed by traditional methods.
3. Constrained evaluation scope: The evaluation is limited to the in-distribution setting. It is uncertain whether PRSformer generalizes to unseen disease traits or can serve as a foundation model across different tasks.
4. Architectural constraint: The use of neighbourhood attention based on chromosomal position restricts the model to learning only local genic interactions, potentially overlooking long-range dependencies.
5. Unaddressed practical trade-offs: The paper does not compare PRSformer with LDPred2 and linear models on training/inference cost or interpretability, where simpler methods may still hold advantages.

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-06-19

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper applies a transformer deep learning model to biobank data, learning to predict disease risk from extremely high dimensional genomic input.  The primary goal is to avoid the assumptions of linear models and models that are based on summary statistics.  The method works by using a dilation-style approach to capture long-range dependencies without undue computational cost.  Empirical results using a large, private cohort suggests that the model substantially outperforms state-of-the-art methods.\n\nSTRENGTHS AND WEAKNESSES:\nQuality\n\nThe model design is sensible, particularly the application of neighborhood attention.  The novelty here is not so much on the ML side but in the application of these methods to this particular problem.\n\nI appreciate the careful attention to avoiding leakage during the feature selection step.\n\nThe experiments are very well designed and executed, and I appreciated the logical flow from the initial tests, to the learning curve, to multi-modal analysis, to cross-ancestry analysis.\n\nClarity\n\nThe paper is very clear.\n\nSignificance\n\nThis paper addresses one of the key questions in personal genomics and as such is quite significant.\n\nOriginality\n\nI don't think the paper is particularly original.  Anyone with access to this dataset and a background in ML would likely come up with a similar idea.  But it's still important work.\n\n\nMinor: What statistical test was used for the p-values in Figure 2.  We are only told that it is one-sided.\n\nQUESTIONS:\nHow did you ensure that your hyperparameter tuning did not end up biasing your results in your favor on this particular dataset?\n\nI am surprised that you opted not to include positional encodings, since the NA only captures position with a resolution of ~100kb.  Did you try using positional encodings, and how well did it work?\n\nDid you do any filtering to remove close relatives between the train and test sets?\n\nHow many parameters in your model, and what is the evidence that it converged after only two epochs?\n\nLIMITATIONS:\nA limitation here is the racial diversity of the cohort, which is limited to European genetic ancestry, though this is somewhat mitigated by the experiments in which this cohort is combined with data from other ancestral groups.\n\nAnother big problem is that the data is all private, so there is no way for anyone else to reproduce and validate the claims made here.  I think the statement in the paper about why the data is not public -- \"Recognizing the sensitive nature of genomic data, the dataset cannot be shared publicly\"  -- is misleading.  There are mechanisms available for making genetic data public but with controlled access.  I assume that there are IP or other issues that preclude making this particular dataset public via such a mechanism.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces PRSformer, a Transformer-based architecture using local attention for multi-task disease prediction. The authors evaluate PRSformer on a unique large-scale private dataset of 5 million individuals, far exceeding public cohorts, and a set of 18 autoimmune diseases. The main finding is that this non-linear architecture outperforms linear baselines only when a large dataset is available, beyond the scale of the UK Biobank. All models are trained and tested on a split from the same cohort. The authors also train a model on a multi-ancestry cohort and show it improves performance in non-European populations.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n* To my knowledge, this is the first time a deep learning genetic risk model has been applied to a dataset of this massive scale.\n\n* The key finding that large datasets are required for non-linear methods to outperform linear baselines is an important contribution to genomics and deep learning based genomic risk prediction.\n\n* Using LDPred2 with summary statistics from the same cohort is a convincing approach to test the authors' main claim.\n\n* Local attention is a reasonable architectural choice to handle the high dimensionality of genomic data.\n\n* The writing is clear and easy to follow and the method is well motivated. \n\nWeaknesses\n\n* No external validation: All experiments were only performed on their own private dataset. A validation of how well these models perform in a cohort such as the UK Biobank would be appropriate to assess model generalization and compare with previous literature.\n\n* No Control for Kinship: The paper does not mention any control for cryptic relatedness between the training and testing sets. In genetic studies, this is a critical step, as failure to remove relatives can lead to inflated and misleading performance metrics due to the model learning family-specific markers instead of generalizable disease signals.\n\n* In both the PRSformer model and the linear baselines, based on the massively inflated parameter counts in Table E10 and the illustration in Figure 1, it appears that the authors simply flatten the sequence of embeddings (dimension L, d) before the final linear layer, such that the input of the linear layer has dimension L x d. This is a very unconventional modeling choice as it massively increases the number of parameters and increases risk of overfitting. For the Transformer architecture, a more appropriate choice would be to use mean pooling or a dedicated CLS token. For the linear baseline this choice is even more puzzling, as this model simply reduces to a massively overparameterized logistic regression model.\n\n* No LD Pruning in GWAS based variant selection: This likely leads to strongly correlated variants in their selected variant sets.\n\n* The fixed-size attention window (k=385) does not account for actual genomic distance or LD structure. Because variant selection did not include LD pruning, many of the 385 variants in a window could be in high LD and physically very close, while others are far apart. The model cannot distinguish these cases and may even attend across chromosome boundaries where the input sequences are concatenated. A more biologically appropriate approach might be to define attention dynamically based on variable-sized LD blocks.\n\n* The main models were trained without including crucial covariates like age, sex, and genetic principal components. This is problematic as it can introduce confounding from population structure and other biases. While some results with covariates are in the appendix, genetic PCs are essential and were not included at all.\n\n* Ablation Studies Needed: Several key design choices lack empirical justification through ablation studies:\n\n    * Does explicitly modeling unknown (UNKN) genotypes with a separate embedding actually improve performance?\n    * What is the impact of omitting standard positional encodings?\n\n* The argument that summary-statistic methods are limited by LD reference panel mismatch seems weak. Since the proposed method requires individual-level data from the cohort, a perfectly matched LD panel can be calculated and used.\n\n* Inductive Bias of Attention: The paper claims attention provides a \"powerful inductive bias for genomics\" but does not sufficiently justify why this is the case, and should be compared to more biologically-informed inductive biases (e.g., using Gene Ontology, gene-gene or protein-protein interaction networks).\n\n* Insufficient Comparison to relevant baselines:\n\n    * It does not compare against meta-GWAS based PRS methods, which are the primary way the field achieves million-plus sample sizes with summary statistics. One large advantage of summary stats based PRS methods is that they can operate on meta-GWAS, whereas their method does require individual level data. This limitation should at the very least be discussed.\n\n    * Especially given some of the more questionable modelling choices, a comparison to a simple MLP on raw dosages baseline would be interesting to see.\n\n    * A comparison to established linear models that also make use of individual level data (BOLT-LMM, GEMMA, ..) would be appropriate.\n\n    * Phenformer (https://arxiv.org/abs/2501.07737) would be a very interesting deep learning baseline to test against.\n\n* Loss Function: The loss function is presented as a sum over the entire dataset, omitting the mini-batching used in practice.\n\n--- \n\nIn summary, I believe the main finding of this work is very exciting, but probably more suited to a genomics journal rather than a ML conference. The method the authors propose is reasonable, apart from the problems discussed above, but does not advance the field of ML.\n\nQUESTIONS:\n* Validation and Confounding: Can the authors comment on why standard and essential validation steps like external cohort validation and kinship control were omitted? A positive response would require, at minimum, an analysis of kinship in the dataset and a discussion of its potential impact.\n\n* Linear Baseline and Output Aggregation: Can the authors justify the unconventional choice of flattening the embedding matrix for the linear baseline, which makes it an overparameterized logistic regression? To strengthen the paper's claims, I would recommend re-running this baseline with a standard, efficient architecture.\n\n* Genomic Input Representation: The fixed-size attention window seems to ignore LD structure and chromosome boundaries. Can the authors comment on the limitations of this approach and the feasibility of using a more biologically informed attention mechanism based on LD blocks?\n\n* Imbalance Handling: The definition of T(i) in the loss function is ambiguous. Could the authors clarify if it refers to diseases with a known status or only diseases a person has, and explain how summation \"implicitly handles imbalance\"?\n\n* Figure 3 Interpretation: In Figure 3, why does the non-linear PRSformer model outperform the linear baseline at the smallest data size (0.15M), which seems to contradict the main finding? Also, why is the LDPred2 baseline not shown for the down-sampled data?\n\n* Overfitting and Regularization: The appendix suggests overfitting is a significant issue. Was any regularization used beyond weight decay? Given that a smaller NA window size performed nearly as well as the best parameters, were even smaller window sizes tested?\n\n* Why was no cross-validation performed? Given the rather small increases in performance for most diseases, this would increase trust in the findings compared to the bootstrap based significance testing\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work introduces PRSformer, a scalable supervised transformer-based model for predicting disease risk from genotype. The scalability of PRSformer lies in the use of linear neighborhood attention layers. When supervised on large (N > 1M, L > 100K) datasets to predict 18 diseases, it outperforms traditional state-of-the-art methods on test data.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1) Improved performance over state-of-the-art methods on large (>1M individuals) datasets.\n2) Shows benefits of multitasking for disease risk prediction, for both PRSformer and linear models (Figure 4).\n\nWeaknesses:\n1) Performance seems slightly worse than linear model on smaller datasets (<1M) for some diseases (Figure 3).\n2) No insights offered into the improved performance of the model -- only the hypothesis that it must come from the non-linear effects. I am inclined to believe this, but empirical evidence that the model has learnt any interactions at all would be helpful.\n3) No code provided for review.\n4) Not necessarily a weakness (arguably a strength!), but this seems to me more of an application paper where an (arguably simple?) deep model is applied to a new (very relevant) task and outperforms traditional models. The secret ingredient seems to be the (proprietary) data. Because of this I am rating the originality as \"fair\" but the significance as \"excellent\".\n\nQUESTIONS:\n1) As seen in Figure 3, why do you think PRSformer is sometimes doing slightly worse than the linear model on smaller datasets (<1M individuals)? Is PRSformer overfitting? How many parameters does it have compared to traditional (e.g. linear) methods, and how does regularization compare, if at all?\n2) Following up on the question above, I'm curious how Figure 3 looks like for the training data: is there evidence that PRSformer is overfitting in the smaller data regime?\n3) I can understand the data not being public and the code release being promised upon acceptance -- but why have you not provided the code *for review*? It seems to me like you only hurt your submission by doing this.\n4) Regarding the benefits of training on a diverse cohort for generalization: is it not expected that performance will go up by doing this? Morevoer, this is not specific to PRSformer: it is expected about any other model? I am puzzled about the contribution here. It would make more sense if the point were that PRSformer benefits *more than traditional methods* by using a diverse cohort -- but this is not what is being shown.\n5) What non-linear interactions is PRSformer learning? Do you have any concrete examples? Are there any known ones (from, say, domain knowledge), and does the model recover them?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 2, "DATE": "2025-07-03"}