TRACE: Transformer-based Risk Assessment for Clinical Evaluation

TL;DR

TRACE addresses clinical risk assessment with highly multi-modal, often incomplete tabular data by introducing a Transformer-based architecture that unifies continuous, categorical, and checkbox features into a single embedding space processed by a Transformer encoder. A strong non-negative MLP baseline (nnMLP) is proposed for comparison, highlighting interpretability through monotonic risk contributions. TRACE explicitly handles missing values and provides attention-based explanations via feature maps, achieving competitive or state-of-the-art performance across six diverse clinical datasets with favorable parameter efficiency. The work demonstrates robust handling of class imbalance through focal loss and ablation studies confirm the value of checkbox embeddings and the proposed data representations. Overall, TRACE offers practical, interpretable, and efficient risk assessment suitable for clinical decision support, with publicly available code for reproducibility.

Abstract

We present TRACE (Transformer-based Risk Assessment for Clinical Evaluation), a novel method for clinical risk assessment based on clinical data, leveraging the self-attention mechanism for enhanced feature interaction and result interpretation. Our approach is able to handle different data modalities, including continuous, categorical and multiple-choice (checkbox) attributes. The proposed architecture features a shared representation of the clinical data obtained by integrating specialized embeddings of each data modality, enabling the detection of high-risk individuals using Transformer encoder layers. To assess the effectiveness of the proposed method, a strong baseline based on non-negative multi-layer perceptrons (MLPs) is introduced. The proposed method outperforms various baselines widely used in the domain of clinical risk assessment, while effectively handling missing values. In terms of explainability, our Transformer-based method offers easily interpretable results via attention weights, further enhancing the clinicians' decision-making process.

Figures (6)

• Figure 1: Non-negative MLP architecture. The weight matrices $W^{1,2,3 (+)}$ are constrained to non-negative values, and the biases $b^{1,2 (-)}$ of the two hidden layers to negative values. Bias $b^3$ is left unconstrained.
• Figure 2: The proposed TRACE model. The model supports three types of input data: continuous, categorical, and "checkbox" features. Continuous data are processed through a two-layer MLP with a weight masking mechanism on its outputs. Categorical and "checkbox" data are embedded through their respective embedding layers. All embeddings are concatenated and fed into a transformer encoder block followed by a linear head.
• Figure 3: Attention map visualization generated by the TRACE model's final layer, for the SBMS (left column) and BRFSS 2022 (right column) datasets. The top row demonstrates the attention weights for each feature. For the BRFSS 2022 dataset, 100 random samples are selected from the validation set. The bottom row depicts the relationship between queries (rows) and keys (columns) (Best viewed zoomed-in)
• Figure 4: Performance of TRACE in terms of F1-Score and Balanced Accuracy in relation to the ratio of missing values on the SBMS dataset.
• Figure 5: Balanced Accuracy vs. number of trainable parameters on SBMS (left) and BRFSS-2022 (right) datasets. Results are shown for the proposed nnMLP and TRACE models against TabNet, DANET, FT-Transformer, GANDALF and GATE.