腾讯安全大数据实验室 | 通过代码动静态信息来微调大模型
2023-9-22 18:30:53 Author: mp.weixin.qq.com(查看原文) 阅读量:6 收藏

原文标题:Understanding Programs by Exploiting (Fuzzing) Test Cases  
原文作者:Jianyu Zhao;Yuyang Rong;Yiwen Guo;Yifeng He;Hao Chen
发表会议:In 61st Annual Meeting of the Association for Computational Linguistics (ACL 23)
原文链接:https://arxiv.org/pdf/2305.13592.pdf
主题类型:代码大模型
笔记作者:梅瑞@iFLYTEK
主编:黄诚@安全学术圈

程序语义理解一直是软件工程领域和系统安全领域的热点问题,近年来由AI驱动的代码智能(Code Intelligence)引起了大量关注,特别是代码表示学习(Code Representation Learning)的发展为代码智能奠定了坚实的基础,使得基于AI的方法能够应用到许多下游任务,包括代码相似性检测、代码分类、克隆检测等。

受到大语言模型(LLM)在自然语言理解方面取得成功的启发,一些研究将编程语言视为另一种自然语言,并在程序代码语料库上训练LLM,已经取得了巨大的进展(代表性工作有CodeBERT,CuBERT,CodeGPT,CodeLlama)。然而,程序和文本有着巨大的不同,程序是高度结构化和语法严格的,特别地,程序及其基本单元(即函数和子程序)的输入输出对(Input/Output Pair)代表了函数的语义进而可以表示整个程序的语义。

本文提出将程序及基本单元的输入和可能的输出之间的关系纳入到预训练大语言模型(Pre-trained LLM)的持续训练中,以实现对程序的更深入的语义理解。为了获得有代表性的输入并足以触发不同代码分支的执行,作者提出一种基于模糊测试(Fuzzing)的副产品——输入输出对生成的Prompt来进行LLM调优以提升程序理解能力的方法。

为了激发对程序或函数/子程序的语义和逻辑行为的更深入理解,文章试图利用它们的输入和可能的输出之间的关系,以实现对程序代码的更好理解,直觉上该方式类似于人类工程师理解代码的方式。

文章提出了方法FuzzT(Fuzz Tuning),一种通过提取程序的静态信息和动态信息,对预训练代码大模型进行微调的方法,强化代码大模型对程序语义的理解,并在下游任务中取得了超过现有工作的效果。

2.1 代码静态信息提取:

本文采用此前工作UniXcoder的方法提取程序的静态信息,用以作为代码大模型持续训练的数据之一。

首先生成程序中函数的抽象语法树AST,然后采用如下算法从AST中提取代码序列,以表示代码的静态信息。比起将源代码直接作为输入以训练大模型的方式,基于AST的代码序列可以消除如程序注释等不相关信息在训练中带来的偏差。

2.2 代码动态信息提取:

尽管在软件测试和漏洞挖掘中广泛采用模糊测试Fuzzing,但是在AI下游任务中很少采用。近年来以AFL为代表的基于代码覆盖率的Fuzzing表现出极好的性能,通过使用不同的输入变异来执行程序,并监测每次执行的行为。作为Fuzzing的副产品,一个模糊器可以产生大量的输入,每个输入都会触发被测程序的一个新的行为(输出)。

文章针对四种主流语言(C,C++,Java、Python)实现了基于AFL++的模糊测试工具,分别介绍如下:

  • C/C++

对于C和C++,本文方法把它们看作C++文件编译执行。为了修复程序中的一些语义无关的错误,文章设计了一个编译器插件,自动修改这些错误:

1)缺少头文件:在每个程序的开头,包含最常用C++库的头文件。

2)返回类型和/或参数不正确:例如,一个main函数被定义为"int main()",但是没有提供返回值,那么在程序的末尾自动添加"return 0;"来解决此类问题。

3)标准库中关键字的误用:为每个被误用关键字的变量添加“fixed_”前缀,以避免关键字冲突。

4)不正确的结构体定义:许多结构体在定义时没有在后面添加分号,这里自动在结构体定义后面添加分号。

5)未声明的标志符:许多程序使用静态变量保存常数值,但是该值有时会超出作用范围,这里通过分析常量的用法,并为它们插入定义。

  • Java

对于Java程序,文章使用已有工作Kelinci将其编译成字节码,并对它们进行Fuzzing。由于并非所有程序都命名为Main.java,因此为了编译Java程序,将包含Main类的文件重新命名为Main.java,以便自动编译和Fuzzing。

  • Python

Python是最难处理的语言。首先,许多词法错误不像C/C++和Java那样容易修复。例如,如果程序混合了制表符和空格,则很难推断出预期的缩进是什么。为了解决这个问题,文章使用autopep8转换程序。另一个挑战是Python2和Python3不能轻易区分,因此无法判断使用哪个解释器进行Fuzzing。这里首先使用python3解释程序,若无法执行,则使用2to3工具进行转换。此外,针对Python程序,使用py-afl-fuzz来监测程序的执行。

说明:本节所有工作仅仅是为了构建自动化的模糊器以获得对大模型训练语料代码的大量输入输出对。

2.3 模型

尽管可以使用本文方法提取的静态和动态信息从头开始训练LLM基模型,然而本文选择基于一个预训练的代码大模型实现FuzzT的方法(在本文评估中基于CodeBERT和UniXcoder两个代码大模型),使得本文工作在最多两张V100卡的基础设施中即可以实现。

2.4 提示词生成 Prompting

原始格式的Fuzzing测试用例一般是字节序列,通过编码,可以获得一系列Unicode字符串,这符合LLM最初应用中的自然语言的特征。为了帮助LLM更好的理解这些测试用例,本文引入了此前的工作完形提示(Cloze Prompt),将每一测试用例的输入和输出填入如下格式的Prompt中,借助Prmopt会话所构建的思维链(Chain of Thought,CoT)来调优代码大模型。

  • 数据集

评估数据集采用POJ104CodeNet两种,其中,POJ104包含104个问题(分类),每个问题包含500个C/C++实现代码,从CodeNet中提取Java250、Python800和C++1000三个数据集,详情如下表所示:

  • 代码克隆检测

代码克隆检测任务的目的是测量代码片段或程序的相似度,用于知识产权保护、恶意代码关联分析等。在本评估中,给定一个代码片段或程序的源代码作为查询(Query)即Prompt输入,模型应该输出测试集中语义上类似的实现。

下图是POJ104数据集上的代码相似性检测平均精度,与现有采用传统精调(Fine Tuning,FineT)技术的方法相比,表现最优。

  • 代码分类

代码分类任务要求模型识别出为解决相同问题并实现相同目标的程序,为每个问题的不同实现分配相同的标签。代码分类问题是基于代码语义理解上的相似性问题,本文提出的FuzzT通过将大量输入输出对加入到大模型调优训练,使得大模型有更多机会可以学习到代码的语义。

示例:在将字符串转换为大写字符时,有不同的实现,例如使用"toupper(c)"函数、通过ASCII码偏移量计算"c-'A'+'a'"、静态映射"caseMap[c]"等方式,利用FuzzT的输入输出对可以让模型学到语义的相似性。

在评估中,代码分类的错误率(Error Rate)比现有工作也更优秀,在POJ104数据集上最低只有1.38%。

3.3 案例研究

下图是关于排序算法的案例,可以看到:(a)和(b): 文本相似的函数可以有大相径庭的行为,如降序/升序的快速排序;(b)和(c): 文本完全不同的函数可以有一样的行为,如不同的升序排序算法。因此,Fuzzing带来的动态程序数据可以帮助LLM更好的学习到程序行为上的区别。

  1. 本文首次提出了采用Fuzzing中的副产品——未触发异常的输入输出对作为代码大模型持续训练和调优的训练数据,并构建了输入输出对的Cloze Prompt,增强训练的效果。评估表明,与传统代码大模型仅使用源代码或其他静态信息如AST的方法相比,本文方法提高了大模型对代码语义理解的能力。

  2. 鲁棒性:尽管FuzzT表现的很好,但Fuzzing策略(如AFL-like技术)的优化目标是为了更多的覆盖代码的执行路径,根据Fuzzing对种子样本的变异策略,针对一个程序的两次Fuzzing测试过程可能生成两组完全不同的输入输出对集合,这对于大模型的训练可能带来不确定的影响。

[1] Zhao J, Rong Y, Guo Y, et al. Understanding Programs by Exploiting (Fuzzing) Test Cases[J]. arXiv preprint arXiv:2305.13592, 2023.
[2] Guo D, Lu S, Duan N, et al. Unixcoder: Unified cross-modal pre-training for code representation[J]. arXiv preprint arXiv:2203.03850, 2022.

Jianyu Zhao,来自腾讯安全大数据实验室,研究方向为AI安全。https://sbdlab.tencent.com/

Peter Rong,UC Davis计算机科学专业博士研究生,研究方向包括模糊测试,LLM软件分析,编译器测试等。https://peterrong.netlify.app/

郭怡文,独立研究员,通信作者,研究兴趣包括机器学习、计算机视觉、自然语言处理和网络安全。https://yiwenguo.github.io/

贺一峰,UC Davis计算机科学专业博士研究生,研究兴趣为程序分析、以及它们在LLM代码理解上的应用。https://eyh0602.github.io/

陈浩,UC Davis教授,研究方向包括机器学习安全、软件安全、移动和无线安全等。http://www.cs.ucdavis.edu/~hchen/

安全学术圈招募队友-ing 
有兴趣加入学术圈的请联系 secdr#qq.com

文章来源: https://mp.weixin.qq.com/s?__biz=MzU5MTM5MTQ2MA==&mid=2247489608&idx=1&sn=aff1efae8e9d0733fb91dbc4d6568e5c&chksm=fe2ee7c3c9596ed58dea22a2c7f3efe938e322b3dad3668a8fdd3ad7b1833ccd3b208d3a4def&scene=58&subscene=0#rd
如有侵权请联系:admin#unsafe.sh