Compositional Embeddings Using Complementary Partitions for Memory-Efficient Recommendation Systems (1909.02107v2)

Abstract: Modern deep learning-based recommendation systems exploit hundreds to thousands of different categorical features, each with millions of different categories ranging from clicks to posts. To respect the natural diversity within the categorical data, embeddings map each category to a unique dense representation within an embedded space. Since each categorical feature could take on as many as tens of millions of different possible categories, the embedding tables form the primary memory bottleneck during both training and inference. We propose a novel approach for reducing the embedding size in an end-to-end fashion by exploiting complementary partitions of the category set to produce a unique embedding vector for each category without explicit definition. By storing multiple smaller embedding tables based on each complementary partition and combining embeddings from each table, we define a unique embedding for each category at smaller memory cost. This approach may be interpreted as using a specific fixed codebook to ensure uniqueness of each category's representation. Our experimental results demonstrate the effectiveness of our approach over the hashing trick for reducing the size of the embedding tables in terms of model loss and accuracy, while retaining a similar reduction in the number of parameters.

• The paper presents a quotient-remainder trick that splits large embedding tables into two smaller tables, reducing memory usage from O(|S|×D) to O(√|S|×D).
• It extends this idea with complementary partitions, achieving scalable compression with parameter counts around O(k|S|^(1/k)×D) and mitigating hash collisions.
• Experimental results on DCN and Facebook DLRM architectures demonstrate up to a 15-fold reduction in model size with negligible impact on accuracy.

Compositional Embeddings Using Complementary Partitions for Memory-Efficient Recommendation Systems

The paper "Compositional Embeddings Using Complementary Partitions for Memory-Efficient Recommendation Systems" proposes an innovative approach to reducing the memory footprint of embedding tables in recommendation systems. The authors address the challenge posed by the high dimensionality and large category set sizes in modern deep learning-based recommendation models (DLRMs), which result in significant memory usage during both training and inference.

Theoretical Contribution and Methodology

The paper introduces a novel strategy for generating unique embeddings through the use of complementary partitions, enabling a significant reduction in the size of embedding tables. Embedding tables represent categorical data as dense vectors in a high-dimensional space and often form the primary memory bottleneck in DLRMs. Standard practices, such as the hashing trick, reduce embedding size by mapping categories to a smaller indexed space but can lead to collisions where different categories share the same embedding, degrading model quality.

Key contributions include:

1. Quotient-Remainder Trick: This method employs both quotient and remainder functions to derive two separate, smaller embedding tables instead of one large table. Embeddings from these tables are combined to form unique vectors, reducing memory usage from $O(|S| \times D)$ to a potential upper limit of $O(\sqrt{|S|} \times D)$.
2. Complementary Partitions: The paper generalizes the quotient-remainder trick to include multiple complementary partitions of the category set. These partitions, which need each category to be distinct according to at least one partition, enable a more robust and scalable embedding compression, potentially reducing parameter counts to $O(k |S|^{1/k} \times D)$.
3. Compositional Embeddings: The notion of composing embeddings through different operations (such as addition, element-wise multiplication, or concatenation) further exemplifies the paper’s emphasis on efficient model parameterization. The element-wise multiplication operations, in particular, yielded scalable and effective results in empirical tests.

Experimental Results

Experiments conducted on DCN and Facebook DLRM architectures using the Kaggle Criteo Ad Display Challenge dataset demonstrated the efficacy of this approach. The compositional method showed superior performance in preserving model accuracy compared to the hashing trick, with up to a 15-fold reduction in model size without notable loss in accuracy. The paper reports consistent improvements in model performance, emphasizing the potential for broader application in large-scale recommendation systems where embedding memory costs are prohibitive.

Implications and Future Work

The findings have significant implications for the design of memory-efficient deep learning models in recommendation systems. This approach facilitates the scaling of complex models without incurring substantial memory costs, opening avenues for more extensive online personalization and real-time data processing. The ability to maintain high-quality recommendations with reduced computational overhead ensures broader accessibility and practical deployment at scale.

Future directions could involve exploring adaptive or data-driven partitioning strategies that could dynamically adjust to the underlying data distribution, potentially enhancing both performance and efficiency. The exploration of path-based compositional embeddings in greater detail might yield even better memory and computation trade-offs, considering various transformation functions and their impacts on model efficacy.

Overall, this paper presents a marked advancement in the field of model compression strategies, offering a structured method to handle the prohibitive memory requirements typical of current DLRMs. By leveraging complementary partitions, this work lays a foundation for further research into embedding optimization and the dynamic balancing of computational resources against predictive performance in recommendation systems.