v1 roadmap

[jeremiedb]

• Table / DataFrame
• Hybrid features support (specify numeric, bool and cat data)
• Online / streaming input
• Custom loss function (loss passed as argument, autodiff handles the rest)
• Cleanup loss name convention (mse rather than linear)
• Ranking
• GPU performance

[Moelf]

Custom loss function

how far are we from having this?

[jeremiedb]

I had only initially experimented with Zygote, but would acutally opt for an Enzyme-based approach if I were to move forward with it. I haven't pushed much on this as I haven't had much use cases where custom loss support was desired.

Would it be a useful feature to you? Is the use case for sngle parameters estimation (ex: MSE or logloss/logistic regression) or multi-parameters (ex: Gaussian Likelihood with both mu and sigma estimators)?

[Moelf]

yeah I think having something like https://dmlc.github.io/XGBoost.jl/dev/features/#Setting-a-Custom-Objective-Function is a good start.

For my application, I was initially looking at https://arxiv.org/abs/1806.00322 to try to use a different loss function but then it turns out to not be very useful under common settings:

[Moelf]

another application has come up: https://arxiv.org/pdf/2001.05310

Basically, I'd like to add a term to the classifier loss function to encourate the model output to be de-correlated with some other variables

[jeremiedb]

That's an interesting case, if I'm understanding correctly, the problematic is not just about defining a new loss function as f(pred, y) (a function of only a model prediction and a target variable), but as f(pred, y, control_vars), where control_vars represent variables for which we want to model's predictions to not be correlated.

Such formulation discussed in DisCo Fever much aligned with the problematic of model fairness; where a form is fairness is described as an absence of relationship between of model prediction and some characteristics deemed unacceptable to be discriminated against:

• FairXGBoost: https://arxiv.org/pdf/2009.01442
• FairGBM: https://openreview.net/pdf?id=x-mXzBgCX3a

The particularity I'm seeing for of such formulation is that the derivatives of the loss function cannot be computed pointwise since it's dependend on global statistics (the distance between the model's preds and the control vars).

From an implementation perspective, although one could conceive such penalty as a distinct loss; my initial thinking would be to view it as a form of regularization applicable to any Regressor loss function.

Also, the update_grads! would have to support an additional argument, the matrix of control variables.

EvoTrees.jl/src/loss.jl

Line 60 in 2435c1c

function update_grads!(∇::Matrix{T}, p::Matrix{T}, y::Vector{T}, ::Type{MSE}, params::EvoTypes) where {T}

For simplicity, considering that the control is generic correlation on a single contrl vector pseudo-julia could look like:

function loss(p, y, ctlr)
   base =  sum((p .- y) .^ 2))
   reg = cor(p, y)
   loss = base + L2 * reg
end

function update_grads!(∇::Matrix{T}, p::Matrix{T}, y::Vector{T}, ctrl::Vector{T}, ::Type{MSE}, params::EvoTypes) where {T}
   ∇[1, i] .= f'(loss(p, y, ctlr, params.L2_ctrl), p)
   ∇[2, i] .= f''(loss(p, y, ctlr, params.L2_ctrl), p)
end

Although the above loss is monolithic, I think it could be convenient to express the loss in a composable form where loss is a sum:

• loss_base: the vanilla loss defined by the Regressor loss param (MSE, logloss, poisson, ...)
• loss_reg: regularization term for (distance) correlation

Does the above seems like a sound approach? I'd see the above as a relevant development direction, as such kind of de-biasing against control features can have application in multiple fields.

[jeremiedb]

I also noted the efficiency issue raised regarding currently available Julia implementation of distance correlation:
pfarndt/EnergyStatistics.jl#2

What I understood from a quick look at dcov calculation is that it can rely on NxN distance matrix, where N==sample size.
In the context of ML application, where over 1M observation is reasonalbe to expect, this seems problematic. However, from the notation used in the Disco paper and the wiki definitions: https://en.wikipedia.org/wiki/Distance_correlation, the measure is expressed as using prime terms ||X - X'||. Not sure about it, but are these prime terme like X', Y', Y'' simply a permutations/reshuffling of the original X, Y vectors?
Assuming so, then fairly effective implementation should be possible. Otherwise, I'm afraid proper calculation would be quite inefficient, even without the allocation of the distance matrix.

[Moelf]

if I'm understanding correctly, the problematic is not just about defining a new loss function as f(pred, y) (a function of only a model prediction and a target variable), but as f(pred, y, control_vars), where control_vars represent variables for which we want to model's predictions to not be correlated.

this can always be done as a closure from user-side I think, you're right that (e.g. XGBoost.jl's) interface limited to scalar derivative function can't be used at all for this purpose.

SymbolicRegression.jl takes a different approach, their custom loss function looks like this:

 function my_loss(tree, dataset::Dataset{T,L}, options)::L where {T,L}
      prediction, flag = eval_tree_array(tree, dataset.X, options)
      if !flag
          return L(Inf)
      end
      return sum((prediction .- dataset.y) .^ 2) / dataset.n
  end

This is an API that can be used to implement distance correlation kind of loss term. dataset in EvoTrees.jl's context would be a Matrix or a DataFrame including the de-correlation target variable.

[jeremiedb]

Regarding the interface, the suppport of kwargs in fit such as feature_names, weight_name etc. could be expanded to support something like decorrelated_feature_names.

The kind of direct support I'd bee contemplating would be something similar to:

using Enzyme
using Statistics: mean, std, cor

function mse_ctrl(x, y, z)
    loss_mse = sum((x .- y) .^ 2)
    loss_corr = cor(x, z) .* length(x)
    return loss_mse + loss_corr
end

x = [0.0, 1.0, 2.0, 3.0, 4.0]
y = copy(x)
z = [0.0, 1.0, 1.0, 1.0, 2.0]

# test for 1st order grads
dx = zeros(5)
autodiff(Reverse, mse_ctrl, Active, Duplicated(x, dx), Const(y), Const(z))
dx

# test for 1st and 2nd order grads
function grad_ctrl(x, dx, y, z)
    Enzyme.autodiff(Reverse, mse_ctrl, Active, Duplicated(x, dx), Const(y), Const(z))
    return nothing
end
dx = zeros(5)
vx = ones(5)
hess = zeros(5)
Enzyme.autodiff(Enzyme.Forward, grad_ctrl,
    Enzyme.Duplicated(x, vx),
    Enzyme.Duplicated(dx, hess),
    Const(y),
    Const(z)
)
dx
vx
hess

I'm however quite not clear if loss definition such as the above would be well behaved (ie: be convex so as to be compatible with gradient boosting mechanism). A constraint for second order to be > 0 could be enforced (or simply error if any second order derivative is <= 0).

I'm unsure about the proper formulation for second order with Enzyme as I just revisited their docs. It behaves as expected with vanilla mse. But I wanted to illustrate how I'd imagine the kind of support for only specifying an arbitrary loss function and let the autodiff handles the rest.

It could also provides an elegant fundation to extend the the support for multi-targets: a custom loss take could take a matrix y rather than a simple vector, etc. Such custom gradient may not be as efficient as the dedicated multi-threaded inplace grads calculation, but could at least provide a functional basis to perform more experimental work.

Note however that this only works for the gradient boosted approach, and wouldn't generalize to other tree-split criterias.

[Moelf]

I'm unsure about the proper formulation for second order with Enzyme as I just revisited their docs. It behaves as expected with vanilla mse. But I wanted to illustrate how I'd imagine the kind of support for only specifying an arbitrary loss function and let the autodiff handles the rest.

I don't know about the details of what kind of correlation-inspired loss is well-behaved, I assumed the Disco Fever paper has math that I glanced over. But I may have some input on this auto diff API usage front, I think we should try to use DifferentiationInterface.jl as much as possible, similar to how Optimization.jl has approached this. Specifically, 2nd-order backend is quite simple [ref]:

so_backend = SecondOrder(AutoForwardDiff(), AutoZygote())
gradient(f_sparse_scalar, so_backend, x)

Enzyme.jl as the outer or inner method may or may not work, but if it works or if maybe Mooncake.jl becomes more viable, using them through DI.jl seems to give us the best flexibility.

[jeremiedb]

My understanding is that DisCo didn't have to bother about higher order derivatives as it was trained as a regular neural network, and as such only needed first-order grad along gradient descent (they used Adam solver). This is an area where classical GBT differs with their dependency on second-order grad.

That makes me think though that the loss function described in the DisCo papaer should be fairly trivial to implement in NeuroTreeModels.jl.

Thanks for the pointer on second-order in DifferentiationInterface; I agree it would be the preferable approach for integrating autodiff into EvoTrees.