PerfRL: A Small Language Model Framework for Efficient Code Optimization

TL;DR

PerfRL introduces a framework combining small language models and reinforcement learning to automate code optimization, achieving faster results with less computational resources and fewer errors.

Contribution

This paper presents PerfRL, a novel framework that enhances code optimization efficiency using small language models and reinforcement learning with environment feedback.

Findings

PerfRL outperforms existing models in speed and resource efficiency.

It reduces logical and syntactical errors in code optimization.

Achieves comparable or better optimization results with shorter training times.

Abstract

Code optimization is a challenging task requiring a substantial level of expertise from developers. Nonetheless, this level of human capacity is not sufficient considering the rapid evolution of new hardware architectures and software environments. In light of this, recent research proposes adopting machine learning and artificial intelligence techniques to automate the code optimization process. In this paper, we introduce PerfRL, an innovative framework designed to tackle the problem of code optimization. Our framework leverages the capabilities of small language models (SLMs) and reinforcement learning (RL), facilitating a system where SLMs can assimilate feedback from their environment during the fine-tuning phase, notably through unit tests. When benchmarked against existing models, PerfRL demonstrates superior efficiency in terms of speed and computational resource usage, attributed to its reduced need for training steps and its compatibility with SLMs. Furthermore, it substantially diminishes the risk of logical and syntactical errors. To evaluate our framework, we conduct experiments on the PIE dataset using a lightweight large language model (i.e., CodeT5) and a new reinforcement learning algorithm, namely RRHF. For evaluation purposes, we use a list of evaluation metrics related to optimization quality and speedup. The evaluation results show that our approach achieves similar or better results compared to state-of-the-art models using shorter training times and smaller pre-trained models.