Hierarchical Retrieval: The Geometry and a Pretrain-Finetune Recipe

We thank the reviewer for the thoughtful review and for finding our paper "well organized and informative”, the feasibility proof valuable, and the empirical results compelling.

How to Apply the Method without DAG Distances

This is an excellent point, and we apologize for the lack of clarity. The key point, which we will clarify in the revision, is that our recipe does not require the full DAG or precise path lengths. It only requires a proxy for distance that allows us to partition the training data into two sets: a "short-distance" (or "easy") set for pretraining, and a "long-distance" (or "hard") set for finetuning. This proxy can be readily available in practice or easy to obtain: On the ESCI shopping dataset, the dataset's labels directly provide this split: we treat "Exact" matches as short-distance and "Substitute" matches as long-distance. In some applications, one may be able to use a simple heuristic (e.g., direct parent-child links vs. all other ancestors) to create the finetuning set. Human annotation of the training dataset into short vs long distance pairs is another option and such a binary annotation is significantly easier to obtain than constructing the full DAG.

Intuition for Why Pretrain-Finetune Works Better

As our experiments show, mixed/rebalanced training forces the model to learn two different structures simultaneously—the local structure of exact matches and the global structure of distant ancestors. This creates a difficult optimization trade-off, where improving on one often comes at the expense of the other (as seen in Figure 3b and Table 1). Our two-stage recipe decouples these objectives. The pretraining stage grounds the model by learning robust embeddings for the local, short-distance structure. The finetuning stage then acts as a gentle adjustment, adapting these well-formed embeddings to the more challenging long-distance pairs without catastrophically forgetting the initial structure.

On Prerequisites and Limitations

This is a valid concern and an important limitation. Our recipe implicitly assumes that there is sufficient "short-distance" data to learn a meaningful initial representation during pretraining. If this data is extremely sparse, pretraining may be ineffective, and a mixed training approach might indeed perform better. We thank the reviewer for this insight. We will explicitly add this to our section 5.2, stating that the effectiveness of our recipe depends on having a reasonable amount of data for the initial pretraining stage.

We are very grateful to the reviewer for the thorough review and the positive assessment of our paper's quality, clarity, and originality. We appreciate that you found our theoretical analysis "rigorous" and our contribution of identifying a failure mode and offering a simple fix to be significant.

Generalization to Different Hierarchies

While our synthetic experiments use perfect trees for controlled analysis, our main real-world experiment is on WordNet, which is a large, complex, and irregular directed acyclic graph (DAG), not a simple tree. The fact that our method yields significant improvements on WordNet (e.g., boosting recall on long-distance pairs from 19% to 76%) demonstrates its effectiveness on real-world, non-uniform hierarchies.

Computational Overhead

At inference time, our method has zero overhead. The final model is a standard dual encoder, and its retrieval latency is identical to that of a model trained via a single-stage process. The overhead exists only during training, as our method involves two stages. On WordNet, pretraining takes 50k training steps while finetuning only takes 700. On ESCI, pretraining and finetuning takes 30k and 50k training steps, respectively.

Statistical Significance

We acknowledge in our paper checklist that we did not report error bars or formal statistical significance tests. We agree this would strengthen the paper. However, we would like to point out that the performance improvements shown are very large and consistent across multiple datasets and embedding dimensions. For example, in Table 1 (d=64), our method improves the minimum recall from 19.4% to 75.7% and the overall recall from 71.4% to 92.3%. We believe these substantial gaps provide strong evidence for the effectiveness of our approach.

Novelty of Pretrain-Finetune

While pretrain-finetune paradigm is common in ML, it is used primarily in the context of pretraining a large model followed by adapting such a model to a particular downstream task. It has never been used to address the lost-in-the-long-distance issue. We see our primary contributions not as the invention of this paradigm, but as: (1) discovering and articulating the "lost-in-the-long-distance" phenomenon as a key failure mode for DEs in hierarchical retrieval, and (2) identifying that this recipe offers a simple and effective solution, especially when contrasted with the more intuitive but much less effective rebalancing strategy.

We thank the reviewer for the positive feedback and for finding our problem variant "interesting" and our experiments well-designed.

Extending to Complex, Real-World Scenarios

This is an excellent question that gets to the heart of practical applicability. Our approach is highly adaptable precisely because the dual encoder (DE) architecture is modular. The core of our contribution is the pretrain-finetune training recipe, which is agnostic to the specifics of the encoders. In our work, we use lookup tables for the WordNet experiments and modern Transformer-based text encoders for the ESCI shopping dataset. For a multi-modal product catalog, these encoders could be readily replaced with architectures that process both text and images (e.g., CLIP-style models) without changing the fundamental training strategy: pretrain on "exact" or "primary" matches, then finetune on "substitute," "broader," or other "long-distance" matches. The ESCI experiment is a direct step toward showing this real-world applicability.

Alternate Training Formulations

We thank the reviewer for this insightful question about alternate training formulations. It is indeed possible, as the reviewer suggests, to frame our Hierarchical Retrieval (HR) task as an Extreme Multi-Label Classification (XMLC) problem by "flattening" the document hierarchy and treating each document as an independent label. However, this approach could be suboptimal as many common XMLC methods are not well-suited for the practical constraints of large-scale retrieval systems, for reasons beyond just the loss of the hierarchical structure.

• Modern retrieval systems require millisecond-level latency. Alternate XMLC methods, such as One-vs-All (OVA) makes a prediction by running the query through a unique classifier for every document in the corpus. The linear cost of this operation is too slow for any real-time application.

• Moreover, real-world document collections are constantly growing. With our embedding-based approach, adding a new document is trivial: we generate its embedding and add it to the ANN index without retraining the model. In contrast, adding a new document (i.e., a new "label") to a tree-based or OVA-based XMLC model is far more complex and would likely require significant retraining or model restructuring.

Therefore, while one could apply XMLC methods by flattening the hierarchy, this would not only discard the crucial hierarchical information that our method is designed to capture but would also likely lead to a system that is too slow and inflexible for a practical retrieval setting. We chose our approach because it both models the hierarchical relationships and aligns with the high-performance architecture of modern retrieval systems.

We sincerely thank the reviewer for the detailed and insightful feedback. We appreciate the accurate summary of our work and the challenging questions, which will help us improve the paper significantly. We address the noted weaknesses and questions below.

Comparison with Hyperbolic Models

Hyperbolic models are not our direct competitors. While they are indeed well-suited for hierarchical data, as we note in Section 1.2, they are "NOT suited for retrieval due to a lack of an efficient nearest neighbor search that hinders their applicability to large scale data". Our goal was to investigate the capabilities and limitations of the widely-used Euclidean dual encoder (DE) architecture and provide a practical method to improve it without sacrificing its key advantage: scalability via efficient ANN search. We will revise Section 1.2 to make this motivation clearer and to better contextualize our work against these important baselines.

That said, we are happy to include additional experimental results with Hyperbolic models to situate our method into a broader context for interested readers.

• On HyperLex (see Table 2), Poincare embedding obtained a Spearman’s rho of 0.512 according to ref. [24], which is better than our method (with rho = 0.415).
• On WordNet (see Table 1), we conducted an additional experiment with Poincare embedding using the reparameterization approach in ref. [12] for embedding dimension = 16. We obtained an overall recall of 45.4% which is better than standard training of DEs (43.0%), but is worse than our pretrain-finetune approach (60.1%). To obtain more insights into the behavior of Poincare embeddings, we plotted the recall at varying distances between query and document, and observed that the model struggles to obtain a good balance between pairs with short vs long distances. That is, the model first learns to retrieve pairs with short distances, i.e. with distance 0 and 2. As it starts to retrieve longer distance pairs, the quality of the short distance pairs drops. We will include these results in the revised version.

Novelty of the Pretrain-Finetune Recipe

While pretrain-finetune paradigm is common in ML, it is used primarily in the context of pretraining a large model followed by adapting such a model to a particular downstream task. It has never been used to address the lost-in-the-long-distance issue. We see our primary contributions not as the invention of this paradigm, but as: (1) discovering and articulating the "lost-in-the-long-distance" phenomenon as a key failure mode for DEs in hierarchical retrieval , and (2) identifying that this recipe offers a simple and effective solution, especially when contrasted with the more intuitive but much less effective rebalancing strategy.

Intuition on Pretrain-Finetune Recipe vs Rebalancing

Our experiments show that rebalancing improves long-distance recall at the severe cost of short-distance performance (Figure 3b , Table 1). We hypothesize this is because the model is forced to solve a difficult trade-off between two distinct objectives simultaneously. Our two-stage recipe decouples these objectives. Pretraining first grounds the embeddings in the local, short-distance structure. The finetuning stage then gently adjusts these embeddings to accommodate the global hierarchy, improving long-distance recall while preserving the already-learned local structure (Figure 3c).

Practicality of Theorem 3.1

The theorem's purpose is to provide a formal existence proof that answers our first key question, "Q1: Does there exist Dual Encoders that solve Hierarchical Retrieval?". By proving feasibility, it motivates the subsequent empirical investigation into whether such embeddings can be learned from data, which is the practical focus of Sections 4-6.

Retrieval time and latency

At inference time, the model is a standard DE. Therefore, there is no change in retrieval latency compared to a conventionally trained model, which is a key practical advantage. The computational difference is in the training, which is a two-stage process. We will clarify this in the paper.

Ablations

In response to this great comment, we conducted sensitivity analysis on the effect of learning rate and temperature during finetuning for our HR on WordNet experiments.

• For varying learning rates, we observe that the averaged recall across all distances are within 1% difference for a wide range of learning rate choices in the range of 1e-4 to 3e-3 (main results in the paper are reported for 1e-3).
• For varying temperatures, we observe that the averaged recall across all distances are within 1% difference for a wide range of temperature choices in the range of 500 to 2000 (main results in the paper are reported for 500).

We will include this discussion and the corresponding plots in the revised version.