Entity Embdedding

Most MLJFlux models support categorical features by learning enitity embeddings. By wrapping such an MLJFlux model using EntityEmbedder, the learned embeddings can be used in MLJ pipelines as transformer elements. In particular, these embeddings can be used for supervised models that are not neural networks and require features to be Continuous. See the example below.

MLJFlux.EntityEmbedder — Type

EntityEmbedder(; model=supervised_mljflux_model)

Wrapper for a MLJFlux supervised model, to convert it to a transformer. Such transformers are still presented a target variable in training, but they behave as transformers in MLJ pipelines. They are entity embedding transformers, in the sense of the article, "Entity Embeddings of Categorical Variables" by Cheng Guo, Felix Berkhahn.

Training data

In MLJ (or MLJBase) bind an instance unsupervised model to data with

mach = machine(embed_model, X, y)

Here:

embed_model is an instance of EntityEmbedder, which wraps a supervised MLJFlux model, model, which must be an instance of one of these: MLJFlux.NeuralNetworkClassifier, NeuralNetworkBinaryClassifier, MLJFlux.NeuralNetworkRegressor,MLJFlux.MultitargetNeuralNetworkRegressor.
X is any table of input features supported by the model being wrapped. Features to be transformed must have element scitype Multiclass or OrderedFactor. Use schema(X) to check scitypes.
y is the target, which can be any AbstractVector supported by the model being wrapped.

Train the machine using fit!(mach).

Examples

In the following example we wrap a NeuralNetworkClassifier as an EntityEmbedder, so that it can be used to supply continuously encoded features to a nearest neighbor model, which does not support categorical features.

using MLJ

# Setup some data
N = 400
X = (
  a = rand(Float32, N),
  b = categorical(rand("abcde", N)),
  c = categorical(rand("ABCDEFGHIJ", N), ordered = true),
)

y = categorical(rand("YN", N));

# Initiate model
EntityEmbedder = @load EntityEmbedder pkg=MLJFlux

# Flux model to do learn the entity embeddings:
NeuralNetworkClassifier = @load NeuralNetworkClassifier pkg=MLJFlux

# Other supervised model type, requiring `Continuous` features:
KNNClassifier = @load KNNClassifier pkg=NearestNeighborModels

# Instantiate the models:
clf = NeuralNetworkClassifier(embedding_dims=Dict(:b => 2, :c => 3))
emb = EntityEmbedder(clf)

# For illustrative purposes, train the embedder on its own:
mach = machine(emb, X, y)
fit!(mach)
Xnew = transform(mach, X)

# And compare feature scitypes:
schema(X)
schema(Xnew)

# Now construct the pipeline:
pipe = emb |> KNNClassifier()

# And train it to make predictions:
mach = machine(pipe, X, y)
fit!(mach)
predict(mach, X)[1:3]

It is to be emphasized that the NeuralNertworkClassifier is only being used to learn entity embeddings, not to make predictions, which here are made by KNNClassifier().

source