An Embarrassingly Simple Approach to Enhance Transformer Performance in Genomic Selection for Crop Breeding

1 Introduction

Grain production security serves as the cornerstone of human existence, exerting a pivotal role in shaping the health, stability, and prosperity of global society. Addressing global hunger not only aligns with the United Nations Sustainable Development Goals (SDGs) [United Nations, 2023a] of Zero Hunger, but also adheres to the Leaving No One Behind Principle (LNOB) [United Nations, 2023b]. However, according to the report by the Food and Agriculture Organization of the United Nations in 2023 [United Nations, 2023c], around 735 million people worldwide were suffering from hunger. Therefore, utilizing technological advancements to boost grain production is vital for eliminating hunger by 2030. Crop breeding stands as a fundamental approach to enhancing grain production. Continuously improving crop varieties to increase yield and resilience effectively boosts grain production, thereby meeting the escalating demand for food.

Genomic selection (GS), proposed by [Meuwissen et al., 2001], is regarded as a promising breeding paradigm [Boichard et al., 2012; Garc´ıa-Ruiz et al., 2016]. GS usually involves predicting the phenotypes of polygenic traits in plants or crops by using high-density markers covering the entire genome. By implementing early screening of candidate populations, the genetic process can be accelerated, reducing generation intervals and significantly shortening breeding cycles. Unlike phenotype selection [Siepielski et al., 2013; Kingsolver and Pfennig, 2007] and marker-assisted selection [Ribaut and Hoisington, 1998; Xu and Crouch, 2008], GS can utilize single nucleotide polymorphism (SNP) data obtained from organisms. SNP sequences are loci sequences with two or more alleles extracted from DNA sequences, which are the most widely adopted genetic markers in GS. 1

1 Stripping away segments that exhibit the same patterns from DNA sequences to obtain SNP sequences is advantageous for analyzing genotype-phenotype associations, as phenotypic differences

A mainstream solution to GS is to mine the statistical knowledge of SNP data. GBLUP [VanRaden, 2008] and SSBLUP [Goddard et al., 2011] utilize kinship matrices to weight and aggregate the phenotypic values of different individuals, thereby obtaining their estimated breeding values (EBVs) but may not fully explore and utilize genomic information. Another category of statistical methods involves estimating marker effects on a reference population and then predicting marker effects on a candidate population, followed by linear aggregation to obtain EBVs for the candidates [Meuwissen et al., 2001; Tipping and Faul, 2003; De Los Campos et al., 2009; K¨arkk¨ainen and Sillanp¨a¨a, 2012]. These methods typically demand substantial computational resources and lack parallelization capabilities. Their practicality in time-sensitive breeding scenarios is constrained. In addition, two common limitations of statistical methods are the strong statistical priors (e.g., Gaussian distribution) and linear relationship assumption between markers. Though these two assumptions simplify the statistical computation, statistical models suffer a setback when the underlying data distribution differs.

In order to capture the non-linear relationships between markers, efforts have been made by applying deep learning to process SNP sequences, especially with convolutional neural networks (CNN) [Ma et al., 2017; Liu et al., 2019; Xie et al., 2023] and Transformers [Wu et al., 2024]. However, two major data challenges threaten the robustness of deep learning models. First, the number of samples for each crop species is limited. This can easily cause overfitting for deep learning models, while statistical models are commonly deemed as robust solutions. Second, SNP are long sequences. Besides, plant genomes, unlike those of animals, demonstrate a higher frequency of long insertions and deletions attributed to the activity of transposable elements [Ramakrishnan et al., 2022], leading to a greater abundance of SNP data among different plant genomes. The locality of convolution introduces a strong inductive bias for local dependencies, thus cannot catch the long-range interactions between markers. On the contrary, while Transformers can better tackle long-range dependencies in sequences, attention's quadratic complexity becomes the major bottleneck for long sequence [Dao et al., 2022].

To mitigate the effect of long sequences, an intuitive way is to statistically pre-process the data before building deep learning models. However, the limitations of statistical methods still remain. For example, DNNGP utilizes principal component analysis (PCA), a statistical dimensionality reduction technique, and condenses the sequence to only several hundred dimensions for CNN modeling [Wang et al., 2023]. However, the viability of this hinges on fulfilling the assumption of PCA which states that there exist linear correlations among variables in the dataset, but this might not be true for the majority of SNP datasets. Thus, the linear dimensionality reduction techniques may result in information loss regarding loci that influences phenotypes. Another statistical way to shorten the sequence is traditional feature selection, i.e., selecting important features out of hundreds of thousands of SNPs [Manor and Segal, 2013]. For example, a promising method is to employ genome-wide association studies (GWAS) to identify major effect loci within gene sequences [Liu et al., 2019]. However, these methods heavily rely on the accuracy of GWAS and are not conducive to modeling and analyzing quantitative traits.

often arise from distinct segments between two DNA sequences.

To address the aforementioned challenges of statistical and deep learning methods, we propose a simple yet effective Transformer-based method that supports end-to-end training. The proposed method leverages simple tricks such as k-mer tokenization, and random masking to achieve robust predictive performance from genotypes to phenotypes on crop breeding datasets. It is worth mentioning that, though similar natural language processing (NLP) techniques have been investigated on DNA data due to the similarities between DNA sequences and text sequences [Ji et al., 2021], there is no such application on SNP data yet. To the best of our knowledge, we are the first to successfully adapt these simple NLP techniques to address the GS problem. Compared with statistical methods, our method does not require rigid statistical priors or linear assumptions, and thus can better capture the non-linear relationships. Supported by GPU, our method has a significantly shorter inference time. Compared with existing deep learning methods, our method can better leverage the attention mechanism [Vaswani et al., 2017] to assist contextual understanding. We summarize the differences between our method and seminal methods in Tab. 1.

We extensively evaluate the robustness of our method on two crop datasets, rice3k [Wang et al., 2018] and wheat3k. rice3k is a public dataset on rice with SNP data. On the rice3k dataset, we outperform the current best-performing method (a hybrid method of GWAS and Transformer) on average over 1.05% in the accuracy metric. wheat3k is a private cereal grain dataset to be released, which contains 3032 SNP sequences following a similar setup of rice3k. Our method also consistently outperforms the seminal baselines on wheat3k.

• We propose an end-to-end Transformer-based framework for genomic selection that enables capturing nonlinear relationships between genotypes and phenotypes.
• We show that in genotype-to-phenotype prediction tasks, using k-mer tokenization and random masking can effectively reduce the data complexity while enhancing model prediction performance.
• We conduct extensive experiments on the rice3k and wheat3k datasets, where our method achieves the stateof-the-art performance on both datasets against mainstream methods.

2 RelatedWork

2.1 Analysis-based Genomic Selection

Currently, analysis-based genomic selection methods are mainly classified into two families: best linear unbiased prediction (BLUP) methods and Bayesian methods [Gualdr´on Duarte et al., 2020]. BLUP-based methods evaluate random effects using a linear model. GBULP [VanRaden, 2008] constructs a pedigree relationship matrix using genomic information, incorporating it as a random effect into the model to estimate genetic values or predict trait values of individuals. RRBLUP [Endelman, 2011] builds upon GBLUP by incorporating the concept of ridge regression. SSBLUP [Goddard et al., 2011] integrates both the genomic relationship matrix and pedigree relationship matrix, along with phenotypic data, into a unified mixed model. For Bayesian methods, such as [Tipping and Faul, 2003; De Los Campos et al., 2009; K¨arkk¨ainen and Sillanp¨a¨a, 2012], the marker effects are assumed to follow different Gaussian distributions. For example, BayesA [Meuwissen et al., 2001] assumes that each marker has its own distribution and variance. Though Bayesian methods can reveal the relationship between genotype and phenotype to some extent, they are constrained by strong statistical priors and linear assumptions. In addition to the two families of methods, statistical learning methods [Ke et al., 2017] may also improve prediction accuracy on specific datasets. However, this improvement might not generalize, especially when the number of samples is significantly smaller than the number of feature dimensions.

2.2 Deep Learning-based Genomic Selection

Fueled by the recent success of deep learning in vision and language tasks, researchers have attempted to integrate deep learning into the Genome Selection (GS) field. DeepGS [Ma et al., 2017] proposes a genome-wide selection framework based on CNN, while DualCNN [Liu et al., 2019] utilizes a dual-stream CNN to predict quantitative traits. ResGS [Xie et al., 2023] uses ResNet [He et al., 2016] to extract gene sequence information without using pooling layers. DNNGP [Wang et al., 2023], a seminal work, performs PCA to reduce the high dimensionality of gene data to extract effective information. However, the limited convolutional receptive field hinders the model's ability to capture linkage disequilibrium loci in long SNP sequences.

Transformer [Vaswani et al., 2017] has exhibited excellent performance in modeling global interactions within sequences. GPformer [Wu et al., 2024] applies Informer [Zhou et al., 2021], a model from the long-term time series prediction field, to aggregate periodic information and capture region-specific features in SNP sequences. However, the computational complexity of Transformer also has a quadratic relationship with the sequence length. The current Transformer-based methods [Wu et al., 2024] are unable to handle the whole genetic SNP sequences for most crop species. There are hybrid methods [De Los Campos et al., 2009; Liu et al., 2019] that use GWAS to pre-process the raw SNP sequence, but the results heavily depend on the results of GWAS and exhibit unstable performance across species. Instead, we propose to leverage tokenization and optimization acceleration techniques to achieve full attention computation on long SNP sequences, thereby mining long-range interactions and obtaining more accurate and robust predictive models.

2.3 Sequence Representation

A straightforward method to handle SNP data is to directly model loci using additive encoding, which is widely adopted by seminal methods [VanRaden, 2008; Aguilar et al., 2010;

Mittag et al., 2015]. But, it is overly rigid to represent genotypes by the number of non-reference alleles.

An emerging research topic is to leverage language models to model DNA data. However, unlike human language, where each word carries rich meaning, DNA's lexical units consist of only four basic nucleotide bases (A, T, C, G) with relatively ambiguous meanings, making them lower-level abstractions. For example, DNABERT [Ji et al., 2021] and DNAGPT [Zhang et al., 2023] have shown that encoding DNA sequences into k-mer patterns can improve the performance on downstream tasks without losing information. DNABERT-2 [Zhou et al., 2023] uses random masking for unsupervised pre-training on DNA sequences. However, as there are fundamental differences between DNA and SNP sequences, it is still unclear whether these NLP techniques work on SNP data. In this work, we present the first empirical study.

3 ProblemFormulation

We formalize the genomic selection problem as a mapping learning task. Given a genotype-phenotype pair (x, y), the final goal is to learn a mapping (or function) f : x ↦→ y. In this work, the genotype data x is the raw SNP sequence. Let S = {s1, s2, . . . , sNl } denote the SNP sequence, where Nl is the sequence length. We have si ∈ {A, T, C, G, N, Y, K, W, R, S, M}, where different letters represent different combinations of alleles [Johnson, 2010]. In this work, the definition of each letter is presented in Tab. 2. The phenotype data y represents the corresponding phenotype values, which can either be discrete or continuous, depending on the task of interest (i.e. classification or regression).

Let Ns denote the number of samples for the species of interest in a crop dataset. In contrast to common machine learning tasks, we have Ns ≪ Nl. This poses a non-trivial challenge for overfitting-prone methods such as Transformer, and thus Transformer cannot be directly used.

4 Methodology

The focus of our framework is to mitigate the limitations of Transformer and capture the contextual information and language patterns of SNP sequence. To this end, we first preprocess raw SNP sequence by an initial mapping which does not change its length. After that, we implement k-mer tokenization and random masking to convert SNP sequence to token ID sequence. Then, embedding vectors are generated by the embedding layer and positional encoding [Vaswani et al., 2017]. Finally, the embedding vectors are mapped to phenotype predictions by a neural network. The overall end-toend training architecture is elaborated in Fig. 1.

4.1 Pre-processed SNP Sequence

In the candidate list of the sequence, {A, T, C, G} is the majority and {N, Y, K, W, R, S, M} is the minority. To avoid the difficulty of excessively large k-mer vocabularies contributed by the minority and improve the computational efficiency of subsequent self-attention, we generate the pre-processed SNP sequence Sp by the mapping rule:

4.2 Sequence Tokenizer

Considering the biological significance of SNP sequences and mitigating the attention computational cost on long sequences, we choose non-overlapping k-mer as a tokenization technique on pre-processed SNP sequence Sp . Let T = {. . . , ti, ti+1, ti+2, . . . } denote the token ID sequence processed by k-mer operation. Here is an illustration for a 4-mer operation with simulated sequence and token IDs.

mer = (·, k) is a function that treats the substring of length k as a whole entity and converts it to a unique ID. Although k-mer shortens the SNP sequence to 1/k of original length Nl, the problem of Ns ≪Nl still exists.

We introduce randomness by random masking on token ID sequence to reduce the risk of overfitting:

Mask(·) is a function that has the pre-defined probability p to convert the original token ID to mask id, which is defined separately. For the ith position in T m, we have

The final embedding vector is the sum of the embedded masked token ID sequence and the positional encoding.

Embed(·) denotes the embedding layer, where the vocabulary size is 5k +1. The "5k " represents the number of all possible combinations and "1" denotes the specific Mask ID for unknown cases. Epos ∈ R1×⌊ Nl k ⌋×Dw is the sinusoidal positional encoding [Vaswani et al., 2017], where Dw denotes embedding width.

4.3 Phenotype Learning

We use Transformer encoder to learn the contextual information of encoded SNP sequences, which consists of L layers of multi-head self-attention (MSA) [Vaswani et al., 2017] and multi-layer perceptron (MLP) blocks.

LN(·) is the layer normalization [Vaswani et al., 2017].

After computing the attention, a linear projector LP(·) is used to map the output to a lower-dimensional space, thus lowering the computational complexity in the regressor:

Finally, the phenotype prediction yp can be output as:

where the projected vector is flattened and inputted into MLP layer to predict the phenotype.

There are two types of learning tasks in predicting phenotype from genotype information: variety recognition and trait prediction, corresponding to classification and regression tasks, respectively. For the classification task, we use the cross entropy (CE) as the loss function:

where yl denotes the label.

For the regression task, we define the loss function by Mean Squared Error (MSE):

5 Experiments

5.1 Setup

Datasets We use two crop datasets: rice3k [Wang et al., 2018] and wheat3k, representing the two most significant staple crops worldwide. Each dataset includes raw SNP sequences and multiple traits of interest (phenotype).

rice3k is a public dataset that comprises genomic data and the corresponding phenotype data for a total of 3024 samples collected from 89 countries. For genomic data, the SNP sequence for each sample has a length of 404388. rice3k encompasses multiple quality traits, i.e. discrete variables. We select six valuable and well-structured phenotypes for experimentation, including APCO REV REPRO, CUST REPRO, LPCO REV POST, LSEN, PEX REPRO, and PTH. The definitions of these phenotypes are described in Tab. 3.

wheat3k is a private dataset to be released. wheat3k is collected following the setup of rice3k to provide a comprehensive study on wheat. It comprises genomic data and the corresponding phenotype data for a total of 3032 samples. Each SNP sequence for each sample has a length of 201740. In contrast to rice3k, wheat3k focuses on quantitative traits, i.e. continuous variables. We select six valuable and wellstructured phenotypes for experimentation, including COLD, FHB, LODGING, STERILSPIKE, TKW, and YIELD. The definitions of these phenotypes are described in Tab. 4.

Implementation We implement the proposed framework using PyTorch and conduct experiments on NVIDIA 3090 GPU. There are a total 80 of epochs with early stops. The dimension of the sequence embedding width after embedding layer is set as 32 (i.e., Dw = 32). Note that the vocabulary size of embedding layer is 5k + 1, where 1 additional ID for unknown words. Following the settings in DNABERT [Ji et al., 2021], we use 6-mer (i.e., k = 6) and the default random masking ratio is 15% (i.e., p = 0.15). We use 3 Transformer encoder blocks in the learning module and an Adam optimizer [Kingma and Ba, 2014] with an initial learning rate of 0.0001 and weight decay of 0.01. All experiments are repeated for five-fold cross-validation, following the standard practice in GS. The reported numbers are the average metrics and standard deviations. Following [Wang et al., 2023; Wu et al., 2024], each phenotype of interest is considered an independent task for training.

Evaluation Metrics Following the setups of [Wang et al., 2023; Wu et al., 2024], we adopt accuracy (ACC) as the evaluation metric for the classification task (rice3k dataset) and Pearson correlation coefficient (PCC) to evaluate the regression performance (wheat3k dataset) of our model. PCC is defined as.

Baselines To demonstrate the robustness of our method, we select four seminal methods for comparison, including GBLUP [VanRaden, 2008; Yin et al., 2023], RidgeClassifier [Quyet Nguyen et al., 2024], BayesA [Meuwissen et al., 2001; Lilin et al., 2022], DNNGP [Wang et al., 2023], and DLGWAS [Liu et al., 2019]. GBLUP is a representative BLUP-based method and performs well across various regression tasks. To suit the classification tasks of rice3k, we replace GBLUP with RidgeClassifier, which shares similar statistical assumptions. BayesA is a robust Bayesian method with simple assumptions. We conduct 1500 iterations to ensure convergence. DNNGP and DLGWAS are two state-of-the-art deep learning methods. We use the source code of DNNGP 2 . We pre-process SNP sequences using PCA and retain 3000 principal components to preserve 95% variance information. We implement DLGWAS as an integration of GWAS and Transformer. GWAS uses the mixed linear model [Muhammad et al., 2021] to analyze the correlation of SNP loci, retaining the top 10% of loci based on effect size. After locus screening, the relative order of loci retained in the original sequence is preserved for Transformer modeling. The model configuration information is kept as consistent as possible for all baselines.

5.2 Results

Performance onClassificationTasks Tab. 5 elaborates the classification results of different models on rice3k dataset. On the phenotypes of APCO REV REPRO, CUST REPRO, LPCO REV POST, LSEN, PEX REPRO, and PTH, our model achieves the best performance compared to the existing models in ACC metric. On LPCO REV POST, we also obtain competitive results. Specifically, we surpass the current best-performing method DLGWAS over 1.05% across these 6 traits on average. The robustness and advancement of our simple approach to enhancing Transformer are evidenced by the overall elevation in mean values and reduction in standard deviation. These outcomes position our method as a pioneering benchmark in the field.

2 https://github.com/AIBreeding/DNNGP

Performance on Regression Tasks We further analyze each baseline's capacity on wheat3k dataset, tabulated in Tab. 6. Similar to the outstanding performance on the rice3k dataset, our model also achieves the best performance compared to existing methods, where the higher PCC metric and lower standard deviation clearly prove the superiority. On YIELD task, which is important as it signifies the meaning of wheat yield, our method is 1.63% superior to the current state-of-the-art DLGWAS, which holds the potential in the application of genomic selection for crop breeding.

5.3 Ablation Studies

To investigate the contribution of each key component of our proposed method, we conduct a series of ablation experiments on the wheat3k dataset.

Analysis of Pre-Processing We conduct ablation experiments to have a deep insight into key components in the tokenization. Note, we choose the model without pre-processing module as baseline to study the impact of each component in Tab. 7. It is found that k-mer plays a more significant role than random masking, which enhances the performance of baseline model by 1.24%, 8.42%, and 3.81% in PCC on three sub-tasks of wheat3k dataset. This suggests that using k-mer in pre-processing can assist attention mechanisms to better capture contextual information. Another finding is that when utilizing k-mer as tokenizer, random masking can yield greater improvements than using random masking on baseline alone. The difference lies in whether the smallest unit of the masking is 1 bit or k bits, which relates to the influence of masking on context. This proves the importance of context extraction on k-mer. Compared with DLGWAS, Transformer without k-mer and random masking shows inferior performance and our model enhances the performance. This proves the validity of key components.

Analysis of k-mer Tokenization We study the influence of k in k-mer tokenizer in Fig. 2. Note, when k = 5, our model obtains the best results on the FHB and YIELD traits, but the best results on the STERILSPIKE occur at k = 6. We hypothesize that the value of k might be task-oriented.

Analysis on Random Masking We conduct experiments on random masking proportion to study the effect of randomness in the tokenization. In Tab. 9, the PCC value of three subtasks improves as the random masking proportion increases, reaching its peak value when the proportion equals to 30%. Through randomly masking tokens, the data diversity is enhanced, thereby promoting the model's generalization ability and alleviating overfitting issues arising from data scarcity, especially in our problem formulation (Ns ≪ Nl). When the masking ratio becomes too high, i.e., 45%, it can lead to diminishing returns and worsen the effectiveness of the model because excessively masking a large portion of tokens may disrupt the coherence of the input sequence to make it harder for the model to learn meaningful representations.

Choice of Tokenizer We compare our methodwith seminal tokenization algorithms to demonstrate the effectiveness of kmer and random masking, shown in Tab. 8. Character-level tokenizer treats each letter in SNP sequence as an independent feature, where we use {A, T, C, G, X} and one additional ID for unknown words as the vocabulary. Byte-Pair Encoding builds a sub-word vocabulary of size 8000. We also implement two learnable tokenizers: one MLP layer and one Transformer encoder, where the input is SNP sequence and the embedding width of encoding output is 20. Compared with the second-best tokenizer (Learnable Transformer), our method achieves 1.20%, 2.22%, and 1.08% PCC improvements on FHB, STERILSPIKE, and YIELD, respectively, providing better SNP sequence representation for learning module.

Computational Cost Lastly, to comprehensively assess a model, computational cost is considered. Take wheat3k as an example, the original sequence length is 201740. Tabulated in Tab. 10, we utilize parameter count (Parameters), GPU memory consumption (Memory), and inference time (Time) to measure all models. GBLUP shows fast inference time as a lightweight linear model. But it sacrifices the degree of freedom, thus leading to lower accuracy in Tab. 6. After PCA, DNNGP only retains 3000 principal components for CNN, thus generating low computational costs. However, DNNGP reports the lowest performance on both two datasets. For vanilla Transformer without tokenization, the original SNP sequence of length 201740 is input into attention mechanism. For DLGWAS and Ours (6-mer), the inputs for both Transformer encoder consist of 33623 dimensions. For Ours (5-mer), this value is 40348. Compared with the statistical method BayesA, our method provides a quicker inference time, and compared with vanilla Transformer, our method efficiently reduces GPU memory cost.

6 Conclusion

In this study, we present a simple yet effective approach to enhance Transformer's performance in the realm of genomic selection for crop breeding. Specifically, we propose the preprocessing module composed of k-mer tokenizer and random masking to assist contextual understanding of SNP sequence. Experiments on the rice3k and wheat3k datasets demonstrate promising performance on genotype-to-phenotype prediction. The empirical findings of this work not only suggest the potential of DL to handle long sequences but also pose a new research direction on end-to-end genomic selection.

Meanwhile, we shall notice that though tokenization plays an important role in genomic selection, a comprehensive understanding is still needed in future work.

Acknowledgements

This work was supported by Shanghai Artificial Intelligence Laboratory and the Strategic Priority Research Program of the Chinese Academy of Sciences under Grant No.XDA0450203.

Contribution Statement

Renqi Chen, Wenwei Han, and Haohao Zhang have equal contributions. Nanqing Dong is the corresponding author.