Source-linked AI summary
PalmTree: Learning an Assembly Language Model for Instruction Embedding
Xuezixiang Li, Qu Yu, Heng Yin
TL;DR
Binary analysis models need fixed-length instruction representations, but existing schemes miss assembly’s internal structure and rely on noisy control-flow context. PalmTree pre-trains a BERT-based assembly language model with three self-supervised tasks, achieving the best intrinsic performance and outperforming other embedding schemes across downstream tasks.
Problem
Existing instruction embedding schemes overlook complex instruction structures and can rely on noisy control-flow context affected by compiler optimizations.
Method
PalmTree is a BERT-based assembly language model pre-trained on large-scale unlabeled binary corpora with MLM, CWP, and DUP tasks.
Results
PalmTree has the best performance in intrinsic evaluations and outperforms other instruction embedding models across downstream binary analysis applications.
Takeaways & Limitations
PalmTree produces general-purpose instruction embeddings that downstream models can directly use, including via a static lookup table when computing resources are limited.
Takeaways & Limitations
The paper does not aim to learn a language model across multiple CPU architectures.
Abstract
from arXiv · showhide
Deep learning has demonstrated its strengths in numerous binary analysis tasks, including function boundary detection, binary code search, function prototype inference, value set analysis, etc. When applying deep learning to binary analysis tasks, we need to decide what input should be fed into the neural network model. More specifically, we need to answer how to represent an instruction in a fixed-length vector. The idea of automatically learning instruction representations is intriguing, however the existing schemes fail to capture the unique characteristics of disassembly. These schemes ignore the complex intra-instruction structures and mainly rely on control flow in which the contextual information is noisy and can be influenced by compiler optimizations. In this paper, we propose to pre-train an assembly language model called PalmTree for generating general-purpose instruction embeddings by conducting self-supervised training on large-scale unlabeled binary corpora. PalmTree utilizes three pre-training tasks to capture various characteristics of assembly language. These training tasks overcome the problems in existing schemes, thus can help to generate high-quality representations. We conduct both intrinsic and extrinsic evaluations, and compare PalmTree with other instruction embedding schemes. PalmTree has the best performance for intrinsic metrics, and outperforms the other instruction embedding schemes for all downstream tasks.
1 INTRODUCTION
PalmTree addresses the challenge of representing assembly instructions by pre-training a general-purpose assembly language model that captures instruction structure and dependencies. Its embeddings achieve strong intrinsic and downstream performance and can be reused efficiently.
- Deep learning supports binary analysis tasks, but models require a fixed-length representation for each instruction.
- Automatically learned instruction representations avoid manual feature design and can capture higher-level features for downstream binary analysis.
- Existing schemes overlook complex instruction formats and rely on control-flow context that compiler optimizations can make noisy.
- PalmTree uses BERT-based pre-training with three tasks targeting instruction formats, contextual control flow, and data-flow dependencies.
- PalmTree provides high-quality general-purpose embeddings that downstream models can directly use, including through a precomputed lookup table.
- PalmTree achieves the best intrinsic-evaluation performance and outperforms other instruction embedding models in downstream binary analysis tasks.
2 BACKGROUND
Instruction embedding approaches range from raw-byte and manual encodings to learned representations. The background identifies shortcomings in semantic coverage, instruction-structure modeling, and control-flow-based context.
- Raw-byte Encoding: Raw-byte encoding is simple and efficient because it avoids computationally expensive disassembly, but it provides little semantic information about instructions and operands.
- Manual Encoding of Disassembled Instructions: Manual encoding extracts disassembly features such as opcodes, registers, addresses, and offsets, but does not provide higher-level semantic information without substantial expert knowledge.
- Learning-based Encoding: Learning-based methods use word2vec or PV-DM to learn continuous instruction vectors from disassembled code, sometimes separating opcode and operand embeddings.
- Complex and Diverse Instruction Formats: x86 instructions have diverse internal formats, including zero to three operands, complex memory expressions, prefixes, and implicit operands.
- Complex and Diverse Instruction Formats: Word2vec treats an entire instruction as one word, while Asm2Vec models at most one opcode and two operands and omits several instruction complexities.
- Noisy Instruction Context: Control-flow context can become unreliable because compiler optimizations reorder instructions and insert irrelevant instructions, while data-flow dependencies may span basic blocks.
3 DESIGN OF PALMTREE
PalmTree is a BERT-based assembly language model that learns instruction embeddings by combining fine-grained tokenization with three tasks targeting instruction structure, control-flow context, and data-flow dependencies.
- PalmTree is a BERT-based instruction embedding scheme that automatically learns a language model for assembly code.
- Input Representation: Instructions are decomposed into basic tokens so the model can capture complex internal formats, including opcodes, registers, memory expressions, constants, and symbols.A sample memory instruction is split into tokens such as “mov”, “rax”, “qword”, brackets, “rsp”, “+”, and “0x58”.
- Context Window Prediction: Context Window Prediction learns whether two instructions co-occur within a sliding window in control flow and preserves their ordering.The task uses binary classification rather than predicting the entire following instruction.
- Def-Use Prediction: Def-Use Prediction trains the model to predict whether two instructions have a data-dependency relation.The sampled dependencies include registers, memory locations, function-call arguments, and implicit EFLAGS dependencies.
- Masked Language Model: Masked Language Model training predicts randomly masked, corrupted, or unchanged tokens within instructions.For selected tokens, 80% are masked, 10% are replaced with another vocabulary token, and 10% remain unchanged.
- Instruction Representation: PalmTree generates instruction embeddings by mean-pooling hidden states from the transformer's second-last layer.The model supports frozen embedding generation or fine-tuning; static lookup tables can improve efficiency on resource-limited devices.
4 EVALUATION
The evaluation compares PalmTree with baseline embedding schemes through intrinsic tests, downstream binary-analysis tasks, and runtime measurements. PalmTree generally achieves the strongest accuracy and can improve efficiency through precomputed lookup tables, although its transformer encoding is slower than simpler approaches.
- Evaluation methodology: The framework evaluates instruction embeddings intrinsically through outlier detection and basic block search, and extrinsically through binary similarity, function type-signature, and value-set analysis.Intrinsic tests use opcode and operand semantics or source-code debug information; extrinsic tests assess embeddings within downstream models.
- Datasets: 3,266 binaries and 2.25 billion instructions form the pre-training corpus, with testing binaries drawn from different families and compilers.The training process includes about 2.36 billion positive and negative instruction-pair samples.
- Outlier detection: PalmTree ranks first in both opcode and operand outlier detection, while all three pre-training tasks contribute positively and DUP provides a substantial accuracy boost.The result indicates that the learned representation captures semantic differences in opcodes and operands and benefits from def-use relations.
- Basic block search: All three PalmTree configurations achieve higher basic-block-search AUC than the baselines, with consecutive improvements across configurations.Word2Vec performs worst, Instruction2Vec exceeds it, and Asm2Vec remains below the PalmTree variants.
- Downstream tasks: PalmTree outperforms other instruction-embedding approaches in each extrinsic evaluation, including tests on substantially different data.The complete model’s three training tasks support generalization, while the results also report faster downstream convergence for DeepVSA.
- Runtime efficiency: GPU acceleration lets PalmTree encode 107 binaries in about 70 seconds, while lookup tables can provide word2vec-like speed for frequently used instructions.PalmTree is slower than word2vec and Asm2Vec during direct encoding, and doubling embedding size makes it 1.7 to 1.9 times slower.
5 RELATED WORK
Prior work applies NLP-inspired representation learning to instructions and programming languages, but existing instruction embeddings have important limitations in preserving instruction structure and semantics.
- NLP representation learning progressed from word2vec to contextual models such as ELMo, transformers, BERT, and ELECTRA.
- Assembly instructions have clear grammar and syntax, enabling NLP models to process them as natural language.
- Instruction representations support binary analysis tasks, including code similarity search, prototype inference, and other applications.
- Instruction2Vec manually designs instruction representations, whereas InnerEye and related methods use word2vec with each instruction treated as a word.
- Word2vec-based instruction embeddings can lose internal instruction information, miss instruction-level semantics, and encounter out-of-vocabulary problems.
- Asm2Vec uses PV-DM and token-level language modeling to preserve instruction boundaries while jointly generating instruction and function embeddings.
6 DISCUSSION
The paper evaluates PalmTree primarily on x86 and discusses extending the approach to other architectures, while leaving cross-architecture modeling and richer input pairings for future work.
- PalmTree is trained for one instruction set or architecture, with the paper’s evaluation focused particularly on x86.The authors state that the technique can also apply to ARM and MIPS.
- Cross-architecture language modeling is outside the paper’s scope, although it could map semantically similar instructions from different architectures nearby in embedding space.The authors identify cross-architecture vulnerability and bug search as a potential use case and leave it for future work.
- Feeding code-segment, basic-block, or function pairs could better capture long-term instruction relations than instruction pairs.The current approach samples from the context window and data-flow graph; richer pairings are proposed as future work with potential performance benefits.
7 CONCLUSION
The paper concludes that PalmTree learns general-purpose instruction embeddings through self-supervised assembly-language pre-training and performs strongly across intrinsic and downstream evaluations.
- PalmTree is proposed to address instruction-representation challenges by generating general-purpose instruction embeddings.
- PalmTree uses self-supervised training on large-scale unlabeled binary corpora and three tasks: MLM, CWP, and DUP.The model is based on BERT and its tasks target inherent characteristics of assembly language.
- PalmTree achieves the best intrinsic-evaluation performance among the compared instruction-embedding models.
- PalmTree outperforms all baseline models in extrinsic evaluations and significantly improves downstream application performance.
A OPCODE AND OPERAND TYPES FOR OUTLIER DETECTION
The appendix categorizes opcodes and operand types for outlier detection, organizing operand combinations and listing the types of individual operands.
- Table 8 categorizes opcodes, while Table 9 categorizes operand types and their combinations.
- The operand-combination column distinguishes instructions with no operands, three operands, or two operands.Examples include “none” for retn and “tri” for three-operand instructions.
- The second and third columns list the types of the operands in each combination.The “reg-reg” example denotes two register operands.
B MORE FIGURES IN EVALUATIONS
Figures 15 and 16 report EKLAVYA’s training loss and accuracy in the Function Type Signature Inference task, alongside opcode and operand type tables.
- Figures 15 and 16 show EKLAVYA’s loss and accuracy curves during training.
- Figure 15 presents the loss value during training.
- Figure 16 presents accuracy during training, while Tables 8 and 9 organize opcode and operand types.
C.1 Embedding sizes
The embedding-size experiment evaluates PalmTree with four instruction-vector sizes and finds better intrinsic performance at larger sizes, subject to hardware-efficiency considerations.
- PalmTree is evaluated with instruction sizes 64, 128, 256, and 512 on intrinsic tasks.
- Performance improves as embedding size increases, with the largest size best across all three intrinsic metrics.
- The authors recommend choosing embedding size according to hardware capacity and use 128 with a single GTX 2080Ti GPU.
C.2 Output layer configurations
The output-layer experiment compares five transformer-layer configurations at embedding size 128 and selects the second-last layer based on intrinsic evaluation results.
- PalmTree is evaluated using the last, second-last, third-last, and fourth-last transformer layers as output configurations.
- The experiment fixes the embedding size at 128 while comparing output layers on intrinsic tasks.
- The second-last layer performs best for opcode outlier detection and basicblock similarity search, so the paper uses it as output.
C.3 Context window for CWP
The CWP context-window experiment compares four window sizes, where a window of 2 uses two preceding and two following instructions as context.
- A context window of 2 includes instructions n−2, n−1, n+1, and n+2 around instruction n.
- The experiment evaluates context window sizes 1, 2, 3, and 4 in the CWP task.
- Table 12 reports intrinsic metrics for PalmTree trained with different context-window configurations.