PaperlayerView exact revision on arXiv

Source-linked AI summary

Entity Embeddings of Categorical Variables

Cheng Guo, Felix Berkhahn

arXiv:1604.06737v1cs.LG

TL;DR

Categorical variables are difficult for neural networks because one-hot encoding is costly for high-cardinality features and ignores relationships among categories. The paper learns entity embeddings during supervised neural-network training, placing similarly behaving categories near one another. These embeddings improve neural-network and other machine-learning performance, support visualization and clustering, and warrant testing on more datasets.

  • Problem

    Categorical variables may lack continuity, while one-hot encoding is resource-intensive for high-cardinality features and treats category values independently.

  • Method

    A neural network learns multidimensional entity embeddings for categorical variables during supervised training.

  • Results

    The learned embeddings place categories with similar effects near one another, improve tested machine-learning methods, and resemble geographic structure in German-state visualizations.

  • Takeaways & Limitations

    Entity embeddings provide a distance measure for categorical variables that can be used to understand, visualize, and cluster categorical data.

  • Takeaways & Limitations

    The method should be tested on more datasets, especially those with many high-cardinality features and sparse data.

Abstract

from arXiv · show

We map categorical variables in a function approximation problem into Euclidean spaces, which are the entity embeddings of the categorical variables. The mapping is learned by a neural network during the standard supervised training process. Entity embedding not only reduces memory usage and speeds up neural networks compared with one-hot encoding, but more importantly by mapping similar values close to each other in the embedding space it reveals the intrinsic properties of the categorical variables. We applied it successfully in a recent Kaggle competition and were able to reach the third position with relative simple features. We further demonstrate in this paper that entity embedding helps the neural network to generalize better when the data is sparse and statistics is unknown. Thus it is especially useful for datasets with lots of high cardinality features, where other methods tend to overfit. We also demonstrate that the embeddings obtained from the trained neural network boost the performance of all tested machine learning methods considerably when used as the input features instead. As entity embedding defines a distance measure for categorical variables it can be used for visualizing categorical data and for data clustering.

I. INTRODUCTION

Neural networks are powerful for unstructured data but less suitable for structured data with categorical variables because categorical values lack an obvious continuity. The paper presents entity embeddings as a learned representation that captures relationships among categories and supports neural-network and other machine-learning methods.

  • Neural networks have driven major advances in computer vision, speech recognition, and natural language processing.
  • Structured-data problems commonly favor tree-based methods, unlike the unstructured domains where neural networks are prominent.
  • Categorical variables challenge neural networks because their values may lack continuity, making naive integer representations ineffective.
  • One-hot encoding is costly for many high-cardinality features and ignores informative relationships by treating category values independently.
  • Entity embeddings automatically map categorical values into multidimensional spaces where values with similar effects are close together.
  • The paper extends earlier distributed representations from relational data and word embeddings to general function approximation.

III. TREE BASED METHODS

Decision trees model structured-data functions by partitioning the feature space into regions and assigning region-specific predictions. Because exhaustive partition search is infeasible, trees use recursive greedy splitting and complexity control.

  • Decision trees partition the feature space into M subspaces and model the function using region-specific values.
  • The tree formulation uses an indicator function to assign each input to its corresponding region.
  • For each region, the optimal prediction parameter is the average of the target values assigned to that region.
  • Finding the globally optimal partition is computationally infeasible because the number of possible partitions grows exponentially with feature-space size.
  • A greedy algorithm recursively selects local splits and applies the same procedure to the resulting half-planes.
  • Tree size controls model complexity: large trees may overfit, whereas small trees may miss important dependencies.

B. Random forests

Random forests reduce the variance of decision trees by averaging an ensemble of trees built from bootstrap samples and randomized feature selection at each split.

  • Random forests address the high variance of individual decision trees by averaging a committee of identically distributed trees.
  • The ensemble outputs all N trained trees, and regression predictions are formed from their linear average.
  • Each forest tree is trained on a bootstrap sample drawn from the training data.
  • At each split node, only a random subset of features is considered for the best split.

C. Gradient boosted trees

Gradient tree boosting approximates a function by adding trees sequentially, with each new tree trained using information from its predecessors. Unlike random forests, it also targets reduction of ensemble bias.

  • Gradient tree boosting approximates the target function as a sum of trees.
  • The nth tree is grown on a quantity computed from its n − 1 predecessor trees for a generic loss function.
  • For squared-error loss, the next tree fits residuals yi − fn−1(xi) left by the preceding trees.
  • Each successive classifier attempts to correct its predecessors, so boosting reduces bias as well as variance.

IV. STRUCTURED DATA

Structured data consists of tabular samples with continuous or discrete variables, but discrete labels may lack meaningful numerical relationships. Entity embeddings map discrete values into a multidimensional space where values with similar function outputs are close.

  • Structured data is organized in tables, with columns representing features or targets and rows representing samples.
  • Discrete variables may be nominal labels with no meaningful numerical order, or ordinal values whose ordering may still be unhelpful for the task.
  • Figure 1 illustrates entity embedding layers as extra layers placed above each one-hot encoded input.
  • Entity embedding maps discrete values to a multidimensional space where values with similar function outputs are close.

V. ENTITY EMBEDDING

Entity embeddings learn vector representations for categorical states through neural-network training, replacing one-hot inputs with learned features that are concatenated before deeper processing. Their dimensions are predefined hyperparameters, and the paper notes that stronger theoretical selection rules remain desirable.

  • Each state of a discrete variable is mapped to a vector, making the representation equivalent to a linear layer on top of one-hot encoding.
  • The embedding vectors are the weights connecting one-hot states to embedding units and are learned with standard neural-network parameters.
  • Embedding outputs and continuous inputs are concatenated, then passed to ordinary neural-network layers trained by standard backpropagation.
  • Embedding dimensions Di are predefined hyperparameters bounded between 1 and mi −1, where mi is the number of categorical values.
  • The paper chooses dimensions empirically, increasing them for more complex entities or starting with mi −1 when complexity is unclear.
  • The authors identify the lack of theoretical guidelines for choosing embedding dimensions as an open issue related to finite metric-space embedding.

A. Relation with embedding of finite metric space

The paper relates entity embeddings to finite metric spaces by defining distances between categorical values according to their effects on the target function. Learned store embeddings are not isometric to this metric, but their distances show a bounded linear relationship.

  • Entity embedding seeks to place categorical values with similar effects closer together, linking the method to finite metric-space embedding.
  • For each categorical variable, the paper defines a finite metric space whose elements are the variable’s possible values and whose metric represents their similarity.
  • The metric is based on averaging over the other function parameters when comparing two values of the categorical variable.
  • Different values that always produce the same output may violate the metric’s identity condition, so they can be merged as a redundant category.
  • The store metric space cannot be isometrically embedded in Euclidean space because its associated matrix was not positive definite.
  • For 10000 random store pairs, learned-embedding distance and metric distance show a linear relation with well-defined upper and lower boundaries, rather than an isometric mapping.

VI. EXPERIMENTS

The experiments use a simplified subset of the Rossmann sales-prediction data to compare machine-learning methods under both shuffled and temporally ordered train-test splits. The design preserves a future-time test setting for assessing generalization.

  • The experiment uses the Kaggle Rossmann competition dataset, whose goal is accurate daily sales prediction for each store.
  • The data contain about 2.5 years of daily sales for 1115 stores, totaling 1017210 records, plus additional store details.
  • The paper uses only a small feature subset and no feature engineering because it compares machine-learning methods rather than optimizing the competition score.
  • Training uses 90% of the data and testing 10%, with both temporally ordered and randomly shuffled splits.
  • In the unshuffled split, future-time test data assess generalization from the training data because their temporal distribution is not exactly sampled during training.

A. Neural networks

The paper compares neural networks and tree-based methods on shuffled and temporally ordered Rossmann data, using entity embeddings alongside one-hot representations. Entity embeddings perform better than one-hot encoding on non-shuffled data and also improve other tested machine-learning methods.

  • Model comparison: Neural networks use either one-hot encoding or learned entity embeddings, with two fully connected layers and a sigmoid output.The embedding-based and one-hot models share the same main network structure for comparison.
  • Model comparison: The experiments compare KNN, random forests, gradient boosted trees, and neural networks using model-specific input representations.KNN receives one-hot-encoded features, while the tree-based methods and neural networks use their specified feature inputs.
  • Evaluation: The benchmark uses 200,000 randomly sampled training examples and evaluates MAPE instead of competition RM-SPE.MAPE was selected because it was considered more stable with outliers.
  • Results: For shuffled data, gradient boosted trees with entity embeddings and neural networks achieve comparable results, while one-hot neural networks perform slightly better than entity embeddings.The shuffled split uses a test set sharing the training distribution, whereas the non-shuffled split tests later data.
  • Results: For non-shuffled data, neural networks perform best, and entity embedding is clearly better than one-hot encoding.The non-shuffled evaluation uses the latest 10% of data as test data to assess generalization from earlier observations.
  • Results: Feeding learned entity embeddings into other machine-learning methods significantly improves all methods tested.The embedded features are used as inputs to methods beyond the neural network that learned them.

C. Distribution in the embedding space

The paper examines how learned categorical embeddings organize geographically and functionally related stores and states. Visualizations show meaningful sales structure along principal components, approximately Gaussian marginal densities, and weak correlation among feature-specific subspaces.

  • German state embeddings: A t-SNE projection of learned German state embeddings places states in relative positions resembling the real German map.Eastern states Sachsen, Thueringen, and Sachsen Anhalt form a cluster, while several western states occupy another cluster.
  • German state embeddings: The embedding places states with similar feature distributions close together, linking geographic proximity to shared economic and cultural environments.The geographic pattern is presented as an outcome of similarity in the learned feature distributions.
  • Store sales distribution: Across all 1115 stores, sales vary continuously along the first principal component of the store embedding, unlike along two random directions.The first two principal components and two random directions are compared in the visualization.
  • Store sales distribution: Other embedding directions show no direct sales correlation but may encode additional store properties that matter when combined with other features.Their possible predictive role appears in deeper network layers rather than through direct univariate sales correlation.
  • Embedding density: The first four principal-component densities are approximately Gaussian individually, but their joint distribution is not multivariate Gaussian.The D’Agostino’s K2 tests are statistically significant, with p-values below 0.05, and the Mardia test rejects multivariate Gaussianity.
  • Concatenated space: The concatenated embedding space has no strong correlation between individual categorical subspaces, supporting separate analysis of those subspaces.The neural network receives the direct product of the entity-embedding subspaces.

VII. FUTURE WORK

The paper identifies several directions for future exploration, including broader dataset testing, deeper theoretical analysis, extensions beyond categorical variables, and activation-function comparisons.

  • Entity embeddings should be tested on more datasets, especially sparse datasets with many high-cardinality features.The authors expect this setting to reveal the method’s full strength compared with other methods.
  • Future work could investigate whether directions in some learned embeddings have interpretable meanings.The paper specifically mentions directions illustrated by Eq. (2) and Eq. (3).
  • A deeper study of entity embeddings’ relation to finite metric spaces could clarify optimal embedding dimensions and neural-network behavior.
  • Similar methods may improve approximation of continuous non-monotone functions by discretizing continuous variables into categorical variables.
  • The authors also propose systematically comparing different activation functions in the entity embedding layer.
Loading 1604.06737v1…