PRSformer: Disease Prediction from Million-Scale Individual Genotypes
Payam Dibaeinia, Chris German, Suyash Shringarpure, Adam Auton, Aly A. Khan
Abstract
Predicting disease risk from DNA presents an unprecedented emerging challenge as biobanks approach population scale sizes (N > 10 6 individuals) with ultra-highdimensional 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 ≈ 5M) for predicting 18 autoimmune and inflammatory conditions using L ≈ 140k 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 > 1M), and that a multi-ancestry trained model improves generalization, establishing PRSformer as a new framework for deep learning in population-scale genomics.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Builds on6
- Swin Transformer: Hierarchical Vision Transformer using Shifted WindowsZe Liu, Yutong Lin, Yue Cao, Han Hu et al.ICCV 2021 · 31,683 citations
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessTri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra et al.NeurIPS 2022 · 5,493 citations
- Big Bird: Transformers for Longer SequencesManzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie et al.NeurIPS 2020 · 3,159 citations
- Rethinking Attention with PerformersKrzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song et al.ICLR 2021 · 122 citations
- Faster Neighborhood Attention: Reducing the O(n^2) Cost of Self Attention at the Threadblock LevelAli Hassani, Wen-Mei Hwu, Humphrey ShiNeurIPS 2024 · 23 citations
Related papers
- Training Flexible Models of Genetic Variant Effects from Functional Annotations using Accelerated Linear AlgebraAlan Nawzad Amin, Andres Potapczynski, Andrew Gordon WilsonICML 2025
- Neuroformer: Multimodal and Multitask Generative Pretraining for Brain DataAntonis Antoniades, Yiyi Yu, Joseph Canzano, William Yang Wang et al.ICLR 2024 · 21 citations
- Neural Pharmacodynamic State Space ModelingZeshan M. Hussain, Rahul G. Krishnan, David A. SontagICML 2021 · 12 citations
- Nyströmformer: A Nyström-based Algorithm for Approximating Self-AttentionYunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan et al.AAAI 2021 · 675 citations
- Proxy-Bridged Game Transformer for Interactive Extreme Motion PredictionYanwen Fang, Wenqi Jia, Xu Cao, Peng-Tao Jiang et al.ICCV 2025
