PaperlayerView exact revision on arXiv

Source-linked AI summary

Unmasking Trees for Tabular Data

Calvin McCarter

arXiv:2407.05593v5cs.LGstat.ML

TL;DR

Tabular imputation and generation remain challenging, with traditional methods continuing to perform strongly on imputation and conditional distributions potentially being multimodal. The paper introduces tree-based autoregressive feature unmasking and BaltoBot for probabilistic prediction, with TabPFN as a possible base learner. On 27 small tabular datasets, UnmaskingTrees leads imputation and achieves state-of-the-art generation with missingness, while remaining competitive for generation without missingness.

  • Problem

    Tabular imputation and generation require modeling conditional distributions that may be multimodal, while traditional methods have continued to win on imputation benchmarks.

    Source: §1 · p. 1

  • Method

    UnmaskingTrees incrementally autoregresses over features with gradient-boosted trees, while BaltoBot models conditional outputs with a balanced tree of boosted classifiers.

    Source: §1 · p. 2

  • Results

    On a benchmark of 27 small tabular datasets, UnmaskingTrees led imputation, was state-of-the-art for generation with missingness, and was competitive for generation without missingness.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The benchmark results show that tree-based autoregressive modeling performs strongly for imputation and generation when tabular data contain missingness.

    Source: §1 · p. 2, §6 · p. 13

    The approach is primarily applicable to small tabular datasets or settings where inference time matters more than training time.

    Source: §4 · p. 11

Abstract

from arXiv · show

Despite much work on advanced deep learning and generative modeling techniques for tabular data generation and imputation, traditional methods have continued to win on imputation benchmarks. We herein present UnmaskingTrees, a simple method for tabular imputation (and generation) employing gradient-boosted decision trees which are used to incrementally unmask individual features. On a benchmark for out-of-the-box performance on 27 small tabular datasets, UnmaskingTrees offers leading performance on imputation; state-of-the-art performance on generation given data with missingness; and competitive performance on vanilla generation given data without missingness. To solve the conditional generation subproblem, we propose a tabular probabilistic prediction method, BaltoBot, which fits a balanced tree of boosted tree classifiers. Unlike older methods, it requires no parametric assumption on the conditional distribution, accommodating features with multimodal distributions; unlike newer diffusion methods, it offers fast sampling, closed-form density estimation, and flexible handling of discrete variables. We finally consider our two approaches as meta-algorithms, demonstrating in-context learning-based generative modeling with TabPFN.

1 Introduction

The paper addresses tabular imputation and generation, especially when data contain missing values, with tree-based methods for incrementally unmasking features. On a 27-dataset benchmark, UnmaskingTrees performs strongly on imputation and generation with missingness.

1 Introduction

  • ForestDiffusion showed leading generation results but struggled with imputation, where MissForest outperformed it.The paper identifies this as a shortfall its approach addresses.Source: §1 · pp. 1–2
  • UnmaskingTrees trains gradient-boosted trees to unmask features autoregressively in random order, conditioning on observed values to impute missing features.The method uses permutation language modeling and tree predictors that can handle varying subsets of missing features.Source: §1 · p. 2
  • Autoregressive generation requires modeling each feature’s full conditional distribution, which may be multimodal and makes mean prediction or parametric approaches a poor fit.The paper motivates probabilistic prediction as a subproblem of autoregressive imputation and generation.Source: §1 · p. 2
  • BaltoBot uses a balanced tree of boosted classifiers to partition each output space, enabling fast sampling, closed-form density estimates, and discrete-outcome modeling.Its partitions use kernel density integral quantization.Source: §1 · p. 2
  • UnmaskingTrees and BaltoBot can combine with probabilistic binary classifiers such as TabPFN to form generative models, using only a few TabPFN forward passes per sample.This contrasts with TabPFGen’s many iterative data updates.Source: §1 · p. 2
  • 27 tabular datasets formed the out-of-the-box benchmark, where UnmaskingTrees led imputation, was state-of-the-art for generation with missingness, and was competitive for generation without missingness.The benchmark concerns small tabular datasets.Source: §1 · p. 2

2 Method

The method models joint tabular distributions by randomly ordering and incrementally unmasking features, and models continuous conditional outputs with a balanced tree of boosted classifiers. It also compares computational costs and adapts the frameworks to TabPFN.

2.1 UnmaskingTrees for tabular joint distribution modeling

  • UnmaskingTrees maximizes expected log-likelihood across feature-order permutations, training one predictor per feature on randomly permuted, incrementally masked samples.Categorical features use softmax classification; the continuous-feature approach is BaltoBot.Source: §2.1 · p. 3

2.2 BaltoBot for tabular probabilistic prediction

  • BaltoBot recursively partitions continuous outputs using KDI and boosted classifiers, then samples along one tree path; its design supports multimodal distributions, discrete outcomes, and closed-form conditional density estimates.KDI adaptively interpolates between uniform and quantile quantization; singleton bins retain frequent discrete values.Source: §2.2 · p. 4

2.3 Computational complexity

  • With T ~50 and H ~4, each training sample reaches DH classifiers versus ForestDiffusion’s DT regressors, while inference requires DH rather than DT predictions.The authors report a large training speedup and note inference savings are especially impactful for multiple imputation.Source: §2.3 · p. 5

2.4 In-context learning-based generation with BaltoBoTabPFN and UnmaskingTabPFN

  • TabPFN serves as a base learner for probabilistic prediction and generation; NanTabPFN wraps it to support missing values in training and test features.Its query-specific row and column selection supplies TabPFN with a training subset without NaNs.Source: §2.4 · p. 5

3 Results

Across the 27-dataset benchmark, UnmaskingTrees leads on imputation and generation with missing data, while remaining competitive on complete-data generation. The section also reports case-study and probabilistic-prediction results for BaltoBot.

3.1 Case studies on Two Moons and Iris datasets

  • On Two Moons, UnmaskingTrees imputes cleanly while competing methods show out-of-distribution values; on Iris it conditions better on species and yields more diverse imputations than MissForest.Source: Figure 1 · p. 6, Figure 2 · p. 6, §3.1 · p. 7

3 Results

  • Across the benchmark, UnmaskingTrees leads on imputation and generation from incomplete data, while remaining competitive on fully observed generation.Source: §3 · p. 7

3.2 Benchmarking UnmaskingTrees on 27 tabular datasets

  • UnmaskingTrees wins 3/9 imputation metrics, beats ForestDiffusion on 8/9, and beats MissForest head-to-head 5–4, with average rank 3.2 versus 3.5.Source: Table 1 · p. 8, §3.2 · p. 8
  • Ablations show progressive improvements from the UnmaskingTrees framework, KDI over k-Means quantization, and BaltoBot in the full solution.Source: Table 2 · p. 8, §3.2 · p. 8
  • With incomplete data, UnmaskingTrees ranks first on 5/9 generation metrics and beats TabDDPM 5–4 and Forest-Flow 6–3 head-to-head; on complete data, it loses to three methods but beats the others.Source: §3.2 · p. 8, Table 3 · p. 9, Table 4 · p. 9

3.3 Evaluating BaltoBot on synthetic probabilistic prediction case studies

  • On 5,000 wave samples, BaltoBot matches Treeffuser's predictions visually, samples about 7× faster, and provides a closed-form pdf estimate.Source: §3.3 · p. 9, Figure 3 · p. 10
  • On Poisson-distributed data, BaltoBot models the count-type distribution, whereas Treeffuser produces a negative-valued outlier and non-integer samples.Source: §3.3 · p. 9, Figure 4 · p. 10

3.4 Sales forecasting with uncertainty

  • On M5 sales forecasting, tuned BaltoBot ties Treeffuser on CRPS but outperforms it on RMSE and MAE; untuned BaltoBoTabPFN strongly outperforms Deep Ensembles on CRPS.Source: Table 5 · p. 11, §3.4 · p. 11

4 Limitations and Future Work

The methods have distinct boundaries: UnmaskingTrees does not win every imputation metric or vanilla-generation benchmark, while its training cost limits its fit to smaller datasets or settings prioritizing inference. BaltoBoTabPFN also remains impractical beyond the smallest datasets on the reported machine, and the paper identifies several routes for future work.

4 Limitations and Future Work

  • MissForest still outperformed UnmaskingTrees on Wasserstein-distance metrics, Forest-Flow won vanilla generation, and whether one method can win across scenarios and metrics remains unresolved.These are exceptions to UnmaskingTrees' overall lead on the imputation benchmark.Source: §4 · p. 11
  • BaltoBoTabPFN was very slow and ran out of memory as an UnmaskingTabPFN subroutine on the benchmark, requiring improvement for practical use beyond the smallest datasets.The out-of-memory result was observed on the authors' machine.Source: §4 · p. 11
  • UnmaskingTrees training scales cubically with feature dimensionality, versus quadratic complexity for ForestDiffusion and single-order autoregression, limiting its fit to small datasets or inference-prioritized settings.Single-order autoregression supports generation but not arbitrary imputation.Source: §4 · p. 11
  • The approach is aimed at smaller datasets and out-of-the-box use by people with limited compute or no hyperparameter-optimization resources; larger-data users may do better with deep-learning methods such as Tabsyn (Zhang et al., 2024).The authors frame the evaluation around that intended scope.Source: §4 · p. 11
  • Future work includes reducing feature-count training cost, choosing feature orders or multi-output trees more efficiently, analyzing BaltoBot's trade-offs, and extending its binary partitioning to neural models.The authors also propose evaluating uncertainty quantification and probabilistic forecasting, and extending the method to multiple dimensions.Source: §4 · p. 12

5 Discussion and Related Work

The discussion offers speculative reasons for the methods' different performance across missing-data and vanilla-generation settings, and situates autoregressive tabular modeling alongside diffusion and earlier statistical work. It also contrasts where randomness enters the competing generation procedures.

5 Discussion and Related Work

  • The authors speculate that autoregression suits imputation because it learns conditional distributions, while diffusion has a train-inference gap for missing data; diffusion's advantages may favor vanilla generation when these issues are absent.The diffusion gap is that observed features start reverse generation from noise during training but from their actual values during imputation.Source: §5 · p. 12
  • Deep learning has lagged boosted trees on tabular data (Shwartz-Ziv and Armon, 2022; Jolicoeur-Martineau et al., 2024b), and sharper inductive priors may help; the authors anticipate a future TabPFN variant could replace XGBoost.Breejen et al. (2024) suggests synthetic random-forest data can further improve TabPFN classification.Source: §5 · p. 12
  • MissForest remains competitive and exploits missingness fractions; the discussion also connects conditional modeling to earlier statistical work and autoregression to the full-mechanism bootstrap.The authors note they are unaware of machine-learning approaches using missingness-fraction information.Source: §5 · p. 12
  • Flow-matching methods inject randomness at the start, diffusion methods inject it at the start and during reversal, and UnmaskingTrees adds it gradually through ordering, probabilistic leaf sampling, and within-bin sampling.The cited flow-matching work includes Liu et al. (2022), Albergo and Vanden-Eijnden (2022), and Lipman et al. (2022); the cited diffusion work includes Sohl-Dickstein et al. (2015) and Song and Ermon (2019).Source: §5 · pp. 12–13

6 Conclusions

The paper concludes by proposing tree-based autoregressive modeling for tabular data, particularly with missingness, alongside hierarchical partitioning for conditional probabilistic prediction. Its benchmark results are strongest for imputation and generation given missing data.

6 Conclusions

  • The paper proposes tree-based autoregressive modeling of tabular data, with particular focus on data containing missingness.Source: §6 · p. 13
  • For conditional probabilistic prediction, the paper presents hierarchical partitioning with benefits over vanilla quantization and diffusion-based prediction, and uses its methods as meta-algorithms with TabPFN.Source: §6 · p. 13
  • On the out-of-the-box tabular benchmark, the authors report leading imputation results and state-of-the-art generation results given data with missingness.Source: §6 · p. 13

A Ablation experiment with imputation - raw scores

The imputation ablation reports that UnmaskingTrees improves on Forest-VP, KDI quantization outperforms k-Means quantization, and BaltoBot adds further improvement.

A Ablation experiment with imputation - raw scores

B Full dataset-level results

The full dataset-level results cover imputation, generation without missingness, and generation with 20% missingness; the paper also reports runtime for both tasks.

B Full dataset-level results

  • 1976 seconds on Libras (N = 360, D = 90) and 1929 seconds on Bean (N = 13611, D = 16) were the slowest imputation runs.These timings were measured on the authors’ 2015 iMac with 16Gb RAM; ForestVP ran on a more powerful cluster.Source: B Full dataset-level results · p. 18
  • 4346 seconds on Bean and 2987 seconds on Libras were the slowest generation runs for UnmaskingTrees.ForestFlow and ForestVP generation timings are also reported for those datasets.Source: B Full dataset-level results · p. 18
  • The full results are reported separately for imputation, generation without missingness, and generation with 20% missingness.These correspond to Tables 7, 8, and 9, respectively.Source: Table 7 · p. 18, Table 8 · p. 19, Table 9 · p. 19
  • UnmaskingTrees is described as relatively efficient for both imputation and generation, with runtime compared across feature and sample counts.Figure 5 depicts the runtime results, and Table 10 reports runtime for the 27-dataset benchmark.Source: B Full dataset-level results · p. 18, Table 10 · p. 19, Figure 5 · p. 20
Loading 2407.05593v5…