[论文解读] kTrans: Knowledge-Aware Transformer for Binary Code Embedding

Wenyu Zhu, Hao Wang+5 more|arXiv (Cornell University)|Aug 24, 2023

Software Engineering ResearchComputer Science被引用 3

一句话总结

kTrans 是一种知识感知的 Transformer 模型，通过注入显式的汇编语言知识（例如，指令集架构）并利用一种新颖的预训练任务建模隐式依赖关系，从而增强二进制代码嵌入。它在二进制代码相似性检测、函数类型恢复和间接调用识别任务上分别实现了 5.2%、6.8% 和 12.6% 的性能提升，达到当前最先进水平。

ABSTRACT

Binary Code Embedding (BCE) has important applications in various reverse engineering tasks such as binary code similarity detection, type recovery, control-flow recovery and data-flow analysis. Recent studies have shown that the Transformer model can comprehend the semantics of binary code to support downstream tasks. However, existing models overlooked the prior knowledge of assembly language. In this paper, we propose a novel Transformer-based approach, namely kTrans, to generate knowledge-aware binary code embedding. By feeding explicit knowledge as additional inputs to the Transformer, and fusing implicit knowledge with a novel pre-training task, kTrans provides a new perspective to incorporating domain knowledge into a Transformer framework. We inspect the generated embeddings with outlier detection and visualization, and also apply kTrans to 3 downstream tasks: Binary Code Similarity Detection (BCSD), Function Type Recovery (FTR) and Indirect Call Recognition (ICR). Evaluation results show that kTrans can generate high-quality binary code embeddings, and outperforms state-of-the-art (SOTA) approaches on downstream tasks by 5.2%, 6.8%, and 12.6% respectively. kTrans is publicly available at: https://github.com/Learner0x5a/kTrans-release

研究动机与目标

• 解决现有二进制代码嵌入模型中缺乏先验知识整合的问题。
• 通过整合指令集架构（ISA）知识，提升对汇编语言的语义理解能力。
• 通过一种新颖的预训练任务，建模指令之间的隐式依赖关系（如数据流和控制流关系），超越传统的基于 token 的掩码方法。
• 构建一个统一框架，在 Transformer 架构中融合显式与隐式知识。
• 在下游逆向工程任务（如相似性检测、类型恢复和间接调用识别）中展示卓越性能。

提出的方法

• kTrans 将关于指令集架构（ISA）的显式知识（包括操作码类型、操作数关系和寄存器层次结构，例如 rax → eax → ax → al）作为额外输入 token 注入。
• 提出一种新颖的预训练任务——指令边界预测（IBP），用于建模指令之间的隐式依赖关系，如控制流和数据流关系。
• 模型使用增强的知识感知 token 表示的多头自注意力机制，以捕捉汇编序列中的语义和结构关系。
• kTrans 采用两阶段训练流程：首先在大规模汇编语料上进行 IBP 和掩码指令预测的预训练；其次在下游任务上进行微调。
• 通过位置编码和指令边界检测提升对指令级语义的建模能力，避免对 token 序列（如 [rbp, pop] 与 [pop, rbp]）的误解释。
• 该框架支持端到端学习嵌入表示，编码来自 ISA 定义的语法结构与语义知识。

实验结果

研究问题

• RQ1 显式注入 ISA 知识是否能提升二进制代码嵌入的质量？
• RQ2 通过新颖的预训练任务建模指令间的隐式依赖关系，是否能提升下游任务性能？
• RQ3 kTrans 在二进制代码相似性检测、函数类型恢复和间接调用识别任务上与最先进模型相比表现如何？
• RQ4 同时整合显式与隐式知识在多大程度上提升了在多样化逆向工程任务中的泛化能力？

主要发现

• 在二进制代码相似性检测（BCSD）任务上，kTrans 相较于最先进模型提升 5.2%，MRR 达到 0.891。
• 在函数类型恢复（FTR）任务上，性能相比现有方法提升 6.8%，召回率@1 达到 0.382。
• 在间接调用识别（ICR）任务上，kTrans 实现 12.6% 的性能提升，MRR 达到 0.280，召回率@1 达到 0.133（池大小为 32）。
• 该模型在不同池大小下均保持一致的优越性能，展现出良好的鲁棒性与泛化能力。
• 可视化与异常值检测结果表明，kTrans 生成的嵌入比基线模型更具语义一致性与区分度。
• 消融实验表明，显式知识注入与 IBP 预训练任务对性能提升均至关重要。