# Multi-Armed Bandits-Based Optimization of Decision Trees

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 6 of its 28 references that are in the map.
- First posted on arXiv 8 Aug 2025.

## Builds on (most cited first, 6 of 6)

- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [Feature Engineering for Predictive Modeling using Reinforcement Learning](https://mlsanity.com/paper/1709.07150/) (2017)
- [Hyp-RL : Hyperparameter Optimization by Reinforcement Learning](https://mlsanity.com/paper/1906.11527/) (2019)
- [AutoRAG-HP: Automatic Online Hyper-Parameter Tuning for Retrieval-Augmented Generation](https://mlsanity.com/paper/2406.19251/) (2024)
- [MABSplit: Faster Forest Training Using Multi-Armed Bandits](https://mlsanity.com/paper/2212.07473/) (2022)
- [RL-LLM-DT: An Automatic Decision Tree Generation Method Based on RL Evaluation and LLM Enhancement](https://mlsanity.com/paper/2412.11417/) (2024)

## Abstract (arXiv)

Decision trees, without appropriate constraints, can easily become overly complex and prone to overfit, capturing noise rather than generalizable patterns. To resolve this problem,pruning operation is a crucial part in optimizing decision trees, as it not only reduces the complexity of trees but also decreases the probability of generating overfit models. The conventional pruning techniques like Cost-Complexity Pruning (CCP) and Reduced Error Pruning (REP) are mostly based on greedy approaches that focus on immediate gains in performance while pruning nodes of the decision tree. However, this might result in a lower generalization in the long run, compromising the robust ability of the tree model when introduced to unseen data samples, particularly when trained with small and complex datasets. To address this challenge, we are proposing a Multi-Armed Bandits (MAB)-based pruning approach, a reinforcement learning (RL)-based technique, that will dynamically prune the tree to generate an optimal decision tree with better generalization. Our proposed approach assumes the pruning process as an exploration-exploitation problem, where we are utilizing the MAB algorithms to find optimal branch nodes to prune based on feedback from each pruning actions. Experimental evaluation on several benchmark datasets, demonstrated that our proposed approach results in better predictive performance compared to the traditional ones. This suggests the potential of utilizing MAB for a dynamic and probabilistic way of decision tree pruning, in turn optimizing the decision tree-based model.