[논문 리뷰] RepoFusion: Training Code Models to Understand Your Repository

Disha Shrivastava, Denis Kocetkov외 3명|arXiv (Cornell University)|2023. 06. 19.

Software Engineering Research인용 수 6

한 줄 요약

RepoFusion는 Fusion-in-Decoder 아키텍처를 사용하여 인코딩된 컨텍스트 스니펫들을 함께 주시함으로써, 임포트된 파일, 형제 파일, 부모 클래스와 같은 여러 레포지터리 수준의 컨텍스트를 활용할 수 있도록 하는 훈련 프레임워크입니다. 이로 인해 단일 라인 코드 완성 성능이 크게 향상됩니다. 220M 파라미터에 불과함에도 불구하고, 73배나 큰 CodeGen-16B-multi(73× larger)와 유사한 성능을 내며, 15.5B 파라미터를 가진 StarCoderBase 모델과도 성능이 유사합니다.

ABSTRACT

Despite the huge success of Large Language Models (LLMs) in coding assistants like GitHub Copilot, these models struggle to understand the context present in the repository (e.g., imports, parent classes, files with similar names, etc.), thereby producing inaccurate code completions. This effect is more pronounced when using these assistants for repositories that the model has not seen during training, such as proprietary software or work-in-progress code projects. Recent work has shown the promise of using context from the repository during inference. In this work, we extend this idea and propose RepoFusion, a framework to train models to incorporate relevant repository context. Experiments on single-line code completion show that our models trained with repository context significantly outperform much larger code models as CodeGen-16B-multi ($\sim73 imes$ larger) and closely match the performance of the $\sim 70 imes$ larger StarCoderBase model that was trained with the Fill-in-the-Middle objective. We find these results to be a novel and compelling demonstration of the gains that training with repository context can bring. We carry out extensive ablation studies to investigate the impact of design choices such as context type, number of contexts, context length, and initialization within our framework. Lastly, we release Stack-Repo, a dataset of 200 Java repositories with permissive licenses and near-deduplicated files that are augmented with three types of repository contexts. Additionally, we are making available the code and trained checkpoints for our work. Our released resources can be found at \url{https://huggingface.co/RepoFusion}.

연구 동기 및 목표

• 기존 코드 LLM이 임포트, 부모 클래스, 파일 간 관계와 같은 레포지터리 수준의 컨텍스트를 이해하는 데에 한계가 있어 정확하지 않은 코드 완성을 유도하는 문제를 해결하기 위해.
• 추론 중에 레포지터리 전반에서 관련된 여러 코드 스니펫을 통합할 수 있도록 모델을 훈련시켜 코드 완성 정확도를 향상시키기 위해.
• 레포지터리 컨텍스트를 활용해 훈련된 작은 모델이 더 큰 모델을 능가할 수 있음을 보여주어, 모델 규모가 성능을 결정하는 유일한 요소라는 가정을 도전하기 위해.
• 향후 연구를 지원하기 위해 200개의 허용형 라이선스를 가진 자바 레포지터리와 세 가지 유형의 레포지터리 컨텍스트를 포함한 새로운 데이터셋인 Stack-Repo를 공개하기 위해.
• 검색 기반 메커니즘에 종속되지 않고 다양한 코드 생성 작업에 적용 가능한 재사용 가능한 프레임워크를 제공하기 위해.

제안 방법

• RepoFusion는 Fusion-in-Decoder (FiD) 아키텍처를 사용하여 각 컨텍스트를 별도로 인코딩하고, 그들 간의 공동 주시를 디코더에서 수행함으로써 여러 레포지터리 컨텍스트를 통합합니다.
• 각 레포지터리 컨텍스트는 관련 파일(예: 임포트된 파일, 형제 파일, 부모 클래스)에서 검색되며, 문법적 및 의미적 컨텍스트를 유지하기 위해 고정 길이의 주변 코드 윈도우로 보강됩니다.
• 목표 틈새와 그에 해당하는 다중 소스 레포지터리 컨텍스트를 조합하여 다음 토큰 예측 목표로 모델을 피니테이닝합니다.
• 이 프레임워크는 다양한 컨텍스트 유형(예: 임포트, 형제 파일, 부모 클래스 등)을 지원하며, 컨텍스트의 수와 길이는 탐색 실험을 통해 최적화된 하이퍼파라미터입니다.
• RepoFusion는 엔드 투 엔드로 훈련되며, 기반 검색 메커니즘에 대해 무관하므로 다양한 검색 전략과 통합이 가능합니다.
• 이 방법은 모델이 레포지터리에서 유래한 다양한 컨텍스트 신호를 우선순위 정렬하고 융합하도록 학습함으로써 예측 정확도를 향상시킵니다.

실험 결과

연구 질문

• RQ1 작은 코드 모델을 레포지터리 수준의 컨텍스트로 훈련시킬 경우, 그러한 컨텍스트 없이 훈련된 더 큰 모델에 비해 코드 완성 성능이 크게 향상될 수 있는가?
• RQ2 다양한 유형의 레포지터리 컨텍스트(예: 임포트된 파일, 형제 파일 등)는 모델 성능에 어떻게 기여하는가?
• RQ3 모델 성능을 최대화하기 위해 레포지터리 컨텍스트의 최적의 수와 길이는 무엇인가?
• RQ4 초기화 전략과 훈련 목표는 다중 소스 컨텍스트 융합 능력에 어떤 영향을 미치는가?
• RQ5 레포지터리 컨텍스트를 활용해 훈련된 작은 모델이, 그러한 컨텍스트 없이 훈련된 훨씬 더 큰 모델의 성능을 따라하거나 능가할 수 있는가?

주요 결과

• 220M 파라미터를 가진 RepoFusion는 16B 파라미터를 가진 CodeGen-16B-multi 모델보다 단일 라인 코드 완성에서 성능이 뛰어나며, 크기가 73배나 작음에도 불구하고 성능을 뛰어넘음.
• RepoFusion는 더 복잡한 Fill-in-the-Middle 목표로 훈련된 15.5B 파라미터를 가진 StarCoderBase 모델과 성능이 유사함.
• 탐색 실험 결과, 레포지터리 내 다양한 소스(예: 임포트된 파일, 형제 파일 등)의 정보를 활용하는 것이 성능 향상에 필수적임을 확인함.
• 레포지터리 컨텍스트의 수와 길이는 성능에 상당한 영향을 미치며, 체계적인 탐색 실험을 통해 최적의 설정을 도출함.
• 모델의 효과성은 컨텍스트 길이나 컨텍스트 수에 기인한 것이 아니라, 융합하는 데 학습한 컨텍스트 신호의 품질과 다양성에 기인함.
• 200개의 허용형 라이선스를 가진 자바 레포지터리와 세 가지 유형의 레포지터리 컨텍스트를 포함한 Stack-Repo 데이터셋의 공개로, 검색 기반 코드 모델에 대한 재현 가능성과 향후 연구가 가능해짐.