Exploring Software Naturalness throughNeural Language Models
论文概述
软件自然性假说（The Software Naturalness hypothesis）认为，编程语言可以通过自然语言处理中使用的相同技术来理解。作者通过使用预训练好的基于Transformer的语言模型来执行代码分析任务来验证这个假设。
现有的程序分析方法高度依赖于从AST（抽象语法树）中提取到的特征，或者是从CFG（控制流图），DFG（数据流图）。然而，这些方法有如下缺点：
- 需要先编译或者是用解析器解析源代码，在计算上非常消耗资源
- 难以分析不完整的代码片段
而作者研究的是基于Transformer的模型能否根据原始的源代码自动提取AST特征。为此，作者引入了预训练任务：
- sequence labeling task
作者采用的是C语言代码，预训练结果证明基于Transformer的模型在 AST-tagging任务中取得高准确率。进一步，作者将此模型应用到软件漏洞识别的任务中，并且实验结果证明，这样取得到效果与基于图（graph）的效果相当，并且不像基于图的方法依赖于编译器。
论文贡献：
- 研究了基于Transformer的语言模型在复杂源代码分析任务中的应用
- 证明了character-based tokenization和WWM预训练可以消除OOV和稀有词
- 首个研究Tranformer语言模型是否能自动发现AST特征的研究
- 基于Transformer的模型在程序分析任务上性能优于用编译器提取程序特征的基于图（graph-based）的模型。
模型:C-BERT
BERT在源代码中的应用需要重新考虑标记策略（tokenization strategies），在预训练的C语言语料库中有400万token，几乎无限，带来难以解决的OOV问题。作者在此探索了3种subword tokenize策略来减小OOV频率和词表大小。
Tokenizer
作者探索了3种subword tokenize策略：
-
Character (Char)
子词tokenization的极限，直接用ASCII词汇表，由于训练语料库只包含由ASCII码构成的词，加上[CLS]，[SEP]，[MASK]，[UNK]。比如int在Char tokenization下变成了i n t3个子词。 -
Character + Keyword (KeyChar)
大多数编程语言的关键字数量相对较少。通过使用32个C语言关键字来扩充词汇表，C语言关键字会被当作1个token处理，而不会拆分为ASCII子词。 -
Sentence piece (SPE)
NLP模型几乎总是使用子词tokenization，如Sentence Piece或字节对编码（BPE）
BPE是一种解决OOV的办法，流程如下：- 首先将统计text中单词，做成词汇表（单词-频率），然后按照unigram进行分解
- 寻找频率最大的片段（字符），进行组合，将组合片段加入词汇表
- 继续重复上述操作，直到达到设定的阈值（词汇数+操作数）->操作数是唯一的超参数
与其他基于transformer的模型一样，使用sentence piece tokenization，词汇大小选择为5000。这包括词汇表中的所有C关键字，以及许多常见的库函数和变量名。
Transformer Based Language Models
作者使用的模型取名为C-BERT，与BERT的使用过程一样，首先在大型无标记的C语言语料库上进行预训练，然后针对特定任务进行fine-tuning。在训练时同样用[CLS] token表示序列的开始，[SEP] token表示序列的结束。
对每一个输入序列
X
=
[
x
1
,
x
2
,
.
.
.
,
x
T
−
1
,
x
T
]
X = [x_1, x_2,...,x_{T-1},x_T]
X=[x1,x2,...,xT−1,xT] (
x
1
=
[
C
L
S
]
x_1 = [CLS]
x1=[CLS],
x
T
x_T
xT = [SEP])
语言模型输出
H
=
[
h
1
,
h
2
,
.
.
.
,
h
T
−
1
,
h
T
]
H = [h_1, h_2,...,h_{T-1},h_T]
H=[h1,h2,...,hT−1,hT] (
h
1
=
h
C
L
S
h_1 = h_{CLS}
h1=hCLS,
h
T
h_T
hT =
h
S
E
P
h_SEP
hSEP)
其中 H H H 为 T × 768 T \times 768 T×768 矩阵。
作者这里提到了2种预训练任务和2个Fine-tuning任务。
Masked Language Model (MLM) Pre-training Objective
一小部分token被mask，在Masked Language Model 中，语言模型后会接全连接层 W L M ∈ R 768 × ∣ V ∣ W_{LM} \in R^{768 \times |V|} WLM∈R768×∣V∣ ( V V V 表示C-Bert词表) + s o f t m a x softmax softmax 来计算被mask的词的概率分布。
p L M = s o f t m a x ( H . W L M ) ∈ R T × ∣ V ∣ p_{LM} = softmax(H.W_{LM}) ∈ R^{T \times|V|} pLM=softmax(H.WLM)∈RT×∣V∣
Whole Word Masked (WWM) Pre-training Objective
作者从预训练数据集中选择要mask的字符串类型（例如，合法变量或函数名的文本字符串），并编写正则表达式来提取它们。
AST Fine-tuning Objective
在这个Sequence labeling任务中，添加全连接层 W A S T ∈ R 768 × ∣ V A S T ∣ W_{AST} \in R^{768 \times |V_{AST}|} WAST∈R768×∣VAST∣ ( V A S T V_{AST} VAST 表示AST标签语料库)
p A S T = s o f t m a x ( H . W A S T ) ∈ R T × ∣ V A S T ∣ p_{AST} = softmax(H.W_{AST}) ∈ R^{T \times|V_{AST}|} pAST=softmax(H.WAST)∈RT×∣VAST∣
VI Fine-tuning Objective.
漏洞识别任务属于二分类任务，并且这里整个源代码的向量表示用 h C L S ∈ R 768 h_{CLS} \in R^{768} hCLS∈R768 b表示（双向语言模型，不知道 h S E P h_{SEP} hSEP 可不可以）。添加全连接层 w ∈ R 768 w \in R^{768} w∈R768
p V I = s i g m o i d ( h C L S T . w ) p_{VI} = sigmoid(h_{CLS}^T.w) pVI=sigmoid(hCLST.w)
预训练语料库
github top-100 star的C语言项目
预训练
AST Node Tagging
为了研究transformer模型对代码语法的理解程度，作者首先引入了一个新的序列标记任务，它直接探测语言模型对Clang编译器生成的AST的理解程度。
这里使用到的模型已经是经过MLM预训练的。
作者将源代码当作编程语言的token序列。序列的每个token对应的标签有token_kind和cursor_kind。
部分token对应的token_kind和cursor_kind对应如下
这些token_kind和cursor_kind由clang编译器产生。clang总共有5种token_kind和209种cursor_kind。
AST可以视为cursor_kind的有向无循环图（directed acyclic graph），cursor_kind表示AST结点类型。
预测clang编译器的cursor_kind类似于NLP的词义消歧任务。比如运算符 * 可以作为二元乘法运算，也可以作为一元指针运算，而它们对应着不同的cursor_kind和同一个token_kind。
- 二元乘法运算
if ( lensum * 9 / 10 > maxpos )
- 一元指针运算
sizeof(**s->coarse)
在上面2个例子中，不仅 *， >也有不同语义，前者作为大于符号，后者属于取指针成员变量。
数据集
用于Fine-Tuning的数据集是github上FFmpeg和QEMU 2个repo。 FFmpeg包括1751个文件，QEMU包括2152个文件。
测试结果
漏洞识别任务 （Vulnerability Identification）
数据集
以function为单位，function来自FFmpeg和QEMU project，并且打好了non-vulnerable / vulnerable标签。数据集包括full datasets和reduced datasets
- full datasets包括9683个FFmpeg样本和17515个QEMU样本。
- Joern代码分析工具不能正确解析full datasets某些文件，因此基于图的方法的对比实验不能复现，所以从数据集中跳过这部分，组成reduced datsets。reduced datasets包括6169个FFmpeg样本和14896个QEMU样本。
baseline
使用BiLSTM和CNN，CGNN作为baseline。 BiLSTM和CNN直接使用源代码作为token序列的表示方式，并且用Word2Vec进行预训练。 GGNN则使用到了图特征。
测试结果
结果使用accuracy作为评估标准，结果如下表所示