Toward a Well-Calibrated Discrimination via Survival Outcome-Aware Contrastive Learning

A contrastive survival representation learning framework for time-to-event prediction.

---

⟡ Table of Contents

Click to expand

• Overview
• Features
• Project Structure
  • Project Index
• Getting Started
  • Prerequisites
  • Installation
  • Usage
  • Testing
• Roadmap
• Contributing
• License
• Acknowledgments

---

◈ Overview

ConSurv introduces a novel contrastive learning framework for survival analysis that enhances discrimination while maintaining calibration. By weighting contrastive pairs according to survival outcome similarity, ConSurv achieves state-of-the-art performance across multiple clinical datasets without directly manipulating model outputs.

---

◇ Project Structure

ConSurv.git/
├── README.md
├── dataset.py
├── learning.py
├── metric.py
├── model.py
├── run.py
├── utils.py
└── requirements.txt

⟠ Getting Started

⟁ Prerequisites

• Python ≥ 3.8
• CUDA 11.7+ (optional for GPU)

Verified compatible versions:

Package	Version	Purpose
PyTorch	2.0.1	Core deep learning framework
scikit-learn	1.3.2	Data preprocessing and scaling
scikit-survival	0.22.2	Survival metrics (CI, Brier)
pycox	0.2.3	Deep survival baselines (optional)
NumPy	1.24.3	Array operations
Pandas	2.0.3	Tabular data handling

Other installed packages (xgboost, statsmodels, umap-learn, etc.) are not required for the main pipeline.

---

⟒ Installation

git clone https://github.com/dongzza97/ConSurv.git
cd ConSurv.git
pip install -r requirements.txt

---

⟓ Usage

1. Dataset format

CSV file structure (datasets/{DATA_NAME}.csv):

feature_1	feature_2	...	time	label
0.12	5.6	...	128	1
7.80	0.2	...	64	0

• label: 1 = event, 0 = censored
• time: discrete time step (divided by 30 days for some datasets)

2. Train & Validate

python run.py \
  --data ./datasets \
  --data-name METABRIC \
  --epochs 2000 \
  --batch-size 32 \
  --hidden-dim 16 \
  --depth 4 \
  --drop-out 0.2 \
  --version out \
  --sigma 0.75 \
  --quantile 10

3. Output & Checkpoints

After training completes:

• Model checkpoint → ./{DATA_NAME}/NLL_SNCE_{SEED}.pt
• Global metrics → {DATA_NAME}_NLL_NCE_performance.csv
• Time-dependent CI → {DATA_NAME}_NLL_NCE_TD_Cindex.csv
• Time-dependent BS → {DATA_NAME}_NLL_NCE_TD_Brier.csv

---

⌆ Testing

run.py automatically executes the test phase after training. Performance metrics and CSV results are saved in the current working directory.

---

⟲ Roadmap

• Base ConSurv contrastive + hazard framework
• Time-dependent calibration metrics (DDC, DCAL)
• Integrate baselines
• Multi-modal expansion for EHR + imaging data

---

⏣ Contributing

1. Fork the repository
2. Create a feature branch
3. Commit with a clear message
4. Push and open a pull request

Bug reports, improvements, and metric extensions are always welcome.