[논문 리뷰] kTrans: Knowledge-Aware Transformer for Binary Code Embedding

Wenyu Zhu, Hao Wang외 5명|arXiv (Cornell University)|2023. 08. 24.

Software Engineering ResearchComputer Science인용 수 3

한 줄 요약

kTrans는 명시적 어셈블리어 지식(예: 명령어 집합 아키텍처)을 통합하고 새로운 사전 훈련 작업을 통해 암묵적 의존성(예: 데이터 플로우 및 제어 플로우 관계)을 모델링함으로써 이진 코드 임베딩을 향상시키는 지식 인식 Transformer 모델이다. 이는 이진 코드 유사도 검출, 기능 유형 복원, 간접 호출 인식 작업에서 각각 5.2%, 6.8%, 12.6% 향상된 최신 기술 수준(SOTA) 성능을 달성한다.

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는 명령어 집합 아키텍처(ISA)에 대한 명시적 지식을 추가 토큰으로 통합하여 입력으로 제공한다. 이는 옵코드 유형, 오퍼랜드 관계, 레지스터 계층 구조(예: rax → eax → ax → al)를 포함한다.
• 제어 플로우 및 데이터 플로우 관계와 같은 암묵적 의존성을 모델링하기 위해 새로운 사전 훈련 작업인 지시어 경계 예측(IBP)을 도입한다.
• 지식 인식 토큰 표현을 강화한 다중 헤드 자기주의 메커니즘을 사용하여 어셈블리 시퀀스 내의 의미적 및 구조적 관계를 포착한다.
• 이중 단계 훈련 프로세스를 적용한다: 첫 번째 단계로 IBP와 마스킹된 지시어 예측을 수행한 대규모 어셈블리 코퍼스에서의 사전 훈련; 두 번째 단계로 하류 작업에서의 미세조정.
• 지osition 인코딩과 지시어 경계 검출을 활용하여 지시어 수준의 의미를 향상시켜 [rbp, pop]과 [pop, rbp]와 같은 토큰 시퀀스의 오해를 방지한다.
• 구문 구조와 ISA 정의로부터 유도된 의미 지식을 모두 포함하는 임베딩을 엔드 투 엔드로 학습할 수 있도록 프레임워크를 지원한다.

실험 결과

연구 질문

• RQ1 명시적 ISA 지식 통합이 이진 코드 임베딩의 품질을 향상시킬 수 있는가?
• RQ2 새로운 사전 훈련 작업을 통해 지시어 간 암묵적 의존성을 모델링하면 하류 작업 성능이 향상되는가?
• RQ3 이진 코드 유사도 검출, 기능 유형 복원, 간접 호출 인식에서 kTrans는 최신 기술 수준 모델보다 어떻게 비교되는가?
• RQ4 명시적 지식과 암묵적 지식를 모두 통합하면 다양한 역설계 작업에 대한 일반화 능력이 얼마나 향상되는가?

주요 결과

• 이진 코드 유사도 검출(BCSD)에서 kTrans는 최신 기술 수준 모델을 5.2% 초월하여 MRR 0.891을 달성한다.
• 기능 유형 복원(FTR) 성능은 기존 방법보다 6.8% 향상되어 Recall@1 0.382를 기록한다.
• 간접 호출 인식(ICR)에서 kTrans는 12.6% 향상되어 MRR 0.280 및 Recall@1 0.133(풀 크기 32)를 달성한다.
• 모델은 다양한 풀 크기에서 일관된 슈퍼리어리티를 유지하여 강건성과 일반화 능력을 입증한다.
• 시각화 및 이방성 탐지 결과, kTrans는 기준 모델보다 더 의미적으로 일관되고 구분 가능한 임베딩을 생성함을 확인한다.
• 제거 실험 결과, 명시적 지식 통합과 IBP 사전 훈련 작업 모두 성능 향상에 필수적임을 입증한다.