Source-linked AI summary
Neural Input Search for Large Scale Recommendation Models
Manas R. Joglekar, Cong Li, Jay K. Adams, Pranav Khaitan, Quoc V. Le
TL;DR
Large-scale recommendation models use embeddings for categorical features, but their vocabulary and embedding sizes are often selected heuristically. NIS uses reinforcement learning to search these choices and introduces item-varying Multi-size Embedding; experiments report improvements on retrieval and ranking metrics over manually crafted baselines.
Problem
In large-scale recommendation, categorical features have large vocabularies and embedding parameters can dominate model size, yet vocabulary and embedding sizes are often chosen heuristically.
Source: §1 · p. 2Method
NIS uses reinforcement learning to search vocabulary and embedding sizes, while Multi-size Embedding varies vector dimensions across feature items.
Source: §1 · p. 2, §3.2 · p. 3Results
6.8% relative improvement on Recall@1 and 1.8% on ROC-AUC over well established manually crafted baselines were reported across retrieval and ranking problems.
Source: §1 · p. 2Takeaways & Limitations
NIS and Multi-size Embedding achieved high coverage of tail items while keeping accurate representations for head items without increasing total model parameters.
Source: §5 · p. 9
Abstract
from arXiv · showhide
Recommendation problems with large numbers of discrete items, such as products, webpages, or videos, are ubiquitous in the technology industry. Deep neural networks are being increasingly used for these recommendation problems. These models use embeddings to represent discrete items as continuous vectors, and the vocabulary sizes and embedding dimensions, although heavily influence the model's accuracy, are often manually selected in a heuristical manner. We present Neural Input Search (NIS), a technique for learning the optimal vocabulary sizes and embedding dimensions for categorical features. The goal is to maximize prediction accuracy subject to a constraint on the total memory used by all embeddings. Moreover, we argue that the traditional Single-size Embedding (SE), which uses the same embedding dimension for all values of a feature, suffers from inefficient usage of model capacity and training data. We propose a novel type of embedding, namely Multi-size Embedding (ME), which allows the embedding dimension to vary for different values of the feature. During training we use reinforcement learning to find the optimal vocabulary size for each feature and embedding dimension for each value of the feature. In experiments on two common types of large scale recommendation problems, i.e. retrieval and ranking problems, NIS automatically found better vocabulary and embedding sizes that result in $6.8\%$ and $1.8\%$ relative improvements on Recall@1 and ROC-AUC over manually optimized ones.
1 Introduction
Large-scale recommendation models depend heavily on categorical-feature embeddings, but vocabulary and embedding sizes are often chosen heuristically. NIS searches these input configurations and introduces Multi-size Embedding, with reported gains over manually crafted baselines.
1 Introduction
- Categorical recommendation features can make the input component more critical than representation learning because their large vocabularies require many embedding parameters.The cited video-ID example uses 256 million parameters for that feature alone, compared with three fully connected layers in the representation component.Source: §1 · p. 2
- NIS automatically searches vocabulary and embedding sizes for each discrete input feature, using a reinforcement-learning algorithm in a single training run.Its search space consists of Embedding Blocks whose combinations represent different configurations.Source: §1 · p. 2
- Multi-size Embedding assigns larger vectors to common or predictive items and smaller vectors to less common or predictive items, unlike uniform Single-size Embedding.The paper argues this can cover more items at the same model capacity while reducing training-data and computation needs for long-tail items.Source: §1 · p. 2
- 6.8% relative improvement on Recall@1 and 1.8% on ROC-AUC over well established manually crafted baselines were achieved by NIS in one training run.Experiments covered retrieval and ranking problems using data collected from the company’s products.Source: §1 · p. 2
2 Related Work
Prior neural architecture search largely targeted representation learning, especially for image and video tasks. This work identifies the input component of large-scale recommendation models as an underexplored target for automated design.
2 Related Work
- NAS methods sought cheaper architecture search by connecting smaller components and learning a subset in one training run, as in ENAS and related approaches.Other approaches included hyper-networks, differentiable connection weights, and architecture embeddings.Source: §2 · p. 2
- Other NAS work reduced search-space size by searching reusable convolutional cells or hierarchical cell blocks, with some cells transferred across dataset sizes.The cited examples include NASNet transfer learning and MNAS’s separately searched blocks.Source: §2 · pp. 2–3
- Earlier NAS work focused on representation learning for image or video tasks, while recommendation input components with large embeddings were often designed heuristically.The authors describe their work as bringing automated neural network design to the input component for large-scale recommendation problems.Source: §2 · p. 3
3 Neural Input Search
Neural Input Search chooses vocabulary sizes and embedding dimensions for categorical features under a total embedding-memory budget. It also introduces Multi-size Embeddings, which allocate different dimensions to frequent and tail items.
3.1 Definitions and Notations
- A feature vocabulary is frequency-ordered, and an embedding matrix with vocabulary size v and dimension d uses v × d values toward memory budget C.The budget is the total number of floating-point values used by all embedding matrices.Source: §3.1 · p. 3
3.2 Neural Input Search Problems
- Single-size Embeddings make coverage of tail items and representation quality for head items difficult to balance; Multi-size Embeddings allocate larger dimensions to frequent items and smaller ones to tail items.NIS-SE searches vocabulary size and embedding dimension per feature, while NIS-ME searches a Multisize Embedding Spec under the objective and memory constraint.Source: §3.2 · pp. 3–4
3.3 Neural Input Search Approach
- NIS uses an ENAS variant to search in one training run, training selected embeddings with the main model and using validation rewards to train an A3C controller.The controller selects an SE or ME for each feature at each step.Source: §3.3 · p. 5
3.3.1 Search Space
- Embedding Blocks discretize a maximum embedding matrix, letting the controller sample either a subgrid for an SE or a custom subset for an ME.The controller samples SE choices from softmax probabilities; ME choices are a sequence of selections across embedding dimensions.Source: §3.3.1 · pp. 5–6
3.3.2 Reward
- The controller reward combines the task objective with a memory cost-loss; retrieval uses Sampled Recall@1 as a proxy, while ranking uses ROC-AUC.The paper reports Sampled Recall@1 as a good proxy for Recall@1 and says the controller converges to better results when ranking uses ROC-AUC rather than negative cross-entropy.Source: §3.3.2 · pp. 6–7
3.3.3 Training
- The main model trains on training examples while validation examples provide controller rewards, and a warm-up phase trains all Embedding Blocks before alternating model and controller training.The warm-up is intended to ensure that all blocks receive some training before the controller begins choosing among them.Source: §3.3.3 · p. 7
4 Experiments
NIS was evaluated on retrieval and ranking tasks, where its searched Multi-size Embeddings (MEs) outperformed the baseline under the reported memory-budget setup. The experiments also compared searched Single-size Embeddings (SEs) with baselines and used one training run to find optimal configurations.
Query Suggest Retrieval Problem
- The retrieval task used four query-derived features to retrieve among 20 million queries, with a total embedding memory budget of C = 2560M.The controller objective for ME search was Sampled Recall@1 with 5,000 negative examples sampled from the 20 million-query vocabulary.Source: §4 · pp. 7–8
Results
- 6.8% relative improvement on Recall@1 and 6.1% on Recall@5 for retrieval, and 1.8% relative improvement on ROC-AUC for ranking, were achieved by NIS-searched MEs over the baseline.The paper reports these gains for the two recommendation problems under a memory budget.Source: Table 1 · p. 8, Table 2 · p. 8, §4 · p. 8
App Install Ranking Problem
- The ranking task used 20 discrete context and app features with vocabulary sizes ranging from hundreds to millions, and used cross entropy loss.The controller objective for ME search was ROC-AUC calculated on 100 validation examples per controller decision.Source: §4 · p. 8
Implementation
- Each feature used 20 Embedding Blocks, and datasets were split into 70% training, 20% validation, and 10% testing.The block settings used vocabulary fractions [0.1v, 0.2v, 0.2v, 0.2v, 0.3v] and embedding-dimension fractions [0.25d, 0.25d, 0.25d, 0.25d].Source: §4 · p. 8
Results
- NIS-searched SEs outperformed the baselines on both problems, with optimal SE configurations found in one training run rather than separate training for each candidate configuration.The baselines trained a model from scratch for each candidate SE configuration.Source: §4 · p. 8
5 Conclusion
The paper concludes that NIS automatically searches vocabulary and embedding sizes, while ME varies embedding dimensions to represent tail and head items. On large-scale retrieval and ranking problems, the approach reported relative gains of 6.8% on Recall@1 and 1.8% on ROC-AUC in one training run without increasing model parameters.
Conclusion
- NIS automatically searches vocabulary and embedding sizes in the model input component, and ME varies embedding dimensions to give tail items high coverage while retaining accurate head-item representations.Experiments on large-scale retrieval and ranking problems reported relative improvements of 6.8% on Recall@1 and 1.8% on ROC-AUC in one training run, without increasing the model's total parameters.Source: §5 · p. 9