Searcharxiv⌕ Search

arXiv · 2610.24233

Censored Flow Matching for Synthetic Time-to-Event Generation

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.

Explore related subjects

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Jianbin Tan, Xingran Chen, Kevin He, Zhenke Wu, Bhramar Mukherjee, Anru R. Zhang. 2026-10-16. Censored Flow Matching for Synthetic Time-to-Event Generation. https://arxiv.org/abs/2610.24233

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

Reinforcement Learning for Individual Optimal Policy from Heterogeneous Data

Offline reinforcement learning (RL) aims to find optimal policies in dynamic environments in order to maximize the expected total rewards by leveraging pre-collected data. Learning from heterogeneous data is one of the fundamental challenges in offline RL. Traditional methods focus on learning an optimal policy for all individuals with pre-collected data from a single episode or homogeneous batch episodes, and thus, may result in a suboptimal policy for a heterogeneous population. In this paper, we propose an individualized offline policy optimization framework for heterogeneous time-stationary Markov decision processes (MDPs). The proposed heterogeneous model with individual latent variables enables us to efficiently estimate the individual Q-functions, and our Penalized Pessimistic Personalized Policy Learning (P4L) algorithm guarantees a fast rate on the average regret under a weak partial coverage assumption on behavior policies. In addition, our simulation studies and a real data application demonstrate the superior numerical performance of the proposed method compared with existing methods.

stat.ML↗

Generative Flows for Potential Outcomes and Counterfactuals

Average treatment effects summarize how interventions change outcomes, but individualized decisions require predictions and outcome distributions of what would have happened to the same individual under an alternative treatment. We develop PO-Flow, a continuous normalizing flow framework that represents treatment-specific outcome distributions through invertible transport maps. It supports potential-outcome generation, conditional average treatment-effect estimation, factual-conditioned counterfactual prediction, and density evaluation. To address treatment imbalance in observational data, we incorporate cross-fitted inverse-probability weighting into the flow-matching loss. We establish sufficient conditions under which encoding a factual outcome and decoding under an alternative treatment recovers the individual counterfactual. Under shared exogenous noise and a common encoded reference distribution, scalar recovery follows from increasing noise-to-latent maps. For multivariate outcomes, we assume triangularity or a convex-gradient structure. These conditions distinguish fitting outcome distributions from recovering their causal correspondence. Experiments on four scalar benchmarks and seven multivariate or image settings, with dimensions up to 784, compare against ten scalar and five multivariate baselines. Controlled ablations show IPW improves all reported scalar counterfactual RMSE, PEHE, and potential-outcome RMSE. PO-Flow achieves the lowest PEHE and potential-outcome RMSE on all four scalar benchmarks and attains or ties the lowest mean in 26 of 35 multivariate dataset--metric comparisons.

stat.ML↗

Hybrid Topic-Semantic Labeling and Graph Embeddings for Unsupervised Legal Document Clustering

Legal documents pose substantial challenges for unsupervised knowledge discovery due to specialized domain terminology, multi-scale document lengths, and scarce expert annotations. This paper introduces Topic-Graph Spherical Fusion (TGSF), a dual-manifold representation learning framework for unsupervised legal document clustering that unifies dense topic-semantic discovery and relational graph embeddings on a joint unit hypersphere. Instantiating dense topic discovery via Top2Vec and biased second-order random walks via Node2Vec, TGSF constructs a connected multi-relational document-topic graph combining multi-topic centroid bridges with temperature-sharpened semantic k-Nearest Neighbor (k-NN) document edges, guaranteeing a single connected component. The resulting L2-normalized semantic and graph representations are fused via modality-weighted concatenation onto a joint unit hypersphere prior to KMeans partitioning, decomposing squared chordal distance into a convex combination of semantic and structural cosine similarities. Evaluations across three legal benchmarks (ACORD, CUAD, and LEDGAR) show that TGSF outperforms semantic-only, structural-only, Latent Dirichlet Allocation (LDA), Non-Negative Matrix Factorization (NMF), and TF-IDF baselines. Across 216 factorial grid configurations and 240 Bayesian TPE and Multi-Objective NSGA-II Pareto optimization trials per dataset (720 trials total), TGSF improves Normalized Mutual Information (NMI) from 0.606 to 0.653 on LEDGAR (ARI = 0.425), 0.228 to 0.284 on ACORD (ARI = 0.260), and 0.535 to 0.570 on CUAD (ARI = 0.282). Functional ANOVA and Pareto frontier diagnostics further characterize hyperparameter sensitivity for reproducible legal text mining.

stat.ML↗