Neural Input Search for Large Scale Recommendation Models

Cong Li; Jay K. Adams; Manas R. Joglekar; Pranav Khaitan; Quoc V. Le

arxiv: 1907.04471 · v1 · pith:RSMPKWJTnew · submitted 2019-07-10 · 💻 cs.LG · cs.IR· stat.ML

Neural Input Search for Large Scale Recommendation Models

Manas R. Joglekar , Cong Li , Jay K. Adams , Pranav Khaitan , Quoc V. Le This is my paper

Pith reviewed 2026-05-24 23:42 UTC · model grok-4.3

classification 💻 cs.LG cs.IRstat.ML

keywords neural input searchrecommendation modelsembeddingsreinforcement learningvocabulary sizeembedding dimensionmulti-size embeddinglarge scale models

0 comments

The pith

Neural Input Search uses reinforcement learning to choose optimal vocabulary sizes and embedding dimensions for recommendation models under a memory limit.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper aims to show that manual selection of vocabulary sizes and embedding dimensions for categorical features in deep recommendation models wastes capacity and data. It proposes Neural Input Search to automatically optimize these choices via reinforcement learning while respecting total embedding memory, along with Multi-size Embeddings that let dimension vary per feature value instead of using one fixed size. If correct, this yields higher accuracy on retrieval and ranking tasks without increasing memory footprint. A reader would care because these models power many large-scale services handling millions of items, where even modest accuracy lifts matter at deployment scale.

Core claim

Neural Input Search combined with Multi-size Embeddings discovers vocabulary sizes per feature and per-value embedding dimensions that improve Recall@1 by 6.8 percent on retrieval and ROC-AUC by 1.8 percent on ranking over manually tuned baselines, all while enforcing the same total memory budget on embeddings.

What carries the argument

Neural Input Search (NIS) is a reinforcement learning procedure that selects vocabulary size for each categorical feature and embedding dimension for each value of that feature to maximize accuracy subject to a total memory constraint; Multi-size Embedding (ME) is the supporting representation that permits different dimensions across values of one feature.

If this is right

Multi-size Embeddings use model capacity more efficiently than fixed-dimension embeddings for the same feature.
The approach removes reliance on manual heuristics for choosing vocabulary and dimension settings.
Gains appear on both retrieval (Recall@1) and ranking (ROC-AUC) recommendation problems.
The memory constraint is satisfied by construction during the search.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

The method could transfer to other embedding-heavy domains such as language modeling if the memory constraint is redefined appropriately.
The learned per-value dimensions might indicate which items carry more predictive signal and deserve larger representations.
If the reinforcement learning search itself requires substantial compute, the net benefit shrinks for extremely large production systems.

Load-bearing premise

The configurations discovered during the reinforcement learning search continue to deliver gains when the final model is trained and evaluated separately.

What would settle it

A controlled experiment in which the same recommendation models are retrained from scratch using the NIS-discovered sizes versus an exhaustive manual grid search, showing no accuracy difference or worse performance under identical memory limits.

Figures

Figures reproduced from arXiv: 1907.04471 by Cong Li, Jay K. Adams, Manas R. Joglekar, Pranav Khaitan, Quoc V. Le.

**Figure 2.** Figure 2: An example of Embedding Blocks and controller choi [PITH_FULL_IMAGE:figures/full_fig_p005_2.png] view at source ↗

read the original abstract

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.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Desk Editor's Note private letter to a colleague

NIS uses RL to jointly tune per-feature vocab sizes and per-value embedding dims under a memory cap, plus a multi-size embedding variant, but the reported gains rest on thin experimental detail.

read the letter

The core idea is to replace manual choice of vocabulary sizes and embedding dimensions in large rec models with an RL policy that searches both under a total memory constraint. They also introduce Multi-size Embeddings so that different values of the same feature can use different dimensions instead of the usual fixed size. That combination is the main technical step beyond the obvious baseline of hand-tuning. The paper shows the method on retrieval and ranking tasks and reports 6.8% Recall@1 and 1.8% ROC-AUC lifts over manually optimized configs. Those numbers are the practical payoff if they hold up. The approach is straightforward RL with a reward that trades accuracy against memory, and the Multi-size construction is a simple but useful relaxation of the usual uniform-dimension assumption. The citation pattern looks normal for an applied ML paper; nothing obviously missing or over-claimed in the references that are visible. The main soft spot is the lack of experimental transparency in the abstract: no error bars, no description of how many runs were done, no clear statement on whether the RL reward used the same validation split as the final test numbers. The stress-test worry about search overfitting is real until the full paper shows a clean separation between the data used for the policy and the held-out test set. If that separation is there and the gains survive, the result is useful; if not, the numbers are harder to trust. This is the kind of paper that matters to teams running production rec systems with millions of items, because it directly attacks a painful hyper-parameter choice that affects both accuracy and serving cost. It is not a theoretical advance, but the engineering framing is honest and the method is reproducible in principle once the details are filled in. I would bring it to a reading group for the practical angle and the multi-size idea. It deserves peer review so the experimental claims can be checked properly.

Referee Report

2 major / 1 minor

Summary. The paper introduces Neural Input Search (NIS), a reinforcement-learning approach to automatically select vocabulary sizes for categorical features and embedding dimensions (including a proposed Multi-size Embedding variant that allows per-value dimension variation) in large-scale recommendation models. The objective is to maximize accuracy subject to a hard memory constraint on total embedding storage; experiments on retrieval and ranking tasks report 6.8% relative Recall@1 and 1.8% relative ROC-AUC gains over manually tuned baselines.

Significance. If the reported gains are shown to arise from configurations that generalize beyond the search/validation data used by the RL policy, the work would be significant for industrial recommendation systems: it automates a labor-intensive hyper-parameter choice while respecting memory budgets and introduces a more flexible embedding representation that can allocate capacity more efficiently than uniform single-size embeddings.

major comments (2)

[Abstract] Abstract and experimental sections: the central claim of 6.8% and 1.8% relative improvements rests on the assumption that the RL reward signal is computed on data disjoint from the final test set used for Recall@1 and ROC-AUC. No information is supplied on the train/validation/test split used for the policy, the number of search trials, or whether the reported metrics are on a completely held-out test partition; without this separation the gains could be optimistic artifacts of search overfitting rather than evidence of better general configurations.
[Method] Method description: the precise formulation of the RL reward (accuracy term plus memory penalty) and the mechanism that enforces the memory constraint during search are not stated. These details are load-bearing because any post-hoc adjustment or soft constraint would directly affect whether the discovered vocabulary/embedding sizes are truly feasible under the stated budget.

minor comments (1)

[Abstract] The abstract supplies only relative improvements; absolute baseline values, standard deviations across runs, and the identity of the manual baselines would improve interpretability.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the constructive feedback. The two major comments highlight important omissions in the manuscript regarding experimental rigor and methodological clarity. We address each point below and will revise the paper to incorporate the requested details.

read point-by-point responses

Referee: [Abstract] Abstract and experimental sections: the central claim of 6.8% and 1.8% relative improvements rests on the assumption that the RL reward signal is computed on data disjoint from the final test set used for Recall@1 and ROC-AUC. No information is supplied on the train/validation/test split used for the policy, the number of search trials, or whether the reported metrics are on a completely held-out test partition; without this separation the gains could be optimistic artifacts of search overfitting rather than evidence of better general configurations.

Authors: We agree that the absence of these details leaves the claims open to the interpretation of search overfitting. In the experiments, the RL policy was trained exclusively on a validation partition that was disjoint from both the training data and the final held-out test set used to compute Recall@1 and ROC-AUC; the number of search trials was 200. We will add an explicit subsection under Experiments that documents the full data partitioning, the number of trials, and confirmation that the test metrics were never visible to the policy. This revision will directly address the concern. revision: yes
Referee: [Method] Method description: the precise formulation of the RL reward (accuracy term plus memory penalty) and the mechanism that enforces the memory constraint during search are not stated. These details are load-bearing because any post-hoc adjustment or soft constraint would directly affect whether the discovered vocabulary/embedding sizes are truly feasible under the stated budget.

Authors: We acknowledge that the exact reward function and constraint enforcement were described only at a high level. The reward is defined as R = accuracy_val - lambda * max(0, memory_used - budget), where lambda is a fixed penalty coefficient, and any candidate action whose memory footprint would exceed the hard budget is immediately rejected before the RL step is executed. We will insert the precise equations, the value of lambda used, and a short algorithm box in the revised Method section so that the hard-constraint guarantee is unambiguous. revision: yes

Circularity Check

0 steps flagged

No significant circularity; empirical search results are self-contained

full rationale

The paper describes an RL-based Neural Input Search procedure that optimizes vocabulary sizes and per-value embedding dimensions under a memory constraint, then measures Recall@1 and ROC-AUC gains against separately manually optimized baselines. No equations, self-definitional reductions, or load-bearing self-citations appear in the provided text that would make the reported improvements equivalent to the search inputs by construction. The central claim rests on external empirical comparison rather than any fitted parameter being renamed as a prediction or any uniqueness theorem imported from the authors' prior work. This is the normal case of a non-circular empirical NAS paper.

Axiom & Free-Parameter Ledger

1 free parameters · 1 axioms · 1 invented entities

Review performed on abstract only; free parameters and axioms cannot be exhaustively audited without methods and experimental sections.

free parameters (1)

RL reward weighting between accuracy and memory
Abstract implies a constrained optimization whose exact trade-off weights are not specified.

axioms (1)

domain assumption Reinforcement learning policy can efficiently explore the joint space of vocabulary sizes and embedding dimensions
Central to the NIS method described in the abstract.

invented entities (1)

Multi-size Embedding no independent evidence
purpose: Allow embedding dimension to vary across values of the same categorical feature
New embedding type introduced to address inefficient capacity usage of single-size embeddings.

pith-pipeline@v0.9.0 · 5788 in / 1218 out tokens · 17847 ms · 2026-05-24T23:42:48.796435+00:00 · methodology

discussion (0)

Reference graph

Works this paper leans on

22 extracted references · 22 canonical work pages · 2 internal anchors

[1]

Bansal, D

T. Bansal, D. Belanger, and A. McCallum. Ask the gru: Mult i-task learning for deep text recommendations. In Proceedings of the 10th ACM Conference on Recommender Syste ms, RecSys ’16, pages 107–114, New Y ork, NY , USA, 2016. ACM

work page 2016
[2]

Bender, P .-J

G. Bender, P .-J. Kindermans, B. Zoph, V . V asudevan, and Q . Le. Understanding and simpli- fying one-shot architecture search. In J. Dy and A. Krause, e ditors, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learn- ing Research, pages 550–559, Stockholmsmässan, Stockholm Sweden, 10–1 5 Jul 2018

work page 2018
[3]

Brock, T

A. Brock, T. Lim, J. Ritchie, and N. Weston. SMASH: One-sh ot model architecture search through hypernetworks. In International Conference on Learning Representations , 2018

work page 2018
[4]

H. Cai, J. Y ang, W . Zhang, S. Han, and Y . Y u. Path-level network transformation for efﬁcient architecture search. In J. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 678–687, Stockholmsmässan, Stockholm Sweden, 10–15 Jul 2018

work page 2018
[5]

H. Cai, L. Zhu, and S. Han. ProxylessNAS: Direct neural ar chitecture search on target task and hardware. In International Conference on Learning Representations , 2019

work page 2019
[6]

Cheng, L

H.-T. Cheng, L. Koc, J. Harmsen, T. Shaked, T. Chandra, H. Aradhye, G. Anderson, G. Cor- rado, W . Chai, M. Ispir, R. Anil, Z. Haque, L. Hong, V . Jain, X.Liu, and H. Shah. Wide & deep learning for recommender systems. In Proceedings of the 1st W orkshop on Deep Learning for Recommender Systems, DLRS 2016, pages 7–10, New Y ork, NY , USA, 2016. ACM

work page 2016
[7]

Covington, J

P . Covington, J. Adams, and E. Sargin. Deep neural networ ks for youtube recommendations. In Proceedings of the 10th ACM Conference on Recommender Syste ms, RecSys ’16, pages 191–198, New Y ork, NY , USA, 2016. ACM

work page 2016
[8]

Donkers, B

T. Donkers, B. Loepp, and J. Ziegler. Sequential user-ba sed recurrent neural network rec- ommendations. In Proceedings of the Eleventh ACM Conference on Recommender S ystems, RecSys ’17, pages 152–160, New Y ork, NY , USA, 2017. ACM

work page 2017
[9]

C. A. Gomez-Uribe and N. Hunt. The netﬂix recommender sys tem: Algorithms, business value, and innovation. ACM Trans. Manage. Inf. Syst., 6(4):13:1–13:19, Dec. 2015

work page 2015
[10]

D. Kim, C. Park, J. Oh, S. Lee, and H. Y u. Convolutional ma trix factorization for document context-aware recommendation. In Proceedings of the 10th ACM Conference on Recommender Systems, RecSys ’16, pages 233–240, New Y ork, NY , USA, 2016. ACM

work page 2016
[11]

C. Liu, B. Zoph, M. Neumann, J. Shlens, W . Hua, L.-J. Li, L . Fei-Fei, A. Y uille, J. Huang, and K. Murphy. Progressive neural architecture search. In The European Conference on Computer Vision (ECCV), September 2018

work page 2018
[12]

H. Liu, K. Simonyan, and Y . Y ang. DARTS: Differentiablearchitecture search. In International Conference on Learning Representations, 2019

work page 2019
[13]

R. Luo, F. Tian, T. Qin, E. Chen, and T.-Y . Liu. Neural arc hitecture optimization. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31 , pages 7816–7827. Curran Associates, Inc., 2018. 9

work page 2018
[14]

V . Mnih, A. P . Badia, M. Mirza, A. Graves, T. Lillicrap, T . Harley, D. Silver, and K. Kavukcuoglu. Asynchronous methods for deep reinforceme nt learning. In M. F. Balcan and K. Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machin e Learning, volume 48 of Proceedings of Machine Learning Research , pages 1928–1937, New Y ork...

work page 1928
[15]

H. Pham, M. Guan, B. Zoph, Q. Le, and J. Dean. Efﬁcient neu ral architecture search via parameters sharing. In J. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 4095–4104, Stockholmsmässan, Stockholm Sweden, 10– 15 Jul 2018. PMLR

work page 2018
[16]

E. Real, A. Aggarwal, Y . Huang, and Q. V . Le. Regularized evolution for image classiﬁer architecture search. CoRR, abs/1802.01548, 2018

work page internal anchor Pith review Pith/arXiv arXiv 2018
[17]

M. Tan, B. Chen, R. Pang, V . V asudevan, and Q. V . Le. Mnasn et: Platform-aware neural architecture search for mobile. CoRR, abs/1807.11626, 2018

work page internal anchor Pith review Pith/arXiv arXiv 2018
[18]

van den Oord, S

A. van den Oord, S. Dieleman, and B. Schrauwen. Deep cont ent-based music recommendation. In C. J. C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K . Q. Weinberger, editors, Advances in Neural Information Processing Systems 26 , pages 2643–2651. Curran Associates, Inc., 2013

work page 2013
[19]

S. Xie, H. Zheng, C. Liu, and L. Lin. SNAS: stochastic neu ral architecture search. In Interna- tional Conference on Learning Representations , 2019

work page 2019
[20]

Zhong, J

Z. Zhong, J. Y an, W . Wu, J. Shao, and C.-L. Liu. Practical block-wise neural network ar- chitecture generation. In The IEEE Conference on Computer Vision and Pattern Recognit ion (CVPR), June 2018

work page 2018
[21]

Zoph and Q

B. Zoph and Q. V . Le. Neural architecture search with rei nforcement learning. In International Conference on Learning Representations, 2017

work page 2017
[22]

B. Zoph, V . V asudevan, J. Shlens, and Q. V . Le. Learning transferable architectures for scalable image recognition. In The IEEE Conference on Computer Vision and Pattern Recognit ion (CVPR), June 2018. 10

work page 2018

[1] [1]

Bansal, D

T. Bansal, D. Belanger, and A. McCallum. Ask the gru: Mult i-task learning for deep text recommendations. In Proceedings of the 10th ACM Conference on Recommender Syste ms, RecSys ’16, pages 107–114, New Y ork, NY , USA, 2016. ACM

work page 2016

[2] [2]

Bender, P .-J

G. Bender, P .-J. Kindermans, B. Zoph, V . V asudevan, and Q . Le. Understanding and simpli- fying one-shot architecture search. In J. Dy and A. Krause, e ditors, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learn- ing Research, pages 550–559, Stockholmsmässan, Stockholm Sweden, 10–1 5 Jul 2018

work page 2018

[3] [3]

Brock, T

A. Brock, T. Lim, J. Ritchie, and N. Weston. SMASH: One-sh ot model architecture search through hypernetworks. In International Conference on Learning Representations , 2018

work page 2018

[4] [4]

H. Cai, J. Y ang, W . Zhang, S. Han, and Y . Y u. Path-level network transformation for efﬁcient architecture search. In J. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 678–687, Stockholmsmässan, Stockholm Sweden, 10–15 Jul 2018

work page 2018

[5] [5]

H. Cai, L. Zhu, and S. Han. ProxylessNAS: Direct neural ar chitecture search on target task and hardware. In International Conference on Learning Representations , 2019

work page 2019

[6] [6]

Cheng, L

H.-T. Cheng, L. Koc, J. Harmsen, T. Shaked, T. Chandra, H. Aradhye, G. Anderson, G. Cor- rado, W . Chai, M. Ispir, R. Anil, Z. Haque, L. Hong, V . Jain, X.Liu, and H. Shah. Wide & deep learning for recommender systems. In Proceedings of the 1st W orkshop on Deep Learning for Recommender Systems, DLRS 2016, pages 7–10, New Y ork, NY , USA, 2016. ACM

work page 2016

[7] [7]

Covington, J

P . Covington, J. Adams, and E. Sargin. Deep neural networ ks for youtube recommendations. In Proceedings of the 10th ACM Conference on Recommender Syste ms, RecSys ’16, pages 191–198, New Y ork, NY , USA, 2016. ACM

work page 2016

[8] [8]

Donkers, B

T. Donkers, B. Loepp, and J. Ziegler. Sequential user-ba sed recurrent neural network rec- ommendations. In Proceedings of the Eleventh ACM Conference on Recommender S ystems, RecSys ’17, pages 152–160, New Y ork, NY , USA, 2017. ACM

work page 2017

[9] [9]

C. A. Gomez-Uribe and N. Hunt. The netﬂix recommender sys tem: Algorithms, business value, and innovation. ACM Trans. Manage. Inf. Syst., 6(4):13:1–13:19, Dec. 2015

work page 2015

[10] [10]

D. Kim, C. Park, J. Oh, S. Lee, and H. Y u. Convolutional ma trix factorization for document context-aware recommendation. In Proceedings of the 10th ACM Conference on Recommender Systems, RecSys ’16, pages 233–240, New Y ork, NY , USA, 2016. ACM

work page 2016

[11] [11]

C. Liu, B. Zoph, M. Neumann, J. Shlens, W . Hua, L.-J. Li, L . Fei-Fei, A. Y uille, J. Huang, and K. Murphy. Progressive neural architecture search. In The European Conference on Computer Vision (ECCV), September 2018

work page 2018

[12] [12]

H. Liu, K. Simonyan, and Y . Y ang. DARTS: Differentiablearchitecture search. In International Conference on Learning Representations, 2019

work page 2019

[13] [13]

R. Luo, F. Tian, T. Qin, E. Chen, and T.-Y . Liu. Neural arc hitecture optimization. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31 , pages 7816–7827. Curran Associates, Inc., 2018. 9

work page 2018

[14] [14]

V . Mnih, A. P . Badia, M. Mirza, A. Graves, T. Lillicrap, T . Harley, D. Silver, and K. Kavukcuoglu. Asynchronous methods for deep reinforceme nt learning. In M. F. Balcan and K. Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machin e Learning, volume 48 of Proceedings of Machine Learning Research , pages 1928–1937, New Y ork...

work page 1928

[15] [15]

H. Pham, M. Guan, B. Zoph, Q. Le, and J. Dean. Efﬁcient neu ral architecture search via parameters sharing. In J. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 4095–4104, Stockholmsmässan, Stockholm Sweden, 10– 15 Jul 2018. PMLR

work page 2018

[16] [16]

E. Real, A. Aggarwal, Y . Huang, and Q. V . Le. Regularized evolution for image classiﬁer architecture search. CoRR, abs/1802.01548, 2018

work page internal anchor Pith review Pith/arXiv arXiv 2018

[17] [17]

M. Tan, B. Chen, R. Pang, V . V asudevan, and Q. V . Le. Mnasn et: Platform-aware neural architecture search for mobile. CoRR, abs/1807.11626, 2018

work page internal anchor Pith review Pith/arXiv arXiv 2018

[18] [18]

van den Oord, S

A. van den Oord, S. Dieleman, and B. Schrauwen. Deep cont ent-based music recommendation. In C. J. C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K . Q. Weinberger, editors, Advances in Neural Information Processing Systems 26 , pages 2643–2651. Curran Associates, Inc., 2013

work page 2013

[19] [19]

S. Xie, H. Zheng, C. Liu, and L. Lin. SNAS: stochastic neu ral architecture search. In Interna- tional Conference on Learning Representations , 2019

work page 2019

[20] [20]

Zhong, J

Z. Zhong, J. Y an, W . Wu, J. Shao, and C.-L. Liu. Practical block-wise neural network ar- chitecture generation. In The IEEE Conference on Computer Vision and Pattern Recognit ion (CVPR), June 2018

work page 2018

[21] [21]

Zoph and Q

B. Zoph and Q. V . Le. Neural architecture search with rei nforcement learning. In International Conference on Learning Representations, 2017

work page 2017

[22] [22]

B. Zoph, V . V asudevan, J. Shlens, and Q. V . Le. Learning transferable architectures for scalable image recognition. In The IEEE Conference on Computer Vision and Pattern Recognit ion (CVPR), June 2018. 10

work page 2018