大规模输出空间预测开源框架解析

大规模输出空间预测开源框架解析

在互联网时代，许多计算任务涉及在巨大的候选空间中寻找少量解决方案。例如问答系统可以从网络任何地方提取答案，而用于文章主题分类的维基百科分类法包含50万个术语。当然，某商店的产品查询有数百万个潜在匹配项。

这类极端多标签排序（XMR）问题面临两大挑战：规模挑战和稀疏性挑战。这些大型搜索空间中的项目往往具有长尾分布：大多数句子很少作为问题答案；维基百科分类法中的大多数主题很少适用于文本；大多数产品很少被购买等等。这意味着尝试使用机器学习解决XMR问题时往往缺乏足够数据。

某机构开发了名为PECOS的通用框架（用于巨大且相关输出空间的预测），成功应用于产品搜索和推荐的关键项目后，现已公开代码以促进这一重要领域的进一步研究。

三阶段架构

PECOS将XMR问题分解为三个阶段：

1. 语义标签索引：根据语义内容将标签分组
2. 匹配：将输入实例与标签组关联
3. 排序：在每个组中找到最符合输入的标签

该框架通过标签分组而非单个标签匹配，显著减少了匹配步骤的搜索空间，同时帮助解决长尾问题，使排序模型能够利用常见标签与不常见标签之间的语义相似性。

核心技术实现

语义标签索引

包含表示方案和分组算法两个组件。对于基于文本的输入，表示方案可利用Word2Vec或ELMo等预训练文本嵌入；对于基于图的输入，可使用图中相邻节点关系信息。PECOS包含高效实现方案，如正实例索引（PII）、正实例特征聚合（PIFA）和图谱表示。

分组算法采用k-means和球形k-means聚类算法，具有递归B叉分区特性。对于某个B值（通常介于2-16之间），算法首先将标签集划分为B个簇，然后将每个簇再划分为B个簇，依此类推。

内置模型

包含X-Transformer（利用Huggingface的预训练transformer模型提升极端多标签文本分类性能）和XR-Linear线性模型。

XR-Linear递归学习匹配算法：首先学习标签空间的B叉分区，然后为每个现有组学习新的B叉分区来实现匹配器，直到达到所需递归深度。此时为每个分区中的标签学习简单的线性一对多排序器。

训练时每个递归层的完整权重集可同时放入内存，节省了低效存储检索时间。推理时使用波束搜索限制搜索空间，如波束宽度为2时，在递归树的每一层只追踪两个最高权重的连接。

性能表现

在包含280万个标签的数据集测试中，深度学习模型相比递归线性模型将top1准确率提升了10%，但训练时间增加了265倍。用户需根据具体用例评估这种权衡。

该框架有助于零样本学习、极端上下文bandit问题和深度强化学习等领域的进一步研究。代码已在GitHub开源提供下载。