Entity Embeddings for High-Cardinality Features

Welcome back. In our last lesson, we explored CatBoost's ingenious method of "Ordered Target Statistics," a technique that uses the target variable to create powerful categorical features while cleverly avoiding data leakage. This approach is highly effective, but it belongs to a family of target-based encodings.

Today, we shift our perspective entirely. What if we could create rich, numerical representations of categories without looking at the target variable at all? This lesson addresses exactly that, focusing on how to use entity embeddings from neural networks to create dense, low-dimensional representations of high-cardinality features. We will draw a compelling parallel to word embeddings, a concept you're familiar with from your NLP work, and see how a neural network can learn the intrinsic "meaning" of a category in the context of a predictive task. This is a fundamental technique in modern machine learning, especially for datasets riddled with high-cardinality features like user IDs, product codes, or geographical locations.

The Challenge of High Cardinality

As you know, traditional methods like one-hot encoding become untenable when a categorical feature has thousands or millions of unique values. The resulting feature space becomes incredibly sparse and high-dimensional, a classic example of the "curse of dimensionality." This makes it difficult for many models to learn effectively.

Target-based methods, which we've discussed, are one solution. But they have their own complexities, such as managing potential leakage. Entity embeddings offer a completely different path, one that leverages the representation-learning power of deep learning.

The Core Idea: Learning Meaning from Data

The central concept behind entity embeddings is inspired by word embeddings from NLP. Just as Word2Vec or GloVe learn that the vector relationship king - man + woman is close to queen, an entity embedding model can learn that zip_code_A is "closer" to a neighboring zip_code_B than a distant zip_code_C if doing so helps predict the target (e.g., house prices).

An embedding is simply a mapping of a discrete, categorical variable to a dense vector of continuous numbers. The key is that this mapping isn't handcrafted; it is learned by a neural network as part of a supervised training process. The network learns to place similar categories—similar in the context of the prediction task—close to each other in the embedding space.

The Mechanism: An Embedding Layer

So how does a neural network learn these vector representations? The magic happens in a special layer called an Embedding Layer.

Imagine you have a categorical feature, say "Country," with 200 unique values.

1. Integer Encoding: First, you map each country to a unique integer from 0 to 199.
2. Embedding Layer: You create an embedding layer in your neural network. This layer is essentially a lookup table—a weight matrix of shape (number_of_categories, embedding_dimension). For our example, if we choose an embedding dimension of 10, the matrix will be of size (200, 10).
3. Lookup: When you pass the integer for a specific country (e.g., 'Canada' is mapped to 35), the embedding layer simply returns the 35th row of its weight matrix. This row is the 10-dimensional embedding vector for 'Canada'.
4. Training: This vector is then fed into the rest of the neural network. During training, as the network learns to minimize its prediction error, gradients are backpropagated all the way to the embedding layer, updating the values in the weight matrix. The vectors are adjusted to become more useful for the final prediction.

The following article walks through this process beautifully, first from first principles using NumPy and then with Keras. Given your interest in the fundamentals, the NumPy section will be particularly illuminating.

Building a Complete Architecture

In a typical real-world dataset, you'll have multiple categorical features and some numerical features. The standard architecture for handling this with embeddings is straightforward and quite powerful.

1. One Embedding Layer per Categorical Feature: Each categorical column gets its own dedicated embedding layer. The embedding dimension for each can be chosen independently.
2. Concatenation: The output vectors from all embedding layers are concatenated into a single, flat vector.
3. Add Numerical Features: The original numerical features are then concatenated with this vector of embeddings.
4. Dense Layers: This final, wide vector is fed into one or more dense (fully-connected) layers, which learn to find complex interactions between all the features.
5. Output Layer: A final output layer produces the prediction.

The diagram below illustrates this common and effective pattern.

Abhishek Thakur's video provides a great walkthrough of this entire process, from data preparation to building and training the model in Keras.

Making Sense of the Embeddings

One of the most powerful aspects of this technique is that the learned embeddings are not just a black box. They are rich features that can be analyzed to provide deep insights into your data.

Once the model is trained, you can extract the weight matrix from an embedding layer. Each row is a vector representing a category. By using dimensionality reduction techniques like PCA or t-SNE, you can visualize these vectors in 2D or 3D. Often, you will discover meaningful clusters that your model has learned automatically. For example, in a retail dataset, you might find that different types of 'winter jackets' cluster together, separate from 'summer dresses'.

Interview & Business Context Alert

Question: "You have a dataset with a user_id column that has millions of unique values. One-hot encoding is not feasible. How would you represent this feature for a gradient boosting model versus a neural network?"

Answer Framework:

1. For Gradient Boosting Models: "For a tree-based model like XGBoost or LightGBM, I would use a target-based encoding scheme. Methods like target encoding with cross-validation or CatBoost's ordered target statistics are effective. The key consideration is preventing target leakage, so a robust validation strategy is critical."
2. For Neural Networks: "For a neural network, entity embeddings are the ideal approach. I would map each user_id to an integer and use an embedding layer to learn a dense vector representation for each user. The network learns these vectors as part of the end-to-end training process, optimizing them to be predictive of the target."
3. The Synergy (Advanced Answer): "A powerful hybrid approach is to first train a neural network with an embedding layer for the user_id. Then, I can extract the learned embedding vectors and use them as rich, pre-computed features in a separate gradient boosting model. This combines the representation learning power of neural networks with the robust performance of gradient boosting on tabular data."

A Glimpse into the Future: Deep Learning on Tabular Data

The concept of using neural networks to automatically learn features from raw data is a cornerstone of deep learning. While entity embeddings are a well-established method, research in this area is rapidly evolving. Your interest in "out of the box" methods like using CNNs for tabular data is very current.

Architectures like TabNet use sequential attention to perform instance-wise feature selection, mimicking the decision-making process of tree-based models but in a differentiable, deep learning framework.

Other approaches even transform tabular data into images to be processed by CNNs or into graphs to be processed by GNNs. We will explore these advanced architectures in detail in Module 7, but it's valuable to know that entity embeddings are a gateway to this exciting field of deep learning on tabular data.

Conclusion

In this lesson, we've unlocked a powerful, modern technique for handling high-cardinality categorical features. By leveraging the representation-learning capabilities of neural networks, entity embeddings allow us to create rich, dense, and meaningful features automatically.

Key Takeaways:

• Entity Embeddings: A method to map high-cardinality categorical variables into a low-dimensional continuous vector space.
• Learned Representations: Unlike handcrafted features, embeddings are learned by a neural network during training, optimizing them for the specific predictive task.
• Mechanism: An embedding layer acts as a lookup table (a trainable weight matrix) that is updated via backpropagation.
• Architecture: The standard approach involves creating separate embeddings for each categorical feature, concatenating them with numerical features, and feeding them into dense layers.
• Interpretability: Learned embeddings are not a complete black box; they can be visualized and analyzed to reveal hidden relationships between categories.

What's Next?

Entity embeddings provide rich, learned representations but come with the computational cost of training a neural network. What if our primary constraints were memory and speed, and we needed a simpler, albeit less "intelligent," way to wrangle high-cardinality features?

In our next lesson, we will investigate feature hashing, also known as the "hashing trick." You will learn how this technique uses a hash function to map a vast number of categories into a fixed-size feature vector, providing a highly memory-efficient alternative to one-hot encoding, and we will analyze its trade-offs regarding information loss through hash collisions.