PaperlayerView exact revision on arXiv

Source-linked AI summary

Generating and Imputing Tabular Data via Diffusion and Flow-based Gradient-Boosted Trees

Alexia Jolicoeur-Martineau, Kilian Fatras, Tal Kachman

arXiv:2309.09968v3cs.LG

TL;DR

Tabular data can be scarce and contain missing values, complicating generation and imputation. This paper uses XGBoost in diffusion and flow-based models and evaluates them across tabular datasets. The reported results favor the method for generation, while identifying imputation and memory constraints.

  • Problem

    Tabular datasets may be small and contain missing values, and restricting analysis to complete cases can shrink the participant pool and reduce statistical power.

    Source: §1 · p. 1

  • Method

    The paper uses XGBoost instead of neural networks to estimate score functions or vector fields for tabular diffusion and flow models, including training directly on incomplete data.

    Source: §1 · p. 2, §2.1 · p. 3

  • Results

    Forest-Flow is highest-performing across nearly all complete-data generation metrics and performs similarly or better than deep-learning methods for incomplete-data generation.

    Source: §4.1 · p. 7, §4.3 · p. 8

  • Takeaways & Limitations

    The authors conclude that tree-based methods can generate highly realistic tabular data, including when training data contain missing values, without requiring GPUs.

    Source: §5 · p. 9

    MICE-Forest and MissForest remain better for imputation, and GBT memory demand can exceed that of deep-learning methods on large datasets.

    Source: §5 · p. 9

Abstract

from arXiv · show

Tabular data is hard to acquire and is subject to missing values. This paper introduces a novel approach for generating and imputing mixed-type (continuous and categorical) tabular data utilizing score-based diffusion and conditional flow matching. In contrast to prior methods that rely on neural networks to learn the score function or the vector field, we adopt XGBoost, a widely used Gradient-Boosted Tree (GBT) technique. To test our method, we build one of the most extensive benchmarks for tabular data generation and imputation, containing 27 diverse datasets and 9 metrics. Through empirical evaluation across the benchmark, we demonstrate that our approach outperforms deep-learning generation methods in data generation tasks and remains competitive in data imputation. Notably, it can be trained in parallel using CPUs without requiring a GPU. Our Python and R code is available at https://github.com/SamsungSAILMontreal/ForestDiffusion.

1 INTRODUCTION

Tabular datasets face limited sample sizes and missing values, while generative models offer a possible route to expanding datasets and restoring missing entries. The paper introduces XGBoost-based diffusion and flow models for tabular generation and imputation, evaluated on a broad benchmark.

1 INTRODUCTION

  • Small tabular datasets with missing values can lose participants when analysis uses only complete cases, reducing statistical power and potentially introducing bias.Single imputation can also yield biased inferences, while multiple imputation represents uncertainty through several completed datasets.Source: §1 · p. 1
  • Generative models have shown promise for expanding datasets and restoring missing content, motivating their use for limited and incomplete tabular data.The introduction connects tabular imputation to image inpainting, which restores missing image portions.Source: §1 · p. 1
  • The paper introduces diffusion and flow models for tabular generation and imputation that use XGBoost rather than neural networks.The stated contributions include training directly on incomplete data because XGBoost learns splits for missing values.Source: §1 · p. 2
  • The benchmark covers 27 real-world datasets and evaluates generation and imputation across distributional closeness, diversity, prediction, and statistical inference.The authors also report highly realistic synthetic data when training data are clean or contain missing values.Source: §1 · p. 2

2 BACKGROUND

The background presents gradient-boosted trees as tabular-data estimators that can handle missing values, then reviews diffusion and conditional flow matching as generative frameworks. It outlines the paper’s use of I-CFM and tree-based vector-field estimation.

2.1 Gradient-boosted trees and XGBoost

  • XGBoost is an open-source gradient-boosted tree method whose sparsity-aware split finding naturally handles missing values.The paper selects it because its ablation found it performed best.Source: §2.1 · p. 3

2.2 Generative diffusion and conditional flow matching models

  • Diffusion and flow models require only one model, unlike GANs and VAEs, which require differentiating through two models and thus exclude non-differentiable GBTs.The paper notes that deterministic flow models are not suited to inpainting, whereas REPAINT is designed for diffusion models.Source: §2.2 · p. 3
  • Diffusion models learn the score function of noise-perturbed data and generate samples by solving a reverse stochastic differential equation from Gaussian noise.The forward process takes real data toward Gaussian noise over time.Source: §2.2 · p. 3
  • Conditional flow matching estimates a vector field for an ordinary differential equation, avoiding direct regression against an intractable marginal vector field.The framework expresses the probability path as a mixture of conditional paths and permits regression against the conditional vector field.Source: §2.2 · p. 3
  • Forest-Flow duplicates observations, adds distinct noise to each copy, interpolates across time levels, and fits a GBT at each noise level in parallel over CPUs.The construction provides a single precomputed dataset for GBT training, which does not use stochastic gradient descent.Source: Figure 2 · p. 4
  • The paper uses I-CFM with independent noise-data pairs, a Gaussian conditional path, and a vector field equal to the difference between paired data and noise.After training, samples are generated by solving the ODE from t=1 to t=0.Source: §2.2 · pp. 3–4

3 TRAINING DIFFUSION & FLOW MODELS WITH XGBOOST

The method uses XGBoost gradient-boosted trees to estimate score functions or vector fields for diffusion and flow models. It adapts training and imputation to tree-based learning, including separate models by noise level and XGBoost’s handling of missing values.

3 TRAINING DIFFUSION & FLOW MODELS WITH XGBOOST

  • GBT estimators replace neural networks for score-function or vector-field estimation; among tested tree methods, XGBoost performed best and can handle missing data.The authors report comparing Random Forests, XGBoost, LightGBM, and CatBoost.Source: §3 · p. 4, §3.3 · p. 5

3.1 Duplicating the dataset to estimate the expectation in diffusion and flow losses

  • Duplicating each observation nnoise times creates fixed noisy inputs and corresponding targets that approximate expectations over data-noise pairs.They use nnoise = 100 unless memory constraints require reducing it to 50.Source: §3.1 · p. 4

3.2 Training different models per noise level

  • Separate GBTs per noise level avoid weak time-feature splits and train in parallel; nt = 50 is competitive, while more levels increase computational cost.For the illustrated case, each noise-level model has 100 trees and nt = 4.Source: §3.2 · p. 5, Figure 3 · p. 5

3.4 Forward Diffusion/Flow process

  • The method uses the Variance Preserving process for diffusion and deterministic independent-coupling conditional flow matching for CFM.Source: §3.4 · p. 5

3.5 Imputation via diffusion and XGBoost

  • For imputation, REPAINT applies diffusion inpainting while XGBoost trains on incomplete data by learning split directions for missing values.The authors state that deterministic flow models are not suited to inpainting and that they know of no CFM inpainting method.Source: §3.5 · pp. 5–6

3.6 Data processing

  • Categorical variables are dummy-encoded, all variables are min-max normalized to [−1, 1], and generated values are clipped, denormalized, and rounded as needed.Integer variables are rounded, and dummy variables are rounded to the nearest class.Source: §3.6 · p. 6

3.7 Gradient-Boosted Tree hyperparameters

  • The authors set XGBoost’s L2 hyperparameter to 0, leave other settings at defaults, and prioritize avoiding underfitting over overfitting.Source: §3.7 · p. 6

3.8 Training one model per category

  • For categorical outcomes, training a separate XGBoost model per label improves generation performance compared with conditioning one model on the label.The separate-model approach is equivalent to forcing each tree to split by label before building deeper trees.Source: §3.8 · p. 6

4 EXPERIMENTS

The experiments benchmark generation and imputation across 27 real-world datasets using metrics spanning distributional closeness, diversity, prediction, and statistical inference. Forest methods lead on generation, while established tree-based methods remain strongest overall for imputation.

4 EXPERIMENTS

  • The benchmark evaluates complete-data generation, imputation, and incomplete-data generation on 27 real-world datasets, with metrics covering four evaluation quadrants.Datasets are split 80% for training and 20% for testing; methods are compared using average ranks and standard errors.Source: §4 · pp. 6–7

4.1 Generation with complete tabular data

  • Forest-Flow ranks highest across nearly all complete-data generation metrics, followed closely by Forest-VP and TabDDPM, and nearly matches the training-data oracle on Wasserstein distance.Table 2 reports averaged ranks over 27 datasets and three experiments per dataset.Source: §4.1 · p. 7, Table 2 · p. 7

4.2 Imputation on missing tabular data

  • MICE-Forest and MissForest are the best overall imputers; MissForest has low diversity, while Forest methods are more diverse and can rank above or below other methods by metric.Forest methods generally outperform remaining methods on raw scores, but rank comparisons vary across metrics.Source: Table 3 · p. 7, §4.2 · p. 8

4.3 Generation with incomplete tabular data

  • Forest-Flow and Forest-VP are generally the highest-ranked incomplete-data generators, except that TabDDPM is slightly better on training coverage, R-squared, and coverage rate.The comparison uses 20% missing values; other methods receive MissForest-imputed inputs, whereas Forest methods train directly on incomplete data.Source: §4.3 · p. 8, Table 4 · p. 8

4.4 Ablation of Forest-VP and Forest-Flow

  • On Iris, nt=50 gives the best performance for both flow and diffusion; XGBoost and LightGBM outperform other tree approximators, while label-conditioning gains are not significant.XGBoost is selected because LightGBM freezes when multiple models run in parallel, and serial training is much slower than parallel XGBoost.Source: §4.4 · pp. 8–9, Table 5 · p. 8, Table 6 · p. 8

5 CONCLUSION

Forest-based diffusion and flow models generate realistic tabular data, including from incomplete training data, and produce diverse imputations. Their generation results are strong, but imputation and memory demands remain limitations.

5 CONCLUSION

  • Forest-based diffusion and flow models generate highly realistic tabular data, even with missing training values, and perform better or on par with deep-learning methods without GPUs.The paper reports state-of-the-art results on generation, while MICE-Forest and MissForest remain better for imputation.Source: §5 · p. 9
  • MICE-Forest and MissForest outperform the method on imputation, while GBT memory demands can be higher for large datasets.The lack of minibatch training and duplication of rows nnoise times can increase memory use.Source: §5 · p. 9
  • Future work could explore multinomial diffusion, new diffusion processes, classifier-free guidance, minibatch training for tree models, and applications such as data augmentation.The paper also mentions class imbalance and domain translation as possible applications.Source: §5 · p. 9
  • Feature importance across XGBoost models may help explain generation and support variable selection.Figure 4 shows average SHAP feature importance for Forest-Flow on Iris.Source: §5 · p. 9

6 BROADER IMPACTS

The paper cautions that realistic synthetic observations are generated rather than real participants, and urges comparing inference from fake and real data. It also discusses missing-data imputation and resource use.

6 BROADER IMPACTS

  • Forest-Flow’s average SHAP feature importance across XGBoost models on Iris can also be plotted by noise level or predicted variable.The plotted features are sepal length, sepal width, petal length, petal width, and species.Source: Figure 4 · p. 9
  • Synthetic observations are generated rather than real participants, so practitioners are encouraged to compare inference from fake data with inference from real data.The paper cautions that conclusions about fake participants may have real-world impact.Source: §6 · p. 9
  • Imputation is sensible with few missing values per participant or strong correlations between observed and missing values; otherwise, it can fabricate much of an observation.The paper recommends multiple imputations for inference and avoiding imputation when too much data is missing.Source: §6 · p. 9
  • The method can run on a regular laptop with 4–12 CPUs, unlike most deep neural networks that require expensive GPUs.The paper presents this as a resource-conservative alternative to resource-intensive generative models.Source: §6 · p. 9

Checklist

The checklist addresses reproducibility and disclosure, while the appendices describe the algorithms, experimental resources, and flow-matching choices.

Checklist

  • The checklist reports that code, data, and instructions for reproducing the main experiments, training details, and measure definitions are included.It also marks computing-infrastructure description as included.Source: Checklist · p. 15
  • The appendices describe the Forest-VP and Forest-Flow algorithms, datasets and metrics, competitors, and averaged raw scores and plots across 27 datasets.The checklist also marks source-code dependencies as included.Source: Checklist · pp. 15–16

A.1 REPAINT Imputation Algorithm

  • The REPAINT imputation algorithm initializes a missingness mask and noise dataset, reverses the process for missing entries, restores observed values, and periodically repaints.The listed defaults include 100 noise samples, 50 noise levels, 10 repaints, and jump size 5.Source: §A.1 · p. 16

A.2 Forward and Reverse step Algorithms

  • The forward and reverse algorithms select conditional flow matching or VP diffusion through a boolean flag and use β = [0.1, 20].The algorithm descriptions distinguish forward steps from reverse steps.Source: §A.2 · p. 17

A.3 Flow matching

  • Flow matching defines a Gaussian conditional probability path and vector field; its variants condition on a training sample or on an independent source-target pair.The independent-conditional variant uses a linear interpolation of source and target samples and a constant standard deviation; experiments set σ to 0.Source: §A.3 · p. 18

A.4 Why Flow methods cannot be used for imputation or inpainting

The authors use diffusion methods, not flow methods, for imputation because a deterministic flow trajectory cannot be steered to preserve a sample’s observed values.

A.4 Why Flow methods cannot be used for imputation or inpainting

  • A flow maps each Gaussian noise sample along a deterministic trajectory to one data sample, so it cannot ensure the output matches a particular participant’s observed values.Forcing the trajectory toward the true non-missing values diverges from the deterministic path; the authors therefore rely exclusively on diffusion for imputation.Source: §A.4 · p. 18

B.1 Datasets

The benchmark datasets are drawn mainly from prior tabular-data work and public repositories, and their details are listed in Table 7.

B.1 Datasets

  • The benchmark mostly reuses datasets from Muzellec et al. (2020), with bean as an exception, and lists them in Table 7.The datasets come from the UCI Machine Learning Repository or scikit-learn; the passage states that all are openly shared without usage restrictions.Source: §B.1 · p. 19, Table 7 · p. 19

B.2 Metrics

The benchmark uses metrics tailored to mixed-type data and evaluates imputation, generation, computational resources, and baseline configurations. It reports averaged scores and plots across datasets, while the supplied passages do not provide the score values.

B.2.1 Gower distance

  • Gower distance addresses mixed-type variables at different scales; its per-variable costs range from 0 to 1.The paper uses it for distance-based metrics; it can be implemented as L1 distance on min-max normalized continuous variables or one-hot categorical variables divided by two.Source: §B.2.1 · pp. 19–20

B.2.2 Imputation

  • Imputation evaluation includes Wasserstein distance, diversity across five imputations, and both minimum and average MAE.Wasserstein distance is reported against training or testing data only for datasets with fewer than 5,000 training samples; minimum and average MAE balance different concerns about stochastic and deterministic methods.Source: §B.2.2 · p. 20
  • Imputation utility averages F1-score for classification or R2 for regression across four non-deep models, while statistical measures assess inference against complete-data estimates.The four utility models are linear/logistic regression, AdaBoost, Random Forests, and XGBoost; inference measures include percent bias and coverage rate, averaged across regression parameters.Source: §B.2.2 · pp. 20–21

B.2.3 Generation

  • Generation evaluation uses Wasserstein distance, coverage, model utility, a real-versus-fake discriminator, and statistical measures.Wasserstein distance is calculated only for datasets with fewer than 5,000 training samples; lower discriminator F1 indicates better generation quality.Source: §B.2.3 · p. 21

B.3.1 Computational resources used

  • Tree-based models were trained on 10–20 CPUs, while the other models used 8 CPUs and 1 GPU.The reported RAM ranges were 64–256Gb for tree-based models and 48–128Gb for the other models.Source: §B.3.1 · p. 21

B.3.2 Generation

  • Generation baselines include GaussianCopula, TVAE, CTGAN, CTAB-GAN+, STaSy, and Tab-DDPM, alongside the Oracle training data.ForestFlow and ForestVP use nt = 50, nnoise = 100, label conditioning, and dummy-encoded categorical variables; the STaSy configuration was adjusted for small datasets.Source: §B.3.2 · pp. 21–22

B.3.3 Imputation

  • Evaluated imputation methods include GAIN, KNN, ICE, and MICE-Forest; Oracle is the actual training data before missing values are added.Forest-Flow and Forest-VP use nt = 50 and nnoise = 100, except bean uses nnoise = 50 to reduce memory. GAIN rounds only variables named as categorical.Source: §B.3.3 · pp. 22–23

C AVERAGED RESULSTS AND BAR PLOTS

  • The appendix reports averaged raw scores and bar plots across datasets and metrics; the tables cover complete-data generation, imputation, and generation with incomplete data.Table 8 covers complete-data generation, Table 9 imputation, and Table 10 generation with incomplete data; Oracle F1disc results should be ignored because the passage says they are incorrect.Source: §C.1 · p. 23, Table 8 · p. 23, Table 9 · p. 23

C.2.1 Bar plots for generation with complete data

The section presents a UMAP visualization of the red wine quality dataset comparing Oracle, ForestFlow, and TabDDPM.

C.2.1 Bar plots for generation with complete data

  • The UMAP visualization compares Oracle, ForestFlow, and TabDDPM on the red wine quality dataset.Source: §C.3 · p. 39
Loading 2309.09968v3…