【论文笔记】Order Matters (AAAI 20)

注:部分内容参考自GPT生成的内容

论文笔记:Order Matters(AAAI 20)

用于二进制代码相似性检测的语义感知神经网络

论文:《Order Matters: Semantic-Aware Neural Networks for Binary Code Similarity Detection》(AAAI 2020)

笔记参考:AAAI-20论文解读:基于图神经网络的二进制代码分析 | 腾讯科恩实验室官方博客 (tencent.com)

动机

传统方法通常使用图匹配算法,但这些方法慢且不准确。尽管基于神经网络的方法取得了进展(如Gemini),但它们每个基本块都是以人工选择特征的低维嵌入来表示的,通常不能充分捕获二进制代码的语义信息。其次,节点的顺序在表示二进制函数时起着重要作用,而以往的方法并没有设计提取节点顺序的方法。

另外,在Related Work中提到,(Zuo et al 2018) 使用的NLP模型也有缺点。他们通过修改编译器,在每个生成的汇编块中添加一个基本块特殊注释器,该注释器为每个生成的块注释一个唯一ID。这样,可以将来自同一源代码片段编译的两个基本块视为等效。获取相似块对是一个有监督的过程,不同操作系统或硬件架构需要训练不同的模型

方法

提出的模型

模型的输入是二进制代码函数的控制流图(CFGs),其中每个块是带有中间表示的token序列。在语义感知上,模型使用BERT预训练接受CFG作为输入,并预训练token嵌入和块嵌入。在结构感知上,使用带GRU更新函数的MPNN来计算图的语义和结构嵌入gssg_{ss}。在顺序感知上,采用CFG的邻接矩阵作为输入,并使用CNN来计算图的顺序嵌入gog_{o}。最后将它们连接起来,并使用一个MLP层来计算图嵌入gfinal=MLP([gss,go])g_{final} =MLP([g_{ss}, g_{o}])

image-20240121203627839

语义感知模块 (Semantic-aware Modeling)

使用BERT进行预训练,包括四个任务:掩码语言模型任务(MLM)、邻接节点预测任务(ANP)、块内图任务(BIG)和图分类任务(GC)。这些任务帮助模型提取CFG的token级别、块级别和图级别的语义信息。

image-20240121204851147
  • 掩码语言模型任务(MLM):通过在输入层掩盖token并在输出层预测它们来提取块内的语义信息。这是一个自监督任务,模型在训练过程中某些token会被隐藏,模型必须基于其他token提供的上下文来预测缺失的token。
  • 邻接节点预测任务(ANP):因为块的信息不仅与块本身的内容相关,还与其邻近的块相关,ANP任务旨在让模型学习这种邻接信息。它涉及提取图中所有相邻块对,并在同一图中随机抽样多个块对,以预测它们是否相邻。
  • 图内块任务(BIG):与ANP类似,BIG任务旨在帮助模型判断两个节点是否存在于同一图中。它涉及随机抽样可能在同一图中或不在同一图中的块对,并预测它们的关系。这有助于模型理解块与整个图之间的关系
  • 图分类任务(GC):使模型能够基于不同平台、架构或优化选项来分类块,特别是在不同编译条件下。GC任务要求模型区分由于这些不同条件而产生的图和块信息的差异

结构感知模块 (Structural-aware Modeling)

使用消息传递神经网络(MPNN)结合GRU(门控循环单元)更新函数,以提取CFG的全图语义和结构嵌入(the whole graph semantic & structural embedding)gssg_{ss}

image-20240121215922638
  • 消息传递(Message Passing)

    • 公式(2)表示的是消息传递阶段。对于图中的每个节点v,它计算了节点v在时间t+1的消息mvt+1m^{t+1}_v。这个消息是通过聚合节点v的所有邻居节点w的信息(使用消息函数M)来得到的。这里,hvth^t_vhwth^t_w分别代表节点v和它的邻居节点w在时间t的嵌入,而evwe_{vw}是节点v和w之间边的特征。
    • 公式(5)中,论文使用了多层感知机(Multi-Layer Perceptron, MLP)对邻居节点w的嵌入hwth^t_w进行处理。
  • 更新(Update)

    • 公式(3)表示的是更新阶段。在这一步中,节点v的新嵌入hvt+1h^{t+1}_v是通过更新函数U,结合节点v在时间t的嵌入和它在时间t+1收到的消息来计算的。

    • 公式(6)说明了论文中的更新函数是通过GRU实现的,GRU考虑了节点的历史信息hvth^t_v和新的消息mvt+1m^{t+1}_v,来学习图的时序信息。

  • 读出(Readout)

    • 公式(4)定义了读出函数R,它计算 gssg_{ss}。这是通过对图中所有节点v的最终嵌入 hvTh^T_v进行聚合来实现的。
    • 公式(7)中,读出函数是通过对所有节点的初始嵌入hv0h^0_v和最终嵌入hvTh^T_v使用多层感知机(MLP)并进行求和来实现的。这里hv0h^0_v是由BERT预训练得到的初始块嵌入。
image-20240121220012016

顺序感知模块 (Order-aware Modeling)

通过卷积神经网络(CNN)处理邻接矩阵,以提取CFG节点的顺序信息。

image-20240122001434356

如图,CNN能捕获从(a)到(b)的变化信息,当 CNN 看到大量训练数据时,它具有平移不变性(translation invariance)

对于(b)->©,与图像放缩类似,在看到足够多的训练数据后,CNN 也可以学习这种伸缩不变性(scale invariance)。

由于二进制代码函数在不同平台上编译时节点顺序通常不会大改变,CNN能够处理由此引起的添加、删除或交换节点等小变化,优势如下:

  1. 使用CNN直接在邻接矩阵上的操作相比于传统的图特征提取算法要快得多。
  2. CNN可以处理不同大小的输入,这允许模型处理不同大小的图而无需预处理,如填充或裁剪。

使用具有 3 个残差块的 11 层 Resnet,所有的feature map大小均为3*3,最后使用最大池化层来计算图的顺序嵌入go=Maxpooling(Resnet(A))g_o = Maxpooling(Resnet(A))

效果

数据集:

  • 任务1是跨平台二进制代码检测,目的是确认相同的源代码在不同平台上编译成的CFG是否具有较高的相似性得分。

    • 与Gemini模型类似,使用孪生网络(siamese network)来减少损失,并使用余弦距离来计算图的相似性。
  • 任务2是图分类,对图嵌入进行优化选项分类

    • 使用softmax函数并选择交叉熵作为损失函数

由于模型具有三个组成部分:语义感知、结构感知和顺序感知,因此进行了不同的实验来找出每个部分的效果。

image-20240122012435397image-20240122012444904

表中:

  • 第一个分块是整体模型,包括graph kernel,Gemini以及MPNN模型。

  • 第二个分块是语义感知模块的对比实验,分别使用了word2vec[5],skip thought[6],以及BERT,其中BERT2是指原始BERT论文中的两个task(即MLM和ANP),BERT4是指在此基础上加入两个graph-level task(BIG和GC)。

  • 第三个分块是对顺序感知模块的对比实验,基础CNN模型使用3层CNN以及7、11层的Resnet,CNN_random是对训练集中控制流图的节点顺序随机打乱再进行训练,MPNN_ws是去除控制流图节点中的语义信息(所有block向量设为相同的值)再用MPNN训练。

  • 最后是本文的最终模型,即BERT (4 tasks) + MPNN + 11layer Resnet。

MPNN (即加上结构感知模块)在所有数据集上都优于 Gemini,这是因为 GRU 更新函数可以存储更多信息,因此在所有其他模型中都使用 MPNN。

基于NLP的块预训练特征比手动特征好得多,并且顺序感知模块在两个任务上也有很好的结果。

在跨平台二进制代码检测任务中,语义信息比顺序信息更有用。不同的CFG可能具有相似的节点顺序,因此仅使用节点顺序信息是不够的。

最后,最终模型优于所有其他模型。

分开观察各个模块的有效性

image-20240122014527423 image-20240122014536191

“语义感知(Semantic-aware)”:

  • 表中的第二块显示,BERT模型的性能优于word2vec和skip thought模型。这是因为BERT在预训练过程中不仅考虑了块级别的预测,还包括了token级别的预测,并且双向Transformer结构能够提取更多有用的信息。
  • 当BERT模型加入了BIG和GC两个图级任务后,性能有了1%到2%的提升,表明引入图级任务对预训练是有益的。
  • 图6展示了4个控制流图(CFG)的块嵌入可视化,使用K-means算法将预训练后的块嵌入分成四个类别,每个类别用不同颜色表示。从图中可以观察到,同一控制流图中的块倾向于拥有相同的颜色,而不同控制流图的主要颜色也不同。

“顺序感知(Order-aware)”:

  • 表中的第三块显示,基于CNN的模型在两个任务上都取得了良好的效果,其中11层的Resnet略优于3层的CNN和7层的Resnet。
  • 与不含语义信息的MPNN(MPNN_ws)相比,基于CNN的模型表现出更好的性能。
  • 节点顺序被随机打乱后,CNN的效果显著下降,这证明CNN模型确实能够学习到图的节点顺序信息。
  • 图7展示了两个由相同源代码编译而成的CFG变化的例子,尽管左图的节点3在右图中被分成了节点3和4,但其他节点的顺序和边的连接方式保持不变。通过CNN模型的计算,这两个CFG的余弦相似度为0.971,并且在整个平台中的代码检测排名中位列第一。这意味着CNN模型能够从邻接矩阵中有效提取控制流图的节点顺序信息,与假设相符。

结论

这篇论文提出了一个新颖的二进制代码图学习框架,包含了语义感知组件、结构感知组件和顺序感知组件。作者观察到,语义信息和节点顺序信息对于表示控制流图(CFGs)都非常重要。为了捕捉语义特征,作者提出了针对CFGs块的BERT预训练,包括两个原始任务MLM和ANP,以及两个额外的图级任务BIG和GC。然后作者使用MPNN来提取结构信息。作者进一步提出了一个基于CNN的模型来捕捉节点顺序信息。作者在两个任务上使用了四个数据集进行了实验,实验结果表明本文提出的模型超越了当时最先进的方法。

附:部分基础概念解释

由于是第一次精读深度学习相关的技术论文,我翻看了很多基础概念

MLM和NSP

MLM(Masked language model)和NSP(next sentence prediction)是BERT模型中两个重要的训练任务,它们共同帮助BERT学习理解语言的深层次结构和关系。以下是对这两个任务的具体介绍:

掩码语言模型任务(MLM)

  • 目的:MLM旨在使模型能够更好地理解语言本身的规律和结构。它通过在文本中随机掩盖一些单词(即使用特殊的“[MASK]”标记替换),然后要求模型预测这些掩盖单词的原始值来实现。
  • 训练过程:在训练时,BERT模型会尝试根据上下文中的其他单词来猜测被掩盖的单词是什么。例如,在句子“The cat sat on the [MASK]”中,模型需要预测被掩盖的词是“mat”。
  • 作用:这种训练方式使得BERT能够有效地学习单词的上下文关系和语义信息,从而更好地理解语言。

下一个句子预测任务(NSP)

  • 目的:NSP的目标是使模型能够理解句子之间的关系。这对于很多NLP任务(如问答系统、自然语言推理等)至关重要。
  • 训练过程:在训练时,模型被给予一对句子,并需要判断第二个句子是否在原文中紧跟在第一个句子之后。训练集由两种类型的句子对组成:一种是真实的相邻句子对,另一种是随机组合的非相邻句子对。
  • 作用:通过这种方式,BERT学习理解句子之间的逻辑和关系,增强对文本的整体理解能力。

消息传递神经网络(MPNN)

消息传递神经网络(Message Passing Neural Network)是一类图神经网络,它通过在图的节点之间交换信息来学习节点的表示。它们基于以下步骤工作:

  1. 消息传递:每个节点接收其邻居节点的信息,并根据这些信息生成“消息”。
  2. 聚合:将所有接收到的消息聚合成单个表示,这可以通过不同的函数实现,如求和、求平均或更复杂的操作。
  3. 更新:使用聚合的信息来更新节点的状态。

MPNN的核心思想是通过迭代这些步骤来精炼每个节点的表示,从而捕捉图的结构特征和节点之间的关系。

门控循环单元(GRU)

GRU(gated recurrent unit)是循环神经网络(RNN)的一种变体,用于处理序列数据。与传统的RNN相比,GRU通过引入门控机制来解决梯度消失和梯度爆炸的问题,使得网络能够捕捉长距离依赖关系。GRU包含两个门:

  1. 更新门:决定状态信息应该如何更新。
  2. 重置门:决定过去的状态信息在计算新状态时应保留多少。

在每个时间步,GRU可以选择保留旧状态的信息并融入新输入的信息,这使得它在处理具有复杂依赖结构的数据时非常有效。

多层感知机 (MLP)

  • 定义:多层感知机是一种基础的人工神经网络,由一个输入层、若干隐藏层和一个输出层组成。每一层由多个神经元组成,相邻层之间的神经元通过权重连接。
  • 功能:MLP主要用于分类和回归问题,能够识别和建模输入数据中的非线性关系。
  • 工作原理:在MLP中,数据从输入层进入,每个神经元对输入进行加权求和,再加上一个偏置项,最后通过激活函数进行非线性转换。这个过程在每个隐藏层中重复进行,直到输出层。在输出层,数据被转换为最终的输出格式(如分类标签或回归值)。

卷积神经网络 (CNN)

  • 定义:卷积神经网络是一种深度学习网络,特别适用于处理具有网格结构的数据,如图像(2D网格)和声音(1D网格)。
  • 功能:CNN广泛应用于图像和视频识别、图像分类、医学图像分析、自然语言处理等领域。
  • 工作原理:CNN通过一系列卷积层、池化层和全连接层处理数据。卷积层使用卷积核提取空间特征,池化层(如最大池化)则减小特征维度并提供一定程度的位置不变性。最后,全连接层将提取的特征用于分类或回归任务。

最大池化 (MaxPooling)

  • 定义:最大池化是一种池化操作,常在卷积神经网络中使用,用于减小特征图的空间尺寸。
  • 功能:最大池化通过降低参数数量和计算量来减少过拟合,同时保持重要特征。
  • 工作原理:最大池化通过在输入特征图的不同区域上应用一个固定大小的窗口,并从每个窗口中选择最大值来实现。这样做可以提取最显著的特征,并且对小的位置变化保持不变性。

残差网络 (ResNet)

  • 定义:ResNet是一种深度卷积神经网络,通过引入残差学习框架来易于优化,并能够构建更深的网络。
  • 功能:ResNet在图像识别、分类和其他计算机视觉任务中表现优异。
  • 工作原理:在ResNet中,残差块的引入允许输入跳过一些层。每个残差块学习输入和输出的残差(差异),而不是直接学习输出。这帮助网络学习恒等映射,解决了深层网络中的梯度消失问题。