Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Hierarchical Embedding Fusion for Retrieval-Augmented Code Generation

Nikita Sorokin, Ivan Sedykh, Valentin Malykh

TL;DR

Results indicate that hierarchical dense caching is an effective mechanism for low-latency, repository-aware code completion, and introduces a utility-weighted likelihood signal for filtering training contexts and reports ablation studies on pseudo-token budget, embedding models, and robustness to harmful retrieval.

Abstract

Retrieval-augmented code generation often conditions the decoder on large retrieved code snippets. This ties online inference cost to repository size and introduces noise from long contexts. We present Hierarchical Embedding Fusion (HEF), a two-stage approach to repository representation for code completion. First, an offline cache compresses repository chunks into a reusable hierarchy of dense vectors using a small fuser model. Second, an online interface maps a small number of retrieved vectors into learned pseudo-tokens that are consumed by the code generator. This replaces thousands of retrieved tokens with a fixed pseudo-token budget while preserving access to repository-level information. On RepoBench and RepoEval, HEF with a 1.8B-parameter pipeline achieves exact-match accuracy comparable to snippet-based retrieval baselines, while operating at sub-second median latency on a single A100 GPU. Compared to graph-based and iterative retrieval systems in our experimental setup, HEF reduces median end-to-end latency by 13 to 26 times. We also introduce a utility-weighted likelihood signal for filtering training contexts and report ablation studies on pseudo-token budget, embedding models, and robustness to harmful retrieval. Overall, these results indicate that hierarchical dense caching is an effective mechanism for low-latency, repository-aware code completion.

Hierarchical Embedding Fusion for Retrieval-Augmented Code Generation

TL;DR

Results indicate that hierarchical dense caching is an effective mechanism for low-latency, repository-aware code completion, and introduces a utility-weighted likelihood signal for filtering training contexts and reports ablation studies on pseudo-token budget, embedding models, and robustness to harmful retrieval.

Abstract

Retrieval-augmented code generation often conditions the decoder on large retrieved code snippets. This ties online inference cost to repository size and introduces noise from long contexts. We present Hierarchical Embedding Fusion (HEF), a two-stage approach to repository representation for code completion. First, an offline cache compresses repository chunks into a reusable hierarchy of dense vectors using a small fuser model. Second, an online interface maps a small number of retrieved vectors into learned pseudo-tokens that are consumed by the code generator. This replaces thousands of retrieved tokens with a fixed pseudo-token budget while preserving access to repository-level information. On RepoBench and RepoEval, HEF with a 1.8B-parameter pipeline achieves exact-match accuracy comparable to snippet-based retrieval baselines, while operating at sub-second median latency on a single A100 GPU. Compared to graph-based and iterative retrieval systems in our experimental setup, HEF reduces median end-to-end latency by 13 to 26 times. We also introduce a utility-weighted likelihood signal for filtering training contexts and report ablation studies on pseudo-token budget, embedding models, and robustness to harmful retrieval. Overall, these results indicate that hierarchical dense caching is an effective mechanism for low-latency, repository-aware code completion.
Paper Structure (56 sections, 4 equations, 3 figures, 7 tables, 2 algorithms)

This paper contains 56 sections, 4 equations, 3 figures, 7 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Repository-Level Code Completion
  4. Snippet-Injection RAG for Code
  5. Structure-Aware Retrieval
  6. Iterative and Adaptive Retrieval
  7. Pseudo-Token Conditioning and Soft Prompts
  8. Dense Retrieval and Representation Learning
  9. Long-Context Modeling
  10. Fusion-in-Decoder Approaches
  11. Method
  12. Task Definition and Notation
  13. Offline Stage: Hierarchical Cache Construction
  14. Chunking.
  15. Embedding.
  16. ...and 41 more sections

Figures (3)

  • Figure 1: HEF pipeline overview. (A) GitHub repositories are crawled and filtered. (B) Repositories are deduplicated and stored. (C) Files are split into $\le$512-token chunks. (D) UWL scoring filters training contexts by utility. (E) Frozen embedder maps chunks to dense vectors. (F) Fuser recursively merges vectors into file/module/repo hierarchy (cached). (G) At inference, retrieved vectors become pseudo-tokens for the generator.
  • Figure 2: Example of a fused repository tree: chunks merge into file embeddings, which merge into module and repository embeddings.
  • Figure 3: Effect of the number of fused embedding tokens on CodeBLEU. Labels show exact scores at each point.