Table of Contents
Fetching ...

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.

MMbeddings: Parameter-Efficient, Low-Overfitting Probabilistic Embeddings Inspired by Nonlinear Mixed Models

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 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.
Paper Structure (14 sections, 13 equations, 3 figures, 8 tables)

This paper contains 14 sections, 13 equations, 3 figures, 8 tables.

Figures (3)

  • Figure 1: MMbeddings scheme for a single categorical feature, combining an encoder to parameterize the variational posterior $q(\mathbf{b}_j|\mathbf{y}_j)$ and a decoder modeling the nonlinear function $f(\mathbf{x}_{ij}, \mathbf{b}_j)$.
  • Figure 2: Left: Mean runtime per epoch (log-log scale) is shown for varying cardinality $q$ (with sample size $n = 10q$), regression setting. Right: The validation loss plot (normalized) for $q = 10^4$.
  • Figure 3: Predicted vs. true embedding values for each component $\mathbf{b}_1$ to $\mathbf{b}_{10}$, for the $q = 10^3$ cardinality scenario, regression setting. For each true component ($\mathbf{B}$), the most correlated predicted component (column in ($\hat{\mathbf{B}}$) was selected based on Pearson correlation. Axis limits are shared across facets to ensure comparability.