# Learning Optimal Tree Models Under Beam Search

arXiv: https://arxiv.org/abs/2006.15408 · HTML: https://mlsanity.com/paper/2006.15408/

## Key facts

- Cited by 26 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 27 Jun 2020.

## Built on by (most cited first, 26 of 26)

- [Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations](https://mlsanity.com/paper/2402.17152/) (2024)
- [On Missing Labels, Long-tails and Propensities in Extreme Multi-label Classification](https://mlsanity.com/paper/2207.13186/) (2022)
- [Constructing Tree-based Index for Efficient and Effective Dense Retrieval](https://mlsanity.com/paper/2304.11943/) (2023)
- [TGB-Seq Benchmark: Challenging Temporal GNNs with Complex Sequential Dynamics](https://mlsanity.com/paper/2502.02975/) (2025)
- [Approximate Nearest Neighbor Search under Neural Similarity Metric for Large-Scale Recommendation](https://mlsanity.com/paper/2202.10226/) (2022)
- [Revisiting Neural Retrieval on Accelerators](https://mlsanity.com/paper/2306.04039/) (2023)
- [Propensity-scored Probabilistic Label Trees](https://mlsanity.com/paper/2110.10803/) (2021)
- [Lexically-Accelerated Dense Retrieval](https://mlsanity.com/paper/2307.16779/) (2023)
- [SPEED: Streaming Partition and Parallel Acceleration for Temporal Interaction Graph Embedding](https://mlsanity.com/paper/2308.14129/) (2023)
- [Long-Sequence Recommendation Models Need Decoupled Embeddings](https://mlsanity.com/paper/2410.02604/) (2024)
- [Generalized test utilities for long-tail performance in extreme multi-label classification](https://mlsanity.com/paper/2311.05081/) (2023)
- [Large Language Model as Universal Retriever in Industrial-Scale Recommender System](https://mlsanity.com/paper/2502.03041/) (2025)
- [Truncation-Free Matching System for Display Advertising at Alibaba](https://mlsanity.com/paper/2102.09283/) (2021)
- [Set-valued prediction in hierarchical classification with constrained representation complexity](https://mlsanity.com/paper/2203.06676/) (2022)
- [Givens Coordinate Descent Methods for Rotation Matrix Learning in Trainable Embedding Indexes](https://mlsanity.com/paper/2203.05082/) (2022)
- [Practice with Graph-based ANN Algorithms on Sparse Data: Chi-square Two-tower model, HNSW, Sign Cauchy Projections](https://mlsanity.com/paper/2306.07607/) (2023)
- [Context-aware Tree-based Deep Model for Recommender Systems](https://mlsanity.com/paper/2109.10602/) (2021)
- [Tree-based Text-Vision BERT for Video Search in Baidu Video Advertising](https://mlsanity.com/paper/2209.08759/) (2022)
- [CoopHash: Cooperative Learning of Multipurpose Descriptor and Contrastive Pair Generator via Variational MCMC Teaching for Supervised Image Hashing](https://mlsanity.com/paper/2210.04288/) (2022)
- [MISS: Multi-Modal Tree Indexing and Searching with Lifelong Sequential Behavior for Retrieval Recommendation](https://mlsanity.com/paper/2508.14515/) (2025)
- [BEAR: Towards Beam-Search-Aware Optimization for Recommendation with Large Language Models](https://mlsanity.com/paper/2601.22925/) (2026)
- [WSLRec: Weakly Supervised Learning for Neural Sequential Recommendation Models](https://mlsanity.com/paper/2202.13616/) (2022)
- [Deep Uncertainty-Based Explore for Index Construction and Retrieval in Recommendation System](https://mlsanity.com/paper/2408.00799/) (2024)
- [Bidding-Aware Retrieval for Multi-Stage Consistency in Online Advertising](https://mlsanity.com/paper/2508.05206/) (2025)
- [Trie-Aware Transformers for Generative Recommendation](https://mlsanity.com/paper/2602.21677/) (2026)
- [OneShot: Index-in-Ranking with Neural Scoring for Large-Scale Retrieval](https://mlsanity.com/paper/2607.27475/) (2026)

## Abstract (arXiv)

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.