Go to ICLR 2025 Conference homepageMOGIC: METADATA-INFUSED ORACLE GUIDANCE FOR IMPROVED EXTREME CLASSIFICATION
Retrieval-augmented classification and generation models significantly benefit from the early-stage fusion of high-quality text-based auxiliary metadata, often called memory, but they suffer from high inference latency and poor robustness to noise. In classifications tasks, particularly the extreme classification (XC) setting, where low latency is critical, existing methods incorporate metadata for context enrichment via an XC-based retriever and obtain the representations of the relevant memory items to perform late-stage fusion to achieve low latency. With an aim of achieving higher accuracy while meeting the low latency constraints, in this paper, we propose MOGIC, an approach for metadata-infused Oracle guidance for XC tasks. In particular, we train an early-fusion Oracle classifier with access to both query-side and label-side ground-truth metadata in the textual form. The Oracle is subsequently used to guide the training of any existing memory-based XC Disciple model via regularization. The MOGIC algorithm, when applied to memory-based XC Disciple models such as OAK, improves precision@1 and propensity-scored precision@1 by ~2% on four standard datasets, at no additional inference-time costs to the Disciple model. We also show the feasibility of applying the MOGIC algorithm to improve the performance of state-of-the-art memory-free XC approaches such as NGAME or DEXA, demonstrating that the MOGIC algorithm can be used atop any existing XC-based approach in a plug-and-play manner. Finally, we also show the robustness of the MOGIC method to missing and noisy metadata settings. We will release code on acceptance.