AutoSrh: An Embedding Dimensionality Search Framework for Tabular Data Prediction

Prediction over tabular data is often a crucial task in many real-life applications. Recent advances in deep learning give rise to various deep models for tabular data prediction. A common and essential step in these models is to vectorize raw input features in tabular data into dense embeddings. Choosing a suitable dimension for each feature is challenging yet necessary to improve model's performance and reduce memory cost of model parameters. Existing solutions to embedding dimensionality search always choose dimensions from a restricted candidate set. This restriction improves the search efficiency but would produce suboptimal embedding dimensions that hurt model's predictive performance. In this paper, we develop AutoSrh, a flexible embedding dimensionality search framework that can select varying dimensions for different features through differentiable optimization. The key idea of AutoSrh is to relax the search space to be continuous and optimize the selection of embedding dimensions via gradient descent. After optimization, AutoSrh performs embedding pruning to derive the mixed embedding dimensions and retrains the model to further improve the performance. Extensive experiments on five real-world tabular datasets demonstrate that AutoSrh can achieve better predictive performance than the existing approaches with 1.1 similar to 1.6x lower training time cost and reserve model's predictive performance while reducing 50 similar to 95% embedding parameters.