Unmasking Trees for Tabular Data
Authors: Calvin McCarter
Venue: Accepted by TMLR ()

Abstract: 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.

Published in Transactions on Machine Learning Research (7/2025)
Unmasking Trees for Tabular Data
Calvin McCarter mccarter , calvin@gmail, com
BigHat Biosciences"
Reviewed on OpenReview: http:://pppaarreee.ar//////////p
Abstract
Despite much work on advanced deep learning and generative modeling techniques for tabular data generation and imputation, traditional methods have contimued to win on imputation benchmarks. We herein present UnmaskingTress a simple method for tabular imputation (and generation) employing gradient- decision trees which are used to incrementally ummask individual features. On a benchmark for on 27 small tabular datasets, Unmasking Trees offers leading performance imputation; state-of-the-art performance generation given data with missingness; and competitive performance vanilla generation given data without missingness. To solve the conditional generation subproblem, we propose tabular probabilistic prediction method, BaltoBot, which fits a balanced tree of boosted tree classifiers. Unlike older methods, requires no parametric assumption on the conditional distribution, accommodating feat ures 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.
performance
1 Introduction
Given a tabular dataset, is froquently desirable to impute any missing values within that dataset, and to generate new synthetic examples. Due to the prevalence missingness tabular datsets, imputation has been a long-standing task statistics and machine learning (Titttle and Rubin 2019). Tn particular, multiple imputation methods, which produce multiple samples from the estimated conditional distribution of missing features, have proved advantagoous for downstream inferential and prediction tasks (Rubin) 1996). Multiple imputation has also been used as subroutine for counterfactual estimation (Kreindeer and Lumsden 2016
Yoon et al. 2018b) and domain adaptation (Ragab et al. 2023 McCarter 2024h. Synthetic data generation for ta labular data has also seen recent interest, with applications in addressing data imbalance (Van Breugel
et al. 2021 Kim et al. 2022b and in preserving privacy (Kotelnikov et al. 2023 Gulati and Roysdon 2024
Imputation and generation are closely related tasks. Multiple imputation can be seen as a form of conditional generation, where the partitioning between output variables and inpul variables is not known in advance.
Generation is then special case of imputation where the set of observed conditioning variables is empty. Furthermore, due to the aforementioned prevalance of missingness in tabular data, generation methods also frequently need to be able to handle missingness at training time.
In this work, we are primarily focused on generation and imputation methods for users with limited data and computing resources. On data generation, recent work Malicoeur-Martinaaa et al.] (024b) (ForestDiffusion) has shown leading results on data generation using gradient- trees (Chen and Guestrin 2016 trained diffusion flow-matching objectives, outperforming deep learning-based approaches, particularly on smaller datasets. However, this approach tended to struggle tabular imputation tasks, outperformed by
* Work done primarily prior joining BigHat Biosciences. Views expressed the author's 's own.

---

Published in Transactions on Machine Learning Research (7/2025)
MissForest (Stekhoven and Bühhmann 2012), an older multiple imputation approach based on random forests [Breiman 2001].
We address this shortfall by training gradient-boosted trees to autoregressively unmask features in random order, via permutation language modeling (Yang 2019). This autoregressive approach, which we dub UnmaskingTress naturally performs conditional generation (i.e. imputation): we simply fill in and condition on observed values, autoregressively generating the remaining missing values. This contrasts with tabular diffusion modeling, for which the RePaint inpainting algorithm (Lugmayr et al., 2022) is employed to mediocre effect Policoeur-Martines et al. 2024b). Because the predictor for given feature must condition on varying subsets of the other features, the ability of gradient-boosted trees to handle missing features makes them a natural choice for autoregressive modeling. Hence, we maintain the tree-based approach of Jolicoeur- Martineen et al. (2024b), while replacing their tree- based regressors with our novel brsed probabilistic predictors, which we turn to next.
While mean-estinating regression models are satisfactory for diffusion, for autoregression we must inject noise, and hence must estimate the entire conditional distribution of each feature. We therefore revisit the long-studied problem of (tabular) probabilistic prediction (Le et al. 2005 Meinshausen and Ridgeway 2006). Because the conditional distribution is possibly multi-modal, paramctric approaches such as XGBoostLSS März 2019), NGBoost (Duan et al. 2020), and PGBM (Sprangers et. al. 2021 are poor choices for our setting. Meanwhile, quantization of a continuous variable can model its multi-modalityy but at the cost of destroying either low-resolution or high-resolution information. A diffusion-bassd method, Treeffuser Beltran Velez et al. (2024), was recently proposed to address these problems. However, as a diffusion method, it suffers from slow sampling and is unable to provide closed-form density estimates; furthermore, Treeffuser does not naturally model discrete outcomes. To address these problems, we propose BaltoBot, a balanced tree of boosted trees. For each individual variable, we recursively divide its output space with the kernel density integral (KDI) quantizer [McCarter) (223) into a "meta-tree" of binary classifiers, which for us are gradicnt-boosted trecs. This allows us to efficiently generate samples and estimate densitics, because cach sample follows only one path from root to leaf of the mela-tree. Performing regression with hierarchical classification proved successful in computer vision object bounding box prediction (Li et al., 2020), but has been surprisingly underexplored in tabular ML and in generative modeling.
Our two methods are in fact mela-algorithms that, in combination, can create a generative model out of any probabilistic binary classifier. To demonstrate this flexibility, we swap out XGBoost (Chen and Guestrin 2016) for TabPFN (Hollmann et al. 2022). TabPFN is a deep learning model pretrained to perform in-context learning for tabular classification. While it has state-of-the-ar classification benchmark performance (McElfresh et al. (004)), it does not perform regression tasks, nor does it inherently perform generative modeling. Constructing a generative model out of TabPEN (Hollmann al. 2022 was first proposed in TabPFGen (Ma et al. 2024), which approximates the posterior from TabPFN- by iteratively applying stochastic gradient Langevin dynamics (Welling and Teh 2011). But unlike the previous work, ours requires only a few TabPFN for forward-passes for each sample rather than many iterative data updates.
likelihoods
We showcase Unmasking'Tree on two tabular case studies, and on the benchmark of 27 tabular datasets presented by Jolicocur- Martincan et al. (2024b). On this benchmark for "out-of-the-bon performance on small tabular datasets, our approach offers leading performance on imputation and slate-of-the-art performance on generation given data with missingness; and it has competitive performance on vanilla generation without missingness. We also demonstrate that BaltoBot is on its own promising method for probabilistic prediction, showing its advantages on synthetic case studies and on a heavy-tailed sales forecasting benchmark.
Finally, we provide code with an easy-to-use sklearn-style APT at htttps://ii thub. com/ ca calvvnnccatterr unmasking-trees treess In addition to being useful for practitioners, we hope our work sparks study within the tabular ML community about whether diffusion or autoregression is better for tabular data. Previous autoregressive tabular modeling methods, TabMT (Gulati and Roysdon 2024) and DP-TBART (Castellon ct al. 2023), ; use Transformer (Vaswani 2017) models, making them less applicable for the GPU-poor; they also lack publicly-available implementations Our simple, efficient implementations of UnmaskingTrees and BaltoBot contribute to investigating this question.
2

---

Published in Transactions on Machine Learning Research (7/2025)
2 Method
2.1
Unmasking Trees for tabular joint distribution modeling
UnmaskingTrees combines gradient-boosted trees with the training objective of generalized autoregressive language modeling (Yang) 2019), inheriting the benefits of both. Consider a dataset with N samples and features. We learn the joint distribution over D-dimensional sample x by maximizing the expected log-likeldoodd with respect to all possible permutations of the factorization order,
logp(x) = log E [[(row))X(<))))
where a is permutation drawn uniformly from U(G)), the permutation group on D features; xx(<() denotes all features that precede the t-th feature in the permuted sequence of features. If we were to have marginalizod permutations, we would have obtained masked language modeling procedure with a randomly-sampled masking rate ~ U(0,11 1) (Liao et al. 2020: Kitouni al. 2023 2024); such procedure was previously shown to have benefits in combination with tabular Transformer models (Gulati and Roysdon 2024) (TabMT).
For each sample, we generate new training samples by randomly sampling an order over the features, then incrementally masking the features that random order. Given duplication factor K, we repeat this process times with different random permutations, leading to a training dataset with KN samples. Given this, train XGBoost (Chen and Guestrin, 2016) models to predict each unmasked sample given the more-masked sample derived from it, one per feature. We model categorical features via softmax-based classification with cross-entropy loss; our approach for non-categorical features is described in Section 2.2
For both generation and imputation, we generate features of each sample in random order. For imputation rather than generation tasks, we begin by filling in each sample with the observed values, and run inference on the remaining unobserved features. Implementing this very simple: requires about 70 lines of Python code for training, and about. 20 lines for inference. The training algorithm for Unmasking" Trees given Algorithm I
Algorithm 1 Unmasking Trees training
Require: dataset X E Nx:; duplication factor K.
1: {# Build self-supervised training set}
Set Xtrain 0, Y train
for h == 1,..., K do
N do for n = ,,...,
Draw random permulation from ((G
6:
Set x = X and y := X
for 1,...., do 7:
Mask random element (o = [MASK].
Append Xtrain == Xurain U (æ), Y train = Y train {
10: end for
11: end for
12: end for
13: { Train conditional generation models}
14: for d= 1,..,, D do
15: if feature d in X is a categorical feature then
16: Train XGBClassfier on (XXrain):. (Yrran)..d)
17: else
18: Run BaltoBot with (Xurainks [Y Yerainl.d).
19: end if
20: end for
3

---

Published in Transactions on Machine Learning Research (7/2025)
2.2 BaltoBot for tabular probabilistic prediction
In diffusion models, predicting the score function can be framed a regression problem where the model learns to estimate the conditional mean. However, key problem when switching to autoregressively generating continuous data is that this regression approach will attempt to predict the mean of conditional distribution, whereas we would like the model to sample from the mussibly-multmmoda conditional distribution. The simplest solution is to quantize continuous features into bins, because classification over histograms is inherently multimodal; TabMT (Gulati and Roysdon 2024) did this with 1d k-Means clustering (Lloyd 1982). Yet this not only destroys information within bins due to rounding, it also destroys information about the proximity among the ordered bins. Thus, it forces us to choose between a small number of quantization bins, yielding low resolution; to choose large number of bins, risking catastrophic errors due to overfitting and/or clumping of generated samples duc to poor calibration. This not only limits performance, but also necessitates hyperparameter luning (Gulati and Roysdon 2024).
Inspired by this, we propose a general-purpose solution to the tabular probabilistic prediction problem. For each individual rogression output variable, we build a height-H balancod tree of binary classifiers. Consider a ,train
node with height on this "meta-tree", which is fit with (Xirain E RRaxd Ytrain E )R²).. Using kernel density integral quantization (KDI) [McCarter (023), which adaptively interpolates between uniform quantization and quantile quantization, we obtain binarized Yurain E [0, 1]". Thus, the input space to every node is partitioned into two with the splitting point determined by KDI. We train an XGBoost classifier on (Xtrain Feraia) If h > 0, we then recursively pass ((X(), ()()) E (X Vtrain ll 0} to its left. child, and analogously for g(o = to its right child. At a leaf node, h = 0, if given single unique training set output value in a bin, we record this value. At inference time, given a query input X, we descend the tree by obtaining predicted probabilities from each node's XGBoost classifier, then sampling from these. Once we reach a leaf node, we either sample uniformly from its appropriate bin, or we return the lone output value if a singleton bin.
Conceptually; our proposed approach has four advantages. First, training and inference time, each XGBoost. model within the meta-tree only sees samples that fall into its corresponding region of the output space. Thus, for a meta-tree with height. H (and thus 2 models), each sample is only passed as input to H different models. While lower-level classifiers receive less data and are poorer quality, the magnitude of such errors are smaller due to our hierarchical partitioning approach. Second, our singleton-bin technique allows us to adaptively generate discrete and mixed-type variables, if the discrete outcome frequent relative to the total number of training samples and to the size of the meta-tree. (Up to 2" discrete outcomes can be produced by BaltoBot.) Third, our adoption of KDI instead of KMeans for feature quantization is beneficial because tabular data often has features which are irregular highly skewed. KMeans clustering is based mixture modeling with cqual-sized Gaussian components: in constrast, KDI is 1d density- based clustering method looks for local minima after applying a smoothing transformation. KDI was previously shown (McCarter 2023) to be better at discretizing such irregular features than KMeans, so we propose adopting it here for generative modeling. Finally, eschewing diffusion modeling enables us to perform closed-form conditional density estimation.
The training and inference algorithms for BaltoBot are given in Algorithms 2 and 3 respectively.
Algorithm 2 BaltoBot training
Require: dataset (X E NxD y E R^); BaltoBot meta-tree height H; 1: if H = 0 or unique(y) = for some constant C' then Save bounds == (min(y), max(y)).
2:
clsc
Obtain split point p from KDI quantization on y.
Train XGBoost binary classifier on (X,1(y S p)).
6:
Train "left-child" BaltoBot on ((x.) y() E (X. < p), with height H - 1.
Train "right-child" BaltoBot on ((1)) E (X,))ly(0) > ph. with height H - 1.
7:
8:
end if
4

---

Published in Transactions on Machine Learning Research (7/2025)
Algorithm 3 BaltoBot inference
Require: imput query x E RD; trained BaltoBot model. 1: bounds is defined then
Sample uniformly from (bbunds)). Return.
else
Obtain prediction from XGBoost binary classifier.
6: if prediction = left-child then
7: Run inference on "left-child" BaltoBot with input query x.
8: else if prediction = right-child then
9: Run inference on "right-child" BaltoBot with input query x.
10: end if
11:
end if
2.3 Computational complexity
Forest with T diffusion steps and duplication factor K, constructs a training dataset of size TKN x D. Given the same duplication factor K, Unmasking Trees will construct a training dataset of size KND x D. Mcanwhile, ForcstDiiffssin must train DT different XGBoost regression models. We, on the other hand, train different BaltoBot models, one per feature; with BaltoBot meta-tree height of H, we then train total of D2" XGBoost binary classifiers. However, classifiers lower in the BaltoBot meta-tree become progressivcly faster to train. Indeed, each constructed training sample will be seen by DT different XGBoost regressors with ForestDiffusion, but only DH classifiers with our approach. Given that T ~ 50 and H ~ 4, this yields a large speedup for our approach.
The KDI quantizer (McCarter 2023 has negligible contribution to runtime, because it uses the polynomial- exponential kernel density estimator (KDE) (Hofmeyr 2019). which has linear complexity in sample size for 1d data, unlike the quadratic complexity of the Gaussian KDE.
At inference time, each Diffusinn generated sample passes through T steps of the diffusson reverse- process, for total of DT XGBoost predictions. For Unmasking with BaltoBot, each generated sample instead requires only DH XGBoost predictions, because each sample follows only one path from root to leaf of the meta-tree. The resulting speedup is especially impactful for the multiple imputation scenario, where inference time dominates.
We observe that MissForest and MoCe-Forest require training DT models, where number of iterations is determined by stopping criterion that tests for convergence. Assuming that dataset size does not affect the number of required iterations, these methods each have complexity O(ND²).
2.4 In-context learning- based generation with BaltoBoTabPFri and CamnaskingTabPre
Within our flexible frameworks for joint and conditional modeling, TabPFN (Hollmann et al. 2022 can be used base learner for probabilistic prediction and generative modeling. For Unmasking TabPFN joint modeling, a difficulty arises from TabPFN's inability to handle inputs Xtmin with missing values (NaNs). To address this, we developed NanTabPFN, a wrapper for TabPFN that supports missingness in both training and test features. Based on each test query Xtest, we select row indices R and column indices C so that XuaainR has no NaNs, using the following key idea. Consider a particular train sample Xirsin and test query Xtest; with visible (non-missing) features denoted by sets (X and (xxust). We can maximize the number of utilized features, while also ensuring that TabPFN receives no NaNs, by restricting the set of columns to those observed for the test query, C := (Xxtest), then choosing training samples R := {i|C C ((x(()) }}- In practice, our procodure is more complicated, because the above choices may result in either empty C empty R. If R empty, we incrementally non-empty training set. If C is empty, we introd all-1s feature to both X and Xtest-
random features of Xtest to missing until we are able to obtain
5

---

Published in Transactions on Machine Learning Research (7/2025)
3
Results
Forest-VP
UnmaskingTrees Trees
MissForest
Forest-VP
Unmasking Trees
Forest-Flow
TnmaskinggTbbPFN
MICE-Forret
Forest-VP RePaint
Unmasking Tab
data generated
ceta imputed
Figure 1: Rosults on Two Moons case study. Original data is shown in green; generated data is shown in red; imputed data is shown in blue.
MissForest imputed
MICE-Forest Imputed
Original data
Species setosa
petal width (cm)
versicolor
virginica Imputed
False
True
FP generated
Forest-Floo generated
Forest-VV VP imputed
Forest VP imputed w/ RePaint
petal wicth (am)
UnmaskingTrees
UnmaskingTabPFN generated
UnmaskingTrees imputed
UnmaskingTabbPFY imputed
generated
pelal width (am)
5.0 potal length (cm)
2.5
5.0 petal length (cm)
2.5
5.0 potal length (cm)
2.5
2.5
5.0
petal length (cm)
Figure 2: Results on Iris dataset, with species, petal width, and petal length depicted. Original data and syntheticall--gnnraaeed datasets are shown on the left columns. The imputed dataset shown on the right columns, with x symbols highlighting the samples with any missingness that required imputation.

---

Published in Transactions on Machine Learning Research (7/2025)
We evaluate Unmasking Trees on two case studies (Section 3.1 and on tabular benchmark of 27 datasets (Section 3.2). We then evaluate BaltoBot and BaltoBoTabPPNN on tabular probabilistic prediction case studies (Section 3.3) and on a sales forecasting dataset (Section 3.4h. Results were obtained always using our method's default hyperparameters: BaltoBot tree height of 4, and duplication factor K = 50. These hyperparameter values were luned on the Two Moons and Tris case studies, then applied without further tuning to the remaining experiments, because tuning problematic for users with limited computing and/or data resources. (The tree height H was tuned on Two Moons and Iris by increasing H until we saw that the resulting plots stopped showing visible improvement.) XGBoost hyperparameters were set to their defaults. Experiments were performed on iMac (21.5-inch, Late 2015) with 2.8GHz Intel Core i5 processor and 16GB memory.
Overall, UnmaskingTrees (using BaltoBot) has leading performance on imputation and saat-of-the-att performance on generation after training on incomplete data; and it has competitive performance on vanilla tabular generation scenarios. We further demonstrate the benefits of BaltoBot and BaltoBoTabPFN when evaluated in their own right for probabilistic prediction.
3.1 Case studies on Two Moons and Iris datasets
Two Moons dataset We first compare our approach to previous leading methods on the synthetic Two Moons dataset with 200 training samples and noise level W 0.1). We compare Unmasking Troes to MissForest. (Stekhoven and Bühlmann 2012), MICE-Forest. (Van Buuren et al. 1999 Wilson al., 2022 (another popular traditional multiple imputation method), and Foresu with default hyperparameters for all methods. For ForestDiffusson, we evaluate both the dow-matching (Forest-Flow) versions on generation; on imputation, we evaluate Forest-VP with and without RePaint, again using default RePaint. hyperparameters; Forest-Flow does not support imputation.
SDE diffusson (Forest- and
We show results in Figure [ On generation, Forest-VE appears to do best according to visual inspection, while UnmaskingTrees and Forest-Flow perform similarly decently. UnmaskingTabbFF performs poorly, but does capture the overall shape of the distribution. Next, we turn to imputation, wherein we request single imputation for copy of the original training data with the second dimension (y-axis) values masked out. DiffstDDi struggles with and without RePaint, with substantial out-of-distribution imputations, and MissForest and MICE-Foree share this problem to lesser degrees. Meanwhile, 'immaskingTree generates impeccable imputations.
plotting petal length, petal width, and species. We compare both methods on generation, and to compare on imputation, we create another version of the Iris dataset, with missingness completely at random: we randomly select samples with 50% chance have any missingness, and on these samples, we mask the non-species feature values with 50% chance. Visually, ForestDiffusion and Unmasking Trees perform about equally well generation. Mcanwhile, imputation, UnmaskingTrees Trees does better job conditioning on species information than DiffssoDi UnmaskingTrress also produces more diverse imputations than MissForest.
Iris dataset In Figure 2 we show results for the Tris dataset Fisher 1936;
3.2 Benchmarking Unmasking Trees on 27 tabular datasets
Imputation Here, we add EnmaskingTrees to the benchmark of 8 imputation methods 27 public datasets, evaluated according to 9 metrics, developed by Marttnnan et al. (2024b) for evaluating tabular imputation and generation methods. This benchmark primarily contains amaller-sized (with 103- < N < 20,640
and 4 < D= < 90) datasets, which our approach is especially geared towards. Namely, we compare our approach against Forest-VP Jolicoeu-MMartneea et al. (2024b). as woll as k-NN imputation (Tryyankkyya et al. 2001). ) ICE (Buck) 1960
MICE-Forest (Van Buuren et al., 1999 Wilson et al. 2022), MissForest (Stekhoven and Bühhmann 2012), Softimpute (Hastie al. (2015), minibatch Sinkhorn optimal transport (Muzellec et al, 2020 and generative adversarial nets (GAIN) (Yoon et al. 2018a). D We follow Jolicoeu--Nartnnau et al.
(024bb in computing the per-dataset rank of each method relative to other methods, then reporting the
We do not add TabMT (Gulati and Roysdon 2024) and TabPFGen [a et al. 2024 the benchmark because code was provided. We do not add UnmaskingTabbl TabPNN because of out-of-memory errors our machine.
7

---

Published in Transactions on Machine Learning Research (7/2025)
average over 27 datasets. For all methods other than our own, we compute ranks by reusing the raw scores provided in Jolicoeur-Martnnen et al. (2024b) 's code repository.
Results for imputation are shown in Table L Unmasking Trees wins first place on 3/9 metrics, including both metrics based on downstream prediction tasks; and it generally outperforms Forest winning on 8/9 metrics. While MissForest wins first place on 4/9 metrics, Unmasking Trees wins 5-4 head-to-head vs MissForest; UnmaskingTrees has average averaged rank of 3.2 compared to 3.5 for MissForest. Unmasking Trees is also the only method with better than 5th place rank on all metrics.
We report further ablation experiments in Table 2 wherein we run Unmasking Trees without BaltoBot, and instead with vanilla quantization using k-Means clustering (Lloyd] 1982 and KDI quantization (McCarter) 2023). Results showing progressive improvements for the Unmasking Treeiinnggressssss Trees framework, for KDI quantization versus k-Means, and for the BaltoBot method used in our full proposed solution.
Table 1: Tabular data imputation (27 datasets, 3 experiments per dataset, 10 imputations per experiment) with 20% missing. Shown are averaged runk over all datasets and experiments (standard-error). Overall best is highlighted; better of Forest-VP versus ours is boldface blue.
MinMAR
AvgMAR 4 Wirein
Wist +
MAD
n² 4.
F
Pains
Convate
KNN
5.5 (0.5)
6.3 (0.4)
4.9 (0.4)
5.0 (0.4)
8.4 (0)
6.5 (1)
5.7 (1.1)
6.2 (1)
5.4 (0.6)
ICE
6.8 (0.4)
4.7 (0.4)
7.0 (0.5)
7.2 (0.4)
1.8 (0.2)
6.2 (1)
7.0 (0.6)
5.7 (0.9)
5.3 (0.6)
MICE-Forest
3.9 (0.4)
2.5 (0.4)
2.9 (0.2)
3.0 (0.2)
3.6 (0.2)
3.7 (1.4)
3.2 (1)
5.5 (1.2)
4.3 (0.6)
Fisesorest
2.7 (0.5)
4.0 (0.4)
(0.3)
2.0 (0.3)
(0.2)
3.8 (1.4)
2.5 (0.5)
(1.5)
(0.5)
Softimpute
6.7 (0.4)
7.5 (0.4)
6.0 (0.5)
7.3 (0.5)
8.4
6.0 (0.9)
7.8 (0.4)
6.3 (0.9)
6.7 (0.4)
OT
(0.4)
(0.3)
(0.3)
(0.5)
6.8 (0.6)
5.5 (0.8)
GAIN
4.7 (0.4)
6.5 (0.3)
6.0 (0.3)
6.0 (0.2)
6.9 (0.1)
5.7 (0.8)
5.4 (0.8)
4.7 (1)
5.0 (0.6)
Forret- VP
5.3 (0.4)
4.0 (0.5)
5.8 (0.3)
5.1 (0.4)
3.2 (0.4)
4.5 (0.9)
4.6 (0.8)
3.3 (0.6)
5.5 (0.7)
UTTress
(0.5)
(0.5)
(0.4)
(0.5)
38 (0.2)
(0.6)
2.2 (0.6)
2.3 (0.9)
4.7 (0.6)
Table 2: Averaged ranks from ablation study of tabular data imputation (27 datasets, 3 experiments per dataset, 10 imputations per experiment) with 20% missing. Shown are averaged rank over all datasets and experiments (standard-error). Overall best is highlighted; better of Forest VP versus ours is boldface blue. See Table D
for column meanings.
MinMAF AvgMAR
W Wisin
Wish
MAD +
R²
F J
Poies
Convate
KNN
6.8 (0.6)
7.8 (0.6)
6.0 (0.4)
6.1 (0.5)
10.4 (0)
8.2 (1.3)
7.0 (1.5)
7.5 (1.5)
6.5 (0.8)
(0.5)
(0.5)
8.5 (0.6)
8.8 (0.5)
1.9 (0.4)
(1.1)
(0.6)
(1.1)
(0.8)
MICE-Forest
4.8 (0.6)
3.3 (0.6)
3.5 (0.3)
3.4 (0.3)
4.6 (0.4)
4.3 (1.8)
4.3 (1.3)
6.8 (1.))
4.8 (0.7)
MissForest
3.3 (0.7)
5.0 (0.6)
2.2 (0.4)
2.3 (0.4)
7.2 (0.3)
4.7 (1.8)
3.3 (0.9)
6.8 (1.9)
3.8 (0.6)
Softimpute
833 (0.5)
9.3 (0.5)
8.8 (0.6)
8.9 (0.6)
10.4 (0)
7.5 (1.2)
9.8 (0.4)
8.3 (0.9)
7.9 (0.6)
OT
7.2 (0.5)
7.6 (0.4)
7.4 (0.6)
7.4 (0.6)
4,8 (0.4)
8.2 (0.5)
8.8 (0.6)
7.3 (0.7)
5.8 (0.7)
GAIN
(0.5)
8.3 (04))
T,2 (0.5)
7.5 (0.4)
8.9 (0.1)
7.5 (0.8)
(1)
6.1
Vorest-VP
6.4 (0.5)
4.8 (0.6)
7.0 (0.4)
6.1 (0.5)
3.8 (0.5)
6.5 (0.9)
6.6 (0.8)
4.5 (0.8)
6.5 (0.8)
UTrers-MMean
6.0 (0.6)
5,8 (0.5)
6.3 (0.6)
6.1 (0.6)
4.1 (0.3)
4.0 (0.7)
2.9 (0.6)
3.8 (1)
6.0 (0.7
UTr
5.1 (0.5)
5.1 (0.5)
5.4 (0.6)
5.6 (0.3)
4.8 (0.3)
4.5 (0.9)
4.0 (0.5)
3.5 (1.2)
6.4 (0.7)
UTTrees
3.8 (0.5)
3.2 (0.5)
3.8 (0.4)
3.8 (0.5)
5.0 (0.3)
2.7 (0.6)
2.9 (0.8)
3.5 (0.8)
5.8 (0.7)
Generation with and without missingness We next repeat. the experimental setup of Jolicoeur- Martineau et al. (024b) for evaluating tabular generation methods. For tabular generation, using the same 27 datasets, Jolicoeur Martennea et al. (024bb benchmark their methods (Forest-VP and Forest-Flow) against 6 other methods, namely, Gaussian Copula (Joe 2014). tabular variational autoencoding (TVAE) (Xu et al, 2019), two conditional generative adversarial not methods (CTGAN (Xu al.) 2019) and CTAB-GAN+ [Zhao et al.] 2021)), ), and two other tabular diffusion methods (STaSy (Kim et al. 2022al) and TabDDPM Kotelnikov et al.] 2023)). ). These are evaluated with 9 metrics, in the vanilla fully-observed setting and in the synthetically-indeeed 20% missing completely at random (MCAR) setting.
Results for partially-missing data are shown in Table 3 Unmasking Trees is first place on 5/9 metrics; head-to-head, Unmasking Trees beats TabDDPM 5-4, and beats Forest-Flow 6-3. Results for fully-observed data are shown in Table 4 Unmasking Trrees loses head-to- toad-tt- to Forest-Flow, Forest-VP. and TabDDPM, but wins against the other methods.
Raw scores, per-dataset results, and runtimes are provided in the Appendix.
8

---

Published in Transactions on Machine Learning Research (7/2025)
Table Tabular data generation with incomplete data (27 datasets, 3 experiments per dataset, 20% missing values), MissForest is used to impute missing data except in Forest-VP, Forest-Flow, and JumaskingTrocs averaged rank over all dalasets and experiments (standard-error). Overall best highlighted; better of Forest-VP VP versus Forest-Flow versus ours is boldface blue.
Wirein 4 Wsest +
confeain contest $ Raake
Fl + Plaise +
Pains +
compate +
Gaussian Copula
7.0 (0.3)
7.1 (0.2)
7.2 (0.3)
7.1 (0.3)
6.3 (0.4)
6.5 (0.3)
6.7 (0.4)
5.5 (1.0)
7.7 (0.6)
TVAE
5.2
(0.3)
4.9 (0.3)
5.7 (9.3)
5.8 (0.2)
6.0 (1.0)
5.8 (0.5)
5.8 (0.4)
8.0 (0.4)
6.2 (1.0)
CTGAN
(0.2)
8.4 (0.2)
(0.2)
(0.2)
(0.3)
8,4 (0.2)
4.8 (1.2)
7.1 (0.7)
CTABGAN
6.7 (0.4)
6.5 (0.4)
7.1 (0.3)
6.8 (0.3)
7.3 (0.6)
7.1 (0.4)
6.6 (0.3)
7.5 (1.0)
6.1 (0.6)
Stasy
5.9 (0.2)
6.1 (0.3)
5.3 (0.2)
(0.3)
5.8 (0:9)
(0.4)
5.3 (0.4)
3.7 (0.4)
4.6 (1.1)
T&bDDPM
3.0 (0.7)
3.4 (0.7)
2.3 (0.5)
2.9 (0.6)
L7 (0.3)
3.3 (0.6)
3.9 (0.0)
3.8 (1.2)
2.0 (0.5)
Forest-VP
3.7 (0.2)
3.2 (0.3)
3.9 (0.2)
3.8 (0.3)
3.2 (0.3)
2.3 (0.3)
4.2 (0.4)
4.2 (0.8)
(1.1)
Forest-Flow
3.0 (0.3)
2.6 (0.3)
2.6 (0.3)
2.7 (0.2)
3.0 (0.7)
3.7 (0.3)
5.0 (0.5)
3.8 (0.9)
3.2 (0.8)
U'Trees 2.1 (0.2)
2.5 (0.3)
2.5 (0.2)
2.5 (0.2)
3.3 (0.8)
3.5 (0.5)
1.0 (0.0)
3.7 (0.9)
3.7 (1.0)
Table 4: Tabular data generation with complete data (27 datasets, 3 experiments per dataset); averaged rank over all datasots and experiments (standard-crror). Overall best highlightod; better of Forest- VP versus Forest-Flow versus ours is boldface blue.
W Werain 4
Wrest
Covsravn
CoUtest
Ruuee 4
Flyake
Plate
Paiax
Covrate
CoussianCopula
7.1 (0.3)
7.2 (0.3)
7.3 (0.3)
7.4 (0.3)
6.2 (0.2)
6.4 (0.3)
7.0 (0.1)
6.5 (1.1)
7.5 (0.7)
TVAE
5.3 (0.2)
(0.2)
5.7 (0.2)
5.7 (0.2)
6.5 (0.7)
(0.5)
5.5 (0.3)
7.3 (0.6)
6.7 (0.6)
CTGAN
8.4 (0.1)
8.4 (0.2)
8.3 (0.2)
8.1 (0.2)
8.5 (0.2)
8.3 (0.2)
6.7 (0.3)
5.3 (1.1)
7.2 (0.5)
CTAB-GAN+
6.8 (0.3)
6.7 (0.3)
7.2 (0.3)
7.1 (0.3)
6.8 (0.4)
6.9 (0.4)
6.9 (0.3)
7.7 (0.8)
6.7 (0.8)
STaSy
6.1 (0.2)
(0.2)
(0.2)
5.4 (0.2)
(1.2)
(0.3)
(0.3)
4.2
TahDDPM
3.0 (0.7)
3.9 (0.6)
2.8 (0.5)
3.4 (0.5)
1.2 (0.2)
3.8 (0.8)
3.2 (0.4)
3.0 (0.9)
1.4 (0.2)
Porest-VI VP
3.2 (0.2)
(0.2)
3.6 (0.3)
3.3 (0.3)
2.8 (0.3)
2.2 (0.3)
4.3 (0.4)
3.2 (0.9)
3.5 (0.8)
Forest-Flow
1.9 (0.2)
1.5 (0.2)
1.7 (0.2)
1.8 (0.2)
2.3 (0.4)
2.4 (0.3)
4.3 (0.4)
2.8 (0.5)
2.7 (0.4)
UTrress
3.1 (0.1)
3.1 (0.2)
3.1 (0.2)
2.8 (0.2)
4.7 (0.3)
(0.3)
1.0 (0.0)
4.7 (0.7)
5.2 (0.9)
3.3 Evaluating BaltoBot on synthetic probabilistic prediction case studies
Wave dataset We compare our approach with Treeffuser (Beltran- Velez al. 2024) on the "wave"
synthetic dataset from Treeffuser (Beltran-Velez et al., 2024h, which shown in Figure 3 is nonlinear, multimodal, and heteroskedastic. On the raw probabilistic predictions in Figure 3 [),, see That BaltoBot
and BabbPNN are (by visual inspection) able to model the conditional distri buttonnution as well as Treeffuser. Yet this case study illustrates the two advantages of BaltoBot. First, in Figure 3 we show the runtime of the different methods: training, sampling, and total. To train on 5000 samples, Treeffuser took 1.1s and BaltoBot took 2.6s. But to generate 5000 samples, Treeffuser took 5.0s while BaltoBot took 0.72s, for ~ 7x speedup. Second, BaltoBot offers the ability to estimate closed-form probability density function (pdf) of the predictive distribution shown in Figure 3 C); in contrast, Treeffuser only sample from the predictive distribution.
Poisson-distrbutted count data We generate 500 samples of X; ~ Unif(0, 3), Y, Poisson(A VX,), and show probabilistic predictions for Y in Figure 4 Whereas Treeffuser generates spurious negative-valuce outlier and many non-integer Y samples, our approach automatically models the count-type distribution of the data.
3.4 Sales forecasting with uncertainty
We employ the M5 sales forecasting Kaggle dataset (Makridakis and Howard, 2020 to compare BaltoBot with other probabilistic prediction methods. The dataset has five years of sales data from ten Walmart stores, and the task requires predicting the (heavy-tailed) number of units sold given product's 's attributes and previous sales. We use the exact same data preparation used for TreefTuser experiments, which yields 1k products, 120k training samples, and 10k test samples. As in the Treeffuser ovaluation (Beltran-Velee et al., 2024), we evaluate probabilistic predictions with the continuous ranked probability score (CRPS), and evaluate the conditional mean predictions with the mean absolute error (MAE) and root mean-squared error (RMSF)..
Beetran-Velee el al. 2024
9

---

Published in Transactions on Machine Learning Research (7/2025)
(B)
Runtime comparison
Treefluser
Total
Observed
Memood
Training
Treefhuser samples
Baltobol
Sempling
-2
BatobolabPPN
Time (8)
Observed data Ballobot samples
(C)
density
Rallobot
BalloboTabPPN
Observed data Baltobo FabPNN samples
-2
Figure 3: Comparison of TreefTuser and our approach on wave 8ynthetic data with 5000 samples. (A) Probabilistic predictions for Treeffuser (top), BaltoBot (center), and BaltoBoTabPPN (bottom). (B) Runtime comparison for the different methods. (C) Estimated pdf from our methods at X 2, depicted as the vertical dotted line in (A).
Original data
Treeffuser
Baltobot
Baltobo TabPFN
6
4
2
0-
-2
0
2
0
2
0
2
Figure 4: Comparison of Treeffuser, BaltoBot, and BaltoBoTabPPN on Poisson-distributed data. The input variable is the x-axis, while probabilistic predictions are shown the y-axis.
For full comparability, we follow the Treeffuser evaluation setup (Beltran-Velee Velez et al. 2024 and evaluate CRPS by generating 100 samples from our estimators' X(y]X) for each X the testset; and for MAE and RMSE, we estimate the conditional means Efy X using 50 samples. For comparability, for this (and only this) dataset, also evaluate BaltoBot with hyperparameter tuning, using the same setup used for all other methods (10 folds, each with 80%-20% train-validation split, and 25 Bayesian optimization iterations). We also compare Treeffuser, BaltoBot, and BaltoBo TabPFN when run without hyperparameter tuning.
We report results Table addition to ours' and Treeffuser, we report results for Deep Ensembles Lakshminarayanan et al. 2017), IBUG (Brophy and Lowd (022). NGBoost Poisson (Duan al. 2020),
and Quantile Regression Forests (Meinshausen and Ridgeway 2006 For methods other than our own,
We optimize following XGRoost hyperparameter spaces: learning log-unfform0005,.00 0.5), {0, 25, 50), and subsample (00g-
10

---

Published in Transactions on Machine Learning Research (7/2025)
Table 5: Sales forecasting evaluation on M5 dataset. We highlight the best 2 methods for each metric. The best of Treeffuser versus ours (with tuning) is boldface blue; the best of Treeffuser versus ours (without luning) is boldface brown.
Method
CRPS x10-
RMSE x10%(4)
MAE x x10
Deep Ensembles
7.05
2.03
0.97
IBUG
8.90
2.12
1.00
NGBoost Poisson
6.86
2.33
0.99
Quantile Regression Forests
7.11
2.88
1.01
Treelfuser
6.44
2.09
0.99
BaltoBot
6.44
2.07
0.98
Treeffuser (no tuning)
6.62
2.09
0.99
BaltoBot (no tuning)
6.69
2.19
0.98
BaltoBoTabPFN (no tuning)
6.66
2.06
0.97
we report the metrics provided in Table 2 of (Beltran- Velez et al. 2024). Overall, our proposed methods outperform previous methods at combining excellent performance on both conditional distribution prediction and conditional mean prediction. Treeffuser and BaltoBot (both with tuning) tie for first-place according to CRPS, yet BaltoBot outperforms Treeffuser on RMSE and MAE. The winners on conditional mean metrics (RMSE and MAE) are Deep Ensembles and BaltoBoTaabPFN yet BaltoBoTTabPNN (no tuning) strongly outperforms Deep Ensembles CRPS.
4 Limitations and Future Work
Limitations While UnmaskingTres Trees leads overall on the tabular imputation benehmark, MissForest still outperformed on the metrics based on Wasserstein distance to train and test dataset distributions. And Forest-Flow still won on vanilla (i.e. no missingness) generation benchmark. It remains to be seen whether a single method can be developed which wins on all scenarios and metrics.
Our proposed BaltoBot method would benefit from more principled method for selection of the meta-tree height H. Unlike with standard flat quantization where having large number of bins cause to make catastrophically wrong predictions, BaltoBot "knows" proximities among meta-tree leaves. This means that making // bigger tends not to cause major errors, so it's better to err on the side of larger H rather than small H. Still, the main drawbacks with incrcasing H are that (1) training takes longer, and (2) imputations and generations are less diverse. Deeper theoretical analysis of these trade-ofT's and a more principled approach for choosing the height would improve the ease-of-use and potentially the performance of our method.
While BaltoBoTabPFN performed well on probabilistic prediction tasks, when used subroutine in UnmaskingTabPPN, it is very slow and experienced out-of-memory errors on the (Jolicoerr-Martineee el al. 2024b benchmark on our machine. Further improvements either to it, or to how it is employed, are needed to make it practical for all but the smallest datasets.
With regards to scalability, the runtime complexity UnmaskingTrress training is cubic in terms of feature dimensionality D. This compares to quadratic complexity for DoffstDiffusion; this would also be the complexity for single-order autoregressive modeling, which would support only generation rather than arbitrary imputation. This means that approach is primarily applicable for small-sized tabular datasets, or for scenarios where inference time is more important than training time.
Related to the above remarks, we would like to emphasize that our proposed approach is aimed al and evaluated on smaller-sized tabular datasets. It also evaluated via "out-of-the- box" performance, being aimed at users lacking the resources for large deep learning models or hyperparameter optimization. For users with access larger tabular datasets and more extensive computing resources, recent deep learning methods like Tabsyn (Zhang et al. 2024 would be expected to perform better.
11

---

Published in Transactions on Machine Learning Research (7/2025)
Future Work Reducing the training complexity with respect to the number of features is key next step. One possibility would be to use an optimized, rather than random, selection of feature orderings at training time (Shih et al. 2022). Another possibility would be to use multi-output trees to train a single XGBoost model for all BaltoBot tree nodes and all features, similarly to a recently-proposed approach for speeding (Cresswell and Kim 2024). In addition to improving scalability, BaltoBot's core idea of binary partitioning could be combined with deep learning approaches. For example, one might equip a neural network with a hierarchical softmax head (Morin and Bengio, 2005) for modeling continuous outputs without losing proximity information among bins. Finally, future work could theoretically analyze BaltoBot, empirically evaluate it on uncertainty quantification and probabilistic forecasting tasks, and extend it to multiple dimensions.
up Forest
5 Discussion and Related Work
Diffussin modeling has recently gained popularity in tabular ML (Zheng and Charoenphakdee 2022 Jolicoeur-
Martineau et al. 2024b Beltran- Velez et al. 2024 Kotelníkov et al. 2023). Our proposed approach is an instance of the autoregressive discrete diffusion framework (Hoogeboom et al. 2021 ) iistaacces of which have shown success in a variety of tasks (Yang, 2019; Austin al.) 2021 Kitouni et al. 2024 Jolicoeu--Martnneau et al. 2024a). Yet our results call into question whether diffusion is beneficial for tabular conditional generation, whether autoregression is sufficient for our setting. It has been observed that diffusion is autoregression in frequency space, progrossing from low frequencies to high frequencies, which makes it a good match for image data with its power law spectra [Rissaen et al. 2022 Dieleran 2024 Stewart 2024). In tabular datasets without this phenomena, we would expect diffusion modeling to be less advantageous.
Why is ForcstDiffusion better at vanilla generative modcling, while UnmaskingTrees is better on missing data problems? We offer two speculative explanations. First, imputation is a conditional modeling scenario, except. that you do not know the partition of the features into imput features and output features priori. One could address imputation by learning all possible 2 conditional distributions, but this is impractical for large D, so one would prefer to learn a single joint distribution. Both autoregression and diffusion are ways of learning a joint distribution; because autoregression does so by learning conditional distributions, it is more suited to the conditional modeling imputation setting. Second, for missing data, diffusion has a train-inference gap: during training, observed features begin the reverse process from (0 1); during inference for imputation, observed features begin the reverse process at their actual values. On the other hand, the advantages of diffusion modeling (no quantization error, holistic generation, needing only estimated score function rather than well-calibrated conditional distributions) give it superiority when these problems can be avoided.
Despite their strong outperformance on other modalities, deep learning approaches have laboured against gradient-boosted decision trees on tabular data (Shwartz-Ziv and Armon 2022
Jolicocur-Martnncau et al.) 2024b). Previous work (Breejen et al., 2024) suggests that tabular data requires an inductive prior that lavors sharpness rather than smoothness, showing that TabPFN (Hollmann et al. 2022) (the leading deep learning tabular classification method) can be further improved with synthetic data generated from random forests. We anticipate that our XGBoost classifiers may be swapped out for future variant of TabPFN that learns sharper boundaries and handles missingness.
We also note that MissForest (Stekhoven and Bühlmann 2012), hailing from statistical literature on multiple imputation, has yet to be completely dethroned. Future progress in tabular conditional generation may require going back to the well of this traditional literature. As one example, we observe that MissForest exploits fcature missingness fraction information, but we are not awarc of any "machine lcarning" approaches which do so. The statistical literature has also previously explored the value of conditional modeling for joint modeling (Gelman and Raghunathan 2001 Liu et all 2014 Kropko et al.) 2014). Indeed, our Unmasking" Trees approach, and all autoregressive modeling, is presaged by the full-mechanism bootstrap Efron 1994).
Finally, we observe where randomness enters into our generation process compared to previous work. Flow- matching (Liu et al. 2022 Albergo and Vanden-Ei Eijjnen 2022 Lipman et al.) 2022) (used in Forest-Flow) injects randomness solely at the beginning of the reverse process via Gaussian sampling, whercas diffusion modeling (Sohl-Dickstein et al.] 22015] Song and Ermon
2019) (used in Forest VP) injects randomness both at the beginning and during the reverse process. In contrast, because our method starts with fully-masked
12

---

Published in Transactions on Machine Learning Research (7/2025)
sample, it injects randomness gradually during the generation process. First, we randomly generate the order over features for unmasking. Second, we do not "greedily decode" to the most likely leaf in the meta-tree, but instead sample according to predicted probabilities. Third, for contimuous features, having sampled a particular meta-tree leaf bin, we sample from within the bin, treating it as a uniform distribution.
6 Conclusions
We proposed tree-based auloregressive modeling of labular data, especially for data with missingness. For the subproblem of conditional probabilistic prediction of individual variables, we presented hierarchical partitioning method with benefits over vanilla quantization and diffusion-based probabilistic prediction. We then considered each of these as meta-algorithms that enable pure in-context learning- based modeling using TabPFN as base classifier On benchmark for out-of-the box performance tabular data, we showed leading results for imputation and state-of-the-art results for generation given data with missingness.
13

---

Published in Transactions on Machine Learning Research (7/2025)
References
Michael Albergo and Eric Vanden-Eiinden Building normalizing flows with stochastic interpolants. arXiv preprint ar Xiv:2200. 15571, 2022. 12
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems,
34:17981-17993, 2021. T2
Nicolas Beltran- Alessandro Antonio Grande, Achille Nazaret, Alp Kucukelbir, and David Blei. Tre- effuser: Probabilistic predictions via conditional diffusions with gradient-boosted trees. arXiv preprint
arXiv:2406.07668, 2024.
2 10
12
Felix den Breejen, Sangmin Bae, Stephen Cha, and Se- Young Yun. Why in-context learning transformers are tabular data classifiers. arXiv preprint arXiv:2405. 13396, 2024.
Leo Breiman. Random forests. Machine learning, 45:5-32, 2001.
Jonathan Brophy and Daniel Lowd. Instanco-based uncertainty estimation for gradient-boosted regression trees. Advances in Newral Information Processing Systems, 35:11145 11159, 2022. 10
Samuel Buck. A method of estimation of missing values in multivariate data suitable for use with an electronic computer. Journal of the Royal Statistical Society: Series B (Methodological), 22( 306, 1960.
Rodrigo Castellon, Achintya Gopal, Brian Bloniarz, and David Rosenberg. Dp-tbart: A transformer-based autoregressive model for differentially private tabular data generation. arXiv preprint orXiv:2307.10430
2023. 2
Tianqi Chen and Carlos Guestrin. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd acm sighdd international conference on knowledge discovery and data mining, pages 785-794, 2016. 3
Jesso Cresswell and Tacwoo Kim. Scaling up diffusion and flow-based xgboost models. ar preprint
arXiv:2408. 16046, 2024. 12
Sander Dieleman. Diflusion is spectral autoregression, 2024. URL https: http://aaneerraa2224//////// a1/2024/99/02/
sppeccaal-uttoeegeesioon... html
12
Tony Duan, Avati Anand, Daisy Yi Ding, Khanh K Thai, Sanjay Basu, Andrew Ng, and Alejandro Schuler. Ngboost: Natural gradient boosting for probabilistic prediction. In International conference on machine learning, pages 2690-2707 2700. PMLR, 2020. 2 I0
Bradley Efron. Missing data, imputation, and the bootstrap. Journal of the American Statistical Associalion, 89442))4463 475, 1994. 12
Ronald A Fisher. The use of multiple measurements in taxonomic problems. Annals of eugenics, 7(2):179 188, 1936.
Andrew Gelman and Trivellore E Raghunathan. Using conditional distributions for missing-data imputation. Statistical Science, 15:268 69, 2001. 12
Manbir Gulati and Paul Roysdon. Tabmt: Generating tabular data with masked transformers. Advances in Neural Information Processing Systems, 36, 2024.
3
7
Trevor Hastic, Rahul Mazumder, Jason D Lee, and Reza Zadeh. Matrix completion and low-rank svd via fast alternating least squares. The Journal of Machine Learning Research, 16(1):3367 3402, 2015.
David P Hofmeyr. Fast exact evaluation of univariate kernel sums. IEEE transactions on pattern analysis and machine intelligence, 43(2):447-455 2019. 5
14

---

Published in Transactions on Machine Learning Research (7/2025)
Noah Hollmann, Samuel Müller, Katharina Eggensperger, and Frank Hutter. Tabpfn: A transformer that solves small tabular classification problems in a second. or Xiu preprint or arXiv:2007.0188 2022. 2
Emiel Hoogeboom, Alexey A Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. arXiv preprint 3rXiv:2110.02097, 2021. 12
Harry Joe. Dependence modeling with copulas. CRC press, 2014.
Alexia Malicoeur-Marinnauu Aristide Baratin, Kisoo Kwon, Boris Knyazev, and Yan Zhang. Any-property- conditional molecule generation with self-criticism using spanning trees. arXiv preprint arXiv:2407.09357
2024a. 12
Alexia Jolicoeur-Martineau Kilian Fatras, and Tal Kachman. Generating and imputing tabular data via diffusion and dow-based gradient-booste trees. In International Conference on Artificial Intelli- gence and Statistics, pages 1288 1296. PMLR, 2024b. URL htttt:///ttthtii cwmccoo//aaaaaa /aaa//anunngAAAMMttrraaalll ForestDiffusion
18
Jayoung Kim, Chacjcong Lee, and Noscong Park. Stasy: Scoro-based tabular data synthesis. arXiv preprint ar arXiv:2210.04018, 2022a..
Jayoung Kim, Chaejeong Lee, Yehjin Shin, Sewon Park, Minjung Kim, Noseong Park, and Jihoon Cho. Sos: Score-based oversampling for tabular data. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 762 772, 2022b. T
Ouail Kitouni, Niklas Nolto, James Hensman, and Bhaskar Mitra. Disk: A diffusion model for structured knowledge. arXiv preprint Xrx:2233 05253, 2023. 3
Ouail Kitouni, Niklas Nolte, Diane Bouchacourt, Adina Williams, Mike Rabbat, and Mark Ibrahim. The factorization curse: Which tokens you predict underlie the reversal curse and more. arXiv preprint
(rXiu:2406.05183 2024.
3 12
Akim Kotelnikov, Dmitry Baranchuk, Ivan Rubachev, and Artem Babenko. Tabddpm: Modelling tabular data with diffusion models. In International Conference on Machine Learning, pages 17564-17579. PMLR, 2023.
David M Kreindler and Charles J Lumsden. The effects of the irregular sample and missing data in time series analysis. In Nonlinear Dynamical Systems Analysis for the Behavioral Sciences Using Real Data, pages 149 172. CRC Press, 2016. [
Jonathan Kropko, Ben Goodrich, Andrew Gelman, and Jennifer Hill. Multiple imputation for continuous and categorical data: comparing joint multivariate normal and conditional approaches. Polilical Analysis, 22(4), 2014. 12
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017. [0
Quoc V Le, Tim Sears, and Alexander J Smola. Nonparametric quantile regression. Technical report, Technical report, National ICT Australia, June 2005. Available at http://sml .... ..... 2005.
Xiaotian Li, Shuzhe Wang, Yi Zhao, Jakob Verboek, and Juho Kannala. Hicrarchical scene coordinate classification and regression for visual localization. Proceedings of the IEEE/CVE Conference on Compuler Vision and Pattern Recognition, pages 11983-11999 2020. 2
Yi Liao, Xin Jiang, and Qun Liu. Probabilistically masked language model capable of autoregressive generation in arbitrary word order. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 263 274, 2020. 3
Yaron Lipman, Ricky TQ Chen, Heli Ben- Hamu, Maximilian Nickel, and Matt Lc. Flow matching for generative modeling. arXiv preprint arXiv:2210 00.027,10.02777, 2022. 12
15

---

Published in Transactions on Machine Learning Research (7/2025)
Roderick JA Little and Donald B Rubin. Statistical analysis with missing data. John Wiley & Sons, 2019.
Jingchen Liu, Andrew Gelman, Jennifer Hill, Yu-Sung Su, and Jonathan Kropko. On the stationary distribution of iterative imputations. Biometrika, 101 (1):155-173, 2014.
12
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint S1t::2209.030000 2022. 12
Stuart Lloyd. Least squares quantization in pem. IEEE transactions on information theory, 28(2):129 137, 1982. g
Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the TEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 11461 11471, June 2022.
Junwei Ma, Apoorv Dankar, George Stein, Guangwei Yu, and Anthony Caterini. Tabpfgen tabular data generation with tabpfn. arXiv preprint arXiv:2406. (5.::24660 2024.
Spyros Makridakis and Addison Howard. M5 forecasting accuracy, 2020. URL https://kaggle. com/ ataperittooonnn--ffceeetttt
g
Alexander März. Xgboostlss extension of xgboost to probabilistic forecasting.
arXiu preprint
arXiv:1907.03178, 2019. 2
Calvin McCarter. The kernel density integral transformation. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
Calvin McCarter. Towards backwards-compatibb data with confounded domain adaptation. Tronsactions on Machine Learning Research, 2024. ISSN 2835-8856. URL https://openrevieex wtt.errrr-....ettotouumm-------
GSp2WC7q0r
1
Duncan McElfresh, Sujay Khandagale, Jonathan Valverde, Vishak Prasad C, Ganesh Ramakrishnan, Micah Goldblum, and Colin White. When do neural nets outperform boosted trees on tabular data? Advances in Neural Information Processing Systems, 36, 2024. 2
Nicolai Meinshausen and Greg Ridgeway. Quantile regression forests. Journal of machine learning research, (6), 2006. 10
Frederic Morin and Yoshua Bengio. Hicrarchical probabilistic neural network language model. In International workshop on artificial intelligence and statistics, pages 246 252. PMLR, 2005. 12
Boris Muzellec, Julie Josse, Claire Boyer, and Marco Cuturi. Missing data imputation using optimal transport. In International Conference on Machine Learning, pages 7130-7144 PMLR, 2020.
Mohamed Ragab, Emadeldeen Eldele, Min Wu, Chuan-Sheng Foo, Xiaoli Li, and Zhenghua Chen. Source-free domain adaptation with temporal imputation for time series data. In Proceedings of the 29th ACM SIGKDD conference on knowledge discovery and data mining, pages 1989-1998, 2023. I
Severi Rissanen, Markus Heinonen, and Arno Solin. Generative modelling with inverse heat dissipation. arXiu preprint arXiv:2206. 19397, 2022. 12
Donald B Rubin. Multiple imputation after 18+ years. Journal of the American statistical Association, 91 (434):473 489, 1996.
Andy Shih, Dorsa Sadigh, and Stefano Ermon. Training and inference on any-order autoregressive models the right way. Advances in Neural Information Processing Systems, 35:2762-2774 2775, 2022. 12
Ravid Shwartz- Ziv and Amitai Armon. Tabular data: Deep learning is not all you need. Information Fusion 81:84-90, 10, 2022.
12
16

---

Published in Transactions on Machine Learning Research (7/2025)
Jascha Sohl-Dickstein, Eric Weiss, Niru Mahoswaranathan, and Surya. Ganguli. Deep unsupervised learning using nonequilibrium thermodynanics. In International conference machine learning, pages 2256-225622 2265. PMLR, 2015.
12
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019. 12
Olivier Sprangers, Sebastian Schelter, and Maarten de Rijke. Probabilistic gradient boosting machines for large-scale probabilistic regression. In Proceedings of the 27th ACM SIGKDD conference on knowledge discovery & dala mining, pages 1510- 15200 2021. 2
Daniel J Stekhoven and Peter Bühlmann. Misslorest non-parametric missing value imputation for mixed-type data. Bioinformatics, 28(1): 118, 2012.
7
Riley Stewart. trasformers are kiki, diffusion is bouba, and language is pointier than images, 2024. URL http::/ww..ooooo/......ss////aa/
12
Olga Troyanskaya, Michael Cantor, Gavin Sherlock, Pat Brown, Trevor Hastie, Robert Tibshirani, David Botstein, and Russ B Altman. Missing value cstimation methods for dna microarrays. Bioinformatics, 17 (6):520 525, 2001.
Boris Van Breugel, Trent Kyono, Jeroen Berrevoets, and Mihaela Van der Schaar. Decaf: Generating fair synthetic data using causally-aware generative networks. Advances in Neural Information Processing Systems, 14:22221-222333
2021.
Stef Van Buuren, Hendriek C Boshuizen, and Dick L Knook. Multiple imputation of missing blood pressure covariates in survival analysis. Statistics in medicine, 18(6):681 694, 1999.
A Vaswani. Attention is all you need. Advances in Newral Information Processing Systems, 2017. 2
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICMI-11), pages 681-688. Citescor, 2011.
Samuel Von Wilson, Bogdan Cebere, James Myatt, and Samuel Wilson. AnotherSanWilson/ mwwp//orccceee:... 7428632
Release for Zenodo DOI, December 2022. URL
Lei Xu, Maria Skoularidou, Alfredo Cuesta-Infante, and Kalyan Veeramachaneni. Modcling tabular data using conditional gan. Advances in neural information processing systems, 32, 2019.
Z Yang. Xlnet: Generalized autoregressive pretraining for language understanding. ar preprint arXiv:1906. 082337 2019. 2
12
Jinsung Yoon, James Jordon, and Mihaela Schaar. Gain: Missing data imputation using generative adversarial nets. Tn International conference on machine learning, pages 5689-5698. PMTRR, 2018a.
Jinsung Yoon, James Jordon, and Mihaela Van Der Schaar. Ganite: Estimation of individualized treatment effects using generative adversarial nets. In International conference on learning representations, 2018b.
Hengrui Zhang, Jiani Zhang, Zhengyua Shen, Balasubramaniam Srinivasan, Xiao Qin, Christos Faloutsos, Huzefa Rangwala, and George Karypis. Mixed-type tabular data synthesis with score-based diffusion in latent space. In The Twelfth International Conference on Learnång Representations, 2024. URL
11
Zilong Zhao, Aditya Kunar, Robert. Birke, and Lydia Y Chen. Clab-gan: Effective table data synthesizing. In Asian Conference on Machine Learning, pages 97- PMLR, 2021. &
Shuhan Zheng and Nontawat Charoenphakdee. Diffusion models for missing value imputation in tabular
data. arx preprint rrXiw:2210.111.1 2022. I2
17

---

Published in Transactions on Machine Learning Research (7/2025)
A Ablation experiment with imputation raaw scores
Raw scores (shown in Table 60 demonstrate that UnmaskingTross on its own improves upon Forest-VP's diffusion approach. We also see that KDI quantization (with 20 bins) contributes to improvement beyond k-Means (also 20 bins), and that BaltoBot yields even further improvement.
Table 6: Raw scores from ablation study for tabular data imputation (27 datasets, 3 experiments per dataset, 10 imputations per experiment) with 20% missing values. Shown are raw scores mean (standard-error). Overall best is highlighted; better of Forest-VP versus ours is boldface blue. See Table [ for column meanings.
MinMAE
AvgMAE 4
Werain 4
Wrest 4
MAD1 1
R
Flamp t
Potas
Conyatt
RNN
0. (0.03)
0.10 (0.03)
0.42 (0:08)
[0.40]
[0]
0.53 (0.08)
0.75 (0:04)
1:27 (0:25)
0.4 (0.11)
ICE
D.1 (0.01)
0.21 (0.03)
0.52 (0.09)
1.99 (0.49)
0:69 (0.1)
0.5a [0.00]
0.74 (0.04)
1.05 (0.29)
0.39 (0.09)
MICE- Foreet
0:08 (0.02)
0.13 (0.03)
0.34 (0.07)
1.86 [0.48]
0.29 (0.06)
0.61 (0.1)
0.78
[0.04)
0.61 (0.2)
0.75 (0.11)
MiasPorest
0.1 (0.03)
0:12 (0.03)
0:32 (0.07)
1:35 (0.48)
0.1 (0.03)
0.61 (0.1)
0.76
(0.94)
0.62 (0.22)
0.79 (0.05)
Softimpate
0.22 (0.03)
0.22 (0.03)
0.53 (0.07)
1.30 (0.48)
(0)
0.58 [0.09]
0.74 (0.04)
1.18 (0.34)
0.31 (0.09)
OT
0.14 (0.02)
0.19 (0.03)
0.56 (0.1)
1.98 (0.49)
0.28 (0.05)
0.50 (0.1)
0.75 (0.04)
1.00 (0.27)
0.30 (0.12)
GAIN
D.1G (0.03)
0.17 (0.03)
0.40 (0.11)
1.90 [0.51]
0.01 (0)
0.6 (0.1)
0.75 (0.04)
1.04 (6.25)
0.54 (0.12)
Vorest-VP
0.14 (0.04)
0.17 (0.03)
0.55 (0.13)
1.96 (0.5)
0.25 (0.03)
0.61 (0.1)
0.74 (0.04)
0.81 (0.25)
0.57 (0.14)
(TTwwa-MMaan
0.1 (0.02)
0.15 (0.03)
0.43 (0.09)
1.9 (0.5)
0.28 (0.06)
0.61 (0.1)
0.78 (0.04)
0.63 (0.21)
0.72 (0.13)
Utrees KDI
0.1 (0.02)
0.14 (0.03)
0.42 (0.09)
1.89 (0.49)
0.27 (0.06)
0.61 (0.1)
0.76 (0.04)
0.68 (0.24)
0.68 (0.14)
UThees
0.08 (0.02)
(0.03)
0.37 (0.08)
(0.48)
0.27 (0.07)
0.61
0.78 (0.04)
(0.19)
(0.13)
Oracle
(0)
D (0)
(0)
1.87 (0.49)
(0)
0.64 (0.09)
0.78 (0.04)
(0)
(0)
B Full dataset-level results
Full imputation results are in Table 7 Full generation results are in Table 8 Timing results are in Table I0 and depicted in Figure 5 Our method is relatively efficient at both imputation and generation. The datasets on which we are slowest for imputation are Libras (1976 seconds, N = 360, D = 90) and Bean (1929 seconds, N 13611, D 16), on our ancient 2015 iMac with 16Gb RAM. On Libras, VP immputen imputation took 12439 seconds (without RePaint) and 14715 seconds (with RePaint); on Bean, Forest took 898 seconds (without. RePaint) and 1318 seconds (with RePaint), datasets on which we are slowest for generation are also Libras (2987 seconds) and Bean (4346 seconds). On Libras, ForestFlow generation took 9481 seconds and Fo took 9042 seconds; on Bean, ForestFlow Flow took 869 seconds and Forest took 947 seconds, once again on their much more powerful computing cluster.
their cluster of 10-20 CPUs with 64-256Gb of RAM. The
Table 7: Full imputation results for UnmaskingTreee on Molicoeur-Martineuu et al. 2024bb benchmark.
Dataset
MisMAR
AvgMAE
Peiay
W. Wiram
Wisst
Var
MAD (mean)
MAD
R² ~
F1
6.91e-02
0.00++00
2.40e-01
2.65e-03
1.46-011
1.22--01
0.00e+00
9.53--01
9.45m-02
1.31e-01
0.00e4400
0.00-100
3.54e-01
6.64n-03
2.37e-01
1.99-01
0.00-100
9.37--1
parkinsons
6969 02
6.32e 02
0.00c+000
0.00c+00
2.94e-01
1.71e400
2.73e
1.23e-0
1.03c-01
0.00c+00
8.30c
climate
2.35e-07 01
3.35e-01
0.00c+000
0,00c+00
1,22-+00
3,87c400
4.38c 02
7.91e-01 01
6.88c-01
0,00c+00
7.08c-01
concrese
2. 2.33e-122
6.196-02
1.17c4002
2.44c-01
7,95e-02
6.05e-01
6.61e-03
1.59-011
1.30-01
7.55e-01
0.00c4.00
yacht kydrodynamies
2. 2.72e-0/2
6.53e-02
8.78c401
9.62e-01
6.46e-02
6.11e-01
1.22c-02
1.90-01
1.45c-01
8.96ie-1
0.00c+400
airfoil self noise
2 2.422-122
6.64e-02
3.37c+000
1.00c+000
4.60c-12
2.50e-01
1.25e-02
1.29e-01
1.84c-01
7.24c-01
0,00e+00
connectionist sonar
0.86n-02
1.18--01
0.00e4000
0.00--400
1,43e+00
8.51e+400
6. 6.14e-03
2.22+-01
1.88r-01
0.00+++00
7.99e-01
fonoephere
8.52n-02
1.18
0.00-4000
0.00-4400
7.82--01
4.44e-100
1,47e-02
2.54e-01
2.02--01
0.00-+000
0.10--0
gaar bindegradation ion
1.53e-022
2.34m-02
0.00
0.00x-100
1. 1.92e-01
1.39e-000
1.25e-03
5.35e=02
4.36e-02
0.00-+100
8.49c-01
seeds
5.36e 02
8.45e-02
0.00c+00
0.00c+00
1. 1.22c-01
4.78c-01
3.37c-03
1.7de-01
1.48e-01
0.00c+00
8 53e 01
glass
1.96e-02
7.59e-02
0.00c+000
0. 0.00c+000
I. 1.40--01
6.42e-01
5. 5.09c-03
1.4e-011
1.17e-01
0.00e+000
5.15-001
5.15e-112
8.00e-02
0.00c++00
0.80c-100
1. 1.09e-01
4.04c-01
3.80c-03
1.54e-01
1.30c-01
0.00le
6,53--01
yeast
4.35e-122
7.30e-02
0.00c+000
0.80c++00
1.07e-U1
3.19e-01
3.58c-03
1.73e-01
1.50e-01
0.00c+300
4.45c-01
Libras
0/2
3.84e-02
0.00c+00
0.00c+000
6.57e-01
8.93e+00
8.11c.04
8.17e-02
7:06c-02
5.69c-01
planning relax
8.41n-022
0.00++00
blood transfusion
3.44o-1/2
6.69e-02
0.00.-400
0.00+100
3.21e-02
1.126-01
4.80a-03
1.58e-01
1.32-01
0.00+.00
5.87-01
breast
96e-02
16e-00
0.00e+00
0.00c+00
3.10
1,85e+000
1.21e-03
101e-01
6.73e-02
0.00c+00
9.59c
connectionist vowel
3.30--02
9.39e-02
0.00++00
0,00e+00
1.86e-01
7.25e-01
5.53e-03
2.46--01
2. 14e-0
0.00++00
6.6ie-01
concrete slump
1.25e-01
1.85--01
4.76e+01
7.25e-01
2.64e-01
1. 1.16-+000
1.48e-02
3.40e-01
2.75--0
6.75m-01
0.00++00
wine quality red
4. 1888--2
7. 7.16e-02
2 2.01e+01
1.00c+00
1.41e-01
5.17c-01
3.52c-03
L83e-01
1.58c-01
3.06e-01
0.00++000
wine quality white
3 3.41e 02
6.45e 02
7.7
0118c.01
1.40c-01
4,53e-01
3.36c 03
1.91e-01
1.68c-01
3.38c-01
0.00e+00
california
1.97n-02
4.37n-02
2.32++001
5.93--01
0.00+4000
0.00c4000
4.38e-03
1.58a-01
1.40--01
6.55e-01
0.00+400
bean
1.02o-172
2.06e-02
0.00c
0.00
0.00e-100
0.00c-100
1.05e-03
6.42e-022
5.80k-02
0.00
7.82<-01
tictactoe
2.96e-D1
5.11e-01
0.00c+000
0.00c-000
T-76e-01
1.93e4400
2 2.45--02
1.47ei-D0
1,13e400
0.00c
8.23c-11
congrese
1.73e-01
2.77e-01
0.00e--00
0.00++000
6.03e-01
2.38e+000
9. 0.76e-03
5.86e-01
4.33e-01
0.00e+000
9.35e-01
car
4.04e-01
6.85e-01
0.00++400
0.00++000
4.84e-01
1.07++000
2.95e-02
2.06++-00
1.65e400
0.00++000
8.01e-01
18

---

Published in Transactions on Machine Learning Research (7/2025)
Table 8: Full generation results for UnmaskingTroos on Molicoeur-Martinee et al.)
2024b) benchmark.
Datazot
Werain
Weest
covtrain
contest
R
Take
fake
Flais
Poiax
Covyate
Ielr
2.34c-01
3:41-011
5:75c-01
5:10e-01
0:002
9-200-01
4.23c-01
0.00c
0.00c T00
sinc
1.000100 100
1.33c1000
0.09c-01
9.37c-01
0:00=100
0.100-01
8.46c-01
0.00c
0.00-100
parkinsons
1.34++00
1.77c+00
7.49c.01
9.08c-01
0.00.100
7:38--01
3.50c-01
0.00++00
0.00++00
model craskes
yochat compresion
4.63c-01
0.21e-01
5.11c-01
5.10e-01
0.79c.01
0.00c+:00
4.15c-01
1.50c402
2.00c-01
4.20c-01
0.36x-01
6.13c-01
7.89c-01
5.48c.01
0.00c
5.14.-01
4.48c-01
strfall zelf notes
1.93e-01
2. 2:936-01
6.38e-01
K.96e-01
6.08e-01
0.00++000
4.65e-01
1.97++01
4.78e-01
connectiontst beach
7.21++00
8.988+000
6.87e-01
8.59m-01
0.00.400
7.20=-01
5. 3.69e-01
0.00+4-00
0.00+000
joncsphere
0.84c100
4.06c+00
6.11c-01
7.94c-01
0. 0.00=100
8.07c-01
4.22c-01
0:00-100
0 0:00-0000
quar biodegradat
1.34-+00
1.0e+000
4.81c-01
6.19e.01
0. 0.00=+00
8. 8.02
4.44c-01
0.00c+00
0 0.000000
sends
3.51e-01
5 5-48--01
6.95c-01
9-63e-01
0. 0.00++00
5. 09
3. 3.00--01
0.00r+00
0 0.00++00
xlasa
4.52c-01
7.12c-01
6.27c-01
9.35c-01
0.00c
4.41c-01
5.86c-01
0.00c:000
0.00c +00
ccoli
2.50c-01
4. -.36-01
5.99c-01
0.58c-01
0.00c 100
6.16o-01
3. 3.78.-01
0.00c 00
0.00c 00
yeast
2.4-e-011
3 3.19--01
8.33e-01
9.4f--1
0.00+400
3.02-01
4.33e-01
0.00=+000
0.00x+00
librae
1.01++01
1.16++d1
4.65e-01
6.43m001
0.00+400
3.54-011
3.53e-01
0.00e+400
0.00x+00
plansing relax
1.022+000
1.47e+000
9.22e-01
9-98--01
0.00s400
1.56-011
3.07e-01
0.00s-000
0.000+000
blood transfucion
1.00c-01
01
0.63c-01
9.56e-01
9.00++00
6.B0e-01
4.07e-01
0.00c+000
000+00
diagnastic
1.50++00
9.12m001
0.00.+00
0.00+000
connectioniat
8.87c-01
3.04c-01
6.76.-01
3.43.-01
0.00c
0.00c 100
0.24c-01
1.20c| 00
8.71c-01
5.67c-01
0.34c.01
0.00c
3.52.-01
4.57c
5.75c-01
quality
4.30c-01
6.40x001
8.63c.01
9.67c-01
3.46.-01
0.00c
000
4.01c-01
01
7.94c-01
quality white
4.28e-01
4 1.97e-01
8.28e-01
9.55e-01
2.32e-01
0.00++000
4.48e-01
1.8e+002
2.53e-01
californin
0.00s+00
0.00c+000
6.22e-01
9.03e-01
3.05e-01
0.00c4000
4.30c-01
1.7be+002
1.70c-01
bean
0.00++00
0.00++00
3.35e-01
7:53e-01
0. 0.00++00
8.16--01
3.97c-01
0.00-+000
0 00e+000
tletacson
9.49e-01
1.95e+00
8:23e-01
6.25e-01
D. 0.00++00
8.31=001
2.74e-01
0.00r+000
0. 0.-0000
congress
1.38++00
2-16++00
9.11e-01
9-16m-01
D. 0.00=400
9-17e-11
2.87--01
0.00x+000
0. 0.00+000
4.61c-01
1.06c|
5.82c.01
6.20c-01
6.00s 100
7.90-01
5.02c-01
0.00c 00
0.00c| 00
Table 9: Full generation with 20% missingness results for Unmasking Trees on JJlicocur-Martineau et al.
2024bb benchmark.
Datasnt
Weravn
Wicsl
coverain
conteat
x>3
Lake
tuke
Plais
Poioy
Conrete
Teln
2.56e-01
3.54e-01
6:30e-61
5.71e-01
0. 0:00m+000
B-42--01
4. 4.30e=01
0.00e+-00
0. 0.00e+000
wine
1.16++00
1.55e+-00
8.70c-01
8.98e-01
0.00s4000
9.17c-01
5.71e-01
0.00c:000
0.00c-000
parkinsons
1.41c100
1.50c
6.49c-01
8:08c-01
0.00=100
7.0le-01
3.76c-01
0.00c
0.00c 100
climate model crachar
3.333+000
3. 3.88++00
8.58e-01
9-53e-01
0.00+400
5.29--01
3.20e-01
0.00--000
0.000+000
concrate compression
4.71e-01
4.83e-01
7.89e-01
6.57e-01
0.00++000
4.18e-01
1.6le-02
1.48e-01
yacht hodroddnamica
4.18c-01
1.25--01
5.80e-01
5-30e-01
8.50e-01
0.00++000
5.5be-01
1.04e-022
6.00e-01
airtell poiso
2.00c-01
00-200
6.0fi-01
6.02e-01
5.976.01
0.00++:00
4:50c-01
2:07c-01
3.45e-01
connectioniet
6.10e-01
8.65e-01
0.00--000
00-0+00
3.50c:100
4.69c1000
0.00.400
8.61e-01
4.32.-01
quae biedogrodation
1.37c
1.62c
4.54c-01
5.17c-01
0.00=
8.07c-01
4.47c-01
0.00c
0.00c
4.34c-01
0.86x00
8.05c-01
0.38c-01
0.00= 100
6.68.-01
8.37c-01
0.00c
glass
5.05e-01
7 1.33e-01
7.75e-01
8.93e-01
0.006400
166--11
3.87e-01
0.00e+-00
0.00++000
2.51e-01
3. 3.4be-01
8.22e-01
9.44n-01
0.00a400
2.01e-01
4.30e-01
0.00c-000
0.00+4000
lbras
1.00e+001
1e16e+01
4.60e-01
6.51e-01
0 0:00-+00
3. 01
3.63c-01
0.00-+00
0 00e-000
planning relae
1.07++00
1.48++000
9.01e-01
9.84e-01
0 0000++00
1.64m-01
3.37--01
0.00-+00
0 0.000000
blood trandfucion
9. 9.11e-02
1.38=-11
9.62e-01
9-56m-01
0.00++00
5.80-001
4.18--01
0.00++00
0 0.00+000
breast cancer dinguastic
1 00
1.01c| 00
7.48c-01
5.00c-01
0.00= 00
2.40-001
3.60--01
0.00c 00
0.00c 00
connection lionist beach vowel
7.42--01
0.16e-01
2.35c-01
6.01c-01
0. 0.00=100
6.46o-01
3.60.-01
0.00c 00
0. 0:00
concrate slump
7.18e-01
1.24e+000
5.25e-01
6.22e-01
4.10e-01
0.00++000
3.58e-01
5:25e-01
wine quality red
4.39c-01
b.4be-01
8.42e-01
9.64s-01
2:49c-01
0.00++000
4.52e-01
00
7.15e-01
wine quality white
4.28c-01
4.90e-01
6.15c-01
9.64p-01
2.52c.01
0.00c::00
4.48c-01
1.
2.33c-01
0.00++000
0.00++000
6.15e-01
6.99e-01
0.00-+000
4.31e-01
1.37--02
2.52e-01
Miclactoe
1.12:
1.06c
8.04c-01
6.65e-01
0.00=
7.40-01
8.20.-01
0.00c
100
1000
2.40c100
8.52c.01
0.07c-01
0:00=100
9.30--01
3.40c-01
0.00c:-00 00
0000000
5.62c-01
1.07e+000
4.95c-01
5.12e-01
0.00.+00
7.6le-01
3.36c-01
0.00c--00
00e++00
Table 10: Runtime results for Unmasking Trees on benchmark of 27 datasets.
Datasct
Samples
di Features
Impusation time (6)
Generation sime (6)
Ista
150
5.31
10:12
wine
176
13
26.76
49. 06
parkimonn
195
22
58. 96
105. 277
climate model craebax
540
18
103-73
207 70
concreto compresion
1036
47. 10
128. 10
yacht hydrodynamies
306
8.50
29. 36
atrioal aelf noise
1803
29.90
92.14
connectioniat bench sonar
208
GB
440.60
686.94
jonosphere
851
201.81
862.00
qear biodagradation
glase
214
17.12
76
ocoli
336
14.60
56
1451
02
150.00 83
libra
860
90
1975.78
2956.76
plansing velox
182
12
25.25
46.15
bload tranefuzion
748
18:18
37.24
braast cancer diagnostic
550
30
279 31
490 71
connectioniat bench vowel
590
10
79.55
179. 48
concreto elump
103
9.74
49
wine quality nel
1598
10
983. 8/
wine quality white
4508
367.08
800.00 90
california
20640
968.
2764.76 76
beau
13611
16
1929.
4340.00
435
1728
30.18
60.36 31
19

---

Published in Transactions on Machine Learning Research (7/2025)
(A)
(B)
4000
2000
300
100000
20000
50
1000
2000
10000 n_Samples
20000
30 r,dims
2000
4000
fLSamplos
7.4ms
dmt
dmn.maknn
Figure 5: Runtime in seconds compared to number of features and number of samples, for imputation (A) and generation (B) tasks.
20


Reviews:
[{"review_id": "C5xNgF29Gh", "reviewer": "Reviewer_cvRQ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n  - Autoregressive modeling with trees is a novel and interesting idea, and it's shown to be effective here.\n  - The paper includes extensive empirical validation.\n\n### Weaknesses\n  - The exposition in the methods section could be clearer.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "- A permutation-based autoregressive framework, _UnmaskingTrees_, with gradient-boosted trees for tabular data imputation and generation.\n  - A novel approach, _BaltoBot_, to model conditional distributions via a balanced meta-tree of binary classifiers that allows fast sampling and closed-form density estimates.\n  - Extensive numerical experiments that show that the two methods combined achieve SOTA performance on generation with missing data, are comparable to SOTA on imputation with missing data, and are competitive on vanilla generation.\n  - A python package. (Great choice of emojis!)", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "These suggestions are not critical, but they would help clarify the presentation and thus improve the paper.\n\nHere are some passages that I found confusing, and would appreciate if they were clarified/reviewed:\n\n  - In Section 2, I think this could be more effectively framed as a single-method paper. I found that introducing two methods broke the flow and added confusion. My point: in UnmaskingTrees, wouldn’t it make more sense to use BaltoBot for everything that isn’t categorical? BaltoBot works for both continuous and discrete data, so the use of XGBClassifier only makes sense for categorical inputs. That said, I don’t think it’s necessary to restructure the entire exposition.\n\n  - In 2.1, _\"UnmaskingTrees combines the gradient-boosted trees of ForestDiffusion..\"_--isn't that just XGBoost? Aside from using XGBoost, I find no similarities between the proposed autoregressive approach and the diffusion setting in Jolicoeur-Martineau et al. (2024b), so it’s unclear why the comparison is being emphasized.\n\n  - In 2.2, _\"A key problem when autoregressively generating continuous data is that a regression model will attempt to predict the mean of a conditional distribution, whereas we would like it to sample from the possibly-multimodal conditional distribution.\"_ I understand the point, but I find mentioning \"autoregressive modeling\", \"continuous responses\", and \"regression model\" could be confusing here. The core idea, I believe, is this: we want to use trees to sample from a univariate distribution. To do that, we frame the problem as a classification task over a partition of the variable’s support. The key question is then: how are these intervals defined? (Again, this problem doesn't apply only to continuous variables, but to any non-categorical data, such as count data (as shown in 3.3--4).)\n  \n  - In 2.2, worth add a sentence at the end of this section that explains that this can be naturally applied to a response with discrete outcome, as validated in the experiments in sections 3.3 and 3.4.\n\n  - In Table 1--3, why is ForestDiffusion the main benchmark for all tasks (I'm referring to the boldface blue color)?\n\nSome questions:\n\n  - Do you have a sense of why MissForest outperforms the proposed approach on some metrics, while UnmaskingTrees performs much better on others? In which scenarios would you prefer one method over the other?\n\n  - In its essence, BaltoBot estimates a conditional density with a mixture of uniforms over a partition of the variabe's support. The height $H$ of the ``meta-tree'' determines the granularity of this partition, and so seems to be a key parameter. Could you discuss how to choose $H$, and more generally, whether you have a sense of how it affects performance?", "broader_impact_concerns": "N/A.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "z52jpGoSU2", "reviewer": "Reviewer_ogxt", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n1. New Approach:\n\n- The introduction of UnmaskingTrees presents a unique method that enhances traditional tabular data handling by focusing on feature unmasking, showcasing advancements in imputation and generation tasks.\n\n2. Performance:\n\n- Empirical results demonstrate that the proposed method achieves decent performance on various benchmarks. This positions it as a reliable solution for both imputation (handling missing values) and generation (creating synthetic data).\n\n3. BaltoBot's Flexibility:\n\n- The development of BaltoBot provides a robust alternative to conventional imputation methods by avoiding parametric assumptions, allowing it to accommodate complex, multimodal distributions in tabular data.\n\n4. Practical Code Availability:\n\n- The authors have contributed to the community by making their implementation available via a GitHub repository, facilitating further research and practical applications of their methods.\n\n\n---\n**Weaknesses**\n\n1. Existing Literature Limitations:\n\n- The paper acknowledges that traditional methods, such as MissForest, have not been completely surpassed in all metrics, particularly those based on Wasserstein distances. This raises questions about whether the proposed method can universally outperform all existing approaches?\n\n\n2. Generalizability Concerns:\n\n- While the performance on benchmark datasets is promising, it would benefit the authors to include more diverse datasets to test the generalizability of their methods. Specific edge cases or datasets with unique characteristics may reveal shortcomings not currently addressed.\n\n- The paper would gain credibility by extending the benchmark to include medium- and large-sized datasets, as the current evaluation is limited to smaller ones.\n\n\n3. Hyperparameter Tuning:\n\n- The submission mentions optimized hyperparameter spaces for XGBoost but lacks a thorough exploration of how sensitive the performance is to these hyperparameters. Including a detailed analysis of hyperparameter tuning would strengthen the robustness of the methodology.\n\n- For a fair comparison, the authors should perform hyperparameter tuning for each evaluated method, rather than using default settings.\n\n\n4. Clarification on Limitations:\n\n- Further elaboration on the limitations of the UnmaskingTrees approach, particularly in terms of scalability for larger datasets, would provide a more balanced view of its applicability in real-world scenarios.\n\n\n5. Future Work Directions:\n\n- The paper would benefit from a more explicit discussion on future research paths. Specifying potential enhancements or explorations that could be made to address current limitations would provide insight into the broader impact of this work.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The submission, titled \"Unmasking Trees for Tabular Data\", presents a new method addressing challenges in tabular data imputation and generation, with these key contributions:\n\n1. Introduction of UnmaskingTrees:\n\n- A new approach leveraging gradient-boosted decision trees to incrementally \"unmask\" individual features in tabular datasets.\n\n- Achieves good performance in imputation and synthetic data generation, particularly when training data contains missing values.\n\n\n2. Development of BaltoBot:\n\n- A tabular probabilistic prediction method that fits a balanced tree of boosted tree classifiers.\n\n- Unlike traditional methods, BaltoBot does not assume a parametric form for the conditional distribution, making it suitable for handling features with multimodal distributions.\n\n- Offers fast sampling, closed-form density estimation, and a flexible approach to discrete variables, distinguishing it from contemporary diffusion methods.\n\n\n3. Meta-Algorithm Approach:\n\n- The authors consider both UnmaskingTrees and BaltoBot as meta-algorithms.\n\n- They illustrate the application of in-context learning-based generative modeling with TabPFN.\n\n\n4. Practical Implementation and Benchmarking:\n\n- The code for UnmaskingTrees is available at GitHub, facilitating broader use within the machine learning community.\n\n- Rigorous benchmarking against 27 tabular datasets, demonstrating competitive performance compared to existing imputation methods, such as MissForest and MICE-Forest, particularly on diversity and accuracy of imputations.\n\n\n5. Addressing Missingness in Tabular Data:\n\n- The method provides new insights into the handling of missing values, showcasing superior imputation capabilities over several traditional and contemporary methods.\n\n- Recognizes limitations in current literature, suggesting that traditional statistical methods like MissForest still hold considerable ground in specific scenarios.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Regarding the notation, in instances such as \"Consider a dataset with N examples,\" would it be more appropriate to utilize the term \"samples\" for consistency within the machine learning literature?\n\n- Concerning Algorithm 1, the presence of three nested iterative structures raises questions about the computational complexity when applied to large-scale datasets. Could the authors provide an analysis of the running time and discuss the feasibility of execution on such datasets?\n\n- In Algorithm 2, could the authors elaborate on the criteria and methodology employed for the selection of the hyperparameter H?\n\n- Within Section 2.3, \"Computational complexity,\" the current comparison is limited to ForestDiffusion. To provide a more comprehensive analysis, would the authors consider including a comparative evaluation against an additional two to three relevant baseline methods?\n  \n- In Section 2.4, addressing the handling of missing values (NaNs), could the authors comment on the potential efficacy and implications of replacing NaNs with a specific numerical value, such as -1?\n\n- As noted in the \"Weaknesses\" section of the review, a comprehensive hyperparameter tuning procedure for all evaluated methods is anticipated in Section 3 to ensure a fair and rigorous comparison.\n\n- In Section 3.1, the inclusion of quantitative results, potentially in tabular form, would significantly enhance the clarity and facilitate a direct comparison of the proposed method's performance.\n\n- With respect to Section 3.2, the observation that \"MissForest achieves the top rank on 4 out of 9 metrics\" suggests a performance level comparable to the proposed method. To statistically validate any perceived superiority, the application of a Nemenyi post-hoc test, as referenced in [1], might be beneficial.\n\n- Regarding Table 3, given that Tabsyn [2] currently represents the state-of-the-art, its inclusion in the comparative analysis would provide a more accurate assessment of the proposed method's position within the field.\n\n- In Table 4, the results indicate that Forest-Flow exhibits the best performance. Could the authors provide a more detailed explanation of the underlying reasons for this observation?\n\n- Within Section 3.3, concerning the \"Wave dataset\" experiments, could the authors justify the limited comparison to only Treeffuser and consider including results from other relevant methods for a more thorough evaluation?\n\n- For Figure 4, enhancing the visual comparison by including the results of a couple of additional relevant methods alongside Treeffuser and the proposed approaches would provide a more comprehensive understanding of the performance landscape.\n\n- Regarding the manuscript structure, specifically Section 5, it is suggested that the \"Related Work\" section be moved to the beginning of the paper and that the \"Discussion\" section be integrated with the \"Results\" section to improve the logical flow and clarity for the reader.\n\n- Given the results presented, where the proposed method does not consistently outperform other approaches, it is recommended that the authors moderate the claim of achieving \"state-of-the-art\" performance.\n\n---\n[1] Janez Demšar. Statistical comparisons of classifiers over multiple data sets. Journal of Machine Learning Research, 7:1–30, 2006. ISSN 15337928.\n\n[2] Zhang, H., Zhang, J., Shen, Z., Srinivasan, B., Qin, X., Faloutsos, C., Rangwala, H., & Karypis, G. (2024). Mixed-Type Tabular Data Synthesis with Score-based Diffusion in Latent Space. In The Twelfth International Conference on Learning Representations.", "broader_impact_concerns": "n.a.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "CS0eGKrnWw", "reviewer": "Reviewer_L1KV", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. Multimodal prediction of continuous values are important. The authors aims to solve an important problem\n2. The authors provide an interesting solution with the combination of KDI and Xgboost.\n\nWeaknesses:\n1. The description of the KDI part needs to add more details. Why do you choose KDI? Is it a mathematically justified approach to split the space by KDI and use Xgboost to predict the result from KDI? What's the motivation behind this?\n2. The authors menetion TabMT in the paper multiple times. It is worthwhile to compare with the TabMT method in the experiment.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose a new continuous value estimation method, BaltoBot, with application in missing value imputation. The BaltoBot is a new alternative of a regression tree, predicting the split of the node with KDI estimation and Xgboost.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Give more motivation or mathematical foundation of the BaltoBot method\n2. Compare with TabMT", "broader_impact_concerns": "No need", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)