Skip to content

Latest commit

 

History

History
56 lines (42 loc) · 4.41 KB

File metadata and controls

56 lines (42 loc) · 4.41 KB

ContraBin 思路总结

ContraBin 的思想比较有意思,是通过语义建立二进制和源码之间的桥梁。阅读的思路总结如下:

它的训练过程是一个多阶段的 pipeline,核心在于通过对比学习将源代码和注释的语义"注入"到二进制代码的表示中。整个训练过程可以清晰地分为以下四个主要阶段: ContraBin

第一阶段:数据准备

在开始训练之前,需要构建一个特定的数据集。

  • 源代码:从 AnghaBench 数据集中获取约100万个单函数C语言代码片段。
  • 二进制代码:使用 Clang/LLVM 编译器将每个C源代码片段编译成中间表示。LLVM IR 被视为一种平台无关的"二进制代码"。
  • 注释:作者发现真实世界中人工编写的注释质量参差不齐,所以使用一个预训练的 CodeT5 模型为每个源代码片段自动生成一句简洁的功能性描述作为注释。

第二阶段:主对比学习

目标是初步对齐三种不同表示的嵌入空间。

  1. 输入:从一个批次的三元组中,随机选择两个表示。
  2. 编码
    • 锚定编码器:一个不更新参数的预训练模型,用于编码源代码和注释。这是因为源代码和注释的语义已经很丰富,直接用预训练模型提取其特征即可。
    • 可训练编码器:一个需要更新参数的编码器,用于编码二进制代码。这是本次训练的重点,目标是让它学会产出有意义的二进制代码表示。
  3. 对比损失计算:计算两种表示嵌入之间的批量相似度矩阵。
    • 损失函数的设计灵感同样来源于 CLIP,目标是让匹配的表示对(如同一个程序的源代码和二进制代码)在嵌入空间中更接近,而不匹配的表示对(如不同程序的源代码和二进制代码)更远离。
  4. 目的:让"可训练的二进制编码器"初步学会将二进制代码映射到一个与其他两种表示相关的语义空间中,为更复杂的训练打下基础。

第三阶段:次级对比学习

这是 ContraBin 最核心、最创新的阶段,它在主对比学习的基础上,引入了单纯形插值来生成中间表示,进行更精细的语义对齐。

  1. 输入:不再随机选择两个,而是固定地选择源代码和注释的嵌入表示(H_s 和 H_c)作为插值的基础。
  2. 单纯形插值:使用两种方法对 H_s 和 H_c 进行插值,以生成一个"中间视图" H_i。
    • 线性插值:H_l = λ_l · H_s + (1 - λ_l) · H_c
    • 非线性插值:H_n = λ_n · H_s + (1 - λ_n) · H_c。其中 λ_n 是一个由另一个小型神经网络预测的与嵌入向量同形的矩阵,允许进行更细粒度、特征级别的插值。
  3. 中间对比学习
    • 目标:将这个生成的中间表示 H_i(它融合了源代码和注释的语义)与对应的二进制代码表示 H_b 拉近。
    • 损失函数:使用标准的对比学习损失(InfoNCE)。公式的核心是:
      损失 = log( exp(相似度(H_i, H_b)) / (exp(相似度(H_i, H_b)) + 所有负样本的相似度之和) )
      
    • 这个损失会使得 H_i 和 H_b 在语义空间中越靠越近,同时远离批次中其他程序的表示。

目的与逻辑: 插值过程模拟了人类在代码和注释之间转换的"思考过程",等于是 comment 作为连接二者之间的一个桥梁。通过强制要求二进制代码表示与这个"思考的中间产物"对齐,模型被逼着去理解并融合源代码和注释中的高层语义,从而生成质量更高、信息更丰富的二进制代码嵌入。

第四阶段:下游任务微调

在预训练完成后,得到了一个高质量的二进制代码编码器。

  1. 任务:为了验证预训练得到的嵌入表示的有效性,作者在四个下游任务上对模型进行了微调:
    • 二进制功能分类:给定二进制代码,判断其算法功能(如排序、哈希)
    • 函数名恢复:给定二进制函数,预测其原始的函数名称
    • 代码摘要:给定二进制代码,生成描述其功能的人类语言摘要
    • 逆向工程:给定二进制代码,尝试恢复出等价的C语言源代码
  2. 方法:在预训练的编码器之后,接上针对特定任务的任务头(例如,分类器或解码器),然后在特定任务的数据集上进行有监督的微调。