Attention-enabled Explainable AI for Bladder Cancer Recurrence Prediction

TL;DR

This work proposes an interpretable deep learning framework that integrates vector embeddings and attention mechanisms to improve NMIBC recurrence prediction performance and identifies new important factors influencing recurrence, such as surgical duration and hospital stay, which had not been considered in existing NMIBC prediction models.

Abstract

Non-muscle-invasive bladder cancer (NMIBC) is a relentless challenge in oncology, with recurrence rates soaring as high as 70-80%. Each recurrence triggers a cascade of invasive procedures, lifelong surveillance, and escalating healthcare costs - affecting 460,000 individuals worldwide. However, existing clinical prediction tools remain fundamentally flawed, often overestimating recurrence risk and failing to provide personalized insights for patient management. In this work, we propose an interpretable deep learning framework that integrates vector embeddings and attention mechanisms to improve NMIBC recurrence prediction performance. We incorporate vector embeddings for categorical variables such as smoking status and intravesical treatments, allowing the model to capture complex relationships between patient attributes and recurrence risk. These embeddings provide a richer representation of the data, enabling improved feature interactions and enhancing prediction performance. Our approach not only enhances performance but also provides clinicians with patient-specific insights by highlighting the most influential features contributing to recurrence risk for each patient. Our model achieves accuracy of 70% with tabular data, outperforming conventional statistical methods while providing clinician-friendly patient-level explanations through feature attention. Unlike previous studies, our approach identifies new important factors influencing recurrence, such as surgical duration and hospital stay, which had not been considered in existing NMIBC prediction models.

Figures (5)

• Figure 1: 3D visualization of vector embeddings for categorical features, showing recurrence risk patterns learned by the AI model. In the IntravesicalTreatment plot (bottom), "BCG Induction Only" is farthest from the origin, indicating the highest recurrence risk. "None" is at the origin, suggesting it as the baseline risk, while "BCG Induction and Maintenance" shows the lowest recurrence risk. In the SmokingStatus plot (top), "Never" smokers are farthest from the other categories, indicating the lowest risk, while "Current" smokers are closest to the origin, suggesting the highest risk. "Previous" smokers lie in between, reflecting an intermediate risk level.
• Figure 2: Neural Network Architecture with Vector Embeddings and Self-Attention used in this work
• Figure 3: Snippet of the patient-level attention heatmap, where each row represents a patient and each column corresponds to a unique clinical or demographic variable. Darker cells represent stronger attention weights, indicating where the model "looks" more intently for each individual’s recurrence risk. The red box in the Recurrence column designates that these patients are predicted to experience recurrence. Notably, although they share the same outcome, each patient’s attention weights differ across various features, highlighting the individualized nature of the model’s analysis. This transparent focus mimics an experienced clinician's reasoning to pinpoint the most influential predictors for any given patient.
• Figure 4: Training (blue) and validation (orange) accuracy across 250 epochs for the proposed neural network. The model achieves a final validation accuracy of about 70%, with training accuracy levelling at around 80%.
• Figure 5: Feature importance rankings from the attention-based neural network for NMIBC recurrence prediction. The gold stars indicate newly recognized variables—such as Surgery Time, Total Number of Cigarettes Smoked, and Total Days in Hospital—which were not prioritized in earlier NMIBC studies. This novel ordering highlights the capacity of modern attention mechanisms to uncover clinically significant factors beyond those identified by conventional statistical or machine learning models.