[Paper Review] kTrans: Knowledge-Aware Transformer for Binary Code Embedding

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

Software Engineering ResearchComputer Science3 citations

TL;DR

kTrans is a knowledge-aware Transformer model that enhances binary code embedding by injecting explicit assembly language knowledge (e.g., instruction set architecture) and modeling implicit dependencies via a novel pre-training task. It achieves state-of-the-art performance, outperforming prior methods by 5.2%, 6.8%, and 12.6% on binary code similarity detection, function type recovery, and indirect call recognition, respectively.

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

Motivation & Objective

• To address the lack of prior knowledge integration in existing binary code embedding models.
• To improve semantic understanding of assembly language by incorporating instruction set architecture (ISA) knowledge.
• To model implicit dependencies—such as data-flow and control-flow relationships—beyond token-level masking.
• To develop a unified framework that fuses explicit and implicit knowledge within a Transformer architecture.
• To demonstrate superior performance on downstream reverse engineering tasks like similarity detection, type recovery, and indirect call recognition.

Proposed method

• kTrans injects explicit knowledge about instruction set architecture (ISA), including opcode types, operand relationships, and register hierarchies (e.g., rax → eax → ax → al), as additional input tokens.
• It introduces a novel pre-training task—Instruction Boundary Prediction (IBP)—to model implicit dependencies such as control-flow and data-flow relationships between instructions.
• The model uses a multi-head self-attention mechanism enhanced with knowledge-aware token representations to capture semantic and structural relationships in assembly sequences.
• kTrans employs a two-stage training process: first, pre-training on a large-scale assembly corpus with IBP and masked instruction prediction; second, fine-tuning on downstream tasks.
• It leverages positional encoding and instruction boundary detection to improve modeling of instruction-level semantics, avoiding misinterpretation of token sequences like [rbp, pop] vs. [pop, rbp].
• The framework supports end-to-end learning of embeddings that encode both syntactic structure and semantic knowledge from ISA definitions.

Experimental results

Research questions

• RQ1 Can explicit injection of ISA knowledge improve the quality of binary code embeddings?
• RQ2 Can modeling implicit dependencies between instructions through a novel pre-training task enhance downstream performance?
• RQ3 How does kTrans compare to state-of-the-art models in binary code similarity detection, function type recovery, and indirect call recognition?
• RQ4 To what extent does incorporating both explicit and implicit knowledge improve generalization across diverse reverse engineering tasks?

Key findings

• kTrans outperforms state-of-the-art models by 5.2% on Binary Code Similarity Detection (BCSD), achieving an MRR of 0.891.
• It improves Function Type Recovery (FTR) performance by 6.8% over existing methods, with a recall@1 of 0.382.
• On Indirect Call Recognition (ICR), kTrans achieves a 12.6% improvement, reaching an MRR of 0.280 and Recall@1 of 0.133 with a pool size of 32.
• The model maintains consistent superiority across varying pool sizes, demonstrating robustness and generalization capability.
• Visualization and outlier detection confirm that kTrans generates more semantically coherent and distinct embeddings than baseline models.
• Ablation studies show that both explicit knowledge injection and the IBP pre-training task are critical for performance gains.