Mlj.jl: computing `UnivariateFinite` matrix seems to be substantially slow

[aviatesk]

To Reproduce

using MLJ, LightGBM
using BenchmarkTools

n = 100000
X = DataFrame(A = rand(n), B = categorical(rand(1:10, n)))
y = categorical(rand(0:2, n))

@load LGBMClassifier
@pipeline LGBM(clf = LGBMClassifier()) prediction_type = :probabilistic
lgbm_machine = machine(LGBM(), X, y)
MLJ.fit!(lgbm_machine)

lgbm_model = LGBMClassification(num_class = length(y.pool))
LightGBM.fit!(lgbm_model, Array(X), collect(y))

@btime MLJ.predict(lgbm_machine, X) # 1.391 sec
@btime LightGBM.predict(lgbm_model, Array(X)) # 49.767 ms

Expected behavior: never so slow, ideally

Versions

Pkg.status()

data) pkg> st
  [336ed68f] CSV v0.6.1
  [a93c6f00] DataFrames v0.20.2
  [1313f7d8] DataFramesMeta v0.5.0
  [682c06a0] JSON v0.21.0
  [7acf609c] LightGBM v0.2.1
  [30fc2ffe] LossFunctions v0.6.1
  [add582a8] MLJ v0.11.1
  [a7f614a8] MLJBase v0.13.2
  [6ee0df7b] MLJLinearModels v0.3.2
  [e80e1ace] MLJModelInterface v0.2.2
  [d491faf4] MLJModels v0.9.9
  [1914dd2f] MacroTools v0.5.5
  [91a5bcdd] Plots v0.29.9
  [2913bbd2] StatsBase v0.33.0
  [f3b207a7] StatsPlots v0.14.5

Thanks !

[tlienart]

Posting what I posted on Slack here as well:

@ablaom my hunch is that it’s the array comprehension [UnivariateFinite(classes, probs) ...] that’s retarded and slow, we should have a single UnivariateFinite(classes, all_probs) where all_probs is a matrix of size n x c that does this in an optimized way, specifically ther's stuff with conversion to LittleDict etc which is quite inefficient to do for every thing I would think...

Btw it’s used everywhere so if we do find this to be the case we’ll have to fix it in a few places… (edited)

[ablaom]

@aviatesk Thanks for flagging this.

@tlienart Yes, your hunch is probably right, and your suggestion to improve the situation sounds reasonable.

[ablaom]

@aviatesk Side remark: You shouldn't need to define num_classes in the MLJ wrap of LightGBM, as my understanding is that this is automatically determined (and I'm surprised the field is still exposed):

lgbm_model = LGBMClassification(num_class = length(y.pool))

cc: @yalwan-iqvia

[aviatesk]

hm, I just tried again but doesn't seem to work:

n = 100000
X = DataFrame(A = rand(n), B = categorical(rand(1:10, n)))
y = categorical(rand(0:2, n))
lgbm_model = LGBMClassification()
LightGBM.fit!(), Array(X), collect(y))

this code ends up with

LightGBM.fit!(lgbm_model, Array(X), collect(y))
[LightGBM] [Fatal] Label must be in [0, 2), but found 2 in label
ERROR: call to LightGBM's LGBM_BoosterCreate failed: Label must be in [0, 2), but found 2 in label
Stacktrace:
 [1] error(::String, ::String, ::String, ::String) at ./error.jl:42

May I open an issue about the above error into your repository separate from this issue, since it's not related to MLJ, @yalwan-iqvia ?

[yalwan-iqvia]

There seems to have been a bit of confusion here.

@ablaom when you requested it wasn't required for the MLJ interface I did so, but the regular interface still leaves that control to the user

LGBMClassification is the regular interface model.
LGBMClassifier is the MLJ interface model.

I don't consider this to be a bug, because the LightGBM.fit! method does not infer information about the number of classes. However, @aviatesk if you feel strongly about this, please do raise an issue into our repository for us to discuss further 😄

[aviatesk]

sounds fair. I just found XGBoost.jl's xgboost function (which is equivalent to LightGBM's LGBMClassification initialization and then fit! step) also only does the "unwrapped" (or in other word, "basic") task:

xgboost(
    Array(X), 100; label = collect(y), 
    objective = "multi:softprob", # <- we need to manually specify `multi:hoge` objective for multi classification task
    num_class = length(levels(y)) # <- XGBoost library also requires an user to specify this manually
)

imho, MLJ as an higher level API for various models is better to handle those data-conversions/task-specifications, while each model package may be better to stay to provide "bare" functionality by default.

[ablaom]

@aviatesk @yalwan-iqvia Sorry for the confusion.

Regarding the issue as initially raised:

There is nothing in the MLJ API requiring a probabilistic model to return a vector of distributions. It need only be an abstract vector, in the present case any instance of AbstractVector{<:UnivariateFinite}. In the long term, I more-or-less expected we would need to replace vectors with more efficient wraps, as we do already for categorical data. So, even if it is not absolutely necessary for speed improvement, this is the path I recommend to address this issue.

Starting with the case at hand, we define a new type UnivariateFiniteArray, whose "elements" (the distributions returned by indexing):

• share a common categorical pool
• have support (classes with non-zero probability) collectively contained in some subset we expect is generally "small". This is relevant, because the size of this upper bound on the support is one of the dimensions of our probability array.

So something like:

# L - categorical element type
# U - categorical reference type (<:Unsigned)
# T - type of raw probablities 
# N - array dimension
struct UnivariateFiniteArray{C,U,T<:ReaL,N} <: AbstractArray{UnivariateFinite{L,U,T}, N}
    probs::Array{T,N}
    support::CategoricalVector{C,U, other stuff to make concrete??}
    p::Array{T,1 + N}}     # first dimension is for classes in support
end

This is, naturally, a bit of work, because we should implement the full AbstractArray interface (which probably why it has not been done until now 😄 )

[tlienart]

So I tried to go some of the way towards this; it's probably not worth a PR but I'll leave it here because some stuff might be of interest for a cleaner implementation. I went the route of only defining a UnivariateFiniteVector because I don't really see the use of an Array (multi-target classification? seems weird).

struct UnivariateFiniteVector{P<:Arr{<:Real},S<:Vec{<:CategoricalElement}} <: Vec{UnivariateFinite}
    scores::P
    support::S
    function UnivariateFiniteVector(sc::Vec{<:Real}, su::Vec{<:CategoricalElement})
        @assert length(su) == 2 "Number of classes don't match."
        @assert all(0 .<= sc .<= 1) "Scores must be between [0,1]."
        scores = hcat(sc, 1 .- sc)
        new{typeof(scores),typeof(su)}(scores, su)
    end
    function UnivariateFiniteVector(sc::Arr{<:Real}, su::Vec{<:CategoricalElement})
        @assert size(sc, 2) == length(su) "Number of classes don't match."
        @assert all(sum(sc, dims=2) .≈ 1.) "Scores must sum to one."
        new{typeof(sc),typeof(su)}(sc, su)
    end
end

Note that I didn't go the route of subtyping everything, I think that's unnecessary. Concrete types are good if you're going to use them in for loops etc, but here it's just a glorified container.

Automatic classes if the user only provides scores:

# Used only here
const UFV = UnivariateFiniteVector

# Auto classes
function UFV(sc::Vec{<:Real})
    su = classes(categorical([:class_0, :class_1])[1])
    return UFV(sc, su)
end
function UFV(sc::Arr{<:Real})
    su = classes(categorical([Symbol("class_$i") for i in 1:size(sc, 2)])[1])
    return UFV(sc, su)
end

Some of the methods for Arrays with fallbacks:

# Note that a user could change the scores to invalid scores...
for method in (:size, :setindex!)
    ex = quote
        Base.$method(u::UFV, a...; kw...) = $method(u.scores, a...; kw...)
    end
    eval(ex)
end
Base.length(u::UFV) = size(u.scores, 1)

function Base.getindex(u::UFV, i::Int) # cast back to UnivariateFinite
    C = eltype(u.support)
    P = eltype(u.scores)
    prob_given_class = LittleDict{C,P}(u.support[i] => u.scores[i,j] for j in eachindex(u.support))
    MMI.UnivariateFinite(prob_given_class)
end
# for slices etc.
Base.getindex(u::UFV, I) = getindex(u.scores, I)

function Base.show(io::IO, ::MIME"text/plain", u::UFV)
    write(io, "UnivariateFiniteVector: \n")
    write(io, "  Support: "); Base.show(io, u.support);
    write(io, "\n  Length:  $(length(u))");
end

Some thoughts

• would need to hack around the broadcasting so that we can do pdf.(u::UFV, c). The default broadcasting cannot be used because it would create all these individual UnivariateFinite objects which would be slow. Basically you want to be able to do something like:

function Base.broadcast(pdf, u::UFV, c)
  i = findfirst(c, u.support)
  @assert !isnothing(i) "Unrecognised class element 'c'"
  return  u.scores[:, i]
end

I don't know whether you can write specific broadcasting mechanisms on a per function basis, I had a look but it seemed a bit complicated so there's probably a better way to do this.

Conclusion

I think the stuff you might want to keep from this is the fact that the vector is an implicit container, and that it only really creates the UnivariateFinite object when you select individual elements so that typeof(ufv[i]) == UnivariateFinite.

[OkonSamuel]

@tlienart. This is great. Just have one question must we do pdf.(u::UFV, c) can we just overload pdf as pdf(u::UFV, c).
PS:
I think this resemble Distributions.jl DiscreteMultivariateDistribution ??

[tlienart]

I'll take a look thanks for the pointer! I'm reluctant for the pdf(u::UFV, ...) because UFV is not a distribution.

I've been thinking about this issue a bit more today and I think another approach to the whole business of storing scores that effectively backs the UnivariateFinite vector might be the way to go. I'll try to come up with a full solution, it's a bit tricky but it seems much needed.

[OkonSamuel]

I just though it might be easier to treat as a distibution (with each column representing the scores for a given class). Then implement the Distributions interface ommiting some methods??
Sorry if am disturbing. Lol

[ablaom]

For clarification, UnivariateFinite is indeed a subtype of Distributions.Distribution. I think @tlienart knows this but has just forgotten 😄 :

https://github.com/alan-turing-institute/MLJBase.jl/blob/5aec436a948509639dd370ebf10194e424445e48/src/interface/univariate_finite.jl#L4

[tlienart]

sorry what I meant in my previous comment is that the whole object UnivariateFiniteVector is not a distribution, I'm aware that UnivariateFinite is a distribution :). Anyway the comment is outdated now with the current PR; pdf(u::UnivariateFiniteVector, ...) with or without broadcasting works.

[yalwan-iqvia]

I'll take a look thanks for the pointer! I'm reluctant for the pdf(u::UFV, ...) because UFV is not a distribution.

Its a vector of distributions, right? Anything against creating a pdfs function, which has described signature? It would be semantically correct and serve the same purpose (apart from any deliberate name aliasing I am unaware of)

[tlienart]

@yalwan-iqvia take a look at https://github.com/alan-turing-institute/MLJBase.jl/pull/317 and the gist in it. It would be trivial to have pdfs but I'm not convinced it's necessary, people would probably do pdf(u,...) or (if they care) pdf.(u, ...) in both case it falls back on an efficient computation of the individual pdfs.

[yalwan-iqvia]

I didn't realise PR was already in progress -- only reason why I suggested it was so that we can avoid hijacking a bunch of base functions just for the sake of getting "efficient broadcasting" when a semantically correct alternative name could do the same job without the overloading

[ablaom]

Update.

This issue will be resolved by these steps:

• [x] Merge https://github.com/alan-turing-institute/MLJBase.jl/pull/327 which introduces a performant subtypes of AbstractArray{<:UnivariateFinite} with performant implementations of broadcast for pdf and mode, and implements these for the probabilistic classifier metrics, like auc.

• [x] Get all models to predict the abstract version of UnivariateFinite vectors, instead of the vanilla vectors

[tlienart]

huge effort by Anthony! thanks a lot :) I can do some of the second part :)

[yalwan-iqvia]

Is there a concise guidance on what precisely implementers need to do? I'll deal with it for LightGBM.jl

[tlienart]

@yalwan-iqvia it's dead simple, basically you'll need to change the comprehension in predict [UnivariateFinite...] for a direct call to UnivariateFinite passing the scores:

So this line

https://github.com/IQVIA-ML/LightGBM.jl/blob/87240e65215d7c6e0c649fd30bdd2ed1c47eef28/src/MLJInterface.jl#L251

would, in your case, become exactly

return MLJModelInterface.UnivariateFinite(classes, predicted)

as predicted is always explicit (same number of columns as classes).

[yalwan-iqvia]

Ok, thanks, and what version of the package do I need to upgrade to?

[tlienart]

For now don't do anything last I heard from Anthony there were issues with Registrator so let's first make sure everything is properly in.

@ablaom I realise that you've added this as patch release in MMI and MLJBase, I think it should be a minor release since it's a change in the API requiring a change in model interfaces => (0.3 / 0.14).

[ablaom]

@ablaom I realise that you've added this as patch release in MMI and MLJBase, I think it should be a minor release since it's a change in the API requiring a change in model interfaces => (0.3 / 0.14).

Yeah, good catch. I hadn't thought this through properly. There is a potential for issues. Many thanks!

[ablaom]

Okay, the only models not buying into the performance boost are third party packages providing the MLJ model interface natively. I have raised issues there.

So closing.

[ablaom]

@aviatesk Be great if you get a chance to confirm resolution from your end. You will need to wait for https://github.com/JuliaRegistries/General/pull/16126#issuecomment-641669872 to merge for the XGBoost.jl test and for https://github.com/IQVIA-ML/LightGBM.jl/issues/56 to be resolved to test LightGBM.jl

[yalwan-iqvia]

I've tested locally using small variation of code supplied by @aviatesk above -- the situation has improved. I'll merge the PR soon and release new version with this and other features.

(amusingly, @btime tells me LightGBM.predict is slower than MLJ.predict, reliably and consistently, which is of course amusing because MLJ.predict calls LightGBM.predict, nevertheless it isn't something I am massively concerned about)

@ablaom thanks for the hard work.

[aviatesk]

cool, thanks for all involved with this.
I just confirmed the original code I posted now yields 29.398 ms which is >50x times faster than before.
I do appreciate your great works.

fwiw, this is my current project status:

  [336ed68f] CSV v0.6.2
  [a93c6f00] DataFrames v0.21.2
  [1313f7d8] DataFramesMeta v0.5.1
  [682c06a0] JSON v0.21.0
  [7acf609c] LightGBM v0.3.0
  [add582a8] MLJ v0.11.4
  [6ee0df7b] MLJLinearModels v0.5.0
  [1914dd2f] MacroTools v0.5.5
  [91a5bcdd] Plots v1.4.0
  [2913bbd2] StatsBase v0.33.0
  [f3b207a7] StatsPlots v0.14.6
  [009559a3] XGBoost v0.4.3