PaperlayerView exact revision on arXiv

Source-linked AI summary

MedFlow: Class-Aware Multi-Scale Generation for Medical Time-Series Synthesis

Yanhao Huang, Shibo Feng, Wanjin Feng, Peilin Zhao, Chunyan Miao

arXiv:2609.04804v1cs.AI

TL;DR

MedFlow targets the limited downstream utility of synthetic medical time series when temporal patterns occur at multiple scales and minority classes are underrepresented. It combines multi-scale vector-quantized representations with class-aware flow matching and Token Marginal Guidance. Across four public EHR, EEG, and ECG datasets, it reports improved utility, fidelity, and sampling efficiency over strong baselines.

  • Problem

    Existing synthesis methods may match aggregate distributions and temporal dynamics without preserving multi-scale clinical patterns and minority-class characteristics needed for downstream prediction.

  • Method

    MedFlow uses a vector-quantized multi-scale tokenizer, class-conditioned flows in token space, and Token Marginal Guidance based on class-conditional token statistics.

  • Results

    Across four public EHR, EEG, and ECG datasets, MedFlow consistently improves synthetic-data utility, fidelity, and efficiency over strong generative baselines, including 5.8% higher average AUPRC and 88.6% lower Context-FID on EHR benchmarks.

  • Takeaways & Limitations

    Class-aware generation in structured multi-scale token space is a promising direction for improving synthetic medical time-series utility and fidelity under clinical data imbalance.

  • Takeaways & Limitations

    MedFlow currently focuses on structured medical time series conditioned on class labels and does not incorporate unstructured clinical text or other complementary modalities.

Abstract

from arXiv · show

Synthetic medical time-series generation can alleviate data scarcity and support the development of reliable clinical prediction models. However, existing methods mainly focus on matching the overall distribution and temporal dynamics of real data, which does not necessarily ensure strong downstream utility on imbalanced medical datasets. Clinically informative patterns often occur at heterogeneous temporal scales, while rare minority-class characteristics can be obscured by dominant population patterns. To address these challenges, we propose MedFlow, a class-aware multi-scale flow matching framework for medical time-series synthesis. MedFlow employs a vector-quantized multi-scale tokenizer to represent medical sequences at complementary temporal resolutions, capturing both coarse clinical trends and fine-grained dynamics. We further introduce Token Marginal Guidance, which incorporates class-conditional token statistics directly into the flow matching process to steer generation toward class-specific regions of the learned tokens. This mechanism strengthens minority-class patterns, while preserving the global and tail distributions of real data. Experiments on four public datasets covering electronic health records, EEG, and ECG signals demonstrate that MedFlow consistently outperforms recent state-of-the-art diffusion-based baselines across downstream prediction tasks. On average, it improves AUPRC by 5.8%, reduces Context-FID by 88.6%, and achieves 3.8$\times$ higher sampling throughput.

Introduction

MedFlow addresses medical time-series scarcity and the need to preserve both temporal structure and class-specific predictive patterns under imbalance. It combines multi-scale tokenization, class-conditioned flow matching, and Token Marginal Guidance, with gains across utility, fidelity, and efficiency.

  • Medical time-series synthesis can expand data availability when clinical cohorts are costly, inaccessible, or sensitive.
  • Distributional fidelity and downstream utility are jointly essential because synthetic sequences must preserve temporal dependencies and class-specific predictive patterns.
  • Clinically relevant patterns span multiple temporal scales, while class imbalance can cause generators to underrepresent minority-class characteristics.
  • MedFlow decomposes sequences into complementary coarse-to-fine tokens and applies class-conditioned flows to generate multi-scale medical time series.
  • 5.8% average AUPRC improvement and 88.6% Context-FID reduction were reported on EHR benchmarks, alongside faster sampling across four public datasets.
  • Token Marginal Guidance uses class-conditional token statistics to steer generation toward class-relevant regions and improve minority-class preservation.

Methodology

MedFlow learns residual multi-scale token representations and scale-specific flows in normalized codebook-embedding spaces. Class-aware source construction and Token Marginal Guidance steer endpoint generation toward class-consistent token usage before parallel decoding.

  • MedFlow uses a two-stage framework: a residual multi-scale tokenizer followed by scale-specific class-conditional flows.
  • Residual scales encode shorter coarse-structure tokens and longer fine-scale variation, with the tokenizer frozen before flow training.
  • Discrete token indices are transported through continuous normalized codebook-embedding spaces rather than directly through index space.
  • A learned low-rank source bank preserves token-space geometry and generates class-specific source states through within-class interpolation.
  • The flow predicts categorical endpoint tokens and derives the continuous marginal vector field analytically from the endpoint posterior.
  • Token Marginal Guidance adds a fixed class-conditional logit bias from smoothed token frequencies, without trainable guidance parameters or extra endpoint evaluations.
  • Scale-specific flows are integrated concurrently with shared source coupling, then requantized and combined through residual multi-scale decoding.

Experiments

Across EHR and physiological-signal benchmarks, MedFlow improves downstream utility and distributional fidelity while retaining efficient sampling. Ablations and feature-level analyses support the roles of multi-scale tokenization, structured sources, and Token Marginal Guidance.

  • EHR downstream utility: MedFlow achieves the best synthetic-data AUPRC and AUROC across four EHR prediction tasks, improving AUPRC by 5.8% and AUROC by 2.1% over TarDiff on average.The larger AUPRC improvement is particularly relevant to the imbalanced clinical tasks.
  • Physiological-signal utility: MedFlow achieves the best synthetic-data AUPRC and AUROC on both APAVA EEG and PTB ECG, including an 18.45% AUROC improvement over TarDiff on APAVA.On PTB, it retains consistent improvements despite strong performance from existing methods.
  • Global distributional fidelity: MedFlow reduces EHR PS by 8.8% and C-FID by 88.6% on average relative to TarDiff, while obtaining the best DS on three of four tasks.Together with TSTR results, these findings indicate stronger temporal and contextual agreement with real EHR data.
  • Feature-space and tail fidelity: MedFlow achieves the best result on nine of twelve feature-distribution statistics and reduces lower- and upper-tail gaps over TarDiff by 80.5%–96.8%.The statistics cover overall moments, tail quantiles, and positive-class features in MIMIC-III and eICU mortality cohorts.
  • Feature-space and tail fidelity: MedFlow’s SBP estimates at the 1st and 99th percentiles are 77.71 and 174.24 mmHg, versus real values of 77.00 and 178.00 mmHg and TarDiff values of 93.00 and 269.75 mmHg.The smaller deviations, especially in the upper tail, support improved tail calibration.
  • Sampling efficiency: MedFlow improves sampling throughput by 3.8× over TarDiff while remaining competitive with lightweight GAN/VAE-based samplers.The efficiency advantage comes from incorporating class-specific guidance without additional guidance-time computation.
  • Ablation study: Replacing multi-scale tokenization with a single-scale variant sharply lowers AUPRC, while Gaussian sources worsen DS and other utility and fidelity metrics.On MIMIC-III, AUPRC drops from 0.5068 to 0.1753; on APAVA, it drops from 0.6891 to 0.4932.
  • Ablation study: Token Marginal Guidance consistently improves class-conditional utility, although CFG and the no-TMG variant slightly improve some global fidelity metrics on APAVA.MedFlow nevertheless retains the strongest downstream utility in that comparison.

Conclusion

MedFlow combines multi-scale tokenization with Token Marginal Guidance to preserve class-specific patterns under medical data imbalance. Across EHR, EEG, and ECG datasets, it improves downstream utility and distributional agreement with real data.

  • Conclusion: MedFlow combines multi-scale tokenizers with Token Marginal Guidance to preserve class-specific patterns under medical data imbalance.The framework performs class-aware generation in structured token space.
  • Conclusion: MedFlow improves downstream utility and distributional agreement with real data across EHR, EEG, and ECG datasets.The conclusion presents this as the main empirical outcome.
  • Conclusion: Class-conditional token-frequency calibration distinguishes Token Marginal Guidance from conditioning strategies that do not directly calibrate empirical class-wise token frequencies.The guidance uses fixed class-conditional token statistics to address minority-pattern underrepresentation.

B Additional Method Details

The additional method details describe residual multi-scale flow training, source-bank regularization, class-preserving mixing, imbalance-aware training, and parallel class-conditional sampling.

  • Tokenizer optimization: Straight-through estimation sends reconstruction gradients through selected codes to encoders and decoders, while codebooks update by exponential moving averages.The commitment term keeps encoder outputs close to selected frozen-gradient codes.
  • Source-bank regularization: The source-bank regularizer controls global location and scale while matching mini-batch pairwise geometry without reconstructing individual endpoints.The reported configuration uses r = 128 and (λµ, λσ, λstr) = (0.1, 0.1, 10.0).
  • Flow training: Class-preserving interpolation uses within-class permutation and a shared coefficient across scales to smooth source distributions and categorical endpoint targets.The reported settings are pmix = 0.5 and α = 1.0.
  • Class-imbalance handling: MIMIC-III, eICU, and PTB use equal-class mini-batch sampling with square-root inverse-frequency endpoint-loss weights.APAVA retains empirical class sampling and uses an unweighted endpoint loss.
  • Class-imbalance handling: MIMIC-III and eICU mortality training additionally expands the positive-class pool fourfold through within-class mixup, without changing synthetic evaluation cohort composition.This augmentation is disabled for ICU-stay, APAVA, and PTB tasks.
  • Sampling: Sampling advances all token positions and scales in parallel, applies endpoint guidance during integration, then decodes and sums the scale components.Each scale maintains its own token state, codebook, and endpoint model.

C.1 Dataset Statistics and Splits

The appendix reports patient-independent cohort construction, preprocessing, class composition, and evaluation metrics for EHR and physiological-signal tasks.

  • Splits: All dataset splits are subject- or patient-independent, and imputation and normalization use only training-set statistics.Figure 4 computes class composition and imbalance severity from original training splits.
  • EHR cohorts: MIMIC-III uses first-24-hour hourly sequences from seven vital signs, with a subject-level stratified 80/10/10 split.Mortality and prolonged ICU stay are defined from hospital death and length of stay greater than three days, respectively.
  • EHR cohorts: eICU uses three vital signs on a 5-minute grid of 288 time steps, requiring at least 280 observed positions per retained variable.It uses a unique-patient stratified 80/10/10 split.
  • Physiological signals: APAVA contains 16-channel EEG segments of length 256, while PTB contains 15-channel ECG segments of length 288 under subject-independent benchmark protocols.The supplied dataset table identifies class-1 counts and sample shapes after preprocessing.
  • Metrics: AUPRC and AUROC evaluate downstream classification utility, while DS, PS, and C-FID assess synthetic-data fidelity and temporal structure.Higher AUPRC indicates stronger positive-class recovery under imbalance; lower PS and C-FID indicate better fidelity-related performance.
  • Metrics: DS measures real–synthetic distinguishability, PS measures one-step prediction transfer, and C-FID compares TS2Vec embedding distributions.For DS, values closer to zero indicate that real and synthetic sequences are harder to distinguish.

C.4 Evaluation Protocols

Evaluation trains downstream models on synthetic cohorts and tests them on untouched real data, using standardized architectures, stopping rules, and dataset-specific configurations.

  • Evaluation protocol: Under TSTR, evaluators train only on generated training cohorts and test on untouched real test splits, with synthetic cohorts matching real training size and class composition.The Real Data reference trains the same evaluator on the real training split.
  • Configurations: MedFlow uses dataset-specific configurations alongside shared architecture, optimization, and sampling hyperparameters.Tables 8 and 9 summarize these configurations and shared settings.
  • Evaluation protocol: The TimesNet evaluator uses fixed architecture and optimization settings across methods, with early stopping after ten validation epochs without improvement.AUPRC and AUROC are computed from test-set probabilities.

C.5 Architecture and Training Configuration

MedFlow uses three residual scales and specified codebook sizes, with fixed training and evaluation settings. Its guidance comparison contrasts TMG with a label-dropout-trained CFG variant while holding the remaining pipeline components matched.

  • Configuration: MedFlow uses three residual scales with codebook sizes 128/512/512.These settings define the multi-scale token representation used in the experiments.
  • Configuration: All runs use seed 42, and baseline and downstream evaluations use matched data splits.
  • Guidance comparison: The CFG-ready variant uses label dropout probability 0.1, while TMG is disabled and conditional endpoint logits are replaced during sampling.The comparison keeps the tokenizer, structured source, flow architecture, solver, synthetic-cohort size, and downstream evaluator otherwise unchanged.

D.2 Sampling Complexity and Recorded Latency

TMG reduces sampling overhead relative to CFG while retaining strong downstream utility across the evaluated datasets. The ablation evidence also shows that coupled latent–endpoint mixup materially supports utility and fidelity, especially on MIMIC-III.

  • Sampling complexity: CFG evaluates the endpoint model twice per ODE step, whereas TMG uses one conditional evaluation followed by a fixed logit correction.With M solver steps and S token scales, CFG therefore doubles endpoint-model evaluations, while shared overheads prevent wall-clock latency from exactly doubling.
  • Guidance comparison: TMG provides the highest downstream utility across MIMIC-III and APAVA while using half as many endpoint-model evaluations as CFG.
  • Recorded latency: 3.9% and 2.5% are TMG’s latency increases over unguided sampling on MIMIC-III and APAVA, respectively.CFG is 1.60× and 1.45× slower than TMG on the same datasets because it evaluates both conditional and unconditional branches at each solver step.
  • Mixup ablation: Removing coupled latent–endpoint mixup substantially degrades utility and global fidelity on MIMIC-III.On APAVA, several fidelity and discrimination metrics deteriorate even when AUPRC changes only slightly.

E Subject-Aware Privacy Analysis

The privacy analysis evaluates black-box, class-conditional membership inference at both segment and subject levels on APAVA. MedFlow reduces membership distinguishability relative to TarDiff while improving predictive utility under this audit.

  • Audit protocol: The black-box APAVA audit gives attackers labeled synthetic data and preprocessing information but no model parameters or additional generator queries.Training segments are members, held-out test segments are non-members, and validation subjects provide the auxiliary reference distribution.
  • Privacy metric: DOMIAS uses class-conditional density comparisons, with a 16-dimensional PCA projection of log-magnitude Fourier features.The PCA basis uses synthetic and validation-reference features, while KDE bandwidth selection uses a validation-only holdout.
  • Privacy metric: DOM = 0.5 denotes chance-level membership inference, and lower values indicate weaker membership distinguishability.The reported risk is the excess direction-invariant attack AUROC above chance.
  • Subject-aware evaluation: Subject-aware estimates average 1,000 class-balanced draws, each sampling nine segments per participant and aggregating the 90th-percentile attack score.Segment-level estimates use 2,000 subject-cluster bootstrap replicates; the split contains 15 training, four validation, and four test subjects.
  • Results: 0.1899 to 0.1813 and 0.1656 to 0.1239 are MedFlow’s reductions in segment- and subject-level MIR relative to TarDiff.The APAVA results combine improved predictive utility with reduced membership distinguishability under this audit.

F Limitations

MedFlow is currently limited to structured, class-conditioned medical time series and does not incorporate unstructured clinical text or other complementary modalities.

  • Scope boundary: MedFlow currently focuses on structured medical time series conditioned on class labels.It does not yet incorporate clinical notes, reports, or other clinical modalities that could provide additional temporal context.

G Visualization of Synthetic Data

Figures 5–8 compare class-conditional synthetic trajectories across four EHR tasks and seven generators against held-out real data. MedFlow more closely follows real trajectories across both classes, preserving cross-variable relationships, long-horizon trends, and local physiological variation.

  • Baseline comparisons: TimeGAN often produces excessive fluctuations or limited trend variation, whereas TimeVAE and TimeVQ-VAE generate overly smooth sequences.
  • Baseline comparisons: Diffusion-TS and BioDiffusion recover more variation but can exhibit abrupt excursions or extended flat segments.
  • MedFlow comparison: MedFlow more closely follows held-out real trajectories across both classes, preserving cross-variable relationships, long-horizon trends, and local physiological variation.TarDiff better captures nonstationary trends, although some traces remain noisy or show implausible level changes.
  • Figure overview: Figures 5–8 show representative class-conditional trajectories for MIMIC-III and eICU mortality and ICU-stay tasks, comparing seven generators with held-out real data.Rows distinguish task classes, columns compare methods, and colors encode physiological variables.
Loading 2609.04804v1…