PRSformer: Disease Prediction from Million-Scale Individual Genotypes
Authors: Payam Dibaeinia
Chris German
Suyash Shringarpure
Adam Auton
Aly A Khan
Venue: NeurIPS 2025 poster ()
Keywords: Multitask Learning
Transformer Architectures
Neighborhood Attention
Deep learning for Healthcare
Primary area: machine_learning_for_sciences
TL;DR: PRSformer: A scalable Transformer using neighborhood attention for multitask disease prediction from million-scale individual genotypes, showing non-linear modeling benefits at large N.

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.

PRSformer: Disease Prediction from Million-Scale Individual Genotypes
Payam Dibaeinia' aeeeeiiia@gmmli. com
Chris German' chris chrisg@23andme. com
Suyash Shringarpure
suyashss@gmail com
Adam Auton! aauton@23andme, com
Aly A. K2an
aakhan@uchi cago edu
23andMe, Palo Alto, CA, USA
University of Chicago, Chicago, IL, USA
Abstract
Predicting disease risk from DNA presents unprecedented emerging challenge as biobanks approach population scale sizes (N > 10th individuals) with ultra-high- dimensional features (L > 102 genotypes). Current methods, often linear and re- liant 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 atten- tion, achieving linear (LL) complexity per layer, making Transformers tractable for genome-scale inputs. Crucially, PRSformer utilizes 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 learn complex, potentially non-linear and long-range inter- actions directly from raw genotypes. We demonstrate PRSformer's effectiveness using unique large private cohort (N 5M) for predicting 18 autoimmune and inflammatory conditions using L x= 140k variants. PRSformer significantly outper- forms highly optimized linear models trained on the same individual- -eveel data and frate-of-the-ee art summary-stin -tattaartstttttt---ateeeeee saadittic--ased methods (LDPred2) derived from the same cohort, quantifying the beneñits of non-linear modeling and multitask learning scale. Furthermore, experiments reveal that the advantage of non- non-linearity emerges primarily at large sample sizes (N 1M), and that aulti-ancestry trained model improves generalization, establishing PRSformer new framework for deep learning population-scale genomics.
1 Introduction
Learning predictive models from high-dimennioonll complex structured data a fundamental machine learning challenge. This challenge is acutely relevant in modern genomics, where biobanks are rapidly scaling towards million-sample sizes (N > 1, 000, 000) and individual genomes are characterized by hundreds of thousands to millions of genetic variants (c.g.. Single Nucleotide Polymorphisms, SNPs), yielding a regime of ultra-high dimensionality (L > 100, 000). Effectively leveraging
"Correspondigg author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

individual-leved genomic data at this unprecedented N x L scale is critical for unlocking deeper insights into complex trait genetics. such as predicting disease susceptibility 22.
Current state-of-the-art methods for disease risk prediction primarily rely on Polygenic Risk Scores (PRS) derived from Genome-Wide Association Study (GWAS) summary statistics 3
5 6 While
effective to a degree, these methods are predominantly linear, capturing additive genetic effects. Furthermore, by operating on summary statisties, they discard potentially valuable individual-Ievel information and struggle to model non-additive genetic interactions (epistasis) 7 8 9 10]. These
limitations may cap predictive performance, especially as dataset sizes scale towards the million- sample regime where subtle interaction effects might become detectable.
Transformer architectures have revolutionized sequence modeling in other domains by capturing com- plex, long-range dependencies via self-attention [11]. We hypothesize that the attention mechanism provides a powerful inductive bias for genomics, enabling more effective modeling of pairwise linear and non-linear interactions between genetic loci compared to traditional architectures or inherently linear models. However, critical barrier persists: the prohibitive OLL²) computational complexity of standard self-attention renders its direct application to genome-scale sequences (L > 100, 000) computationally infeasible. Addressing this scalability bottleneck is critical to harnessing the power of Transformers for large-scale genomics.
Here we introduce PRSformer, novel Transformer-basee architecture specifically engineered for scalable, end-to-end, multitask disease risk prediction directly from ultra- hutra-high-dimensionn (T>> 100k) individual-level genotypes (N 1M). PRSformer's core innovation lies in its scalability, achieved by incorporating neighborhood attention (NA) 12,, an efficient attention mechanism restricting computations local genomic windows, resulting in (L) complexity per layer. This design aligns with the biological structure of the genome, which we treat as a series of linked regions, called linkage disequilibrium (LD) blocks, where genetic variants are often inherited together. PRSformer stacks NA layers, which first model interactions within LD blocks, then progressively integrates information between neighboring LD blocks in deeper layers, capturing larger-scale genetic patterns influencing disease.
To evaluate PRSformer's 's ability to harness the shared genetics underlying multiple, related traits, trained and validated it in a multitask setting across 18 autoimmune and inflammatory conditions. This trait set provides ideal testbed for multitask learning due to immune-mediated inflammatory diseases frequently exhibiting shared inflammatory pathophysiology and overlapping genetic factors 13. The multitask formulation allows us test PRSformer's ability exploit shared genetic associations while learning disease-specific patterns, aiming to enhance predictive performance through a shared representation.
We provide rigorous empirical validation using data from a unique large private cohort (N = 3.8M European ancesttr individuals) for training, validation, and evaluation across D = 18 autoimmune and inflammatory conditions using L Rs 140k variants. The scale of this cohort significantly exceeds current public biobanks [14, providing critical real-world testbed for methods designed to handle genomic data of the million-scale magnitude. We conduct stringent comparisons against: (i) highly optimized linear models (regularized logistic regression with learnable embeddings) trained on the exact same individual-level genotype data, isolating the benefit of PRSformer's non-linea architecture; and (ii) state-of-the-art bunmary-statisiic-aasss from the exact same cohort, enabling a direct comparison between end-to-end and summary-statistic- based approaches. Our experiments demonstrate statistically significant performance gains for PRSformer.
PRS methods (LDPred2 (5)) derived
Our main contributions are:
Scalable deep learning for genomics at population scale: Transformer architecture applied to population-scale data (N es 5M) with ultra-high- dimensional features (L & 140K), establishing a blueprint for tackling other genome sequence prediction tasks.
We present an efficient multitask
Critical scaling law for non-linear models: We empirically establish and quantify key scaling law demonstrating that the predictive advantage of non-linear models over linear methods emerges primarily at the million-sample scale for complex immune-related condi- tions. Our analysis quantifies this effect, showing that performance gains grow consistently the training set size increases beyond one million individuals.

---

Multitask learning improves genomic prediction: We show that multitask training across related traits consistently outperforms the standard single-task paradigm, demonstrating the benefit of learning shared genetic representation across complex immune-mediated inflammatory diseases.
Improved cross-ancestry generalization: We show that training PRSformer on multi- ancestry data, including an additional 1.1M non-European individuals, markedly improves prediction accuracy for held-out non-European individuals compared to model trained only on European-ancestry data, offering a path toward more equitable genomic prediction.
2
Related work
This work is situated at the intersection of statistical genetics, genomics, and deep learning. We specifically advance upon prior work in three key areas: polygenic risk prediction, the application of deep learning to genomic data, and the development of efficient Transformer architectures for ultra-long sequences.
2.1 Polygenic risk score methods
Traditional complex trait prediction relies heavily on PRS derived from GWAS summary statistics 6
Early methods often involved simple thresholding and summing of SNP effects [3 15 More recent Bayesian approaches, such as LDpred2 (S] and PRS-CS [I6), explicitly model LD patterns and utilize shrinkage priors to improve predictive accuracy. These methods represent the current state-of-the-art for prediction from summary statistics. However, PRS methods based on summary statistics are fundamentally limited in several ways.
First, by discarding individual-level genotype and haplotype information, these methods cannot capture LD structure and must instead rely on LD estimates that are typically imputed from external reference panels, which can introduce biases due population mismatches [17 [8]. Second, they cannot capture variant-variant interactions such as epistasis as they are restricted to using marginal variant effects. Third, the use of precomputed summary statistics constrains these models largely linear architectures, precluding the discovery of complex multi-locut or hierarchical genetic patterns. Recent summary-statistie approaches to leverage non-additive signal remain constrained by the lack of individual-level haplotype context [19]. Taken together, these limitations may cap predictive performance, particularly as biobank-scale datasets grow large enough to enable the detection of more subtle and nonlinear genetic effects.
Alternative individual-level approaches, such as BOLT-LMM [20) and GEMMA [21], estimate SNP effect sizes under a linear mixed model framework to account for population structure and polygenic background effects. However, these methods are computationally demanding at our study's 's scale (N=3.8M, D=18 traits): GEMMA's cubic complexity in N renders it intractable, while BOLT- LMM, though more scalable, operates on single trait at a time, requiring 18 separate runs. Prior work has shown that LDPred2 achieves predictive performance comparable to BOLT-LMM across multiple traits
22 23
supporting its use as a strong linear baseline for comparison.
2.2 Deep learning in genomics and trait prediction
Deep learning has been successfully applied to various supervised genomic prediction tasks. Much work has focused on modeling sequence-level information (DNA base pairs) to predict molecular phenotypes like transcription factor binding 224 25 26), chromatin accessibility (27). or gene expression (28 29). These approaches have predominantly utilized Convolutional Neural Networks (CNNs) or Transformers incorporating CNN-style tokenization, which are well-suited to capturing biologically meaningful motifs and local patterns at base-pair resolution. However, this paradigm is less intuitive when modeling the influence of genetic variants (e.g.. SNPs) on complex traits, as causal variants can be spread across the genome and may interact over long distances, often without strong local sequence motifs defining their impact.
The application of deep learning to predict complex traits (like disease status) directly from individual- level genotype data (i.e., SNP arrays) remains relatively underexplored, particularly at the population scale addressed in this paper [30 B1) This is largely due to the challenges of higra-high& dimensionality
3

---

(L) and, until recently, the limited statistical power of publicly available cohorts with both individual- level genotypes and phenotypes (N). Prior work in this specific domain has often relied on: (i) tree ensemble models such as gradient boosting or simple neural networks trained on reduced feature sets (e.g., using LASSO feature selection); (ii) smaller cohorts where complex interactions are difficult to detect; and (iii) models operating on precomputed PRS or summary statistics rather than raw 35 36). Recently, Phenformer (37] proposed a multi-scale Transformer that predicts disease risk from DNA sequences by linking genetic variation, gene expression, and phenotype through a pretrained sequence-to-expresiio backbone. While conceptually similar to our end-to-end genotype-to Sanotype-o--pheootpee goal, Phenformer operates DNA sequences covering approximately z 3% of the genome and is trained on ~ 150K individuals, whereas our approach models variant-level genotype data and scales to millions of individuals, enabling systematic analysis of how nonlinearity interacts with data scale in complex trait prediction.
genotype data 8 32 33 34
2.3 Efficient transformer architectures
Applying standard Transformers to genome-scale data (L > 100k) is computationally prohibitive due to the ( complexity of self-attention 38). A wide range of efficient attention mechanisms have been proposed to address this limitation. including sparse attention patterns (e.g.. Longformer 39. BigBird 40), low-rank approximations (e.g.. Linformer 44]),, and kernel-based methods
(e.g., Performer 42) Other architectures exploit locality through sliding windows or blockwise [4] ) to reduce complexity while capturing local dependencies. In our work, we adopt neighborhood attention [12] variant of self-attention in which each query attends only to a fixed-size local window of neighboring tokens, rather than the full sequence. This inductive bias aligns well with the block-like correlation structure of genomic data driven LD. By limiting attention to a neighborhood of size k < L, NA reduces both computational and memory complexity to OLL k) -achieving linear scaling in sequence length. We employ the optimized GPU implementation provided by the NATTEN library [12 44, which supports scalable training on long sequences while maintaining the expressiveness of content- based attention.
mechanisms (e.g., Swin Transformer
3 Methods
3.1
Problem definition
We aim to predict susceptibility to multiple (D = 18) autoimmune and inflammatory conditions from individual-level genotypes. Formally, given a dataset of N individuals, the input for individual i is their genotype profile x; {0. 1,2, UNKN] /, representing genotypes of L pre- pre-selected genetic variants, where UNKN indicates missing data. The target output is vector Y; E fo, 1, UNKN) D. representing the binary status (case/control) for D diseases, where UNKN indicates unrecorded status. Our goal is to learn multitask function f: Z-- + (0, 1]P that predicts the probability of each disease y; fxx)). The primary challenges lie in the ultra-high dimensionality (7. ~ 1 140k in this work), the need to capture potentially non-linear and long-range interactions, and leveraging the statistical power of population-scale datasets (N x 5M total used in this study). We propose Transformer- based architecture adapted for this task, leveraging efficient attention mechanisms for scalability and multi-task learning for joint prediction across diseases.
3.2 PRSformer architecture
PRSformer adapts the Transformer architecture for disease prediction from ultra-long (L Rs 140k) individual-level genotype sequences using the following key designs:
Scalability via Neighborhood Attention: Standard O(L?) self-attention is computationally infeasi- ble. We replace it with Neighborhood Attention (NA) [12. restricting each query token's 's attention to symmetric local window of size k. This reduces complexity to O(L k), enabling efficient processing of the L = 137, 245 input variants used in this study. We use A = 385, chosen via hyper- parameter tuning (Section 3 Supplementary Table F8), which conceptually aligns with capturing dependencies within local LD blocks (Figure TTA) and corresponds to roughly =100 kilobases along the genome (45)
Genome-ordered input without explicit positional encodings: Input variants are ordered by their chromosomal position and then concatenated from Chrl to Chr22. This fixed order is used for all

---

Add
Food Farvard
cen
MIP 12d
Oceene Flanda
MLP d
MLP id 2d0
SBE
Vekart
Layer
Layer Norm
Murti-head Ye Neighboroodd Allenton
Variat Emmbedder
Variatt
Embeddoer
Learnalee Variant Embedding
Layer Noret
desage UMcc
desage 012
Inzut Vanani
Dosage
Input Emboodingg,
(A)
(B)
(C)
(D)
Figure Schematic overview of PRSformer and its core components. (A) A preselected set of variants across the 22 chromosomes, sorted by genomic position, allows each query variant (c.g., blue or red) to attend within local block (k: 385) via neighborhood attention. (B) Model architecture with imput, Transformer blocks, and output layer. (C) Variant embedding layer, which encodes each variant and its corresponding genotype (0, 1, 2, or UNKN) into d4-dimensional representation. (D) Transformer block with pre-layer normalization, neighborhood attention, and GELU activation.
individuals. We omit standard positional encodings (c.g., sinusoidal or learned absolute). The fixed genomic order provides implicit relative positional information that NA inherently leverages within its local attention windows. We also experimented with learned positional encodings, but did not observe measurable improvement in performance.
The overall data flow of PRSformer proceeds as follows (Figure
Learnable variant embedding layer (Figure [c): Each variant in the input sequence of variant genotypes is mapped to a dmoderdimmnnionnll vector. For each variant j and individual an observed genotype Tay E {0, 1, represented as E, - Tag using learned variant-specific embedding Ej. A missing genotype (xig UNKN) is represented by separate learned variant-specific embedding My. This allows the model to distinguish missingness from observed genotypes distinctly for each variant. We use dmodel = 64.
2. Transformer blocks (Figure ID): The embedded sequence is processed by NNuyess
Transformer blocks. Fach block applies pre-layer normalization [46] followed by multi-head NA (Nnents and feed-forward network with GELU activation [47]. The choice Nuyers = was based on validation performance, where deeper models did not show significant improvement for this task (Supplementary Table F7), suggesting two stacked NA layers provide sufficient receptive field to capture interactions between adjacent LD blocks.
3. Output layer: The normalized sequence representation from the last block is flattened and optionally concatenated with covariates such and age, and passed a fully- fully-connected layer generating D 18 independent disease likelihood predictions. We also evaluated mean pooling and dedicated [CLS] tokens alternatives simple flattening of normalized representations, but neither outperformed the flattening- based design (see Supplementary Table F10).
Key architectural hyperparameters (dmodels Nhends, Ninyrrs NA window size k) were optimized based on validation set performance (Supplementary Tables F5-10).
5

---

3.3 Datasets
We utilized data from a large, private biobank consisting of individuals who consented to participate in research under IRB-approved protocol. Starting from internal GWAS data freeze timestamped 08-2021 (used to prevent information leakage, see Section 3.4). we identified individuals genotyped on the same platform and excluded all pairs of individuals related by more than 700 cM (i.e., first cousins or closer), thereby minimizing the risk of learning simple familial signals. Individuals included had recorded phenotypes (i.e., self-reported status) for at least one of the D 18 autoimmune and inflammatory conditions considered. We also excluded individuals under age 10 who did not have a case diagnosis. This resulted in a training set of Nuain = 3, 838, 5449 individuals of European genetic ancestry (throughout this work, ancestry was determined via an internal classifier 48 ).
We constructed temporally distinct validation and test sets using individuals who enrolled and consented after the 08-2021 data freeze date and up to 12-2024, applying the identical filtering criteria. The validation dataset was used for hyperparameter tuning, while the test dataset was used to report final performance metries. This yielded Naal = 525, 448 and Nostt = 494, 265 individuals. To assess whether models capture familiil relationships versus causal genotype-phenotyy constructed a conship-contollled European test set (N = 148k) by subsetting test individuals related to any training sample by no more than 300 cM and ensuring that no pair within the subset is related by more than 700 cM. In total, the European dataset comprised N = 4,858,302 individuals across training, validation, and test sets, with an additional Rs 1.1M non-European individuals included for cross-ancestry training (see Section 3.7 Case/control counts per disease and further details regarding cohort construction (including differences from the subset used for GWAS computations) are provided in Appendix A and Supplementary Tables F2-4
associations, we
3.4 Variant selection
To define the input feature space (L), we selected variants associated with at least one of the D = 18 diseases based on internal GWAS summary statisties (European ancestry cohort, computed prior to an 08-2021 1 data freeze to prevent information leakage into model training). For each disease, variants passing standard QC, located on autosomal chromosomes, had a genotyping rate > 0.95, MAF >> 0.001. and exhibited nominal association with the disease (GWAS p-value < 1 x 10 10- The final PRSformer input set was the union across all 18 diseases, resulting in = 137, 245 variants. Further details on GWAS procedures, exploration of variant's 's pruning by LDD and per-disease variant counts are in Appendix B and Supplementary Table F9)
3.5 Training
Given the training data D = ((x., Mi we trained PRSformer by minimizing the following loss, summed over individuals and their available (non-UNKN) disease labels / E T(1):
Numa,
Niple
(un,e log(ūr.) + (1 - ys.e) log(1 (c,+)l
EE)
where T(i) denotes the set of recorded disease statuses for individual i. We also evaluated focal loss loss [50], and standard averaged cross-entropy. all of which were outperformed by the proposed loss function in terms of validation AUROC. We used the AdamW optimizer [5T] (3₁ = 0.9, 82 = 0.999, weight decay=0.05) with an initial learning rate of 5 x 10-1, decreased via Cosine Annealing scheduler. Training was performed efficiently on this large-scale dataset for 2 epochs consisting of x 120,000 gradient updates in an effective batch size of 64 across four NVIDIA A100 GPUs, leveraging Distributed Data Parallel and mixed-precision (FP16) training (training duration was tuned based on validation AUROC across diseases (Table F5): most models showed signs of overfitting beyond two epochs). Hyperparameters, including architecture choices (Nunees, dmodet: Nueads: k). were selected based on optimal AUROC on the validation set after extensive searches (e.g., See Supplementary Tables F5-10), following standard ML best practices to minimize overfitting to validation data and ensure that test performance provides an unbiased estimate of generalization. To isolate the contribution of genotype to model performance, all models presented in the main text were trained without including covariates such as or age.
[4]. fask-uncertainty-weegghee

---

3.6 Baseline models
To comprehensively evaluate PRSformer and validate our main claims regarding the utility of end- to-end non-linear modeling on large-scale individual-level genotypes, established three rigorous baselines. These baselines are specifically designed to: (1) compare against the current state-of-the-an using conventional summary-statiste inputs (LDPred2), (2) benchmark against an enhanced version of this state-of-thh-arr (Stacked LDPred2), and (3) isolate the specific performance gains attributable to PRSformer's non-linear Transformer architecture via carefully matched linear counterpart.
1. LDPred2: state-of-the-arf summary-statistic method. We selected LDPred2 [5] due to its strong empirical performance and widespread adoption in the field 36 52 To ensure the most direct comparison possible, we configured L.Pred22 meticulously:
Matched data source: LDPred2 was applied to GWAS summary statisties derived from the same European training data freeze used for PRSformer's data and variant selection.
Cohort-specific LD: An LD reference panel from our research cohort was used.
Standard QC: Input variants (~445K per disease) passed standard GWAS QC and LDPred2-specific filtering [18].
Tuning: LDPred2 hyperparameters (p, 1²) were extensively tuned (up to 100 models per disease) by maximizing AUROC on the same validation set used for PRSformer (details in Appendix C).
2. Stacked LDPred2: enhanced summary-statisti baseline. To create a stronger summary- statistic baseline, we ensembled the converged LDPred2 models from the hyperparameter search using elastic-net regression trained via cross-validation on the validation set (Ap- pendix C).
Since Stacked LDPred2 uses the validation set for training ensemble weights, develop and compare it against PRSformert, which is the final PRSformer model retrained on the combined training and validation datasets, ensuring parity in total data usage (Supplementary Figure E2).
PRSformert:
3. Linear model: direct architectural ablation. This crucial baseline isolates the contribution of PRSformer's non-linear Transformer architecture. It mirrors PRSformer precisely except for omitting the Transformer blocks:
Identical data, inputs & training: Uses the exact same L ~ 140k input variants and individual-level train/validation/eest splits. Employs the same multitask framework
(D = 18), loss function, AdamW optimizer, and training schedule (Section 3.5). Uses
the same embedding layer (Figure[[]) for genotypes and missingness.
Architecture difference: The input embeddings are fed directly to the final linear output layer, bypassing the Transformer blocks (Figure IP)).
This provides a multitask linear model on the same large-scale individual data, allowing direct assessment of the performance gain from PRSformer's 's non-lineat processing.
3.7 Cross-ancestry experiments
To assess generalization, we developed PRSformer-ME (Multi-Etthnic). We performed ancestry- specific GWAS (African American (AFR), European (EUR), Latino (L.AT), East Asian (EAS), and South Asian (SAS): determined by internal classifier) using the same 08-2021 data freeze and variant selection criteria (Section 3.4, p-value < 1 x 10-2, QC) where sample sizes permitted (Supplementary Table F12). We defined an expanded input set (L 251, 538 variants) as the union of selected variants across all available disease-
pairs (including Europeans). We constructed multi- ancestry training set (~ 5M total individuals) by combining the European training set (Section (3.3) with N = 1, 136, 746M non-European individuals meeting the same filtering criteria. PRSformer- ME was trained on this combined dataset using the same architecture and hyperparameters as the European-only PRSformer, without additional ancestry tuning. Evaluation was performed on a combined test set including the European test set and held-out non-European individuals processed identically (Supplementary Table F3).
7

---

4 Experiments and results
We present results evaluating PRSformer's performance against baselines, analyzing the impact of non-linearity and sample scale, assessing the benefit of multitask learning, and testing cross-ancestry generalization using AUROC as the primary metric unless otherwise stated.
4.1 PRSformer outperforms state-of-the-art baselines
We first benchmarked PRSformer against the highly optimized linear model and the state-of-the- summary-statistie method, I.Pred2,, on the European test set (Niest R 494k). As shown in Figure 2 PRSformer consistently achieves higher AUROC scores than its linear counterpart across all 18 autoimmune and inflammatory conditions. This comparison, using identical data and training setups except for the Transformer blocks, directly quantifies the predictive benefit derived from PRSformer's non-linear architecture.
P-C 10-4
Sord
8.027
p-r.
BV704
1.000.044
0.79
0.567
0P
0.632
0.591
3.900042
0-6655
0.44
LOPeE2 Wilere m9
12 Mo
Fec Receere 30
Zorder Sem
Abordd Arerss
0.987
0.905
0121
Figure Benchmarking PRSformer against baseline methods using AUROC. PRSformer consistently outperforms the linear model (trained on identical individual-level data) and LDPred2 (state-of-the- art summary statistic method derived from the same cohort). Error bars: 95% CI (10k bootstraps). p-values: estimated using one-sided paired bootstrap test (10,000 replicates), sampling with replacement from the test set and comparing AUROCs of PRSformer and LDPred2 on identical sample pairs. The v-auuee reflects the fraction of replicates where LDPred2's s AUROC PRSformer's
Crucially, PRSformer also significantly outperforms LDPred2 (using summary statistics derived from the same cohort) on 16 out of 18 diseases, with 1 differences being statistically significant (p < 0.05, one-sided paired bootstrap test). This demonstrates the advantage of end-to-end modeling on individual-level data compared to state-of-the-art methods relying on summary statistics. Consistent improvements were also observed in area under the precision-recall curve (Supplementary Figure E1) and explained variance (Supplementary Table FI1), as well as when comparing against an enhanced Stacked LDPred2 baseline (PRSformer+, Supplementary Figure E2), confirming the robustness of PRSformer's ' advantage.
We also evaluated PRSformer and LDPred2 on the kinship-controlledd test set, reproducing similar trends (Supplementary Figure E3): PRSformer outperformed LDPred2 in 14 of 18 diseases, main- taining its lead in 13 of the 16 and newly improving Alopecia Areata, with six remaining statistically significant (p < 0.05). The smaller number of significant improvements is expected given the reduced power of the kinship-controlledd test set (~14 vs. ~494k). These results confirm that PRSformer's s advantage is not driven by familial confounding and persists under stringent kinship control, reinforcing the validity of our findings.
4.2 Benefit of non-lincarity emerges at million callion--samle scale
To understand when the non-linear modeling capabilities of PRSformer become advantageous, we compared its performance against the linear baseline across varying training dataset sizes (down- sampling the N 3.8M training set). Figure reveals a critical insight: at smaller sample sizes, comparable to current large public cohorts like UK Biobank [53] (up to N ss 1M), the performance
8

---

Callac
T1D
Graves
Sjogren's
0
0.85
0.75
AUROC
0.75
UKBS
scale
UK38
stale
0.62
0.80
UKB8
0.70
0.13
0.63
0.56-
0.65
0.70
0.54
0.55
1.30
Training data size
Training data $22 1e6
Training data Pestormer
Training data s2z0
MT: non linear
MT: Inear
LDPred2
PRStormer
Figure 3: Impact of training scale on non-lineat model advantage. Performance (AUROC) of PRSformer (non-linear) and the linear baseline across downsampled training sets (multitask setting). The benefit of non- oon-iinerrty becomes apparent only at N > 1M scale.
of PRSformer is similar to the simpler linear model. However, as the training size exceeds one million individuals, a clear advantage for the non-linear PRSformer emerges and progressively widens. This trend holds across multiple diseases (Supplementary Figure E4) and persists even when using appro- priately subsetted variant sets for smaller scales (Supplementary Figure E5). These results indicate a scaling law: the benefits of non-linear architectures like PRSformer manifest only when sample sizes sufficient resolve higher-order genetic interactions. Below this threshold, linear models remain competitive, whereas beyond the million-sample regime, PRSformer achieves measurable gains (although with higher computational cost in FLOPSs per sample).
4.3 Multitask learning consistently improves performance
We investigated the benefit of PRSformer's multitask design by comparing it against single-task (ST) models trained independently for each disease. Figure[ssho shows that multitask (MT) training consistently yields superior AUROC compared to ST training for both PRSformer and the linear baseline, across different data scales (see Supplementary Figure E6 for other diseases). This improve- ment was robust even when ST models used disease-specific optimized variant sets (Supplementary Figure E7). Thus, the gain stems from leveraging shared information across related immunc- meddatee inflammatory diseases allowing shared model components (variant embeddings and Transformer blocks in PRSformer, and variant embedding in the linear baseline) to be optimized more effectively. By training these shared layers across multiple related diseases, the model can capture generalizable representations that enhance performance beyond what is achievable with isolated, ST training.
Collac
T1D
Figure 4: Multitask (MT) vs. Single-Task (ST) training for Celiac disease (left) and TID (right). MT consistently out- performs ST for both non-linear (based PRSformer) and the linear baseline across different scales (X-axis: ST sizes / (MT sizes)).
0.85
0.75
AUROC
Ntt ren-iree
5tt restree
It ree-lesrr
0.70
0.60
restimer
1.53--05
5.06-05 05
275e-05
L.teeoee 353e-30 training date
3.30x00
training deta
4.4 Improved cross-ancestry generalization via multitask multi-ancestry training
com- bined multi-ancestry cohort (~ 5M individuals, including ~ 1. non-Europeans) using an expanded variant set (L R 252k, see Section 3.7 Methods). We evaluated its performance a held-out test set containing individuals from European (EUR), African American (AFR), Latino (LAT), East Asian (EAS), and South Asian (SAS) ancestries, comparing it to the original PRSformer trained only on EUR individuals.
Recognizing the need for equitable genomic prediction [54), we trained PRSformer-
As summarized in Table I PRSformer-ME. demonstrates significantly improved generalization to non-EUR populations. achieves substantially higher AUROC scores across most diseases in AFR, LAT, EAS, and SAS individuals compared to the EUR-only model. Importantly, this gain in non-EUR

---

populations is achieved with minimal to no degradation in performance on EUR individuals. These results indicate that training on diverse, aggregated individual-Ievel data allows PRSformer-ME to capture both shared and ancestry-speciff genetic signals, leading to more accurate and potentially more equitable predictions across populations compared to models trained on a single ancestry group (incorporating covariates such as sex and age further improves predictive performance, see Supplementary Table F14). This is an important finding since state-of-the-art methods generally rely on single-ancestry summary statistics, preventing them from jointly training on individual-level data across multiple ancestries and from leveraging shared pross-populltion genetic signals.
Table 1: AUROC of EUR-only PRSformer vs. test set. PRSformer-Me shows improved performance in non-EUR ancestries often without sacrificing EUR performance. Bold font denotes the higher AUROC in each pairwise comparison.
PRSformer-ME on the anclity
Disesses
Elahhmoto
Rbeumazodd
Muitiple
Codon's
Polymyalgia
Rivormtiicn
Puliatic
Villlgo Spandyy
Migmtory Cilookii
Alopecia Arcala
Son
Anoes.
Model
PRSformer
0.7933 0.8457 0.7661 (16357 (1666H
0.7504
ULSEEN
U.60RAL
0,7184
0.6885 0.4048
016422
(16660 0.5802 0.6212
0.7164
0:6.27 0.A0S]
EUR
PRSTormer-ME
0,7867 0.8431 0,7674 0.6404
0.6683
0,7487
0.6458
0.6069
0,7157
0.6981
0.3995
0.6447
0.6654
0,5818
0.6222 0,7146
0.6246
0.6125
AFR
PRSTormer-ME 0,6269 0.6986 0.6997 0.6399 0,5822
PRSformer
0.6219 0.6391 0.5676
(1555U
0,6939
0.6140
0,5660
0.6397
0,5447
0.6217
0.5902
0.6285
(15928
0,5410
0,5605
0,5550
0.5913
0.6026
0.5844
PRSformer
0,7437 0,7459
(1611 (16449
(16216
0.5579
PRStormer-ME 0,7521 0,7732 0.7568
0.6685 0.6596 0,7581
0.6536
0.5979
0,7256
0,6842
0.6533
0.6291 0.6539 0,6938 0,6166 0,7270 0.6397 0.6366
SAS
PRSformer-ME
PRSformer
0,8002 0.6852 0,7612
0,8190 0.6640 0.7667
(163NS 0,6209
0.6885 0.6269 0,6907 0.6196
0,5974
0,7503
0,7543
0.0630
0,7061
0.6974
6969 0.6858 O. 0.528 0.6009 0,5146 O 0,5813 0,5972
(16423 O. 0,5999 5693 0.5421 O. 0.5999 0.5754
PRSformer
0.6.UT 0.6957 0,577U
(L6167 U.6SE5
U4614
0.540
0.7451
0.5902 0.5184 115463 O 0,6809
0.5846
PRStormer-ME
0,7121 0.7487 0.6176 0.6355 0,6876 0,6573 0,6085
0,50
0.6256 0.6973
0.6545 0,6518 0,5860 0,5602 0,7085 0,5929 0.6366
Conclusion
We introduced PRSformer, a scalable Transformer architecture leveraging neighborhood attention to enable end-to-end, multitask disease prediction from population-scale individual genotypes (N 5M. L = 140k). Our rigorous evaluation on unique large private cohort, conducted under IRB and using consented research participant data, demonstrates that PRSformer significantly outperforms strong linear and state-of-the of--te-artt-t art summary-statiste baselines (LDPred2) derived from the same cohort.
A key finding of this work is that the benefit of PRSformer's non-linear modeling for complex meddnte- immune- mediated inflammatory diseases emerges primarily at the million- sample scale (N 1M). This advantage varies across diseases, with traits like celiac disease and type 1 diabetes benefiting substantially from non-linear modeling to explain disease risk variance (55). This scaling law, alongside our findings that multitask training improves performance and multi-ancestry data enhances generalization, establishes a new framework for genomic prediction.
While PRSformer advances predictive accuracy, its gains come with higher computational demands that may limit immediate clinical scalability. Furthermore, future work is required to develop interpretation methods to understand the learned non-linear interactions, which is essential for biological hypothesis generation and experimental validation. A key future direction is to extend the framework beyond single disease domain to a phenome-scale setting spanning thousands of traits. This approach is motivated by widespread genetic pleiotropy, where a single variant can influence multiple, seemingly disparate conditions. A unified model could therefore capture the shared genetic underpinnings linking diverse biological systems, such as the contribution of immune pathways to neurodegeneration and cancer.
Our research prioritizes fairness across diverse populations and the responsible deployment of genomic models. Recognizing the sensitivity of genomic data, we have taken steps to balance transparency with participant privacy: we provide detailed methodological deseriptions and have released our implementation code at http:://gghbb..ommmmmmmmmcPccoot trained models are not publicly available.
however, the data and
Taken together, our results and scaling analyses position PRSformer as foundation for phenome- scale genetic risk modeling that can fully leverage genetic pleiotropy to improve prediction and generalization at population scale.
10

---

Acknowledgments and Disclosure of Funding
The authors thank the past and present employees and research participants of 23andMe for making this work possible. We are grateful to Akele Reed, Teague Sterling, David Hinds, Steve Pitts, Wei Wang, Bertram Koelsch, Michael Holmes, Stella Aslibekyan, Cordell Blakkan and Barry Hicks for their valuable contributions and insightful comments on the manuscript, and to Ali Hassani for helpful discussions on employing Neighborhood Attention. The authors also gratefully acknowledge the support of AWS for providing GPU computing resources and credits. A. A. Khan is supported in part by a Chan Zuckerberg Investigator Award.
References
[1] Arnór I Sigurdsson, Joannis Louloudis, Karina Banasik, David Westergaard, Ole Winther, Ole Lund. Sisse Rye Ostrowski, Christian Erikstrup, Ole Birger Vesterager Pedersen, Mette Nyegaard, DBDS Genomic Consortium, Søren Brunak, Bjarni J Vilhjálmsson, and Simon Rasmussen, Deep integrative models for large-scale human genomics. Nucleic Acids Research, 51(12):e67-e67, 05 2023.
[2] Erping Long, Peixing Wan, Qingyu Chen, Zhiyong Lu, and Jiyeon Choi. From function to translation: Decoding genetic susceptibility to human diseases via artificial intelligence. Cell Genomics, 36)):003320, 2023.
[3] Naomi R Wray, Michael E Goddard, and Peter M Visscher. Prediction of individual genetic risk disease from genome-wide association studies. Genome research, 17(10):1520-155: 2007.
[4] Ali Torkamani, Nathan E Wineinger, and Eric Topol. The personal and clinical utility of polygenic risk scores. Nature Reviews Genetics, 19(9):511-590, 2018.
[5] Florian Privé, Julyan Arbel, and Bjarni Vilhjálmsson. Ldpred2: better, faster, stronger. Bioinformatics, 36(22-23):5424 5431, 2020.
[6] Cathryn M Lewis and Evangelos Vassos. Polygenic risk scores: from research tools to clinical instruments. Genome medicine, 12(1):44, 2020.
[7] Jana Schwarzerova, Martin Hurta, Vojtech Barton, Matej Lexa, Dirk Walther, Valen- tine Provaznik, and Wolfram Weckwerth. A perspective on genetic and polygenic risk advane-sdvances and limitations and overview of associated tools. Briefings bioinformatics, 25(3):bbae240, 2024.
[8] Michael Elgart, Genevieve Lyons, Santiago Romero-Brufann Nuzulul Kurniansyah, Jennifer Brody, Xiuqing Guo, Henry J Lin, Laura Raffield, Yan Gao, Han Chen, et al. Non-linear machine learning models incorporating snps and prs improve polygenic prediction in diverse human populations. Communications biology, 5(1):856, 2022.
[9] Pankhuri Singhal, Yogasudha Veturi, Scott Dudek, Anastasia Lucas, Alex Frase, Kristel Van Steen, Steven Schrodi, David Fasel, Chunhua Weng, Rion Pendergrass, et al. Evidence of epistasis regions of long-range linkage disequilibrium across five complex diseases in the uk biobank and emerge datasets. The American Journal of Human Genetics. 110(4):575-591, 2023.
[10] Juannan Zhou, Mandy S Wong, Wei-Chia Chen, Adrian R Krainer, Justin B Kinney, and David M McCandlish. Higher-order epistasis and phenotypic prediction. Proceedings of the National Academy of Sciences, 1190399)c2204233119 2022.
[11] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[12] Ali Hassani, Steven Walton, Jiachen Li, Shen Li, and Humphrey Shi. Neighborhood attention transformer. In Proceedings of the IEFF/CVF conference on computer vision and pattern recognition, pages 6185-6194, 2023.

---

[13] Alexandra Zhernakova, Cleo C Van Diemen, and Cisca Wijmenga. Detecting shared pathogene- sis from the shared geneties of comune-related diseases. Nature Reviews Generies, 10(1):43-55, 2009.
[14] Clare Bycroft, Colin Freeman, Desislava Petkova, Gavin Band, Lloyd T Elliott, Kevin Sharp, Allan Motyer, Damjan Vukcevic, Olivier Delaneau, Jared O O'Connell, et al. The uk biobank resource with deep phenotyping and genomic data. Nature, 562(7726):203-2000 2018.
[15] Florian Privé, Bjarni Vilhjälmsson, Hugues Aschard, and Michael GB Blum. Making the most of clumping and thresholding for polygenic scores. The American journaa of human genetics, 105(6):1213-12211 2019.
[16] Tian Ge, Chia- Yen Chen, Yang Ni, Yen-Chen Anne Feng, and Jordan W Smoller. Polygenic prediction via bayesian regression and continuous shrinkage priors. Nature communications, 10(1):1776, 2019.
[17] Bjarni/ Vilhjálmsson, Jian Yang, Hilary K Finucane, Alexander Gusev, Sara Lindström, Stephan Ripke, Giulio Genovese, Po-Ru Loh, Gaurav Bhatia, Ron Do, al. Modeling linkage disequilibrium increases accuracy of polygenic risk scores. The american journal of human genetics, 9764)5576 2015.
[18] Florian Privé, Julyan Arbel, Hugues Aschard, and Bjarni J Vilhjälmsson. Identifying and correcting for misspecifications in gwas summary statistics and polygenic scores. Human Genetics and Genomics Advances, 3(4), 2022.
[19] Rikifumi Ohta, Yosuke Tanigawa, Yuta Suzuki, Manolis Kellis, and Shinichi Morishita. A polygenic score method boosted by non-additive models. Nature Communications, 15(1):4433, 2024.
[20] Po-Ru Loh, George Tucker, Brendan K Bulik-Sullivan Bjami J Vilhjálmsson, Hilary K Finu- cane, Rany M Salem, Daniel Chasman, Paul M Ridker, Benjamin M Neale, Bonnie Berger, al. Efficient bayesian mixed-model analysis increases association power in large cohorts. Nature genetics, 47(3):284-290, 2015.
[21] Xiang Zhou and Matthew Stephens. Genome-wide efficient mixed-moder analysis for associa- tion studies. Nature genetics, (((7):2118244 2012.
[22] Guiyan Ni, Jian Zeng, Joana A Revez, Ying Wang, Zhili Zheng, Tian Ge, Restuadi Restuadi, Jacqueline Kiewa, Dale R Nyholt, Jonathan RI Coleman, al. A comparison of ten polygenic score methods for psychiatric disorders applied across multiple cohorts. Biological psychiatry, 90(9):61 1-620, 2021.
[23] Ruilin Li, Christopher Chang, Yosuke Tanigawa, Balasubramanian Narasimhan. Trevor Hastie, Robert Tibshirani, and Manuel Rivas. Fast numerical optimization for genome sequeneing data in population biobanks. Bioinformaties, 37(22):4148-411 2021.
[24] Babak Alipanahi, Andrew Delong, Matthew T Weirauch, and Brendan J Frey. Predicting the sequence specificities of dna-and ma-binding proteins by deep learning. Nature biotechnology. 33(8):831-88 2015.
[25] Daniel Quang and Xiaohui Xie. Danq: hybrid convolutional and recurrent deep neural network for quantifying the function of dna sequences. Nucleic acids research, 44(1 14111))1177-1107, 2016.
[26] Žiga Avsec, Melanie Weilert, Avanti Shrikumar, Sabrina Krueger, Amr Alexandari, Khyati Dalal, Robin Fropf. Charles McAnany, Julien Gagneur, Anshul Kundaje, et al. Base-resolution models of transcription-factor binding reveal soft motif syntax. Nature genetics, 53(3):354-366, 2021.
[27] David R Kelley, Jasper Snoek, and John L Rinn. Basset: learning the regulatory code of the accessible genome with deep convolutional neural networks. Genome research, 26(7):990-999 2016.
12

---

[28] Ziga Avsec, Vikram Agarwal, Daniel Visentin, Joseph R Ledsam. Agnieszka Grabska- Barwinska, Kyle R Taylor, Yannis Assael, John Jumper, Pushmeet Kohli, and David R Kelley. Effective gene expression prediction from sequence by integrating long-range interactions. Nature methods, (8(10):196-1203) 196-1203, 2021.
[29] Johannes Linder, Divyanshi Srivastava, Han Yuan, Vikram Agarwal, and David R Kelley. Predicting raa-seq coverage from dna sequence as a unifying model of gene regulation. Nature Genetics, pages 1-13, 2025.
[30] Pau Bellot, Gustavo de Los Campos, and Miguel Pérez-Enciso. Can deep learning improve genomic prediction of complex human traits? Genetics, 210(3):809-8) -19 2018.
[31] Rostam Modollahi-Arpanaal Daniel Gianola, and Francisco Petagaricaao.. Deep learning versus parametric and ensemble methods for genomic prediction complex phenotypes. Genetics Selection Evolution, 52:1 15, 2020.
[32] Arnór I Sigurdsson, Ioannis Louloudis, Karina Banasik, David Westergaard, Ole Winther, Ole Lund, Sisse Rye Ostrowski, Christian Brikstrup, Ole Birger Vesterager Pedersen, Mette Nyegaard, et al. Deep integrative models for large-scale human genomics. Nucleic Acids Research, 51(12)x667-667, 2023.
[33] Adrien Badré, Li Zhang, Wellington Muchero, Justin C Reynolds, and Chongle Pan. Deep neural network improves the estimation of polygenic risk scores for breast cancer. Journal of Human Genetics, 66(4):359-369, 2021.
[34] Clara Albiñana, Zhihong Zhu, Andrew Schork, Andrés Ingason, Hugues Aschard, Isabell Brikell, Cynthia M Bulik, Liselotte V Petersen, Esben Agerbo, Jakob Grove, et al. Multi-pgs enhances polygenic prediction by combining 937 polygenic scores, Nature communications, 14(1):4702, 2023.
[35] Han Li, Jianyang Zeng, Michael P Snyder, and Sai Zhang. Prs-net: Interpretable polygenic risk scores via geometric learning. In International Conference on Research in Computational Molecular Biology, pages 377-380. Springer, 2024.
[36] Zijie Zhao, Tim Gruenloh, Meiyi Yan, Yixuan Wu, Zhongxuan Sun, Jiacheng Miao, Yuchang Wu, Jie Song, and Qiongshi Lu. Optimizing and benchmarking polygenic risk scores with gwas summary statistics. Genome Biology, 256)):260, 2024.
[37] Frederik Träuble, Lachlan Stuart, Andreas Georgiou, Pascal Notin, Arash Mehrjou, Ron Schwessinger, Mathieu Chevalley, Kim Branson, Bernhard Schölkopf, Cornelia van Duijn, et al. Multi-megabase scale genome interpretation with genetic language models. arXiv preprint arXiv:2501.077 2025.
[38] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems,
2022.
[39] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
[40] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santi- ago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283- 2020.
[41] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint 2tt:::2000004776.. 2020.
[42] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et Rethinking attention with performers. arXiv preprint arXiv:2009.14994 2020.
13

---

[43] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEF/CVF international conference on computer vision, pages 10012-10022, 2021.
[44] Ali Hassani, Wen-mei Hwu, and Humphrey Shi. Faster neighborhood attention: Reducing the o(n"2) cost of self attention at the threadblock level. In A. Globerson, L. Mackey, D. Bel- grave, A. Fan, U. Paquet, J. Tomczak. and C. Zhang, editors, Advances in Neural Information Processing Systems, volume 37, pages 64717-64734. Curran Associates, Inc., 2024.
[45] Tomaz Berisa and Joseph K Pickrell. Approximately independent linkage disequilibrium blocks in human populations. Bioinformatics, 32(2):283, 2015.
[46] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprint vrxiv:1909 190900803,, 2019.
[47] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 2023.
[48] Eric Y Durand, Chuong B Do, Peter R Wilton, Joanna L Mountain, Adam Auton, G David Poznik, and J Michael Macpherson. A scalable pipeline for local ancestry inference using tens of thousands of reference haplotypes. bioRxiv, pages 2021-01, 2021.
[49] Tsung-Yi Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár. Focal loss for dense object detection. Proceedings of the IEEE international conference on computer vision, pages 2980-2988, 2017.
[50] Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 7482-7491, 2018
[51] Iya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 1.Xiv:1711.05101, 2017.
[52] Oliver Pain, Kylie Glanville, Saskia P Hagenaars, Saskia Selzam, Anna E Fürtjes, Héléna A Gaspar, Jonathan RI Coleman, Kaili Rimfeld, Gerome Breen, Robert Plomin, al. Evaluation of polygenic prediction methodology within a referrenee-ttandardied framework. PLoS genetics, 2021.
[53] Hannah Taylor, Melissa Lewins, M George B Foody, Oliver Gray, Jelena Bešević, Megan C Conroy, Rory Collins, Ben Lacey, Naomi Allen, and Lucy Burkitt- Burkitt-Gray. Uk biobank-a a unique resource for discovery and translation research on genetics and neurologic disease. Neurology: Genetics, 1 11())x200026, 2025.
[54] Alicia R Martin, Masahiro Kanai, Yoichiro Kamatani, Yukinori Okada, Benjamin M Neale, and Mark J Daly. Clinical use current polygenic risk scores may exacerbate health disparities. Nature genetics, 51(4):544559) 2019.
[55] Tobias L Lenz, Aaron Deutsch, Buhm Han, Xinli Hu, Yukinori Okada, Stephen Eyre, Michael Knapp, Alexandra Zhernakova, Tom WJ Huizinga, Goncalo Abecasis, et al. Widespread non- additive and interaction effects within hla loci modulate the risk of autoimmune diseases. Nature genetics, 47(9): 2015.
[56] The M3 Research Team. 23andme technical white paper: Overview of 23andMe GWAS release: 18 18. Technical report, 23andMe, Inc., 2023.
[57] Christopher C Chang, Carson C Chow, Laurent CAM Tellier, Shashaank Vattikuti, Shaun M Purcell, and James J Lee. Second-generatio plink: rising to the challenge of larger and richer datasets. Gigascience, 1(1):s13742-015, 2015.
14

---

A Details of the GWAS runs
The internal ancestry classifier assigns individuals to one of five major genetic ancestry groups African American (AFR), European (EUR), East Asian (EAS), South Asian (SAS), or Latino (LAT) based on local ancestry inference [48]. To reduce confounding introduced by population structure, GWAS analyses were stratified by these genetically inferred ancestry groups.
Principal component analysis (PCA) was conducted separately within each ancestry group using a subset of <100,000 high-quality genotyped variants shared across all internal platforms. A randomly selected subset of individuals was used for each group: 513K for AFR, 398K for EAS, 1M for EUR, 1M for LAT, and 111K for SAS [48].
For each disease and ancestry group, GWAS was performed using logistic regression with additive allelic effects predictors. Covariates included age, sex, genotype platform (to adjust for batch effects), and top principal components specifically, the top 5 PCs for EUR, EAS, and SAS; the top 6 for AFR; and the top for LAT. Association p-values were derived using likelihood ratio test, comparing a reduced model fitted using covariates only to a full model fitted with both additive genetic effects and covariates [56]
B Exploration of LD-based variant pruning
We additionally experimented with training PRSformer on subset of variants that had been LD- pruned using PLINK 2.0 [57]. LD pruning removes highly correlated SNPs to retain approximately independent markers. In this procedure, sliding window is moved across the genome, pairwise linkage disequilibrium (1-) is computed among variants, and SNPs exceeding a specified correlation threshold with nearby variants above that threshold. Starting from a union variant set constructed similarly to that in Section 3.4 (but with slightly adjusted filtering thresholds), we applied PLINK 2.0 with window size of 6,000 kb (6 Mb), step size of one variant, and an y² threshold of 0.5. Supplementary Table F11 compares two models from the hyperparameter tuning round trained with and without LD-based variant pruning. Interestingly, despite reducing multicollinearity among variant features, LD pruning led to lower validation performance, suggesting that PRSformer benefits from leveraging the local correlation structure within LD blocks to capture causal signals more effectively.
iteratively removed until pair within each window remains
C Details of LDpred2 runs
For each disease we used the Gibbs sampler LDpred2 software (5] on the summary statistics with an internal LD panel. LDD matrix computation included variants with minor allele frequency greater than 0.1%, and genotype call rate greater than 90%. Variants greater than 5cM apart were assumed to be independent. Summary statistics were filtered to keep variants that had a minor allele frequency greater than 0.1% and had a genotype call rate greater than 95%. This consisted of variant sets with roughly 445,000 variants. We estimated posterior SNP-effect sizes using the grid option with a set of 100 combinations of hyperparameters, leading to up to 100 sets of polygenic risk scores (PRS) per disease (depending on convergence). The hyperparameters that L.pred2 takes are an estimate for the proportion of causal variants, p. and trait heritability, 2 We used LD score regression to estimate h then used a grid of the 2 estimate multiplied by 0.6, 0.8, 1, 1.2, and 1.4. For p. we used sequence of values equally spaced on a logarithmic scale from 10-5 to 1. The best hyperparameters for each disease were selected based on validation AUROC leading to the final LDPred2 PRS models. For Stacked LDPred2, however, we ensembled all of the converged PRSs per disease (up to 100) by training elastic net on the validation data using 5-fold cross validation.
D Subsetting variants for down-sampled experiments
When training on smaller datasets, we may not have access to the same high-powered variant selection as in the full-data setting. To account for this, we repeated variant selection using GWAS summary statisties adjusted to reflect the reduced sample size of each downsampled dataset. For each downsampled dataset, we estimated GWAS p-values under the reduced sample size using the original
15

---

GWAS summary statistics:
B
Nde N
==
Zds = Z X SE'
Pds = ( [dsl (-
where / and SE denote the effect size and standard error from the original GWAS, N is the original sample size, Nas is the downsampled sample size, and Φ is the standard normal cumulative distribution function. Subsequently, variant selection was performed independently for each disease and each downsampled dataset using the estimated p-values, applying threshold of < 1e-2. Multi- task model training was then conducted using the union of the selected variant sets across diseases at each downsampled scale, following the same procedure as in the full-data experiments. Additional details on variant sets and data sizes and model configurations are provided in Supplementary Tables F1,F2 and F13.
E Supplementary figures
T10
20p 1
INPIPIO Thydidss
6:009
8:044
0.529
0,41
5.224
2.21
0.117
0.002
LSA P7P0
Ap
Migubaa Clansh
Escerta Zari
0:311
6:001
2.00
0.221
MI
Figure EI: Benchmarking PRSformer against baseline methods using AUPRC as the evaluation Metiic Numbers above the bars indicate test set AUPRC values; error bars denote 95% confidence intervals estimated via bootstrapped test samples. The reported p-values reflect the one-sided statistical significance of PRSformer outperforming LDPred2.
DURSS
5.47 5355
719 0333
8444 0.577 0.40
Breake
0.555 SA2Z
2.41
3.41
0.547
4734 672
L205
0.50
052 0
052
02
Figure E2: Benchmarking of models trained on the combined training and validation datasets. Numbers above the bars indicate test set AUROC; error bars represent 95% confidence intervals computed via bootstrapped test samples. The two sets of p-values reflect the one-sided statistical significance of PRSformer+ outperforming stacked LDPred2, and PRSformer outperforming non- stacked LDPred2 the latter being the same as those reported in Figure 2 of the main text.
16

---

T10
Diveet
Sieer's
993
Hodurall Trydldl
$.00
p-0.0
0:116
9.03
0.48
6221
0.112
0,07
0111
12
USh
Gad Spa
gz
Markhyylmidds
Rigan Anez
2 6.00
0.301
0.37
1:11
5.708
5.1
Figure E3: Comparison of PRSformer and LDPred2 the kinship-controlledd European test evaluated by AUROC. Numbers above the bars indicate AUROC values, and error bars represent 95% confidence intervals estimated from 10,000 bootstrapped samples. The reported p-values reflect the one-sided statistical significance of PRSformer outperforming LDPred2.
SAppor'
Periab
Irelo Enpradto
AURO2
Bn
053
3.0
7.50
1.5
035
2.0
6.8
150
530
1.5
0.5
138
2.50
C55
174
1xe
3ne
INoiatii:
Andd Stantyhettiito
Pigratery
Cariorr Soru
Alagso Areata
3.53
9.53 610
ThE
Balting dabs den
Traring C3C1
TanRio des 020
Traninn OLD d2
Figure E4: Prediction performance acToss different training scales for the non-lineas model and linear baseline, both trained in multitask (MT) setting using the same input variant set as PRSFormer. For most diseases, performance improves with more training data, with the non-linear model surpassing the linear baseline at larger scales. Fluctuations in performance for Polymyalgia Rheumatica and Axial Spondyloarthritis likely stem from the former's rarity and the latter's relatively small training size (see Supplementary Table F2).
17

---

Genvee
Sparens
Profbels
Hovd
4.00
LOPRNA)
1.0
150
3.0
2.0
Lupu5
Rieunon Amires
M.Rtde Strengs
Chnters
Plymayggl Mhrumatics
Leteraatin Callss
03.5
615
1.52
3.50
Li6
3.00
Li0
Adc
925
1.5
0.25
8.50
Pediatic 8D:
PAAM Sponsyoatttisss
Migiitory Gloontt
Chrise Soon
Alsper Aaaats
6.000
9.34
0.55
3.29
3.50
0.3
8.5
0.15
A/6
6.9
252
3.00
025
3.39
3.50
350
Figure E5: Prediction performance across different training scales for the non-linear model and linear baseline, both trained in a multitask (MT) setting using subsetted variant sets at each down-sampled scale (see Supplementary section D and Table F13). For most diseases, performance improves with increasing training data, with the non- non-linear model outperforming the linear baseline at larger scales. These trends are consistent with those observed using fixed input variant set across scales (see Supplementary Figure E4).
18

---

T1D, Non-linear Models
T1D, Linear Models
0.85
0.80
AUROC
0.75
0.70
NT: non-ineer ST: non-lincrr
0.70
NT: Linear STo Linwar
PnStormer
2.70M-05
9-30w (1.0e+00)
1.06-000 (2.0a+00)
2.5 (3.0440))
2.70e105 (3.0e405)
9.20um (1.0wetbh)
1.85e105 02. (2.0e+06)
7.5/e
13.00
Non-linear Models
Graves', Madels
0.75
AUROC
0.70
0.65
0.60
3.00=-55
9.60e+05
1.03e+06
3.71e-36 (3.56406)
2.90e-05 (3.06405)
9.60e- 03 (1.0e105)
3.71e+06 [2.00/06]
1.03e-06
(3.
(1.0es
(2 (06)
Lupus, Non-linear Models
Lupus, Linear Models
AUROC
0.60
0.56
0.54
0.54
9 6.76++05
1.75e=06
2.03+405
9.76-=01 1.90esdr 3.75e-06 training data size
training data size
Figure E6: Comparison of multitask (MT) versus single-task (ST) training for three additional diseases across different training scales, all using the same input variant set as PRSformer: X-axis values outside parentheses indicate ST training sizes, and those inside indicate corresponding MT training sizes. Across all tested diseases, MT training outperforms ST training for both the non-linear model (left) and the linear bascline (right).
19

---

Celiac, Non-linear Models
Celiac, Linear Models
0.78
0.76
AUROIC
0.76
0.74
0.72
0.74
0.70
(3.0c
(1.0a105)
(2.0cn0d)
(3.5c
(3.0ard5)
(2.0c105)
(2.0cr0e)
(3 (3.0c106)
T1D, Non-linear Models
T1D, Linear Models
0.83
0.82
AUROC
0.82
0.81
0.80
0.79
0.79
9.30w4DS
3:00
3.37==26
9:30w=05
(3.0c+05)
(3.0c-05
(2.0cr0)
Graves', Non-linear Models
Graves', Linear Models
0.74
0.76
AUROC
0.72
0.70
0.68
0.66
2.90s-05
9.66e+05 [1 06)
1.83e+06 -06)
3.71e=06 (3.5e
2,90e-d5 (3.0e+05)
9.66c=05
1.83e-06
3.71e-06 (3.0e000
(3.
(1.0wa
[2.0e408]
Lupus, Non-linear Models
Lupus, Linear Models
0.62
0.58
AUROC
0.60
0.56
0.58
0.54
0.56
0.52
2.936105
1.05e106
3.75e
2.936105
1.050106
3.75e
(3.06-05)
(1.0e800)
(2.06-06)
(3.56406)
(3.06-05)
(2.06-06]
(3.96+06)
training data size
training data size
Figure E7: Comparison of multitask (MT) versus single-task (ST) training for four tested diseases across different training scales. Both MT and ST models were trained on downsampled datasets using subsetted variants; additionally. ST models used disease-specific variant sets (see Supplementary section and Table F13). X-axis values outside parentheses indicate ST training sizes, while those inside indicate the corresponding MT sizes. Across all tested diseases, MT training outperforms ST training for both the non-linear model (left) and the lincar baseline (right). These results are consistent with those observed using fixed input variant set for both MT and ST across scales (sec Supplementary Figure E6).
20

---

F
Supplementary tables
Table Fl: Number of selected variants per disease across European datasets
Disease
EUR dataset
3M
2M
1.5M
1M
0.5M
0.3M
0.15M
Celiac
5576
3380
2722
TID
16176
12607
7865
6424
5323
4033
3371
2530
Graves'
14104
6614
5231
3725
2984
2096
Sjögren's
11949
9381
4692
3502
2633
1954
1439
845
Psoriasis
24685
19621
11943
9088
6770
4529
3478
2148
Hashimoto Thyroiditis
29806
25103
15141
11314
7869
4677
3256
2030
Lupus
13618
10252
4840
3305
2384
1512
943
494
Rheumatoid Arthritis
18426
13683
6303
3980
2497
1484
891
373
Multiple Sclemsis
13347
10124
5587
4309
3301
2328
1683
1226
Crohn's
14140
10444
5126
3397
2099
739
3012
90
Polymyalgia Rheumatica
8113
5818
2440
1674
1135
558
267
78
Ulcerative Colitis
14138
10381
5081
3487
2300
1042
547
164
Pediatri IBD
8677
5679
1954
1095
650
204
72
24
Axial Spondylsurthritis
6352
2215
6688
316
122
2U
4
Vitiligo
12980
9720
5231
3926
2900
1695
1092
525
Migratory Glossitis
14309
8692
4516
3327
2394
1457
1019
610
Canker Sore
18418
6181
3204
2126
1322
485
206
34
Alopecia Areata
9741
7117
3160
2087
1325
680
312
59
Union
137245
94176
38018
23397
15060
9181
6829
4863
Table F2: Training sample sizes (case control) per disease across European and multi-ancestry
datasets
Dhesse
Malti-nnzodyy Danaci
Full Dateet
3M
2M
1.5M
IM
0.SM
0.IM
0.15M
Crtoe
18851/244488
1310:/100
10851
1947/251818
10897151699
3550017
8248/
231414566474
1626/228712
54377/4767527
28822/778555
2231/289748
Sptrre''s
23480/4914499
19546/371 37
7562/
0688336
25881484220
1497/250448
Pasinis
3550125
118647/3
59461/ 13910012
39530
(463485
278004
139104
Hashinoo Thenddii
130911
83363/
41581 4449004
27953 964094
144010
Lugus
27795/3747935
1243 1976554
2144/ 292815
24232214
1182141
00647//411177
11411177
27066 941039
13610/470156
141143
Miltilh
17015/3738896
14594074
2170 456400
Polsmodggi
7403/ 4623991
$392/ 9912/900500
Ulueative Coliii
4636460
44568/
EYLY CISTE
THAL 1100T
Pediankk
10034J
1571-3571512
Milligo
Spoalyeoathelss
38206/ 4556376
336188
2930M/ 35501 86
6845/29500 222019
15844 121089
15561 (100665
1-961 EA00E
Mignatry Gloeatie
1312511
22832/1001 1O01
(8852) 08851160554 19657/97T59
T0M (SWNT1
Cercer
S74878/19444
491951 1899611
5440000 401661
251536 599174
299056/241400
139865/168 (14900
64644 64012
14630 4E0T4
Alopezis
38915/4496767
22506/352266
19305 2751491
11820/1835462
6740-117671
5085 1915775
0888/45899
1925/225141
9019 199684
21

---

Table F3: Test data sample sizes (case/control) per disease across ancestry groups
Disease
EUR
AFR
LAT
EAS
SAS
Celiac
2160/243495 243495
80/46973
630/ 172120
80/8072
TID
2720/ 457259
615/ 105176
1391 /093332
77 55295
48/17309
Graves'
3388 / 469283
914/110197
1766/317074
422/5560
47947946
Sjögren's
2816 472870
520/111322
1310/319124
157/5611
50/ 18187
Psoriasis
21015 431 120
2521 101166
9615/292208
1530/50308
569 16048
Hashimoto Thyroiditis
45846
804/ I 10307
6093/31274
367 /56697
317 17674
Lupus
3936 473772
1137/111165
249000
225/56313
72/ 18237
Rheumatoid Arthritis
3280/ 101623
7071 /297862
507
188 16702
Multiple Sclerosis
2025 446886
475/ 101646
780/ 297888
36/51393
29 16457
Crohn's
3005 455308
513/ 105068
1033 305040
67/53528
66/ 16929
Polymyalgia Rheumatica
708/435628
124/ 287923
10/49222
Ulcerative Colitis
5298 / 452158
796/ 104667
2518 2118/33312 303129
1
147 16845
Pediatric IBD
1147 / 449546
191 104216
546/ 3021
33/52611
37 16782
Axial Spondyloarthritis
1086 36301
47/3144
363/ 16630
53/2073
22/503
Vitiligo
5109 / 446859
1246 102158
3199/ 298421
298/51520
214/16422
Migratory Glossitis
(092/3647)
108 4993
18007
39/2133
11/565
Canker Sore
28275/ 8842
2442/2463
16662/178
1579/596
284/27/
Alopecia Areata
3263/ 445400
2198/ 100221
3527 2948599
606/ 50802
358/ 15913
Table F4: Validation dataset sample sizes (case/control) per disease (only European individuals)
Disease
EUR Validation Size
Celiac
2096 262416
TID
3088 0088/488880 488880
Graves"
3469/ 506279
Sjögren's
2964/ 509623
Psoriasis
22883 467419
Hashimoto Thyroiditis
14555 4951 193193
Lupus
4173/51100
Rheumatoid Arthritis
14315 /779237
Multiple Sclerosis
2203 /84295
Crohn's
3449/ 493246
Polymyalgia Rheumatica
662/47291
Ulcerative Colitis
5966 489957
Pediatric IBD
1285 / 487101
Axial Spondyloarthritis
1153 39350
Vitiligo
5320/ 484047
Migratory Glossitis
1374 46592
Canker Sore
36835 / 11355
Alopecia Areata
3567 482771
22

---

Table F5: Tuning of the training steps based on validation AUROC
Hashimoto Thyrodiitis
Rheumatoid Anthritis
Mulriple Selerosis
Celiac
Crohn's 's
Graves'
Lupus
Psoriasis
Sjögren's
TID
Model
Epoch
model-1059
0.7732
0.6843
0.7621
0.7412
0.6295
0.7165
0.6593
0.5911
0.6401
0.8392
model-10559
3.0
0.7368
0.6592
0.7323
0.7225
0.6087
0.6829
0.6415
0.5780
0.6164
0.8162
model-115
1.8
0.7775
0.6869
0.7623
0.7417
0.6279
0.7137
0.6581
0.5901
0.6357
0.8398
model-t16
1.9
0.7785
0.6895
0.7632
0.7435
0.6345
0.7173
0.6600
0.5924
0.6371
0.8404
model-r17
2.0
0.7784
0.6904
0.7648
0.7446
0.6341
0.7218
0.6613
0.5939
0.6405
0.8424
model-t18
2.1
0.7463
0.6717
0.7386
0.7235
0.6147
0.6905
0.6410
0,5782
0.6164
0.8173
modcl-119
2.2
0.7430
0.6666
0.7341
0.7193
0.6099
0.6829
0.6376
0.5758
0.6132
0.8132
Note. Most models (e.g., model-(059) showed signs of overfitting beyond 2 training epochs. We also explored training durations around this point (1.8-2.2 epochs) and selected epochs as the final configuration.
Table F6: Tuning of attention heads and model dimension based on validation AUROC
Hashimoto Thyrodiiis
Rheumatoid Arthrins
Multppe Selorosis
Crohn's is
Gravess
Psoriasis
Sjogren's
Colinc
Lapus
Model
#Atten. Head #d model
Meun
Model-117
0.7784
0.6904
0.7648
0.7446
0.6341
0.7218
0.6613
0.5939
0.6405
0.8424
0.7072
Model-120
32
0.6794
0.7568
0.7418
0.6584
0.5902
0.6343
0.6992
Model-121
0.6892
0.7656
0.7464
0.7208
0.5454
Model-122
0.6909
0.7640
0.7449
0.7192
0.6612
0.5927
0.7070
Model-123
0.7714
0.6819
0.7527
0.7404
0.6259
0.7136
0.6558
0.5897
0.6315
0.7002
Model-124
72
0.7667
0.6865
0.7649
0.7448
0.6289
0.7131
0.6611
0,5932
0.6403
0.8397
0.7039
Model-125
80
0,7659
0.6871
0.7643
0.7446
0.6278
0.7127
0.6613
0.5941
0.6406
0.8405
0.7039
Model-126
5
40
0.7741
0.6874
0.7602
0.7426
0.6309
0.7145
0.6591
0.5912
0.6377
0.8401 0.7038
Note. We selected 4 attention heads with dmnste = 64 for their competitive performance despite smaller model dimension compared Model-t21.
Table F7: Tuning of attention dilation and number of transformer blocks based on validation AUROC
Haltimoto Thyroidiiis
Rheummatodd Anthrits
Maltiple Sclerois
Psoriasis
Sjögren's
Celinc
Coobn's
Graves
Lupus
Model
#Trunsfommer Blocks NA Dilation
Meun
Model-99
(1-1)
0.7807
0.6938
0.7772
0.7508
0.6422
0.7297
0.6654
0.5987
0.6501
0.8425
0.7131
M12el-112
2
(1-2)
0.7806
0.6904
0.7766
0.7507
0.6401
0.7292
0.6639
0.5991
0.6497
0.8435
0.7126
Model-113
(1-2-3)
0.7808
0.6921
0.7749
0.7498
0.6398
0.7291
0.6650
0.5982
0.6459
0.8427
0.7118
Note. Increasing the number of transformer blocks and applying dilated attention (e.g.. (1-2),
(1-2-3)) led mild overfitting. We selected 2 transformer blocks without dilation (1-1) as the final configuration.
23

---

Table F8: Tuning of the window size of neighborhood attention based on validation AUROC
Hashimoto Thyroiditis
Rheumateid Arthritis
Muleiple lerrosis
Celiae
Crohn's 's
Graves
Lap
Psoriusis 0.6625
Sjögren's
TID
Model
#Atten. Head
#NA Window
Model-t75
0.7806
0.6944
0.7700
0.7415
0.7265
0.5975
0.6311
0.8418
Model-1127
4
129
0.7796
0.6955
0.7699
0.7410
0.6310
0.7265
0.6618
0.5963
0.6322
0.8412
Model-t128
4
257
0.7821
0.6922
0.7679
0.7403
0.6299
0.7246
0.6610
0,5963
0.6303
0.8410
M2del-1129
4
513
0.7798
0.6945
0.7684
0,7412
0.6292
0.7252
0.6619
015969
0.6305
0.8413
M0del-r130
4
641
0.7787
0.6940
0.7691
0,7411
0,6308
0.7233
0,6625
0,5959
0.6342
0.8405
Model-t131
2
129
0.7799
0.6927
0.7670
0.7398
0.6298
0.7252
0.6615
0,5954
0.6297
0.8403
Model-t132
2
257
0.7785
0.6936
0.7676
0.7408
0.6303
0.7234
0.6617
0,5962
0.6315
0.8404
Model-t133
0.7799
0.6946
0.7694
0.7413
0.6297
0.7256
0.6619
0.5961
0.6319
0.8410
Model-t134
2
641
0.7807
0.6923
0.7668
0.7404
0.6286
0.7247
0.6618
0.5948
0.6307 0.8407
Note. We selected 4 attention heads with Neighborhood Attention's window size of 385 as the final configuration.
Table P9: Exploring the impact of LD-based variant set pruning (r = 0.5) on validation AUROC.
Hashimoto Thyrosdios
Rhcumatoid Arthrits
Multiple Solerosis
Celiac
Crohn's
Graves'
Lupus
Psoriasis
Sjögren''s
T1D
Model
LD-pruned
model-t031
<444
No
0.7640
0.6103
0.7334
0,7335
0.6192
0.6717
0.6572
0.5755
0.6340
0.8328
model-t034
232K
Yes
0.7564
0.5951
0.7154
0.7179
0.6143
0.6595
0.6466
0.5732
0.6300
0.8238
Note. Interestingly, removing correlated variants via LD pruning lowers AUROCs, indicating that PRSformer benefits from the underlying LD structure when identifying causal signals.
Table F10: Exploring different output heads based on validation AUROC.
Hashimoto Tbyroidííis
Rheumaroid Arthrits
Multiple Selerosis
Psariasis
Sjögren's
Celinc
Crohn's
Graves'
Lupus
Model
Output Head
model-15 Flatten-FC
0.7799
0.6896
0.7773
0.7497
0.6385
0.7301
0.6655
0.5981
0.6491
0.8448
model-t10
0.7343
0.5678
0.6656
0.6682
0.5840
0.6410
0.6036
0.5590
0.6007
0.8004
model-t11
0.5820
0.5348
0.5955
0.5822
0.5516
0,5187
0.5504
0.5292
0.5601
0.5996
Note. In model-t10, we introduced [CLS] tokens into the input sequence and vocabulary. each with learnable embeddings and full attention over all variant tokens. The flattened, normalized representations of these [CI.S] tokens were passed to a linear layer producing an 18-dimensional output. In model-t11, 1, average-pooling layer was applied to the normalized representations from the last transformer block (reducing tensors from B x L x d to d), followed by a linear layer mapping d to the 18-dimensional output.

---

Table F11: Comparison of explained variance between PRSformer and LDpred2 across diseases on the European test set
Tlyrokciess
P-lyay Ricunanie
Spondyloattriis
Muntgle Soloodis
Polunk IBD
Cluker Sure
Ukzraaive
Colle
GTanns
Pomani
Leges
Villlge
Model
PRStormer
0.1269
0.3189
0,0855
0.0224
0,0570
0.3039
0U1231
0.0163
0J866
0.0922
0.226
0,0288
0,0364
0,0179
0,8239
0.0998
0,0055
0.0165
LDpod22
0.0874
0.0545
0.1236
0.887
00178
0.441
0.275
0,8891
Note. Predicted probabilities from both models were calibrated on the test data.
Table F12: Number of selected variants per disease across non-European ancestries
Disease
AFR
LAT
EAS
SAS
Celiac
0
7515
0
0
TID
7067
8546
Graves'
6964
8889
6611
0
Sjögren's
5926
6645
0
0
Psoriasis
6211
10545
6524
6041
Hashimoto Thyroiditis
6342
11037
5780
Lupus
6172
7413
Rheumatoid Arthritis
6342
8104
4453
Multiple Sclerosis
5936
6730
0
Crohn's
5177
4959
Polymyalgia Rheumatica
0
0
0
Ulcerative Colitis
5254
5743
Pediatric IBD
4862
Axial Spondyloarthritis
0
0
Vitiligo
5599
7017
Migratory Glossitis
6927
Canker Sore
5732
6812
4296
5021
Alopecia Areata
5705
7521
4626
Table F13: Characteristics of multitask models trained on down-sampled datasets with subsetted input variants
Model
Input Variants Train Data Size # Model Parameters
Non-linear downsampldddd
Non-Linear_downsanpha.
941 76
3000000
M
Con-linear_downsamge
38018
2000000
48.73M
Non-linear_dae len-inearr_dwnsappleed
23397
1500000
30.01M
Non-lincar dowanampees
15060
1000000
19.34M
Non-lincar_done nne-lincar_downsmpplel
9181
500000
11.82M
Non-lincar_dowwe sincar_downnampled
6829
300000
8.81M
Non-lincar Non-lincar_downssnne downsampled
4863
150000
6.29M
Linear downsampled
94176
3000000
M
Lineur dowdeaple
38018
2000000
48.66M
Lincar_downsen dowesppeow
23397
1500000
29.95M
Lineur dowwnsmpeed
15060
1000000
19.28M
Linear downsumpled
9181
500000
11.75M
Lineur dowwsampldd
6829
300000
8.74M
Linear_ downsampldd
4863
150000
6.22M
25

---

Table F14: AUROC comparison across diseases under different covariate settings (sex and age). For each disease and ancestry group, bold font denotes the best performance.
ddcanand Ariiiii
(olmpatgs Ruuamaic
Attal Spooah lardbiis
Hadiinnt Thynollii
Mignarry Cloonte
Alppsia Amaza
Ulrraine Collie
BID
Orlerr Scm
Graves
Malple
Ooka''s
Rantaas
Lapes
Ancety
Model
Contriatn
PRStomerr
0.761
0.7154
10004
02045
0.3502
0.7504
0,000
0,661
01441
0.337
PRSomee
0,855
0,79917
0,139
0.7199
0.7W87
0.6684
0.626
0.6295
0.7245
Misommer
scs+age
0.6077
0/2213
PRSoomes
0.595
0.6458
0.1166
PRSomee
0.7972
0.6476
0.7197
PRSomer
mo0e
0.219
03675
015341
0.5726
0.559
0.5401
0,0038
PRSomer
06546
06153
0.5674
sas+age
0.9917
PRSGomer
8.6866
07431
0.6215
0.9771
0,5720
0,7007
0.5460
0.0155
0.567
0.6070
06343
0.0021 0,205
0,6777
PRSSomee
sesrage
06640
07883
0.7150
0,6344
0,6217
0,5973
0,6894
0.6633
PiSiomas
mo0e
I1227
07700
0.832
133112
0.3610
0.540
033579
0.991
PRSGonee
56%
0,7999
0.7455
0.6261
0,0042
0664
0.0200
0.5689
0.6674
0.6512
0.6299
0.5751
0,6262
0.7088
PRSGumer
xa+eag
0/791
0
0.1875
0.216
0/7816
0,020
0 02235
0.6770
0.6677
08855
0.5001
0.609
0.5981
0.6651 02611
0.1881
MR
mome
0.7511
0.7561
0.7256
0.9979
06533
08516
0.6038
0.6842
0.5905
08099
0.6366
0.221 0.083
0.7220
PRSlamar-ME
8.7649
0.7915
0.6545
8.8000
0,200
06225
02235
0.5063
0.6827
0.8081 0.5001
01417
0.638)
0.024
6.7299
LAT
PRStomee-
sex+age
0/764 0.7792 6.170
1L7W5K 8AN19
IL7671 1M09 6.559
0,8700
BLOMRI 018222
0
066426
0,522
0-7241
8A5
PRSGomer
2002
0.110 0.6640
0.7667 0.7341
0.6423
0,0052
05550
05556
0.5999
0.6630
0,6703
0,5209
0,5999
0.5744
0,007
0.562]
8A5
PRStomer
50%
8,8183
0,650
0.7755
0.7496
0,6330
1.455
05500
0.5557
0.6618
0,2377
0,6224
86055
03583
01885
11225
8A5
PRStomer
sestage
0.8422 0.683
0.8185 0,7536
0.7044
1.7479
(55500) 0,5193
0.6684
0.2479
0.6186
06446
03567
0/682 112593
8AS
PRStomer ME
no0e
01002
0A853
0.7612
017505 0.AS4
119744
10000
(6079) 0.1224
0.7061
0.800 0.6369
U3413
0.597
10080
015140
SAS
PRSGomer ME
sea
0L8108
0,6743
0.7784
0.7813
0.6563
3.6666
05955
027053
0.6221
0.6732
0,2452 0,5230
0.6089
0.6842
0.6653
(1243)
0.5992
SAS
PRSGomer ME
sesiage
0.5096
0,6796
0,5231
0,7737
0.7561
0.7632
0,60
022977
0,6559
(U6248
0,7580 0,6325
0,602
0,681
0.6851
0,770L
015548
EA3
Rssummer
1000
06011
1.550 0.3WI
015586
015453 03611
06644
0.1181
0.653
03161
0.1556
010204
(1377)
EAS
MRSiomer
X23
0.657 0.3357
Bml36 0.6106 0.5906
0.5425 0.7208
0.803
<0.3
0.5928
0.7355
0.5181
0.5642
0.5995 (16027
0.6541
0.5659
EAS
PRStomer
ses+age
0.658 0.7515 1.6612 0A151
117759
015512 0.8856
0.7164
0.5640
(15032
0.7491 0,0245
U.5773
06151
0.6607
017667
0,0660
RAS
MESomma-MII mone
0.7121 077877 01:0119 BASIS 11MIS 1 015010
0.691
0.1660
(16256
0.6876 0,6555
11507M
0.6545
0.6176
RAS
PRSiomar-MII v2a
(17001 0.7674 1 9955 0.6765 0.8100
11
0.6661
0.7704
0.6061
(1 (15660
0.7518 0.6511
1L9973
(16545
0.6918
0.6015
E45
PRitomee ME sesragge
0.7157 0.7%64
0.6455
0.6222 0.7975
0,6193
0.558T
0.7450
0.65%
0,0001
0,7729
0,6449
0.5972
0,6662
0,0I50
0,7742
0,7041
26

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction claim the development of scalable Transformer (PRSformer) for disease prediction from large-scale genotypes, outperforming linear and summary-statistee methods, demonstrating benefits of non-linearity at scale and multitask learning, and improving rross-ancestry generalization. These claims are supported by the experimental results presented in Section 4 and associated figures/tables.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Section explicitly discusses limitations, namely the inability share the private dataset and the need for future work on interpretability methods.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
27

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: The paper presents an empirical study based novel architecture and its experimental validation. does not include theoretical results, theorems, or formal proofs.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
The paper provides substantial detail on the model architecture (Section 3.2, Fig 1), datasets and filtering (Section 3.3, Appendix A). variant selection (Section 3.4, Appendix A-B, and D), training procedure (Section 3.5), bascline implementations (Section 3.6, Appendix C), and hyperparameter selection (Section 3.5, and Appendix Supplementary Tables F5-F10). The GitHub repository containing the model architecture is also referenced in Section 5.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
28

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: The paper states the dataset is from large, private biobank and cannot be shared publicly (Section 3.3, Section 5). The implementation code has been released with the accepted version (Section 5), although reproducing the main experimental results remains infeasible without access the private data.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ ptaaps.cc/puulcc/gu.d.oo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Section provides details on data splits (3.3), variant selection (3.4), training procedure including optimizer and scheduler (3.5), and model architecture (3.2). Appendix F and associated supplementary tables provide extensive details on hyperparameter tuning and selection.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
29

---

Answer: [Yes]
Justification: Figures 2 and E1-E3 report performance metrics (AUROC, AUPRC) with error bars explicitly defined 95% confidence intervals calculated via 10k bootstraps. P-values for significance testing (one-sided paired bootstrap test) comparing PRSformer to LDPred2 are also provided (Section 4.1, Figure 2 caption).
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative етror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Section 3.5 mentions training was performed on four NVIDIA A100 GPUs using Distributed Data Parallel and FP16 for epochs (ss 120, 000 gradient updates). While total wall-clock time isn't explicitly stated, the key hardware and training duration metrics are provided.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: The research fully conforms to the NeurlPS Code of Ethics. It uses consented data under IRB approval (Section 3.3. 5) and discusses ethical considerations like data privacy, potential misuse, and fairness (Section 5).
Guidelines:
30

---

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Section 5 discusses potential positive impacts (disease prediction/discoverr) and negative impacts/risks (genomic data misuse, discrimination). It also touches upon fairness across populations and responsible deployment.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
The paper uses sensitive genomic data. Safeguards mentioned include using data from a private biobank under IRB approval (Section 3.3, Section 5) and explicitly stating that the trained models are withheld to mitigate misuse potential (Section 5).
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
31

---

12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: The primary asset is the private biobank data obtained under IRB protocol. The methods use standard libraries/techniqes (e.g.. PyTorch, AdamW) or cited methods (e.g,, LDPred2 [5], NATTEN [2 44)) for which explicit license discussion within the paper text is not standard practice or necessary. No other major external assets requiring specific license attribution appear be used.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. www.coo///aaaaee ham/datasete has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The paper introduces new model architecture (PRSformer), which is well- documented (Section 3.2, Figure 1, and Appendix F). Furthermore, the code is available GitHub (Section 5). However, no new publicly released datasets are provided with the submission (data is private and trained models are withheld).
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the datasct/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The research existing data from biobank where individuals previously consented to participate in research (Section 3.3 and 5). The study did not involve new recruitment, direct interaction with human subjects, crowdsourcing.
Guidelines:
32

---

The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [Yes]
Justification: The paper explicitly states that the data was utilized under IRB IRB-approoved protocol (Sections 3.3) and references the use of consented research participant data (Section 5). Given it's secondary analysis of existing, likely de-identified data under IRB over- sight, specific risks incurred by this study are minimal and typically covered by the initial consent/IRB review; the paper focuses broader ethical considerations of genomic data (Section 5).
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core methodology involves custom Transformer architecture (PRS- former) with neighborhood attention, trained directly on genotype data. No LLMs were used as a core component of the method development experiments.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
33


Reviews:
[{"review_id": "tLQH5FRr8a", "reviewer": "Reviewer_V3P1", "summary": "This 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.", "questions": "1) 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?", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 2, "strengths_and_weaknesses": "Strengths:\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\".", "was_revised": true, "final_justification": "I am not an expert in the field and recommend a weak accept (4) with low confidence (2) based on potential impact. To the Area Chair: I was disappointed that this submission contains no code whatsoever for peer review, which in my opinion hampers transparency and reproducibility. According to the authors they did this to \"preserve the integrity of the double-blind review process\", which is an odd argument. Perhaps this is the first time they are submitting to an ML conference. Also, it will be interesting to see the final verdict of reviewer 8raR who seems to be an expert in the field.", "rebuttal": "## **Response to Reviewer V3P1**\nWe 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.\n\n---\n\n### 1. On performance on smaller datasets and overfitting (Figure 3)\n\n>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?\n\nThis 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.\n\nA 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.\n\n>I'm curious how Figure 3 looks like for the training data.\n\nWhile 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.\n\n**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.\n \n---\n\n###  2. On the contribution of cross-ancestry generalization\n>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.\n\nThis 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.**\n\nCurrent 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.\n\n**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.\n\n**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. \n\n---\n\n###  3. On interpreting learned non-linear interactions\n>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?\n\nWe agree that model interpretability is a critical long-term goal. However, scientifically we deliberately frame this research direction as a two-stage process:\n\n- **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.\n- **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.\n\nTo 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.\n\n**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.\n\n---\n\n### 4. On providing code for review\n>No code provided for review.\n\nWe 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.\n\n**Action:** We include code with publication.\n\n---\n\nWe 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!"}, {"review_id": "4dHKLGDDEv", "reviewer": "Reviewer_8raR", "summary": "This 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.", "questions": "* 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", "limitations": "Yes", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "After detailed discussion with the authors and other reviewers, all major concerns have been addressed:\n\n* Kinship and Confounding: The authors performed a new kinship-controlled analysis, constructing a stringent test set with no close relatives between train and test cohorts. Results confirm that PRSformer still outperforms LDpred2, supporting the claim that non-linear and interaction effects explain the improved performance.\n* Baselines: The rationale for using LDpred2 as a baseline over BOLT-LMM/GEMMA is convincing given prior benchmarking and computational constraints at this scale.\n* Clarity of Contributions: The main finding that deep, non-linear models can outperform linear methods at scale is now robustly supported. The new evidence justifies this as a substantive contribution to the ML community.\n\nRemaining Issues:\nSome concerns about the chosen model architecture and external validation remain, but these do not undermine the main result.\n\nRecommendation:\nBased on the new kinship-controlled analysis and the overall strength of the findings, I now recommend acceptance. The paper establishes a new benchmark and scaling law for genetic ML models at unprecedented scale, and all key reviewer concerns have been satisfactorily addressed.", "rebuttal": "## **Response to Reviewer 8raR**\nWe 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.\n\n---\n\n### 1. On Validation, Kinship, and LD Pruning\n>No external validation...No Control for Kinship...No LD Pruning\n\nThese 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.\n\n- **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.\n\n- **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.**\n\n- **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.\n\n**Action:** We will add a section to Methods clarifying our prediction-focused rationale and include results from our negative LD pruning experiments.\n\n---\n\n### 2. On the Flattened Output Layer and Baseline\n>...simply flatten the sequence of embeddings... This is a very unconventional modeling choice... \n\nThis is an excellent question. We did in fact test standard alternatives and found our \"flattened\" architecture superior for this prediction task.\n\n* **Ablation of Pooling/CLS token vs. Flattening:** We evaluated both **mean pooling** and dedicated **CLS tokens** against our current architecture. \n  * 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.\n  * 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. \n\n* 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).\n\n- **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.\n\n**Action:** We will add a concise paragraph to the Methods summarizing these ablation results to justify our architecture.\n\n---\n\n### 3. On Other Technical Points\n> Fixed-size attention... No covariates... Imbalance Handling... Figure 3... No cross-validation...\n\n- **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.\n\n- **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.\n\n- **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. \n\n- **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. \n\n- **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.\n\n- **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.\n\n- **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.\n\n- **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.\n\n**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.\n\n---\n\n### 4. On Other Methods\n\n- **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. \n\n- **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.**\n\n- **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.\n\n**Action:** We will add a discussion of Phenformer to our Related Work section, contextualizing our work with respect to its scale and findings.\n\n---\n\n### 5. Suitability for NeurIPS\n\nWe 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:\n\n- **A scalable framework (PRSformer) :** The first viable architecture to make end-to-end learning on million-scale individual genotypes feasible.\n- **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.\n\nOur 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."}, {"review_id": "s38R6CTDIK", "reviewer": "Reviewer_8Lrv", "summary": "This 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.", "questions": "How 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?", "limitations": "A 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.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 2, "strengths_and_weaknesses": "Quality\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.", "was_revised": true, "final_justification": "The authors did a very good job of responding to the various critiques raised in the reviews.  Despite the relatively straightforward methodology, the results are noteworthy and I believe the paper should be accepted.", "rebuttal": "## **Response to Reviewer 8Lrv**\n\nWe 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.\n\n---\n\n### 1. On Experimental Details\n\n>What statistical test was used for the p-values in Figure 2? \n\nThank 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.\n\n**Action:** We will add this clear definition to the Figure 2 caption.\n\n>How did you ensure that your hyperparameter tuning did not end up biasing your results?\n\nThis 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.\n\n**Action:** We will add a sentence to the Methods to clarify this standard practice.\n\n---\n\n### 2. On Design Choices\n>I am surprised that you opted not to include positional encodings... Did you try using positional encodings?\n\nYes, 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.\n\n**Action:** We will add a brief sentence to the Methods clarifying this design choice.\n\n>Did you do any filtering to remove close relatives between the train and test sets?\n\nWe 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.\n\n**Action:** We will clarify this prediction-focused rationale in the Methods.\n\n---\n\n### 3. On Model Scale and Data Access\n\n>How many parameters in your model, and what is the evidence that it converged after only two epochs?\n\n- **Model Parameters:** The main PRSformer model has ~180M parameters. A detailed breakdown of the other model’s variants is provided in **Supplementary Table E10.**\n\n- **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.\n\n**Action:** We will add a sentence in the main text referencing Table E5 as evidence for our chosen training duration.\n\n>Another big problem is that the data is all private... I think the statement... is misleading.\n\nThank 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.\n\n**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.\n\n---\n\nWe 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."}, {"review_id": "hjtm1rMQyE", "reviewer": "Reviewer_f3oc", "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.", "questions": "N/A", "limitations": "N/A", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths\n\n0. The empirical results are promising.\n\nWeaknesses\n1. 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.\n2. 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.\n3. 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.\n4. 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.\n5. 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.", "was_revised": true, "final_justification": "All my concerns have been well-addressed.", "rebuttal": "## **Response to Reviewer f3oc**\n\nWe 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.\n\n---\n\n### 1. On Novelty and Significance\n\n>Limited technical novelty: The performance advantage of Transformers over linear models... is well-established... making the core contribution unsurprising.\n\nWe 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*. \n\nOur contributions are:\n- **A Scalable Framework:** We developed the first viable architecture to successfully tackle disease prediction from million-scale individual-level genotypes.\n\n- **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.\n\nWe view this discovery as a central contribution of our study.\n\n**Action:** We will revise the Introduction to state this contribution more forcefully, framing the paper around the discovery of this fundamental scaling principle.\n\n---\n\n### 2. On Biological Insight \n>Lack of biological insight... neighbourhood attention... restricts the model to learning only local genic interactions.\n\nWe would like to highlight two key insights our work provides, which were only discoverable at this scale, as well as clarify our architectural choices.\n\n- **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*).\n\n- **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. \n\n**Action:** We will add a discussion on these computational trade-offs to the Methods section.\n\n\n---\n\n### 3. On Evaluation Scope and Practical Trade-offs\n>Constrained evaluation scope... Unaddressed practical trade-offs...\n\n- **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.\n\n- **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.\n\n**Action:** We will add a brief sentence to the Discussion acknowledging the computational cost trade-off, and implications for clinical practice.\n\n---\n\nWe 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."}]

Author rebuttal:

Meta-review: This paper proposes PRSformer, a scalable Transformer architecture with neighborhood attention for predicting disease risk from million-scale genotypes, and its central contribution is the demonstration of a fundamental scaling law: non-linear models only begin to outperform strong linear baselines once we have millions of samples. The reviewers initially noted concerns about technical novelty, modeling choices, kinship confounding, lack of external validation, and the absence of code in the submission, but these were largely addressed through rebuttal and additional analyses, particularly a rigorous kinship-controlled evaluation that confirmed the robustness of the main findings. While the architecture itself is not highly novel and some reviewers argued that the work may be more natural for a genomics venue, the consensus converged on the significance of the empirical result, which establishes an important new benchmark for genetic risk prediction and highlights an ML insight with broad impact. The reviewers appreciated the authors’ clarifications, additional experiments, and reframing around the scaling law, and ultimately raised their assessments to support acceptance. Given the strong empirical contribution, the demonstrated scaling phenomenon, and the potential to spur follow-up work on more data-efficient architectures, I recommend to accept this paper.
Decision: Accept (poster)