Category Embedding

A Keras-based neural encoder for categorical variables in tabular machine learning.
It learns dense vector representations (embeddings) for categorical features and outputs a clean numeric DataFrame that integrates seamlessly with gradient-boosted tree models.

This library is designed for ML engineers who want the benefits of deep-learning-based embeddings without replacing their GBM models.
It is sklearn-compatible, deterministic, and production-ready.

🚀 Features

Multi-column categorical embeddings
Each categorical feature receives its own learned embedding matrix.
Smart default embedding dimensions
Uses a simple, interpretable rule:
- If n_cat ≤ 10: dim = n_cat - 1
- Else: dim = max(10, n_cat // 2)
- Always capped at 30.
Per-column embedding dimension overrides
Pass a list of integers to control embedding size manually.
Robust categorical handling
- Missing values (NaN/None) → mapped to a dedicated _MISSING_ token with its own trainable embedding.
- Unseen categories at inference → mapped to a separate _UNKNOWN_ token with its own trainable embedding.
Residual MLP architecture
LayerNorm + GELU + Dropout + skip connections for stable training.
Flexible numeric feature output
Numeric columns are always imputed and scaled internally for stable neural training. Use numeric_output to control what appears in transform() output:
- 'raw' (default): Return original numeric values unchanged—ideal for GBMs.
- 'processed': Return imputed + scaled values—ideal for linear models.
- None: Return only categorical embeddings.
Configurable numeric imputation
num_imp_mode chooses between 'mean' or 'median' for internal imputation. If your data has no missing numeric values, this parameter has no effect.
Optional log-scaling of regression targets
For regression tasks, the target can be optionally transformed using log(y + 1e-6) during training (and inverse-transformed at prediction time) by setting log_target=True.
Imbalance-aware training via Focal Loss
For binary classification tasks, set focal_gamma to replace standard binary cross-entropy with Focal Loss. This down-weights easy majority-class examples and concentrates gradient updates on hard, uncertain minority-class examples — producing better embedding geometry when class distributions are skewed. The stronger the imbalance, the more this matters for downstream GBM quality.
Supports regression and binary classification
The neural head is used only for training/tuning; GBMs remain the final predictor.
Optional external validation set
Enables clean early stopping and stable embedding learning.
Sklearn-compatible API
Implements fit, transform, predict, and get_feature_names_out.
Outputs a pandas DataFrame
Perfect for LightGBM, XGBoost, CatBoost, or any sklearn model.
Optional raw categorical passthrough
Use return_raw_categoricals=True to include original categorical values alongside embeddings. This gives GBMs both exact matching (for frequent categories) and similarity signals (for rare/unseen categories).

📦 Installation

pip install category-embedding

Requires:

Python ≥ 3.9
TensorFlow ≥ 2.12
scikit-learn ≥ 1.2
pandas ≥ 1.5

🔧 Quick Start

import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from lightgbm import LGBMRegressor
from category_embedding import CategoryEmbedding

df = pd.DataFrame({
    "country": ["DE", "FR", "DE", "US", "US"],
    "device": ["mobile", "desktop", "tablet", "mobile", "desktop"],
    "age": [25, 40, 31, 22, 35],
})
y = [10.5, 20.1, 15.3, 8.7, 18.0]

categorical = ["country", "device"]
numeric = ["age"]

# CategoryEmbedding acts as a transformer inside a sklearn pipeline
preprocess = ColumnTransformer(
    transformers=[
        ("emb", CategoryEmbedding(
            task="regression",
            categorical_cols=categorical,
            numeric_cols=numeric,
            epochs=20,
            batch_size=32,
            numeric_output='raw',      # Return original numeric values for LightGBM
            num_imp_mode='median',     # Internal imputation strategy (always applied)
        ), categorical + numeric)
    ],
    remainder="drop",
)

model = Pipeline([
    ("prep", preprocess),
    ("lgbm", LGBMRegressor(n_estimators=300, learning_rate=0.05)),
])

model.fit(df, y)
preds = model.predict(df)

Binary classification with imbalanced data

For classification tasks with skewed class distributions, set focal_gamma to activate Focal Loss. This shapes the embedding space to better represent the minority class, which directly benefits the downstream GBM:

from lightgbm import LGBMClassifier

preprocess = ColumnTransformer(
    transformers=[
        ("emb", CategoryEmbedding(
            task="classification",
            categorical_cols=categorical,
            numeric_cols=numeric,
            epochs=30,
            batch_size=256,
            focal_gamma=2.0,           # Activate Focal Loss for imbalanced data
            numeric_output='raw',
        ), categorical + numeric)
    ],
    remainder="drop",
)

model = Pipeline([
    ("prep", preprocess),
    ("lgbm", LGBMClassifier(n_estimators=300, scale_pos_weight=pos_weight)),
])

model.fit(df, y)

Tip: a focal_gamma in the range [1.0, 2.0] works well for moderate imbalance (5:1–10:1). For more extreme ratios, consider values up to 3.0. The GBM's own imbalance handling (e.g. scale_pos_weight in LightGBM) remains complementary and should still be set.

🧠 Why Use Category Embedding?

Traditional encoders struggle with:

High-cardinality categorical features
Sparse interactions
Noisy or rare categories
Missing or unseen values at inference

Neural embeddings solve this by learning dense, continuous representations that capture similarity structure between categories.

This library gives you:

The power of deep learning
The simplicity and performance of GBMs as final predictors
A clean sklearn interface
Deterministic, production-ready behavior
Automatic handling of missing/unseen categoricals via dedicated trainable tokens
Flexible numeric output: raw values for trees, scaled for linear models, or embeddings-only
Imbalance-aware embeddings via Focal Loss for binary classification

⚙️ API Overview