Applying RLAIF for Code Generation with API-usage in Lightweight LLMs (2024)

Abstract: Reinforcement Learning from AI Feedback (RLAIF) has demonstrated significant potential across various domains, including mitigating harm in LLM outputs, enhancing text summarization, and mathematical reasoning. This paper introduces an RLAIF framework for improving the code generation abilities of lightweight (<1B parameters) LLMs. We specifically focus on code generation tasks that require writing appropriate API calls, which is challenging due to the well-known issue of hallucination in LLMs. Our framework extracts AI feedback from a larger LLM (e.g., GPT3.5) through a specialized prompting strategy and uses this data to train a reward model towards better alignment from smaller LLMs. We run our experiments on the Gorilla dataset and meticulously assess the quality of the modelgenerated code across various metrics, including AST, ROUGE, and Code-BLEU, and develop a pipeline to compute its executability rate accurately. Our approach significantly enhances the fine-tuned LLM baseline’s performance, achieving a 4.5% improvement in executability rate. Notably, a smaller LLM model (780M parameters) trained with RLAIF surpasses a much larger fine-tuned baseline with 7B parameters, achieving a 1.0% higher code executability rate.

heading: 1 Introduction
text: LLMs have demonstrated unprecedented natural language understanding and generation capabilities in recent times (Brown et al., 2020; Chowdhery et al., 2022; OpenAI, 2023; Anil et al., 2023; Jiang et al., 2023; Touvron et al., 2023). Reinforcement Learning with Human Feedback (RLHF) is a key contributor to this success. RLHF is a fine-tuning approach that uses human feedback to train models by incorporating human evaluations into the reward signal. This method improves model performance on complex tasks by aligning the model’s
† Work done as a part of an internship at Apple.
behavior with human preferences. However, this technique is expensive due to the requirement for high-quality human feedback. RLAIF (Bai et al., 2022; Lee et al., 2023) has emerged as a promising alternative to replace human feedback with AI feedback, making the fine-tuning more scalable. Concurrently, there is growing research interest in teaching LLMs how to use external tools (APIs) (Schick et al., 2024; Nakano et al., 2021; Patil et al., 2023; Qin et al., 2023; Li et al., 2023; Zhuang et al., 2024; Hao et al., 2024). However, the focus on lightweight models (<1B parameters) is limited. In this work, we propose an RLAIF framework to enhance lightweight LLMs’ capability to generate code and effectively integrate API calls. Following Patil et al. (2023), we consider the task of generating Python codes that include suitable API calls given instructions across a wide array of applications. The authors published the Gorilla dataset and showed that fine-tuned LLaMA-7B (Touvron et al., 2023) on this dataset outperforms non-finetuned LLMs like GPT-4 (OpenAI, 2023) in terms of understanding a natural language request and mapping it to API calls. Using our RLAIF framework, we fine-tune GPT-2-large (Radford et al., 2019) (780M parameters), which not only demonstrates comparable API call correctness to (Patil et al., 2023) but also surpasses its code generation performance.
Code Generation. Although extensively studied since the early days of AI research, code generation (Waldinger and Lee, 1969; Budinsky et al., 1996; Svyatkovskiy et al., 2020; Li et al., 2022) remains a challenging problem. In recent years, the community has explored ways to apply RL in training machine learning models for code generation tasks. For instance, Seq2SQL (Zhong et al., 2017) proposed a neural network trained through RL for generating SQL queries given a text description. During training, a generated query is executed
39
against a database, and the result is utilized as the reward in the RL algorithm. Le et al. (2022) developed CodeRL, a sequence-to-sequence language model fine-tuned through an actor-critic RL approach for program synthesis. The code-generator LM is treated as the actor during the training, and the critic model, which is trained to predict unit test results, provides the reward for a generated code. Another work (Shojaee et al., 2023) similar to the above, proposed using feedback from code execution and a ground truth target code to compute the reward. While these approaches may perform well on classical programming tasks (e.g., writing SQL queries, solving competitive/interviewlevel coding problems, etc.), they are inapplicable on Gorilla-like (Patil et al., 2023) code generation where the program is required to load and execute ML models using the correct API. The bottleneck comes from the fact that the above-mentioned techniques require execution of the generated code to either compute the reward directly or train the critic model, but running thousands of such programs is prohibitively expensive.

heading: Reinforcement Learning with AI Feedback.
text: Bai et al. (2022) introduced the concept of Reinforcement Learning with AI Feedback (RLAIF), which combines preferences labeled by LLMs with human-labeled preferences to optimize for helpfulness and harmlessness. Since then, many studies have explored the usefulness of AI-generated feedback as an alternative to expensive human annotations in various tasks. For instance, Luo et al. (2023) proposed WizardMath, which enhances the mathematical reasoning abilities of Llama-2 using AI feedback in the training process. In another work (Zhang et al., 2023), researchers used realworld data along with RLAIF to improve LLMs as medical consultants. Prior research has also explored AI evaluation for improving factual correctness in LLM-generated medical summaries (Mishra et al., 2023). Kwon et al. (2023) explored the usefulness of LLMs in the reward design for RL agents in Ultimatum Game, matrix games, and the DealOrNoDeal negotiation task. Recently, Lee et al. (2023) demonstrated that RLAIF can achieve human-level performance in summarization and helpful and harmless text generation. However, the possibility of using RLAIF to improve the code generation and API usage ability in small models (<1B parameters) is under-explored. We demonstrate that even with a few model parameters, AI feed-
back significantly improves code generation quality over simple fine-tuning baselines. Moreover, we found RLAIF applied on smaller 780M parameter GPT-2-large model outperforms LLaMA-7B fine-tuned models, which has nine times more parameters.

heading: 2 Dataset
text: We applied our proposed method to the Gorilla dataset published by Patil et al. (2023). The Gorilla dataset consists of three parts - HuggingFace, TensorFlow, and PyTorch. In this work, we only focus on the HuggingFace dataset, which is the most extensive among the three, featuring over 925 unique APIs. These APIs belong to 37 different domains (e.g., Multimodal Text-to-Image, Computer Vision Image Classification, Audio Text-to-Speech, etc.), and for each API, there exist ten unique instructions. Each instance of the data contains an instruction (task description), domain, API call (a single code line), explanation (how to solve the task using the API), and a complete code (Python script) to accomplish the task. Here, we highlight some key differences between Gorilla and the traditional code generation datasets. Most of the problem statements and corresponding code snippets present in the benchmark datasets including CodeSearchNet (Husain et al., 2019), XLCoST (Zhu et al., 2022), APPS (Hendrycks et al., 2021) and MBPP (Austin et al., 2021) are related to traditional software engineering tasks, representative of common interview questions, require minimal computational resources to execute and do not require internet connection. On the contrary, the Python scripts in the Gorilla dataset focus on AIrelated tasks and require an internet connection and significant computing resources (storage and processing power) to execute. The scripts are expected to download ML models hosted on HuggingFace, load them in memory, and run inference. So, techniques where code execution or unit test outcomes are treated as feedback (Zhong et al., 2017; Le et al., 2022; Shojaee et al., 2023) become inapplicable.
While Patil et al. (2023) focused only on generating the API call, we demonstrate the effectiveness of our approach both on API call correctness and the ability to use that API in a complete code.

heading: 3 Methodology
text: Our framework follows a similar pipeline to RLHF (Ouyang et al., 2022). However, instead of asking
human annotators to rank the generated responses, we employ a bigger LLM by using a novel prompting strategy. More specifically, for a given instruction and generated code (containing an API call), we ask multiple binary (yes/no) questions that capture different aspects of the generated code (and API call) to determine its quality. Our intuition is, that while generating code from natural language might be still challenging for LLMs, providing binary (yes/no) answers guided by few-shot exemplars is a much easier task. These feedbacks in turn could be aggregated as a preference ground truth to train the reward model in the RLHF (Ouyang et al., 2022) process. Thus our approach eliminates the need for expensive human annotation cost. We describe the proposed framework (Figure 1) in detail.
• Step 1: Training a base model The first step in the pipeline is to fine-tune a language model on the dataset to get a base model. We choose GPT-2-large and train it on the Gorilla dataset using the supervised fine-tuning technique for causal language models. We denote the fine-
tuned model by MSFT . • Step 2: Training a reward model using LLM feedback
Instead of human feedback from annotators, we employed a bigger LLM to generate the labels for the reward model. We realized that human graders while judging the correctness of a response, consider different aspects of the generated output. Based on this intuition, we created multiple prompts (Pi) that ask different questions (Qi) for the same input-output pair. More specifically, we created a set of 8 questions which We feed as prompts to a state-of-the-art language model (GPT3.5) to get a binary response. Each of these questions addresses a different desired quality (free of bugs, correct imports, no undefined variables, correct syntax, etc.) of the output relevant to the task. Step 2 in Figure 1 presents a sample prompt made using one of the questions. The appendix contains the complete list of prompts. As the questions are binary (yes/no) in nature, we simply count the number of yes replies by MGPT 3.5 to score each input-
output pair. More formally, given a task t, generated output o, and question set {Qi} the prompt set is defined as P (t, o) = {Pi | Pi = [Qi, t, o]}. The corresponding score (S) is given as:
S(t, o) =
∑ Pi∈P (t,o) I(MGPT 3.5(Pi) = yes)
|P (t, o)|
where I is the indicator function and MGPT 3.5(Pi) is the reply from MGPT 3.5 for the prompt Pi. We use this score to prepare the training data for Mreward in the following way. For each instruction in the training data, we generate two outputs from MSFT by varying the generation parameters (topk, temperature, etc.). Then they are scored using the method described above and labeled (accept or reject) based on this score. These tuples of {input instruction, accepted output, rejected output} are then combined to form the dataset for Mreward. In the training phase, Mreward learns to classify whether a machine-generated code is acceptable (or not) for a given input instruction. We append a classifier head on top of MSFT and use this as the starting point of Mreward and train for three epochs. • Step 3: Reinforcement Learning
Finally, in the RL step, we fine-tune MSFT using the proximal policy optimization (PPO) algorithm (Schulman et al., 2017). The reward in this step is given by Mreward’s logit scores. We denote our final fine-tuned model by MRL.

heading: 4 Results and Discussions
text: We compute the code generation quality using multiple metrics by comparing the generated output with the ground truth. The reported ROUGE is the average of ROUGE-1, ROUGE-2, ROUGE-L, and ROUGE-sum metrics introduced in (Lin, 2004). CodeBLEU (Ren et al., 2020) was specifically designed for evaluating code synthesis. Ren et al. (2020) defined CodeBLEU as the weighted average of standard BLEU (Papineni et al., 2002), the weighted n-gram match (BLEUweight), the syntactic AST match (Matchast), and the semantic dataflow match (Matchdf). CodeBLEU = α · BLEU + β ·
BLEUweight + γ · Matchast + δ · Matchdf . We set α = β = γ = δ = 0.25 to give equal importance to all the components. The AST sub-tree-matching metric was proposed in (Patil et al., 2023) to capture the correctness of the API calls. In addition to that, we also report the successful execution rate of the generated code (Executability Rate). It is worth noting that running this amount of machinegenerated programs that download and use large AI models is challenging. We created a pipeline to automatically run the machine-generated codes in an isolated environment.
Table 1 compares the performance of the proposed model with MGorilla (finetuned LLaMA-7B) (2023). The results clearly show that the proposed MRL boosts the performance of the supervised fine-tuned MSFT in terms of CodeBLEU (1.6 points abs), AST (0.66% abs) and Executability Rate (4.5% abs). We also note that MRL outperforms the MGorilla despite having only 1/9-th of the parameters. It is also reflected in the Executability Rate of the generated code. Figure 2 shows an instance where our framework helps in fixing a common error present in MGorilla and MSFT generations.

heading: 5 Ethics statement
text: This work adheres to the ethical guidelines and principles set out in the ACM Code of Ethics and followed by the broader research community. The dataset used in this paper was originally collected from public repositories hosted on HuggingFace. The authors are aware of the growing literature on jailbreaking language models to generate unsafe content. We hope the community will use the proposed models responsibly and only for the intended use cases.

heading: 6 Limitations
text: One of the common limitations faced by similar fine-tuned models is the presence of biases inherited from the pre-trained model. We anticipate that the biases present in the chosen base model (GPT-2-large) also exist in the final model MRL which might lead to the generation of biased code comments.
Another limitation of this work is the lack of diversity in programming language. The public dataset we considered contains only Python code. Future work should consider expanding this approach to encompass additional programming languages such as C++, Java, JavaScript, etc. Besides, we have not analyzed the performance between more frequent APIs (head) and infrequent APIs (tail). There might be some scope for improvements by focusing on tail APIs more.
Lastly, the learning methodology applied in this study is offline. Given the rapid evolution and proliferation of machine learning models and the corresponding APIs for specific tasks, the model may not leverage more suitable APIs that emerge posttraining. To address this, periodic updates to the model are necessary. Our framework’s reliance on machine-generated feedback significantly reduces the resource intensity associated with the RLHF process, making these updates more feasible and less costly than a human feedback-based approach.

heading: A Prompts
text: Table 2 lists all the prompts used in accessing the quality of the generated codes.

heading: B Experimental details

heading: B.1 Dataset
text: The HuggingFace part of the Gorilla dataset (Patil et al., 2023) consists of over 9k instruction-output pairs. We trained our model on 90% of the data and kept the rest for evaluation.

heading: B.2 Model and implementation details
text: MSFT, Mreward and MRL all have 780M parameters. While training MSFT and Mreward we used a learning rate of 5×10−4 and 5×10−5 respectively. In the RL step (PPO algorithm), we set the learning rate to 6× 10−6. We did not perform any hyperparameter search. The results are reported by taking the mean of three inference runs. We implemented the training pipeline using the following Python libraries: transformers (Wolf et al., 2020) and TRL (von Werra et al., 2020).

heading: B.3 Computational cost
text: We used a cluster of NVIDIA A100 40GB GPUs for our experiments. We spent in total ∼ 60 GPU hours for all of the experiments.

heading: Future Work
text: Future work should consider expanding this approach to encompass additional programming languages such as C++, Java, JavaScript, etc. Besides, we have not analyzed the performance between more frequent APIs (head) and infrequent APIs (tail). There might be some scope for improvements by focusing on tail APIs more. Lastly, the learning methodology applied in this study is offline. Given the rapid evolution and proliferation of machine learning models and the corresponding APIs for specific tasks, the model may not leverage more suitable APIs that emerge posttraining. To address this, periodic updates to the model are necessary. Our framework’s reliance on machine-generated feedback significantly reduces the resource intensity associated with the RLHF process, making these updates more feasible and less costly than a human feedback-based approach