title: Optimizing conditional multiclass softmax objective function in XGBoost
body: <p>I have successfully implemented a custom multiclass softmax function function in XGBoost based on <a href="https://github.com/dmlc/xgboost/blob/c024c42dce0a2c4d01881751f0e3dc7c7691ebf8/demo/guide-python/custom_softmax.py" rel="nofollow noreferrer">this</a> tutorial. The reason for customization is that the classes I want to predict are conditional on some data inputs - i.e. of the 24 possible classes being predicted, only a certain subset are valid. <code>valid_transitions</code> are lists of indices corresponding to classes we want to make predictions on and <code>invalid_transitions</code> are the inverse set of indices.</p>
  <p>I have implemented <code>.fit()</code> and <code>.predict_proba()</code> such that they take <code>valid_transitions</code> and <code>invalid_transitions</code> as arguments which tells <code>softprob_obj()</code> and <code>softmax()</code>which classes to null out during training and prediction.</p>
  <pre><code>def softmax(x, valid_transitions, invalid_transitions):
      for i in range(len(x)):
          e = np.exp(x[i,valid_transitions[i]])
          x[i, valid_transitions[i]] = e/np.sum(e)
          x[i, invalid_transitions[i]] = 0
      return x
  
  def softprob_obj(labels, predt, data, valid_transitions, invalid_transitions):
      '''Loss function.  Computing the gradient and approximated hessian (diagonal).
      Reimplements the `multi:softprob` inside XGBoost.
      '''
      kRows = len(data)
      kClasses = len(np.unique(labels))
      # The prediction is of shape (rows, classes), each element in a row
      # represents a raw prediction (leaf weight, hasn't gone through softmax
      # yet).  In XGBoost 1.0.0, the prediction is transformed by a softmax
      # function, fixed in later versions.
      assert predt.shape == (kRows, kClasses)
  
      eps = 1e-6
      # compute the gradient and hessian, slow iterations in Python, only
      # suitable for demo.  Also the one in native XGBoost core is more robust to
      # numeric overflow as we don't do anything to mitigate the `exp` in
      # `softmax` here.
      probs = softmax(predt, valid_transitions, invalid_transitions)
      labels = labels.astype(int)
      hess = np.maximum((2.0 * probs * (1.0 - probs)), eps)
      probs[np.arange(len(probs)),labels] -= 1
      # Right now (XGBoost 1.0.0), reshaping is necessary
      grad = probs.reshape((kRows * kClasses, 1))
      hess = hess.reshape((kRows * kClasses, 1))
      return grad, hess
  </code></pre>
  <p>This works, but is pretty slow in training, presumably because the core xgboost functions I'm replacing are not written in python. I made some attempts to try to vectorize the calculation in numpy to avoid the for loop in <code>softmax()</code>, but ran into some issues with the ragged arrays that <code>valid_transition</code> and <code>invalid_transition</code> create. Was wondering if anyone had any ideas on how to optimize this within python. Thanks!</p>