# Learning Decision Trees as Amortized Structure Inference

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 46 of its 83 references that are in the map.
- First posted on arXiv 10 Mar 2025.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Revisiting Deep Learning Models for Tabular Data](https://mlsanity.com/paper/2106.11959/) (2021)
- [TabNet: Attentive Interpretable Tabular Learning](https://mlsanity.com/paper/1908.07442/) (2019)
- [Self-Normalizing Neural Networks](https://mlsanity.com/paper/1706.02515/) (2017)
- [Tabular Data: Deep Learning is Not All You Need](https://mlsanity.com/paper/2106.03253/) (2021)
- [TabTransformer: Tabular Data Modeling Using Contextual Embeddings](https://mlsanity.com/paper/2012.06678/) (2020)
- [Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation](https://mlsanity.com/paper/2106.04399/) (2021)
- [GFlowNet Foundations](https://mlsanity.com/paper/2111.09266/) (2021)
- [Trajectory balance: Improved credit assignment in GFlowNets](https://mlsanity.com/paper/2201.13259/) (2022)
- [When Do Neural Nets Outperform Boosted Trees on Tabular Data?](https://mlsanity.com/paper/2305.02997/) (2023)
- [Classification-Based Anomaly Detection for General Data](https://mlsanity.com/paper/2005.02359/) (2020)
- [CatBoost: gradient boosting with categorical features support](https://mlsanity.com/paper/1810.11363/) (2018)
- [Neural Oblivious Decision Ensembles for Deep Learning on Tabular Data](https://mlsanity.com/paper/1909.06312/) (2019)
- [Why do tree-based models still outperform deep learning on tabular data?](https://mlsanity.com/paper/2207.08815/) (2022)
- [Bayesian Structure Learning with Generative Flow Networks](https://mlsanity.com/paper/2202.13903/) (2022)
- [COPOD: Copula-Based Outlier Detection](https://mlsanity.com/paper/2009.09463/) (2020)
- [DROCC: Deep Robust One-Class Classification](https://mlsanity.com/paper/2002.12718/) (2020)
- [GFlowNets and variational inference](https://mlsanity.com/paper/2210.00580/) (2022)
- [Optimal Sparse Decision Trees](https://mlsanity.com/paper/1904.12847/) (2019)
- [Generalized and Scalable Optimal Sparse Decision Trees](https://mlsanity.com/paper/2006.08690/) (2020)
- [Better by Default: Strong Pre-Tuned MLPs and Boosted Trees on Tabular Data](https://mlsanity.com/paper/2407.04491/) (2024)
- [Benchmarking Distribution Shift in Tabular Data with TableShift](https://mlsanity.com/paper/2312.07577/) (2023)
- [Generative Flow Networks as Entropy-Regularized RL](https://mlsanity.com/paper/2310.12934/) (2023)
- [MurTree: Optimal Classification Trees via Dynamic Programming and Search](https://mlsanity.com/paper/2007.12652/) (2020)
- [Bayesian Model Selection, the Marginal Likelihood, and Generalization](https://mlsanity.com/paper/2202.11678/) (2022)
- [Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network](https://mlsanity.com/paper/2305.19366/) (2023)
- [Discrete Probabilistic Inference as Control in Multi-path Environments](https://mlsanity.com/paper/2402.10309/) (2024)

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

- [Bayesian Symbolic Regression with Entropic Reinforcement Learning](https://mlsanity.com/paper/2608.09617/) (2026)
- [MotherTree: Meta-learning on synthetic data improves decision tree training](https://mlsanity.com/paper/2610.10832/) (2026)

## Abstract (arXiv)

Building predictive models for tabular data presents fundamental challenges, notably in scaling consistently, i.e., more resources translating to better performance, and generalizing systematically beyond the training data distribution. Designing decision tree models remains especially challenging given the intractably large search space, and most existing methods rely on greedy heuristics, while deep learning inductive biases expect a temporal or spatial structure not naturally present in tabular data. We propose a hybrid amortized structure inference approach to learn predictive decision tree ensembles given data, formulating decision tree construction as a sequential planning problem. We train a deep reinforcement learning (GFlowNet) policy to solve this problem, yielding a generative model that samples decision trees from the Bayesian posterior. We show that our approach, DT-GFN, outperforms state-of-the-art decision tree and deep learning methods on standard classification benchmarks derived from real-world data, robustness to distribution shifts, and anomaly detection, all while yielding interpretable models with shorter description lengths. Samples from the trained DT-GFN model can be ensembled to construct a random forest, and we further show that the performance of scales consistently in ensemble size, yielding ensembles of predictors that continue to generalize systematically.