Censored Flow Matching for Synthetic Time-to-Event Generation

Jianbin Tan; Xingran Chen; Kevin He; Zhenke Wu; Bhramar Mukherjee; Anru R. Zhang

stat.MLOct 16, 2026

Censored Flow Matching for Synthetic Time-to-Event Generation

Authors: Jianbin Tan, Xingran Chen, Kevin He, Zhenke Wu, Bhramar Mukherjee, Anru R. Zhang

Organizations: Department of Biostatistics & Bioinformatics, Duke University · Department of Biostatistics, University of Michigan · Department of Biostatistics, Yale University · Department of Biostatistics & Bioinformatics and Department of Computer Science, Duke University

Abstract

Generating time-to-event data is increasingly important in survival analysis, but it remains statistically challenging because event times are often censored and thus prevent flexible generative modeling. To fill this gap, we propose a novel method called censored flow matching (CFM), aiming at providing a flexible flow-based generative mechanism for synthesizing time-to-event data as well as learning their corresponding probabilistic densities. Unlike existing methods, CFM does not rely on survival likelihoods to handle censoring, but instead exploits the monotonicity property of differential flows to address censoring via pullback sampling. Under this framework, we implement an alternating procedure between pullback sampling and CFM steps, leading to a flow-based expectation--maximization (EM) algorithm tailored for time-to-event inference. Under mild distributional conditions, we establish the convergence of CFM to the true event-time distribution as the EM iterations proceed. We then apply CFM to generate time-to-event data and estimate their probabilistic densities, where the observed censoring time can be incorporated to provide truncated sampling or truncated density estimation for survival inference. Through extensive simulation studies, we demonstrate the benefits of CFM for complex and large-scale data generation. Finally, we apply CFM to analyze data from the Organ Procurement and Transplantation Network (OPTN) and the Medical Information Mart for Intensive Care IV (MIMIC-IV), showcasing its benefits for synthetic and predictive survival analyses in terms of data fidelity, predictive performance, and empirical privacy characteristics.

Figures & tables

Figure 1: An illustration of the iterative procedure in censored flow matching. The purple, yellow, red, and blue solid curves represent the density of a base variable (e.g., an exponential random variable), the observed time for non-censored subjects $(\Delta=1)$ , the observed time for censored subjects $(\Delta=0)$ , and the imputed time for censored subjects based on the flow matched in the previous iteration, respectively. The flow at each iteration is obtained by matching the interpolation paths from the base to the target mixture represented by the blue (or red) and yellow solid curves. The red dashed curve denotes the latent time density for censored subjects, which is approached by the imputed density as the iterations proceed.

Figure 2: A pipeline for censored flow matching. Panel (A) illustrates the input data including uncensored and right-censored observations. Panel (B) illustrates the CFM step, where pullback transport guides forward transport to generate event times beyond the observed time $Y$ for a censored subject. Panel (C) presents the outputs of CFM for synthetic event time generation and individual survival-time density estimation.

\widetilde{T}_{b}^{(k)}=\Delta_{b}Y_{b}+(1-\Delta_{b})T_{b}^{(k)},\qquad\widetilde{T}_{b,U_{b}}^{(k)}=(1-U_{b})Z_{b}+U_{b}\widetilde{T}_{b}^{(k)}.

Algorithm 1 Censored Flow Matching (CFM)

\widehat{\mathcal{L}}_{\mathrm{CFM},B}^{(k,h)}(\theta):=\frac{1}{B}\sum_{b=1}^{B}\left[v_{\theta}\left(U_{b}^{(h)},T_{U_{b}^{(h)},i_{b}}^{(k,h)},\boldsymbol{X}_{i_{b}}\right)-V_{b}^{(k,h)}\right]^{2}.

Algorithm 2 Censored Flow Matching with Minibatch

Figure 3: Mean averaged conditional Wasserstein errors over 100 simulation replications across simulation settings and sample sizes.

Figure 4: Mean truncated conditional Wasserstein errors across truncation levels for the nonlinear settings, with 95% confidence intervals computed from 100 simulation replications.

Figure 5: Survival curves estimated using the KM method for the real and synthetic data. The solid and dashed black curves represent the survival curves estimated from the real training and testing sets, respectively. For each synthetic method, the shaded band represents the pointwise minimum and maximum of the KM estimates across 200 independently generated synthetic datasets.

Figure 6: Mean squared error (MSE) between the synthetic and real KM curves within subgroups jointly defined by race/ethnicity and recorded gender. For each method, 200 synthetic datasets, each containing 5000 samples, are independently generated. Within each subgroup, we compute the average MSE across the 200 synthetic datasets, where, for each synthetic dataset,

\mathrm{MSE}=\frac{1}{500}\sum_{j=1}^{500}\{\widehat{S}^{\mathrm{syn}}(t_{j})-\widehat{S}^{\mathrm{real}}(t_{j})\}^{2}.

Here,

\widehat{S}^{\mathrm{syn}}(t_{j})

denotes the KM estimate at time

t_{j}

within the corresponding subgroup,

\widehat{S}^{\mathrm{real}}(t_{j})

denotes the corresponding KM estimate from the real testing data within the same subgroup, and

t_{j}

is defined on a uniform grid of 500 time points over

[0,365]

days.

Gender	Asian	Black	Hispanic	White
Female	$-0.004\;(-0.077,0.070)$	$0.001\;(-0.006,0.008)$	$-0.007\;(-0.023,0.009)$	$0.002\;(-0.004,0.008)$
Male	$-0.029\;(-0.061,0.003)$	$0.001\;(-0.002,0.004)$	$-0.008\;(-0.024,0.008)$	$-0.002\;(-0.004,0.001)$
Overall	$9.2\times 10^{-5}\;(-0.0002,0.0004)$

Table 1: Privacy evaluation for CFM based on censoring-adjusted nearest-neighbor distances. Each entry reports the average of

\widehat{D}_{\mathrm{train}}-\widehat{D}_{\mathrm{test}}

and its 95% confidence interval, computed across 50 replications.

Figure 7: Analysis of the MIMIC-IV survival data. (A) Boxplots of the test IBS over 100 random train–test splits for CoxPH, CoxTime, Hazard-Neural, and CFM. The dashed red line indicates the median IBS of CFM. (B) Estimated patient-specific truncated event-time densities for three representative censored patients. For each patient, the dashed vertical line indicates the observed censoring time

Y

in days, the shaded region represents the estimated probability mass over the subsequent 60 days, and the red point indicates the mode of the estimated truncated density.

Synthetic Data Generation

Synthetic data generation creates artificial examples, simulations, and datasets to train, test, or protect machine-learning systems.

See all 609 papers

A Filtered Mixture-of-Generators for Fully Synthetic Survival Training
Jun 30, 2026Niccolò Maria Rizzi, Eugenio Lomurno, Alberto Archetti +1Survival Analysis Synthetic Data Generation
Politecnico di Milano, Milan, Italy
Smooth Flow Matching for Synthesizing Functional Data
Aug 19, 2025Jianbin Tan, Anru R. ZhangFunctional Data Analysis Flow Matching
Department of Biostatistics & Bioinformatics, Duke University, Durham, NC, USA · Department of Computer Science, Duke University, Durham, NC, USA
SDPM: Survival Diffusion Probabilistic Model for Continuous-Time Survival Analysis
May 21, 2026Stanislav R. Kirpichenko, Andrei V. Konstantinov, Lev V. UtkinSurvival Analysis Conditional Diffusion Models
Peter the Great St.Petersburg Polytechnic University

Jun 30, 2026J/K move · Enter open · S save

A Filtered Mixture-of-Generators for Fully Synthetic Survival Training

Niccolò Maria Rizzi, Eugenio Lomurno, Alberto Archetti, Matteo Matteucci

Politecnico di Milano, Milan, Italy

Survival Analysis Synthetic Data Generation Synthetic Data Evaluation Synthetic Data Selection Synthetic Population Generation

Open paper AskPDF arXiv

Survival analysis models time-to-event data, but in clinical settings training data are costly and scarce: events accrue over years of follow-up, cohorts are small, and privacy regulations restrict sharing across institutions. Tabular generative models promise augmentation and privacy-preserving cohort sharing, yet are themselves data-hungry -- on the small cohorts typical of survival analysis, a single generator rarely characterizes the population well enough for downstream models trained on its output to match real-data performance. FoGS (Filtered Mixture-of-Generators for Survival analysis) reframes synthetic-data construction as sample selection rather than generation. A candidate pool is drawn from four architecturally distinct tabular generators, and each sample is scored by an ensemble of seven survival models trained on real data, using proper scoring rules as a per-sample plausibility proxy. A two-level pipeline optimizes, in its outer loop, a selection policy -- generator quotas, scorer weights, a random complement, and stratified balancing on event time and censoring -- against held-out downstream performance, while an inner loop tunes the downstream model (XGBoost-Cox). On 16 public datasets under train-on-synthetic, test-on-real (C-index and IBS,

0

100

scale), FoGS yields mean improvements of

+2.17

in C-index and

+0.67

in IBS, improving both metrics on 9 of 16 datasets and at least one on 13 (one-sided Wilcoxon

p=0.039

and

p=0.035

). It matches or exceeds real-data training on most cohorts, with no significant change in nearest-neighbour privacy margin relative to unfiltered sampling. Sample filtering over a heterogeneous generator pool is thus a viable substitute for real-data training in privacy-restricted clinical settings.

Abstract

Synthetic Data Generation

Explore similar work

A Filtered Mixture-of-Generators for Fully Synthetic Survival Training

Smooth Flow Matching for Synthesizing Functional Data

SDPM: Survival Diffusion Probabilistic Model for Continuous-Time Survival Analysis