계층형 임베딩 융합을 통한 검색 기반 코드 생성

초록

**
HEF(Hierarchical Embedding Fusion)는 저장소 전체를 두 단계의 밀집 벡터 계층으로 압축하고, 온라인 단계에서는 소수의 벡터를 의사 토큰으로 변환해 코드 생성기에 입력한다. 이를 통해 수천 개의 검색 토큰을 고정된 의사 토큰 수로 대체하면서도 저장소 수준 정보를 유지한다. RepoBench와 RepoEval에서 1.8 B 파라미터 파이프라인이 기존 스니펫 기반 검색과 동등한 정확도를 보이며, A100 하나로 서브초 지연을 달성한다. 또한 그래프·반복 검색 시스템 대비 13~26배 빠른 응답 속도를 기록한다.

**

상세 분석

**
본 논문은 대규모 코드 저장소를 활용한 코드 자동완성·생성 작업에서 발생하는 두 가지 핵심 문제, 즉 검색 토큰 규모에 따른 추론 비용 증가와 긴 컨텍스트가 모델에 주는 노이즈를 해결하기 위해 새로운 두 단계 파이프라인인 Hierarchical Embedding Fusion(HEF)을 제안한다.

1. 오프라인 캐시 단계

   • 저장소를 일정 길이(chunk)로 분할하고, 각 청크를 사전 학습된 임베딩 모델(예: Qwen‑3‑Embedding‑8B)로 고차원 벡터화한다.
   • 작은 fuser 모델(수백만 파라미터)로 청크 벡터들을 계층적으로 집계한다. 구체적으로, 청크‑레벨 임베딩을 클러스터링·트리 구조(예: HNSW, IVF)로 조직해 상위 레벨(문서·패키지) 임베딩을 생성한다. 이 과정은 완전 offline이며, 저장소 전체를 수 GB 수준의 dense cache(수백만 벡터)로 압축한다.
   • 계층 구조는 다중 스케일 정보를 보존한다. 하위 레벨은 세부 구현을, 상위 레벨은 모듈·패키지 수준의 의미를 담아, 검색 시 다양한 granularity를 활용할 수 있다.

2. 온라인 인터페이스 단계

   • 질의(프롬프트)와 함께 가벼운 검색(예: top‑k=5~10)만 수행한다. 검색은 위에서 만든 계층적 인덱스를 이용해 빠르게 수행되며, 반환되는 것은 청크‑레벨 임베딩이 아니라 압축된 벡터이다.
   • 반환된 벡터들을 pseudo‑token으로 매핑한다. 이 매핑은 작은 변환 네트워크(예: 2‑layer MLP)로 수행되며, 각 벡터를 고정된 길이(예: 32) 토큰 시퀀스로 변환한다.
   • 변환된 pseudo‑token은 기존 코드 생성 모델(예: 1.8 B 파라미터 코더‑디코더) 입력에 추가적인 토큰 시퀀스로 삽입된다. 모델은 기존 토큰(프롬프트·코드 컨텍스트)과 pseudo‑token을 동시에 처리해, 저장소 수준 정보를 직접 활용한다.

3. 학습 및 필터링

   • 훈련 시 utility‑weighted likelihood를 도입해, 검색된 컨텍스트가 실제 코드 생성에 기여하는 정도를 정량화한다. 즉, 높은 utility를 가진 컨텍스트에 가중치를 부여해 손실을 계산함으로써, 노이즈가 많은 검색 결과를 억제한다.
   • 또한, harmful retrieval(보안·버그가 포함된 코드) 방지를 위해 컨텍스트 필터링 모델을 별도로 학습한다. 이는 악성 코드 패턴을 사전에 탐지해 캐시에서 제외한다.

4. 실험 및 결과

   • RepoBench와 RepoEval 두 벤치마크에서 HEF는 기존 스니펫 기반 Retrieval‑Augmented Generation(RAG)과 동등하거나 약간 우수한 exact‑match 정확도를 기록했다.
   • 지연(latency) 측면에서, HEF는 동일 하드웨어(A100 40 GB)에서 중간값 0.8 s 이하를 달성했으며, 이는 그래프 기반(예: CodeGraph) 및 반복 검색(예: ReCoSa) 시스템 대비 13~26배 빠른 수치다.
   • pseudo‑token budget(예: 64128 토큰)와 임베딩 모델(BERT‑style vs. Transformer‑style) 변화에 대한 ablation 실험에서, pseudo‑token 수가 64 이하일 때도 성능 저하가 미미했으며, 임베딩 차원 256512가 최적의 trade‑off를 제공한다는 결과를 얻었다.

5. 시사점 및 한계

   • HEF는 저지연·고정 메모리 환경(IDE 플러그인, 클라우드 IDE)에서 실시간 코드 보조에 적합하다.
   • 현재 구현은 정적 저장소에 최적화돼 있어, 빈번히 업데이트되는 레포(예: CI/CD 파이프라인)에서는 캐시 재생성 비용이 문제될 수 있다. 향후 incremental update 메커니즘이 필요하다.
   • pseudo‑token 변환이 선형 구조라 복잡한 논리 관계(예: 함수 호출 그래프)를 완전히 전달하지 못할 가능성이 있다. 향후 graph‑aware 변환기 도입이 기대된다.

**