ARGformer: learning on ancestral recombination graphs with transformers

Abstract

Recent advances in inference of the ancestral recombination graph (ARG), which describes how segments of chromosomes trace back through recombination and shared lineages, have made it possible to reconstruct genome-wide genealogies for large cohorts, but it remains difficult to summarize and use this information for population genetic analyses. We present ARGformer, an encoder-only transformer that learns context-dependent embeddings with a self-supervised masked objective finetuned with contrastive learning for downstream retrieval tasks. We train ARGformer on genealogies from coalescent simulations and on genealogies inferred from ancient and present-day Homo sapiens genomes. Using only these learned embeddings, without access to genotype matrices, ARGformer captures patterns of global population structure and supports ancestry inference through clustering and nearest-neighbor retrieval. On genealogies that include archaic hominins, ARGformer can highlight Denisovan-derived segments in Oceanian genomes and reveals Oceanian-like ancestry in South American Indigenous populations. ARGformer is available at https://github.com/AI-sandbox/ARGformer.