Hierarchical Embedding Fusion for Retrieval-Augmented Code Generation

이 논문은 대규모 코드 저장소의 정보를 고정된 수의 의사 토큰으로 압축하여 지연 시간을 획기적으로 줄이면서도 저장소 수준의 맥락을 유지하는 계층적 임베딩 융합 (HEF) 방식을 제안합니다.

핵심

이 논문은 **"방대한 코드 도서관에서 필요한 정보를 찾아내어 코딩을 도와주는 AI"**에 대한 연구입니다. 제목은 '계층적 임베딩 융합 (HEF)'이라고 하는데, 이를 쉽게 풀어서 설명해 드리겠습니다.

🏗️ 비유: 거대한 도서관과 똑똑한 사서

상상해 보세요. 전 세계의 모든 코드가 담긴 거대한 도서관이 있다고 칩시다. 지금 작성 중인 코드가 어떤 기능을 해야 할지 모를 때, 이 도서관에서 관련 정보를 찾아와야 합니다.

기존의 방식들은 다음과 같은 문제가 있었습니다:

1. 책장 전체를 가져오는 방식 (Snippet Injection): 필요한 정보를 찾기 위해 관련 책장 (코드 조각) 100 권을 모두 가져와서 책상 위에 펼쳐놓습니다. 책상이 너무 꽉 차서 (컴퓨터 메모리 부족) 작업이 느려지고, 쓸데없는 책들이 섞여 있어 집중이 안 됩니다.
2. 복잡한 지도를 그리는 방식 (Graph-based): 책과 책 사이의 관계를 복잡한 지도로 그려가며 찾아갑니다. 정확하긴 하지만, 지도를 그리는 데 시간이 너무 오래 걸려서 답이 늦게 나옵니다.

✨ HEF 의 새로운 아이디어: "요약된 지도"와 "기억의 열쇠"

이 논문에서 제안한 HEF는 이 문제를 아주 영리하게 해결합니다. 두 단계로 나뉩니다.

1 단계: 오프라인 작업 (도서관 정리)

• 작업: 도서관에 있는 모든 책 (코드 파일) 을 미리 읽어서, 핵심 내용만 추려낸 요약본을 만듭니다.
• 계층적 구조:
  • 작은 책장 (파일) → 층 (모듈) → 건물 전체 (저장소) 순서로 내용을 압축합니다.
  • 마치 책 한 권의 내용을 "이 책은 A, B, C 를 다룬다"는 한 줄 요약으로 바꾸는 것과 같습니다.
  • 이 요약본들은 밀집된 벡터 (숫자 덩어리) 형태로 저장되어, 나중에 바로 꺼내 쓸 수 있게 됩니다.
• 효과: 수천 페이지의 코드를 읽을 필요 없이, 이 요약본만 보면 전체적인 맥락을 파악할 수 있습니다.

2 단계: 온라인 작업 (실시간 코딩)

• 질문: 개발자가 "이 함수를 어떻게 구현하지?"라고 물어보면, AI 는 도서관에서 가장 관련 있는 요약본 32 개만 찾아옵니다.
• 가상 토큰 (Pseudo-tokens): 이 32 개의 요약본을 AI 가 이해할 수 있는 **마법의 열쇠 (가상 토큰)**로 바꿉니다.
• 결과: 원래는 수천 개의 단어를 입력해야 했을 것을, 32 개의 마법 열쇠만 입력하면 됩니다.
  • 속도: 책장 전체를 가져오는 게 아니라 열쇠 몇 개만 끼워 넣으니, 0.6 초 만에 답이 나옵니다. (기존 방식은 10 초 이상 걸림)
  • 정확도: 요약본이 핵심을 잘 담고 있어서, 책장 전체를 가져왔을 때와 거의 똑같은 정확도를 냅니다.

🚀 왜 이것이 중요한가요?

1. 속도 vs 정확도의 딜레마 해결: 보통 "정확하려면 느려지고, 빠르려면 정확도가 떨어진다"고 생각했는데, HEF 는 둘 다 잡았습니다.
2. 소음 제거: 관련 없는 코드가 섞여 들어와서 AI 가 헷갈리는 현상 (노이즈) 을 줄여줍니다. 요약본만 있기 때문입니다.
3. 저렴한 비용: 거대한 컴퓨터가 아니라, 일반적인 서버 하나에서도 빠르게 돌아갑니다.

💡 핵심 요약

이 기술은 **"방대한 코드 도서관을 미리 요약해서 '요약 카드'로 만들어두고, 코딩할 때 필요한 카드 몇 장만 뽑아내어 AI 에게 보여주는 방식"**입니다.

• 기존 방식: "이 책 100 권 다 읽어봐." (느리고 지루함)
• HEF 방식: "이 책들의 핵심 내용 32 줄 요약본만 봐." (빠르고 정확함)

결론적으로, 이 연구는 AI 가 방대한 프로젝트의 맥락을 이해하면서도, 개발자가 기다리는 시간을 1 초 미만으로 줄여주는 초고속 코딩 비서를 만드는 방법을 제시했습니다.

기술 요약

Hierarchical Embedding Fusion (HEF) 기술 요약

본 논문은 저지연 (Low-latency) 환경에서 저장소 수준 (Repository-level) 의 코드 생성을 위한 새로운 접근법인 **계층적 임베딩 융합 (Hierarchical Embedding Fusion, HEF)**을 제안합니다. 기존 검색 증강 생성 (RAG) 방식이 직면한 긴 컨텍스트 노이즈와 높은 온라인 비용 문제를 해결하기 위해, 오프라인에서 저장소 정보를 압축된 계층적 벡터로 변환하고, 온라인에서는 이를 고정된 수의 '의사 토큰 (pseudo-tokens)'으로 변환하여 생성 모델에 주입하는 두 단계 파이프라인을 제시합니다.

---

1. 문제 정의 (Problem)

저장소 수준의 코드 완성 (Repository-level Code Completion) 은 단일 파일 내의 코드뿐만 아니라, 프로젝트 전체에 걸친 다른 파일의 클래스, 함수 시그니처, 타입 정의 등을 참조해야 정확한 생성이 가능합니다.

• 기존 방식의 한계:
  • Snippets Injection (조각 삽입): 검색된 원본 코드 조각을 프롬프트에 직접 붙이는 방식은 검색된 토큰 수에 비례하여 온라인 지연 시간 (Latency) 이 증가하고, 관련 없는 노이즈가 컨텍스트 윈도우에 유입되어 성능을 저하시킵니다.
  • 구조 인식 (Structure-aware) 방식: 그래프 기반이나 반복적 검색 (Iterative retrieval) 방식은 정확도는 높일 수 있으나, 매 쿼리마다 그래프 탐색이나 여러 번의 모델 호출이 필요하여 실시간 응답이 어렵습니다.
  • 긴 컨텍스트 모델: 전체 저장소를 입력으로 주는 방식은 계산 비용이 매우 큽니다.

2. 방법론 (Methodology)

HEF 는 오프라인 (Offline) 단계와 온라인 (Online) 단계로 구성된 2 단계 파이프라인입니다.

A. 오프라인 단계: 계층적 캐시 구축 (Hierarchical Cache Construction)

저장소 전체를 한 번만 처리하여 재사용 가능한 밀집 벡터 (Dense Vectors) 계층을 생성합니다.

1. Chunking: 소스 파일을 최대 512 토큰 단위의 의미 있는 조각 (Chunk) 으로 분할합니다.
2. Embedding: 고정된 (Frozen) 인코더 (Qwen3-Embedding-8B) 를 사용하여 각 조각을 고밀도 벡터로 변환합니다.
3. Hierarchy Construction (Fuser):
   • 조각 벡터들을 파일, 모듈, 저장소 단위로 계층화합니다.
   • 작은 **Fuser 모델 (Qwen-2.5-Coder-0.5B)**이 자식 벡터들을 재귀적으로 병합하여 상위 레벨 (파일/모듈/저장소) 의 벡터를 생성합니다.
   • 이 과정에서 수천 개의 원본 토큰이 계층적 벡터 트리로 압축됩니다.
4. Indexing: 생성된 모든 계층 노드를 HNSW 인덱스로 저장합니다.

B. 온라인 단계: 쿼리 처리 및 생성 (Query Processing)

사용자가 코드를 입력하면 실시간으로 다음과 같은 과정을 거칩니다.

1. Retrieval: 입력 프롬프트의 끝부분을 인코딩하여 계층적 캐시에서 상위 K(예: 32) 개의 관련 노드를 검색합니다.
2. Pseudo-Token Projection: 검색된 벡터들을 **Projector (MLP)**를 통해 생성 모델이 이해할 수 있는 **의사 토큰 (Pseudo-tokens)**으로 변환합니다.
   • 이 과정은 수천 개의 원본 토큰을 고정된 수 (예: 32 개) 의 벡터로 대체하여 프롬프트 길이를 제한합니다.
3. Generation: 변환된 의사 토큰을 소프트 프롬프트 (Soft-prompt) 형태로 생성기 (Generator) 에 주입하여 코드를 완성합니다.

C. 데이터 필터링 및 학습 전략

• UWL (Utility-Weighted Likelihood): 학습 데이터에서 실제 완성 품질을 높이는 컨텍스트만 선별하기 위해 UWL 점수를 활용하여 불필요한 데이터를 필터링합니다.
• 학습 regimes:
  1. Contrastive Pre-training: Fuser 모델을 InfoNCE 손실 함수로 사전 학습시켜 계층적 벡터의 품질을 높입니다.
  2. End-to-End Optimization: Fuser, Projector, Generator 를 함께 미세 조정하여 최적의 성능을 달성합니다.

3. 주요 기여 (Key Contributions)

1. 계층적 밀집 캐시 및 의사 토큰 인터페이스: 저장소 크기와 온라인 프롬프트 길이를 분리 (Decouple) 하여, 저장소 규모가 커져도 지연 시간이 일정하게 유지되도록 합니다.
2. 학습 regimes 분석: 대비 학습 (Contrastive) 과 종단간 (End-to-End) 학습의 효과를 분석하고, 데이터 필터링 (UWL) 의 중요성을 입증했습니다.
3. 성능 및 효율성 입증: RepoBench 와 RepoEval 벤치마크에서 기존 고지연 모델들과 유사하거나 더 높은 정확도를 유지하면서, 지연 시간을 획기적으로 단축했습니다.

4. 실험 결과 (Results)

• 정확도 (Accuracy):
  • RepoBench: HEF (End-to-End) 는 **61.3%**의 Exact-Match 정확도를 기록했습니다. 이는 1.8B 파라미터 모델로, 16B 파라미터를 사용하는 GraphCoder(64.1%) 와 비교해도 2.8%p 차이에 불과합니다.
  • RepoEval: 42.7% 의 정확도를 기록하여 기존 저지연 모델 (RepoFusion 33.2%) 보다 10%p 이상 우위를 점했습니다.
• 지연 시간 (Latency):
  • 단일 A100 GPU 에서 중앙값 0.68 초의 응답 시간을 기록했습니다.
  • 기존 고지연 모델 대비 13 배~26 배 빠릅니다 (GraphCoder: 17.5 초, DRACO: 11.0 초).
• 강건성 (Robustness): 관련 없는 해로운 컨텍스트가 검색되더라도, 원본 텍스트를 직접 주입하는 방식보다 HEF 가 성능 저하를 훨씬 덜 겪는 것을 확인했습니다.

5. 의의 및 결론 (Significance)

HEF 는 **정확도와 지연 시간 사이의 트레이드오프 (Trade-off)**를 효과적으로 균형 잡은 솔루션입니다.

• 실용성: 대규모 코드베이스를 실시간으로 처리해야 하는 IDE 자동 완성 등 저지연이 필수적인 시나리오에서 고비용의 그래프 기반 검색을 대체할 수 있는 실용적인 대안을 제시합니다.
• 효율성: 저장소 전체를 매번 처리할 필요 없이, 오프라인에서 한 번만 압축된 계층을 구축하면 되므로 온라인 비용이 극도로 낮습니다.
• 미래 방향: 심볼릭 프로그램 구조와 컴팩트한 연속적 캐시를 결합한 하이브리드 방식 등 향후 연구의 기초를 마련했습니다.

결론적으로, HEF 는 원본 코드 조각을 직접 주입하는 방식의 한계를 극복하고, 압축된 계층적 벡터 표현을 통해 저장소 수준의 맥락을 유지하면서도 초단위 응답이 가능한 코드 생성 파이프라인의 새로운 표준을 제시합니다.