Anytime-Budget Embedding Allocation for DLRMs via Continuous-Capacity Multiple-Choice Knapsack

Abstract

Embedding tables dominate the memory footprint of deep learning recommendation models (DLRMs), yet existing compression methods usually assume that the deployment memory budget is fixed before training. This is inconvenient in production, where the same recommender may be served across hardware tiers, multi-tenant resource conditions, or changing A/B-test operating points. We study anytime-budget embedding allocation: train one DLRM once, then export a hard budget-feasible embedding layout for any memory target in a specified range. We propose Continuous-Capacity Entropy-Smoothed MCKP (CC-MCKP), a differentiable allocation layer that integrates the real-valued budget directly into the multiple-choice knapsack partition function through a smooth feasibility gate. Differentiating this continuous-capacity partition function yields per-table action marginals for training and an analytic budget shadow price for sensitivity analysis, while final deployment still uses a hard MCKP solve with conservative discretization to guarantee feasibility. We provide smoothness, marginal, shadow-price, hard-limit, finite-gate, and grid-consistency results for the layer. Across Criteo Kaggle and Criteo Terabyte, one CC-MCKP-trained model approaches independently trained per-budget specialists (within about 10-3 AUC on average) and provides the strongest average single-run accuracy, with its clearest gains at intermediate and unseen budgets and small reversals at the minimum-budget corner, where the feasible layout is unique. Hard export preserves zero budget violations, while one training run costs about 2× a plain model instead of the 13× of a per-budget specialist suite.