Auto-Treinamento Iterativo para Geração de Código via Re-Ranqueamento Reforçado

Iterative Self-Training for Code Generation via Reinforced Re-Ranking

April 13, 2025

Autores: Nikita Sorokin, Ivan Sedykh, Valentin Malykh

cs.AI

Resumo

Gerar código de alta qualidade que resolva tarefas de programação complexas é um desafio, especialmente com os modelos baseados em decodificadores atuais, que produzem saídas altamente estocásticas. Na geração de código, até mesmo erros menores podem facilmente comprometer toda a solução. Aproveitar múltiplas soluções amostradas pode melhorar significativamente a qualidade geral da saída. Uma maneira eficaz de aprimorar a geração de código é emparelhar um modelo de geração de código com um modelo de reranking, que seleciona a melhor solução entre as amostras geradas. Propomos uma nova abordagem iterativa de auto-treinamento para modelos de reranking usando Otimização de Política Proximal (PPO), com o objetivo de melhorar tanto a precisão do reranking quanto o processo geral de geração de código. Diferentemente das abordagens tradicionais de PPO, onde o foco está em otimizar um modelo generativo com um modelo de recompensa, nossa abordagem enfatiza o desenvolvimento de um modelo robusto de recompensa/reranking. Esse modelo melhora a qualidade do código gerado por meio do reranking e aborda problemas e erros que o modelo de recompensa pode ignorar durante o alinhamento de PPO com o reranker. Nosso método refina iterativamente o conjunto de dados de treinamento reavaliando as saídas, identificando exemplos negativos de alta pontuação e incorporando-os no loop de treinamento, impulsionando o desempenho do modelo. Nossa avaliação no conjunto de dados MultiPL-E demonstra que nosso modelo de 13,4 bilhões de parâmetros supera um modelo de 33 bilhões em qualidade de geração de código, sendo três vezes mais rápido. Além disso, ele alcança desempenho comparável ao GPT-4 e o supera em uma linguagem de programação.

English

Generating high-quality code that solves complex programming tasks is challenging, especially with current decoder-based models that produce highly stochastic outputs. In code generation, even minor errors can easily break the entire solution. Leveraging multiple sampled solutions can significantly improve the overall output quality. One effective way to enhance code generation is by pairing a code generation model with a reranker model, which selects the best solution from the generated samples. We propose a novel iterative self-training approach for self-training reranker models using Proximal Policy Optimization (PPO), aimed at improving both reranking accuracy and the overall code generation process. Unlike traditional PPO approaches, where the focus is on optimizing a generative model with a reward model, our approach emphasizes the development of a robust reward/reranking model. This model improves the quality of generated code through reranking and addresses problems and errors that the reward model might overlook during PPO alignment with the reranker. Our method iteratively refines the training dataset by re-evaluating outputs, identifying high-scoring negative examples, and incorporating them into the training loop, that boosting model performance. Our evaluation on the MultiPL-E dataset demonstrates that our 13.4B parameter model outperforms a 33B model in code generation quality while being three times faster. Moreover, it achieves performance comparable to GPT-4 and surpasses it in one programming language.