🤖 AI

SurvDiff: A Diffusion Model for Generating Synthetic Data in Survival Analysis

This paper introduces SurvDiff, the first end-to-end diffusion model specifically designed to generate synthetic survival data by jointly modeling mixed-type covariates, event times, and censoring mechanisms to ensure high distributional fidelity and downstream task performance.

Original authors: Marie Brockschmidt, Maresa Schröder, Stefan Feuerriegel

Published 2026-07-27

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Marie Brockschmidt, Maresa Schröder, Stefan Feuerriegel

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a mystery about how long people stay healthy before a specific event happens, like a disease returning or a treatment stopping working. In the world of medical research, this is called survival analysis. It's not just about counting how many people are alive; it's about timing. But here's the tricky part: often, the story doesn't have a clear ending. Some patients stop showing up for checkups, some move away, or some are still doing well when the study ends. In the detective world, we call this "censoring"—it's like a witness who disappears before they can tell you the whole truth. Because of this missing information, the data looks messy and incomplete.

Now, imagine you want to train a super-smart AI to predict these outcomes, but you can't use real patient data because of privacy laws. You need to create fake, "synthetic" data that looks and acts exactly like the real thing. This is where things get really hard. If you just make up random numbers, your AI will learn the wrong lessons. It needs to understand not just the patient's age or weight, but also the timing of events and the reasons why some stories end abruptly (censoring). If the fake data gets the timing wrong, the AI will be useless in the real world. This paper steps into that messy, high-stakes corner of science to build a better tool for creating these fake stories.

The Paper: SURVDIFF

The authors of this paper, Marie Brockschmidt and her team, are introducing a new tool called SURVDIFF. Think of it as a master forger who doesn't just copy a painting; they understand the brushstrokes, the drying time of the paint, and even the cracks in the canvas.

The Problem with Old Tools
Previously, scientists tried to make fake survival data using tools like GANs (Generative Adversarial Networks). Imagine GANs as two robots playing a game: one tries to draw a fake picture, and the other tries to spot the forgery. They get better at the game, but they often crash and burn, getting stuck in a loop where they only draw the same boring picture over and over (a problem called "mode collapse"). Other methods tried to build the fake data in separate steps—first making up the patient's age, then making up the time of the event, then deciding if the event was "censored." The authors argue this is like building a car by assembling the engine, the wheels, and the seats in three different factories and hoping they fit together perfectly. It often leads to errors and a car that doesn't run right.

The New Solution: A Diffusion Model
SURVDIFF uses a different approach called a diffusion model. To understand this, imagine a cup of clear coffee.

The Forward Process (The Mess): You slowly pour in milk, then more milk, then more, until the coffee is completely white and cloudy. You've added "noise" until the original pattern is gone.
The Reverse Process (The Magic): Now, imagine you have a super-smart AI that knows exactly how to un-mix the milk. It looks at the cloudy cup and slowly removes the milk, step by step, until the coffee is clear again.

Most diffusion models are great at making up pictures or standard lists of numbers. But they don't know how to handle the "censoring" mystery. If you just feed them survival data, they might forget that some patients disappeared before the event happened, or they might mess up the timing.

What SURVDIFF Does Differently
SURVDIFF is the first tool designed specifically to handle this "censoring" puzzle from start to finish. Instead of building the fake data in separate steps, it generates everything at once: the patient's details (like age), the time until an event, and whether that event was actually observed or if the patient was "censored" (disappeared).

The secret sauce is a special loss function (a set of rules the AI follows to learn). The authors designed this rulebook to pay extra attention to the timing. They realized that in real medical data, early events are common and easy to see, but late events are rare and often get lost in the noise. So, SURVDIFF's rules tell the AI: "Don't worry too much about the rare, long-term events that are hard to find; focus on getting the common, early patterns right." This keeps the training stable and ensures the fake data looks realistic.

What They Found
The team tested SURVDIFF on three real-world medical datasets: one about HIV/AIDS patients, one about breast cancer patients, and a large international breast cancer study. They compared their new tool against the best existing methods, including the old GANs and other diffusion models.

The results suggest that SURVDIFF is a strong contender.

Better Fake Data: The synthetic patients it created had characteristics (like age and tumor size) that matched the real patients much better than the other tools.
Better Timing: The fake data preserved the correct timing of events and the correct pattern of who got "censored."
Better AI Training: When they used the fake data to train a new survival model and then tested that model on real patients, the model performed very well. In fact, on datasets with a lot of missing data (high censoring), SURVDIFF helped the AI learn better than any other method they tried.

The Bottom Line
The paper suggests that SURVDIFF is a significant step forward because it is the first "end-to-end" diffusion model built specifically for survival data. It doesn't just copy numbers; it learns the complex dance between patient details, time, and missing information. While the authors note that their method works best with the specific type of "right-censoring" found in most medical studies, they show that it can generate high-quality synthetic data that helps researchers train better AI without needing to touch sensitive real-world patient records. It's a promising new way to keep medical research moving forward while keeping patient privacy intact.

Technical Summary: SURVDIFF

Problem Statement

Survival analysis is a fundamental tool in clinical research for modeling time-to-event outcomes, such as disease relapse or patient death. Unlike standard tabular data, survival data is characterized by right-censoring, where the event of interest is not observed for some individuals due to dropout, loss to follow-up, or study termination. This creates a unique challenge for synthetic data generation: a model must not only reproduce the distribution of patient covariates but also faithfully capture the time-to-event distribution and the censoring mechanism.

Existing generative models for tabular data (e.g., GANs, VAEs, standard diffusion models) fail in this context because they are not designed to handle the interplay between covariates, event times, and censoring indicators. Naïve applications of these models often result in biased downstream clinical results. Furthermore, existing methods tailored for survival data, such as SurvivalGAN and the Ashhad framework, suffer from significant limitations:

Instability: SurvivalGAN relies on a GAN backbone, which is prone to mode collapse and unstable training.
Decomposition: Both existing methods decompose the generation task into separate stages (covariates vs. event times/censoring) rather than learning a single joint model. This multi-stage pipeline is prone to error propagation and prevents end-to-end optimization.

To date, there has been no end-to-end diffusion model explicitly designed to generate synthetic survival data while preserving censoring mechanisms.

Methodology: SURVDIFF

The authors propose SURVDIFF, a novel end-to-end diffusion model specifically tailored for generating synthetic survival data. SURVDIFF jointly generates mixed-type covariates, event times, and right-censoring indicators.

Architecture and Process

SURVDIFF follows a standard diffusion framework adapted for mixed-type data, comprising three main components:

Forward Diffusion Process:
- Continuous Variables: For continuous covariates and event times, the model employs a Variance-Exploding (VE) Stochastic Differential Equation (SDE), perturbing data with Gaussian noise.
- Discrete Variables: For discrete covariates and the binary event indicator, the model uses a masked diffusion process. Discrete values are progressively replaced by a "mask" state with a scheduled probability, forcing the model to learn the original values during the reverse process.
Reverse Diffusion Process:
- A neural network (parameterized by $\theta$ ) learns to denoise the data iteratively, reconstructing the original covariates, event times, and indicators from pure noise.
- The network outputs predictions for both the continuous/denoised features and the discrete/masked features.
Survival-Tailored Loss Function:
The core innovation of SURVDIFF is a custom loss function that guides the generative process to respect survival-specific structures. The total loss is a weighted sum of the diffusion loss and a survival loss:
$\mathcal{L}_{total}(\theta) = \mathcal{L}_{diff}(\theta) + \lambda_{surv}\mathcal{L}_{surv}(\theta)$
- Diffusion Loss ( $\mathcal{L}_{diff}$ ): Standard reconstruction losses for continuous (MSE) and discrete (cross-entropy) variables.
- Survival Loss ( $\mathcal{L}_{surv}$ ): An extension of the Cox partial negative log-likelihood. It maps denoised covariates to a risk score and optimizes the ranking of event times.
- Sparsity-Aware Weighting: A critical design choice addresses the issue of sparse risk sets at late time points. The authors introduce a weighting scheme $w_i$ that downweights rare, long-duration events (where $t_i > \tau$ ) using an exponential decay. This stabilizes training by preventing the model from being dominated by unstable gradients from small risk sets at the tail of the distribution.
The weight $\lambda_{surv}$ is adaptively calibrated during a warm-up period to ensure the survival term contributes a target fraction of the total objective, balancing the scales of the diffusion and survival losses.

Key Contributions

First End-to-End Diffusion Model for Survival Data: SURVDIFF is the first diffusion model explicitly designed to generate synthetic survival data, addressing the gap in existing literature where diffusion models have been limited to general tabular data.
Joint Modeling: Unlike previous approaches that separate covariate generation from event-time generation, SURVDIFF jointly models covariates, event times, and censoring in a single end-to-end framework, avoiding error propagation.
Censoring-Aware Objective: The method introduces a novel loss function that encodes time-to-event structures and explicitly accounts for censoring via a sparsity-aware weighting scheme, ensuring realistic event-time distributions.
State-of-the-Art Performance: Extensive experiments demonstrate that SURVDIFF outperforms existing baselines (including SurvivalGAN, Ashhad, and generic tabular diffusion models like TabDiff) in both distributional fidelity and downstream survival analysis performance.

Experimental Results

The authors evaluated SURVDIFF on three medical datasets: ACTG (AIDS), GBSG2 (Breast Cancer), and METABRIC (Breast Cancer).

Evaluation Metrics

Covariate Fidelity: Measured via Jensen-Shannon (JS) distance and Wasserstein distance. SURVDIFF consistently achieved the lowest JS distances across all datasets and competitive Wasserstein distances, outperforming SurvivalGAN significantly.
Survival-Specific Fidelity: Evaluated using Event-Time Divergence (ETD) and temporal distribution plots. SURVDIFF successfully reproduced the temporal patterns of censored and uncensored events, whereas baselines like TabDiff failed to preserve the censoring mechanism.
Downstream Performance: Synthetic datasets were used to train five different survival models (DeepHit, CoxPH, Weibull AFT, Random Survival Forest, XGBoost), which were then tested on real data.
- SURVDIFF achieved the highest Concordance Index (C-index) and lowest Brier Score on the AIDS and GBSG2 datasets, particularly in settings with strong right-censoring.
- On the METABRIC dataset, SURVDIFF remained competitive with TabDiff in downstream metrics while maintaining superior survival-specific fidelity.

Additional Findings

Small Sample Sensitivity: In sensitivity studies with downsampled datasets, SURVDIFF showed substantial improvements over baselines, highlighting its robustness in data-scarce medical scenarios.
Differential Privacy: The authors demonstrated that SURVDIFF can be extended to a differentially private variant (DP-SURVDIFF) using DP-SGD, outperforming a DP-GAN baseline while maintaining privacy guarantees.

Significance and Claims

The paper claims that SURVDIFF represents a significant advancement in the generation of synthetic clinical data. By faithfully reproducing both patient characteristics and the complex censoring mechanisms inherent in survival data, SURVDIFF enables the creation of high-quality synthetic datasets for:

Augmenting small medical datasets.
Mitigating bias and improving fairness in clinical research.
Promoting data accessibility in low-resource healthcare settings without compromising patient privacy.

The authors emphasize that their method is the first to successfully integrate the unique challenges of survival analysis (time-to-event and censoring) into the stable, high-fidelity framework of diffusion models, offering a reliable alternative to unstable GAN-based approaches and fragmented multi-stage pipelines.