SDPM: Diffusion Models Meet Survival Analysis Without Time Discretization

Paper: 2605.22776 Authors: Stanislav R. Kirpichenko, Andrei V. Konstantinov, Lev V. Utkin Categories: cs.LG, cs.AI, stat.CO, stat.ML

The Gap

Survival analysis predicts when events happen from censored data (you know some patients left the study, but not when they’d have the event). Existing methods hit a dilemma: parametric models (Cox proportional hazards, Weibull) impose rigid functional forms on the hazard rate, while flexible neural approaches (DeepSurv, DeepHit) discretize continuous time into bins, introducing approximation errors and losing resolution.

The core tension: you want distributional flexibility without sacrificing the continuous nature of time. Tree-based methods (Random Survival Forests) and recent neural models either assume specific hazard shapes or chop time into intervals. No one’s treating survival prediction as a generative modeling problem where you directly sample from P(T, δ | x) in continuous time.

Problem: Predict time-to-event from censored data
    |
    v
Existing approaches split:
    |
    +---> Parametric (Cox, Weibull)
    |         |
    |         v
    |     Rigid functional form
    |     (can't capture complex patterns)
    |
    +---> Flexible neural (DeepHit, MTLR)
              |
              v
          Time discretization
          (approximation errors, resolution loss)
    
Gap: No continuous-time generative model
    |
    v
SDPM: Diffusion model on (log T, δ) space
    |
    v
Evidence: Competitive C-index/AUC/Brier + better distribution recovery
    |
    v
Conclusion: Generative modeling viable for survival analysis

The Increment

One sentence: Before, survival models either assumed parametric forms or discretized time; after, you can generate continuous survival outcomes directly using diffusion models.

Core Mechanism

SDPM treats survival prediction as conditional generation: given covariates x, generate samples of (T, δ) where T is event time and δ is the censoring indicator (1 = event observed, 0 = censored). The model uses a denoising diffusion probabilistic model (DDPM) that learns to reverse a noise process.

The key architectural choice: operate in a transformed space. Raw event times T get log-transformed and standardized to z-scores. The binary censoring indicator δ becomes a continuous variable through a Gaussian mixture representation (δ = 0 maps to one Gaussian mode, δ = 1 to another). This creates a smooth target space where the diffusion process can operate effectively.

During training, the model learns to denoise corrupted (log T, δ) pairs conditioned on x. At inference, it starts from pure noise and iteratively denoises to generate plausible survival outcomes. Multiple samples per patient give you a distribution, which you feed into the Kaplan-Meier estimator to get survival curves—this step assumes conditionally independent censoring (censoring doesn’t depend on unobserved event times given x).

Training:
  Input: (x, T, δ)
      |
      v
  Transform: log T -> z-score, δ -> Gaussian mixture
      |
      v
  Add noise at timestep t: (z_t, δ_t)
      |
      v
  Neural net predicts: noise(z_t, δ_t | x, t)
      |
      v
  Loss: MSE between predicted and actual noise

Inference:
  Input: x, random noise (z_T, δ_T)
      |
      v
  Iterative denoising: T steps
      |
      v
  Output: clean (z_0, δ_0)
      |
      v
  Inverse transform: z_0 -> T
      |
      v
  Repeat N times -> \{(T_i, δ_i)\}_i=1^N
      |
      v
  Kaplan-Meier on samples -> S(t|x)

Think of SDPM as a sculptor working with clay. Traditional survival models are like using cookie cutters (parametric) or building with LEGO blocks (discretized time)—you’re constrained by predefined shapes. SDPM is like having a lump of clay (noise) and gradually sculpting it into the final form (survival outcome) by removing material (denoising). The sculptor’s skill (trained neural network) knows what the final sculpture should look like given the subject (covariates x). The clay’s consistency matters—raw event times are lumpy (heavy-tailed, skewed), so you knead them into a workable form (log-transform, standardize) before sculpting. The censoring indicator is like a binary switch, but you melt it into a continuous dial (Gaussian mixture) so the sculptor can work smoothly. Each sculpture (sample) is slightly different, but together they reveal the true distribution of possible outcomes.

Key Concepts

• Denoising Diffusion Probabilistic Models (DDPM): Imagine teaching a student to draw by showing them progressively blurrier versions of a picture, then asking them to reverse the process. DDPM does this with data: during training, you gradually add Gaussian noise to real data over T steps until it’s pure noise. The model learns to predict the noise at each step. At inference, you start with random noise and iteratively remove predicted noise, reconstructing a plausible data sample. For survival analysis, “data” is (log T, δ) pairs, and “plausible” means consistent with the patient’s covariates x. The magic: by learning to denoise, the model implicitly learns the data distribution without parametric assumptions.

• Conditionally Independent Censoring: This is the survival analysis equivalent of “missing at random.” It means that whether a patient drops out (gets censored) doesn’t depend on their unobserved event time, once you account for their covariates x. Concretely: if two patients with identical x have different censoring times, the one censored earlier isn’t systematically more or less likely to have an event soon. This assumption lets you use the Kaplan-Meier estimator on generated samples—without it, censored samples would bias your survival curve estimates. Most survival methods make this assumption; SDPM inherits it.

• Target Space Transformation: Raw survival data is messy: event times span orders of magnitude (days to years), and the censoring indicator is binary. Neural networks struggle with such heterogeneous, skewed inputs. SDPM transforms the target: log(T) compresses the time range and symmetrizes the distribution; z-score standardization centers it around zero. The censoring indicator δ ∈ {0,1} becomes a continuous variable via a Gaussian mixture: δ = 0 samples from N(μ_0, σ²), δ = 1 from N(μ_1, σ²). This creates a smooth, well-behaved space where diffusion can operate. After generation, you reverse the transform: threshold the continuous δ back to binary, exponentiate z-scores back to times. The paper shows this transformation is critical—without it, the model generates invalid times and poor calibration.

Framework Shift

Before (mainstream approach):          After (this paper):

Covariates x                           Covariates x
    |                                      |
    v                                      v
Model: f(x) -> hazard λ(t|x)          Model: p(noise | x, t)
    |                                      |
    v                                      v
Integrate hazard                       Iterative denoising
    |                                      |
    v                                      v
Survival function S(t|x)               Samples \{(T_i, δ_i)\}
                                           |
[Single deterministic curve]              v
[Parametric or discretized]            Kaplan-Meier estimator
                                           |
                                           v
                                       Survival function S(t|x)
                                       
                                       [Distribution from samples]
                                       [Continuous, nonparametric]

From hazard-centric modeling to outcome-centric generation, the core shift is treating survival prediction as sampling from a learned distribution rather than computing a functional form.

Expert Assessment

Problem choice: Real gap. Survival analysis has been stuck between parametric rigidity and discretization artifacts for years. Bringing generative modeling (specifically diffusion) to this domain is a natural next step given diffusion’s success in images and tabular data, but no one had worked out the details for censored data. The problem sits at the intersection of two mature fields, which is often where interesting work happens.

Method maturity: Solid engineering, not a fundamental breakthrough. The core insight—transform the target space and apply standard DDPM—is straightforward once stated. The Gaussian mixture trick for the censoring indicator is clever but not deep. The reliance on Kaplan-Meier for final estimates feels like a patch (you’re using a generative model but then falling back to a classical estimator). A more mature approach might model the survival function directly or use a different aggregation method. That said, the execution is competent and the ablations are thorough.

Experimental integrity: Baselines are fair (Random Survival Forests, DeepHit, Cox, etc.), and the datasets are standard benchmarks. The synthetic Cox-Weibull experiment is a nice touch—it shows SDPM can recover ground truth distributions better than RSF when you generate enough samples. However, the gains over baselines are modest (often within error bars), and the paper doesn’t discuss computational cost. Diffusion models are expensive at inference (multiple denoising steps × multiple samples per patient). For a clinical deployment, you’d want to know if the marginal performance gain justifies 10-100× slower predictions. Also, the paper claims “competitive” performance, which is code for “we don’t consistently beat the best baseline.” That’s honest but not exciting.

Writing quality: The paper is dense and assumes familiarity with both survival analysis and diffusion models. The related work section is thorough but reads like a literature dump. The method section would benefit from a clearer motivation for each design choice (why Gaussian mixture for δ? why log-transform T?). The ablation study (Section 4.4) is the strongest part—it systematically validates each transformation. The results section buries the lead: the key finding (SDPM recovers distributions better on synthetic data) is in a subsection, while the main tables show “competitive” performance. Rewriting the results to lead with the distribution recovery story would elevate the paper.

Verdict: Weak accept — competent application of diffusion to a real problem with thorough ablations, but modest empirical gains and no conceptual breakthrough. Useful for practitioners exploring generative survival models, not a landmark paper.

Takeaways

Steal the target transformation recipe: If you’re applying neural models to survival data, the log-transform + z-score + Gaussian mixture trick is immediately useful. It’s not specific to diffusion—any neural survival model would benefit from operating in this transformed space. The ablation study proves it improves calibration and reduces invalid outputs.

Generative models for uncertainty quantification: The key advantage of SDPM over discriminative models isn’t point prediction (C-index is similar), it’s that you get a distribution of outcomes per patient. If you need confidence intervals or risk stratification beyond a single survival curve, generating samples is a principled way to do it. This idea transfers to any regression problem with censored or missing data.

When to use diffusion for tabular data: SDPM works because survival outcomes are low-dimensional (2D: time + censoring) and the transformation makes them smooth. If your tabular problem has similar structure (low-dimensional targets, heavy-tailed distributions that benefit from log-transform), diffusion might be worth trying. But if your targets are already well-behaved or high-dimensional, simpler models (boosting, neural nets) will likely outperform with less computational cost.

The Kaplan-Meier crutch: SDPM generates samples but then uses a classical estimator to aggregate them. This is pragmatic but suggests the generative model isn’t fully trusted. If you’re building a generative model for a structured prediction task, think carefully about whether you’re truly modeling the target distribution or just using generation as a feature engineering step. Sometimes the latter is fine, but it’s worth being explicit about where the model ends and the post-processing begins.