Enabling Efficiency-Precision Trade-offs for Label Trees in Extreme Classification

TL;DR

This paper introduces an information theory-inspired algorithm that allows flexible trade-offs between efficiency and accuracy in label tree construction for extreme classification, improving latency without sacrificing accuracy.

Contribution

The authors propose a novel algorithm enabling interpolation between statistical performance and latency in label tree design for XMC, addressing a key challenge in the field.

Findings

Reduces expected latency by up to 28% on Wiki-500K dataset

Improves latency by up to 20% on e-commerce datasets

Maintains accuracy while optimizing latency

Abstract

Extreme multi-label classification (XMC) aims to learn a model that can tag data points with a subset of relevant labels from an extremely large label set. Real world e-commerce applications like personalized recommendations and product advertising can be formulated as XMC problems, where the objective is to predict for a user a small subset of items from a catalog of several million products. For such applications, a common approach is to organize these labels into a tree, enabling training and inference times that are logarithmic in the number of labels. While training a model once a label tree is available is well studied, designing the structure of the tree is a difficult task that is not yet well understood, and can dramatically impact both model latency and statistical performance. Existing approaches to tree construction fall at an extreme point, either optimizing exclusively for statistical performance, or for latency. We propose an efficient information theory inspired algorithm to construct intermediary operating points that trade off between the benefits of both. Our algorithm enables interpolation between these objectives, which was not previously possible. We corroborate our theoretical analysis with numerical results, showing that on the Wiki-500K benchmark dataset our method can reduce a proxy for expected latency by up to 28% while maintaining the same accuracy as Parabel. On several datasets derived from e-commerce customer logs, our modified label tree is able to improve this expected latency metric by up to 20% while maintaining the same accuracy. Finally, we discuss challenges in realizing these latency improvements in deployed models.