StepCoder: Mejora la Generación de Código con Aprendizaje por Refuerzo a partir de Retroalimentación del 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

Resumen

El avance de los modelos de lenguaje de gran escala (LLMs, por sus siglas en inglés) ha impulsado significativamente el campo de la generación de código. Trabajos previos integraron el aprendizaje por refuerzo (RL, por sus siglas en inglés) con retroalimentación del compilador para explorar el espacio de salida de los LLMs y mejorar la calidad de la generación de código. Sin embargo, el código extenso generado por los LLMs en respuesta a requisitos humanos complejos convierte la exploración mediante RL en un desafío. Además, dado que las pruebas unitarias pueden no cubrir el código complicado, optimizar los LLMs utilizando estos fragmentos de código no ejecutados resulta ineficaz. Para abordar estos desafíos, presentamos StepCoder, un novedoso marco de RL para la generación de código, que consta de dos componentes principales: CCCS aborda el desafío de la exploración dividiendo la tarea de generación de secuencias largas de código en un Currículo de Subtareas de Completado de Código, mientras que FGO optimiza el modelo únicamente enmascarando los segmentos de código no ejecutados para proporcionar una Optimización de Grano Fino. Además, construimos el conjunto de datos APPS+ para el entrenamiento de RL, el cual ha sido verificado manualmente para garantizar la corrección de las pruebas unitarias. Los resultados experimentales muestran que nuestro método mejora la capacidad de explorar el espacio de salida y supera a los enfoques más avanzados en los benchmarks correspondientes.

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.