StepCoder: 컴파일러 피드백을 통한 강화 학습으로 코드 생성 향상

StepCoder: Improve Code Generation with Reinforcement Learning from Compiler Feedback

February 2, 2024

저자: Shihan Dou, Yan Liu, Haoxiang Jia, Limao Xiong, Enyu Zhou, Junjie Shan, Caishuang Huang, Wei Shen, Xiaoran Fan, Zhiheng Xi, Yuhao Zhou, Tao Ji, Rui Zheng, Qi Zhang, Xuanjing Huang, Tao Gui

cs.AI

초록

대형 언어 모델(LLMs)의 발전은 코드 생성 분야를 크게 촉진시켰다. 기존 연구에서는 컴파일러 피드백과 강화 학습(RL)을 통합하여 LLMs의 출력 공간을 탐색함으로써 코드 생성 품질을 향상시키고자 하였다. 그러나 복잡한 인간의 요구사항에 대응하여 LLMs가 생성하는 긴 코드는 RL 탐색을 어렵게 만든다. 또한, 단위 테스트가 복잡한 코드를 모두 커버하지 못할 수 있기 때문에, 실행되지 않은 코드 스니펫을 사용하여 LLMs를 최적화하는 것은 비효율적이다. 이러한 문제를 해결하기 위해, 우리는 코드 생성을 위한 새로운 RL 프레임워크인 StepCoder를 소개한다. StepCoder는 두 가지 주요 구성 요소로 이루어져 있다: CCCS(Curriculum of Code Completion Subtasks)는 긴 시퀀스의 코드 생성 작업을 코드 완성 하위 작업으로 나누어 탐색 문제를 해결하고, FGO(Fine-Grained Optimization)는 실행되지 않은 코드 세그먼트를 마스킹하여 모델을 최적화한다. 또한, RL 훈련을 위해 APPS+ 데이터셋을 구축하였으며, 이 데이터셋은 단위 테스트의 정확성을 보장하기 위해 수동으로 검증되었다. 실험 결과는 우리의 방법이 출력 공간 탐색 능력을 향상시키고, 해당 벤치마크에서 최신 접근법을 능가함을 보여준다.

English

The advancement of large language models (LLMs) has significantly propelled the field of code generation. Previous work integrated reinforcement learning (RL) with compiler feedback for exploring the output space of LLMs to enhance code generation quality. However, the lengthy code generated by LLMs in response to complex human requirements makes RL exploration a challenge. Also, since the unit tests may not cover the complicated code, optimizing LLMs by using these unexecuted code snippets is ineffective. To tackle these challenges, we introduce StepCoder, a novel RL framework for code generation, consisting of two main components: CCCS addresses the exploration challenge by breaking the long sequences code generation task into a Curriculum of Code Completion Subtasks, while FGO only optimizes the model by masking the unexecuted code segments to provide Fine-Grained Optimization. In addition, we furthermore construct the APPS+ dataset for RL training, which is manually verified to ensure the correctness of unit tests. Experimental results show that our method improves the ability to explore the output space and outperforms state-of-the-art approaches in corresponding benchmarks.