# Unmasking Trees for Tabular Data

arXiv: https://arxiv.org/abs/2407.05593 · HTML: https://mlsanity.com/paper/2407.05593/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 48 of its 73 references that are in the map.
- First posted on arXiv 8 Jul 2024.

## Builds on (most cited first, 30 of 48)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [RePaint: Inpainting using Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2201.09865/) (2022)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Modeling Tabular data using Conditional GAN](https://mlsanity.com/paper/1907.00503/) (2019)
- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [Tabular Data: Deep Learning is Not All You Need](https://mlsanity.com/paper/2106.03253/) (2021)
- [TabDDPM: Modelling Tabular Data with Diffusion Models](https://mlsanity.com/paper/2209.15421/) (2022)
- [GAIN: Missing Data Imputation using Generative Adversarial Nets](https://mlsanity.com/paper/1806.02920/) (2018)
- [When Do Neural Nets Outperform Boosted Trees on Tabular Data?](https://mlsanity.com/paper/2305.02997/) (2023)
- [Autoregressive Diffusion Models](https://mlsanity.com/paper/2110.02037/) (2021)
- [Mixed-Type Tabular Data Synthesis with Score-based Diffusion in Latent Space](https://mlsanity.com/paper/2310.09656/) (2023)
- [Generative Modelling With Inverse Heat Dissipation](https://mlsanity.com/paper/2206.13397/) (2022)
- [CTAB-GAN: Effective Table Data Synthesizing](https://mlsanity.com/paper/2102.08369/) (2021)
- [DECAF: Generating Fair Synthetic Data Using Causally-Aware Generative Networks](https://mlsanity.com/paper/2110.12884/) (2021)
- [Matrix Completion and Low-Rank SVD via Fast Alternating Least Squares](https://mlsanity.com/paper/1410.2596/) (2014)
- [Missing Data Imputation using Optimal Transport](https://mlsanity.com/paper/2002.03860/) (2020)
- [NGBoost: Natural Gradient Boosting for Probabilistic Prediction](https://mlsanity.com/paper/1910.03225/) (2019)
- [Hierarchical Scene Coordinate Classification and Regression for Visual Localization](https://mlsanity.com/paper/1909.06216/) (2019)
- [STaSy: Score-based Tabular data Synthesis](https://mlsanity.com/paper/2210.04018/) (2022)
- [Diffusion models for missing value imputation in tabular data](https://mlsanity.com/paper/2210.17128/) (2022)

## Built on by (most cited first, 4 of 4)

- [TabTreeFormer: Tabular Data Generation Using Hybrid Tree-Transformer](https://mlsanity.com/paper/2501.01216/) (2025)
- [NRGBoost: Energy-Based Generative Boosted Trees](https://mlsanity.com/paper/2410.03535/) (2024)
- [Democratizing Tabular Data Access with an Openunicodex2013Source Syntheticunicodex2013Data SDK](https://mlsanity.com/paper/2508.00718/) (2025)
- [Marformer: A Transformer for Predicting Missing Data Distributions](https://mlsanity.com/paper/2610.12379/) (2026)

## Abstract (arXiv)

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.