Custom binary decision tree module. How to calculate gradient and propagate loss?

pumplerod: <p>I know there are a few torch modules out there which deal with tabular data and emulate a binary decision tree.  However, for one reason or another they don’t quite fit for my purpose or require a conflicting version of other modules.  And, really, I’d like to better my understanding of torch in general so I’ve been trying to create my own binary decision tree which attempts to learn the split value for the branches being applied to the feature columns.</p>
<p>I’m not even sure if what I’m trying to do is possible, but it “feels” to me like there should be a way to do this and I’ve been pounding on it so long that it’s time I ask for help.</p>
<p>There are a number of inefficiencies with the code I’m sure, however, what I’m looking to get here is just the basics of being able to adjust the <code>split_value</code> attribute of my <code>DecisionNode</code> by way of calculating a gradient and passing back through during the <code>_distribute_loss()</code> method.</p>
<p>Here is what I have…</p>
<p>for just some features and labels to be able to pass into the final module I’ve generated random values (5 different class labels from 0-1), and I’m creating 5 features for 20 different samples, all of similar values.</p>
<pre><code class="lang-auto">import torch
import numpy as np

classes = [ 0.0, 0.25, 0.5, 0.75, 1.0]
feature_data = torch.tensor( np.random.choice( classes, (20,5))).type( torch.float32)
target_data = torch.tensor( np.random.choice( classes, (20))).type( torch.float32)
</code></pre>
<p>The <code>DecisionNode</code> Class module, which takes all features as input but will split the samples into a true_branch (<code>node._tb</code>) and false_branch (<code>node._fb</code>) using the values from its assigned feature column (<code>node.col</code>) and its <code>node.split_value</code> attribute (which is what I am trying to optimize).</p>
<pre><code class="lang-auto">class DecisionNode( torch.nn.Module):
    def __init__(self, col:int, tb=None, fb=None):
        super().__init__()
        self.col = col

        #
        # Do these all need to be Prameters, do they all need to retain_grad()? I have no idea.
        self.split_value = torch.nn.Parameter( torch.rand(1, requires_grad=True).type( torch.float32))
        self.split_value.retain_grad()
        self.value = torch.nn.Parameter( torch.empty( 1))
        self.value.retain_grad()
        self.entropy = torch.nn.Parameter( torch.tensor( 0.0).type( torch.float32))
        self.entropy.retain_grad()
        self.variance = torch.nn.Parameter( torch.tensor( 0.0).type( torch.float32))
        self.variance.retain_grad()

        self._tb = tb
        self._fb = fb
    
    @property
    def device(self):
        return next(self.parameters()).device


    def forward( self, feature_value):
        if feature_value.dim() &gt; 1:
            results = torch.nn.Parameter( torch.zeros( feature_value.size(0), requires_grad=True).to( feature_value.device))
            results.retain_grad()
        else:
            results = torch.nn.Parameter( torch.zeros( 1,1, requires_grad=True).to( feature_value.device))
            results.retain_grad()

        if (self._tb is None) and (self._fb is None):
            with torch.no_grad():
                results[ :] = torch.nn.Parameter( torch.tensor( [ self.value]))
        else:
            with torch.no_grad():
                if self._tb is not None:
                    results[ feature_value[:, self.col].ge( self.split_value)] = self._tb( feature_value[ feature_value[:, self.col].ge( self.split_value)])
                if self._fb is not None:
                    results[ feature_value[:, self.col].lt( self.split_value)] = self._fb( feature_value[ feature_value[:, self.col].lt( self.split_value)])
        results.retain_grad()
        return results
</code></pre>
<p>The <code>TorchDecisionTree</code> Class module, which builds a tree of <code>DecisionNode</code> objects.  the <code>.fit()</code> method must be called before any training or prediction in order to assign the <code>node.value</code> attribute of the tree’s DecisionNodes.  It would then have to be called again after each adjustment to the <code>node.split_value</code> attribute.</p>
<p>I am calculating <code>entropy</code> and <code>variance</code> error values for the nodes, because I assume I will want to use this information to determine the gradient and apply the loss back into each node.  Here is where I am quite stuck, and my great lack of understanding for gradients and their application has me stabbing in the dark.</p>
<pre><code class="lang-auto">class TorchDecisionTree( torch.nn.Module):
    def __init__(self, max_depth:int, n_features:int, max_split_value=1.0, min_split_value=0.0):
        super().__init__()
        self.max_depth = max_depth
        self._root_node = self._build_tree( n_features, self.max_depth)
        self._max_split_value = None
        self._min_split_value = None
        
    """
        Function: device() - utility function so I can make sure the module is on the correct device
    """
    @property
    def device( self):
        return next( self.parameters()).device
    
    """
        Entropy function.
    """
    def _entropy( self, labels):
        values, counts = torch.unique( labels, return_counts=True, sorted=False)
        p = counts / labels.size( 0)
        ent = -(p * self._log2( p)).sum()
        return ent
    
    """
        Variance function.
    """
    def _variance( self, values):
        var = torch.var( values)
        if var.isnan():  # This really shouldn't happen but I'll try and worry about this later.
            var = torch.tensor( 0.0).to( self.device)
        return var
    
    """
        Function: log2()
    """
    def _log2( self, x:torch.Tensor):
        return torch.log( x) / torch.log( torch.tensor( [ 2.0]).to( x.device))
    
    """
        Private recursive function used to build the tree.
    """
    def _build_tree( self, n_features, depth):
        node = DecisionNode( col=np.random.randint( n_features))
        if depth &gt; 0:
            node._tb = self._build_tree( n_features, depth - 1)
            node._fb = self._build_tree( n_features, depth - 1)
        return node
    
    """
        Private function used to set value of leaf node.
    """
    def _set_node_value( self, node, feature_value, labels):
        if labels.size(0) == 0:
            return
        
        if ( node._tb is None) and ( node._fb is None):
            values, counts = torch.unique( labels, return_counts=True, sorted=False)
            node.value = torch.nn.Parameter( torch.tensor( [(values * counts).sum() / counts.sum()]))
            node.value.retain_grad()
            node.entropy.data = self._entropy( labels)
            node.variance.data = self._variance( labels)

        else:
            if node._tb is not None:
                self._set_node_value( node._tb, feature_value[ feature_value[:, node.col].ge( node.split_value)], labels[ feature_value[:, node.col].ge( node.split_value)])
                node.entropy.data = node.entropy + node._tb.entropy
                node.variance.data = node.variance + node._tb.variance

            if node._fb is not None:
                self._set_node_value( node._fb, feature_value[ feature_value[:, node.col].lt( node.split_value)], labels[ feature_value[:, node.col].lt( node.split_value)])
                node.entropy.data = node.entropy + node._fb.entropy
                node.variance.data = node.variance + node._fb.variance
                
        return
    
    """    PLACEHOLDER!!
        TODO:
        functions to distibute loss back and adjust node.split_value based on error
    """
    def _distribute_node_loss( self, node, mean_grad):
        node.split_value.data -= mean_grad
        node.split_value.data = torch.maximum( torch.minimum( node.split_value.data, self._max_split_value), self._min_split_value)
        if node._tb is not None:
            self._distribute_node_loss( node._tb, mean_grad)
        if node._fb is not None:
            self._distribute_node_loss( node._fb, mean_grad)
            
        return

    """    PLACEHOLDER!!
        Function: distribute_loss()
            pass the loss through all branches of the tree  
    """
    def _distribute_loss( self, mean_grad):
        self._distribute_node_loss( self._root_node, mean_grad)
            
        return
    
    """
        Function: show_splits() 
            - utility function just so I can look at the split values for the nodes.  Just for debugging
    """
    def show_splits( self, node):
        print( f"Col: {node.col}\tSplit: {node.split_value.item()}")
        if node._tb is not None:
            self.show_splits( node._tb)
        if node._fb is not None:
            self.show_splits( node._fb)
        return
    
    """
        Function: show_values() 
            - utility function just so I can look at the entropy variance values for the nodes.  Just for debugging
    """
    def show_values( self, node, level=0, branch='root'):
        print( f"Level: {level}\tBranch: {branch}\tEntropy: {node.entropy.item():0.5f}\tVariance: {node.variance.item():0.5f}")
        if node._tb is not None:
            self.show_values( node._tb, level=level+1, branch='_tb')
        if node._fb is not None:
            self.show_values( node._fb, level=level+1, branch='_fb')
        return
    
    """
        Function: fit()
            - set node.value based on training features and labels
    """
    def fit( self, features, labels):
        if self._max_split_value is None:
            self._max_split_value = features.max()
        if self._min_split_value is None:
            self._min_split_value = features.min()

        self._set_node_value( self._root_node, features, labels)
        return
    
    """
        Function: predict()
            - only for compatibility
    """
    def predict( self, features):
        return self.forward( features)
    
    """
        Function: forward()
            - return values based on feature inputs
    """
    def forward( self, features):
        X = self._root_node( features)
        return X

</code></pre>
<p>then to test the code…</p>
<pre><code class="lang-auto">tree = TorchDecisionTree( n_features=feature_data.size(1), max_depth=2)
</code></pre>
<p>which creates a tree and random <code>split_value</code> for each node, which can be seen via…</p>
<pre><code class="lang-auto">tree.show_splits( tree._root_node)
</code></pre>
<p>You can see that the same column is often selected for multiple nodes, which may not be the best approach, but it should still allow the model to run and train.</p>
<pre><code class="lang-auto">Col: 1	Split: 0.8780428171157837
Col: 4	Split: 0.5944340229034424
Col: 1	Split: 0.8516061305999756
Col: 3	Split: 0.8720638751983643
Col: 3	Split: 0.7830216884613037
Col: 3	Split: 0.570353627204895
Col: 1	Split: 0.11497271060943604
</code></pre>
<p>To <code>fit</code> the model to the training data…</p>
<pre><code class="lang-auto">tree.fit( feature_data, target_data)
</code></pre>
<p>which will now generate predictions on features passed in…</p>
<pre><code class="lang-auto">preds = tree( feature_data)
print( preds)
</code></pre>
<p>showing…</p>
<pre><code class="lang-auto">Parameter containing:
tensor([0.5833, 0.6667, 0.3750, 0.3750, 0.5833, 0.5833, 0.2500, 0.5833, 0.5833,
        0.2500, 0.5833, 0.6667, 0.5833, 0.5833, 0.5833, 0.6667, 0.2500, 0.5833,
        0.5833, 0.5833], requires_grad=True)
</code></pre>
<p>and just to check that my <code>entropy</code> and <code>variance</code> attributes for the nodes are being set (not that I’m doing the right thing with regard to summation of branch nodes, I’m not yet sure what the correct approach is here but, again, the values are affected and I should be able to use them to adjust the <code>split_value</code>).  Ideally, for the leaf nodes, I’m looking for low entropy and low variance (if I understand binary trees correctly, which I may not).</p>
<pre><code class="lang-auto">tree.show_values( tree._root_node)
</code></pre>
<p>which shows…</p>
<pre><code class="lang-auto">Level: 0	Branch: root	Entropy: 5.28742	Variance: 0.82860
Level: 1	Branch: _tb	    Entropy: 1.91830	Variance: 0.61458
Level: 2	Branch: _tb	    Entropy: 1.00000	Variance: 0.28125
Level: 2	Branch: _fb	    Entropy: 0.91830	Variance: 0.33333
Level: 1	Branch: _fb	    Entropy: 3.36912	Variance: 0.21402
Level: 2	Branch: _tb	    Entropy: 1.58496	Variance: 0.06250
Level: 2	Branch: _fb	    Entropy: 1.78416	Variance: 0.15152
</code></pre>
<p>The main questions…</p>
<ol>
<li>How do I calculate/accumulate the gradient for <code>node.split_value</code> w.r.t entropy or variance loss.</li>
<li>How do I then apply the gradient in order to adjust the <code>node.split_value</code>?</li>
<li>Can this be done with the help of <code>torch.autograd</code>, or am I completely on my own when needing to distribute the loss?</li>
</ol>
<p>If you’ve read and followed this far then I want to thank you very much.  Any help is greatly appreciated.</p>