PRSformer: Disease Prediction from Million-Scale Individual Genotypes

Response to Reviewer f3oc

We thank Reviewer f3oc for their feedback. We believe the concerns about novelty and significance stem from a framing issue. Our core contribution is not simply applying a Transformer, but the introduction of a principled, scalable architecture that reveals the fundamental relationship between genetic data scale and model complexity, which has never been accomplished before in this high-impact domain.

1. On Novelty and Significance

Limited technical novelty: The performance advantage of Transformers over linear models... is well-established... making the core contribution unsurprising.

We respectfully disagree that the contribution is unsurprising. While Transformers outperform linear models in many domains, it was a completely open and highly debated question in genomics whether this held true, especially given the extraordinary strength of linear baselines. Our work is the first to resolve this debate by discovering the fundamental scaling law governing this transition. The key question was never if Transformers could outperform linear models, but under what conditions, at what scale, and by how much.

Our contributions are:

• A Scalable Framework: We developed the first viable architecture to successfully tackle disease prediction from million-scale individual-level genotypes.

• The Discovery of a Fundamental Scaling Law: Our key finding that the advantage of non-linear models only emerges at a million genetic data scale is a non-obvious and impactful ML insight. This was impossible to discover previously and provides a crucial guideline for the genetics, population health and precision medicine fields.

We view this discovery as a central contribution of our study.

Action: We will revise the Introduction to state this contribution more forcefully, framing the paper around the discovery of this fundamental scaling principle.

2. On Biological Insight

Lack of biological insight... neighbourhood attention... restricts the model to learning only local genic interactions.

We would like to highlight two key insights our work provides, which were only discoverable at this scale, as well as clarify our architectural choices.

• Biological Insight: Our work provides two key high-level biological insights: 1) we quantitatively demonstrate that shared genetic signals exist across related diseases (evidenced by the success of multitask learning), and 2) we show that multi-ancestry training improves predictive equity across diverse populations (Table 1).

• Architectural Constraints: The neighborhood attention is not a limitation but a deliberate choice that makes learning from million-variant inputs computationally tractable. Crucially, by stacking these layers, the effective receptive field grows, allowing the model to learn longer-range dependencies. The fact that this simple, scalable bias is sufficient to unlock significant performance gains at scale is itself an important finding.

Action: We will add a discussion on these computational trade-offs to the Methods section.

3. On Evaluation Scope and Practical Trade-offs

Constrained evaluation scope... Unaddressed practical trade-offs...

• Evaluation Scope: We agree that PRSformer is a supervised prediction model, not a foundation model for unseen diseases. Our goal was to establish a new state-of-the-art for the specific, high-impact task of polygenic risk prediction.

• Practical Trade-offs: This is a fair point. There is a classic trade-off between PRSformer's higher computational cost and its significant gain in predictive accuracy. For clinical applications where accuracy is paramount, this trade-off is often acceptable. Interpretability is a key limitation we discuss in our Future Work section; our paper's focus was to first prove that a predictive signal worth interpreting exists at scale.

Action: We will add a brief sentence to the Discussion acknowledging the computational cost trade-off, and implications for clinical practice.

We hope this response clarifies why we believe our work is more than a simple application paper, but offers a foundational insight for the ML community. Our choice to submit to NeurIPS was deliberate. The goal is to introduce this high-impact domain to an interdisciplinary ML community with the rigor it deserves, providing the first scalable baseline and discovering a fundamental scaling law that will guide future research. We hope that these clarifications help convey the paper’s core ML novelty and address the reviewer's concerns more directly.

Response to Reviewer 8Lrv

We are grateful to Reviewer 8Lrv for their exceptionally positive review and for finding our work "sensible," "well designed," and the experiments "very well designed and executed." We appreciate the opportunity to clarify the few remaining questions and reinforce the core contributions that we believe make this work a strong fit for NeurIPS.

1. On Experimental Details

What statistical test was used for the p-values in Figure 2?

Thank you for pointing out the ambiguity. We used a one-sided paired bootstrap test (10,000 replicates). For each replicate, we sampled with replacement from the test set and computed the AUROC for both PRSformer and LDPred2 on the same sample pairs. The p-value is the fraction of replicates where LDPred2's AUROC was greater than or equal to PRSformer's.

Action: We will add this clear definition to the Figure 2 caption.

How did you ensure that your hyperparameter tuning did not end up biasing your results?

This is an excellent point. We followed standard ML best practices by performing all hyperparameter tuning on a single, dedicated validation set, which was kept entirely separate from the final test set. This standard procedure is designed to prevent overfitting to the validation set (e.g. by avoiding extensive optimization of hyper-parameters) and ensures that our final reported performance is an unbiased estimate of generalization.

Action: We will add a sentence to the Methods to clarify this standard practice.

2. On Design Choices

I am surprised that you opted not to include positional encodings... Did you try using positional encodings?

Yes, we experimented with standard sinusoidal positional encodings but found they offered no performance improvement. We hypothesize this is because the fixed, genome-wide ordering of variants already provides sufficient implicit positional information for the neighborhood attention mechanism to leverage, making explicit encodings redundant.

Action: We will add a brief sentence to the Methods clarifying this design choice.

Did you do any filtering to remove close relatives between the train and test sets?

We did not filter for relatives. This was a deliberate choice that aligns with our work's focus on prediction, a point also discussed with Reviewer 8raR. While filtering relatives is critical for GWAS discovery studies, our goal was a fair predictive comparison between architectures. Since all models were trained and tested on the exact same data splits, any potential effect from relatedness is consistent across all models. This ensures our reported relative improvements remain robust and fairly evaluated, which is the central claim of the paper.

Action: We will clarify this prediction-focused rationale in the Methods.

3. On Model Scale and Data Access

How many parameters in your model, and what is the evidence that it converged after only two epochs?

• Model Parameters: The main PRSformer model has ~180M parameters. A detailed breakdown of the other model’s variants is provided in Supplementary Table E10.

• Convergence Evidence: We treated the total number of training epochs as a hyperparameter and tuned it based on validation AUROC. As shown in Supplementary Table E5, validation AUROC consistently peaked at 2 epochs, with performance declining thereafter (indicating overfitting). This provides clear empirical evidence that 2 epochs was the point of optimal convergence for generalization.

Action: We will add a sentence in the main text referencing Table E5 as evidence for our chosen training duration.

Another big problem is that the data is all private... I think the statement... is misleading.

Thank you for this feedback. We agree our wording could be more precise. You are correct that there are mechanisms for controlled access to sensitive data. The primary barrier in our case is related to the specific terms of our data use agreement, which currently preclude sharing via such public mechanisms.

Action: We will revise the "Limitations" section to state that the data cannot be shared publicly due to the terms of our specific data use agreement, removing the more general statement.

We sincerely thank you again for your strong support and for helping us improve the paper. Your positive assessment of our work as "quite significant" and as addressing "one of the key questions in personal genomics" reinforces our belief that this paper will be an impactful contribution to the NeurIPS community.

Response to Reviewer 8raR

We sincerely thank Reviewer 8raR for providing an exceptionally detailed and thoughtful review. Your comments resulted in lively discussion amongst the authors. The questions you raise are especially important as they highlight the fundamental distinction between the goals of GWAS (discovery of genetic associations) and large-scale ML prediction. Our work is firmly framed within the ML prediction paradigm. We hope our response clarifies this crucial distinction.

1. On Validation, Kinship, and LD Pruning

No external validation...No Control for Kinship...No LD Pruning

These are all critical best practices for a GWAS discovery study. However, for a pure prediction-focused ML study, they are either unnecessary or detrimental to performance.

• LD Pruning is Information-Destroying for Prediction: In GWAS, pruning is essential to isolate independent signals. For prediction, it discards valuable information. We confirmed this experimentally: applying LD pruning at standard thresholds consistently and substantially reduced predictive accuracy on our validation set. PRSformer is designed to leverage the complete polygenic signal, which resides in correlated LD blocks.

• Kinship's Effect on Relative Performance is Minimal: While kinship can inflate absolute performance, its effect on a comparative study is negligible. Since all models (PRSformer and baselines) were trained and tested on the exact same data splits, any inflation from relatedness is present equally across all models. Therefore, the relative performance gains of PRSformer over the baselines remain robust and fairly evaluated.

• External Validation Scope: We agree this is an important future direction. Our paper's central goal is to conduct a controlled comparison of model architectures. Validating on an external cohort like UKBB would introduce confounding variables (eg. different phenotyping standards, etc) that would make it impossible to isolate the effect of the model architecture itself.

Action: We will add a section to Methods clarifying our prediction-focused rationale and include results from our negative LD pruning experiments.

2. On the Flattened Output Layer and Baseline

...simply flatten the sequence of embeddings... This is a very unconventional modeling choice...

This is an excellent question. We did in fact test standard alternatives and found our "flattened" architecture superior for this prediction task.

• Ablation of Pooling/CLS token vs. Flattening: We evaluated both mean pooling and dedicated CLS tokens against our current architecture.

  • We tested an average pooling layer after the last transformer block to aggregate embeddings across tokens (converting B×L×d tensors to B×d), followed by a linear layer mapping d to the output dimension.
  • We tested adding 1, 10, and 20 CLS tokens to the input sequence and vocabulary, each with learnable embeddings and full attention to all variant tokens.

• The flattened embedding approach consistently outperformed these alternatives across nearly all diseases on our validation set. We hypothesize that mean pooling dilutes signals from large-effect variants, while CLS tokens are difficult to train effectively without massive unsupervised pre-training (a promising future direction for a genotype foundation model).

• Linear baseline rationale: We note that our linear baseline was designed as a direct architectural ablation of the Transformer blocks, not as a vanilla logistic regression. It includes a learnable embedding layer, which already introduces non-linearity into the parameterization. A standard logistic regression with one-hot inputs performed substantially worse. The comparison fairly isolates the benefit of our stacked neighborhood attention layers.

Action: We will add a concise paragraph to the Methods summarizing these ablation results to justify our architecture.

3. On Other Technical Points

Fixed-size attention... No covariates... Imbalance Handling... Figure 3... No cross-validation...

• Attention Window: Our architecture compensates by stacking these neighborhood attention layers, which hierarchically increases the effective receptive field, allowing the model to integrate information across multiple local blocks and thus learn longer-range dependencies. A dynamic, LD-aware attention mechanism is an excellent idea for future work. In this paper, our goal was to establish a strong, scalable baseline that balanced performance with computational feasibility.

• Covariates (Age, Sex, PCs): Omitting covariates was essential to isolate the specific contribution of the genotype model itself, which is the focus of our paper. In the future, an "all-in" clinical prediction model would certainly include these covariates, but this would confound the evaluation of the underlying PRS model itself.

• UNKN Embedding & Positional Encodings: We assert that a separate UNKN embedding lets the model learn from missingness and is more flexible than simple mean-imputation. We tested and found that standard sinusoidal positional encodings offered no benefit, likely because the fixed genome-wide order provides implicit position.

• On LD Mismatch: Our point in the introduction is a more general one about the limitations of summary-statistic methods in practice, where researchers often must rely on external, potentially mismatched reference panels. Our central claim is that even when a SOTA method is given a matched LD panel, our end-to-end, non-linear architecture still provides a significant performance gain. This isolates the advantage to PRSformer's ability to learn from individual-level data and capture non-additive effects.

• Loss Function (T(i)): T(i) refers to the set of diseases for which individual i has a known (non-missing) status. By summing, rather than averaging the loss, tasks with more samples naturally contribute more to the gradients, implicitly handling the imbalance between disease prevalences. We also tested focal loss, task‐uncertainty weighted loss (Kendall et al., 2018) and standard averaged cross entropy --all of which were outperformed by the current PRSformer loss function based on validation AUROCs.

• Figure 3 Interpretation: The apparent outperformance at 0.15M is likely due to variability from using a single downsampled dataset. We expect this difference to diminish if the analysis were repeated across multiple replicates.

• Regularization & Overfitting: We tuned weight decay based on validation AUROCs. Also, we found dropout offered no benefit. For the NA window size, 129 was the smallest we tested, and as the linear baseline (equivalent to a window size of 0) performed worse, this highlights the benefit of NA.

• Cross-Validation: Given the immense computational cost of training on this scale, we used a single, large, temporally-split validation set instead of K-fold CV, which is a standard practice in large-scale deep learning.

Action: We will add brief, clear justifications for these design choices in the Methods and Appendix. We also will revise the introduction to cite: Vilhjálmsson et al., Am J Hum Genet 2015; Weissbrod et al., Nat Genet 2021 demonstrating that reliance on mismatched LD panels can introduce bias in PRS models.

4. On Other Methods

• BOLT-LMM / GEMMA: These are excellent tools for GWAS discovery (finding significant variants), not for polygenic prediction. Our chosen baseline, LDPred2, is the established state-of-the-art for the specific prediction task we address.

• Phenformer: Thank you for highlighting this exciting concurrent work. The key difference is scale and scientific finding. Phenformer was applied to ~150k individuals. Our work, at the vastly larger scale of ~5M individuals, allowed us to uncover the fundamental scaling law that the benefit of these complex non-linear architectures only becomes apparent when N > 1M. This is a core conclusion that was impossible to draw from previous studies and represents our primary contribution.

• Gene Ontology / interaction networks: We made a deliberate choice to first establish a strong, data-driven baseline using genomic locality, as integrating biological priors like gene networks is a major research challenge in itself. It requires solving difficult, open problems in mapping non-coding SNPs to genes and developing novel multi-modal architectures. As such, we believe it is a logical and important next step, but one that was beyond the scope of this initial benchmark-setting paper.

Action: We will add a discussion of Phenformer to our Related Work section, contextualizing our work with respect to its scale and findings.

5. Suitability for NeurIPS

We appreciate the acknowledgement that our work is already suitable for the genomics community. However, our choice to submit to NeurIPS was deliberate. We believe this paper serves as a vital bridge between the ML and genetics communities by introducing a new grand-challenge problem for high-dimensional, structured data and delivering two foundational ML contributions:

• A scalable framework (PRSformer) : The first viable architecture to make end-to-end learning on million-scale individual genotypes feasible.
• The Discovery of a Fundamental Scaling Law for Genetics: The non-obvious relationship between data scale, model complexity, and predictive power for this high-impact problem is a core ML insight.

Our work establishes the rigorous benchmark upon which future ML research in this critical domain can be built and extended. We are confident this paper offers the kind of impactful, benchmark-setting research that is highly valued by the NeurIPS community. Thank you again for your detailed and constructive engagement; we hope this increases your support and confidence in this exciting direction of work.

Response to Reviewer V3P1

We sincerely thank Reviewer V3P1 for their thoughtful review and for recognizing our work’s "excellent" significance. We are encouraged by this positive assessment. Your questions allow us to clarify and reinforce the core contributions that we believe make this work a strong fit for NeurIPS.

1. On performance on smaller datasets and overfitting (Figure 3)

Performance seems slightly worse than the linear model on smaller datasets (<1M) for some diseases... why do you think PRSformer is sometimes doing slightly worse...? Is PRSformer overfitting?

This is an excellent observation that points directly to a central and significant finding of our paper. PRSformer's higher model capacity makes it more prone to overfitting than the linear baseline when statistical power is limited.

A key result of our work is this critical scaling law: the predictive advantage of complex, non-linear models in genomics only emerges at the million-sample scale. This finding is a crucial, data-driven guideline for the ML and genomics communities.

I'm curious how Figure 3 looks like for the training data.

While we cannot include new figures in this rebuttal, we can confirm that PRSformer outperforms the linear model at all sample sizes. The performance gap between the train and validation sets is largest at smaller scales and narrows as N increases, which is classic evidence of overfitting and directly supports our discovery of this scaling law.

Action: We will revise Section 4.2 to more forcefully frame this as a key finding, explicitly discussing the trade-off between model complexity and data scale.

2. On the contribution of cross-ancestry generalization

Regarding the benefits of training on a diverse cohort... is it not expected that performance will go up? ... It would make more sense if the point were that PRSformer benefits more than traditional methods.

This is an insightful question that allows us to clarify a core methodological innovation in the paper. While more data is generally better for any model, our contribution is providing the first scalable framework that makes joint, end-to-end training across diverse ancestries possible at this scale.

Current SOTA methods (e.g., LDPred2) operate on summary statistics, which are typically computed within a single ancestry group. These methods fundamentally cannot perform joint training on individual-level data to learn shared genetic signals across populations. The standard paradigm is to train separate, stratified models, which exacerbates performance disparities and fails to leverage shared biology.

PRSformer’s end-to-end design is a fundamental departure from this limiting, stratified approach. Our results (Table 1) are the first to demonstrate that a single, unified model trained on diverse individual-level data can significantly improve predictive equity for non-European populations, often without sacrificing European performance. This was a major unsolved challenge in the field.

Action: We will sharpen the language in Section 4.4 to explicitly contrast our unified, end-to-end approach with the limitations of the standard stratified paradigm, making the novelty and impact of our work clear.

3. On interpreting learned non-linear interactions

No insights offered into the improved performance... empirical evidence that the model has learnt any interactions at all would be helpful... What non-linear interactions is PRSformer learning?

We agree that model interpretability is a critical long-term goal. However, scientifically we deliberately frame this research direction as a two-stage process:

• Stage 1 (This Paper): First, we prove the phenomenon. The primary contribution here was to rigorously establish that a predictive advantage from non-linearity exists at scale, which is a significant and previously unproven hypothesis. Discovering the N > 1M scaling law provides the essential foundation for all future work.
• Stage 2 (Future Work): Then, we characterize the phenomenon. A full analysis of what the model learned in various immune diseases is a major undertaking that requires its own dedicated study. This involves significant computational analyses, hypothesis generation and experimental validation.

To attempt this analysis within this paper would detract from the clarity of our core finding of a scaling law. Our work provides the community with the first viable tool to begin investigating genetic interactions of complex immune diseases.

Action: We will expand the "Future Work" section to articulate this two-stage approach and outline our concrete plan to use analysis of attention patterns to interpret the model's predictions in subsequent work.

4. On providing code for review

No code provided for review.

We apologize for this omission. It was done solely to preserve the integrity of the double-blind review process. We hope the simplicity and efficiency of our model provided necessary insight during review and will include the code with publication.

Action: We include code with publication.

We hope these clarifications address your concerns and strengthen your confidence in the originality and significance of our paper. Our work is the first to demonstrate that deep learning on individual-level genomes can surpass highly optimized SOTA methods and, crucially, uncovers the fundamental scaling law governing this transition. Thank you for your time and engagement!