[논문 리뷰] Planning with Large Language Models for Code Generation

Shun Zhang, Zhenfang Chen외 4명|arXiv (Cornell University)|2023. 03. 09.

Software Testing and Debugging Techniques인용 수 13

한 줄 요약

이 논문은 Planning-Guided Transformer Decoding (PG-TD)을 소개합니다. 사전 학습된 코드 생성 Transformer와 협력하는 계획 기반 디코딩 방법으로, 계획자와 테스트 케이스 보상을 활용한 롤아웃 탐색을 수행하여 더 높은 품질의 프로그램을 생성하고, 합격 비율을 향상시키며 미세조정 없이도 제어 가능한 코드 생성을 가능하게 합니다.

ABSTRACT

Existing large language model-based code generation pipelines typically use beam search or sampling algorithms during the decoding process. Although the programs they generate achieve high token-matching-based scores, they often fail to compile or generate incorrect outputs. The main reason is that conventional Transformer decoding algorithms may not be the best choice for code generation. In this work, we propose a novel Transformer decoding algorithm, Planning-Guided Transformer Decoding (PG-TD), that uses a planning algorithm to do lookahead search and guide the Transformer to generate better programs. Specifically, instead of simply optimizing the likelihood of the generated sequences, the Transformer makes use of a planner to generate candidate programs and test them on public test cases. The Transformer can therefore make more informed decisions and generate tokens that will eventually lead to higher-quality programs. We also design a mechanism that shares information between the Transformer and the planner to make our algorithm computationally efficient. We empirically evaluate our framework with several large language models as backbones on public coding challenge benchmarks, showing that 1) it can generate programs that consistently achieve higher performance compared with competing baseline methods; 2) it enables controllable code generation, such as concise codes and highly-commented codes by optimizing modified objective.

연구 동기 및 목표

• 전통적인 Transformer 디코딩(빔 탐색/샘플링)이 코드 생성을 위한 한계, 고토큰 매칭이 있어도 컴파일되지 않거나 잘못된 프로그램이 나올 수 있다는 점을 지적합니다.
• PG-TD를 제안합니다. 이 프레임워크는 테스트 케이스 보상으로 안내된 후보 프로그램을 탐색하기 위해 플래너를 사용하는 계획 기반 디코딩을 제시합니다.
• 미세조정이나 추가 샘플 수집 없이도 표준 코드 생성 Transformer에 모델-독립적으로 통합 가능하도록 만듭니다.
• PG-TD가 더 높은 합격률을 얻고 보상 설계를 통해 응답 길이 축소나 주석 추가와 같은 제어 가능한 목표를 지원하는지 보여줍니다.

제안 방법

• 코드 생성을 문제 설명과 부분적 프로그램을 이어붙인 상태를 상태로, 다음 토큰을 행동으로 가지는 MDP로 공식화합니다.
• Transformer가 제공하는 다음 토큰 확률을 가이드로 사용하는 Monte-Carlo 트리 탐색에서 영감을 받은 플래너(P-UCB 선택 포함)를 적용하여 심층 탐색을 수행합니다.
• 상위-k 개의 높은 확률 다음 토큰으로 노드를 확장해 후보 연속 값을 형성합니다.
• 부분 상태를 Transformer 빔 검색으로 완성하고 공개 테스트 케이스로 평가하여 보상을 얻고, 이 보상은 트리에서 역전파됩니다.
• 트리 구조와 생성 시퀀스의 캐싱 메커니즘을 도입해 Transformer와 플래너 간 정보를 공유하고 효율성을 향상시킵니다.
• APPS와 CodeContests 데이터셋에서 GPT-2 및 GPT-Neo 백본으로 테스트하여 모델-독립적 적용 가능성을 입증합니다.
• 플래너 보상 함수를 수정하여 제어 가능한 목표(예: 더 짧은 코드, 더 많은 주석)를 탐구하고 PG-TD로 생성된 샘플을 사용해 미세조정의 가능성을 보여줍니다.

실험 결과

연구 질문

• RQ1 계획 알고리즘을 사전 학습된 코드 생성 Transformer와 통합하면 표준 디코딩 또는 샘플링 기반 비교 Baseline보다 더 나은 프로그램을 생성할 수 있을까?
• RQ2 PG-TD가 다양한 Transformer 백본(모델-독립적일 수 있음)에서 효과적일까?
• RQ3 정보 공유 및 캐싱을 통한 계획 기반 디코딩이 샘플 효율성 및 계산 시간 향상에 기여할까?
• RQ4 미세조정 없이도 PG-TD가 제어 가능한 코드 생성 목표를 지원할 수 있을까?
• RQ5 PG-TD로 생성된 샘플을 사용해 Transformer를 미세조정해 성능을 further 향상시킬 수 있을까?

주요 결과

• PG-TD는 APPS 및 CodeContests에서 빔 검색 및 샘플링+필터링 기반 기준보다 일관되게 높은 합격률을 보였습니다.
• PG-TD는 GPT-2 및 GPT-Neo 백본에 적용했을 때 모델-독립적 향상을 보여주었습니다.
• 계획 기반 디코딩은 효과적인 선행 탐색을 제공하여 Sequential Monte-Carlo-Guided TD (SMCG-TD)보다 우수하고 같은 계산 예산 하에서 합격률을 달성하거나 초과했습니다.
• 트리 구조 및 시퀀스 캐싱과 같은 캐싱 메커니즘은 계산 시간과 중복된 Transformer 호출을 크게 줄였습니다.
• PG-TD는 플래너 보상 함수를 변경해 제어 가능한 코드 생성을 가능하게 하며, 생성된 샘플은 Transformer를 추가로 미세조정해 성능을 더욱 향상시킬 수 있습니다.