Collaborative Research: EAGER: End-to-end Neural Training for Very Large Output Spaces

Modern machine learning models often need to make predictions with an enormous amount of choices. For example, on the internet, search engines need to predict the most relevant candidate for a given query from billions of potential candidates. There are similar prediction problems that are ubiquitous in many search, retrieval and recommendation systems in our daily lives. It is challenging for a machine learning algorithm to deal with a large output space in both the training and inference phases, as any linear scan through all candidates is computationally prohibitive. This project aims to develop a family of scalable and reliable algorithms to tackle the problem of predicting in a large output space. To develop an end-to-end solution, we will tackle the problem of designing novel architectures, and accompanying training and inference procedures that jointly optimize inference speed and prediction accuracy. These efforts will eventually produce a comprehensive toolkit for learning with large output spaces, thus enabling its application in both practical systems and future research activities. The project will also support students and train them in conducting research activities in collaboration with application domains. Existing approaches for dealing with a large output space split the prediction task into two separate components: a neural network encoder and an approximate nearest neighbor search module. The neural network encoder encodes queries and items into a latent space, while the nearest neighbor search module finds the closest vectors in the database for a given query vector. This two-stage approach simplifies the development of each module, but this splitting of components is not focused on end-to-end prediction performance, and thus compromises accuracy and efficiency. The core challenging technical direction of this project is to create algorithms that allow the two components to be aware of each other and thus develop an end-to-end model and training algorithm to handle very large output space. This research direction will be addressed through the development of a novel end-to-end neural network architecture that contains both encoders and trainable search modules. The end-to-end training process will enable direct optimization for precision and efficiency in a single step, instead of requiring two separate steps. This award reflects NSF's statutory mission and has been deemed worthy of support through evaluation using the Foundation's intellectual merit and broader impacts review criteria.