Applying RLAIF for Code Generation with API -usage in Lightweight LLM s (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., GPT-3.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 model-generated 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.

## introduction llms have demonstrated unprecedented natural language understanding and generation capabilities in recent times @xcite @xcite @xcite . 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. ## dataset we applied our proposed method to the gorilla dataset published by @xcite . 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 code-searchnet @xcite , @xcite , apps @xcite and mbpp @xcite 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 @xcite @xcite 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. ## methodology our framework follows a similar pipeline to rlhf @xcite . however, instead of asking training data prompt given an input task and a python code, determine if the code imports all the necessary classes/modules for execution. task: code: prompt given an input task and a python code, determine if the code imports all the necessary classes/modules for execution. task: code: {instruction, output} pairs prompt given an input task and a python code, determine if the code imports all the necessary classes/modules for execution. task: code: accumulate gpt-3.5 responses over all the prompts to score and rank the generated codes 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 @xcite process. thus our approach eliminates the need for expensive human annotation cost. we describe the proposed framework (figure 1 ) in detail. ## results and discussions figure 2 : example code generated by different models for the same instruction. in the generations of m gorilla and m sft the variable russian_text is undefined and hence will result in an error. whereas m rl defines the variable text before using it. bleu weight + γ • match ast + δ • match df . we set α = β = γ = δ = 0.25 to give equal importance to all the components. the ast sub-tree-matching metric was proposed in @xcite 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. ## ethics statement 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. ## limitations 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 m rl 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. ming zhu, aneesh jain, karthik suresh, roshan ravindran, sindhu tipirneni, and chandan k. reddy. 2022. xlcost: a benchmark dataset for cross-lingual code intelligence. yuchen zhuang, yue yu, kuan wang, haotian sun, and chao zhang. 2024. toolqa: a dataset for llm question answering with external tools. advances in neural information processing systems, 36.