Zong, L.; Friedler, S. A.; Mathieson, S.

Understanding natural selection can help shed light on the genetics underpinning adaptive evolution. The widespread availability of large-scale human genetic variation data has led to the development of data-driven methods for detecting signatures of selection, many of which are based on deep learning. However, these methods often fail to generalize well to the diversity of selection signatures across a broad range of evolutionary scenarios. We propose a novel transformer-based model, Popformer, for learning encodings of general patterns of genetic variation. Popformer includes site-wise and haplotype-wise attention, allowing us to capture variation among both genetic positions and individuals. It additionally learns relative positional embeddings for inter-SNP distances. The model is pre-trained with an analog of the masked language modeling objective across a range of real human genomic data, similar to the task of genetic imputation. Using dimensionality reduction, we show that the pre-trained model learns meaningful embeddings of genomic windows that correspond with population structure. We also show that the model can accurately perform genotype imputation. We further fine-tune the model on selection classification and demonstrate that our model is more accurate than other selection classification methods, on selection simulations of both well-specified and mis-specified demographic models. Using a novel real data validation approach, we apply Popformer to human data from the 1000 Genomes Project and reveal its ability to generalize from simulated data to the diversity of real data. Overall, our method provides a new direction for population genetic inference methods, with future fine-tuning applications including the inference of recombination rates, introgression, and local ancestry.