[论文解读] Learning Optimal Tree Models Under Beam Search

Jingwei Zhuo, Ziru Xu+5 more|arXiv (Cornell University)|Jun 27, 2020

Advanced Image and Video Retrieval Techniques被引用 7

一句话总结

本文提出了一种新颖的训练算法——最优树模型（Optimal Tree Model, OTM），通过在训练过程中显式考虑束搜索（beam search），有效消除了基于树的模型在大规模检索任务中训练与测试之间的差异。OTM 引入了束搜索感知的子采样策略和最优伪目标定义，实现了最先进性能，在 Amazon Books 上实现了 29.8% 的相对召回率提升，在 UserBehavior 上实现了 6.3% 的提升，优于先前方法。

ABSTRACT

Retrieving relevant targets from an extremely large target set under computational limits is a common challenge for information retrieval and recommendation systems. Tree models, which formulate targets as leaves of a tree with trainable node-wise scorers, have attracted a lot of interests in tackling this challenge due to their logarithmic computational complexity in both training and testing. Tree-based deep models (TDMs) and probabilistic label trees (PLTs) are two representative kinds of them. Though achieving many practical successes, existing tree models suffer from the training-testing discrepancy, where the retrieval performance deterioration caused by beam search in testing is not considered in training. This leads to an intrinsic gap between the most relevant targets and those retrieved by beam search with even the optimally trained node-wise scorers. We take a first step towards understanding and analyzing this problem theoretically, and develop the concept of Bayes optimality under beam search and calibration under beam search as general analyzing tools for this purpose. Moreover, to eliminate the discrepancy, we propose a novel algorithm for learning optimal tree models under beam search. Experiments on both synthetic and real data verify the rationality of our theoretical analysis and demonstrate the superiority of our algorithm compared to state-of-the-art methods.

研究动机与目标

• 解决基于树的模型在推理阶段使用束搜索时，尽管训练阶段表现最优，但测试性能仍下降的训练-测试差异问题。
• 通过引入贝叶斯最优性与束搜索下的校准性，理论分析该问题，提出新的最优性准则。
• 设计一种通用的训练算法，解决该差异问题，并在多种数据集上提升检索性能。
• 在合成数据集与真实世界数据集（包括极端多标签分类任务）上，验证所提方法的优越性。

提出的方法

• 提出束搜索下的贝叶斯最优性概念，形式化定义在使用束搜索时的最优检索性能。
• 提出“束搜索下的校准性”作为训练目标，使模型预测与束搜索结果对齐。
• 设计一种束搜索感知的子采样策略，基于节点在束搜索中被扩展的可能性来选择训练节点。
• 通过束搜索路径重新定义伪目标，将其转化为最优检索目标，确保训练与推理的一致性。
• 使用一种损失函数联合优化节点得分器，目标是提升束搜索性能，而非仅概率估计。
• 实现一种端到端的训练算法（OTM），在束搜索约束下联合优化树结构与节点得分器。

实验结果

研究问题

• RQ1 我们能否在推理阶段使用束搜索时，形式化定义树模型的最优性？
• RQ2 现有树模型在训练与测试之间存在性能差距的根本原因是什么？该差距能否被理论刻画？
• RQ3 是否可能设计一种训练过程，完全消除基于束搜索检索的训练-测试差异？
• RQ4 束搜索感知的子采样与最优伪目标分别在提升检索性能中起到何种作用？
• RQ5 所提方法是否能在不同树结构与节点得分器架构上实现良好泛化？

主要发现

• 与最先进方法 JTM 相比，OTM 在 Amazon Books 上实现了 29.8% 的相对召回率提升，在 UserBehavior 上实现了 6.3% 的提升。
• 束搜索感知的子采样策略对性能提升的贡献大于最优伪目标定义。
• 树结构中层级间的数据不平衡程度影响性能提升幅度，关键节点集中度越高，性能增益越大。
• OTM 在 Tesla P100 GPU 上每批次训练时间为 0.671 秒，虽为 PLT 的 3.6 倍，但仍在同一数量级，适用于分布式训练。
• 理论分析表明，TDM 与 PLT 均非束搜索下的最优模型，且推导出最优树模型的充分条件。
• 实验验证了移除约束 |I_x|=1（即允许每个实例有多个相关目标）可提升性能，HSM 表现优于 PLT。