Tab-semiSL: Tabular Data-Driven Semi-Supervised Learning to Identify Factors Associated with Immune-Related Adverse Events

1 Introduction

The high incidence and mortality rates of cancer threaten people all over the world. Although current treatment modalities include surgery, chemotherapy, radiotherapy, targeted therapy, poor prognosis remains a challenge [Haefner et al.,

∗Corresponding author

2015; Smyth et al., 2020]. While various therapies are currently available in clinical practice, immunotherapy, particularly Immune Checkpoint Inhibitors (ICIs), is regarded as one of the most promising approaches [Lee et al., 2016]. ICIs restore anti-tumor T-cell activity by blocking inhibitory pathways such as Programmed cell Death protein 1 (PD1)/Programmed cell Death Ligand 1 (PD-L1) and Cytotoxic T Lymphocyte-associated Antigen-4 (CTLA-4), thereby exerting therapeutic effects [Postow et al., 2018]. However, studies have shown that the use of ICIs can introduce autoimmune inflammation in multiple organs, collectively referred to as immune-related Adverse Events (irAEs) [Ramos-Casals et al., 2020]. These irAEs significantly affect patient prognosis and survival rates. Therefore, there is an urgent need to establish a reliable and clinically applicable predictive method for irAEs identification. Current predictive methods can be classified into two main categories: statistical analysis methods which usually applied meta-analysis, logistic regression, etc. [Jing et al., 2021; Fujisawa et al., 2017; Wang et al., 2024], and machine learning approaches, e.g., Support Vector Machines (SVM), Random Forests (RF), eXtreme Gradient Boosting (XGBoost), etc. These machine learning models are often combined with interpretability algorithms for feature importance ranking and biomarker identification [Gong et al., 2023; Zhou et al., 2022]. However, these methods are based on retrospective hospital collected data, which is often small and prone to overfitting. Although some researchers have utilized the public datasets FDA Adverse Event Reporting System (FAERS) to identify irAEs biomarkers [Bomze et al., 2019; Jing et al., 2020]. However, FAERS is a spontaneous reporting system with a large amount of missing data [Sakaeda et al., 2013], where critical clinical information such as cancer severity, laboratory and radiological findings is unavailable [Raschi et al., 2019]. Recently, [Wang et al., 2023] achieved significant improvements using SSL with a small amount of labeled CT data and a large amount of unlabeled data. The success of this study indicates that SSL has potential in medicine.

While SSL has achieved notable success with imaging data, adapting it to tabular data (e.g., lab values, gene expression) in irAEs research faces challenges. Most SSL methods involve the process of data perturbation, such as image rotation [Sohn et al., 2020] or two images and their convex combination(s) [Zhang et al., 2017] and these perturbation methods are not suitable for application to tabular data [Yoon et al., 2020]. VIME [Yoon et al., 2020] is the first SSL method specifically designed for tabular data. These SSL methods applied to tabular data mainly have two problems: 1) The data perturbation is not effective enough. 2) Lack of sample selection mechanisms, where all unlabeled samples are used throughout training regardless of prediction confidence, allowing noisy and low-confidence unlabeled samples to degrade model performance.

In addition, several studies have extended self-supervised learning to tabular data (SubTab [Ucar et al., 2021] and SCARF [Bahri et al., 2021]). Unlike SSL, self-supervised learning first leverages unlabeled data for pre-training, followed by fine-tuning exclusively on labeled data. This approach, however, has some limitations, e.g., the absence of explicit supervision limits the model's ability to obtain task-specific representations that are optimal for downstream tasks. When only a limited number of labeled data is available for downstream task fine-tuning, the resulting improvements in model performance tend to be marginal [Li et al., 2021].

In this study, we retrospectively collected the clinical laboratory tabular data of 237 cancer patients treated with ICIs and 4,817 unlabeled cancer patients without the irAEs label, as illustrated in Figure 1. We introduced Tab-semiSL, a novel SSL method specifically designed for tabular data. Tab-semiSL can be classified into two stages: In the pretraining stage, to enhance the effectiveness of perturbation for unlabeled samples, we used Variational Auto-Encoder (VAE) [Kingma et al., 2013] combined with marginal distribution sampling to perturb unlabeled samples, then using original unlabeled samples and perturbed samples jointly trained the auto-encoder with masked reconstruction and contrastive learning. In the SSL stage, labeled samples undergo standard supervised learning. For unlabeled samples, we enforce the results of the perturbed samples to be consistent with the pseudo-labels, only the class with the maximum predicted probability was selected as the true labels [Lee and others, 2013]. Considering that the reliability of pseudo-labels generated for unlabeled samples fluctuates during the SSL training process, we set a specific threshold for each class and only regarded samples whose maximum confidence belongs to a class and is higher than the threshold of class-specific as pseudo-labels. Then, we proposed a self-adaptive dynamic threshold update mechanism to update the threshold according to different stages of SSL training. Once the model was obtained, we used Area Under the Curve (AUC), accuracy, recall and f1 score metrics for evaluation and SHapley Additive exPlanations (SHAP) values [Lundberg and Lee, 2017] for interpretability analysis and feature ranking, identifying factors associated with irAEs. Our contributions can be summarized as follows

• We designed Tab-semiSL, a novel semi-supervised learning method specifically for clinical tabular data processing, which was benchmarked against conventional classification models, SSL and self-supervised learning methods, almost achieving the best performance.
• We introduce a novel perturbation framework for tabular data that combines a VAE with marginal distribution sampling. Furthermore, masked reconstruction and contrastive learning are jointly incorporated in the pretraining stage to improve representation quality.
• We set a specific threshold for each class and proposed a self-adaptive and dynamic threshold update mechanism to update the threshold. Only unlabeled samples whose maximum confidence belongs to a class and is higher than the threshold of class-specific are regarded as pseudo-labels during different stages of SSL training.
• We collected and validated on real-world ICIs-treated patient data, and identified interpretable relevant factors of irAEs.

2 RelatedWork

SSL aims to enhance model performance by using a small amount of labeled samples and a large amount of unlabeled samples. Compared with self-supervised learning, SSL combines labeled samples and unlabeled samples during the training process, which can effectively capture the relevant information of both [Yang et al., 2022]. SSL can be roughly divided into three methods: The first is consistency regularization [Samuli and Timo, 2017; Tarvainen and Valpola, 2017; Xie et al., 2020]. The second is self-training [Tai et al., 2021; Rizve et al., 2021]. The third combines both approaches. Representative methods include FixMatch [Sohn et al., 2020], FlexMatch [Zhang et al., 2021], and FreeMatch [Wang et al., 2022]. Most mainstream methods rely on strong and weak perturbations to the same sample: pseudo-labels are generated from the weakly perturbed version, and the strongly perturbed version is used to compute the loss based on those pseudo-labels. Initially, the majority of SSL research was concentrated in the domain of image data. VIME is the first method to apply SSL to tabular data. VIME constructs 'pretext tasks' to train the unlabeled samples, simultaneously generating perturbed samples by sampling from empirical marginal distributions. Then, consistency regularization is used to encourage the output consistency of different perturbed samples.

3 Methodology

3.1 Preliminaries

We defined an SSL scenario for tabular datasets with both labeled and unlabeled samples.Let X = {(xi, yi)}Ni=1 represents a label dataset of size N, where xi ∈ RK and yi ∈ {1, · · · ,C} denote labeled training samples with feature dimension of K and labels with C categories, respectively. U = {uj}N+M j=N+1 denotes an unlabeled dataset of size M, where uj ∈ RK . Generally, N << M. Let Y = {y1, · · · , yN}, the task objective of SSL is to train mapping relationships X → Y using both Dlab and Dunlab simultaneously.

3.2 Pre-Training Stage of Tab-SemiSL

Samples Generation using VAE

In the pre-training phase (see Figure 2), to achieve data perturbation, we first train a VAE using unlabeled samples, and then samples were generated by probabilistically modeling the latent space and sampling from it. Specifically, the encoder of VAE ψϕ encodes the raw samples U into posterior distributions qϕ(z|U) of the latent representation z and utilizes Kullback-Leibler(KL) divergence to construct the loss function LKL = KL[qϕ(z | U)∥p(z)] to constrain the posterior distribution, encouraging it to approximate a standard normal distribution, here ϕ denotes the encoder parameters of VAE, p(z) = N(0, I). Latent representations z = µ(U) + σ(U) ⊙ ε are sampled via the reparameterization trick, here µ(U) and σ(U) represent the mean and standard deviation of U, respectively, ε ∼ N(0, I), ⊙ denotes element-wise product. Moreover, the VAE also encourages the samples generated by the decoder to be similar to the original samples, with the loss function formulated as:

where pθ( ¯U) represents the output distribution of the decoder, represents the reconstructed samples, θ denotes the VAE decoder parameter. Therefore, the loss function of VAE can be expressed as:

Thus, the generated samples was represented as ˆU = {ˆuj}N+M j=N+1 , where ˆuj can be expressed as:

where φθ(·) represents the decoder of VAE with sigmoid activation, µj, log σ2 j = ψϕ(uj). See Supplementary Material Section 2 for more details.

Samples Perturbation

Samples perturbation were achieved by stochastically replacing feature values of generated samples ˆU through sampling from their marginal distributions. More precisely, we generated a mask matrix M ∈ {0, 1}M×K , here Mj,k ∼ Bernoulli(pm), pm ∈ [0, 1] denotes a hyperparameter for setting the masked rate. IfMj,k = 1 indicates that the value is retained, andMj,k = 0 indicates that the value is masked. Then for the perturbed samples ˜U = {˜uj}N+M j=N+1 , the j - th row and the d - th column can be expressed as:

Training the Model of Pre-Training Auto-Encoder

The objective is to obtain a pre-training encoder for feature representation during the pre-training stage. We define the pre-training encoder and decoder as fδ(·) and f∆(·), respectively. The pre-training auto-encoder used the original unlabeled samples U and their perturbed samples ˜U as inputs, integrating the collaborative training strategy of masked reconstruction and contrastive learning, where masked reconstruction needs to complete two 'pretext tasks' - mask matrix reconstruction and feature reconstruction. The loss in the pretraining stage is given by the sum of mask matrix reconstruction loss Lm, feature reconstruction loss Lf and contrastive loss Lcon:

where Lf and Lm both computed via mean squared error and binary cross-entropy, respectively, α and β are hyperparameters. When the feature is categorical, the mean squared error is replaced by the cross entropy loss. Further details can be found in Supplementary Material Section 3.

For an unlabeled sample uj ∈ RK , we treat its perturbed sample ˜uj as a positive pair and all other samples ut and ˜ut, t ̸= j as negative pairs to construct the InfoNCE contrastive loss [Oord et al., 2018]. The objective is to encourage the latent representations of positive pairs to be close to each other, while pushing those of negative pairs farther apart. Thus, contrastive loss Lcon is defined as follows:

where zj = g(fδ(uj)), ˜ zj = g(fδ(˜uj)), zt = g(fδ(ut)), sim (zj , ˜ zj) = zT j ˜ zj ∥zj∥2 ·∥˜ zj ∥2 , τ denotes temperature parameter, g(·) is a projection head of two layers of neural networks, and activation functions are applied after all network layers.

3.3 SSL Stage of Tab-SemiSL

Training Predictor Model

In the SSL stage (see Figure 3), labeled samples X, unlabeled samples U and unlabeled perturbed samples ˜U will be encoded by the encoder fδ(·) which retained from the pretraining stage and then input into a predictor P(·) composed of two neural networks for training, where labeled samples undergo standard supervised learning and the unlabeled samples are trained via unsupervised learning. The objective loss in the semi-supervised stage is:

where γ is the hyperparameter for balancing supervised loss and unsupervised loss. The supervised loss Ls for the multiclass case is given as follows:

where ˆ zi = fδ(xi), N represents the number of label samples, and C represents the number of categories. When the task is binary classification, the binary cross entropy loss will be adopted or the mean square error loss will be used when it is a regression task.

Self-Adaptive Dynamic Threshold Update Mechanism

The objective of unsupervised loss is to keep the pseudolabels of unlabeled samples U output consistent with the output of its perturbed samples ˜U. However, in the early stage of training, the poor performance of the model led to low-quality pseudo-labels generated from the low-confidence unlabeled samples. We set a specific threshold for each class to select samples with the maximum confidence higher than threshold and have proposed a self-adaptive dynamic threshold update mechanism to update the specific threshold for each class. Unlike FreeMatch, which estimates thresholds from the global prediction distribution with class fairness adjustment, our method computes class-specific thresholds based on the confidence statistics of progressively selected highconfidence samples. In addition, we introduce an iterationdependent sample selection strategy and EMA-based threshold smoothing to stabilize the pseudo-label generation process. Our method for dynamically calculating the threshold is as follows: Based on the number of iterations, selected the high-confidence samples, and simultaneously calculate the average value of the class corresponding to the maximum confidence of selected samples as the threshold for this class (see Supplementary Material Algorithm 1). Specifically, for unlabeled samples U = {uj}N+M j=N+1 , its output confidence is denoted as pj = P(fδ(uj)) and pj = {pj,1, · · · , pj,C}, pj,c ≥ 0, ∀c ∈ {1, · · · ,C}, ∑Cc=1 pj,c = 1. The sample selected at the t - th iteration can be represented as:

where T represents the total number of iterations sort denotes ordering the elements in the set in descending order based on the maximum predicted confidence of each sample, such that pmax N+1 ≥ · · · ≥ pmax N+M , pmax N+1 =max{pN+1,1, · · · , pN+1,C} represents the maximum predicted confidence for sample uN+1. The role of F(·) is to select the top-t/T samples from the sorted set. Then, calculate the average value based on the class corresponding to the maximum confidence value of each sample in Ds, which is regarded as the threshold of each class. Meanwhile, the threshold for each iteration is maintained by a dynamically updated Exponential Moving Average (EMA). Let the initial thresholds all be η0 =∞, and the threshold for the k - th class in the t - th iteration (t > 0) can be expressed as:

When ηt-1(k) ̸= ∞, we employ EMA to maintain the threshold, which can be expressed as:

where 1(·) denotes Indication function, ps = P(fδ(us)), us ∈ Ds, λ denotes the momentum decay of EMA, and λ = λ0 + (0.99 - λ0) tT , λ0 ∈ (0, 1). Therefore, the unsupervised loss in t - th iterations can be expressed as:

where H denotes cross entropy loss, ˆ yj = argmax c (pj,c) is a pseudo-label, ˜ pj = P(fδ(˜uj)).

4 Experiments and Results 4.1 Experimental Settings

Datasets

We used two publicly available datasets to evaluate the effectiveness of our method: the Adult Income dataset for binary classification, and the MNIST dataset for multi-class classification. For the public datasets, we select a small portion of the data as labeled samples, while discarding the labels of the remaining samples to treat them as unlabeled samples. Meanwhile, to determine factors associated with irAEs, this study retrospectively collected clinical laboratory test data from 237 cancer patients who had received ICIs treatment and 4,817 cancer patients without irAEs labels from partner hospitals, including 87 characteristics. For cancer patients who had received ICIs treatment (labeled samples), we collect the most recent examination values of the patients before they use ICIs. According to the medical records, the occurrence of irAEs in patients was assessed. Patients (N =58 or 24%) were identified as having experienced irAEs, whereas patients (N=179 or 76%) did not experience irAEs. For cancer patients without irAEs labels (unlabeled samples), we collect the latest examination results of the patients. The advantage of adding a large amount of unlabeled samples lies in alleviating the model overfitting and poor generalization caused by the limited number of labeled samples. Detailed clinical laboratory test indicators can be found in Supplementary Material Section 4. For the collected samples, missing values are filled using the KNN imputer from the scikit-learn package. To address label imbalance, we apply the Synthetic Minority Oversampling Technique (SMOTE) [Chawla et al., 2002] to the labeled data. All datasets are then split into training, validation, and test sets in a 7:1:2 ratio.

Implementation Details

Tab-semiSL consists of the following modules: 1) VAE module. 2) Pre-training auto-encoder module. 3) Predictor module. Our approach was developed in Python utilizing the PyTorch framework, with computations on NVIDIA RTX 4080 GPU and AMD Ryzen 9 7945HX. The epochs for training the VAE module and the pre-training auto-encoder module are 50 rounds and 30 rounds respectively. In the SSL stage, the predictor module was trained for a maximum of 5,000 iterations, with an initial warm-up phase of 50 iterations using labeled data. All learning rates in the pre-training stage are lr = 1 × 10-3 . the learning rate of the predictor module trained using the Adult income dataset is lr = 1 × 10-5 , and the learning rates trained employing both the MNIST dataset and the private clinical dataset are lr = 1 × 10-3 .

Comparative Methods

The contrastive methods used in this study can be divided into two categories: 1) Supervised learning-based methods, including XGBoost [Chen and Guestrin, 2016], CatBoost [Prokhorenkova et al., 2018], and Lightgbm [Ke et al., 2017]. For these methods, we only used labeled samples for training. 2) Semi-supervised and self-supervised learning-based methods, including VIME, SubTab, SCARF, FlexMatch and FreeMatch. For the FlexMatch and FreeMatch methods, we conducted experiments using the same data perturbation strategy adopted in this study. For the self-supervised methods, we first performed pre-training using unlabeled samples, followed by fine-tuning only with labeled samples. All baseline methods were implemented in strict accordance with the experimental settings described in their original papers or by directly employing their official open-source implementations.

Quantitative Results

We evaluate different methods on MNIST and Adult income datasets, reporting the mean ± standard deviation over 10 runs with different random seeds (see Table 1). 'Label Amount'

refers to the number of labeled samples. Accuracy is used as the metric for MNIST, and AUC for Adult Income. On the MNIST dataset, our method almost achieves the best scores under different numbers of labels. For the Adult Income dataset, Catboost performs best when the number of labeled samples exceeds 256. This phenomenon may be attributed to the relatively simple classification task of this dataset, which enables the model to quickly reach performance saturation even with a limited number of labeled samples. However, as the number of labeled samples decreases, the performance gap between our model and other baseline methods gradually widens, highlighting the effectiveness of our SSL method.

4.2 Ablation Study

To evaluate the contribution of each module to the overall performance, we conducted ablation experiments on the MNIST dataset under the same experimental conditions. The key components of Tab-semiSL include the pre-training stage and the SSL stage. We have set up several variants of TabsemiSL.

Ablation Experiments in the Pre-Training Stage

1. We replaced the VAE-generated samples with samples
2. drawn from a Gaussian distribution, denoted as w/o VAE .
3. We removed pre-training stage, and the raw data are used directly without feature representation learning, denoted as w/o pre-training.

Ablation Experiments in the SSL Stage

1. To verify the effectiveness of our proposed self-adaptive dynamic threshold update mechanism, we conducted experiments using different unlabeled sample selection strategies.

Specifically, we employed two approaches: using a fixed value (0.95) as the experimental threshold or not setting a threshold at all, denoted as threshold 0.95 or w/o threshold. 2) We removed the SSL stage and directly fed the feature representations learned during pre-training into a multilayer perceptron to obtain the final predictions, denoted as w/o SSL.

We conducted ablation experiments with 500 and 1000 labeled samples, and the results are shown in Figure 4.

4.3 Parameter Sensitivity Analysis

We conduct a parameter sensitivity analysis using the MNIST dataset, where 6,000 labeled samples are selected. The purpose is to analysis how different parameter settings influence the model's performance. The hyperparameter ranges are respectively pm ∈ {0.1, . . . , 0.9}, α ∈ {0.1, 0.5, 0.7, 1, 2, 5}, β ∈ {0.1, 0.5, 0.7, 1, 2, 5} and γ ∈ {0.1, 0.5, 0.7, 1, 2, 5}. The analysis results are presented in Supplementary Material Figure 2. Ultimately, pm = 0.5, α = 2.0, β = 0.1 and γ = 0.1.

4.4 Determine the Factors Associated with irAEs and Analysis

Determine the Factors

We classified the collected clinical laboratory test data using Tab-semiSL and other contrastive methods respectively. Tab-semiSL still achieved the best results (see Table 2). We then used SHAP to rank the feature importance of the pretraining encoder-predictor combined model after training (see Figure 5). SHAP is a game-theoretic interpretability algorithm that determines the marginal contribution of each feature by computing its shapley value [Shapley and others, 1953] when added to a combination of other features. This value fairly attributes the contribution of each feature to the model output. We ranked the features according to shapley values and visualized the top-10 features contributing most to the model output in Supplementary Material Section 5. Ultimately, the 10 important factors associated with irAEs that we identified are: white blood cell count (WBC), the percentage of lymphocytes in the total number of WBC (Lymph), blood platelet count (PLT), basophils as a percentage of total WBC (Baso), free thyroxine (FT4), red blood cell count (RBC), hemoglobin (Hb), potassium ion (K), alpha fetoprotein (AFP), squamous cell carcinoma antigen (SCC).

Analysis

Lymphocytes, as a subset of leukocytes, mainly include T cells, B cells, and Natural Killer (NK) cells. They play a crucial role in adaptive immunity and are also involved in the systemic inflammatory response [LaRosa and Orange, 2008; Buonacera et al., 2022]. Platelets have been confirmed to be closely associated with inflammation and immune responses and can interact with various cells involved in innate and adaptive immunity. Through the release of chemokines, platelets guide lymphocytes, neutrophils, and monocytes to migrate to sites of inflammation [Sonmez and Sonmez, 2017; Semple et al., 2011]. Additionally, intercellular communication exists between platelets and different types of leukocytes, influencing inflammatory and immune responses [Totani and Evangelista, 2010]. Basophils are considered to induce allergic symptoms by releasing soluble pro-inflammatory mediators and have subsequently been found to be involved in inflammatory and immune responses [Falcone et al., 2011; Harvima et al., 2014]. These lines of evidence suggest that the identified factors associated with irAEs have the potential to serve as predictive biomarkers for irAEs.

5 Conclusion and FutureWork

In conclusion, we proposed a novel semi-supervised learning framework entitled 'Tab-semiSL' for clinical tabular data, and demonstrate its ability to predict irAEs by jointly leveraging 237 ICIs-treated patients with labeled irAEs and 4,817 unlabeled oncology cases. The results of Tab-semiSL outperforms conventional supervised, self-supervised, and semisupervised methods for almost all metrics. Future work will focus on expanding the labeled dataset with additional realworld clinical samples and integrating multimodal inputs, including imaging, genomics, to enhance robustness and uncover further predictive factors for irAEs.

Acknowledgments

This work was supported by the National Natural Science Foundation of China (Grant No. 62462022, Grant No. 62572157).