[Paper Review] Language Models are Few-shot Learners for Prognostic Prediction

Zekai Chen, Mariann Micsinai Balan+1 more|arXiv (Cornell University)|Feb 24, 2023

Radiomics and Machine Learning in Medical Imaging9 citations

TL;DR

The paper demonstrates that large language models can be fine-tuned or prompted to predict immunotherapy prognosis from clinical and molecular tabular data, achieving superior or competitive accuracy in few-shot settings across multiple cancer types compared to traditional ML baselines.

ABSTRACT

Clinical prediction is an essential task in the healthcare industry. However, the recent success of transformers, on which large language models are built, has not been extended to this domain. In this research, we explore the use of transformers and language models in prognostic prediction for immunotherapy using real-world patients' clinical data and molecular profiles. This paper investigates the potential of transformers to improve clinical prediction compared to conventional machine learning approaches and addresses the challenge of few-shot learning in predicting rare disease areas. The study benchmarks the efficacy of baselines and language models on prognostic prediction across multiple cancer types and investigates the impact of different pretrained language models under few-shot regimes. The results demonstrate significant improvements in accuracy and highlight the potential of NLP in clinical research to improve early detection and intervention for different diseases.

Motivation & Objective

• Motivate prognostic prediction in cancer immunotherapy and address data scarcity by exploring transformer-based and language model approaches.
• Evaluate whether transformers improve prognostic prediction over conventional ML on real-world MSK-IMPACT data.
• Assess how different pretrained language models and knowledge sources influence downstream clinical prediction under few-shot conditions.

Proposed method

• Serialize clinical tabular features into natural language token sequences consumable by LLMs.
• Leverage an end-to-end tuning approach where a frozen LLM is augmented with a trainable encoder and a multi-task loss Lf combining multiple endpoints.
• Use domain-specific and domain-agnostic LLMs pretrained on various corpora to assess knowledge transfer under few-shot regimes.
• Adopt ClinTaT, an enhanced TabTransformer-based model, incorporating a continuous embedding layer for continuous features and concatenating embeddings before transformer processing.
• Train with AdamW and a learning-rate warmup with cosine annealing; use a multi-task loss across OS, PFS, and BOR with task-specific losses (CrossEntropy for BOR, CoxPH for OS/PFS).
• Evaluate robustness and generalization via AUC for treatment response and C-index for OS/PFS across cancer types.

Experimental results

Research questions

• RQ1 To what extent can transformer-based models improve clinical prognostic prediction over conventional machine learning with limited labeled data?
• RQ2 How do different pretrained language models, especially domain-specific versus domain-agnostic, influence prognostic prediction in few-shot settings?
• RQ3 Can LLM-based approaches produce accurate OS/PFS predictions prior to immunotherapy administration across cancer types?
• RQ4 Is a multi-endpoint, joint learning objective beneficial for cancer prognostic prediction with heterogeneous clinical outcomes?

Key findings

• ClinTaT achieves superior AUC for treatment response prediction across cancer types compared to logistic regression, random forest, and XGBoost when using all samples (and generally outperforms traditional baselines across most settings).
• Transforms can yield significant C-index improvements for OS and PFS at the pan-cancer level (ClinTaT C-index 0.724 vs XGBoost 0.688; p<0.05).
• Under few-shot settings, domain-specific pretrained models (notably PubMedBERT) provide the largest gains, with up to 9.4% absolute improvement in very small samples (4-shot) for treatment response prediction.
• An additional transformer encoder on top of LLM embeddings consistently improves performance over a single linear finetuning layer.
• PubMedBERT pretrained from scratch often yields the best average improvement across k-shot settings, outperforming other encoders in multiple configurations.
• In very few-shot regimes, LLM-based finetuning shows notable gains (e.g., 9.4% in 4-shot), indicating effective knowledge transfer to clinical tasks.