Navigating Extremes: Dynamic Sparsity in Large Output Spaces

It would be good to add training time estimates in addition to memory usage.

The training times of the proposed dynamic sparse approach is close to LightXML, and about 1.5x compared to using a dense last layer (corresponding to the Renee architecture). In order to have a fair timing comparison, all these experiments were run on an A100 GPU, despite the fact that our method would work with cheaper hardware.

Training time in hours:

---

What if DST is applied to fixed embeddings generated from NGAME, DEXA etc.? Why is there a need to train these models end-to-end?

Further results with fixed embeddings obtained from CascadeXML features:

The results demonstrate that training the model end-to-end is beneficial in terms of prediction performance as it provides noticeable improvements over fixed embeddings. Furthermore, the embeddings obtained from other pipelines are dataset specific, and in a realistic application scenario, if the dataset is not derived from any of the pre-existing ones from the Extreme Classification Repository, the only option is to learn the model from scratch in an end-to-end manner.

---

DEXML(Dual-Encoders for Extreme Multi-Label Classification) is a new encoder-based approach that seems to perform at par with XMC models. Please compare the proposed approach to DEXML on a couple of datasets with label features.

Below, we compare our method with DEXML for the LF-AmazonTitles-131K dataset. Please refer to the subsequent answer for the comparison with the LF-WikiSeeAlso-320K dataset.

It would be good to add a couple more datasets with LF features.

We have added below the comparison on another dataset with label features (LF-WikiSeeAlso-320K). These demonstrate that performance at par with state-of-the-art methods DEXML and NGame can be achieved at a fraction of GPU memory consumption and training time.

Please note that the peak memory and training time comparisons are conducted with the same batch size. For the newly added LF-WikiSeeAlso-320K dataset, we used a batch size of 64. The DEXML model, with all negative label settings on the LF-WikiSeeAlso-320K dataset, required more than 160GB of memory (causing OOM on our 4xA100 node server). Therefore, we used the 10 hard negatives settings as described in Table 3 of the DEXML paper.

---

It would be good to compare to Renee with smaller sized encoders.

Thank you for highlighting this point. We realize that our explanation on line 234 could benefit from additional clarity. Our dense baseline is similar to Renee but with a smaller encoder. The key differences are: (i) We use squared hinge loss instead of BCE, (ii) We employ Adam for the extreme layer, whereas Renee uses SGD. These modifications result in a slight performance improvement over Renee. We included this enhanced dense baseline for comparison in the manuscript. For your reference, we provide the results of Renee with a smaller (base) encoder in the table below.

Note that for Wiki500K, the encoder architecture remains the same as in the Renee paper, with the sequence length reduced to 128 for a fair comparison.

---

It would be good to add sensitivity of the model's performance to the new hyperparameters.

Please refer to Table 1 in the Author Rebuttal for the impact of different sparsity levels in conjunction with the use of auxiliary loss, and to Table 2 for a detailed ablation study of hyperparameters related to the auxiliary loss. The ablation study concerning the intermediate layer size is presented in Appendix D (line 568). Also, please note that, while the method introduces hyperparameters such as sparsity levels, it on the other hand does not require any hyperparameters related to hard-negative mining.

---

It would be good to comment on the scalability of the proposed approach given that the real-world datasets have much more than 3M labels

To the best of our knowledge, Amazon3M is the largest dataset that is publicly available; if you know of a dataset we could use, we'd be happy to run these additional experiments.

The authors claim, in line 142, that 2:4 structured sparsity results in deteriorated model accuracy, but there was no reference (other than the 2:4 whitepaper showing only high model qualities), and I couldn't find evidence in the submission that 2:4 sparsity behaved poorly for XMC tasks.

We should have been more precise here. The argument should be three-fold:

• First, the more constaints are put on the sparsity pattern, the less representational capacity a model has; 2:4 sparse models are a tiny subset of all possible sparsity patterns with 50% sparsity.
• Second, despite having been introduced with Ampere, 2:4 sparsity so far hasn't found wide-spread use in model training. Given the difficulty in publishing negative results, it is hard to pinpoint the exact reasons. Note, however, that a recent blogpost by the pytorch team introducing 2:4 sparse training support shows degraded performance of the 2:4 trained model that they need to fix by a short period of fully-dense training in the end.
• Third, for the goal of memory-reducion of this paper, 50% sparsity as induced by 2:4 is not sufficient.

Further, 2:4 sparsity has also proven to be challenging to learn in the DST framework. Leading SOTA methods for learning 2:4 sparsity such as SR-STE [1] require dense gradients every optimization step, thereby eliminating any potential memory savings such as those we achieve in this work. Further, [1] also acknowledged that for N:M sparsity, generalization performance increases as the size of M increases. Fixed fan-in sparsity can be viewed as a particular type of N:M sparsity where M is simply the dense fan-in of each neuron.

The improvement over past work is generally only in memory savings; there are often past techniques that give a higher quality model (though with higher memory requirements). It'd be more compelling if the method matched or exceeded the accuracy of past work with a smaller memory requirement.

In XMC setups, there is always a trade-off between memory/compute and accuracy. Previous methods mainly address the memory consumption problem by projecting to smaller embeddings before the classification layer, e.g., LightXML projects BERT's 768-dimensional embeddings down to 300 dimensions; For Renee, in their experiments with 100M labels, they use only 64 dimensions. In our experiments, we found that DST consistently outperformed bottleneck approaches with matching memory consumption.

What are the concrete novel aspects of your work?

The reviewer is correct to observe that semi-structured sparsity in the last layer is not a novel contribution, and section 2.2 could be moved as part of 2.1. The main contribution of our work are (i) identification on the need of an auxiliary objective to stabilize the gradient flow with a semi-structured sparse layer, (ii) key insight on the role of label clusters (random or k-means) beyond its conventional role in XMC as a negative sampling technique for speeding up training, (iii) armed with this insight, using the meta-branch for achieving gradient stabilization, and (iv) further tuning of the network to achieve training with the largest publicly available dataset on a single commodity GPU with close to state-of-the-art results, which otherwise requires multiple A100s.

Please note that, as indicated in Table 1 of the common response, we initially observed a very low base performance of 5% for the Amazon670K dataset with 96% sparsity. However, by employing the auxiliary branch, the performance increased significantly to 38.4%.

Can you give me some idea of the importance of XMC as a task?

XMC refers to any (deep) learning setting in which the size of the output layer is extremely large. This could be an instance of text-classification[2] or multi-modal learning [3]. Apart from other applications in health [4], the recent works from Amazon and Microsoft, other tasks that can be reduced to an XMC task are product/item recommendation[5], and prediction of Related Searches[6], and dynamic search advertising [7].

In addition to the current application areas, we think that XMC techniques may become relevant for resource-efficient LLM settings. Recent models have shown an increase in vocabulary size, from 32k in LLama2 to 128k in LLama3, and even 256k in Gemma2. That means that for Gemma2 2B with an embedding dimension of 2304, about 590M --about a quarter-- of the weights are used just for token embeddings.

[1] A. Zhou et al., “Learning N:M Fine-grained Structured Sparse Neural Networks From Scratch,” Apr. 18, 2021, arXiv: arXiv:2102.04010. doi: 10.48550/arXiv.2102.04010.

[2] K. Dahiya, D. Saini, A. Mittal, A. Shaw, K. Dave, A. Soni, H. Jain, S. Agarwal and M. Varma. DeepXML: A deep extreme multi-Label learning framework applied to short text documents. In WSDM 2021

[3] A. Mittal, K. Dahiya, S. Malani, J. Ramaswamy, S. Kuruvilla, J. Ajmera, K. Chang, S. Agrawal, P. Kar and M. Varma. Multimodal extreme classification. CVPR 2022.

[4] Mario Almagro, Raquel Martínez Unanue, Víctor Fresno, and Soto Montalvo. ICD-10 coding of Spanish electronic discharge summaries: An extreme classification problem. IEEE Access 8 (2020), 100073–100083.

[5] H. Jain, V. Balasubramanian, B. Chunduri and M. Varma, Slice: Scalable linear extreme classifiers trained on 100 million labels for related searches, in WSDM 2019.

[6] Y. Prabhu, A. Kag, S. Harsola, R. Agrawal and M. Varma, Parabel: Partitioned Label Trees for Extreme Classification with Application to Dynamic Search Advertising in WWW, 2018.

[7] Chang, W. C., Jiang, D., Yu, H. F., Teo, C. H., Zhang, J., Zhong, K., ... & Dhillon, I. S. (2021, August). Extreme multi-label learning for semantic matching in product search. In Proceedings of the 27th ACM SIGKDD conference on knowledge discovery & data mining (pp. 2643-2651).

Limited novelty in core techniques: The main components (DST, semi-structured sparsity, auxiliary objectives) are existing methods, though their combination and application to XMC is novel. More discussion and intuition on how these components interact and complement each other in an overview would be beneficial.

We agree that each of the above-noted components are not novel in their own right; however, integrating these components is a non-trivial contribution. The application of DST to the XMC task is of significant interest since the memory overhead of XMC tasks remains a challenging constraint. The application of DST methods to this setting showcases that the theoretical justifications used to motivate much of the DST literature do in fact result in real-world improvements when memory overhead is a concern. Further, the effect of the auxiliary objective on convergence of higher sparsity DST experiements may, in turn, spur the broader DST community to incorporate similar mechanisms to improve the generalization performance of DST in general. In our opinion, the integration of these various components and our detailed empirical results represent a significant contribution to both the XMC and DST literature.

The paper is primarily empirical and lacks rigorous theoretical justification for why the proposed modifications work. Some insights into the underlying mechanisms that make the proposed approach effective would be valuable.

While our results are primiarly derived from our empirical observations, we believe their publicaiton will lead to further interest in the application of DST methods to the XMC setting. In future work, we hope to further examine our findings and develop a theoretical framework to help describe why this approach appears to be so effective.

How sensitive are the results to the various architectural hyperparameters like fan-in degree, decay schedule for auxiliary loss? Some ablation experiments exploring these choices besides the intermediate layer size would be informative.

We have provided additional hyperparameter ablations in the common response of the rebuttal.