8.2 Embeddings for Categorical Features

When One Hot Encoding Stops Scaling

Structured or tabular datasets often contain categorical columns such as country, product id, user id, or weekday. A common baseline is one hot encoding, where each category becomes a separate 0 or 1 feature. This works for small cardinalities, but it becomes inefficient when a column has thousands or millions of unique values, because the input vector becomes huge and sparse.

Embeddings replace a large sparse one hot vector with a small dense vector that is learned during training. Conceptually, you can think of an embedding as a table that maps each category index to a learned feature vector. If a category is represented by an integer $i$, and the embedding dimension is $d$, then the embedding layer returns a vector in $\mathbb{R}^d$ for that category.

The Embedding Lookup in PyTorch

In PyTorch, categorical embeddings are usually implemented with torch.nn.Embedding(num_embeddings, embedding_dim). If num_embeddings = K, the layer contains a learnable weight matrix $E \in \mathbb{R}^{K \times d}$. Given an index $i$, the output is the row $E[i]$.

For a batch of indices shaped (batch_size,), the output has shape (batch_size, d). For multiple categorical features, you typically keep one embedding layer per column. Each column has its own vocabulary size and may use a different embedding dimension.

A minimal example for a single categorical column looks like this:

import torch
import torch.nn as nn
K = 1000          # number of categories
d = 16            # embedding size
emb = nn.Embedding(K, d)
x_cat = torch.tensor([3, 10, 999], dtype=torch.long)   # (batch,)
z = emb(x_cat)                                         # (batch, d)
print(z.shape)  # torch.Size([3, 16])

Preparing Categorical Columns: Vocabulary, Unknowns, and Missing Values

Real datasets contain categories that were not seen during training, plus missing values. A practical approach is to build a vocabulary mapping from category value to an integer id and reserve special ids.

A common convention is to reserve index 0 for unknown or missing, then map known categories to 1 through $K-1$. In that case, num_embeddings is the size of the final index space, including the reserved token.

If you are using pandas, you can convert a column to categorical codes, but you still need a stable mapping saved from training time. The key idea is that the model must see the same integer encoding at train and inference.

Multiple Categorical Features in a Feedforward Model

In tabular models, you usually concatenate all embeddings and numeric features, then feed them into an MLP. Suppose you have three categorical columns with embedding dimensions $d_1, d_2, d_3$ and $m$ numeric features. After embedding lookups, you concatenate into a single vector of size $d_1 + d_2 + d_3 + m$.

In PyTorch, one simple pattern is to keep embeddings in an nn.ModuleList and concatenate along the last dimension:

class TabularMLP(nn.Module):
    def __init__(self, cat_sizes, emb_dims, num_numeric, hidden=128):
        super().__init__()
        self.embs = nn.ModuleList([
            nn.Embedding(K, d) for K, d in zip(cat_sizes, emb_dims)
        ])
        in_dim = sum(emb_dims) + num_numeric
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, 1)
        )
    def forward(self, x_cats, x_num):
        # x_cats: (batch, num_cat_cols) long
        # x_num:  (batch, num_numeric) float
        embs = []
        for j, emb in enumerate(self.embs):
            embs.append(emb(x_cats[:, j]))
        x = torch.cat(embs + [x_num], dim=1)
        return self.mlp(x)

Choosing Embedding Dimensions

Embedding size is a hyperparameter. If it is too small, the model may not capture useful differences between categories. If it is too large, you can overfit and waste memory. There is no universally correct rule, but practical choices are usually small, especially for moderate cardinalities.

A common heuristic is to choose $d$ that grows slowly with the number of categories $K$, such as $d \approx \min(50, \lceil \sqrt{K} \rceil)$ or similar. The main point is that embeddings are meant to be compact, you rarely need hundreds of dimensions for ordinary tabular categories.

Regularizing Embeddings

Embeddings can memorize rare categories, especially when a category appears only a few times. Two simple regularization ideas are common in practice.

One approach is embedding dropout, which randomly drops embedding vectors during training, forcing the network to rely less on any single categorical id. Another approach is weight decay on the embedding parameters, which can help prevent large embedding weights.

You can apply weight decay through the optimizer. If you want different weight decay for embeddings and the rest of the network, you can use parameter groups, but the general optimizer setup is handled in the optimizer chapter.

Common Shape and Dtype Pitfalls

The most frequent embedding errors in beginner tabular models are dtype mismatches and unexpected shapes. The categorical tensor should be two dimensional as (batch_size, num_cat_cols) and torch.long. Numeric tensors should typically be floating point and already scaled as discussed in the feature scaling chapter.

Another common issue appears when batching data. If you store categories as strings in the dataset and encode them inside __getitem__, you must ensure the DataLoader collates them into a tensor of integers, not a list of Python objects.

Handling High Cardinality and Rare Categories

When cardinality is extremely high, such as user ids, embeddings are still useful, but you should expect many categories to be rare. In such cases, it helps to cap the vocabulary by keeping only the most frequent categories and mapping the rest to unknown. This reduces memory and encourages generalization.

A simple frequency cutoff produces an effective vocabulary size $K$ that is manageable. The unknown bucket then represents all rare categories.

What to Expect During Training

Embeddings are learned end to end like any other parameters. They start as random vectors and gradually organize so that categories that are useful in similar ways produce embeddings that lead to similar downstream activations. You do not need to pretrain them for basic tabular tasks.

In practice, if your model trains but validation performance is poor, embeddings can be one contributor, either because the dimensions are too large and overfit, too small and underfit, or because your preprocessing mapping is inconsistent across splits. The debugging chapter covers how to sanity check this by overfitting a small batch and verifying that the model can learn when the data is simple.