Hierarchical Embedding Fusion for Retrieval-Augmented Code Generation

O artigo apresenta a Hierarchical Embedding Fusion (HEF), uma abordagem de duas etapas que comprime repositórios de código em uma hierarquia de vetores densos e os integra ao gerador via pseudo-tokens, permitindo uma geração de código assistida por recuperação com latência ultrabaixa e precisão comparável a métodos baseados em snippets.

O essencial

Imagine que você é um programador tentando escrever uma nova função para um projeto gigante, como o sistema de um banco ou um jogo complexo. O problema é que esse projeto tem milhares de arquivos. Para escrever o código certo, você precisa lembrar de como outras partes do sistema funcionam, quais nomes de variáveis foram usados lá fora e como as peças se encaixam.

Se você tentar ler todos os arquivos do projeto de uma vez para se lembrar, seu computador vai travar (é muito pesado e lento). Se você não ler nada, seu assistente de IA vai inventar coisas que não existem (alucinar).

Os métodos antigos tentavam resolver isso de duas formas ruins:

1. Jogar tudo na mesa: Colocar milhares de linhas de código na frente da IA. Funciona, mas é lento e confuso (muito ruído).
2. Ficar procurando em mapas: Criar gráficos complexos de como o código se conecta. Funciona bem, mas demora muito para processar cada vez que você pede ajuda.

Aqui entra a solução proposta neste paper: HEF (Fusão de Embeddings Hierárquicos).

A Analogia do "Resumo de Biblioteca"

Imagine que o repositório de código é uma biblioteca gigante com milhões de livros.

• O Método Antigo (Injeção de Snippets): É como pedir para o bibliotecário trazer todos os livros da estante para sua mesa, na esperança de que você encontre a informação que precisa. Sua mesa fica cheia, você se perde entre as páginas e demora horas para achar o que quer.
• O Método HEF: É como ter um bibliotecário superinteligente que, antes mesmo de você chegar, leu todos os livros e criou um índice hierárquico (um resumo).
  • Ele não guarda o texto inteiro. Ele guarda "ideias" compactas.
  • Ele organiza esses resumos: primeiro por página, depois por capítulo, depois por livro, e finalmente por toda a biblioteca.

Como o HEF funciona (Passo a Passo Simples)

O sistema funciona em duas etapas principais:

1. A Etapa Offline (O "Preparo do Índice")

Antes de você começar a programar, o sistema varre todo o projeto e cria esse "índice mágico".

• Ele pega pedaços de código (como capítulos) e os transforma em vetores densos (imagina que são como "impressões digitais" matemáticas do significado do código, não o texto em si).
• Um pequeno modelo chamado "Fuser" (Fundidor) pega essas impressões digitais e as mistura. Ele pega 8 pedaços de código e os funde em 1 resumo maior. Depois, pega esses resumos e funde em resumos de pastas inteiras, e assim por diante, até ter um resumo de todo o projeto.
• Resultado: Em vez de ter milhões de linhas de texto, você tem uma árvore de resumos compactos que cabe na memória do computador.

2. A Etapa Online (Quando você pede ajuda)

Quando você digita uma linha de código e pede para a IA continuar:

• A IA olha o que você escreveu e pergunta ao "índice": "Qual resumo do projeto se parece mais com isso?"
• O sistema busca os 32 resumos mais relevantes (em vez de buscar 3.000 linhas de texto).
• Esses 32 resumos são transformados em "pseudo-tokens" (palavras mágicas que só a IA entende, mas que carregam o peso de milhares de linhas de código).
• A IA recebe essas "palavras mágicas" e escreve o código.

Por que isso é incrível?

1. Velocidade Relâmpago: Como a IA não precisa ler milhares de linhas de texto, ela responde em menos de 1 segundo (0,68s no teste). É como pedir um resumo rápido em vez de ler um livro inteiro.
2. Precisão: Mesmo sendo rápido, ela sabe exatamente o que está acontecendo no projeto, porque os "resumos" (vetores) guardam o significado das conexões entre os arquivos.
3. Resistência ao Ruído: Se o sistema buscar um resumo errado, ele não "polui" a memória da IA com texto inútil. Como são apenas resumos matemáticos, a IA consegue ignorar o que não serve melhor do que se estivesse lendo texto bruto.

O Resultado Final

Os autores testaram isso em benchmarks reais (como o RepoBench) e descobriram que:

• O sistema HEF é 26 vezes mais rápido que os métodos baseados em gráficos complexos.
• Ele é tão preciso quanto os métodos lentos que jogam todo o código na tela.
• Com um modelo pequeno (1,8 Bilhão de parâmetros), ele bate de frente com modelos gigantes (16 Bilhões) que são muito mais lentos.

Em resumo: O HEF é como ter um assistente que leu todo o projeto, fez anotações inteligentes e organizou tudo em um "mapa do tesouro" compacto. Quando você precisa de ajuda, ele não te entrega a montanha de papelada; ele te entrega o mapa exato, permitindo que você programe rápido e sem erros.

Resumo técnico

Título: Fusão Hierárquica de Embeddings para Geração de Código Aumentada por Recuperação (HEF)

1. O Problema

A conclusão de código em nível de repositório (Repository-Level Code Completion) exige que o modelo preveja o próximo trecho de código considerando não apenas o arquivo atual, mas também contexto cruzado de todo o projeto (classes importadas, definições de tipos, APIs compartilhadas).

As abordagens atuais de Geração Aumentada por Recuperação (RAG) para código enfrentam dois desafios principais:

1. Custo Online e Latência: Métodos baseados em "injeção de snippets" recuperam e concatenam trechos de código brutos no prompt. Isso acopla o custo de inferência ao tamanho do repositório e ao número de tokens recuperados, introduzindo ruído e latência significativa.
2. Complexidade de Estrutura: Métodos que utilizam grafos ou recuperação iterativa melhoram a relevância, mas exigem múltiplas chamadas de modelo ou travessias de grafos complexas durante a consulta, tornando-os inviáveis para aplicações de baixa latência.

O objetivo é criar um sistema que mantenha a precisão da recuperação de contexto rico, mas com latência subsegundo, desacoplando o tamanho do prompt online do tamanho do repositório.

2. Metodologia: Hierarchical Embedding Fusion (HEF)

O HEF propõe uma representação de repositório em duas etapas: um cache offline e uma interface online baseada em "pseudo-tokens".

A. Etapa Offline: Construção do Cache Hierárquico
O repositório é processado uma única vez para criar uma hierarquia densa de vetores:

1. Fragmentação (Chunking): Arquivos são divididos em blocos de até 512 tokens (limites semânticos ou sintáticos).
2. Embedding: Um codificador (Embedder) congelado (Qwen3-Embedding-8B) converte cada bloco em um vetor denso.
3. Fusão Hierárquica: Um modelo "Fuser" pequeno (Qwen-2.5-Coder-0.5B) funde recursivamente os vetores dos blocos irmãos em vetores de nível superior (arquivo → módulo → repositório).
   • Isso cria uma árvore onde cada nó representa uma agregação semântica de seus filhos.
   • A hierarquia é indexada para busca rápida.

B. Etapa Online: Processamento de Consulta
Durante a inferência:

1. Recuperação: O prefixo do código é usado para buscar os $K$ nós mais relevantes na hierarquia (usando HNSW).
2. Projeção para Pseudo-tokens: Os vetores recuperados são projetados por um Projector (MLP) em "pseudo-tokens" contínuos.
   • Em vez de injetar milhares de tokens de texto, o gerador recebe um número fixo de vetores contínuos (ex: 32 pseudo-tokens).
3. Geração: O modelo gerador (Decoder) recebe o prefixo original + os pseudo-tokens condicionantes para gerar o código.

C. Treinamento e Filtragem de Dados

• Regimes de Treinamento: O sistema explora pré-treinamento contrastivo do Fuser e otimização conjunta (End-to-End) do Fuser, Projector e Gerador.
• Filtragem UWL (Utility-Weighted Likelihood): Um sinal não supervisionado filtra os dados de treinamento, mantendo apenas contextos que aumentam a verossimilhança da conclusão correta, evitando treinar com ruído.

3. Principais Contribuições

1. Método HEF: Uma arquitetura que desacopla o tamanho do repositório do comprimento do prompt online, substituindo snippets brutos por uma interface de pseudo-tokens fixa.
2. Pipeline End-to-End: Integração de um codificador forte, um cache hierárquico denso, um modelo de fusão leve e uma interface de condicionamento, tudo otimizado para conclusão de código.
3. Análise de Regimes de Treinamento: Estudo comparativo entre pré-treinamento contrastivo e otimização conjunta, demonstrando que a otimização conjunta traz ganhos significativos de precisão.
4. Robustez: Demonstração de que a fusão hierárquica é mais robusta a contextos recuperados irrelevantes ou prejudiciais do que a injeção direta de texto.

4. Resultados Experimentais

Os resultados foram avaliados nos benchmarks RepoBench e RepoEval em um único GPU A100.

• Precisão vs. Latência:
  • O HEF (End-to-End) atinge 61,3% de precisão (Exact-Match) no RepoBench, superando o baseline sem recuperação (Qwen-2.5-Coder-1.3B) em 12,2 pontos e o modelo de baixa latência RepoFusion em 22,5 pontos.
  • A latência mediana é de 0,68 segundos, sendo 13x a 26x mais rápido que sistemas baseados em grafos (DRACO, GraphCoder) e recuperação iterativa (RLCoder), que levam de 8s a 17s.
• Eficiência de Tokens: Substitui milhares de tokens recuperados por apenas ~32 pseudo-tokens, mantendo o acesso à informação de nível de repositório.
• Ablações:
  • Tamanho do Fuser: Aumentar o modelo de fusão além de 0,5B parâmetros traz retornos decrescentes, confirmando que sua função é compressão, não raciocínio complexo.
  • Quantidade de Tokens: Cerca de 30-40 pseudo-tokens capturam a maior parte da informação necessária; aumentar além de 60 degrada o desempenho.
  • Robustez: Em cenários com contexto recuperado prejudicial, o HEF degrada menos a performance do que a injeção de snippets brutos.

5. Significado e Conclusão

O trabalho demonstra que é possível realizar conclusão de código em nível de repositório com baixa latência sem sacrificar significativamente a precisão. A chave é a agregação hierárquica de informações em representações densas (vetores) em vez de manter o texto bruto no contexto.

O HEF preenche uma lacuna importante no trade-off entre precisão e latência, oferecendo uma alternativa prática para cenários onde a responsividade é crítica (ex: IDEs em tempo real), sem a necessidade de processamento pesado de grafos ou injeção massiva de contexto. O artigo fornece uma "receita" prática para substituir a injeção de snippets por representações fundidas compactas, validando que a maioria das informações relevantes de um repositório pode ser destilada em uma hierarquia vetorial compacta.