장문 맥락 활용 코드 완성 AI의 비밀 COLA

초록

레포지토리 수준 코드 완성에서 LLM이 긴 교차 파일 컨텍스트 정보를 제대로 활용하지 못하는 문제를 해결하기 위해, COLA(Code Long-context Alignment) 방법론을 제안한다. COLA는 장문 교차 파일 컨텍스트와 컨텍스트 인식 코드 생성을 요구하는 대규모 데이터셋(COLA-132K)을 구축하고, 두 단계 훈련 파이프라인을 통해 LLM이 컨텍스트 내 관련 정보를 찾는 능력을 학습하도록 한다. 실험 결과, COLA는 aiXcoder-7B 등 여러 LLM의 코드 완성 성능을 크게 향상시키고, 새로운 언어로의 일반화 능력과 컨텍스트 활용 능력도 증진시킨다.

상세 분석

본 논문은 레포지토리 수준 코드 완성(Repo-level Code Completion)이라는 실용적인 과제에서 대규모 언어 모델(LLM)이 직면한 근본적인 한계를 정확히 지적하고, 데이터 중심의 우아한 해법을 제시한다. 핵심 문제는 모델의 사전 학습 방식과 실제 작업 요구사항 사이의 ‘미스얼라이먼트(Misalignment)‘다. 사전 학습은 대체로 근접한 토큰에 의존해 다음 토큰을 예측하는 방식으로 이루어지나, 레포지토리 완성에서는 수천 토큰에 달하는 다른 파일의 API나 유사 코드 패턴을 찾아 참조해야 한다. 이 간극으로 인해 모델은 긴 컨텍스트 창을 제공받더라도 그 안의 핵심 정보를 효과적으로 활용하지 못한다.

논문의 핵심 기여인 COLA는 이러한 미스얼라이먼트를 해소하기 위해 모델에게 ‘명시적으로’ 교차 파일 컨텍스트를 활용하도록 가르치는 것이다. 이를 위한 핵심 장치는 COLA-132K 데이터셋이다. 기존 데이터셋과의 차별점은 두 가지다. 첫째, 평균 12K 토큰, 최대 128K 토큰에 이르는 매우 긴 교차 파일 컨텍스트를 포함한다. 둘째, 단순한 코드 조각이 아닌, 컨텍스트에 의존적인 ‘컨텍스트 인식 코드’를 정답으로 요구한다. 이는 크게 두 유형으로, (1) 교차 파일에 정의된 API를 호출하는 코드와, (2) 교차 파일 컨텍스트에 등장하는 코드 스팬과 유사한 패턴의 코드를 생성하는 것이다. 모델은 이러한 정답을 생성하기 위해 반드시 긴 컨텍스트를 샅샅이 훑어 관련 정보를 찾아내야 하므로, 자연스럽게 원하는 능력을 학습하게 된다.

학습은 두 단계로 진행된다. 첫째, 지도 학습을 통해 COLA-132K 데이터로 기본적인 컨텍스트 활용 패턴을 미세 조정한다. 둘째, 정답 코드 생성에 대한 보상을 최대화하는 방향으로 모델을 조정하는 강화 학습 단계를 거쳐 성능을 한층 강화한다. 저자들은 이 방법론을 aiXcoder-7B, DeepSeek-Coder-6.7B, CodeLlama-7B 등 여러 인기 LLM에 적용해 그 유효성을 입증했다. 특히 aiXcoder-7B에 적용한 모델은 aiXcoder-7B-v2로 명명되었으며, Exact Match 점수에서 최대 19.7%p 향상되었고, 파라미터가 훨씬 큰 DeepSeek-Coder-33B보다도 우수한 성능을 보였다. 더욱 주목할 만한 점은 COLA로 학습된 능력이 훈련 데이터에 없던 새로운 프로그래밍 언어(C#, TypeScript)로도 일반화되었다는 것이다. 이는 모델이 단순히 데이터를 외우는 것이 아닌, ‘컨텍스트에서 정보를 찾는 방법’이라는 보다 추상적인 스킬을 습득했음을 시사한다.

결론적으로, 이 연구는 LLM의 장문 컨텍스트 처리 능력을 향상시키는 기존 연구들이 주로 모델 아키텍처나 위치 인코딩에 집중한 반면, 데이터와 학습 목표의 정렬이라는 근본적이고 실용적인 측면에서 접근했다는 점에서 의미가 크다. COLA는 복잡한 모델 수정 없이도 기존 LLM의 레포지토리 이해 및 협업 코딩 보조 능력을 획기적으로 개선할 수 있는 가능성을 제시한다.