KMNet: Discrete-Time Survival Analysis with Deep Learning

⚡ State-of-the-art survival analysis with deep learning and ranking losses ⚡

Documentation • Features • Installation • Quick Start • Demo Notebook • Citation

🎯 Use Cases & Applications

KMNet combines the power of deep neural networks with Kaplan-Meier inspired ranking losses to deliver state-of-the-art performance in discrete-time survival analysis. Perfect for clinical research, reliability engineering, and time-to-event prediction.

It extends standard neural survival models by incorporating a novel Kaplan-Meier inspired rank loss, allowing the model to learn not just from local hazard rates but also from global ranking constraints inherent in survival data.

This library provides a high-performance implementation using PyTorch JIT to speed up custom loss calculations, making it suitable for large-scale survival datasets.

🚀 Key Features

📊 Why KMNet?

📦 Installation

pip install kmnet

Or install the latest development version from source:

git clone https://github.com/yuvrajiro/KMNet.git
cd KMNet
pip install -e ".[dev]"

Requirements

• Python >= 3.8
• PyTorch >= 1.7.0
• NumPy, Pandas
• torchtuples
• pycox
• numba

⚡ Quick Start

Here is a complete example of how to use KMNet on a synthetic dataset. You can also check out the interactive demo notebook and the full documentation.

1. Data Preparation

KMNet requires the target variable (time and event) to be discretized. We use label_transforms from pycox for this.

import numpy as np
import pandas as pd
from kmnet.model import KMNet

# Generate synthetic data
def make_data(n=1000):
    X = np.random.randn(n, 5).astype('float32')
    T = np.random.exponential(1 / (0.1 * np.exp(0.5 * X[:, 0])))
    C = np.random.exponential(1 / 0.05, size=n)
    time = np.minimum(T, C)
    event = (T <= C).astype('float32')
    return X, time, event

X, time, event = make_data()

# Discretize time into 20 bins
num_durations = 20
labtrans = KMNet.label_transform(num_durations)
get_target = lambda df: (df['duration'].values, df['event'].values)

df = pd.DataFrame({'duration': time, 'event': event})
y = labtrans.fit_transform(*get_target(df))

2. Define the Neural Network

You can use any PyTorch network architecture. The output dimension must match the number of time bins.

import torch.nn as nn

in_features = X.shape[1]
out_features = labtrans.out_features

net = nn.Sequential(
    nn.Linear(in_features, 32),
    nn.ReLU(),
    nn.Linear(32, 32),
    nn.ReLU(),
    nn.Linear(32, out_features)
)

3. Train the Model

Initialize KMNet and fit it to the data.

# Initialize model with the discretized time grid
model = KMNet(net, duration_index=labtrans.cuts)

# Train
batch_size = 64
epochs = 10
model.fit(X, y, batch_size, epochs, verbose=True)

4. Prediction and Visualization

Predict survival functions and plot them.

import matplotlib.pyplot as plt

# Predict survival probabilities for the first 5 samples
surv_df = model.predict_surv_df(X[:5])

# Plot
plt.figure(figsize=(10, 6))
for col in surv_df.columns:
    plt.step(surv_df.index, surv_df[col], where="post")
plt.ylabel("Survival Probability")
plt.xlabel("Time")
plt.title("Predicted Survival Curves")
plt.show()

🔬 Mathematical Background

KMNet models the discrete conditional survival $p(t | x)$. The survival function is given by:

$$ \boxed{S(t | x) = \prod_{k=0}^{t} p(k | x)} $$

The loss function $\mathcal{L}$ is a weighted sum of two components:

$$ \boxed{\mathcal{L} = \alpha \mathcal{L}{bce} + (1 - \alpha)\lambda \mathcal{L}{Rank}} $$

📊 Performance & Optimization

⚡ Speed Benchmarks

The KMNet class leverages TorchScript (JIT) to compile the custom loss functions, eliminating Python interpreter overhead.

🔥 JIT-Optimized: 1.38s
📈 Speedup: 1.6x faster
💾 Memory: ~15% reduction

🛠️ Advanced Configuration

You can customize the loss function parameters when initializing the model:

1. base: 'nll' (Negative Log-Likelihood) or 'bce' (Binary Cross-Entropy)
2. rank_mode: 'full' (CDF-based) or 'conditional' ranking
3. rank_penalty: 'softplus' or 'exp' for the ranking loss formulation

from kmnet.model import KMLoss

loss = KMLoss(
    alpha=0.5,          # Weight for NLL vs Rank loss
    sigma=0.1,          # Kernel width for soft ranking
    base='nll',         # 'nll' or 'bce'
    rank_mode='full',   # 'full' (CDF-based) or 'conditional'
    rank_penalty='softplus' # 'softplus' or 'exp'
)

model = KMNet(net, loss=loss, duration_index=labtrans.cuts)

💬 Getting Help

• 📖 Documentation: yuvrajiro.github.io/KMNet (quickstart, method, API reference)
• 💡 Issues: Found a bug? Open an issue
• 💬 Discussions: Have questions? Start a discussion
• 📧 Email: For collaboration inquiries: yuvrajiro@gmail.com

🤝 Contributing

We welcome contributions! KMNet is an open-source research project.

📄 License

Distributed under the MIT License. See LICENSE for more information.

📚 Citation (Under Consideration)

The paper is under review. Until it is published, please cite the software:

@software{goswami_kmnet,
  title  = {KMNet: Discrete-Time Survival Analysis with a Kaplan-Meier Inspired Rank Loss},
  author = {Goswami, Rahul},
  url    = {https://github.com/yuvrajiro/KMNet},
  year   = {2025}
}

🏆 Acknowledgments

KMNet builds upon excellent work from:

• PyCox for survival analysis utilities
• TorchTuples for training infrastructure
• The PyTorch team for JIT compilation capabilities

_{Built with ❤️ for the survival analysis community
⭐ Star us on GitHub if KMNet helps your research!}