Abstract
Genomic prediction models often fail to transfer across institutions due to differences in sequencing panels, leading to structural feature missingness at deployment. Existing approaches typically restrict analysis to shared genes, exclude patients with incomplete profiles, or rely on test-time imputation, which can reduce robustness and limit the use of multi-center data.
We propose Survival prediction Handling Incomplete Features using Transformer (SHIFT), a missingness-aware survival model that directly predicts from incomplete genomic inputs without test-time imputation. SHIFT represents each genomic feature separately and employs masked self-attention along with a feature-availability mask, ensuring that predictions are based only on observed inputs. Additionally, we introduce variable-rate feature masking during training to enhance robustness to heterogeneous missingness patterns.
We evaluate the approach on glioblastoma and lung squamous cell carcinoma with external validation across multiple cohorts, including a challenging setting with severe cross-cohort panel mismatch. SHIFT demonstrates strong generalization and compares favorably with standard survival baselines and imputation-based methods while utilizing a single model across differing feature sets. Furthermore, we find that incorporating patients from incomplete cohorts during development can enhance performance on external data, indicating that partially observed cohorts need not be excluded from model building. These results support missingness-aware modeling as a practical strategy for multi-center survival prediction in precision oncology.
Blogger's Review: The SHIFT model effectively addresses the challenges of incomplete genomic data through its innovative masked self-attention mechanism, showcasing broad applicability in multi-center survival prediction. This approach not only enhances model robustness but also opens new avenues for research in precision medicine. By not excluding partially observed data, SHIFT has the potential to drive advancements in clinical research.