코드 생성의 장기 관점: 라인 단위 MCTS와 자기 정제

초록

본 논문은 대형 언어 모델(LLM)의 코드 생성 시 토큰 단위가 초점이 되는 단기 의존성 문제를 해결하고자, 코드 라인을 기본 처리 단위로 삼아 라인‑레벨 몬테카를로 트리 탐색(LSR‑MCTS)을 제안한다. 각 라인 노드에서 자체 정제(self‑refine) 메커니즘을 적용해 오류를 교정하고, 전체 프로그램을 전역 최적화된 경로로 선택한다. HumanEval·MBPP·Code Contests 등 세 벤치마크에서 기존 Beam‑Search, Top‑p, Refexion, PG‑TD 등에 비해 pass@k 점수가 크게 향상됨을 실험적으로 입증한다.

상세 분석

본 연구는 LLM 기반 코드 생성에서 “단기‑편향”이라는 근본적인 한계를 정확히 짚어낸다. 기존 토큰‑단위 디코딩은 각 토큰이 직전 몇 개 토큰에만 의존하도록 설계돼, 프로그래밍 언어가 요구하는 긴 범위의 구조적 일관성(예: 변수 선언‑사용, 함수 정의‑호출, 제어 흐름) 을 충분히 포착하지 못한다. 저자들은 LLM 내부 어텐션 히트맵을 분석해, 라인 종료 토큰(‘\n’)이 다른 토큰에 비해 현저히 높은 어텐션 스코어를 갖는 것을 발견하고, 이를 “요약 토큰(summary token)”이라 명명한다. 이러한 토큰은 이전 라인의 핵심 정보를 압축해 다음 라인 생성에 영향을 미치므로, 라인 자체를 기본 단위로 삼는 것이 자연스럽다.

LSR‑MCTS는 전통적인 MCTS의 네 단계(Selection‑Expansion‑Evaluation‑Backpropagation)를 라인‑레벨에 맞게 재구성한다. 각 노드는 (1) context: 현재까지 선택된 라인들의 누적 코드 블록, (2) line: 현재 라인 자체, (3) supplement: 라인 완성을 위해 필요한 보조 코드(예: import 문, 함수 정의의 나머지 부분) 로 구성된다. Selection 단계에서는 UCT(Upper Confidence bound for Trees) 공식을 적용해 현재까지 가장 높은 기대 보상을 가진 라인 경로를 탐색한다. Expansion 단계에서는 LLM을 호출해 후보 라인 m개(논문에서는 m=3)와 각각에 대한 보완 코드를 동시에 생성한다. 여기서 핵심은 self‑refine 메커니즘이다. 선택된 라인 중 점수가 낮은 경우 혹은 현재 경로의 마지막 노드에 대해 추가적인 “정제” 라인을 한 번 더 생성한다. 이 정제 라인은 기존 후보와 독립적인 새로운 자식 노드로 삽입되어, 탐색 공간을 인위적으로 제한하면서도 놓칠 수 있는 유망한 브랜치를 보완한다.

Evaluation 단계에서는 생성된 코드 블록을 공개 테스트 케이스에 실행해 통과율(pass rate)을 보상으로 사용한다. 보상은 단순히 pass@k 점수가 아니라, 라인‑레벨에서의 부분 점수도 반영해, 부분적으로 올바른 라인이 있더라도 전체 프로그램이 실패할 경우 그 영향을 완화한다. Backpropagation 단계에서는 이 보상이 트리 상위 노드에 누적돼, 이후 선택 단계에서 더 높은 확률로 좋은 라인들을 선택하도록 유도한다.

실험 설계는 세 가지 공개 데이터셋(HumanEval, MBPP, Code Contests)과 네 종류의 모델(코드‑특화 LLM인 CodeLlama‑7B‑Instruct, aiXcoder‑7B, 일반 LLM인 GPT‑4, Llama‑3‑8B‑Instruct)을 사용해 포괄적으로 진행되었다. Baseline은 Beam‑Search, Top‑p, Refexion(자기‑반성 기반 정제), PG‑TD(토큰‑레벨 MCTS) 등이다. 결과는 모든 모델·데이터셋 조합에서 LSR‑MCTS가 가장 높은 pass@1/3/5 점수를 기록했으며, 특히 토큰‑레벨 MCTS인 PG‑TD 대비 평균 3~5%p 상승을 보였다. 이는 라인‑레벨 탐색이 코드의 전역 구조를 더 잘 파악하고, self‑refine이 로컬 오류를 효과적으로 교정함을 의미한다.

한계점으로는 (1) 라인‑단위 토큰화가 언어마다 다를 수 있어, 비‑Python 언어에 대한 일반화가 아직 검증되지 않았으며, (2) self‑refine 단계에서 추가적인 LLM 호출이 연산 비용을 증가시킨다. 향후 연구에서는 라인‑레벨 어텐션을 사전 학습에 통합하거나, 정제 단계의 비용‑효율성을 높이는 경량화 전략을 탐색할 필요가 있다.