Doubly-stochastic mining for heterogeneous retrieval

Information retrieval concerns finding documents that are most relevant for a given query, and is a canonical real-world use case for machine learning [Manning et al., 2008]. The simplest incarnation of retrieval models involves learning a real-valued scoring function that ranks, for each example, the set of possible labels it may be matched to. A core challenge is scalability: there may be billions of examples (e.g., user queries) and labels (e.g., videos in a recommendation system), each of whose scores naïvely needs to be updated at every training iteration. Effective means of addressing both problems have been widely studied [Mikolov et al., 2013, Jean et al., 2015, Reddi et al., 2019]. A distinct challenge is heterogeneity: the distribution over examples is often a mixture of diverse subpopulations (e.g., queries may arise from geographically disparate user bases). Naïve training on such data may lead to models that perform disproportionately well on one subpopulation at the expense of others; e.g., if queries originate from multiple countries, the retrieval model may only perform well on queries from the dominant country. Such behaviour is clearly undesirable.