[論文レビュー] 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 は、命令セットアーキテクチャ（ISA）などの明示的なアセンブリ言語知識を統合し、新しい事前学習タスクを通じて暗黙的な依存関係をモデル化することで、バイナリコード埋め込みを向上させる知識認識型 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）知識を統合することで、アセンブリ言語の意味的理解を向上させること。
• トークンレベルのマスキングを超えて、データフローおよび制御フロー関係などの暗黙的な依存関係をモデル化すること。
• Transformer アーキテクチャ内に明示的および暗黙的知識を統合する包括的なフレームワークを構築すること。
• 類似度検出、型回復、間接呼び出し認識などの下流リバースエンジニアリングタスクにおける優れた性能を実証すること。

提案手法

• kTrans は、オペコードタイプ、オペランド関係、レジスタ階層（例：rax → eax → ax → al）などを含む、命令セットアーキテクチャ（ISA）に関する明示的知識を追加の入力トークンとして統合します。
• 制御フローおよびデータフロー関係などの暗黙的依存関係をモデル化するため、新しい事前学習タスク「命令境界予測（IBP）」を導入します。
• 知識認識型トークン表現を強化したマルチヘッド自己注意メカニズムを用い、アセンブリシーケンス内の意味的および構造的関係を捉えます。
• 2段階の学習プロセスを採用：まず、IBP およびマスクされた命令予測を用いた大規模アセンブリコーパスでの事前学習；次に、下流タスクでの微調整。
• 位置エンコーディングおよび命令境界検出を活用して、命令レベルの意味をより正確にモデル化し、[rbp, pop] と [pop, rbp] のようなトークンシーケンスの誤解を回避します。
• フレームワークは、ISA 定義からの構文的構造および意味的知識を統合したエンドツーエンドの埋め込み学習を可能にします。

実験結果

リサーチクエスチョン

• RQ1 明示的な ISA 知識の統合が、バイナリコード埋め込みの品質向上に寄与するか？
• RQ2 新しい事前学習タスクを通じて命令間の暗黙的依存関係をモデル化することで、下流タスクの性能が向上するか？
• RQ3 kTrans は、バイナリコード類似度検出、関数型回復、間接呼び出し認識において、最先端モデルと比較してどのように差をつけるか？
• RQ4 明示的および暗黙的知識の両方を統合することで、多様なリバースエンジニアリングタスクにおける一般化能力がどの程度向上するか？

主な発見

• kTrans は、バイナリコード類似度検出（BCSD）において、最先端手法を 5.2% 上回り、MRR が 0.891 を達成しました。
• 関数型回復（FTR）においては、既存手法を 6.8% 上回り、recall@1 が 0.382 に達しました。
• 間接呼び出し認識（ICR）では 12.6% の改善を達成し、MRR が 0.280、pool size 32 の条件下で recall@1 が 0.133 に達しました。
• モデルは、さまざまな pool size において一貫した優位性を示し、耐障害性および一般化能力の高さを裏付けています。
• 可視化および外れ値検出の結果、kTrans はベースラインモデルよりも意味的に整合性が高く、明確に分離された埋め込みを生成することが確認されました。
• アブレーションスタディの結果、明示的知識の統合と IBP 事前学習タスクの両方が、性能向上に不可欠であることが示されました。