리포지토리 수준 코드 완성을 위한 정렬형 검색기 AlignCoder

초록

AlignCoder는 코드 완성 시 검색 질의를 다중 후보 완성으로 강화하고, 강화학습 기반의 AlignRetriever를 통해 검색 단계에서 추론 정보를 활용한다. 이를 통해 기존 RAG 방식의 질의‑목표 코드 불일치와 검색기의 학습 부족 문제를 해결하고, CrossCodeEval와 RepoEval 벤치마크에서 평균 18.1%의 EM 향상을 달성하였다.

상세 분석

본 논문은 리포지토리 수준 코드 완성이라는 특수한 시나리오에서 기존 코드 LLM이 겪는 두 가지 근본적인 한계를 정확히 짚어낸다. 첫 번째는 “질의‑목표 불일치”이다. 기존 RAG 방식은 미완성 코드(Cu)를 그대로 질의로 사용해 전체 저장소에서 상위 k개의 스니펫을 검색한다. 그러나 Cu는 아직 구현되지 않은 토큰을 포함하지 않으며, 목표 코드(Ct)의 핵심 키워드가 질의에 드러나지 않아 의미적 격차가 발생한다. 두 번째는 “검색기 학습 부재”이다. 현재 사용되는 BM25와 같은 희소 검색기나 UniXcoder 기반의 밀집 검색기는 LLM의 추론 결과를 반영하도록 설계되지 않아, 생성된 후보 완성과 질의 사이의 연관성을 학습하지 못한다.

AlignCoder는 이러한 문제를 두 단계로 해결한다. ① 질의 강화(Query Enhancement): 샘플러를 이용해 다수의 후보 완성(예: n=10)을 생성하고, 각 후보를 미완성 코드와 결합해 확장된 질의를 만든다. 이 과정은 “키 토큰”을 후보 완성에 포함시켜 질의 공간을 목표 코드와 더 가깝게 만든다. ② 강화학습 기반 검색기(AlignRetriever): 강화학습(RL) 프레임워크를 도입해 검색기를 학습시킨다. 강화학습의 보상은 검색된 코드 스니펫을 이용해 목표 코드를 생성했을 때의 퍼플렉시티(PPL) 감소량으로 정의한다. 즉, 검색된 컨텍스트가 목표 코드 생성에 얼마나 기여했는지를 직접 측정해 보상으로 활용함으로써, 검색기가 “추론 정보를 활용하는 방법”을 스스로 학습한다.

기술적으로는 다음과 같은 흐름을 가진다. (1) 미완성 코드 Cu를 입력으로 다중 샘플링을 수행해 후보 집합 {Ĉ₁,…,Ĉₙ}을 얻는다. (2) 각 후보를 Cu와 연결해 강화된 질의 Q̂ = concat(Cu, Ĉ₁,…,Ĉₙ) 를 만든다. (3) AlignRetriever는 Q̂을 임베딩하고 전체 저장소 R의 코드 스니펫을 벡터화한 인덱스와 비교해 상위 m개의 스니펫 Ŝ를 반환한다. (4) 반환된 스니펫을 LLM에 컨텍스트로 제공하고, 목표 코드 Ct̂ 를 생성한다. (5) 생성된 Ct̂ 의 로그우도 기반 퍼플렉시티를 계산해 보상 r = -PPL(Ct̂ | Q̂, Ŝ) 로 정의하고, REINFORCE 방식으로 검색기 파라미터 θ를 업데이트한다. 이 루프는 전체 데이터셋에 대해 반복되며, 검색기와 LLM이 공동 최적화된다.

실험에서는 다섯 가지 최신 코드 LLM(DeepSeekCoder‑1B, StarCoder‑Base 등)을 백본으로 사용하고, 두 개의 공개 벤치마크(CrossCodeEval, RepoEval)에서 EM(Exact Match) 점수를 측정했다. 특히 Python 영역에서 기존 최첨단인 RepoCoder 대비 18.1% 절대적 향상을 보였으며, Java, API, 라인 수준 평가에서도 일관된 개선을 기록했다. 또한, 다양한 프로그래밍 언어와 모델 규모에 대해 동일한 프레임워크를 적용했을 때 성능 저하가 거의 없다는 “범용성”을 입증하였다. Ablation study를 통해 (i) 질의 강화 없이 단일 샘플만 사용했을 때 성능이 급격히 떨어짐을, (ii) 강화학습 없이 기존 BM25/UniXcoder를 그대로 사용했을 때도 개선 효과가 제한적임을 확인했다. 마지막으로, 샘플 수 n과 검색 후보 수 m에 대한 민감도 분석을 통해 적절한 n≈10, m≈5가 실용적인 트레이드오프임을 제시한다.

이 논문의 핵심 기여는 (1) “다중 후보 완성을 질의에 통합”함으로써 의미적 격차를 구조적으로 해소한 점, (2) “퍼플렉시티 기반 보상”을 이용해 검색기가 LLM의 추론 정보를 학습하도록 만든 점, (3) 다양한 백본 모델과 언어에 대해 일관된 성능 향상을 보인 범용성을 입증한 점이다. 향후 연구는 (a) 후보 완성의 품질을 사전 필터링하는 메커니즘, (b) 검색 결과의 가중치를 동적으로 조정하는 어텐션 기반 모듈, (c) 대규모 기업용 리포지토리에서 실시간 서비스 적용을 위한 효율적인 인덱싱/재학습 전략 등을 탐색할 여지를 남긴다.