The deep learning framework (Figure 1) proposed in this paper consists of four main modules, which are separately feature extraction module, data processing module, classification model and result integration module. The relevant codes can be found on GitHub: https://github.com/wr-sky/ARGs/tree/main/Code.

Firstly, we utilize two different protein language models [ProtBert-BFD (Elnaggar et al., 2021) & ESM-1b (Rives et al., 2021)] to extract features from proteins sequences, which can facilitate both data augmentation and prediction accuracy compared with the single-language model. Secondly, by cross-referencing two protein language models, we designed a novel data augmentation method to enhance less prevalent ARGs examples during training process, which makes the training set more balanced. Thirdly, we used two semantic-based encoding models (LSTM & MH-LSTM) to classify the embedding results separately from ProtBert-BFD and ESM-1b. Finally, our framework provides a 16-dimension vector by integrating two classification results above. The position with the maximal value will be chosen and its corresponding ARGs type is the final prediction result.

To ensure the authority and comparability of the data, this study primarily uses data from DeepARG and HMD-ARG as the basic ARGs dataset (Arango-Argoty et al., 2018; Li et al., 2021). Besides, 2,000 non-resistant genes reported in HyperVR (Ji et al., 2023) were included for related experiments. Protein sequences in three datasets (Table 1) are compared with blastp, removing completely identical sequences (identity = 100% & coverage = 100%). For ease of reference and result reproduction, detailed information of the protein sequence in Table 1 are uploaded to GihHub: https://github.com/wr-sky/ARGs/tree/main/Data.

All ARGs in Table 1 are categorized into 16 groups. Aside from the ARGs categorized as “other”, some resistance gene groups (marked in black) are more abundant, particularly those associated with bacitracin and beta-lactam resistance. In contrast, the remaining resistance gene groups (marked in red) are less prevalent. Each protein sequence, including the resistance or non-resistance genes, will be taken as initial input for our proposed framework.

ARGs often contribute to bacterial metabolism through specific structures (both 2D and 3D) of the proteins they encode, resulting in resistance by degrading, obstructing, or expelling antibiotics (Darby et al., 2023; Kakoullis et al., 2021). We have analyzed existing protein language models and their characteristics (Supplementary material) (Elnaggar et al., 2021; Lin et al., 2023; Rao et al., 2020; Meier et al., 2021; Hsu et al., 2022). Based on the specifics of this study, we selected two pre-trained protein language models as upstream feature extractors to embed information carried by protein sequences. The ProtBert-BFD model extracts embedding vectors that capture key information from protein sequences and is also used in downstream tasks such as secondary structure prediction (Elnaggar et al., 2021). The ESM-1b model, through logistic regression and linear projection, encodes embedding features containing the secondary and tertiary structural information of protein sequences (Rives et al., 2021). Thus, this step employs these two models as feature extraction methods, embedding the sequence and structural features of the target proteins from different dimensions.

Both models take the protein sequence as input. ProtBert-BFD encodes each amino acid as a 30-dimensional vector. Each protein sequence is encoded into a 30,720-dimensional vector by padding 0 or vector truncation (1,024 amino acids). Similarly, ESM-1b encodes each amino acid as a 1,280-dimensional vector and each protein sequence is encoded into a 1,310,720-dimensional vector (1,024 amino acids).

In Natural Language Processing (NLP) tasks, data augmentation for small datasets is a crucial strategy to enhance model performance and generalization (Chen et al., 2023). However, due to differences between resistance gene protein sequences and natural language features, traditional NLP data augmentation strategies cannot be directly applied to this task. Therefore, For the first time, we designed a new data augmentation method for the limited antibiotic resistance data (Figure 2). This method exponentially increased the limited amount of resistant gene data (marked as red in Table 1), making the input data for each type of resistant gene more balanced.

We utilize Principal Component Analysis (PCA) (Hasan and Abdulazeez, 2021) to decrease ESM-1b's embedding results of each amino acid from 1,280 dimensions to 32 dimensions. The overall dimension of a protein sequence's (1,024 amino acids) embedding results is 32,768, which will further be truncated to 30,720 dimensions. For the embedding results of ProtBert-BFD, we use a 43-fold linear concatenation to directly extend the feature vector of the entire protein to 1,320,860 dimensions, which is also truncated to 1,310,720 dimensions. For each input, the green line forms one training process, and the orange line forms another training process, which can double the examples in less prevalent ARGs (red groups in Table 1).

The overall test results, which will be illustrated in the “Results” section below, provide evidence to support the feasibility of such a transformation between the two embedding spaces.

LSTM is well-suited for handling long-range dependencies in language data and can effectively capture contextual information in sequential data, making it particularly suitable for processing the temporal nature of linguistic data (Yu et al., 2019). On the other hand, MH-LSTM introduces a multi-head mechanism that allows for the parallel processing of multiple types of contextual information, further enhancing the model's ability to understand the complex syntax, semantics, and ambiguity in language. Therefore, both LSTM and MH-LSTM are ideal choices for processing linguistically encoded data, as they are better at capturing the complex dependencies and multiple layers of contextual information in language, ultimately improving the model's performance and expressive power. In this paper, we employed Multi-Head Attention LSTM (MH-LSTM) and LSTM to extract effective information while reducing the dimensionality of the protein embedding vector (Figure 3). By mixing different models, MH-LSTM can fully extract features from high-dimensional embedding results (ESM-1b), while LSTM can avoid over-abstraction of relatively low-dimensional embedding results (ProtBert-BFD).

The input feature size is 30,720 dimensions for LSTM and 1,310,720 dimensions for MH-LSTM. Both models have hidden layers and output layers of size 512. The final classification is performed by a linear layer with GELU activation, which takes input sizes of 512, 1,024, and 2,048 across three layers (Table 2). These architecture parameters were determined through experimentation and preliminary trials. The output of the linear layer is a 16-dimensional vector, with each dimension representing a group of ARGs (Table 1).

By using ensemble learning, combining the predictions of two classification models can effectively reduce overfitting (Ying, 2019). At the same time, it increases the diversity and robustness of the model, thereby potentially improving prediction accuracy. In our proposed framework, we process the results of the two models using linear integration or probability maximization integration methods (Figure 4).

Especially, for linear integration (Figure 4a), the probability values output by the two models are linearly combined, as described by Equations 1, 2.

In this context, logits1 represents the probability vector output by the LSTM network, and logits2 represents the probability vector output by the MH-LSTM. a and b are constant coefficients, and their sum equals 1. Both a and b are initialized to 0.5 and dynamically adjusted during the training process. The argmax() function returns the index y_hat corresponding to the maximum probability in the logits, which indicates the final predicted label.

For probability maximization integration (Figure 4b), the probabilities corresponding to their predicted labels are compared, and the label associated with the higher probability is selected as the final prediction, as described by Equations 3, 4.

max_prob_1 and max_prob_2 represent the maximum probability values in the probability vectors output by the two models. max_index_1 and max_index_2 indicate the predicted label types corresponding to these maximum probability values. The max() function computes the maximum value and its index in the probability vectors. Then, the sizes of max_prob_1 and max_prob_2 are compared, and the label associated with the larger probability is chosen as the final predicted type.

Although appropriately increasing the number of models in ensemble learning can usually further improve prediction performance, it is also necessary to consider computational resources and effective fusion methods. Taking these factors into account, we adopt two models (ProtBert-BFD with LSTM & ESM-1b with MH-LSTM) in ensemble learning.

Based on the aforementioned structures, we developed two overall architectures of the prediction model: Linear-integration-based architecture (LSTM_MH-LSTM_LINEAR) and Probability-maximization-based architecture (LSTM_MH-LSTM_MAX). We constructed datasets for training, validation, and test purposes. The training process includes data processing step, and the other two processes exclude the step (Figure 5).

We used the train_test_split function from the scikit-learn machine learning toolkit in Python to partition the data into training, validation, and test sets with a ratio of 0.6/0.2/0.2. We also selected Adam optimizer with a learning rate of 2e-4 for network training. To avoid overfitting, the optimizer includes a training termination mechanism: if the accuracy does not improve after 30 iterations, the training process will be terminated early, and the model weights will be saved. The final experimental results are obtained by testing the model on the test set. Performance evaluation is conducted using four metrics: accuracy, precision, recall, and F1 score, which provide a comprehensive assessment of the model's performance. Their mathematical expressions are shown in Equations 5–8, respectively.

In this context:

Precision represents the proportion of samples predicted as positive that are actually positive. Recall represents the proportion of actual positive samples that are correctly predicted as positive. The F1 Score considers both precision and recall, providing a comprehensive measure of the model's performance. These four metrics collectively account for both false-negative and false-positive scenarios. When all metrics approach 1, it indicates superior model performance with minimized misclassification rates for both negative and positive samples.

Predicting the presence of ARGs fundamentally indicates whether a bacterial strain has the potential to develop a resistant phenotype. Comparatively, direct prediction of phenotypic resistance can more effectively guide clinical antibiotic selection for bacterial infections. However, defining the antimicrobial resistance (AMR) phenotype of a target bacterial strain solely based on the aggregate features of all its resistance genes would lead to substantially elevated false-positive prediction rates. Therefore, in this subsection, we further try to adapt the architecture of our previously proposed model to predict the whole-bacterial antimicrobial resistance (AMR) phenotypes.

In critical care, combination antibiotic therapy is often employed to achieve the most rapid therapeutic effect. Therefore, false-positive predictions should be rigorously minimized, as they may lead to the avoidance of first-line antibiotics that would have been most effective. In contrast, false-negative predictions in the context of multi-antibiotic regimens typically have less detrimental impact on overall therapeutic efficacy. As traditional best-hit approach has a low false-positive rate (Arango-Argoty et al., 2018), we incorporated this approach [CARD with blast (Alcock et al., 2023)] into our model as a whole-genome screening tool at the bacterial species level. The incorporation of CARD serves dual purposes: primarily filter out both negative and false-positive genes to reduce false-positive AMR predictions, while concurrently reducing computational load for downstream AI networks to accelerate the entire prediction pipeline.

The training and validation processes are omitted and only the test process is committed. The experimental steps (Figure 6) are as follows.

Based on the original data from DeepARG-DB, we test the LSTM and MH-LSTM performance on the ProtBert-BFD (Figure 7a) and ESM-1b (Figure 7b) embedding results respectively. For the ProtBert-BFD, LSTM with 3 layers network structure achieved the best performance with a clear advantage compared with other layers and MH-LSTM structures. This is likely because the low-dimensional data (30-dimensional ProtBert-BFD encoding per amino acid) is more suitable for small-scale, shallow networks, while the multi-head attention mechanism in MH-LSTM may lose some critical information. The embedding results of ESM-1b, which encodes each amino acid into 1,280 dimensions, is significantly higher than ProtBert-BFD's results. In this case, MH-LSTM with 6 layers network structure achieved the best performance. This improvement is due to the increased overall data dimensionality, which favors deep MH-LSTM networks for effective key data abstraction and extraction, while reducing interference from noisy information. In conclusion, proteins encoded by ProtBert-BFD are more suitable for LSTM structures with fewer layers (D3), while proteins encoded by ESM-1b is better suited for MH-LSTM structures with relatively more layers (D4~D6). However, an excessive number of LSTM layers (D7 as an example) increases the model's parameter count, leading to overfitting, which in turn causes a significant decline in test performance.

Based on the overall dataset (Table 1), we evaluated the prediction performance of different architectures under the optimal structures (Table 2). The optimal structure is determined through repeated experimental iteration with different structural combinations, which also follows the conclusion above. The evaluation was carried out using four metrics: accuracy, precision, recall, and F1-score. We compared our proposed methods with related works, including traditional sequence alignment methods like CARD and machine learning methods such as HMD-ARG and DeepARG. By applying these methods to ARGs identification (Figure 8a) and ARGs classification (Figure 8b) tasks, our proposed model consistently achieved superior performance almost across all evaluation metrics. The only exception is the CARD method under “perfect” criteria, the higher precision of which is due to its more lenient criteria for identifying resistant genes. This kind of criteria will result in a lower false-positive rate but a higher false-negative rate (Arango-Argoty et al., 2018), and consequently, performs the worst under the other three evaluation metrics. Besides, the framework including the data augmentation process usually provides better results compared with the framework without data augmentation technique.

We manually screened 262 bacterial strains from NCBI which include both complete antibiotic susceptibility test (AST) results and whole-genome sequences. The strains involve S.enterica, E.coli, K.pneumoniae, C.freundii, S.marcescens, and etc. Each strain was annotated based on their AST results and our proposed 16 resistance labels (Table 1). For ease of replication studies, details of each stain can be found on GitHub: https://github.com/wr-sky/ARGs/blob/main/Data/AST_NCBI_id.txt.

In our application pipeline, we tested both “perfect” and “strict” screening standards (CARD RGI) as pre-screening tools for each protein sequence. Basically, the “strict” standard is relatively more lenient than the “perfect” standard, allowing us to optimally preserve high-fidelity resistance genes.

To comprehensively demonstrate predictive performance, we separately quantified model outputs encoded by ESM-1b (right panel of Figure 5 with MH-LSTM) and ProtBert-BFD (left panel of Figure 5 with LSTM). From the perspective of screening criteria, the test results show that datasets filtered by the RGI strict criteria achieve higher prediction accuracy (Equation 5) both for the ESM-1b and ProtBert-BFD embedding results (Figure 9). Compared to the “perfect” standard, the “strict” criteria effectively eliminate both negative and false-positive genes, reducing false-positive (FP) prediction probability and consequently enhancing overall prediction accuracy. From the perspective of embedding models, the ESM-1b model demonstrated superior prediction accuracy for label 2 (other), 4 (quinolone), and 12 (sulfonamide), whereas ProtBert-BFD achieved higher precision for label 3 (tetracycline) and 5 (aminoglycoside). Notably, both models attained 100% accuracy in predicting label 7 (β-lactam) and 11 (rifampin). When combining these two models' results, the prediction accuracy can theoretically exceed 90%, with peak performance reaching 100% for specific antibiotic classes. Overall, our model demonstrates exceptional performance in predicting specific resistance phenotypes (e.g., labels 5 (aminoglycoside), 7 (β-lactam), and 11 (rifampin)). However, prediction accuracy remains suboptimal for smaller datasets, particularly glycopeptide (label 9), chloramphenicol (label 10), and polymyxin (label 14) resistance categories, indicating areas for future improvement.