Optimization in Machine Learning: Decision Trees, Sparse Linear Models and LLM Pruning
Name
meng-mengx-phd-orc-2026-thesis.pdf
Size
2.51 MB
Format
Adobe PDF
Checksum (MD5)
1b379ee0885905e6016e3885d76b2f1b
Author(s)
Meng, Xiang
Advisor(s)
Mazumder, Rahul
Date Issued
May 2026
Publisher
Massachusetts Institute of Technology
Abstract
This thesis develops optimization algorithms for three families of discrete optimization problems arising in machine learning and statistics: building decision trees, large language model pruning, and sparse regression. Depending on the problem structure and scale, we exploit problem structure to design effective algorithms including both exact methods that certify global optimality and scalable heuristics that prioritize computational tractability.
In the first part, we study optimal decision trees, a predictive tool with a flowchart-like structure that allows users to derive interpretable decisions. Learning an optimal decision tree is NP-hard, and greedy heuristics such as CART may lead to suboptimal predictive performance. Recent advances in mixed-integer programming (MIP) and dynamic programming have enabled computation of optimal trees, but MIP formulations have limited scalability and most algorithms are restricted to binary features. We develop a novel branch-and-bound algorithm for optimal trees on datasets with continuous features. The algorithm uses quantiles of the feature distribution to decompose the search space and computes upper and lower bounds to eliminate sub-optimal regions. Our approach is the first standalone method for optimal classification and regression trees that directly handles continuous features, providing orders of magnitude speedup over existing methods on datasets with tens of thousands of instances.
In the second part, we present four optimization-based frameworks for neural network pruning, targeting one-shot pruning of large language models, one-shot structured pruning, FLOP-aware pruning, and transposable N:M sparsity. ALPS formulates one-shot pruning of large language models (LLMs) as an ell_0-constrained optimization problem solved via operator splitting (ADMM), followed by preconditioned conjugate gradient to compute optimal weights on the support. OSSCAR proposes a combinatorial optimization framework for one-shot structured pruning using a layer-wise reconstruction objective and structure-aware reformulation, scaling to models with up to 30 billion parameters on a single GPU. FALCON introduces FLOP-aware pruning by jointly handling sparsity and FLOP constraints through an integer linear program, achieving significantly better accuracy-FLOPs-sparsity tradeoffs. TSENOR formulates transposable N:M sparse mask generation as optimal transport problems with capacity constraints, solved via entropy regularization and Dykstra's algorithm, enabling hardware acceleration of both forward and backward passes during training.
In the third part, we develop a GPU-accelerated branch-and-bound (BnB) framework for sparse linear regression with ell_0 and ell_2 regularization. Solving this problem to global optimality requires traversing a massive BnB search tree, and existing algorithms are designed for CPU architectures. We develop a specialized first-order method whose updates are parallelizable across decision variables and across BnB nodes, enabling large speedups on GPUs. Our approach keeps lightweight BnB logic on the CPU while offloading computationally expensive bound computations to the GPU, significantly outperforming both commercial solvers and prior custom methods.
MIT Department
Massachusetts Institute of Technology. Operations Research Center
Sloan School of Management
Terms of Use
In Copyright - Educational Use Permitted
Copyright retained by author(s)
Persistent DSpace Link
