MOGIC: Metadata-infused Oracle Guidance for Improved Extreme Classification

We thank the reviewer for their detailed review. Please find our response to your comments below.

1. The two-phase training process involving Oracle training and Oracle-guided disciple training might introduce additional complexity.

• Response: We agree with the reviewer that the two-stage training framework introduces some additional complexity compared to single-stage baselines. However, the improvements in performance and the increased robustness achieved through regularization in this two-stage approach present a reasonable tradeoff for the added complexity. Moreover, this complexity is limited to the training phase (a one-time cost), while the deployed disciple models remain highly efficient at inference.

2. Although it shows good performance on the tested datasets, its scalability to extremely large-scale or rapidly evolving datasets remains to be seen.

• Response: The MOGIC approach can be readily scaled to larger datasets. However, there are no high-quality rapidly-evolving dataset for XC, to the best of our understanding. We believe this is a strong area for future research. The Wikipedia-500k is one of the largest datasets used by the XC community, comprising 1.8M training samples and 501K labels, while many other XC datasets such as EURLex-4.3K, Bibtex or AmazonCat-13K have significantly fewer labels (such as 4.3K in the EURLex-4.3K dataset or 13K in the AmazonCat-13K dataset). While LF-AmazonTitles-1.3M is another viable choice, we reported results on Amazon-131k, since metadata in this setting is not readily available, and must be generated using GPT-based approaches. In light of your comments, we are also working to include experimentation on the LF-AmazonTitles-1.3M dataset. While we will include these results in the final version of the manuscript, given the additional resources involved in generating the metadata and subsequently training the models, we do not have the results ready to report at this time, but are striving to incorporate them before the end of the rebuttal discussion phase.

Thank you for your detailed review. Below are our responses to your comments.

1. Comparisons with LLM-based approaches

• Response: We have already included comparisons of MOGIC against LLaMA and Phi-based Oracles without disciple models, when LoRA-finetuned for label generation (XC task) in Tab 3 (and row 1-3 below). Generative LMs cannot be directly employed for the XC task, and therefore, for these experiments we carry out classification by comparing the embeddings associated with final token of the query and the labels from the last layer of the SLM. Results reported in Tab. 3 of the paper summarizes this experiment (cf. Lines L367-380, col 1). We observe that the (65M sized) MOGIC disciple, trained with a DistilBERT/LLaMA-based oracle outperforms the 7B-scale fine-tuned SLMs used directly without any disciple. We attribute this worse performance to the fact that these SLMs were pre-trained for text generation and then adapted for label generation. However despite their performance, their embeddings, when used for regularization in the MOGIC framework, result in improved performance of the disciple models. We will improve the clarity of this discussion.

• To further investigate performance comparisons against language models, as part of the rebuttal, we also include comparisons against LLaMA, Phi, and GPT when used directly for label generation (row 4-6). We observe that oracle models trained specifically for the XC task perform better than much larger LMs. Also, for the XC task, using the embedding from LMs is better than mapping their generations to the label set since the generations often contain labels which do not belong to the label set.

Oracle metrics on LF-WikiSeeAlsoTitles-320K

2. Metadata dependency during training

• Response: As discussed in Tab. 8 (cf. Sec 4.3), disciples trained with an Oracle are more robust to noisy/missing metadata in comparison to their Oracle counterparts during inference. This robustness is not limited to noise at inference, but also noise present in the metadata during training. To simulate a training scenario wherein reliable metadata linkages are not available, we consider the following setting: We train MOGIC on LF-WikiSeeAlsoTitles-320K dataset wherein 50% of the metadata is replaced with a randomly selected metadata from the corpus. The following table summarizes the linker, the Oracle and disciple models' performance. We observe that MOGIC (OAK) is more robust to noisy metadata, even during training with negligible drop in performance. The MOGIC model trained with noisy metadata, continues to be SOTA, outperforming other baselines.

Linker metrics with 50% noise-added metadata, where 50 % metadata for each query is replaced by random metadata.

Oracle metrics with 50%noisy metadata

Performance of MOGIC (OAK) and OAK with 50%noisy metadata during training

3. Sensitivity of Loss Balancing Hyperparameters

• Response: We now include a sensitivity analysis over the $𝛼$ and $β$ hyperparameters, considering all combinations of $𝛼\in\{0.1,1,10\}$ and $β\in\{0.1,1,10\}$. We performed a sensitivity analysis prior to our experimentation and found that $(𝛼,β) = (1,0.1)$ leads to the best performance, which are the values we use for the experiments in the paper. We now present the ablations on the choice of $𝛼$ and $β$ on the LF-WikiSeeAlsoTitles-320K dataset, with all other hyperparameters as defined in Appendix H. This table summarizes the results. We observe that $(𝛼,β) = (1,0.1)$ results in the best performance. However, the performance of MOGIC is generally robust to the choice of $(𝛼,β)$ with a variance of 0.108 in P@1, when comparing across all the choices considered. We will improve the clarity of this hyperparameter choice in the final version of the manuscript, and include this ablation in the Appendix.

Thank you for your detailed review. Please find our response to your comments below.

1. Clarity on Rademacher constants

• Response: The Rademacher complexity constants $R_q$ and $R_l$ in Theorem 1 are scalar values which quantify the complexity or capacity of the hypothesis classes corresponding to the query tower and label tower, respectively. We use the standard definitions of Rademacher constants from the Statistical Learning Theory for Binary Classification problems (ref: https://www.cs.cmu.edu/~ninamf/ML11/lect1117.pdf).
• Explanation on Rademacher complexity: Mathematically, Rademacher complexity constants $R_q,R_l$ are estimated as the average empirical loss of minimizing the hypothesis class on a data sample with randomly annotated labels i.e. labels are generated by a purely random Bernoulli distribution with probabilities 0.5 to be either positive or negative. In intuitive terms, the smaller the values of $R_q, R_l$, the less prone are the query tower and label tower to overfit the finite training data, and consequently the accuracy on the test set is expected to be better.

2. Novelty of disciple training and clarity of oracle architecture

• Response: The disciple training in MOGIC is a novel variant of knowledge distillation (KD). Compared to standard KD, it differs in multiple ways. First, the metadata is provided as early concatenation in textual form for the oracle, but is used to train free parameters in the disciple's memory via a novel regularization framework. This can be viewed as KD from an early-concatenation model to a two-tower model. While all disciple models benefit from this framework, the additional parameters present in the memory-based disciples such as MOGIC (OAK) demonstrate the most gains. Furthermore, the oracle has access to ground-truth metadata, which is privileged information that is not available to the disciple model (unlike standard KD). In this MOGIC framework, the oracle can either be larger, or of the same size as the disciple. Table 3 of the main manuscript presents comparisons between larger (LLM-based, 2.7B/7B sized) oracle and a disciple-sized (65M sized), DistilBERT oracle (cf. Lines 367-380, column 1).

3. Clarifying XML Labels and Data Scarcity

• Response: We will update the introduction of the manuscript accordingly. We will also include a detailed example, as suggested by you, in the appendix in the revised draft.

4. Testing on Large-Scale XML Datasets like WikiLSHTC-325K or Amazon-670K

• Response: Datasets such as WikiLSHTC-325K (which contains 325K labels) do not contain raw text or label features associated with labels, and therefore have not been directly used for training encoder-based models such as NGAME, DEXA, DEXML, ANCE etc. This is why we chose to report on LF-WikiTitles-500k and LF-Wikipedia-500K dataset (which has 501K labels), which are from the same distribution and have the same task as WikiLSHTC (category prediction), but contain a larger label set, with label features/text available. Similarly, to be aligned with the baseline methods, we use LF-Amazon-131k dataset instead of Amazon-670K. To demonstrate MOGIC's effectiveness on larger-scale datasets, we are working to include experiments on the LF-AmazonTitles-1.3M dataset in the final version of the manuscript. While the results are not yet available due to the resource-intensive process of generating metadata and training models, we aim to include them before the conclusion of the rebuttal phase.

5. Comparison against SLEEC.

• Response: SLEEC does not use label-text for classification, and therefore performs poorer than methods such as Parabel [1] on multiple standard datasets. A direct comparison against SLEEC is challenging, as we were unable to find an up-to-date implementation of the algorithm. Therefore, we choose to compare MOGIC against Parabel, which has been shown to outperform SLEEC across datasets[1]. We present these comparisons on LF-AmazonTitles-131K, LF-WikiSeeAlsoTitles-320K and LF-WikiTitles-500K. These results are reported in the table below, with performance numbers for Parabel obtained from ECLARE[2]:

• Performance of MOGIC(OAK) and Parabel on different benchmark datasets

[1] Prabhu, Yashoteja, et al. "Parabel: Partitioned label trees for extreme classification with application to dynamic search advertising." World Wide Web Conference 2018.

[2] Mittal, Anshul, et al. "ECLARE: Extreme classification with label graph correlations." theWebConf 2021.