Extreme Classification in Log Memory using Count-Min Sketch: A Case Study of Amazon Search with 50M Products

Abstract

In the last decade, it has been shown that many hard AI tasks, especially in NLP, can be naturally modeled as extreme classification problems leading to improved precision. However, such models are prohibitively expensive to train due to the memory blow-up in the last layer. For example, a reasonable softmax layer for the dataset of interest in this paper can easily reach well beyond 100 billion parameters (>400 GB memory). To alleviate this problem, we present Merged-Average Classifiers via Hashing (MACH), a generic K-classification algorithm where memory provably scales at O(logK) without any strong assumption on the classes. MACH is subtly a count-min sketch structure in disguise, which uses universal hashing to reduce classification with a large number of classes to few embarrassingly parallel and independent classification tasks with a small (constant) number of classes. MACH naturally provides a technique for zero communication model parallelism. We experiment with 6 datasets; some multiclass and some multilabel, and show consistent improvement over respective state-of-the-art baselines. In particular, we train an end-to-end deep classifier on a private product search dataset sampled from Amazon Search Engine with 70 million queries and 49.46 million products. MACH outperforms, by a significant margin,the state-of-the-art extreme classification models deployed on commercial search engines: Parabel and dense embedding models. Our largest model has 6.4 billion parameters and trains in less than 35 hours on a single p3.16x machine. Our training times are 7-10x faster, and our memory footprints are 2-4x smaller than the best baselines. This training time is also significantly lower than the one reported by Google's mixture of experts (MoE) language model on a comparable model size and hardware.

Published at NeurIPS 2019. arXiv admin note: text overlap with arXiv:1810.04254

References in corpus (6)

• Adam: A Method for Stochastic Optimization
• Distributed Representations of Words and Phrases and their Compositionality
• Sequence to Sequence Learning with Neural Networks
• TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems
• Multi-Label Prediction via Compressed Sensing
• Logarithmic Time Online Multiclass prediction

Cited by in corpus (9)

• DeepXML: A Deep Extreme Multi-Label Learning Framework Applied to Short Text Documents
• DECAF: Deep Extreme Classification with Label Features
• On Missing Labels, Long-tails and Propensities in Extreme Multi-label Classification
• Multi-modal Extreme Classification
• Extreme Multi-label Classification from Aggregated Labels
• Deep Retrieval: Learning A Retrievable Structure for Large-Scale Recommendations
• Online probabilistic label trees
• Learning with Holographic Reduced Representations
• CARLS: Cross-platform Asynchronous Representation Learning System