[论文解读] LightXML: Transformer with Dynamic Negative Sampling for High-Performance Extreme Multi-label Text Classification

Ting Jiang, Deqing Wang+4 more|arXiv (Cornell University)|Jan 9, 2021

Text and Document Classification Technologies参考文献 23被引用 14

一句话总结

LightXML 提出了一种基于 Transformer 的端到端轻量化模型，用于极端多标签文本分类（XMC），通过生成式协作网络实现动态负样本采样。通过在一个统一框架中联合训练文本表征、标签召回和标签排序，LightXML 在五个基准数据集上实现了最先进性能，模型大小最大减少 72%，计算成本显著低于先前方法（如 AttentionXML 和 X-Transformer）。

ABSTRACT

Extreme Multi-label text Classification (XMC) is a task of finding the most relevant labels from a large label set. Nowadays deep learning-based methods have shown significant success in XMC. However, the existing methods (e.g., AttentionXML and X-Transformer etc) still suffer from 1) combining several models to train and predict for one dataset, and 2) sampling negative labels statically during the process of training label ranking model, which reduces both the efficiency and accuracy of the model. To address the above problems, we proposed LightXML, which adopts end-to-end training and dynamic negative labels sampling. In LightXML, we use generative cooperative networks to recall and rank labels, in which label recalling part generates negative and positive labels, and label ranking part distinguishes positive labels from these labels. Through these networks, negative labels are sampled dynamically during label ranking part training by feeding with the same text representation. Extensive experiments show that LightXML outperforms state-of-the-art methods in five extreme multi-label datasets with much smaller model size and lower computational complexity. In particular, on the Amazon dataset with 670K labels, LightXML can reduce the model size up to 72% compared to AttentionXML.

研究动机与目标

• 解决现有 XMC 方法依赖多阶段训练和静态负样本采样所导致的低效与高计算成本问题。
• 在保持或提升极端多标签场景下分类准确率的同时，减小模型大小和计算复杂度。
• 通过将标签召回与排序整合到单一联合优化框架中，实现基于 Transformer 的 XMC 模型端到端训练。
• 设计一种在训练过程中可自适应调整的动态负样本采样策略，提升模型鲁棒性与收敛性。
• 仅使用单张 GPU 在大规模数据集（如 Amazon-670K）上实现高性能，且训练时间极短。

提出的方法

• LightXML 使用多层 Transformer 编码器，从输入文本中生成丰富且分层的文本表征。
• 引入生成式协作网络，包括标签生成器（用于召回正负标签）和判别器（用于将标签分类为正或负）。
• 在训练过程中通过将相同的文本表征输入生成器，实现动态负样本采样，使判别器能从多样化且持续演化的负样本集中学习。
• 模型采用端到端训练，联合优化文本表征、标签召回与标签排序，避免了分阶段训练模型的需要。
• 标签生成器利用标签聚类信息以提升召回质量，而判别器则用于区分真实正样本与召回样本。
• 该框架支持高效推理，在标签数超过 50 万的数据集上，单样本预测时间低于 5 毫秒。

实验结果

研究问题

• RQ1 基于 Transformer 的 XMC 模型端到端训练是否能在准确率与效率上优于多阶段训练方法？
• RQ2 在生成式-判别式协作框架中学习的动态负样本采样，是否相比静态采样能提升模型泛化能力与收敛性？
• RQ3 单一统一模型是否能在减小模型大小与计算成本的同时，实现在极端多标签数据集上的最先进性能？
• RQ4 Transformer 编码器的多层文本表征对 XMC 中的收敛速度与最终准确率有何影响？
• RQ5 与 SOTA 方法（如 AttentionXML 和 X-Transformer）相比，LightXML 在模型大小与推理时间上能减少多少？

主要发现

• 在 Amazon-670K 数据集上，与 AttentionXML 相比，LightXML 将模型大小减少了 72%，在内存使用显著降低的同时达到相同或更优的性能。
• 在 Amazon-670K 上，LightXML 实现了每样本 4.09 毫秒的预测速度，相比 AttentionXML 的 8.59 毫秒提升了 52%。
• 在 Amazon-670K 上，LightXML 将训练时间缩短至 28.75 小时，仅使用单张 GPU，而 AttentionXML 为 26.10 小时。
• 使用多层文本表征相比仅使用最后一层 [CLS] 标记，使最终 P@5 准确率提升超过 1%。
• 尽管静态采样变体（$S$ 和 $BS$）训练时间更长，但动态负样本采样在准确率与训练效率上均优于它们。
• LightXML 在五个基准数据集（包括 Wiki-500K 和 Amazon-670K）上均实现了最先进性能，计算开销极低，且完全兼容单 GPU 训练。