Navigating Extremes: Dynamic Sparsity in Large Output Spaces

Summary

**[Edited: My overall score has increased from 6 to 7 (Accept) after the authors' rebuttal.]** Dynamic Sparse Training (DST) holds the promise of more efficient training and model robustness, but remains largely impractical due to the lack of inherent structure and severe amplification of training steps needed to converge to high quality. The authors apply DST to extreme multi-label classification (XMC) problems and build on past work to avoid losing important gradient signals and achieve a practical speedup. They use a structured fixed fan-in sparsity at the final, extremely large, classifier layer to provide benefit in practice, and an auxiliary classifier to enhance gradient signals to layer before the semi-structured sparse classifier. This classifier is designed to provide useful signals early in training, but not become so important that it cannot be removed after that. The proposed modifications to DST in XMC are evaluated across a variety of data sets and against a seemingly comprehensive set of baseline techniques, and the results suggest that the method reduces memory consumption compared to the most baselines and gives a model with comparable quality to the best-performing baseline. The authors also show that the auxiliary objective not only helps with initial training, but also with high sparsity in the final classifier.

Strengths

**Originality** There seem to be two novel aspects: applying a past semi-structured fixed fan-in sparsity to the final classifier, and moving the low-rank intermediate layer to become an auxiliary classifier, solving a standard label shortlisting task, that can be removed by gradually decaying its importance over time. The authors also provide new insight into the importance of the commonly-used meta-classifier and demonstrate their approach is effectively robust to label imbalances. **Quality** I believe the claims have been supported by experimental results in all cases. The ablation study in Figure 3 makes the benefit of the auxiliary classifier clear. The authors also provided a sweep of rewiring intervals to support their experimental settings, and the baselines seem fair to me; kudos to the authors for thinking of including a "smaller, dense" version of the fixed fan-in sparse classifier with the same number of parameters. **Clarity** In general, I found the paper easy to read with a fine layout, and I appreciate the effort the authors spent in explaining the XMC problem for readers, like me, who are not familiar with its unique aspects. **Significance** Though I'm not familiar with XMC, if it is as widespread as some other tasks, successfully training a sparse model from scratch will be a huge advancement.

Weaknesses

**Originality** It really seems like [32] used the same type of semi-structured sparsity on their classifier. What is the difference? Section 2.2 makes it appear as though this were your contribution. **Quality** 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. **Clarity** Though the authors provided context for XMC, it still took me a couple reads to really understand the task, but once I did, the difficulty of the problem the authors set out to solve was clear. I'm confused about the authors' intent to claim novelty for applying semi-structured fixed fan-in sparsity to the final classifier. **Significance** This is my first time hearing about the XMC problem, but I've heard of many other tasks and problems that I do not work on directly. This makes me think that perhaps the applicability of the work may be limited, but I'll be happy to be wrong. 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. It's hard to judge the benefit of one of your primary contributions, the python bindings for CUDA kernels, without seeing them.

Questions

- What are the concrete novel aspects of your work? - Can you give me some idea of the importance of XMC as a task?

Rating

7

Confidence

2

Soundness

3

Presentation

2

Contribution

2

Limitations

Yes, the authors have discussed the limitations and potential societal impact of their work.

Summary

This paper investigates the application of Dynamic Sparse Training (DST) methods to the domain of extreme multi-label classification (XMC), where the label space can be very large, on the order of millions. The authors propose several enhancements to standard DST approaches to address the challenges posed by the highly skewed label distributions and scarcity of training data typical in XMC datasets. These include using semi-structured sparsity with fixed fan-in connections, adding an intermediate layer between the encoder and sparse classifier, and incorporating an auxiliary training objective to improve gradient flow, especially in the early phases of training. Empirical results on various large-scale datasets demonstrate that the proposed approach can significantly reduce GPU memory requirements while maintaining competitive performance compared to dense models and specialized XMC methods.

Strengths

- The paper is well-organized and clearly written, providing a comprehensive overview of the challenges in XMC and the proposed modifications to DST to address them. The technical contributions are well-motivated and supported by empirical results. - The paper tackles an important problem of scaling DST to the domain of XMC, characterized by enormous label spaces, label imbalance, and label noise. It demonstrates DST's applicability beyond the typical benchmark datasets used in sparsity research. - The authors present a well-motivated set of modifications to standard DST algorithms to handle XMC-specific issues. The semi-structured sparsity, intermediate layer, and auxiliary loss are supported by empirical analysis showing their impact on performance. - The results show substantial memory savings on large-scale datasets with minimal loss in predictive performance. This enables end-to-end training of XMC models on commodity hardware.

Weaknesses

- 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. - 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.

Questions

- 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.

Rating

7

Confidence

2

Soundness

3

Presentation

4

Contribution

3

Limitations

Yes.

Summary

The paper proposes the application of DST(Dynamic Sparse Training) to XMC(Extreme Multi-label classification) by employing an intermediate layer and adding an auxiliary training objective to enable end-to-end training of XMC problems with millions of labels on commodity hardware.

Strengths

1. The paper is well-written and easy-to-follow. 2. The paper solves an important problem of democratizing access to the state-of-the-art XMC model training. 3. The paper provides a thorough analysis of challenges and solutions in applying DST to XMC. 4. The conclusion is supported by experimental results on 4+1 datasets.

Weaknesses

1. The paper addresses the memory efficiency, however the training time still remains a concern, specifically for time-sensitive real world application datasets. 2. The paper does not present results on datasets with label features except one(LFAT-131K). Many modular techniques such as NGAME, DEXA etc. perform much better on these datasets and are memory efficient by using negative sampling and optimizers such as SparseAdam. 3. The introduction of new hyperparameters, such as sparsity levels, intermediate layer sizes, and auxiliary objectives, adds to the complexity of the model. The sensitivity of the model's performance to these hyperparameters is not thoroughly explored.

Questions

1. It would be good to add training time estimates in addition to memory usage. 2. In the introduction, the authors refer to [32] by Schultheis and Babbar which assumes the existence of pre-trained embeddings and claim that DST can give substantial memory savings. 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? 3. It would be good to add a couple more datasets with LF features. 4. It would be good to compare to Renee with smaller sized encoders. 5. 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. 6. It would be good to add sensitivity of the model's performance to the new hyperparameters.

Rating

6

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

1. The authors have discussed some limitations. 2. 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.

I appreciate these responses

Thank you, authors, for your thoughtful and clear responses to my concerns. Assuming you'll add the clarifications provided here, and the results provided to other reviewers, into a revised version, I've increased my rating to Accept.

Thank you for the responses

I thank the authors for detailed responses to my questions. The rebuttal response has addressed most of my concerns, specifically related to the hyper-parameter sensitivity, comparison with Renee, DEXML. After carefully reading the responses, I would like to keep my score.