Embedding optimization for machine learning models

Abstract

Methods, systems, and computer programs are presented for determining parameters of neural networks and selecting embedding dimensions for the feature fields. One method includes an operation for initializing parameters of a neural network and weights for embedding sizes for each feature associated with the neural network. The parameters of the neural network and the weights are iteratively optimized. Each optimization iteration comprises training the neural network with current parameters of the neural network to optimize a value of the weights, and training the neural network with current values of the weights to optimize the parameters of the neural network. Further, the method includes operations for selecting embedding sizes for the features based on the optimized values of the weights, and for training the neural network based on the selected embedding sizes for the features to obtain an estimator model. A prediction is generated utilizing the estimator model.