Overview of Popformer method

Popformer is an encoder-only transformer model which takes haplotype matrices as input and outputs a vector representation (embedding) of each position of the matrix (further details can be found in the Methods). The encoder uses an implementation of axial attention: where attention is computed both along the variants for each haplotype, and along the haplotypes for each variant (Fig 1) [31,32]. This allows the model to dynamically weight particular haplotype variants based on the full population and window context. Distances between variants are included as a learned positional attention bias, allowing the model to learn patterns of variant density.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 1. Overview of the architecture and training strategies for Popformer.

a, overview of the architecture and pre-training strategy of Popformer. Haplotype matrices and distance vectors are input to the model. The model consists of multiple blocks of axial attention. Each block consecutively applies SNP-wise attention across every haplotype, haplotype-wise attention for every SNP, and a fully-connected feed-forward layer. Normalized residual connections are applied after each layer. Pre-training is self-supervised, using real genomic windows. Random tokens are replaced with a mask token, and the model is tasked with unmasking the token. b, Popformer-base is first trained using self-supervised unmasking. Using it, we evaluate three different selection training recipes. For Popformer-lp, the pre-trained base model is frozen and the head is trained on the static embeddings. For the Popformer-ft, gradients from training also flow through the base model. For Popformer-no-pretrain, we ablate the pre-training step, starting from a newly initialized model. The gray background indicates where gradients are computed.

https://doi.org/10.1371/journal.pcbi.1014492.g001

During pre-training, the embedding representation is learned by training the model to unmask random positions in matrices derived from real human populations in the 1000 Genomes dataset [33]. We refer to this model as Popformer-base. We perform this pre-training primarily to learn the patterns of genetic variation in real data - population structure, patterns of LD, etc. The aim of the pre-training is to produce high-quality representations of this real variation that will enable better generalization performance on the downstream task of selection detection. To this end, we simulate a set of diverse genomic regions under human-inferred demographies, including both neutrally evolving regions and those evolving under various selection strengths. We use these simulations to fine-tune a simple linear model (the classification head) on top of Popformer-base under several different training recipes. We evaluate:

1. A linear probe (Popformer-lp), where the base model (encoder) is frozen and only the linear head is allowed to update.
2. Full fine-tuning (Popformer-ft), where both the base model and the linear head are allowed to update.
3. A pre-training ablation (Popformer-no-pretrain), where a newly initialized encoder is used in place of the pre-trained base model. Here, both the encoder and the head are allowed to update.

Once trained, these models are verified on held-out test simulations, where the ground truth is known. Additionally, they can be used to perform windowed inference along real genomes.

Learned population embeddings capture informative genomic variation

Popformer-base is trained with an analogue of the masked language modeling task (see Methods) [27]. In this modeling scheme, the input sequences to be masked are haplotype matrices (S1 Fig). In order to ensure that our model accurately learns complex dependencies between SNPs, we first test the actual performance of our method on its training objective by computing the unmasking accuracy on a test set of 75%-uniformly-masked samples from 1000 Genomes populations. On this test, Popformer-base outperforms two simple baselines and performs quite well. The nearest neighbor baseline simply chooses the allele using the most similar windowed haplotype to the masked haplotype, and the allele frequency baseline chooses the most frequent unmasked allele (Table 1).

Download:

• PNG
  larger image
• TIFF
  original image

Table 1. Denoising accuracy for Popformer-base and two naive baselines.

https://doi.org/10.1371/journal.pcbi.1014492.t001

The task of unmasking haplotypes is also conceptually similar to the task of genetic imputation, whereby a ‘target’ panel of individuals has missing variant data imputed using a ‘reference’ panel. Our masked training differs from genotype imputation primarily in the nature of the missing data - in our masked training, there is not a set of complete haplotypes used as the ‘reference’. Instead, all entries in the haplotype matrix were masked uniformly at random (see Methods and S1 Fig). To compare to methods specifically designed for genotype imputation, we formulated an imputation-like test to further validate our pre-training objective performance. 64 haplotypes are used as an unmasked reference panel, and 64 other haplotypes are used as a target panel, where a percentage of SNPs are randomly masked in the whole target panel (S2 Fig). This test setup is identical to classic genotype imputation, allowing us to compare to a widely used imputation method, IMPUTE5 [34].

We calculated the squared correlation R² using the posterior imputed allele probabilities for IMPUTE5 and Popformer by comparing the imputed allele at each site to the true allele (Fig 2) [35]. We additionally calculated error rate using the maximum posterior imputed allele. Comparing R², Popformer-base performs as well as IMPUTE5 [34], and continues to outperform both baselines. We find a similar result for the error rate, where our method imputes exact haplotypes comparably to IMPUTE5. A breakdown of the results by binned minor allele frequency demonstrates that all methods perform worse at very low allele frequency, with IMPUTE5 performing the best (S3 Fig). Of course, IMPUTE5 is designed for biobank-scale panels, and is therefore orders of magnitude more efficient than our method (S4 Fig). These results are quite promising given that our self-supervised training data has no special examples with a reference-vs-target masking pattern, and demonstrate that the pre-training strategy effectively learns patterns of genomic variation.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 2. Imputation results for Popformer and other baselines.

a, imputation error rate (lower is better) and b, squared correlation R² (higher is better) at four different imputation levels (20%, 40%, 60%, and 80%). We compare performance for Popformer, IMPUTE5, and the nearest neighbor baseline. The column frequency baseline is omitted, performing significantly worse than the other methods.

https://doi.org/10.1371/journal.pcbi.1014492.g002

While trained on unmasking, Popformer is primarily designed for applications in downstream tasks. In order to achieve good performance on these tasks, it is important that the model’s pre-training objective has produced an informative learned embedding. To inspect these representations, we first examined whether the embeddings adequately represent real genomic variation between populations at least at a continent-level scale. Individual embeddings were calculated for each individual in the 1000 Genomes dataset by concatenating mean-pooled window-level embeddings for sliding windows across chromosome 1 [33].

We performed a principal component analysis (PCA) on the these mean-pooled embeddings and used the first two components to visualize them (Fig 3a). Given that even raw allele frequency patterns reflect continental-level population structure, we expect that Popformer’s learned embeddings should also represent this structure. Indeed, Popformer-base’s embeddings match commonly reported results in which dimensionality-reduced genomic variation differentiates continental populations [36,37].

Download:

• PNG
  larger image
• TIFF
  original image

Fig 3. Visualization of Popformer’s embeddings for real and simulated data.

a, PCA reduction of region-averaged learned embeddings for chromosome 1 for all 1000 Genomes populations, with colors indicating superpopulations (continent-level variation). Points represent chromosome 1 embeddings for diploid individuals. Population codes: EUR (Europe), EAS (East Asia), AMR (America), SAS (South Asia), AFR (Africa). b, PCA reduction of window-level learned embeddings for simulated selection windows at a range of selection strengths. Points represent individual haplotype matrices.

https://doi.org/10.1371/journal.pcbi.1014492.g003

Secondly, we aimed to identify whether the pre-training objective aligned at all with the task of selection detection. To do this, we calculated window-level embeddings for our selection simulations, which span a selection coefficient of 0-0.1 (Fig 3b). We observe that there is a clear gradient between low and high selection strengths. While not definitive, this indicates that the model’s pre-training task is partially aligned with the selection classification objective.

Trained Popformer models can accurately detect selection

In the standard application of a supervised selection model, the model is trained on simulations that are designed to match the downstream real data. In our case, we trained all of the supervised models on a balanced mix of human European-inferred (CEU), East Asian-inferred (CHB), and African-inferred (YRI) demographic simulations (for details on the simulations, other models, and training see Methods). In this study, we only aimed to detect positive selection, but training data encompassing negative or balancing selection could also be used. Each model classifies these simulated haplotype windows as either containing a selected-for mutation or not. We additionally include shouldering windows from selection simulations as negatives, aiming to improve localization power.

We report AUC curves on a held-out test set of CEU, CHB, and YRI simulations, and find that our method, in all training configurations, outperforms other models (fine-tuned, CEU: AUC = 0.93) (Fig 4a). We additionally stratify the metrics by selection strength and by whether windows are shoulders (to a window containing a selected allele) and find that stronger selection is easier to detect, except for when the window is in a shoulder region of a stronger selected allele. These shoulders of strong selection are likely a strong confounding signal (S5 Fig).

Download:

• PNG
  larger image
• TIFF
  original image

Fig 4. Selection detection results on simulated human demographies.

a, ROC curves for our Popformer models, FASTER-NN CNN model, resnet34 CNN, and a summary-statistic (Tajima’s D) score, evaluated on a held-out test set of CEU-inferred demographic simulations. b, for a test set of CHB-inferred demographies. c, for a test set of YRI-inferred demographies.

https://doi.org/10.1371/journal.pcbi.1014492.g004

While these results demonstrate the effectiveness of our model’s architecture at classifying selection from in-distribution simulations, the true benefits of the pre-training strategy are in the method’s ability to generalize to out-of-distribution (OOD) data. There are several types of distribution shifts that may occur in the simulation-based inference pipeline. First, the demographic history used to simulate the training data may not match the real data. Second, the real data may have various effects (genotyping error, phasing error, population structure) that are not typically simulated.

We performed several experiments to test the mis-specified demographic history case. In order to isolate the effects of the pre-training, we evaluated all methods on a set of simulations with varying bottleneck strengths and a set of simulations with different constant population sizes. Instead of using Popformer-base (pre-trained on real data) as the pre-trained model, we pre-train new base Popformer models. Then, we train on the selection task. This experiment therefore parallels the real-data pipeline, but with entirely simulated data. For the bottleneck experiment, we pre-train a base Popformer model on a 10% bottleneck, train all models (including fine-tuning and linear probe variants) on no bottleneck, and evaluate on varying strengths of bottlenecks (Fig 5a). For the constant population size, we pre-train on N = 10000, train all models on 1) N = 1000 and 2) N = 10000, and evaluate on varying constant population sizes (Fig 5b).

Download:

• PNG
  larger image
• TIFF
  original image

Fig 5. Selection detection results for models trained on varying demographic scenarios.

Solid vertical lines indicate training demographic configuration and dashed lines indicate pre-training demographic configuration (for popformer-ft and popformer-lp). For all results, bootstrapping with N = 100 resamples was applied. The bootstrapped mean and 1SD error bars for the area under the ROC curve (AUC) are shown. In most cases, the error bars are smaller than the points. a, varying strengths of bottlenecks, models trained on no bottleneck. b, varying constant population sizes, models trained on N = 100000.

https://doi.org/10.1371/journal.pcbi.1014492.g005

We find the bottleneck example to be especially informative. Since bottlenecks can leave similar signatures to natural selection in the genome, we expect that regions in stronger bottlenecks are harder to classify. We see this decrease in performance for all models. The linear probe (with the pre-trained model from data from a 10% bottleneck), though, generalizes better both to the matching 10% bottleneck data and to more extreme bottlenecks. This result also follows in the varying constant sized simulations experiment, with the linear probe model generalizing nearly as well as Tajima’s D. This finding strongly indicates that the pre-training assists in generalization ability.

The fine-tuned model only clearly outperforms the pre-training ablation in the bottleneck experiment, and does not outperform the linear probe model. We hypothesize that allowing the encoder to update, as in the fine-tuned model, likely leads to some ‘forgetting’ of the pre-training demography. In the varying constant population size experiment, the order of results is not clear. For example, the orders of the Popformer fine-tuned and pre-training ablations are inconsistent. These inconclusive results may be indicative of instability in training (only one model was trained for each configuration) or that the models are ‘overfit’ to a particular feature of the simulations (e.g., if lower diversity is the only signal needed in the training dataset, the model may not learn to use haplotype structure).

Overall, Popformer models outperform the other methods we tested on in-distribution simulations, and pre-trained models demonstrate improved generalization on out-of-distribution tests. These results show that the pre-training of our models is beneficial in the case that the distribution shift is related to mis-specified demography. We additionally hypothesized that pre-training might also be beneficial in a limited-labeled-data regime, and find this to be true at very low numbers of training examples (S6 Fig). The linear probe and fine-tuned Popformer models outperform the other deep learning methods when only 10 simulated windows are used to train the models. In practice, though, more simulated data should be used to train the models.

Validation on known selected regions

Though we display promising results on simulated tests, ultimately, real evolutionary effects encompass a wider range of variation then we can simulate. As mentioned above, real evolution encompasses both categories of distribution shifts: we are both 1) unable to know the true demography (to create ideal simulations) and 2) unable to capture the range of variation caused by other effects (technical, biases, etc). Despite this, methods are still validated on real data by inferring selection on a list of known or potentially novel sites, and examining whether predictions match the literature. We follow this tradition and analyze the commonly used LCT/MCM6 region (S7 Fig), which has a strong signature of positive selection in European populations [38].

However, this approach is prone to bias, in the form of cherry-picking a set of regions that a method identifies well, and does not provide any information on the false positive rate of a method. This approach is also uninformative on how powerful a model is, as any signals which are already well known are likely strong. We propose that selection detection methods, which are potentially most well suited for discovery (hypothesis generation), should focus heavily on controlling the false positive rate. The major difficulty with this is the fact that it is impossible to know the full ground truth for selection in the human genome.

Here, we introduce two approaches to validating selection detection methods on real data which can be used to compare methods on both power and false positive rate. First, new ancient DNA data and methodology has allowed for selection inference using real allele frequency trajectories [39]. Using these results, we gather a list of putatively neutral regions (X < 2) from Akbari et al. [29] to use as ‘true negatives’. By using the regions from Grossman et al. [28] as ‘true positives’, we can generate a pseudo-ROC-curve with the goal of comparing how well methods can differentiate between selected regions and putatively negative regions. Crucially, in this framework, we treat the Grossman et al. as a correct, but incomplete list. It is then possible that some methods may predict other regions correctly that are missing from our positive set, thus leading to a misleading comparison.

Along these lines, the composite of multiple signals (CMS) test used by Grossman et al. presents a circular validation issue. Since the CMS method uses specific patterns of variation (long haplotypes, higher differentiation, lower diversity), a method which can better recover the Grossman et al. list might merely be making predictions which correlate with the CMS method. This circularity especially affects Tajima’s D, which indeed correlates directly with some of the CMS signals.

For most populations, which lack sufficient ancient DNA, we can instead compare [28]’s true positives with the fraction of the genome predicted to be under selection, under the assumption that a majority of the genome is neutrally evolving. Both of these approaches can be interpreted as analogous to a standard ROC curve, allowing for simple comparisons. These approaches are also independent of binary prediction threshold, which is important for deep learning methods, whose outputs are rarely calibrated.

Using these approaches, we find that the deep learning methods’ strong performances on simulated data do not translate well to performance on real data (Fig 6). Notably, Tajima’s D is the method which typically aligns best with the list of selection regions from [28], with the deep learning methods falling behind. This is a worrying result for deep learning methods, which are significantly more computationally expensive relative to the simpler summary statistics. Of the deep learning methods, however, Popformer generally recovers the known regions the best. Crucially, our Popformer method outperforms all other methods, including Tajima’s D, at recovering known positive regions in the East Asian (CHB) and Yoruban (YRI) populations. This result only holds in the fine-tuned Popformer training recipe, not in the pre-training ablation, suggesting that real pre-training does provide an improvement when applying downstream models to real genomes.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 6. Validation results of selection detection on real data.

a, all models evaluated on recovery of Grossman et al. [28] known selection regions in the CEU population versus fraction of the genome predicted to be under selection. b, models evaluated on recovery of Grossman et al. CEU regions versus false discovery of Reich et al. negative regions. c, same as a, but on the CHB population d, same as a, but on the YRI population.

https://doi.org/10.1371/journal.pcbi.1014492.g006

The linear probe, though, performs worse than the other Popformer variants, which conflicts with the pattern of results from the varying simulated evaluation. Although speculation, we believe that for the simulated demographic mis-specification, updating the encoder likely caused some degree of demographic forgetting and causing reduced performance. For the real data, the simulations match the real data in demography. Instead, they are mismatched due to other effects. In this case, updating the encoder may have been necessary for the model’s ability to generalize across the simulation-real distribution shift without causing any forgetting.

These results are representative of the fact that the selection signatures can manifest differently under different demographic models, and suggest that deep learning methods may be learning specific, not general, signatures. To begin to understand which, we examine correlations between Popformer-ft’s model scores and summary statistics, both for simulated (S8 Fig) and real (S9 Fig) datasets. There are some moderately sized correlations, which are expected (Spearman rho: Tajima’s D = -0.61, # singletons = 0.56).

We also find evidence to support that using non-reported-positive regions as neutral regions is valid. Using inferred neutral alleles from ancient DNA as true negatives strikingly resembles using the fraction of the genome called (Fig 6a, 6b). Thus, if one is confident that their list of selected regions are true, using the rest of the genome as a neutral baseline in a test for recovery is likely to be informative. These analyses do not take into account potentially undiscovered selected regions in the genome – but, given that the reported regions are likely to be strong signals, detection methods should be able to predict them at higher probability than weaker and neutral signals.

Details of the Popformer architecture

Popformer, in a similar manner to other raw-data deep-learning methods, operates on matrices of n haplotype rows by S single nucleotide polymorphism (SNP) columns derived from genomic windows. Only biallelic SNPs are included. These matrices are computed in the form of binary matrices in which 0 represents the more common (major) allele and 1 represents the less common (minor) allele. We also include a vector encoding distances between SNPs, which allows the model to incorporate positional information in its predictions. Unlike most CNN methods, the transformer architecture can handle both a variable number of haplotype rows and SNP columns without additional padding. This renders Popformer more flexible when performing inference on populations that have a different number of individuals.

An overview of the architecture can be found in Fig 1. We adopted the model architecture from the axial attention transformer architecture, originally developed for image learning and later for protein sequence learning by MSA Transformer [31,32]. Attention is a widely-used method that allows each element in an input sequence to be dynamically weighted (attended to) based on every other element in the sequence. For example, text tokens in a sentence can each be weighted based on the context from the full sentence. 2D axial attention, originally designed for image methods, extends the standard mechanism of attention, which operates on 1-dimensional sequences, to operate consecutively across multiple dimensions [31]. Column-wise attention (in our case representing structure across haplotypes) is applied for every SNP with complexity O(Sn²). We use tied row attention, introduced by [32], to share the row-wise self attention maps between every haplotype in the matrix. This forces the model to learn one pattern of SNP dependency, which maps directly between different haplotypes. It also reduces memory usage from O(S²n) to O(S²). Following most transformer architectures, we first use an embedding layer to learn representations of our input tokens (major allele, minor allele, mask, start, end, and pad tokens). We include residual connections around the attention layers [49]. In all, each axial attention block consists of a row-attention block, a column attention block, and a feedforward dense network, allowing us to aggregate information about variant and population structure.

In the standard self-attention implementation, a positional embedding of a number of forms can be used to break standard attention’s positional-invariance. These forms all assume that ‘distances’ between consecutive tokens are equal. Unlike protein and image sequences, adjacent columns of SNPs in the input matrices do not represent consecutive positions within the original genome. The distances between variants provide SNP density information otherwise lost in a pure haplotype matrix. In order to encode the meaningful patterns of inter-SNP distances, we include an altered form of the T5 language model’s learned positional embedding approach [50]. Inter-SNP genomic distances (in base pairs) are first binned into a fixed number of linearly-spaced distance buckets. An embedding is learned for each bucket separately for each row-attention head and for each layer. The embedding directly modifies the attention matrix in the form of an added bias. We adapt T5’s 32 learned positional embeddings, using 31 linearly-spaced discrete distance buckets from 0 to 50k base pairs, with a final 32nd bucket reserved for any larger distances. We do not include a positional embedding for haplotypes, as the positional-invariance in this dimension allows our method to be insensitive to the ordering of the haplotypes.

Pre-training on real genomes

Our pre-trained model (Popformer-base) is trained with an analogue of the masked language modeling objective. Given masked positions using which the masked input are created, we aim to find model parameters which minimize the binary cross-entropy loss:

The model outputs logits for the major and minor alleles for every position, while the loss is computed only for masked tokens. We do not mask d, the distance vector. In our case, this masking is applied uniformly across all SNPs (S1 Fig). We use a high masking rate (75%) to prevent models from learning simple extrapolation from neighboring haplotypes or alleles, as in [51]. We experimented with additionally masking columns of SNPs, blocking the model from using context from different haplotypes, but find that this degrades downstream model performance. We additionally tried to weight rare variants by oversampling minor alleles (1s) during masking. This resulted in poor performance on validation sets as the model learned from an incorrect distribution of major/minor alleles.

The training data for Popformer’s pre-training contained 260,000 haplotype windows, sampled uniformly from all 26 populations of the 1000 Genomes phase 3 release [33]. A random population and start position along the genome are chosen. We find the nearest SNP to the start position, and perform basic quality control by reselecting a new window if less than 50% of the bases in the window are in the 1000 Genomes phase 3 strict accessibility mask. For each window, we take variants within 50,000 base pairs from the start position (only including biallelic SNPs) and take a number of haplotypes uniformly at random between 32 and 128. Thus, haplotypes within a window are guaranteed to be from the same population but their order and number is randomly selected. This pre-training data likely also contains windows with real, unlabeled selective sweeps, which might contribute to the real variation the model learns.

For all Popformer results, we used a model with four hidden layers, with each layer containing four self-attention heads each for row- and column- attention (34M parameters in 16-bit floating-point precision). This model is relatively smaller than similar models used in the vision, protein, and language domains, as we both have less complex data representations and less training data. Cross-entropy loss was used during both pre-training (at the per-SNP level) and during fine-tuning (at the window level). The Adam optimizer was used with a linear learning rate decay and a warmup period of 0.1 epochs. This pre-trained model (Popformer-base) is used as the base model for the downstream selection task by adding a simple linear classification head. We tested two downstream training strategies: full fine-tuning, in which all the weights of the model (including the transformer encoder) are allowed to update, or training only the linear head (with the encoder frozen, referred to as linear probing). We additionally test a model with ablated pre-training, in which we trained a randomly initialized encoder and linear head from scratch on the selection detection task. For these downstream selection-trained results, we trained on windows of length 50kbp (resulting in windows with variable numbers of SNPs) and subsampled to 64 random haplotypes. We note that this data configuration falls within the ranges used during pre-training.

In all, the pre-training took 6 hours to complete on a 4xB200 GPU node with 128 GB of memory. Fine-tuning took 1 hour to complete on a single RTX 5000 GPU (32 GB). Inference over 49,000 64-haplotype 50kbp windows (genome length) took 7 minutes to complete on the same RTX 5000 GPU. During pre-training, selection training, and selection inference, we used effective batch sizes of 32, 8, and 1, respectively. We observe no degradation of performance performing training and inference on different batch sizes. The masking strategy also differs: during pre-training we used 75% random masking. For the imputation comparison, we used an unmasked reference panel and random SNP masking for the target panel. For selection, we used no masking.

Selection simulations

We used the SLiM simulation framework through the high-level stdpopsim library to generate samples for training and testing [22,52,53]. To cover a wide range of demographic parameters, we use demographic parameters randomly sampled from inferred demographies for CEU, CHB, and YRI populations using pg-gan, a GAN-based approach to demographic inference [54]. We filter sets of demographic parameters which have growth rate < 0.01 for ease of simulation. A summary of demographies used is in S1 Table. We sample recombination rate uniformly at random from the HapMap recombination map and sample mutation rate randomly in [55]. We simulated a total of 5,000 neutrally evolving regions and 5,000 regions with a centered beneficial allele, all of length 250kbp. We randomly varied the strength of selection in , and chose an onset time such that the allele would roughly reach fixation by the end of the simulation using Equation A17 of [56].

For each neutrally evolving sample, we randomly choose 10 sub-windows, all of which are labeled as neutral. For each selected sample, we randomly choose 10 sub-windows not containing the selected allele to be labeled as neutral (shoulders). We choose 20 sub-windows containing the selected allele to be labeled as positive. In all, this gives us 100,000 total neutral windows and 100,000 total positive windows. These are split first into populations, and for each population into an 80/20 train/test split. In all, these human-inferred simulations took 5 hours on 20 CPUs and storage of raw genotypes took 10GB on disk.

We used the discoal simulation engine to generate the varying simulations of bottlenecks and constant population sizes. We simulated a total of 10,000 neutral regions and 10,000 regions with a centered beneficial allele of length 50kbp. The strength of selection was randomly varied in and the final selected allele frequency was set to 0.95. For the bottlenecks, we used a ancestral population size of 10000, with a bottleneck of varying strengths from between 2,000 and 1,000 generations before the present. population sizes, we simulated populations of varying constant population sizes. For the training data size test, we used simulations of constant N = 10000 for both pre-training (unlabeled) and training (labeled).

Selection methods

We found it extremely important to implement fair model comparisons. Most machine learning selection classifiers are not designed to be used out-of-the-box, but rather trained from scratch on a target demographic model. Thus, we designed an evaluation harness for training and evaluating window-based classification methods that facilitates making fair model comparisons as we did in this work. This harness was also designed to be extensible, allowing for additional models to be evaluated with minimal effort (besides the effort to implement and re-train), and is released along with the models at https://popformer.leonzong.com.

We compared our method with two deep-learning methods and Tajima’s D. A widely used architecture for selection detection deep-learning models are convolutional neural networks (CNNs). CNNs are naturally suited to the haplotype input data, both due to the input shape and due to the fact that learning convolutional filters aggregates local patterns of diversity in an appropriate way for discovering local signatures of sweeps.

FASTER-NN is a 1D-CNN that takes as input 2 1D vectors - minor allele frequencies and distances [40]. We use the same training recipe as was used in the original FASTER-NN method, except using our training simulations. We additionally compared against a ResNet architecture, which are CNNs commonly used for images. For this model, we adopt the preprocessing steps from [16,57] in which we sort haplotypes based on genetic similarity. This model was trained with a batch size of 16 and learning rate of 1e-4 for around 20 epochs, at which point validation loss stopped improving. For both models, we used a width of 512 SNPs (matching Popformer’s input length) and used zero-padding for inputs shorter than this length. The ResNet model had a maximum height of 256 haplotypes, also using zero-padding.