A novel transformer-based stacking ensemble method with multi-model integration for cancer classification

Computing infrastructure

The experiments are based on a high-performance platform and are submitted by Slurm. The specific configurations are, OS: CentOS Linux release 7.5.1804 (Core), Memory: 256 GB, Number of cores: 51.

The workflow

Base models (such as Random Forest, eXtreme Gradient Boosting (XGBoost), etc.) predict data and generate probability values. These probability values are fed into the Transformer, whose attention mechanism learns the relationships between the prediction results of different models. After multiple layers of processing, the final output is an optimized prediction result (Fig. 1). For example, if a particular base model performs exceptionally well on a specific type of sample, the Transformer will learn to assign a higher weight to that model’s prediction results when processing similar samples.

Data preprocessing

Although the use of gene expression data from the RNA-seq technique greatly facilitates the classification of cancers, it has its limitations due to the small sample size of a particular cancer type with a large number of genes (curse of dimensionality) (Blagus & Lusa, 2010; Yang & Naiman, 2014). In addition, the samples also contain genes that lack information which degrade the classification performance (García-Díaz et al., 2020; Vanitha, Devaraj & Venkatesulu, 2015), Therefore, before performing the analysis, it is necessary to preprocess the gene expression data.

To identify genes that were significantly differentially expressed in different cancer types, we performed analysis of variance (ANOVA) for each gene. After correction for the false discovery rate (FDR), genes with p-values <0.05 were considered differentially expressed and were selected for further analysis.

Data were preprocessed by removing non-informative columns such as sample identifiers and encoding categorical cancer type labels into numeric format using LabelEncoder. The matrix of Spearman correlation coefficients between samples was calculated by using the array-array strength correlation (AASI) method (Fig. 2), after which the average correlation coefficient between each sample and other samples was calculated, a threshold was set (the mean of the average correlation coefficients minus two times the standard deviation) from which potential outlier samples could be identified, and finally, samples of the filtered genes were normalized. Genes with low correlation were filtered out using the calculated threshold, returning genes with sample mean intensities greater than 0.809071442511119 (the mean of the mean correlation coefficient minus two times the standard deviation). In feature selection, a triple filtering mechanism was employed: variance filtering to remove genes with low variability, mutual information screening to retain genes associated with cancer, and correlation filtering to eliminate redundant genes (correlation coefficients > 0.96), ultimately retaining approximately 77% of the key genes. After preprocessing, the risk of dimensionality disaster is effectively reduced, thereby improving the generalization ability of the model.

Base model optimization

We evaluated multiple machine learning models, including RF (Breiman, 2001), Gradient Boosting (GBT) (Friedman, 2001), XGBoost (Chen & Guestrin, 2016), and support vector machine (SVM) (Shmilovici, 2023). Each model underwent extensive hyperparameter tuning using GridSearchCV to determine the optimal parameters. The models were trained and validated using an 80-20 training-test split, with the training data further split into training and validation sets (80-20) to prevent overfitting.

Our framework is based on four core models, each of which has been optimized for hyperparameters using parallel grid search combined with 10-fold cross-validation (Table S2). The optimization process utilizes ProcessPoolExecutor to achieve efficient computation across multiple CPU cores.

Cross-validation (Bates, Hastie & Tibshirani, 2024), also known as rotation estimation, is a model validation technique used to assess the quality of statistical analysis and results, with the goal of enabling the model to generalize to an independent test set (Fig. 3). During model deployment, differential validation is critical for evaluating the predictive accuracy of the model in real-world applications. During cross-validation, the model is typically trained using a dataset of known type.

In contrast, unknown-type data is used to test the model. This helps describe the features of the dataset used to test the model through the validation set during the training phase. There are two types of cross-validation: exhaustive cross-validation, including leave-one-out cross-validation (Rayo & Hoteit, 2012) and p-cross-validation (Liu, 2019); and non-exhaustive cross-validation, including K-fold cross-validation and repeated random subsampling cross-validation.

When using K-fold cross-validation, the entire training dataset can be divided into k disjoint folds, where each fold is a subset of the data used as the test set in one iteration, while the remaining k-1 folds are used as the training set. In each iteration, the model is trained on k-1 folds and tested on the remaining fold. The performance metrics for that iteration are recorded, and the entire cycle is repeated k times. After all iterations are completed, the average performance metrics across all iterations are calculated, and the optimal parameter combination is selected to obtain the optimized model.

Transformer-based meta-learning architecture

We propose an innovative stacked ensemble architecture based on Transformer meta-learning, which can dynamically integrate the prediction results of multiple base classifiers. Our method extends traditional stacked ensemble methods by replacing traditional simple meta-learners with attention-based deep neural networks, which can learn complex relationships between the prediction results of base models. Our ensemble foundation consists of four distinct base classifiers, selected based on their complementary strengths in different aspects of classification tasks. For example, the Random Forest Classifier (RFC) was chosen for its ability to capture nonlinear feature interactions and its robustness to outliers through bootstrap aggregation. For the meta-learners’ input, a clever transformation is applied: the prediction probabilities from all base models are reshaped into a sequence format. For example, when using four base models in a three-class classification problem, the input for each sample is a sequence of length 12 (4 models × 3 probabilities per category). The purpose of this operation is to make the data compatible with the Transformer’s input format, thereby leveraging the self-attention mechanism to learn relationships between prediction values. The meta-learners adopt a complex Transformer architecture, including: (1) Dual multi-head attention layers (eight heads per layer). (2) Layer normalization and Dropout for regularization. (3) Dense layers with gradually decreasing Dropout rates. (4) A Softmax output layer for final classification.

Input layer and feature representation

The meta-learner receives probabilistic predictions from all base models as input. For a classification problem with K classes and M base models, this creates an input matrix of dimension N × (K × M), where N is the number of samples. We reshape this input into a sequence format of dimension (N × 1 × (K × M)) for Transformer processing. This reshaping operation converts the flat probability vector into a format that can be processed by a self-attentive mechanism that treats each model’s prediction as a sequence element.

Multiple attention mechanisms

Two consecutive multi-head attention layers are used, each containing eight attention heads. Each attention head can be imagined as a specific angle of analytical perspective; the attention heads in the first layer may learn that some heads may focus on relationships between similar models (e.g., Random Forest and XGBoost), others may focus on relationships between complementary models, and still others may focus on identifying certain specific types of predictive patterns; the attention heads in the second layer build on the first layer with a more higher level of integration, may learn the reliability of different model combinations in different scenarios, identify the predictive advantages of certain models on specific categories, and discover complex predictive patterns across models.

Unlike traditional voting or simple weighting methods, this framework enables dynamically assign weights to each prediction. For example, if a prediction from a random forest is found to be particularly reliable in a certain situation, the system will automatically increase its influence. With deep neural networks, the system can learn complex integration methods that go far beyond simple weighting. It may discover patterns such as “pay special attention when both Model A and Model B are highly confident, but Model C holds the opposite opinion”. Instead of using fixed rules, the attention mechanism allows the system to adapt the integration strategy to the specifics of the current prediction. The architecture makes it easy to integrate new base models by simply adding the predictions of the new model to the input sequence.

Performance evaluation

To comprehensively assess the classification performance of our machine learning models, we employed multiple evaluation metrics that capture different aspects of model performance. The selection of these metrics is crucial for multi-class cancer classification, as it provides a holistic view of model effectiveness across different cancer types.

• 1.

  Accuracy

Accuracy represents the proportion of correctly classified samples among all samples and serves as the primary metric for overall model performance:

(1) $$Accuracy={\displaystyle \frac{TP+TN}{TP+TN+FP+FN}}$$where: TP = true positives, TN = true negatives, FP = false positives, FN = false negatives.

For multi-class classification, accuracy is calculated as:

(2) $$Accuracy={\displaystyle \frac{{\sum }_{i-1}^{n}Correctpredictionsforclassi}{Totalnumberofsamples}.}$$

• 2.

  Confusion matrix

The confusion matrix provides a detailed overview of the classification performance for each cancer type. For our multi-class problem (STAD, LUNG, BRCA, CESC, LIHC), the confusion matrix is a 5 × 5 table, where rows represent actual cancer types, columns represent predicted cancer types, diagonal elements represent correct classifications, and non-diagonal elements represent misclassifications. The confusion matrix can be used to identify which cancer types are classified most accurately, common misclassification patterns between different cancer types, and performance differences among classifications.

• 3.

  Precision, recall, and F1-score

For each cancer type $i$, we calculated:

Precision (positive predictive value):

(3) $$precisio{n}_i={\displaystyle \frac{T{P}_i}{T{P}_i+F{P}_i}}$$

Precision measures the proportion of samples predicted as cancer type i that are actually cancer type $i$.

Recall (sensitivity or true positive rate):

(4) $$recal{l}_i={\displaystyle \frac{T{P}_i}{T{P}_i+F{N}_i}}$$

Recall measures the proportion of actual cancer type $i$ samples that are correctly identified.

F1-score (harmonic mean of precision and recall):

(5) $${\mathrm{F}1\text{-}\mathrm{s}\mathrm{c}\mathrm{o}\mathrm{r}\mathrm{e}}_{\mathrm{i}}=2\times {\displaystyle \frac{\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{c}\mathrm{i}\mathrm{s}\mathrm{i}\mathrm{o}{\mathrm{n}}_{\mathrm{i}}\times \mathrm{r}\mathrm{e}\mathrm{c}\mathrm{a}\mathrm{l}{\mathrm{l}}_{\mathrm{i}}}{\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{c}\mathrm{i}\mathrm{s}\mathrm{i}\mathrm{o}{\mathrm{n}}_{\mathrm{i}}+\mathrm{r}\mathrm{e}\mathrm{c}\mathrm{a}\mathrm{l}{\mathrm{l}}_{\mathrm{i}}}}$$

For our cancer classification problem, Transformer stacks performed best across all metrics.

Model performance

Stacking is performed and finally the accuracy of the original stacked model is obtained, after which the final probability value is obtained by five-fold cross-validation. Figure 4 shows the accuracy rates of all models (including individual models and ensemble models). By visually comparing the performance of different models, the best model can be quickly identified. The Transformer-based stacked ensemble model consistently outperforms individual base models in terms of accuracy. The Transformer stack achieves an accuracy rate of 0.997, outperforming all individual models. This model effectively combines the strengths of the base models to enhance predictive performance.

Figure 5 shows the heatmap of the confusion matrix, illustrating the comparison between the model’s predictions on the test set and the actual labels. The model’s performance across each category is detailed, identifying which categories are prone to confusion, and through an in-depth analysis of the model’s classification performance, areas for improvement are identified. It can be seen that by stacking the classification and identification of the five types of cancer, the classification performance has significantly improved. Among them, the performance of individual models was the worst, and TabNet performed worse than the Transformer stack (precision: 0.9972) in terms of metrics (precision: 0.9723). In addition, we compared traditional stacking with Transformer stacking (Fig. S1). The results showed that the Transformer-based stacking method further improved classification accuracy. A Friedman test (multiple comparison) was performed, with a p-value <0.05, indicating statistical significance. This proves the necessity of integrating predictions based on Transformer machine learning models.

Figure S2a illustrates how model training scores and cross-validation scores change with the number of training samples. This method can be used to determine whether the model suffers from overfitting or underfitting, and whether adding more training data can improve performance. It is necessary to decide whether to collect more data or adjust the model complexity. Observation of the results in Fig. S2b shows that when the batch size reaches 1,400 to 1,700, most models achieve good performance, and as the batch size increases, the standard deviation gradually decreases, indicating that the model stability has improved.

Table S3 shows the performance of each model during cross-validation. The average precision, recall and F1-score of the best model reached 0.9858, 0.9858 and 0.9858, respectively. while the classification accuracy was 98.58%. Moreover, these results after performing integration learning also show that our proposed model outperforms all individual models and machine learning results. In addition, Figs. 4 and 5 (confusion matrix) show that our proposed model has better classification performance compared to a single model.

SHAP analysis

We used SHAP analysis to identify 20 genes that are most important for cancer classification, meaning that these genes play a key role in model decision-making (Fig. S3a). SHAP (SHapley Additive exPlanations) values represent the marginal contribution of each gene to the final prediction result: the higher the value, the more significant the gene’s influence on classification decisions. SHAP values reflect how changes in gene expression drive model predictions toward specific cancer types. By interpreting the biological significance of important genes, we can better understand cancer changes (Fig. S3b). TSPY1 (SHAP: 0.0148)—Testis-specific protein Y-linked 1. This gene functions in cell proliferation and apoptosis regulation and is abnormally expressed in various tumors, particularly those of the reproductive system, making it a potential biomarker for gender-related cancers. NF525 (0.0126)—Zinc finger protein 525, this gene is a transcription regulator involved in gene expression regulation. The zinc finger protein family is closely associated with tumor suppression and carcinogenic processes. Its clinical significance lies in the fact that transcriptional regulation disorders are an important mechanism in cancer development. Additionally, SHAP analysis enhances model interpretability, enabling a transition from traditional black-box models to SHAP-explained models and transparent decision-making processes, thereby making complex machine learning decision processes understandable. We used functional classification heatmaps (Fig. S3c) and cancer association networks (Fig. S3d) to reveal how each gene influences cancer classification decisions, and used SHAP to quantify the magnitude and direction of gene contributions.

Model complementarity analysis

After analyzing the prediction patterns of different models, we found that they are significantly complementary. As shown in Fig. 5, although the models are similar in overall accuracy, the types of errors they make across the entire dataset are different. This complementarity is reflected in the low correlation between the error patterns of different underlying models, indicating that integrating their prediction results can produce better results.

Analysis of the attention model

An analysis of the attention weights of the Transformer architecture reveals interesting insights about how the model combines predictions: (1) First attention layer: the initial attention heads exhibit specialized behavior, with different heads focusing on different aspects of the base model predictions. We observe that certain heads are responsible for consistently focusing on high-confidence predictions, while others are responsible for focusing on cases where the base model is inconsistent. (2) Second attention layer: the secondary attention mechanism exhibits more context-dependent behavior, effectively weighting the fine-grained representations according to the specific characteristics of each input case. This layer shows particularly strong activation patterns for cases where the base model shows conflicting predictions.