[논문 리뷰] Code Repair with LLMs gives an Exploration-Exploitation Tradeoff

Hao Tang, Keya Hu외 5명|arXiv (Cornell University)|2024. 05. 26.

Scientific Computing and Data ManagementDecision Sciences인용 수 3

한 줄 요약

이 논문은 LLM 기반 프로그램 합성 방법인 REx를 소개한다. REx는 다중 손잡이 슬롯머신 설정에서 톰슨 샘플링을 사용하여 코드 복구를 탐색-이점의 상충 관계로 프레임워크화한다. 높은 성능을 보이는 프로그램을 개선하는 것(이점)과 덜 방문된 프로그램을 탐색하는 것(탐색) 사이의 균형을 통해, REx는 코드 생성, 루프 인variant 합성, 시각적 추론 작업 전반에서 더 많은 문제를 해결하면서도 LLM 호출 수를 1.5배에서 5배까지 줄였다.

ABSTRACT

Iteratively improving and repairing source code with large language models (LLMs), known as refinement, has emerged as a popular way of generating programs that would be too complex to construct in one shot. Given a bank of test cases, together with a candidate program, an LLM can improve that program by being prompted with failed test cases. But it remains an open question how to best iteratively refine code, with prior work employing simple greedy or breadth-first strategies. We show here that refinement exposes an explore-exploit tradeoff: exploit by refining the program that passes the most test cases, or explore by refining a lesser considered program. We frame this as an arm-acquiring bandit problem, which we solve with Thompson Sampling. The resulting LLM-based program synthesis algorithm is broadly applicable: Across loop invariant synthesis, visual reasoning puzzles, and competition programming problems, we find that our new method can solve more problems using fewer language model calls.

연구 동기 및 목표

• 간단한 탐욕스러운 또는 너비 우선 탐색 정책을 사용하는 기존 코드 정련 전략의 비효율성을 해결한다. 이러한 전략들은 근본적인 성과 향상에 기여하지 못한다.
• LLM을 사용한 반복적 코드 복구를 탐색-이점 상충 관계로 재정의한다. 즉, 더 많은 테스트 케이스를 통과하는 프로그램을 정련하는 것(이점)과 정련 횟수가 적은 프로그램을 탐색하는 것(탐색) 사이에서 선택한다.
• 무한 분기 및 확률적 트리에서 몬테카를로 트리 탐색의 금전적 비용을 피할 수 있는 확장 가능한 밴딧 기반 전략을 개발한다.
• 경쟁 프로그래밍, 소프트웨어 검증, 시각적 추론 퍼즐과 같은 다양한 분야에서 효율적이고 효과적인 프로그램 합성을 가능하게 한다.
• 정련 과정에서 탐색-이점 균형을 최적화하여 복잡한 프로그래밍 문제를 해결하기 위해 필요한 LLM 호출 수를 최소화한다.

제안 방법

• 각 프로그램을 '손잡이'로 간주하고 보상으로 통과하는 테스트 케이스 수를 사용하여 코드 정련을 다중 손잡이 슬롯머신 문제로 모델링한다.
• 각 프로그램의 기대 보상에 대한 사후 분포에서 샘플링하여 탐색과 이점을 균형 잡는 톰슨 샘플링을 사용한다.
• 각 LLM 호출 후 베이지안 업데이트를 통해 각 프로그램의 보상 분포에 대한 확률적 믿음을 유지한다.
• 각 단계에서 사후 분포에서 기대 보상을 샘플링하고, 가장 높은 샘플 값의 프로그램을 선택하여 정련한다.
• 선택된 프로그램을 정련함으로써 검색 트리를 동적으로 확장하고, 새로운 LLM 호출을 통해 새로운 노드(정련된 결과)를 생성한다.
• 무한한 정련 트리에서 효율적으로 탐색하기 위해 밴딧 전략을 재귀적으로 적용하며, 할인된 미래 보상 최적화를 통해 깊은 롤아웃을 피한다.

실험 결과

연구 질문

• RQ1 탐색과 이점을 균형 잡는 밴딧 기반 전략이 LLM 기반 코드 생성에서 단순한 탐욕스러운 전략이나 너비 우선 탐색 전략보다 뛰어나게 성능을 발휘할 수 있는가?
• RQ2 반복적 코드 복구에서 탐색-이점 상충 관계가 복잡한 프로그래밍 문제를 해결하기 위해 필요한 LLM 호출 수에 어떤 영향을 미치는가?
• RQ3 톰슨 샘플링이 LLM 기반 프로그램 합성에서 무한 분기 및 확률적 전이를 처리하는 데 얼마나 잘 적응될 수 있는가?
• RQ4 제안된 REx 방법은 루프 인variant 합성, 경쟁 프로그래밍, 시각적 추론 퍼즐과 같은 다양한 프로그래밍 작업에 일반화될 수 있는가?
• RQ5 이 방법은 표준 정련 기반 전략으로는 도달할 수 없는 어려운 문제를 일관되게 해결할 수 있는가?

주요 결과

• REx는 평가된 모든 도메인에서 기준 전략 대비 LLM 호출 수를 1.5배에서 5배까지 줄였다.
• REx는 특히 기존에 해결이 어려운 도전적인 문제들에 대해 탐욕스러운 전략과 너비 우선 전략보다 더 많은 문제를 해결했다.
• 루프 인variant 합성 과제에서는 이전 방법이 실패한 문제들에 대해서도 REx가 정확한 인버리언트를 성공적으로 생성했다.
• ARC 벤치마크의 시각적 추론 퍼즐에 대해서는 REx가 표준 정련 기반 전략보다 더 높은 성공률을 기록하면서도 LLM 호출 수를 줄였다.
• 이 방법은 초기 프롬프트에서 독립적으로 동일하게 분포된 샘플링보다 일관되게 뛰어난 성능을 보이며, 탐색-이점 기반의 구조화된 정련이 무작위 재샘플링보다 더 효과적임을 입증했다.
• 톰슨 샘플링은 무한하고 확률적인 정련 트리에서 탐색과 이점을 효과적으로 균형 잡으며, 비용이 많이 드는 롤아웃 없이도 효율적인 탐색을 가능하게 했다.