Survival Prediction Using Transformer-Based Categorical Feature Representation in Diffuse Large B-Cell Lymphoma

Healthcare (Basel) 2023 AI 8 Explanations View Original
Original Paper (PDF)

Unable to display PDF. Download it here or view on PMC.

Plain-English Explanations
Pages 1-2
The Challenge of Predicting Survival in Diffuse Large B-Cell Lymphoma

Diffuse large B-cell lymphoma (DLBCL) is the single most common subtype of lymphoma, accounting for roughly one-third of all non-Hodgkin lymphomas, which themselves represent approximately 90% of all lymphoma diagnoses. DLBCL is characterized by its aggressive behavior and tendency to be diagnosed at advanced stages, particularly in older patients. Despite standard first-line chemoimmunotherapy, approximately 30 to 40% of patients relapse or prove refractory to initial treatment, making accurate prognosis prediction a critical unmet clinical need.

The limits of the International Prognostic Index: The International Prognostic Index (IPI), developed in 1993, remains the dominant clinical prognostic tool for DLBCL. It incorporates five pretreatment variables: age, clinical stage, lactate dehydrogenase (LDH) level, ECOG performance status, and the number of extranodal involvement sites. While clinically useful, the IPI is increasingly recognized as insufficient because individual patient outcomes are highly heterogeneous even within the same IPI risk group. A single-point scoring system built on five binary variables cannot capture the biological complexity that drives survival differences.

The Deauville score as an on-treatment prognostic signal: The Deauville score (DS) is a five-point visual scale used to interpret F-18 fluorodeoxyglucose (FDG) PET/CT scans. Staging and interim FDG PET/CT plays a critical role in assessing disease extent and treatment response in DLBCL. Multiple studies have demonstrated that the DS obtained at interim PET (performed mid-treatment to assess chemosensitivity) is predictive of long-term patient outcomes. A DS of 1, 2, or 3 is generally considered favorable, while 4 or 5 indicates residual or progressive disease. Because DS values encode disease severity on an ordinal scale rather than as independent categories, they present a specific technical challenge for standard statistical models.

This 2023 paper from Chonnam National University, published in the journal Healthcare (within the Cancers domain), proposes a transformer-based deep learning model called TTSurv to improve survival time prediction in DLBCL. The study is notable for incorporating both pretreatment clinical variables and on-treatment Deauville scores, for conducting experiments across two independent institutional datasets, and for explicitly addressing the challenge of handling ordinal categorical features in clinical survival models.

TL;DR: DLBCL is the most common NHL subtype, with 30-40% of patients relapsing on standard therapy. The IPI (5 variables, developed 1993) cannot capture within-group heterogeneity. The Deauville score from interim PET/CT is a strong prognostic signal but poses a categorical encoding challenge. This paper proposes TTSurv, a transformer-based deep learning survival model, trained on 604 patients and validated on 220 independent patients.
Pages 2-3
From Cox Models to Deep Learning: The Survival Analysis Landscape

Survival analysis in oncology has traditionally been dominated by the Cox proportional hazard (CPH) model, a semiparametric approach that estimates the hazard of an event (such as death or relapse) as a function of patient covariates. The CPH model assumes a linear relationship between covariates and the log hazard, which limits its ability to capture non-linear interactions among clinical variables. Despite this limitation, CPH remains the reference standard in clinical research because of its interpretability and mathematical properties.

Machine learning alternatives: A range of non-parametric and machine learning survival models have been developed to overcome the linearity constraint. Random survival forests (RSF) and oblique random survival forests (ORSF) partition the feature space in tree-based ensembles that capture non-linear covariate effects. Hazard boosting methods leverage gradient boosting to model the hazard function non-parametrically. Bayesian networks combined with the CPH model have been explored to improve both prediction performance and model interpretability.

Deep learning survival models: More recently, deep learning approaches have replaced the linear predictor in Cox-based frameworks with deep feed-forward neural networks. DeepSurv extends CPH with a multi-layer perceptron and demonstrated better C-index performance than standard CPH by modeling complex covariate interactions and supporting treatment recommendations. CoxTime lifts the proportionality constraint of CPH by allowing time-dependent covariate effects, meaning the relationship between a predictor and the hazard can change over the follow-up period. CoxCC is a proportional variant of CoxTime. Both CoxTime and CoxCC are implemented in the pycox Python library and serve as the primary deep learning baselines in this study.

The categorical feature problem: A limitation shared by all of these approaches is their handling of categorical clinical features. Variables like Ann Arbor disease stage (I, II, III, IV) or the Deauville score (1 through 5) are typically encoded as integers or one-hot vectors, neither of which captures the ordinal relationship between categories. Stage IV is not merely "different" from Stage I; it is more severe. Standard encoding methods either treat these as purely continuous (incorrectly assuming uniform intervals) or as fully independent categories (losing the severity ordering). Transformer-based categorical embedding offers a principled solution to this problem, motivating the development of TTSurv.

TL;DR: Standard CPH assumes linear covariate effects. Deep learning alternatives (DeepSurv, CoxTime, CoxCC) add non-linearity through neural networks but still encode categorical variables poorly. Stage and Deauville score values are ordinal, not independent classes. TTSurv addresses this by using transformer-based categorical embedding to capture severity relationships between ordered clinical categories.
Pages 3-4
TTSurv: Architecture of the Transformer-Based Survival Model

TTSurv is a dual-input deep learning survival model designed to handle the mixed data types present in DLBCL clinical datasets. The model accepts two separate input branches: one for continuous numerical features (age and LDH in IU/L) and one for categorical features (sex, ECOG performance status, LDH normal/elevated status, number of extranodal sites, bone marrow involvement, B symptoms, Ann Arbor stage, IPI score, and when available, the Deauville score). The dual-branch design allows each data type to be processed with the most appropriate encoding strategy before being combined for the final survival prediction.

Categorical embedding with column tokens: The categorical branch follows the column embedding method introduced by Huang et al. in the TabTransformer framework. Each categorical feature xn in the set of N categorical variables is embedded into a learned vector with dimension d. To allow the model to distinguish between different categorical features (since tabular data lacks the sequential positional structure present in natural language), special column tokens are added alongside each feature's class embeddings. The total number of embeddings passed to the transformer is n times c plus t, where c is the number of classes for a given feature and t is the number of special tokens. This structure enables the transformer's attention mechanism to learn which categorical features and which class values are most relevant in combination.

Transformer encoder layers: The embedded categorical features are processed through a stack of transformer encoder layers. Each encoder layer consists of multi-head self-attention, layer normalization, residual connections, and a two-layer MLP. The self-attention operation computes attention scores between every pair of categorical feature embeddings, producing context-aware feature representations that encode inter-feature relationships. The multi-head variant performs this attention in parallel across multiple subspaces, enabling the model to simultaneously attend to different aspects of the categorical feature interactions. This mechanism allows TTSurv to learn, for example, that a high Deauville score combined with advanced stage carries different prognostic implications than either factor alone.

Survival prediction head: The categorical embeddings from the final transformer layer are concatenated with the continuous numerical features and passed through a residual dense block for the survival prediction task. TTSurv follows the time-dependent Cox approach of CoxTime, where the relative risk function is h(t|x) = h0(t) times exp[g(t,x)], with the time-dependent risk function g learned by the neural network. The model is trained by minimizing the negative log-likelihood of the Cox partial likelihood. Cumulative hazard and survival functions are estimated from predicted risks using the Breslow estimator.

TL;DR: TTSurv uses dual input branches for continuous (age, LDH) and categorical features (stage, DS, IPI, extranodal sites, etc.). Categorical features are encoded with column-token embeddings following the TabTransformer approach, then processed through multi-head self-attention transformer encoders. The output is concatenated with continuous features and fed into a time-dependent Cox survival head optimized via negative log-likelihood. Survival functions are estimated with the Breslow estimator.
Pages 4-6
Multi-Institutional DLBCL Dataset and Experimental Design

The study uses clinical data collected between 2011 and 2018 from two South Korean hospitals: Chonnam National University Hwasun Hospital (CNUHH, n = 604) and Jeonbuk National University Hospital (JBUH, n = 220). The CNUHH dataset served as the training and validation set, while the JBUH dataset was held out as a completely independent external test cohort to assess cross-institutional generalizability. Both datasets were approved by institutional review boards (CNUHH-2022-095 and CUH 2022-11-013). The log-rank test confirmed that individual clinical covariates were statistically significant predictors of overall survival (p less than 0.005) in both cohorts.

Censoring and follow-up: Both datasets had similar proportions of censored cases: 70.86% in CNUHH and 73.18% in JBUH. Censored patients are those who survived until the end of the observation period or were lost to follow-up without experiencing the event of interest (death). The high censoring rate is typical for DLBCL cohorts given that many patients achieve long-term remission with standard therapy, but it creates a statistical challenge because the exact survival time is unknown for these patients.

Feature set: Only age and LDH (measured in IU/L) were treated as continuous variables. All other covariates were treated as categorical: sex (female/male), ECOG performance status (0, 1, 2, 3, 4), number of extranodal involvement sites (0 through 6), bone marrow involvement (yes/no), B symptoms (yes/no), Ann Arbor stage (I, II, III, IV), and IPI score (0 through 5). The Deauville score (1 through 5 on the PET/CT scale) was included only for on-treatment analysis. The age distributions differed between cohorts: CNUHH ranged from 36 to 81 years, while JBUH spanned 15 to 87 years, reflecting differences in the patient populations at each institution.

Experimental conditions: To rigorously assess the incremental value of each data source, the authors conducted separate experiments for pretreatment and on-treatment feature sets. Pretreatment experiments used only clinical variables available before therapy began. On-treatment experiments added the Deauville score obtained from interim PET/CT. All models were trained using stratified five-fold cross-validation on the CNUHH dataset, with model weights from each fold used for inference on the independent JBUH test set. The final reported JBUH performance is the average across all five fold-trained models, providing a robust estimate of generalizability. All experiments were implemented in Python using PyTorch, with baseline models from the pycox library, run on an Nvidia GeForce RTX 3080Ti GPU with 12 GB of memory.

TL;DR: Training set: 604 DLBCL patients from CNUHH (2011-2018); independent external test set: 220 patients from JBUH. Both cohorts had roughly 71-73% censored cases. Pretreatment features included IPI components plus extranodal sites, bone marrow involvement, and B symptoms. On-treatment analysis added interim PET Deauville score. Five-fold cross-validation on CNUHH; external generalizability assessed on JBUH.
Pages 7-9
Performance Results: C-Index and Mean Absolute Error Comparisons

The study evaluates all survival models using two complementary metrics. The concordance index (C-index) measures how well a model ranks patients by predicted survival time, with 1.0 indicating perfect discrimination and 0.5 indicating no better than random. The mean absolute error (MAE) measures the average difference in days between predicted and actual survival times, computed only for patients who experienced the event (non-censored). The MAE is a clinically meaningful metric because it quantifies how far off the model's survival time estimates are in actual days, not just whether it ranks patients correctly.

C-index results: On the CNUHH training/validation set using pretreatment features, TTSurv achieved a C-index of 0.7245, compared to 0.7213 for DeepSurv, 0.7134 for CoxPH, 0.6929 for CoxTime, and 0.6925 for CoxCC. With the addition of the on-treatment Deauville score, all models improved: TTSurv reached 0.7457, CoxPH rose to 0.7440, DeepSurv to 0.7428, CoxTime to 0.7358, and CoxCC to 0.7055. On the independent JBUH external test set using pretreatment features, CoxPH achieved the highest C-index at 0.7858, followed by TTSurv at 0.7756, and DeepSurv at 0.7403. With on-treatment features, CoxPH led at 0.7990, followed by TTSurv at 0.7950 and CoxTime at 0.7782.

MAE results: TTSurv demonstrated substantially stronger advantages in MAE. On CNUHH pretreatment features, TTSurv achieved an MAE of 995.32 days versus 1047.79 days for CoxCC (the next best), 1092.89 for CoxPH, 1117.52 for DeepSurv, and 1195.63 for CoxTime. On the JBUH external test set using pretreatment features, TTSurv's MAE of 613.51 days represented a dramatic improvement over CoxCC (798.58 days) and CoxPH (822.15 days), a reduction of approximately 185 days over the best-performing existing method. With on-treatment features on JBUH, TTSurv's MAE fell further to 559.80 days, versus 662.04 for CoxPH, 774.65 for CoxTime, and 783.57 for CoxCC.

Impact of the Deauville score: Including the Deauville score as an on-treatment feature consistently improved performance across all models, but the improvement was most pronounced in the MAE metric. For TTSurv on JBUH, adding the DS reduced MAE by 53.71 days. The C-index improvement was 0.02 (from 0.7756 to 0.7950). These quantitative gains reinforce the well-established clinical importance of interim PET response assessment in DLBCL prognostication.

TL;DR: TTSurv achieved a JBUH pretreatment MAE of 613.5 days vs. 798.6 days for the next best model, a 185-day improvement. Adding the Deauville score further reduced MAE to 559.8 days. C-index ranged from 0.7245 to 0.7950 for TTSurv across conditions. CoxPH led on C-index in the external set but TTSurv led decisively on MAE, the clinically more meaningful survival time accuracy metric.
Pages 9-10
Why Both C-Index and MAE Are Needed to Evaluate Survival Models

The C-index is the most widely used evaluation metric in survival analysis and serves as the primary benchmark in most published studies. It measures the probability that, for any two patients where one had a shorter observed survival time, the model assigns a higher risk score to the patient who died sooner. A C-index of 0.75, for example, means that the model correctly ranks 75% of comparable pairs. This metric is insensitive to the absolute magnitude of the survival time estimates; it only cares about relative ordering.

The limitation of C-index alone: A model could rank patients perfectly (high C-index) but predict everyone's survival time as, say, 1000 days regardless of whether their actual survival was 300 or 3000 days. In clinical practice, the absolute survival time estimate matters. A physician advising on treatment intensity or palliative care planning needs a model that estimates, with reasonable accuracy, how many months or years a patient is likely to survive, not just whether patient A has a worse prognosis than patient B. The C-index does not penalize predictions that are correctly ranked but wildly wrong in absolute time.

MAE as a complementary metric: The MAE addresses this gap by computing the average absolute difference in days between the model's predicted median survival time and the actual observed survival time, restricted to non-censored patients where the true outcome is known. The formula involves summing the absolute difference between predicted and true survival times over all event-observed patients and dividing by their count. An MAE of 559 days means the model's predicted median life is off by about 18.6 months on average for patients who died during the study period. Lower MAE directly translates to more accurate survival time estimates for clinical use.

The study's finding that TTSurv leads most clearly on MAE while C-index differences are smaller suggests that the transformer-based categorical embedding primarily improves the model's ability to estimate absolute survival times rather than just the relative ranking of patients. This distinction is clinically meaningful: MAE improvement represents better absolute survival time accuracy, which is more directly useful for individualized treatment planning and patient counseling.

TL;DR: C-index measures only ranking accuracy (0.5 = random, 1.0 = perfect); it ignores whether survival time estimates are numerically accurate. MAE measures average days of error in absolute survival time predictions, restricted to patients with known outcomes. TTSurv's largest advantage is in MAE (559-613 days vs. 798-1117 days for competitors), suggesting transformer embedding improves absolute survival time estimation more than ranking.
Pages 10-11
Interpreting the Results: What Transformer Embedding Adds to Clinical Survival Analysis

The authors' main conclusion is that categorical embedding via transformer architecture improves survival time prediction in DLBCL by learning contextual relationships among categorical feature classes rather than treating them as independent one-hot vectors. The improvement is most pronounced on the external JBUH dataset, suggesting that the features learned by the transformer generalize better to new patient populations than those learned by simpler architectures. This generalizability gain is practically important: clinical AI tools must perform well outside their training environment to be clinically deployable.

Why categorical features matter: Clinical data frequently contain ordinal categorical variables where class values carry severity information. Disease stage (I through IV), Deauville score (1 through 5), IPI score (0 through 5), and ECOG performance status (0 through 4) all encode gradations of disease burden or patient fitness. When these are encoded as integers, the model assumes equal spacing between levels (i.e., the difference between Stage I and II is treated as identical to the difference between Stage III and IV). When one-hot encoded, all ordering information is lost. Transformer-based embeddings learn a dense vector representation for each class value that positions it in a learned feature space where similar severity levels cluster together. This is analogous to word embeddings in NLP, where semantically similar words are geometrically proximate in embedding space.

The prognostic importance of the Deauville score: The consistent improvement across all models when the DS is added confirms its high prognostic value. The DS is obtained only after the interim PET scan (typically after 2 to 4 cycles of R-CHOP), so it is not available at treatment initiation. The finding that even a 0.02 improvement in C-index and a 53.71-day improvement in MAE results from adding a single ordinal feature reinforces the case for interim PET-guided adaptive treatment strategies in DLBCL.

The individual survival curve predictions shown in Figure 4 illustrate both the model's strengths and limitations. For most patients, predicted survival times are close to ground-truth values (e.g., patient JBUH_DLB007 had an actual survival of 1537 days, predicted as 977 days before treatment and 1531 days during treatment). However, patient JBUH_DLB106, with Stage III disease, an LDH of 756, and a DS of 3, had an actual survival of 2127 days while the model predicted 673 days before treatment and 677 days during treatment. This case represents a clinically recognized phenomenon: some patients with moderately severe features survive substantially longer than expected, possibly reflecting tumor biology not captured by the available clinical variables.

TL;DR: Transformer embedding improves survival prediction by learning ordinal relationships within categorical variables (stage, DS, IPI, ECOG). The Deauville score adds 0.02 C-index and 53.71-day MAE improvement when included as an on-treatment feature. Individual predictions are generally accurate but the model underestimates survival for outlier cases with unexpectedly favorable outcomes despite moderate-risk profiles.
Pages 11-12
Limitations of the Current Model and Paths Forward

Reliance on clinical and Deauville score data only: TTSurv is trained exclusively on structured clinical variables and the Deauville score calculated by experienced nuclear medicine physicians. It does not incorporate raw PET/CT image data, radiomic features derived from tumor volumes, histopathological features from tissue biopsies, or molecular profiling data. Patients with similar clinical profiles can have dramatically different tumor biology at the molecular level (e.g., GCB vs. ABC cell-of-origin subtypes, MYC/BCL2 double-hit genetics), and these biological differences drive survival outcomes that clinical variables alone cannot fully capture.

Data censoring: The high censoring rate (approximately 71-73% in both cohorts) represents a fundamental challenge. The MAE metric can only be computed for the approximately 27-29% of patients who experienced the event during the study period. The 70%+ of patients who were censored contribute to training through the Cox likelihood but cannot be used to evaluate absolute survival time accuracy. Larger datasets with longer follow-up periods and more observed events would enable more robust MAE evaluation and model training.

Dataset size and demographic scope: The training cohort of 604 patients, while reasonable for a single-institution study, is modest relative to the number of parameters in deep learning models. The age distribution difference between CNUHH (36-81 years) and JBUH (15-87 years) represents one form of demographic shift, and models may not generalize well to populations with substantially different treatment practices, ethnicity distributions, or follow-up protocols. Validation on broader, more diverse international cohorts is needed before clinical deployment.

Future directions: The authors identify two main directions for future work. First, incorporating radiological imaging data (PET/CT scans) directly into the model to complement the clinician-assigned Deauville score with quantitative imaging features such as total metabolic tumor volume (TMTV), standardized uptake value (SUV), and radiomic texture descriptors. Second, validating the model on larger and more diverse patient populations, including multi-institutional datasets that span different treatment eras, geographic regions, and institutional practices. Both directions reflect the broader trajectory of clinical AI: moving from single-modality, single-center models toward multimodal, multi-institutional systems capable of robust real-world deployment.

TL;DR: TTSurv is limited to clinical and Deauville score inputs only; raw PET/CT imaging and molecular profiling are excluded. High censoring rate (71-73%) restricts MAE evaluation to roughly 27-29% of patients. The training cohort of 604 patients is modest for deep learning. Future work targets direct PET/CT image integration and validation on larger, more diverse multi-institutional cohorts.