Hierarchical Retrieval: The Geometry and a Pretrain-Finetune Recipe (2509.16411v1)

Abstract: Dual encoder (DE) models, where a pair of matching query and document are embedded into similar vector representations, are widely used in information retrieval due to their simplicity and scalability. However, the Euclidean geometry of the embedding space limits the expressive power of DEs, which may compromise their quality. This paper investigates such limitations in the context of hierarchical retrieval (HR), where the document set has a hierarchical structure and the matching documents for a query are all of its ancestors. We first prove that DEs are feasible for HR as long as the embedding dimension is linear in the depth of the hierarchy and logarithmic in the number of documents. Then we study the problem of learning such embeddings in a standard retrieval setup where DEs are trained on samples of matching query and document pairs. Our experiments reveal a lost-in-the-long-distance phenomenon, where retrieval accuracy degrades for documents further away in the hierarchy. To address this, we introduce a pretrain-finetune recipe that significantly improves long-distance retrieval without sacrificing performance on closer documents. We experiment on a realistic hierarchy from WordNet for retrieving documents at various levels of abstraction, and show that pretrain-finetune boosts the recall on long-distance pairs from 19% to 76%. Finally, we demonstrate that our method improves retrieval of relevant products on a shopping queries dataset.

• The paper introduces a theoretical foundation for using Dual Encoders to model hierarchical document relations with logarithmic embedding dimensions.
• It presents a pretrain-finetune recipe that improves recall for long-distance (remote ancestor) retrieval in asymmetric hierarchical settings.
• Empirical studies on synthetic and real-world datasets, including WordNet, validate the method’s effectiveness and scalability.

Hierarchical Retrieval with Dual Encoders: Geometric Foundations and a Pretrain-Finetune Paradigm

Problem Formulation and Motivation

The paper addresses the problem of Hierarchical Retrieval (HR), where the document set is structured as a hidden hierarchy, typically modeled as a directed acyclic graph (DAG). The retrieval objective is, given a query, to return all its ancestors in the hierarchy. This setting is motivated by practical applications such as phrase match in online advertising, where more general keywords (ancestors) are relevant to a specific query.

Dual Encoder (DE) architectures, which independently embed queries and documents into a shared vector space, are widely used for scalable retrieval. However, the standard Euclidean geometry of DEs imposes limitations, particularly in modeling the asymmetric and transitive relationships inherent in hierarchies. The paper investigates both the theoretical feasibility and practical learnability of DEs for HR, and introduces a pretrain-finetune recipe to address empirical deficiencies in long-distance retrieval.

Figure 1: (Left) Example of a document hierarchy and the HR task; (Middle) DAG abstraction for HR; (Right) Illustration of the lost-in-the-long-distance phenomenon in DEs trained for HR.

Theoretical Analysis: Embedding Geometry for Hierarchical Retrieval

The first major contribution is a geometric analysis of the feasibility of DEs for HR. The authors provide a constructive algorithm that, given knowledge of the hierarchy, produces asymmetric query and document embeddings in Euclidean space that provably solve the HR task. The key result is that, for a hierarchy with $m$ documents and maximum ancestor set size $s$, there exist embeddings in dimension $d = O(\max\{s \log m, \epsilon^{-2} \log m\})$ that separate relevant (ancestor) and irrelevant documents by an arbitrary margin $2\epsilon$.

This result is significant in that the required embedding dimension is only logarithmic in the number of documents and linear in the maximum number of ancestors per query, which is typically much smaller than $m$. The construction is efficient, with runtime $\tilde{O}(md + nsd)$, and does not depend on the number of queries $n$.

However, this construction assumes access to the full hierarchy, which is rarely available in practical retrieval scenarios. Instead, DEs are typically trained from query-document pairs sampled from the HR relation, raising the question of whether such embeddings can be learned in practice.

Empirical Study: Learning Dual Encoders for HR

The paper conducts a systematic empirical paper on synthetic tree-structured hierarchies to assess the learnability of DEs for HR via standard training (softmax loss on sampled query-document pairs). The experiments vary the tree height $H$ and width $W$, and measure the minimum embedding dimension $d$ required to achieve high recall.

Figure 2: Required embedding dimension $d$ for successful retrieval as a function of tree height $H$ (for $W=2$).

The results confirm that the empirical scaling of $d$ with $H$ and $W$ matches the theoretical predictions, and that learned DEs can achieve successful HR with much smaller $d$ than the worst-case construction. This validates the practical feasibility of DEs for HR under sufficient embedding capacity.

The Lost-in-the-Long-Distance Phenomenon

A critical empirical finding is the "lost-in-the-long-distance" phenomenon: when the embedding dimension is constrained, DEs trained with standard sampling exhibit high recall for short-distance (close ancestor) pairs but poor recall for long-distance (remote ancestor) pairs. This is attributed to the imbalance in the training data, which is dominated by short-distance pairs due to the hierarchical structure.

Figure 3: Recall at varying query-document distances under regular sampling, rebalanced sampling, and pretrain-finetune training.

Attempts to rebalance the training data by upsampling long-distance pairs improve recall for those pairs but degrade recall for short-distance pairs, indicating a trade-off that cannot be resolved by simple resampling.

Pretrain-Finetune Recipe for Hierarchical Retrieval

To address the above limitation, the authors propose a pretrain-finetune recipe: first pretrain the DE on standard (short-distance dominated) data, then finetune on a dataset focused on long-distance pairs. This approach enables the model to retain high recall for short-distance pairs while substantially improving recall for long-distance pairs, without requiring explicit knowledge of the full hierarchy.

Empirical results on synthetic hierarchies demonstrate that the pretrain-finetune approach achieves near-perfect recall across all distances, outperforming both standard and rebalanced training.

Experiments on Real-World Hierarchies

WordNet

The method is evaluated on WordNet, a large lexical database with a hypernym DAG structure. The HR task is to retrieve all hypernyms (ancestors) of a given synset. The pretrain-finetune recipe yields substantial improvements in recall for long-distance pairs, with overall recall increasing from 43.0% (standard training, $d=16$) to 60.1% (pretrain-finetune). For $d=64$, the improvement is from 71.4% to 92.3%. Notably, the method outperforms both standard Euclidean and hyperbolic embeddings of the same dimension.

Figure 4: Retrieval quality at varying query-document distances on WordNet using hyperbolic embeddings (dimension 16).

ESCI Shopping Dataset

On the ESCI shopping queries dataset, the pretrain-finetune recipe is applied to retrieve both Exact and Substitute product matches. Pretraining on Exact matches and finetuning on Substitute matches enables a single DE to achieve higher recall on both categories compared to joint training.

Figure 5: Recall on Exact matches in the ESCI dataset using different training strategies.

Ablation and Comparative Analysis

The paper provides ablation studies on learning rate and temperature during finetuning, showing that the pretrain-finetune recipe is robust to these hyperparameters. Comparative experiments with hyperbolic embeddings confirm that, while hyperbolic geometry is beneficial for hierarchies, the pretrain-finetune approach in Euclidean space achieves superior recall and is more compatible with large-scale retrieval systems due to efficient nearest neighbor search.

Practical Implications and Future Directions

The findings have several practical implications:

• Scalability: The geometric analysis shows that DEs can scale to large hierarchies with modest embedding dimensions.
• Training Protocols: The pretrain-finetune recipe is simple to implement and does not require explicit knowledge of the hierarchy, only a proxy for long-distance pairs.
• System Integration: The approach is compatible with existing DE-based retrieval systems and can be applied to a wide range of hierarchical retrieval tasks, including product search and knowledge base traversal.

Theoretically, the work clarifies the conditions under which DEs can represent hierarchical relations and highlights the importance of training data distribution in learning such representations. The empirical results suggest that further improvements may be possible by integrating geometric priors (e.g., hyperbolic or box embeddings) with the pretrain-finetune paradigm, or by developing adaptive sampling strategies that dynamically balance short- and long-distance pairs.

Conclusion

This paper provides a rigorous geometric foundation for dual encoder models in hierarchical retrieval, establishes their practical learnability, and introduces a pretrain-finetune training paradigm that resolves the lost-in-the-long-distance problem. The approach is validated on both synthetic and real-world hierarchies, demonstrating strong improvements in recall for long-distance retrieval without sacrificing performance on short-distance pairs. These results have direct implications for the design of scalable, hierarchy-aware retrieval systems and open avenues for further research on geometric and training-based enhancements for structured retrieval tasks.