Image by Author | Midjourney & Canva

Introduction

Survival analysis, at its core, is about modeling the time until something happens, whether that’s equipment failure, patient discharge, or customer churn. Time-to-event modeling helps us ask not just what might happen but when it’s likely to happen while also accounting for incomplete information through censoring.

But things start to get tricky when we leave behind small, clean datasets and move into high-dimensional territory. I’m talking about situations where we’re working with hundreds or even thousands of features and maybe only a few hundred samples.

Here’s what makes high-dimensional survival analysis challenging:

• More features than observations: Traditional models like Cox can’t handle this well without overfitting or blowing up.
• Censoring remains an issue: A large chunk of observations is incomplete, which makes things even more fragile.
• Sparsity: Not every feature is relevant, but finding the signal through the noise gets harder as dimensions grow.

In this article, we’ll move past the textbook Cox model and look at how regularized models, tree-based methods, and deep learning can help survival analysis scale to meet real-world complexity.

By the end of this, you’ll have a working understanding of how to approach time-to-event problems in high-dimensional settings and know exactly which tools to reach for when the old ones just don’t cut it.

Case Study: Predicting ICU Patient Outcomes

To make this more concrete, let’s walk through a case study that reflects a real-world survival analysis scenario: predicting outcomes for patients in the ICU.

Imagine we’re working with an anonymized, high-dimensional clinical dataset. Something like MIMIC-III or a synthetic version of it. Each row represents a patient admitted to the ICU. We’ve got time-to-event information like how long they stayed before discharge or, unfortunately, death. But more importantly, we’ve also got hundreds of features—test results, vitals, medication history, diagnostic codes, demographic info—all recorded within the first 24 hours of admission.

For starters, not every patient has the same kind of outcome. Some make a quick recovery, some stay longer, and others pass away. Then there are those who get transferred or leave against medical advice, which introduces censoring. So we’re dealing with:

• High-dimensional feature space
• Censored survival times
• Class imbalance (most people survive, but we still care deeply about those who don’t)

Our goal is to build a model that can predict, based on those early clinical signals, how long a patient is likely to stay in the ICU or survive. Ideally, this kind of modeling helps hospitals with triage, resource allocation, and early interventions.

We’ll treat this like any real project. First, we’ll simulate or load a high-dimensional dataset that mimics this ICU setting. Then, we’ll walk through preprocessing, modeling, and evaluation, starting with a baseline Cox model and moving on to regularized models, tree ensembles, and deep learning.

Preparing the Dataset

I did some digging and found that the SurvSet repository is a solid choice for our high-dimensional, time-to-event modeling needs. It offers 77 survival datasets, many with genomic-level feature counts (so $p \gg n$), and everything is formatted cleanly for quick experimentation. It’s ideal whether you’re simulating big-data challenges or scaling to full-scale.

Why I chose SurvSet:

• It includes real high-dimensional cases, not just toy demos
• Everything is pre-formatted into pandas-friendly tables with clear survival columns
• It is great for testing scalable models because it can naturally contain from a few hundred to tens of thousands of features

We’ll pick one example dataset that mimics our ICU-style setting: lots of clinical or sensor data, sparse signals, and survival outcomes.

Loading the Dataset

First, install the library and load a dataset:

pip install SurvSet

from SurvSet import SurvLoader
import pandas as pd
from sksurv.util import Surv

# Load a high-dimensional dataset (e.g., 'gse1992')
loader = SurvLoader()
df = loader.load_dataset("gse1992") # gene-expression example

# Assume df has columns: features + ['event', 'duration']
df.head()

gse1992 is one of the high-dimensional gene-expression datasets in SurvSet, perfect for our purpose. If it doesn’t fit, another similarly structured dataset will do.

Handling Missing Values

Let’s take care of missing data next:

from sklearn.impute import SimpleImputer

num_cols = df.select_dtypes(include=["float64", "int64"]).columns.drop(["event", "duration"])

num_imputer = SimpleImputer(strategy="median")
df[num_cols] = num_imputer.fit_transform(df[num_cols])

Since gene-expression (or high-dimensional clinical) data is typically numeric, median imputation works well here.

Define Target and Features

Split the dataset into features X and the survival object y:

X = df.drop(columns=["event", "duration"])
y = Surv.from_dataframe("event", "duration", df)

This gives us our high-dimensional X, and y handles censoring and timing out of the box.

Stratified Train-Test Split

We still want a train/test split that respects the proportion of events vs censored instances:

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,
    stratify=df["event"],
    random_state=42
)

Now we have:

• A clean, imputed high-dimensional dataset
• A proper survival target with censoring information
• A stratified train-test split that preserves the event ratio

Baseline Model: Cox Proportional Hazards

Let’s start with the classic baseline: the Cox Proportional Hazards (CoxPH) model. It’s the go-to for survival analysis for good reason: It’s interpretable, fast, and gracefully handles censored data.

But it has one major drawback: it doesn’t play nicely when the number of features starts to exceed the number of samples. That’s exactly the situation we’re dealing with here, and you’ll see pretty quickly where it starts to crack.

Fitting the Model

We’ll use scikit-survival, which is more scalable and sklearn-compatible compared to lifelines. Here’s how to fit a CoxPH model on our high-dimensional data:

from sksurv.linear_model import CoxPHSurvivalAnalysis
from sksurv.metrics import concordance_index_censored

# Fit baseline Cox model
cox = CoxPHSurvivalAnalysis()
cox.fit(X_train, y_train)

# Predict risk scores on test set
risk_scores = cox.predict(X_test)

# Evaluate using Concordance Index (C-index)
cindex = concordance_index_censored(
    y_test["event"], y_test["duration"], risk_scores
)[0]

print(f"C-index on test set: {cindex:.4f}")

If your feature space is already huge (like thousands of gene expression features), you’ll probably get a C-index around 0.5–0.6, or even warnings about convergence or singular matrices. That’s a red flag: it means the model is either overfitting badly or struggling to compute the coefficients at all.

Why Cox Fails When p > n

CoxPH relies on partial likelihood and assumes a fixed set of non-collinear features. But in high-dimensional settings, that assumption breaks:

• There’s more noise than signal
• The design matrix becomes singular or ill-conditioned
• Coefficients become unstable or get regularized implicitly to near-zero

Basically, CoxPH wasn’t built for this kind of data. It gives us a solid starting point, but it can’t scale without help.

Penalized Cox Models: Shrinking Toward Signal

Since our plain Cox model struggles with high dimensionality, it’s time to bring in regularization. By applying a penalty to the model’s coefficients, we can shrink irrelevant features toward zero, reduce overfitting, and uncover the variables that actually matter.

This is where Lasso (L1) and Ridge (L2) regularization come into play. Lasso performs feature selection by setting many coefficients to exactly zero, while Ridge keeps all coefficients but shrinks them down. Both are solid tools when you’re dealing with hundreds or thousands of features.

We’ll use scikit-survival, which has built-in support for both.

Fitting a Cox Model with L1 (Lasso) Regularization

Let’s start with L1, which is more aggressive about dropping irrelevant features:

from sksurv.linear_model import CoxnetSurvivalAnalysis
import numpy as np

# Alpha controls strength of regularization
alphas = 10 ** np.linspace(-4, 4, 50)

lasso_model = CoxnetSurvivalAnalysis(l1_ratio=1.0, alphas=alphas)
lasso_model.fit(X_train, y_train)

Once the model is fit, we can select the best-performing set of coefficients. We’ll use the one closest to the minimum deviance:

# Extract coefficients for the best alpha
best_index = np.argmax(lasso_model.score(X_test, y_test))
coefs = lasso_model.coef_[best_index]
feature_names = X_train.columns

# Filter to non-zero features
non_zero = coefs != 0
selected_features = feature_names[non_zero]
selected_coefs = coefs[non_zero]

Now we can plot the non-zero coefficients to see which features the model considered relevant.

Plot: Selected Feature Coefficients (Lasso)

import matplotlib.pyplot as plt

plt.figure(figsize=(10, 6))
plt.barh(selected_features, selected_coefs, color="teal")
plt.title("Selected Features from Lasso-Penalized Cox Model")
plt.xlabel("Coefficient")
plt.ylabel("Feature")
plt.grid(axis="x")
plt.tight_layout()
plt.show()

This gives you a clear sense of which features are driving the survival prediction, and how strongly.

So far, we’ve moved from a failing baseline model to a more stable and interpretable version that can survive in high dimensions. But even regularized linear models might not fully capture non-linear effects or complex interactions. That’s where survival forests come in, and we’re heading there next.

Survival Forests: Embracing Nonlinearity

Now that we’ve explored linear models, it’s time to go nonlinear. Tree-based methods, like Random Survival Forests, are especially useful when your data has complex interactions or nonlinear relationships that linear models can’t capture.

The great thing about Random Survival Forests is that they handle:

• High-dimensional feature spaces
• Missing data (often with less preprocessing)
• Censoring
• Nonlinear effects and interactions

They also give us feature importance out of the box, which helps with interpretation even if the model itself is less transparent than Cox. In scikit-survival, this comes neatly packaged as RandomSurvivalForest.

Fitting a Random Survival Forest

from sksurv.ensemble import RandomSurvivalForest

rsf = RandomSurvivalForest(
    n_estimators=100,
    min_samples_split=10,
    min_samples_leaf=15,
    max_features="sqrt",
    n_jobs=-1,
    random_state=42
)

rsf.fit(X_train, y_train)

Once trained, we can evaluate performance just like before using the Concordance Index:

from sksurv.metrics import concordance_index_censored

preds = rsf.predict(X_test)
c_index = concordance_index_censored(
    y_test["event"], y_test["duration"], preds
)[0]

print(f"Random Survival Forest C-index: {c_index:.4f}")

You’ll likely see an improvement over the Cox models, especially in datasets where nonlinearity is strong or features interact in unexpected ways.

Feature Importance Plot

One of the benefits of survival forests is that they can rank features by how much they contribute to reducing error across the ensemble. Let’s visualize the top 20 features:

import matplotlib.pyplot as plt
import numpy as np

importances = rsf.feature_importances_
sorted_idx = np.argsort(importances)[-20:]
plt.figure(figsize=(10, 6))
plt.barh(X_train.columns[sorted_idx], importances[sorted_idx], color="salmon")
plt.title("Top 20 Feature Importances (Random Survival Forest)")
plt.xlabel("Importance")
plt.ylabel("Feature")
plt.grid(axis="x")
plt.tight_layout()
plt.show()

Compared to penalized Cox, survival forests require less tuning to work well and tend to scale nicely with larger datasets. However, they don’t give you hazard ratios or coefficients, so you trade a bit of interpretability for power and flexibility.

Next, we’ll take things a step further and explore how neural networks can be adapted for survival problems using libraries like pycox.

Deep Learning for Survival

So far, we’ve used statistical models and tree ensembles to model survival data. But when you have large datasets with high-dimensional features and potentially complex patterns, deep learning becomes a serious option.

In survival analysis, DeepSurv is one of the most widely used deep learning architectures. It’s basically a feedforward neural network that predicts risk scores using the Cox partial likelihood. If you need a more flexible, non-Cox approach, DeepHit is another powerful model that predicts the full survival function directly.

We’ll use the excellent pycox library, which wraps everything you need to build DeepSurv and DeepHit, and uses torchtuples under the hood for training.

Installing pycox

If you don’t already have it installed:

pip install pycox torchtuples

Building and Training a DeepSurv Model

Let’s start by defining a basic DeepSurv model. We’ll assume your X_train and X_test are still high-dimensional DataFrames.

import torchtuples as tt
from pycox.models import DeepSurv
from pycox.evaluation import EvalSurv
import torch
import pandas as pd

# Convert to numpy
x_train = X_train.values.astype("float32")
x_test = X_test.values.astype("float32")
durations_train = y_train["duration"]
events_train = y_train["event"]

# Define network architecture
net = torch.nn.Sequential(
    torch.nn.Linear(x_train.shape[1], 64),
    torch.nn.ReLU(),
    torch.nn.BatchNorm1d(64),
    torch.nn.Linear(64, 1)
)

# Optimizer + model wrapper
model = DeepSurv(net, tt.optim.Adam)
model.fit(x_train, (durations_train, events_train), batch_size=128, epochs=100, verbose=False)

Evaluating with Time-Dependent AUC

Once the model is trained, we evaluate its survival prediction quality over time using time-dependent AUC.

from pycox.evaluation import EvalSurv
import numpy as np

# Predict survival function
surv = model.predict_surv_df(x_test)

# Convert y_test to durations and events
durations_test = y_test["duration"]
events_test = y_test["event"]

ev = EvalSurv(surv, durations_test, events_test, censor_surv="km")

# Plot AUC over time
time_grid = np.linspace(durations_test.min(), durations_test.max(), 100)
auc = ev.concordance_td("antolini")

print(f"Time-dependent concordance index: {auc:.4f}")

ev.plot_auc(time_grid)

Model Comparison and Performance Discussion

Now that we’ve explored a range of survival models, from classical Cox to deep learning, it’s worth stepping back to compare them across a few key dimensions.

The most important metrics we’ve used throughout are the Concordance Index (C-index) and Time-dependent AUC. But real-world modeling isn’t just about metrics. Things like interpretability, scalability, and feature selection also matter depending on your use case.

Here’s a side-by-side snapshot of what each model gives you:

A few things stand out here:

• CoxPH is still a great educational starting point. But once dimensionality increases, it struggles unless heavily curated.
• Lasso-penalized Cox balances performance with interpretability. You get better generalization and can actually see which features matter.
• Survival Forests offer great performance and scale well, plus they give you intuitive feature importance. They’re a solid middle ground.
• DeepSurv wins on raw predictive power, especially with nonlinear or high-dimensional data. But it’s a black box. You’ll lose visibility into why predictions are made unless you bring in post hoc explanation tools.

From a resource standpoint, Cox is fast and light. DeepSurv, on the other hand, takes longer to train and tune. Survival forests fall somewhere in between, especially if you’re parallelizing across trees.

In the final section, we’ll wrap this up and discuss when to use each approach based on the data and the context you’re working with.

Conclusion

If there’s one thing to take away from all this, it’s that Cox isn’t always enough. It’s a powerful tool for small to medium datasets where assumptions hold and interpretability matters, but once your feature count explodes or relationships get more complex, you have to reach for something more scalable.

Here’s how I’d think about it:

• Stick with Cox if your dataset is relatively small and clean, and you need something interpretable, fast, and easy to explain.
• Use penalized Cox models (Lasso or Ridge) when you’ve got more features than samples and need built-in feature selection without leaving the linear world.
• Try survival forests if you want strong performance without sacrificing interpretability entirely. They’re solid when nonlinear patterns and interactions exist.
• Go for deep learning (like DeepSurv or DeepHit) if you have the data, compute, and flexibility to explore powerful black-box models. These shine when patterns are complex and accuracy is critical.

Whichever route you choose, the key is to adapt your modeling choices to your data, your goals, and your constraints. There’s no one-size-fits-all in survival analysis, especially when high-dimensionality and censoring are involved.

So test these techniques out on your own data. Push past the usual defaults. And see what insights show up when you let your models scale with the problem you’re really trying to solve.