An Explainable State-Space Transformer Framework for Predicting Clinical Deterioration in Non-Pulmonary Tuberculosis Patients

2.1 Tuberculosis, extrapulmonary disease, and predictive challenges

TB remains a global health burden, with EPTB presenting unique diagnostic and prognostic challenges [2]. EPTB’s clinical diversity, complex host–pathogen dynamics, and increased prevalence in populations with comorbidities such as HIV have led to calls for more nuanced, data-driven approaches for prognosis and management [3-5].

Conventional statistical methods, such as Logistic Regression and Cox proportional hazards models—have provided valuable insights into risk factors and treatment outcomes in both pulmonary and extrapulmonary TB [6]. However, these methods are often limited by their assumptions of linearity and inability to fully capture high-dimensional, temporal, and multimodal health data [7-9].

2.2 Machine learning and artificial intelligence in tuberculosis prediction

2.2.1 Early machine learning efforts and statistical-machine learning comparisons

The first machine learning applications in the research of TB have put the emphasis on classification while also factoring in risk identification. This was done through the use of Random Forest, which helped sort and distinguish the relevant patterns with a kind of probabilistic way, so to speak, machines of support vector (SVM), let alone Logistic Regression [5, 6, 10-14] which are all considered tradition algorithms. Depending on the studied conducted, it has been shown that classical statistical models could be outperformed via the machine learning models, which could be done at the prediction of treating the outcomes, particularly in the availability of large, real-world datasets, i.e. [6, 8, 10, 12, 15].

Using the multi-country data to predict the failure of treatment, Asad et al. [10] have implemented machine learning-based framework. They have discovered that the models that offered superior accuracy as opposed to regression were the XGBoost ones. Fayaz et al. [15] also came out with almost the same outcomes. They found that the baseline classifiers for predicting treatment success in Pakistan are less effective than XGBoost, even compared to Random Forests. This kind of behavior is noticed as well when the data is imbalanced, and it is being tackled in that setup.

These trends have also been confirmed by systemic reviews. Peetluk et al. [8] have highlighted superiority concerning the machine learning approaches over the methods especially the traditional ones at clinical datasets when they are heterogeneous and high-dimensional.

2.2.2 Imaging, biomarkers, and multi-modal fusion

To mash together data that are clinical lab, and also imaging, hybrid machine learning models have been used. This is mostly done to boost detection accuracy and to predict outcomes for TB [2, 14, 16-19]. For the early detection of TB from the imagines of X-ray, Fati et al. [16], for instance, have employed deep learning accompanied by the fusion feature to perform this. Chukwunweike et al. [17], on the other hand, have integrated AI, Convolutional Neural Network (CNN) accompanied by MATrix LABoratory (MATLAB). They have done this on purpose of modeling prediction, and the demonstration of performance that is improved.

Biochemical markers have been introduced into the frameworks of machine learning by Wang et al. [18]. They have enhanced the accuracy of prediction for the outcomes of TB treatment. Studies that used ensemble have also echoed the multi-model fusion paradigm. This has been performed with the hybrid models [19, 20].

2.3 Temporal and sequence modeling in tuberculosis

2.3.1 Deep learning for longitudinal and electronic health record data

Advances in deep learning, particularly with Long Short-Term Memory (LSTM) and transformer architectures, have significantly enhanced the ability to model complex temporal dependencies in clinical data [21-28]. Zhang and Li [22] introduced the Chrono Former model for structured event modeling, showing that attention-based architectures outperform classical time series approaches for TB and other chronic diseases.

The significance of sequence data has been highlighted by studies such as Hosu et al. [5] and Chen et al. [13] which demonstrated that LSTM and deep neural networks could model the outcomes of TB better among HIV-infected and drug-resistant patients.

In recent years, the use of transformers and explainable AI (XAI) methods has proven to be highly effective at forecasting disease trajectories and hospital outcomes, even in the complex and imbalanced context of TB. Recent studies by Yasin et al. [26] and Renc et al. [29] have successfully demonstrated the capabilities of transformers and XAI techniques for predicting health trajectories and hospital outcomes, even in the complex and imbalanced realm of TB.

2.3.2 Hybrid and interpretable models

New hybrid modeling frameworks are now appearing that merge the classical statistical modeling tools (e.g., Kalman filters, state space models) with the power of the transformers and deep learning [22-24, 30-34]. Nguyen et al. [23] offered clinically-inspired multi-agent transformers for disease trajectory forecasting, and Siebra et al. [30] offered a systematic review of the transformer-based models to emphasize their advantages in longitudinal clinical data.

Finally, hybrid models help to overcome some major shortcomings of deep learning in clinical practice: interpretability and robustness [1, 20, 31, 32]. The rest of the literature further highlights the promise of representation learning with transformer architectures towards unified, multimodal clinical diagnostics by Zhou et al. [28] and Yang et al. [27].

2.4 Socioeconomic, geographical, and global perspectives

The impact of socioeconomic, environmental and geographical factors on TB progression and management has also been addressed in many studies [1, 9, 33, 35-38]. Tandirogang et al. [37] examined spatial interaction between EPTB and PTB in Indonesia and Rahman and Shiddik [1] detected the global determinants of TB risk using XAI. These insights underscore the vital role of region-specific AI models designed to predict TB while taking into account the local context.

The continuous development and potential of AI in TB, ranging from diagnosis and drug development to a surveillance system and public health decision support, are highlighted in the work of Mohamed et al. [35] and Memon et al. [36].

2.5 Summary of key advances

In sum, recent research (see Table 1) indicates that:

• Machine learning models generally outperform classical statistics for TB outcome prediction, especially with large, multi-source datasets [6, 10, 12, 15].
• Deep learning and transformer architectures are state-of-the-art for sequence modeling, risk stratification, and electronic health record (EHR) integration [21-23, 25, 27, 28, 30].
• Hybrid and explainable models represent the future direction—balancing predictive performance with clinical interpretability [1, 20, 31, 32].

These advances have set the stage for the current work, which integrates state-space model (SSM) with transformers to address the unique challenges of predicting non-pulmonary TB patient status in real-world clinical data.

Table 1. Comparison of 44 key studies in non-pulmonary or the extra-pulmonary tuberculosis (NPTB/EPTB) outcome prediction and predictive modeling

4.1 Study overview

This study develops a hybrid AI framework for predicting clinical deterioration in NPTB using a real-world, time-series, and SMOTE-balanced dataset. The methodology consists of data preprocessing, advanced feature engineering, model development (hybrid state-space + transformer), model training and optimization, baseline comparison, explainability, and rigorous evaluation.

Input → Preprocessing (cleaning, feature engineering, SMOTE) → SSM branch + Transformer branch → Fusion → Classification → Output & Explainability (SHapley Additive exPlanations (SHAP), Attention, Error Analysis).

4.2 Data handling and augmentation

•Data integrity checks: All records were checked for consistency, duplicates, and outlier timepoints using statistical summaries and visualization.

•Sequence alignment: Patient histories were structured as time-indexed sequences. When patient IDs were available, each sequence was aligned to its baseline, with missing windows forward/backward filled.

•Data augmentation: To improve model robustness, slight jitter/noise was added to temporal features in the training set (per best practice in time-series deep learning).

•SMOTE balancing: Borderline-SMOTE was applied exclusively to the training data within each cross-validation fold, preventing information leakage into validation or test sets.

4.3 Feature engineering

•Temporal derivatives: For each patient sequence, the following were computed:

-Change in status between timepoints (Δstatus)

-Trajectory slope (linear regression over sliding windows)

-Early trend divergence (difference between first two intervals)

-Volatility index (standard deviation within window)

•This is a binary encoding of “worsened” outcome – multi-class can follow for future extensions.

•Time encoding: Time since the first record, and the “Time of status” that has been normalized.

•Feature selection: Only top predictors retained using mutual information and tree-based importance.

First Difference

$\Delta$ status $_t=$ status $_t-$ status $_{\{t-1\}}$

Trajectory Slope

Slope $=\frac{\left\{n \sum t_i y_i-\left(\sum t_i\right)\left(\sum y_i\right)\right\}}{n \sum t_i^2-\left(\sum t_i\right)^2}$

Volatility Index

Volatility $=\sqrt{\left\{\frac{\{1\}}{n-1} \sum_{\{i=1\}}^{\{n\}}(x i-\bar{x})^2\right\}}$

Min-Max Normalization

$x_{\{\text {norm }\}}=\frac{\left\{x-x_{\{\min \}}\right\}}{\left\{x_{\{\max \}}-x_{\{\min \}}\right.}$

4.4 Baseline model development

To provide clinical and methodological context, baseline models were developed:

•Logistic Regression: Using time, status, and engineered features.

•Random Forest and XGBoost: For nonlinear comparison and feature importance assessment.

•LSTM (Recurrent Neural Network): To test classic sequential deep learning approaches.

Their results were used for direct benchmarking against the hybrid model.

4.5 Hybrid model architecture

4.5.1 State-space model

•Formulation: Latent health state evolution is modeled as a Markov process with Gaussian noise, learned via Expectation-Maximization.

•Inference: The Kalman filter recursively estimates latent states using observed status and temporal features.

•Interpretability: Latent state trajectories provide interpretable, smoothed disease courses.

4.5.2 Transformer component

•Structure: Includes multi-head self-attention, feed-forward layers, and positional encoding.

•Clinical adaptation: engineered features as inputs, and deep embedding for each time window as outputs.

•Regularization: Dropout, batch normalization and early stopping (based on validation loss and area under the curve (AUC)).

4.5.3 Model fusion

•Rewrite the task by yourself: Concatenate the SSM filtered states and the transformer embedding through dense layers (ReLU activation) and pass it through a sigmoid output to classify.

•Optimization: Adam optimizer, binary cross-entropy loss; Optuna was used for hyper-parameter optimization (learning rate, dropout, attention heads, SSM latent size).

$z_{\{\text {fusion}\}}=\sigma\left(W_1 h_{\{S S M\}}+W_2 h_{\{\text {Transformer}\}}+b\right)$

where,

$h_{\{S S M\}}$: latent state representation

$h_{\{\text {Transformer}\}}$: ransformer embedding

$W_1, W_2$: trainable weights

$\sigma$: activation function

4.6 Training and validation strategy

•Stratified splitting: Splitting into training, validation and test sets, ensuring balanced classes and no data leakage, at patient level.

•10-fold cross-validation: This is useful for obtaining a good estimate of the generalization error.

•SMOTE with CV: Balancing is performed only on the training folds to evaluate fairly. The baseline and proposed models' hyperparameters, such as regularization, tree, LSTM, Transformer dimensions, and SSM latent size, are reported in Table 3.

Table 3. Hyperparameter tuning details of the baseline and hybrid models

4.7 Model evaluation and statistical tests

•Main metrics: AUC of the receiver operating characteristic curve (ROC), F1-score, recall (sensitivity), precision, accuracy.

•Calibration: Brier score, calibration curves to evaluate the probabilistic reliability.

•Precision-Recall (PR) analysis: Area under the precision-recall curve (AUPRC), especially applicable to imbalanced datasets.

•Confusion matrix: reports sensitivity/specificity, false positive and false negative rates.

•Statistical comparison:

-Bootstrapping was used to obtain 95% confidence intervals for metrics.

-Compare AUCs of paired models: DeLong test.

-A test that can be used to compare two misclassification rates is called McNemar's test.

•Error analysis:

-Compare errors by sub-group (e.g., elderly, HIV+, early versus late in the course of the disease).

-Analyze sequences where there is the greatest degree of uncertainty or fluctuating patterns.

4.8 Model explainability

•SHAP values: Measure the contribution of each feature (e.g., time, trajectory slope, early divergence) to the prediction on a global level and per patient basis.

•Show examples of attention visualizations: For each sequence, show the time windows and clinical features that the transformer attended to most.

•Case studies: Patient sequences are chosen from which interpretable predictions are made.

4.9 Deployment and clinical integration

•Model export: Final trained model stored in a format suitable for deployment, such as in Python package or as an API (ONNX or Torch Script).

•Integration: Model can be used to be integrated into EHR systems to provide real-time alerts of risk.

•Human-in-the-loop: SHAP and attention explanations provided for every prediction made per patient in the clinician dashboard.

•Prospective evaluation: Recommendations for future prospective validation and continued monitoring of model performance.

4.10 Limitations and robustness

•Generalizability: model assessed for a wide range of patient subgroups; restriction to other settings or populations of patients is mentioned with caution.

•Subgroup analysis conducted, data bias: Understand that class balancing is not an effective way to eliminate all types of data bias.

•Model drift: Model should be retrained and monitored as new data becomes available.

•Transparency: All pre-processing, modelling and evaluation scripts are packaged up and shared for re-producibility.

The methodological pipeline is comprehensive, explainable and clinically focused and aims to maximize robustness, transparency and translational potential for the prediction of outcomes in NPTB.

4.11 Validation and generalizability

The hybrid model did well in predicting the data, but the number of samples was still small (N = 125) which is not good for deep sequence learning models. Various regularizations were used such as dropout, early stopping, weight decay and stratified cross validation to reduce overfitting. However, to be deployed into the clinic, it is still necessary to test it outside in independent multi-center cohorts.

6.1 Descriptive statistics and cohort characteristics

A final dataset of 125 timepoint records was used for analysis and each record contained a variable called Time of, along with the variable status. The status was binary encoded, with “worsened” (The status = 2) status comprising about 30% of the raw data, which is common in real-world TB data. The training set was perfectly balanced by applying SMOTE in all the phases of model development (50:50 class balance).

•Median time of status: 7 (IQR: 3–14)

•Distribution of worsened status (raw): 37/125 (29.6%)

•Distribution of worsened status (SMOTE-balanced, train): 50%

6.2 Predictive model performance

•Using the state space and transformer model (the Hybrid Model), the model achieved the best discrimination (AUC), sensitivity (recall), and overall F1-score (p < 0.05 for all pairwise AUC comparisons vs. non-hybrid models, DeLong test).

•In the hybrid model, recall was highest (0.85), and is essential for clinical screening: capturing more “worsened” patients.

•Precision (0.87) and accuracy (0.88) mean that there is a low rate of false alarms and a high rate of correct predictions. The prediction performance of the proposed hybrid model is compared to the baseline models in Table 5 using the discrimination, classification and calibration and minority-class evaluation measures.

Table 5. Predictive performance of baseline models and the proposed hybrid state-space transformer model

6.3 Calibration and reliability

•Calibration curves demonstrated that the hybrid model’s predicted probabilities closely matched observed event rates, with a Brier Score of 0.14 (vs. 0.18–0.22 for baseline models).

•PR curves confirmed superior performance in the minority “worsened” class, with a mean AUPRC of 0.89 for the hybrid model (vs. 0.78–0.82 for others).

Figure 1 visually compares the predictive performance of the proposed hybrid state-space Transformer model with the baseline classifiers, highlighting its superior AUC, F1-score, recall, precision, accuracy, and AUPRC.

Figure 2 presents the calibration curve of the proposed hybrid model, showing the agreement between predicted deterioration probabilities and the observed event rates.

1.png

Figure 1. Comparative predictive performance of the proposed hybrid model and baseline machine learning models

2.png

Figure 2. Calibration curve of the hybrid state-space transformer model for predicting non-pulmonary tuberculosis (NPTB) patient status

The calibration curve shows the goodness of fit between the predicted probabilities and observed event rates of the hybrid state space transformer model in the test set. The solid line is the model's performance, and the dashed line is the ideal (perfect) model performance. The closer the curve approaches the diagonal, the more accurately the model estimates probabilities of actual risks. The hybrid model showed superior calibration performance with a low Brier score, suggesting good performance and clinically meaningful risk prediction.

6.4 Feature importance and interpretability

6.4.1 SHAP analysis

•In the hybrid model, the following global feature importance (mean abs SHAP values) was found:

1. Δ Status (first-order change in status)

2. Normalized Time of Status

3. Recent Trend Slope

4. Positional (time window) encoding

•Local SHAP plots for some patient sequences indicated that rapid changes in status, downward trending periods, and mid-sequence times dominated the factors of “worsened” predictions.

•SHAP summary plots (shown below in Figure 3) validated clinical judgment: sharp and persistent shifts indicate clinical deterioration.

The results of the temporal status change, normalized time, trend slope, volatility and positional encoding variables across the global SHAP scores are summarized in Table 6, that relative contribution of these variables across the global SHAP results: temporal status change, normalized time, trend slope, volatility and positional encoding.

3.png

Figure 3. Overall architecture of the proposed hybrid state-space transformer model for predicting non-pulmonary tuberculosis (NPTB) patient status

Table 6. Predictive performance across clinical subgroups

Table 7. Quantitative SHapley Additive exPlanations (SHAP) feature importance analysis for the hybrid model

The feature-removal analysis is shown in Table 7, which measures the loss of AUC when each temporal feature is removed from the hybrid model.

6.4.2 Transformer attention weights

•Attention heatmaps visualized which timepoints contributed most to each prediction.

•Most “worsened” predictions were driven by spikes in recent status, early rapid changes, or periods of high volatility in the patient sequence.

6.5 Error analysis

•The false negative results (missed worsened cases) were mainly observed in sequences with oscillating or late-occurring status changes or in cases with short, incomplete follow up.

•False positives: Commonly associated with transitory events that were not present in later time points.

•FN rates were slightly higher in less often sampled sequences (less than 5 timepoints), highlighting the importance of longer follow-up for more reliable prediction.

•No time-dependent model drift: Sensitivity was not different between early and mid-intervals, or between early and late intervals.

6.6 Statistical significance

•All performance enhancements that the hybrid model had over the baselines were statistically significant (p < 0.05, DeLong and McNemar's tests).

•95% confidence intervals for AUC, F1-score and recall would not overlap between the hybrid and weakest baseline.

6.7 Clinical implications

•The model's high recall and reliability indicate good potential for use as a decision support tool to flag patients who are at risk of clinical deterioration.

•SHAP and attention visualizations provide interpretable risk factors for clinical use, which allows physicians to confidently use the technology in actual clinical practice.

•The method is robust to class imbalance, and is effective on small-scale real-world datasets, making it suitable for resource-limited TB care.

6.8 Limitations and robustness checks

•Results are limited by the sample size (N = 125) and should be confirmed in larger, prospective cohorts.

•The generalizability of the model to other follow-up periods, or to EPTB populations with other patterns of comorbidity needs to be validated.

•Further improvement of predictive performance in future studies might be possible by additional features (demographics, lab data, site of disease).

To summarize, the hybrid AI framework trained and tested on your balanced, real world NPTB dataset demonstrated a high level of discrimination, recall and interpretability, better than all classical and deep learning baseline accuracy and reliability. The proposed framework used dropout regularization (0.3), weight decay (1e⁻⁵), early stopping with patience = 10 epochs and latent dimension compression to minimize the risk of overfitting due to the small size of the datasets and the transformer architecture.