MMbeddings: Parameter-Efficient, Low-Overfitting Probabilistic Embeddings Inspired by Nonlinear Mixed Models
Giora Simchoni, Saharon Rosset
TL;DR
MMbeddings address the challenge of high-cardinality categorical features by reframing embeddings as latent random effects in nonlinear mixed models within a variational autoencoder framework. The encoder outputs per-category means and variances, enabling sampling of embeddings with a reconstruction model that generalizes well while dramatically reducing parameter counts, independent of cardinality. Empirical results on simulated and real data show improved predictive performance and reduced overfitting across tasks like classification and regression, with substantial parameter efficiency compared to standard embeddings. This probabilistic, NLMM-grounded approach promises scalable, robust embeddings for large-scale categorical data and suggests extensions to sequences and knowledge-graph-informed settings.
Abstract
We present MMbeddings, a probabilistic embedding approach that reinterprets categorical embeddings through the lens of nonlinear mixed models, effectively bridging classical statistical theory with modern deep learning. By treating embeddings as latent random effects within a variational autoencoder framework, our method substantially decreases the number of parameters -- from the conventional embedding approach of cardinality $\times$ embedding dimension, which quickly becomes infeasible with large cardinalities, to a significantly smaller, cardinality-independent number determined primarily by the encoder architecture. This reduction dramatically mitigates overfitting and computational burden in high-cardinality settings. Extensive experiments on simulated and real datasets, encompassing collaborative filtering and tabular regression tasks using varied architectures, demonstrate that MMbeddings consistently outperforms traditional embeddings, underscoring its potential across diverse machine learning applications.
