StepCoder: Melhorando a Geração de Código com Aprendizado por Reforço a partir de Feedback do Compilador

StepCoder: Improve Code Generation with Reinforcement Learning from Compiler Feedback

February 2, 2024

Autores: 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

Resumo

O avanço dos modelos de linguagem de grande escala (LLMs) impulsionou significativamente o campo da geração de código. Trabalhos anteriores integraram o aprendizado por reforço (RL) com feedback de compiladores para explorar o espaço de saída dos LLMs, visando melhorar a qualidade da geração de código. No entanto, o código extenso gerado pelos LLMs em resposta a requisitos humanos complexos torna a exploração via RL um desafio. Além disso, como os testes unitários podem não cobrir o código complicado, otimizar os LLMs usando esses trechos de código não executados é ineficaz. Para enfrentar esses desafios, apresentamos o StepCoder, um novo framework de RL para geração de código, composto por dois componentes principais: o CCCS aborda o desafio da exploração ao dividir a tarefa de geração de sequências longas de código em um Currículo de Subtarefas de Conclusão de Código, enquanto o FGO otimiza o modelo apenas mascarando os segmentos de código não executados para fornecer uma Otimização de Granularidade Fina. Adicionalmente, construímos o conjunto de dados APPS+ para treinamento de RL, que é verificado manualmente para garantir a correção dos testes unitários. Os resultados experimentais mostram que nosso método melhora a capacidade de explorar o espaço de saída e supera abordagens state-of-the-art nos benchmarks correspondentes.

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.