Learning Binary Trees via Sparse Relaxation

1 INTRODUCTION . Learning discrete structures from unstructured data is extremely useful for a wide variety of real-world problems ( Gilmer et al. , 2017 ; Kool et al. , 2018 ; Yang et al. , 2018 ) . One of the most computationallyefficient , easily-visualizable discrete structures that are able to represent complex functions are binary trees . For this reason , there has been a massive research effort on how to learn such binary trees since the early days of machine learning ( Payne & Meisel , 1977 ; Breiman et al. , 1984 ; Bennett , 1992 ; Bennett & Blue , 1996 ) . Learning binary trees has historically been done in one of three ways . The first is via greedy optimization , which includes popular decision-tree methods such as classification and regression trees ( CART ) ( Breiman et al. , 1984 ) and ID3 trees ( Quinlan , 1986 ) , among many others . These methods optimize a splitting criterion for each tree node , based on the data routed to it . The second set of approaches are based on probabilistic relaxations ( İrsoy et al. , 2012 ; Yang et al. , 2018 ) . The idea is to optimize all splitting parameters at once via gradient-based methods , by relaxing hard branching decisions into branching probabilities . The third approach optimizes trees using mathematical programming ( MIP ) ( Bennett , 1992 ; Bennett & Blue , 1996 ) . This jointly optimizes all continuous and discrete parameters to find globally-optimal trees.1 Each of these approaches have clear shortcomings . First , greedy optimization is clearly suboptimal : tree splitting criteria are even intentionally crafted to be different than the global tree loss , as the global loss does not encourage tree growth ( Breiman et al. , 1984 ) . Second , probabilistic relaxations : ( a ) are rarely sparse , so inputs probabilistically contribute to branches they would never visit if splits are mapped to hard decisions ; ( b ) they do not have principled ways to prune trees , as the distribution over pruned trees is often intractable . Third , MIP approaches , while optimal , are only computationally tractable on training datasets with thousands of inputs ( Bertsimas & Dunn , 2017 ) , and do not have well-understood out-of-sample generalization guarantees . In this paper we present a new approach to binary tree learning based on sparse relaxation and argmin differentiation . Our main insight is that by quadratically relaxing an MIP that learns the discrete parameters of the tree ( input traversal and node pruning ) , we can differentiate through it to simultaneously learn the continuous parameters of splitting decisions . This allows us to leverage the superior generalization capabilities of stochastic gradient optimization to learn splits , and gives 1Here we focus on learning single trees instead of tree ensembles ; our work easily extends to ensembles . a principled approach to learning tree pruning . Further , we can derive customized algorithms to compute the forward and backward passes through this program that are much more efficient than generic approaches ( Amos & Kolter , 2017 ) . We demonstrate that ( a ) in classification/regression our method , which learns a single tree and a classifier on top of it , is competitive with greedy , probabilistic , MIP-based tree methods , and even powerful ensemble methods ; ( b ) in hierarchical clustering we match or improve upon the state-of-the-art . 2 RELATED WORK . The paradigm of binary tree learning has the goal of finding a tree that iteratively splits data into meaningful , informative subgroups , guided by some criterion . Binary tree learning appears in a wide variety of problem settings across machine learning . We briefly review work in two learning settings where latent tree learning plays a key role : 1 . Classification/Regression ; and 2 . Hierarchical clustering . Due to the generality of our setup ( tree learning with arbitrary split functions , pruning , and downstream objective ) , our approach can be used to learn trees in any of these settings . Finally , we detail how parts of our algorithm are inspired by recent work in isotonic regression . Classification/Regression . Decision trees for classification and regression have a storied history , with early popular methods that include classification and regression trees ( CART ; Breiman et al. , 1984 ) , ID3 ( Quinlan , 1986 ) , and C4.5 ( Quinlan , 1993 ) . While powerful , these methods are greedy : they sequentially identify ‘ best ’ splits as those which optimize a split-specific score ( often different from the global objective ) . As such , learned trees are likely sub-optimal for the classification/regression task at hand . To address this , Carreira-Perpinán & Tavallali ( 2018 ) proposes an alternating algorithm for refining the structure and decisions of a tree so that it is smaller and with reduced error , however still sub-optimal . Another approach is to probabilistically relax the discrete splitting decisions of the tree ( İrsoy et al. , 2012 ; Yang et al. , 2018 ; Tanno et al. , 2019 ) . This allows the ( relaxed ) tree to be optimized w.r.t . the overall objective using gradient-based techniques , with known generalization benefits ( Hardt et al. , 2016 ; Hoffer et al. , 2017 ) . Variations on this approach aim at learning tree ensembles termed ‘ decision forests ’ ( Kontschieder et al. , 2015 ; Lay et al. , 2018 ; Popov et al. , 2019 ) . The downside of the probabilistic relaxation approach is that there is no principled way to prune these trees as inputs pass through all nodes of the tree with some probability . A recent line of work has explored mixed-integer program ( MIP ) formulations for learning decision trees . Motivated by the billion factor speed-up in MIP in the last 25 years , Rudin & Ertekin ( 2018 ) proposed a mathematical programming approach for learning provably optimal decision lists ( one-sided decision trees ; Letham et al. , 2015 ) . This resulted in a line of recent follow-up works extending the problem to binary decision trees ( Hu et al. , 2019 ; Lin et al. , 2020 ) by adapting the efficient discrete optimization algorithm ( CORELS ; Angelino et al. , 2017 ) . Related to this line of research , Bertsimas & Dunn ( 2017 ) and its follow-up works ( Günlük et al. , 2018 ; Aghaei et al. , 2019 ; Verwer & Zhang , 2019 ; Aghaei et al. , 2020 ) phrased the objective of CART as an MIP that could be solved exactly . Even given this consistent speed-up all these methods are only practical on datasets with at most thousands of inputs ( Bertsimas & Dunn , 2017 ) and with non-continuous features . Further , the out-of-sample generalizability of these approaches is not well-studied , unlike stochastic gradient descent learning . Hierarchical clustering . Compared to standard flat clustering , hierarchical clustering provides a structured organization of unlabeled data in the form of a tree . To learn such a clustering the vast majority of methods are greedy and work in one of two ways : 1 . Agglomerative : a ‘ bottom-up ’ approach that starts each input in its own cluster and iteratively merges clusters ; and 2 . Divisive : a ‘ top-down ’ approach that starts with one cluster and recusively splits clusters ( Zhang et al. , 1997 ; Widyantoro et al. , 2002 ; Krishnamurthy et al. , 2012 ; Dasgupta , 2016 ; Kobren et al. , 2017 ; Moseley & Wang , 2017 ) . These methods suffer from similar issues as do greedy approaches to tree learning for classification/regression : they may be sub-optimal for optimizing the overall tree . Further they are often computationally-expensive due to their sequential nature . Inspired by approaches for classification/regression , recent work has designed probabilistic relaxations for learning hierarchical clusterings via gradient-based methods ( Monath et al. , 2019 ) . Our work takes inspiration from both the MIP-based and gradient-based approaches . Specifically , we frame learning the discrete tree parameters as an MIP , which we sparsely relax to allow continuous parameters to be optimized by argmin differentiation methods . Argmin differentiation . Solving an optimization problem as a differentiable module within a parent problem tackled with gradient-based optimization methods is known as argmin differentiation , a particular instance of bi-level optimization ( Gould et al. , 2016 ) . This situation arises in as diverse scenarios as hyperparameter optimization ( Pedregosa , 2016 ) , meta-learning ( Rajeswaran et al. , 2019 ) , or structured prediction ( Stoyanov et al. , 2011 ; Domke , 2013 ; Niculae et al. , 2018 ) . General algorithms for quadratic Amos & Kolter ( 2017 ) and disciplined convex programming ( Section 7 , Amos , 2019 ; Agrawal et al. , 2019a ; b ) have been given , as well as expressions for more specific cases like isotonic regression ( Djolonga & Krause , 2017 ) . By taking advantage of the very specific structure of the decision tree induction problem , we obtain a direct , efficient algorithm . Isotonic regression . Also called monotonic regression , isotonic regression ( Barlow et al. , 1972 ) constrains the regression function to be non-decreasing/non-increasing . This is useful if one has prior knowledge of such monotonicity ( e.g. , the mean temperature of the climate is non-decreasing ) . A classic algorithm is pooling-adjacent-violators ( PAV ) , which optimizes the pooling of adjacent points that violate the monotonicity constraint ( Barlow et al. , 1972 ) . This initial algorithm has been generalized and incorporated into convex programming frameworks ( see Mair et al . ( 2009 ) for an excellent summary of the history of isotonic regression and its extensions ) . Our work builds off of the generalized PAV ( GPAV ) algorithm of Yu & Xing ( 2016 ) . 3 METHOD . Given inputs { xi } ni=1 , our goal is to learn a latent binary decision tree T with maximum depth D. This tree sends each input x through branching nodes to a specific leaf node in the tree . Specifically , all branching nodes TB ⊂ T split an input x by forcing it to go to its left child if sθ ( x ) < 0 , and right otherwise . There are three key parts of the tree that need to be identified : 1 . The continuous parameters θt ∈ Rd that describe how sθt splits inputs at every node t ; 2 . The discrete paths z made by each input x through the tree ; 3 . The discrete choice at of whether a node t should be active or pruned . We describe how to learn each of these below . 3.1 TREE-TRAVERSAL & PRUNING PROGRAMS . Imagine the splitting functions of the tree sθt are fixed . Given this , the following integer linear program ( ILP ) describes how inputs x traverse the tree . The solution zit ∈ { 0 , 1 } indicates if xi reaches node t ∈ T ( for notational simplicity let zi ∈ { 0 , 1 } |T | be the vectorized indicator for xi ) , max z1 , ... , zn n∑ i=1 z > i qi ( 1 ) s.t . ∀i ∈ [ n ] , qit = min { rit , lit } rit = min { sθt′ ( xi ) , ∀t ′ ∈ AR ( t ) , } lit = min { −sθt′ ( xi ) , ∀t ′ ∈ AL ( t ) } zit ∈ { 0 , 1 } . HereAL ( t ) is the set of ancestors of node t whose left child must be followed to get to t , and similarly for AR ( t ) . The quantities qit ( where qi ∈ R|T | is the tree-vectorized version of qit ) describe the ‘ reward ’ of sending xi to node t. This is based on how well the splitting functions leading to t are satisfied ( qit is positive if all splitting functions are satisfied and negative otherwise ) . Notice that the solution is unique so long as sθt ( xi ) 6=0 for all t ∈ T , i ∈ { 1 , . . . , n } ( i.e. , sθt ( xi ) =0 means there is no preference to split xi left or right ) . Further note that the integer constraint on zit can be relaxed to an interval constraint zit ∈ [ 0 , 1 ] without loss of generality . This is because if sθt ( xi ) 6= 0 then zt = 0 if and only if qt < 0 and zt = 1 when qt > 0 ( and qt 6= 0 ) . While the above program works for any fixed tree , we would like to be able to also learn the structure of the tree itself . Let ηt ∈ R be our preference for pruning/keeping node t ( larger ηt indicates node t should be kept ) . To encourage this pruning preference while ensuring connectivity , we introduce an additional optimization variable at ∈ { 0 , 1 } , indicating if node t ∈ T is active ( if 1 ) or pruned ( if 0 ) . We may now adapt eq . 1 into the following pruning-aware mixed integer program ( MIP ) : max z1 , ... , zn , a n∑ i=1 z > i qi + η > a ( 2 ) s.t . ∀i ∈ [ n ] , at ≤ ap ( t ) , ∀t ∈ T \ { root } zit ≤ at zit ∈ [ 0 , 1 ] , at ∈ { 0 , 1 } . Here we have removed the first three constraints in eq . 1 as they are a deterministic computation independent of z1 , . . . , zn , a . Further p ( t ) indicates the parent of node t. The added constraint at ≤ ap ( t ) ensures that child nodes t are pruned if parent nodes p ( t ) are pruned . While the new constraint zit ≤ at ensures that no point xi can reach node t if node t is pruned .

Summary: The given paper describes a novel approach of learning binary trees using a differentiable relaxation of MIP. For that matter, authors describe their approach of tree traversal for the given input x and a complete binary tree of depth D. This problem is formulated as MIP which is relaxed to make it differentiable. This MIP is then reformulated to learn a tree structure (i.e. tree pruning). Finally, a decision node parameters are optimized using backpropagation. The algorithm alternates between tree structure learning using relaxed MIP and tree parameters learning using SGD.