CN116643759A - 一种基于程序依赖图预测的代码预训练模型训练方法 - Google Patents
一种基于程序依赖图预测的代码预训练模型训练方法 Download PDFInfo
- Publication number
- CN116643759A CN116643759A CN202310688989.9A CN202310688989A CN116643759A CN 116643759 A CN116643759 A CN 116643759A CN 202310688989 A CN202310688989 A CN 202310688989A CN 116643759 A CN116643759 A CN 116643759A
- Authority
- CN
- China
- Prior art keywords
- dependency graph
- symbol
- code
- training
- program
- Prior art date
- Legal status (The legal status is an assumption and is not a legal conclusion. Google has not performed a legal analysis and makes no representation as to the accuracy of the status listed.)
- Pending
Links
- 238000012549 training Methods 0.000 title claims abstract description 106
- 238000000034 method Methods 0.000 title claims abstract description 29
- 239000013598 vector Substances 0.000 claims abstract description 49
- 238000013528 artificial neural network Methods 0.000 claims abstract description 7
- 230000003068 static effect Effects 0.000 claims abstract description 5
- 230000006870 function Effects 0.000 claims description 18
- 239000011159 matrix material Substances 0.000 claims description 15
- 230000001419 dependent effect Effects 0.000 claims description 12
- 230000014509 gene expression Effects 0.000 claims description 11
- 238000004364 calculation method Methods 0.000 claims description 7
- 238000013507 mapping Methods 0.000 claims description 7
- 230000000873 masking effect Effects 0.000 claims description 5
- 238000006467 substitution reaction Methods 0.000 claims description 2
- 238000012038 vulnerability analysis Methods 0.000 abstract 1
- 238000011156 evaluation Methods 0.000 description 12
- 238000001514 detection method Methods 0.000 description 8
- 230000011218 segmentation Effects 0.000 description 8
- 238000002474 experimental method Methods 0.000 description 7
- 238000012360 testing method Methods 0.000 description 4
- 238000010586 diagram Methods 0.000 description 3
- 230000000694 effects Effects 0.000 description 3
- 238000003058 natural language processing Methods 0.000 description 3
- 230000008569 process Effects 0.000 description 3
- 238000012795 verification Methods 0.000 description 3
- ORILYTVJVMAKLC-UHFFFAOYSA-N Adamantane Natural products C1C(C2)CC3CC1CC2C3 ORILYTVJVMAKLC-UHFFFAOYSA-N 0.000 description 2
- 239000000284 extract Substances 0.000 description 2
- 230000006872 improvement Effects 0.000 description 2
- 238000012545 processing Methods 0.000 description 2
- 230000002829 reductive effect Effects 0.000 description 2
- 238000010200 validation analysis Methods 0.000 description 2
- 238000012935 Averaging Methods 0.000 description 1
- 241001050985 Disco Species 0.000 description 1
- 230000006978 adaptation Effects 0.000 description 1
- 230000004931 aggregating effect Effects 0.000 description 1
- 230000002776 aggregation Effects 0.000 description 1
- 238000004220 aggregation Methods 0.000 description 1
- 230000009286 beneficial effect Effects 0.000 description 1
- 230000000052 comparative effect Effects 0.000 description 1
- 238000010276 construction Methods 0.000 description 1
- 238000013135 deep learning Methods 0.000 description 1
- 230000007547 defect Effects 0.000 description 1
- 238000013461 design Methods 0.000 description 1
- 238000000605 extraction Methods 0.000 description 1
- 230000003993 interaction Effects 0.000 description 1
- 230000000670 limiting effect Effects 0.000 description 1
- 238000005259 measurement Methods 0.000 description 1
- 238000012986 modification Methods 0.000 description 1
- 230000004048 modification Effects 0.000 description 1
- 238000012544 monitoring process Methods 0.000 description 1
- 238000003062 neural network model Methods 0.000 description 1
- 230000036961 partial effect Effects 0.000 description 1
- 230000009467 reduction Effects 0.000 description 1
Classifications
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06F—ELECTRIC DIGITAL DATA PROCESSING
- G06F8/00—Arrangements for software engineering
- G06F8/40—Transformation of program code
- G06F8/41—Compilation
- G06F8/43—Checking; Contextual analysis
- G06F8/433—Dependency analysis; Data or control flow analysis
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06F—ELECTRIC DIGITAL DATA PROCESSING
- G06F11/00—Error detection; Error correction; Monitoring
- G06F11/36—Preventing errors by testing or debugging software
- G06F11/3604—Software analysis for verifying properties of programs
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06F—ELECTRIC DIGITAL DATA PROCESSING
- G06F21/00—Security arrangements for protecting computers, components thereof, programs or data against unauthorised activity
- G06F21/50—Monitoring users, programs or devices to maintain the integrity of platforms, e.g. of processors, firmware or operating systems
- G06F21/57—Certifying or maintaining trusted computer platforms, e.g. secure boots or power-downs, version controls, system software checks, secure updates or assessing vulnerabilities
- G06F21/577—Assessing vulnerabilities and evaluating computer system security
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06F—ELECTRIC DIGITAL DATA PROCESSING
- G06F8/00—Arrangements for software engineering
- G06F8/40—Transformation of program code
- G06F8/41—Compilation
- G06F8/43—Checking; Contextual analysis
- G06F8/436—Semantic checking
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06N—COMPUTING ARRANGEMENTS BASED ON SPECIFIC COMPUTATIONAL MODELS
- G06N3/00—Computing arrangements based on biological models
- G06N3/02—Neural networks
- G06N3/04—Architecture, e.g. interconnection topology
- G06N3/047—Probabilistic or stochastic networks
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06N—COMPUTING ARRANGEMENTS BASED ON SPECIFIC COMPUTATIONAL MODELS
- G06N3/00—Computing arrangements based on biological models
- G06N3/02—Neural networks
- G06N3/04—Architecture, e.g. interconnection topology
- G06N3/048—Activation functions
-
- G—PHYSICS
- G06—COMPUTING; CALCULATING OR COUNTING
- G06N—COMPUTING ARRANGEMENTS BASED ON SPECIFIC COMPUTATIONAL MODELS
- G06N3/00—Computing arrangements based on biological models
- G06N3/02—Neural networks
- G06N3/08—Learning methods
- G06N3/0985—Hyperparameter optimisation; Meta-learning; Learning-to-learn
-
- Y—GENERAL TAGGING OF NEW TECHNOLOGICAL DEVELOPMENTS; GENERAL TAGGING OF CROSS-SECTIONAL TECHNOLOGIES SPANNING OVER SEVERAL SECTIONS OF THE IPC; TECHNICAL SUBJECTS COVERED BY FORMER USPC CROSS-REFERENCE ART COLLECTIONS [XRACs] AND DIGESTS
- Y02—TECHNOLOGIES OR APPLICATIONS FOR MITIGATION OR ADAPTATION AGAINST CLIMATE CHANGE
- Y02D—CLIMATE CHANGE MITIGATION TECHNOLOGIES IN INFORMATION AND COMMUNICATION TECHNOLOGIES [ICT], I.E. INFORMATION AND COMMUNICATION TECHNOLOGIES AIMING AT THE REDUCTION OF THEIR OWN ENERGY USE
- Y02D10/00—Energy efficient computing, e.g. low power processors, power management or thermal management
Landscapes
- Engineering & Computer Science (AREA)
- Theoretical Computer Science (AREA)
- General Engineering & Computer Science (AREA)
- Physics & Mathematics (AREA)
- Software Systems (AREA)
- General Physics & Mathematics (AREA)
- Computer Hardware Design (AREA)
- Computing Systems (AREA)
- Computational Linguistics (AREA)
- Life Sciences & Earth Sciences (AREA)
- Evolutionary Computation (AREA)
- Artificial Intelligence (AREA)
- Biomedical Technology (AREA)
- Biophysics (AREA)
- Computer Security & Cryptography (AREA)
- Data Mining & Analysis (AREA)
- Health & Medical Sciences (AREA)
- General Health & Medical Sciences (AREA)
- Molecular Biology (AREA)
- Mathematical Physics (AREA)
- Quality & Reliability (AREA)
- Probability & Statistics with Applications (AREA)
- Devices For Executing Special Programs (AREA)
Abstract
本发明公开了一种基于程序依赖图预测的代码预训练模型训练方法,该方法通过在预训练阶段预测代码的程序依赖图,帮助预训练模型感知并捕获代码的语义结构特征,以提升对代码语义结构敏感的下游任务的性能。该方法首先利用静态程序分析工具提取代码的程序依赖图并据此构建预训练监督信息,然后使用神经网络编码器将代码编码为实值特征向量序列,并基于该特征向量序列预测代码的程序依赖图,最后根据预测结果与实际程序依赖图的差异来更新模型的权重。预训练的模型可直接被用于程序依赖分析,依赖GPU等并行处理器对海量代码快速进行程序依赖分析,也可在微调后用于对代码语义结构敏感的下游任务,如漏洞分析相关任务等。
Description
技术领域
本发明涉及深度学习和软件工程领域,尤其涉及一种基于程序依赖图预测的代码预训练模型训练方法。
背景技术
近年来,预训练模型大幅提升了各类自然语言处理(NLP)任务的效果,具有很强的泛化性和通用性。预训练模型通常利用多个预训练任务在海量无监督数据上学习通用知识,然后利用下游任务的有监督数据进行微调(Fine-tuning),以快速适应下游任务。鉴于预训练模型在NLP任务上的优异性能,一批在海量代码数据上预训练的模型被提出,例如CodeBERT、GraphCodeBERT等,并被用于多种代码相关的任务,例如代码搜索、代码克隆检测、代码补全和代码摘要等。这些模型被称为代码预训练模型。
预训练模型的性能与预训练任务密切相关。有效的预训练任务能够指导模型从海量数据中学习通用的、有助于下游任务的知识。一些代码预训练模型直接沿用为自然语言设计的预训练任务,例如替换单词预测(RTD)。而区别于自然语言,代码具有强结构性。因此,一些代码预训练模型提出了考虑代码结构信息的预训练任务以提升模型效果,例如符号类型推断(Fine-Grained Type Inference)、AST节点类型预测、数据流边预测等。然而这些预训练模型要么仅考虑代码的句法结构,忽略了代码的语义结构(例如代码中的控制依赖和数据依赖),从而在下游任务上性能受限;要么需要显式地提供由代码分析工具提取的结构信息作为输入,而非仅以源代码作为输入,限制了其通用性,且无法被用于分析代码结构信息。目前尚无在预训练阶段考虑代码语义结构,且仅以源代码作为输入的代码预训练模型。本发明通过学习代码的语义结构知识,预训练模型有望提升对代码语义结构敏感的下游任务的性能。
发明内容
本发明的目的在于针对现有技术的不足,提供了一种基于程序依赖图预测的代码预训练模型训练方法。
本发明的目的是通过以下技术方案来实现的:一种基于程序依赖图预测的代码预训练模型训练方法,包括如下步骤:
(1)抽取代码数据的程序依赖图作为预训练监督信息,所述程序依赖图包括控制依赖图和数据依赖图;
(2)通过分词器将代码分词为符号序列T,其中T=[tcls,t1,...,tn],tcls为插入在序列最前端的符号,n为符号序列T的长度;
(3)基于步骤(2)得到的符号序列T将程序依赖图转换为语句级控制依赖图和符号级数据依赖图;
(4)通过编码器将分词后的符号序列T编码为符号特征向量序列Ht,其中为插入在序列最前端的符号tcls对应的向量,n为符号序列Ht的长度;
(5)将符号特征向量序列输入到语句级控制依赖图预测任务、符号级数据依赖图预测任务和掩码语言模型任务中,并汇总各个任务的损失值以获取预训练阶段总目标
(6)使用优化器优化总目标并更新编码器和神经网络组件的参数。
进一步地,所述步骤(1)具体为:使用静态程序分析工具分析用于预训练的代码数据,获取每段代码的程序依赖图;通过提取程序依赖图中的所有AST节点和指示控制依赖的边,获取代码的控制依赖图,再通过提取程序依赖图中所有AST节点和指示数据依赖的边,获取代码的数据依赖图。
进一步地,所述步骤(3)中包括以下子步骤:
(3.1)对于控制依赖图,对步骤(1)输出的控制依赖图中的每一条边,提取其起止AST节点对应的代码元素的具体位置,然后以AST节点对应的代码行号作为代码语句序号,将AST节点映射为对应的代码语句,从而将原AST节点的控制依赖关系转化为代码语句之间的控制依赖关系,进而构建出语句级控制依赖图,并存储为控制依赖图矩阵
(3.2)对于数据依赖图,对步骤(1)输出的数据依赖图的中每一条边,提取其起止AST节点对应的标识符;然后通过检查标识符对应的代码字符区间与分词后符号对应的代码字符区间的重叠情况,将AST节点映射为符号序列T中的符号;最后,基于上述符号间的数据依赖关系,构建符号级数据依赖图,并存储为数据依赖图矩阵
进一步地,所述步骤(5)中,预训练总目标由符号特征向量序列Ht计算得到,并定义为语句级控制依赖图预测任务、符号级数据依赖图预测任务和掩码语言模型任务损失值的加权和,有如下表达式:
其中和/>分别代表语句级控制依赖图预测任务、符号级数据依赖图预测任务和掩码语言模型任务产生的损失值,a1、a2和a3分别为语句级别控制依赖图预测任务、符号级数据依赖图预测任务和掩码语言模型任务的损失值在总目标/>中的权重。
进一步地,所述步骤(5)包括如下子步骤:
(5.1)语句级控制依赖图预测任务损失值的计算:首先根据代码符号特征向量序列Ht计算语句级控制依赖图中每个节点qk的特征向量然后使用双线性层预测节点qi控制依赖于另一节点qj的概率/>最后使用交叉熵损失函数来度量预测的概率值/>与语句级控制依赖图矩阵Gc中对应位置的标签值/>的差异,并生成语句级代码控制依赖图预测任务的损失值/>有如下表达式:
(5.2)符号级数据依赖图预测任务损失值的计算:首先根据代码符号特征向量序列Ht计算符号级数据依赖图中每个节点ti的特征然后使用双线性层预测符号ti数据依赖于另一个符号tj的概率/>最后使用经过掩码的交叉熵损失函数来度量预测的概率值与符号级数据依赖图中对应位置的标签值/>的差异,并生成符号级数据依赖图预测任务的损失值/>有如下表达式:
其中,每个符号ti对应的掩码值基于语法分析器对代码符号的类型分析结果来计算,属于标识符类型的代码符号的掩码值mi d=1,而不属于标识符类型的符号的掩码值mi d=0;
(5.3)掩码语言模型任务损失值的计算:首先将输入中的部分符号随机替换为特殊的掩码符号[MASK],然后利用编码器将随机掩码的输入编码为特征向量序列,最后基于掩码位置的特征向量预测输入中被掩码的符号;从分词后符号序列T中随机抽取比例为rsample的符号,这些被抽取到的符号中,rmask%的符号被替换为[MASK]符号,rreplace%的符号被替换为词汇表中的一个随机符号;然后将掩码后的符号序列输入到编码器并获得符号特征向量序列Ht,基于符号特征向量序列Ht使用一个多层感知器模型来预测掩码前的原符号,并根据预测结果与原符号的差异计算交叉熵损失作为掩码语言模型任务损失值有如下表达式:
其中,V为模型的词汇表,表示符号ti预测正确的概率,Mt表示被随机采样的下标集合。
本发明的有益效果是:本发明在预训练阶段通过显式预测代码的程序依赖图引导模型学习代码语义结构相关的知识,能提升对代码语义结构敏感的下游任务的性能。同时该预训练模型在迁移到下游任务时无需利用代码分析工具来抽取代码的结构信息,避免了代码分析工具的通用性问题对预训练模型的通用性造成影响。预训练完成以后,模型不仅能够进行微调以适应下游任务,还能直接被用作一个高效的、基于神经网络的程序依赖分析工具,依赖GPU等并行处理器对海量代码快速进行程序依赖分析。
附图说明
图1是本发明方法的流程图;
图2是本发明中将代码的程序依赖图与分词得到的代码符号序列建立联系的示例图;
图3是本发明中使用的预训练任务的示意图;
图4是本发明中基于语法分析器生成用于缓解符号级数据依赖矩阵稀疏问题的符号掩码值的示例图;
图5是本发明实验评估中内在评估的结果表格;
图6是本发明实验评估中外在评估的结果表格。
具体实施方式
本发明通过静态程序分析工具提取代码的程序依赖图,在预训练阶段预测程序依赖图,计算输出的程序依赖图与实际程序依赖图的差异,获取程序依赖图预测任务的损失值,最后汇总所有预训练任务的损失值,获得预训练总损失值(即总目标/>),并使用优化器通过最小化预训练总损失值/>来更新模型参数。
程序依赖图展现了代码元素之间的控制依赖和数据依赖关系,反映了代码的语义结构信息。程序依赖图中包含控制依赖边和数据依赖边。控制依赖边展现了条件断言与代码语句之间的控制流关系,可用于推断代码语句执行的控制流条件;数据依赖边展现了代码中的定义-使用(def-use)关系。
以下结合附图和实例对本发明进行进一步介绍。
本发明提供的一种基于程序依赖图预测的代码预训练模型训练方法,参见图1,包括以下步骤:
(1)抽取代码数据的程序依赖图(Program Dependence Graph)作为预训练阶段程序依赖图预测任务的监督信息。本发明首先利用静态程序分析工具抽取预训练代码数据集中代码的程序依赖图作为标签,为预训练阶段的程序依赖图预测任务提供监督信息;程序分析工具的输出为一张图,图中的节点为抽象语法树(Abstract Syntax Tree,AST)节点,图中的边具有不同的类型。通过提取程序依赖图所有节点和控制依赖边,可获取定义在AST节点上的控制依赖图;由于数据依赖指变量之间存在定义-引用(def-use)关系,因此通过提取所有节点并合并图中类型为“def”和“use”的边,能够获取定义在AST节点上的数据依赖图。
(2)代码分词。输入模型的代码本质是一段文本;为了使神经网络能够处理代码,首先需要将代码转换为符号序列。本发明采用BPE算法(Byte-Pair Encoding)构建预训练模型的词汇表和分词器,并利用该分词器将输入代码转换为符号序列,以消除”词汇表外问题”(Out-of-Vocabulary Problem),即消除测试数据中存在而训练数据中不存在的符号。分词后获得的符号序列被定义为T=[tcls,t1,...,tn],其中tcls是一个插入在序列最前端的特殊符号,其被模型处理后获得的特征向量常被用作整个代码的表示,n为代码符号序列的长度。另外,由于分词不会删减代码中的字符,因此每个符号都对应于代码中一个字符区间。
(3)构建程序依赖图矩阵。参见图3,本发明将基于符号序列T预测程序依赖图作为预训练任务。但步骤(1)输出的程序依赖图定义在AST节点上,没有与输入神经网络模型的符号序列T进行关联。因此,本步骤需要基于步骤(1)的输出在步骤(2)输出的符号序列上重新构建控制依赖图和数据依赖图。该步骤主要分为以下两个子步骤:
(3.1)对于控制依赖图,由于代码的控制依赖关系通常定义在代码语句之间,因此将重新构建的控制依赖图中的节点qk定义为代码语句。首先,对步骤(1)输出的控制依赖图中的每一条边,提取其起止AST节点对应的代码元素的具体位置(代码行和对应字符区间)。然后以AST节点对应的代码行号作为代码语句序号将AST节点映射为对应的代码语句(代码行),从而将原AST节点的控制依赖关系转化为代码语句之间的控制依赖关系,进而构建出语句级别的控制依赖图,并存储为控制依赖图矩阵其中/>代表第i个代码语句与第j个代码语句之间存在控制依赖,/>代表不存在控制依赖,mc为控制依赖图中的节点数量。最后,由Gc的定义可知,其中每一行/>和每一列/>表示了某代码语句与其他语句之间的控制依赖关系,因此对于符号序列T中的每一个符号ti,通过检查其所在的代码行来将其对应到某一个语句中,实现符号序列T与控制依赖图矩阵Gc的关联。
(3.2)对于数据依赖图,由于数据依赖关系通常存在于变量与变量之间,因此直接定义在代码语句之间的数据依赖关系粒度过粗,可能导致模型无法准确捕获和学习代码中的数据依赖。为此,本发明将代码数据依赖关系定义在符号级,并构建定义在符号序列上的符号级数据依赖图。具体而言,首先,对步骤(1)输出的数据依赖图的中每一条边,提取其起止AST节点对应的标识符。然后通过检查标识符对应的代码字符区间与分词后符号对应的代码字符区间的重叠情况,将AST节点映射为符号序列T中的符号。由于BPE算法可能将一个标识符分割为多个符号,因此一个标识符对应的AST节点可能被映射为多个(一组)符号。以图2为例,代码中的第一个“temp_flag”标识符对应的字符区间为[4,12],通过检查字符区间的重叠情况,其被映射为三个字符,即“_”和“flag”,这三个字符分别对应代码区间[3,7]、[8,8]和[9,12]。由此,AST节点之间的数据依赖关系被转换为一组符号与另一组符号之间的数据依赖关系。为了避免这种“多对多”关系产生大量的无意义的数据依赖边,本发明仅保留数据依赖边起止AST节点对应的两组符号中的第一个,并丢弃掉其他符号。接前述的图2的例子,两个“temp_flag”标识符之间的数据依赖关系被简化为两个标识符的第一个符号/>(符号下标分别为1和13)之间的数据依赖关系。最后,基于上述符号间的数据依赖关系,构建符号级数据依赖图并存储为矩阵/>其中/>代表ti与tj之间存在数据依赖关系,反之/>代表不存在,md为符号序列的长度。
(4)利用编码器将符号序列编码为符号特征向量序列。在完成分词之后,输入的代码被转换为符号序列T,该序列将被输入到神经网络编码器中以获取每个符号的特征向量。后续的预训练任务将基于符号的特征向量序列进行。编码器一般由一个嵌入层和一个特征提取组件组成,嵌入层负责基于词汇表将符号映射为词嵌入(实值向量),特征提取组件则基于符号的词嵌入序列进行符号间的信息交互,以捕获每个符号的特征向量。具体地,符号序列T=[tcls,t1,...,tn]输入到编码器后获得符号特征向量序列被记为其中每个符号被表示为一个向量/>d为特征向量的维度,/>为插入在序列最前端的符号tcls对应的向量,n为符号序列Ht的长度。本发明不对编码器的结构做具体限制,但该编码器至少需具备将代码的符号序列T转换为符号特征向量序列Ht的功能。
(5)基于符号特征向量序列进行包括程序依赖图预测在内的多个预训练任务,计算各个预训练任务的损失值并汇总以获取预训练总损失值(预训练阶段总目标)。该步骤将符号特征向量序列输入到各个预训练任务中,生成各个预训练任务的损失值。本发明中采用的预训练任务共三个:语句级控制依赖图预测任务、符号级数据依赖图预测任务和掩码语言模型任务,这三个预训练任务产生的损失值分别记为和/>预训练总损失值记为/>有以下表达式:
其中a1、a2和a3为超参数,分别用于调整语句级控制依赖图预测任务、符号级数据依赖图预测任务和掩码语言模型任务的损失值在预训练最终损失值中所占的权重。这一组超参数的值通过验证集进行选择。三个预训练任务损失值的计算的具体步骤分别为:
(5.1)对于语句级控制依赖图预测任务,首先需要根据符号特征向量序列Ht构建代码语句的特征表示,用于预测语句间的控制依赖关系。在步骤(3.1)中提到过每个符号都通过检索其所在的代码行对应到语句级控制依赖图中的某一个节点qk上,因此一个节点中通常包含多个符号。本步骤中通过检索并聚合每一个节点qk中包含的所有符号的特征向量来构建节点(代码语句)qk的特征向量
其中|qk|是节点qk(代码语句)中包含的符号数量,MLP是一个多层感知器。MLP主要用于特征降维,因为编码器输出的符号特征向量的维度通常很高(例如代码预训练模型CodeBERT中此维度为768),因此直接使用作为qk的特征向量会消耗大量的内存资源。另外,虽然公式中取节点内所有符号的特征向量的平均值作为节点的特征向量,但本发明并不对此处做限制,任何有序或无序集合的聚合函数均可用于此处。获得代码语句的特征向量后,本发明利用一个双线性层以预测节点(代码语句)qi控制依赖于另一节点(代码语句)qj的概率/>
其中Wc和bc为可学习的参数,σ为Sigmoid函数。此处采用双线性层是因为其能够很好地建模两个输入的关系,且双线性层的两个输入不满足交换律,这与代码中控制依赖关系的有向性保持一致。最后,采用交叉熵损失来度量预测的概率值与控制依赖图矩阵Gc中相应的标签/>的差异,并对获得的结果取平均以生成损失值/>
(5.2)对于符号级数据依赖图预测任务,由于Ht本身就是符号级别的特征向量序列,因此无需再像步骤(5.1)中一样为符号额外构建特征表示。类似地,此处使用一个MLP为符号特征降维,然后使用一个双线性层根据降维后的符号表示预测一个符号ti数据依赖于另一个符号tj的概率值有以下表达式:
其中Wd和bd为双线性层的可学习参数,σ为Sigmoid函数。然而不同的是,一段代码的符号数一般远大于其语句数,这就导致符号级数据依赖图的节点数远大于语句级控制依赖图的节点数,需要预测的矩阵也更大。例如,对于一段包含512个符号、30个代码语句的代码,其需要预测的语句级控制依赖图中包含的元素数量为30×30=900,然而符号级数据依赖图中包含的元素数量却达到了512×512=262144。更糟糕的是,很多代码中仅有几十条乃至几条数据依赖关系,这导致需要预测的符号级数据依赖图高度稀疏,难以直接学习。为此,本发明设计了一种符号掩码方法,将部分符号级数据依赖图中不可能存在数据依赖的节点进行掩码,从而不预测这些位置的符号级数据依赖情况,以缓解数据稀疏问题。具体地,本发明首先基于程序词法分析器将代码分割一个词法符号序列,且词法分析器会为每一个词法符号标注其类型;然后从该序列中筛选出类型为标识符的符号,组成一个词法符号集合S';最后,基于字符区间的重合情况检查符号序列T中的符号ti是否属于词法符号集合S'中任意一个词法符号,为符号ti生成一个掩码值mi:若符号ti属于一个标识符,则mi=1,否则mi=0。本发明利用交叉熵损失度量预测的概率值与实际的符号级数据依赖图中对应位置的标签值/>的差异,并按照未被掩码的元素数量取平均以计算符号级数据依赖损失值/>时,有以下表达式:
其中,如果一对符号中任意一个符号的掩码值为0(该符号不属于标识符),那么这一对符号之间的数据依赖关系的预测损失值就不会被汇总到中,从而减小了需要预测的矩阵大小。如图4所示,这种符号掩码方法使需要预测的矩阵元素减少了70%,大大缓解了矩阵稀疏的问题。
(5.3)掩码语言模型任务最早由BERT模型提出,后被广泛应用于预训练模型中。该任务首先将输入中的部分符号随机替换为特殊的掩码符号[MASK],然后利用编码器将随机掩码的输入编码为特征向量序列,最后基于掩码位置的特征向量预测输入中被掩码的符号。本发明沿用了已有预训练模型中掩码语言模型的设定。具体地,在完成分词后,该任务从分词后符号序列T中随机抽取比例为rsample的符号。这些被抽取到的符号中,rmask%的符号被替换为[MASK]符号,rreplace%的符号被替换为词汇表中的一个随机符号。在将掩码后的符号序列输入到编码器并获得符号特征向量序列Ht之后,该任务基于符号特征向量序列Ht使用一个多层感知器模型来预测掩码前的原符号,并根据预测结果与原符号的差异计算交叉熵损失作为掩码语言模型任务损失值
其中V为模型的词汇表,表示符号ti预测正确的概率,Mt表示被随机采样的下标集合。
(6)使用优化器基于预训练总损失值更新模型参数。在步骤(5)获得预训练阶段的总损失值之后,使用优化器基于该损失值来端到端更新所有模型的参数,包括编码器模型的参数和三个预训练任务的神经网络组件的参数(例如双线性层)。经过预训练后,只有编码器模型将会作为预训练模型被迁移到下游任务上使用,而预训练任务的任务结构(例如双线性层)不会被迁移。
实验验证:
本发明通过两种方式对所提出的训练方法的效果进行了评估:首先验证了经预训练的模型预测程序依赖图的性能,称为内在评估;然后验证了预训练模型经过微调后在代码相关的下游任务上的性能,称为外在评估。
下面分别阐述预训练、内在评估和外在评估的具体实验方法和设定。
预训练:实验使用的预训练数据集包含超过2280000个从开源项目中收集的C/C++函数。该数据集被划分为训练集、验证集和测试集三部分,分别包含约1900000、155000和60400个函数。训练集用于模型预训练,见前述本发明方法的步骤(1)-步骤(6);验证集用于监测训练过程中模型的泛化能力;测试集用于在内在评估中评估模型预测程序依赖图的性能。在评估实验中,本发明提出的预训练模型采用了CodeBERT模型的模型结构、分词器和字典,并使用CodeBERT的权重初始化模型。在预训练过程中,本发明使用Adam作为优化器,其初始学习率设置为0.0001,并需要指定3个参数,分别为:代码的最大行数(语句级控制依赖图中的最大节点数量)mc、分词词后符号序列的最大长度md和训练周期数。本发明涉及的实验中设置mc=50、md=512,训练周期数为10。
内在评估:使用预训练阶段划分的测试集进行实验。评价指标为控制依赖边的F度量值Fcdp、数据依赖边的F度量值Fddp和所有程序依赖边整体的F度量值Foverall。F度量值基于公式:
其中Precision为精准率,Recall为召回率。对于语句级控制依赖图预测任务:TP(真正例)表示真实标签为1且预测结果也为1的语句对数量;FP(假正例)表示真实标签为0但预测结果为1的语句对数量;TN(真反例)表示真实标签为0且预测结果也为0的语句对数量;FN(假反例)表示真实标签为1但预测结果为1的语句对数量。对于符号级数据依赖图预测任务,TP(真正例)表示真实标签为1且预测结果也为1的符号对数量;FP(假正例)表示真实标签为0但预测结果为1的符号对数量;TN(真反例)表示真实标签为0且预测结果也为0的符号对数量;FN(假反例)表示真实标签为1但预测结果为1的符号对数量。对于所有程序依赖边整体,其TP为语句级控制依赖图预测任务的TP值和符号级数据依赖图预测任务TP值之和,其TN、FP和FN指标的计算方式同理。F度量值的范围从0到1,指标值越大表示性能越好。
图5是本发明的内在评估实验结果表格。实验结果显示,采用本发明提出的训练方法预训练的代码模型在预测语句级控制依赖图时能够取得超过99%的F度量值,在预测符号级数据依赖图时能够取得超过93%的F度量值,且程序依赖图整体的F度量值超过了95%,说明本发明提出的预训练方法能够比较精准地预测代码的程序依赖图。
外在评估:本发明所涉及的外在评估实验选择漏洞检测作为下游任务,并且使用了三个现有公开的基准数据集进行验证:ReVeal、Devign和Big-Vul,分别包含约22600、272000和188600个C/C++函数,其中漏洞函数的比例分别为9.7%、5.8%和45.6%。模型在经过预训练后还需在漏洞检测数据集上进行参数微调后方能适应漏洞检测任务。微调的具体方式为:将一个MLP分类器拼接到代码预训练模型之后。该分类器将会以模型输出的第一个符号的表示(即)作为输入来预测代码中存在漏洞的概率。预测概率值大于0.5的样本被预测为包含漏洞的函数。然后,使用交叉熵损失函数度量输出概率与该代码实际是否存在漏洞的标签值的差异以产生微调损失,最后使用Adam优化器联合更新预训练模型和分类器的参数。外在评估的评价指标为漏洞分类的准确率和F度量值。F度量值见内在评估中的描述,其中:TP(真正例)表示存在漏洞的函数被正确预测的函数数量;FP(假正例)表示不存在漏洞的被错误地预测为存在漏洞的函数数量;TN(真反例)表示不存在漏洞的函数被正确预测为不存在漏洞的函数数量;FN(假反例)表示存在漏洞的函数被错误地预测为不存在漏洞的函数数量。准确率accuracy基于以下公式:
采用本发明提出的训练方法以及上述预训练设置预训练的代码模型被称为PDBERT,因为其预训练任务中不仅涉及程序依赖预测(Program Dependence Prediction),而且同时采用了BERT模型提出的掩码语言模型任务。为了进一步验证本发明的有效性,本发明将PDBERT与主流的漏洞检测算法和代码预训练模型进行了对比,包含:5个不基于预训练的漏洞检测模型,即Bi-LSTM、Transformer、VulDeePecker、Devign和ReVeal;3个代码预训练模型,即CodeBERT、GraphCodeBERT和DISCO。图6是本发明的外在评估结果表格。实验结果显示,相比于5个未经过预训练的漏洞检测模型,PDBERT在经过微调后能够取得最佳的性能,在三个基准数据集的对应指标上分别获得平均80.0%、197.1%和24.7%的性能提升。相比于3个代码预训练模型,经过微调后能够取得最佳的性能,在三个基准数据集的对应指标上分别获得平均7.0%、9.5%和6.8%的提升。可以看到,PDBERT在漏洞检测任务的三个数据集上都明显优于其他的对比模型,这表明本发明提出的基于程序依赖图预测的预训练方法是有效的。
最后说明的是,以上实施例仅用以说明本发明的技术方案而非限制,尽管参照较佳实施例对本发明进行了详细说明,本领域的普通技术人员应当理解,可以对本发明的技术方案进行修改或者等同替换,而不脱离本发明技术方案的宗旨和范围,其均应涵盖在本发明的权利要求范围当中。
Claims (5)
1.一种基于程序依赖图预测的代码预训练模型训练方法,其特征在于,包括如下步骤:
(1)抽取代码数据的程序依赖图作为预训练监督信息,所述程序依赖图包括控制依赖图和数据依赖图;
(2)通过分词器将代码分词为符号序列T,其中T=[tcls,t1,...,tn],tcls为插入在序列最前端的符号,n为符号序列T的长度;
(3)基于步骤(2)得到的符号序列T将程序依赖图转换为语句级控制依赖图和符号级数据依赖图;
(4)通过编码器将分词后的符号序列T编码为符号特征向量序列Ht,其中 为插入在序列最前端的符号tcls对应的向量,n为符号序列Ht的长度;
(5)将符号特征向量序列输入到语句级控制依赖图预测任务、符号级数据依赖图预测任务和掩码语言模型任务中,并汇总各个任务的损失值以获取预训练阶段总目标
(6)使用优化器优化总目标并更新编码器和神经网络组件的参数。
2.根据权利要求1所述的基于程序依赖图预测的代码预训练模型训练方法,其特征在于,所述步骤(1)具体为:使用静态程序分析工具分析用于预训练的代码数据,获取每段代码的程序依赖图;通过提取程序依赖图中的所有AST节点和指示控制依赖的边,获取代码的控制依赖图,再通过提取程序依赖图中所有AST节点和指示数据依赖的边,获取代码的数据依赖图。
3.根据权利要求2所述的基于程序依赖图预测的代码预训练模型训练方法,其特征在于,所述步骤(3)中包括以下子步骤:
(3.1)对于控制依赖图,对步骤(1)输出的控制依赖图中的每一条边,提取其起止AST节点对应的代码元素的具体位置,然后以AST节点对应的代码行号作为代码语句序号,将AST节点映射为对应的代码语句,从而将原AST节点的控制依赖关系转化为代码语句之间的控制依赖关系,进而构建出语句级控制依赖图,并存储为控制依赖图矩阵
(3.2)对于数据依赖图,对步骤(1)输出的数据依赖图的中每一条边,提取其起止AST节点对应的标识符;然后通过检查标识符对应的代码字符区间与分词后符号对应的代码字符区间的重叠情况,将AST节点映射为符号序列T中的符号;最后,基于上述符号间的数据依赖关系,构建符号级数据依赖图,并存储为数据依赖图矩阵
4.根据权利要求3所述的基于程序依赖图预测的代码预训练模型训练方法,其特征在于,所述步骤(5)中,预训练总目标由符号特征向量序列Ht计算得到,并定义为语句级控制依赖图预测任务、符号级数据依赖图预测任务和掩码语言模型任务损失值的加权和,有如下表达式:
其中和/>分别代表语句级控制依赖图预测任务、符号级数据依赖图预测任务和掩码语言模型任务产生的损失值,a1、a2和a3分别为语句级别控制依赖图预测任务、符号级数据依赖图预测任务和掩码语言模型任务的损失值在总目标/>中的权重。
5.根据权利要求4所述的基于程序依赖图预测的代码预训练模型训练方法,其特征在于,所述步骤(5)包括如下子步骤:
(5.1)语句级控制依赖图预测任务损失值的计算:首先根据代码符号特征向量序列Ht计算语句级控制依赖图中每个节点qk的特征向量然后使用双线性层预测节点qi控制依赖于另一节点qj的概率/>最后使用交叉熵损失函数来度量预测的概率值/>与语句级控制依赖图矩阵Gc中对应位置的标签值/>的差异,并生成语句级代码控制依赖图预测任务的损失值/>有如下表达式:
(5.2)符号级数据依赖图预测任务损失值的计算:首先根据代码符号特征向量序列Ht计算符号级数据依赖图中每个节点ti的特征然后使用双线性层预测符号ti数据依赖于另一个符号tj的概率/>最后使用经过掩码的交叉熵损失函数来度量预测的概率值/>与符号级数据依赖图中对应位置的标签值/>的差异,并生成符号级数据依赖图预测任务的损失值/>有如下表达式:
其中,每个符号ti对应的掩码值mi d∈{0,1}基于语法分析器对代码符号的类型分析结果来计算,属于标识符类型的代码符号的掩码值mi d=1,而不属于标识符类型的符号的掩码值mi d=0;
(5.3)掩码语言模型任务损失值的计算:首先将输入中的部分符号随机替换为特殊的掩码符号[MASK],然后利用编码器将随机掩码的输入编码为特征向量序列,最后基于掩码位置的特征向量预测输入中被掩码的符号;从分词后符号序列T中随机抽取比例为rsample的符号,这些被抽取到的符号中,rmask%的符号被替换为[MASK]符号,rreplace%的符号被替换为词汇表中的一个随机符号;然后将掩码后的符号序列输入到编码器并获得符号特征向量序列Ht,基于符号特征向量序列Ht使用一个多层感知器模型来预测掩码前的原符号,并根据预测结果与原符号的差异计算交叉熵损失作为掩码语言模型任务损失值有如下表达式:
其中,V为模型的词汇表,表示符号ti预测正确的概率,Mt表示被随机采样的下标集合。
Priority Applications (1)
Application Number | Priority Date | Filing Date | Title |
---|---|---|---|
CN202310688989.9A CN116643759A (zh) | 2023-06-12 | 2023-06-12 | 一种基于程序依赖图预测的代码预训练模型训练方法 |
Applications Claiming Priority (1)
Application Number | Priority Date | Filing Date | Title |
---|---|---|---|
CN202310688989.9A CN116643759A (zh) | 2023-06-12 | 2023-06-12 | 一种基于程序依赖图预测的代码预训练模型训练方法 |
Publications (1)
Publication Number | Publication Date |
---|---|
CN116643759A true CN116643759A (zh) | 2023-08-25 |
Family
ID=87643378
Family Applications (1)
Application Number | Title | Priority Date | Filing Date |
---|---|---|---|
CN202310688989.9A Pending CN116643759A (zh) | 2023-06-12 | 2023-06-12 | 一种基于程序依赖图预测的代码预训练模型训练方法 |
Country Status (1)
Country | Link |
---|---|
CN (1) | CN116643759A (zh) |
Cited By (1)
Publication number | Priority date | Publication date | Assignee | Title |
---|---|---|---|---|
CN117573096A (zh) * | 2024-01-17 | 2024-02-20 | 合肥综合性国家科学中心人工智能研究院(安徽省人工智能实验室) | 一种融合抽象语法树结构信息的智能代码补全方法 |
-
2023
- 2023-06-12 CN CN202310688989.9A patent/CN116643759A/zh active Pending
Cited By (2)
Publication number | Priority date | Publication date | Assignee | Title |
---|---|---|---|---|
CN117573096A (zh) * | 2024-01-17 | 2024-02-20 | 合肥综合性国家科学中心人工智能研究院(安徽省人工智能实验室) | 一种融合抽象语法树结构信息的智能代码补全方法 |
CN117573096B (zh) * | 2024-01-17 | 2024-04-09 | 合肥综合性国家科学中心人工智能研究院(安徽省人工智能实验室) | 一种融合抽象语法树结构信息的智能代码补全方法 |
Similar Documents
Publication | Publication Date | Title |
---|---|---|
Dam et al. | A deep language model for software code | |
CN114610515B (zh) | 基于日志全语义的多特征日志异常检测方法及系统 | |
CN113420296B (zh) | 一种基于Bert模型与BiLSTM的C源代码漏洞检测方法 | |
Yogatama et al. | Learning to compose words into sentences with reinforcement learning | |
CN114169330B (zh) | 融合时序卷积与Transformer编码器的中文命名实体识别方法 | |
WO2019233112A1 (zh) | 一种软件源代码的向量化表征方法 | |
CN114048464B (zh) | 基于深度学习的以太坊智能合约安全漏洞检测方法及系统 | |
CN113434357B (zh) | 基于序列预测的日志异常检测方法及装置 | |
CN109726120B (zh) | 一种基于机器学习的软件缺陷确认方法 | |
CN113312447B (zh) | 基于概率标签估计的半监督日志异常检测方法 | |
CN112560036B (zh) | 一种基于神经网络与深度学习的c/c++漏洞静态检测方法 | |
CN110321430B (zh) | 域名识别和域名识别模型生成方法、装置及存储介质 | |
CN110956309A (zh) | 基于crf和lstm的流程活动预测方法 | |
CN116643759A (zh) | 一种基于程序依赖图预测的代码预训练模型训练方法 | |
CN114942879A (zh) | 一种基于图神经网络的源代码漏洞检测与定位方法 | |
CN113591465A (zh) | 基于关联增强的网络威胁情报多维IoC实体识别方法及装置 | |
CN116305158A (zh) | 一种基于切片代码依赖图语义学习的漏洞识别方法 | |
CN116340952A (zh) | 一种基于操作码程序依赖图的智能合约漏洞检测方法 | |
Lawrence et al. | Explaining neural matrix factorization with gradient rollback | |
Jeyakarthic et al. | Optimal bidirectional long short term memory based sentiment analysis with sarcasm detection and classification on twitter data | |
CN116702160B (zh) | 一种基于数据依赖增强程序切片的源代码漏洞检测方法 | |
CN113158194A (zh) | 一种基于多关系图网络的漏洞模型的构建方法及检测方法 | |
CN113076545A (zh) | 一种基于深度学习的内核模糊测试序列生成方法 | |
Wen et al. | A Cross-Project Defect Prediction Model Based on Deep Learning With Self-Attention | |
Yan et al. | Prompt Engineering-assisted Malware Dynamic Analysis Using GPT-4 |
Legal Events
Date | Code | Title | Description |
---|---|---|---|
PB01 | Publication | ||
PB01 | Publication | ||
SE01 | Entry into force of request for substantive examination | ||
SE01 | Entry into force of request for substantive examination |