Update README.md

f5e14c5 verified about 11 hours ago

8.05 kB

license: eupl-1.2
tags:
  - assembly
  - arm64
  - amd64
  - risc-v
  - i386

Model Description

ASMTransformers is a project to train and use a machine learning model to compare assembly (arm, amd, risc-v, i386) functions to a database of known functions, to aid in the process of reverse engineering.

Status

Not actively maintained after publishing

Relevant links

Github repo ASMtransformers

Version

2026-08-06

Usage

To use this model, use inference.py in our repository.

Intended use

The model has been trained and tested to be used for similarity search of assembly code. It has not been trained/tested on any other languages than arm64, amd64, risc-v or i386, nor has it been tested on other downstream tasks.

Architecture description

This model has been trained similarly to Wang et al.. We did diverge somewhat from Wang et al. their approach, since our BERT model is trained multilingual. We also experimented with different training parameters during the Masked Language Modelling (MLM). For example, we set the --mlm-prob parameter to 0.4 in accordance to Wettig et al., 2023. We have done some experiments with different mlm-probs (namely 0.15, 0.3 and 0.4) and found that 0.4 yielded the best results.

One remarkable aspect of Wang et al.'s paper, was that they came up with the concept of "Jump target prediction". As part of the MLM procedure, the model also has to predict jump targets: whenever there is a jump in the code, the model shouldn't only predict it to be a jump instruction, but it should also predict to which line of code it's supposed to jump. The idea is that this leads to better understanding of the long-term relations in the code, and as Wang et al yielded good results, we have adopted it.

Finetuning happens based on triplet loss: given three functions, of which two are similar (see also Data) and one is deviant, the model has to pick which two functions are similar to each other. We used Batched Semi Hard Triplet Loss to ensure that the training task is not too easy.

Estimated training time:

pretraining: 38 hours on 4 NVIDIA H200s
finetuning: 4 hours on 1 NVIDIA H200

Full architecture:

[
  {
    "idx": 0,
    "name": "0",
    "path": "0_ASMTransformerModule",
    "type": "asmtransformers.models.asmsentencebert.ASMTransformerModule"
  },
  {
    "idx": 1,
    "name": "1",
    "path": "1_Pooling",
    "type": "sentence_transformers.sentence_transformer.modules.pooling.Pooling"
  }
]

Pretraining parameters used:

  {
    "architectures": [
      "BertForMaskedLM"
    ],
    "attention_probs_dropout_prob": 0.1,
    "classifier_dropout": null,
    "hidden_act": "gelu",
    "hidden_dropout_prob": 0.1,
    "hidden_size": 768,
    "initializer_range": 0.02,
    "intermediate_size": 3072,
    "layer_norm_eps": 1e-12,
    "max_position_embeddings": 512,
    "model_type": "bert",
    "num_attention_heads": 12,
    "num_hidden_layers": 12,
    "pad_token_id": 0,
    "position_embedding_type": "absolute",
    "torch_dtype": "float32",
    "transformers_version": "4.12.5",
    "type_vocab_size": 2,
    "use_cache": true,
    "vocab_size": 6161
  }

Preraining parameters used:

  {
    "epochs": 19, 
    "eval_steps": 10000, 
    "batch_size": 512, 
    "gradient_accumulation_steps": 1, 
    "mlm_prob": 0.4, 
    "bf16": True, 
    "tf32": True
  }

Output

The model outputs embeddings of size 768 that can be compared using cosine similarity

Data

The dataset is built on the official Debian Repository. To obtain multiple families of assembly, we used apt to cross-build the same source package to multiple architectures. The idea is that this gives us the same functions for all four architectures. For all four architectures, these functions are compiled with different optimisation: O0, O1, O2, O3, Os and manually selected set with advanced instructions further referenced here as Oc for Optimised-Custom. This results in a maximum of 24 (6 optimisation * 4 architectures) different functions which are semantically similar. (i.e. they represent the same functionality but are written differently) In practise, it was much easier to obtain amd64 functions than riscv64 functions. Thus, not all functions have 24 semantically similar functions.

The dataset is split into a train, test and an evaluation set. This in done on source package, so all binaries and functions belonging to one source package are part of either the train or the test set, not both.

Total amount of functions per architecture

Architecture	# functions
amd64	8 202 164
i386	4 868 531
arm64	4 421 768
riscv64	3 791 434

Amount of functions per architecture per optimization

Architecture	Optimization	# functions
amd64	O0	1762442
amd64	O1	1448046
amd64	O2	1403180
amd64	O3	1308199
amd64	Oc	899892
amd64	Os	1380405
arm64	O0	918402
arm64	O1	824291
arm64	O2	790679
arm64	O3	741057
arm64	Oc	331736
arm64	Os	815603
i386	O0	1081475
i386	O1	873128
i386	O2	845943
i386	O3	816186
i386	Oc	412503
i386	Os	839296
riscv64	O0	845154
riscv64	O1	689670
riscv64	O2	670651
riscv64	O3	641428
riscv64	Oc	259691
riscv64	Os	684840

Preprocessing

Several preprocessing steps have been taken: CFGs are processed to become a flat token list suitable for a tokenizer or vocabulary builder. This procedure differs slightly between architectures, the processes can be inspected here.

Some operands have been normalized, to reduce token explosion caused by raw numeric values. These numerical values are expressed as powers of two, such that 4096 would be 2 ** 12, so #0x2^c, so would 6000, but 9000 would be #0x2^d. As a consequence, small numbers are distinguishable in the vocabulary list, but bigger numbers get one batched representation.

Finally, jump tokens are normalised to be relative JUMP_ADDR_* tokens. (* being the line number of the jump adress in the given function, rather than in the full CFG)

Performance

Performance was measured in two ways: Mean Reciprocal Rank (MRR) and Accuracy@1. These metrics are useful for this problem since we compare a given embedded function to a database of embedded functions, and then rank the database by cosine similarity. Mean Reciprocal rank shows the mean rank (1 the positive example in the database being ranked first, 0.5 when the positive example is ranked second, 0.25 when the positive example is ranked fourth, etc). I.e. an MRR of 0.66 shows that the positive example was usually ranked first or second. Accuracy@1 shows how many times the positive example ranked first. It is a useful metric, but it's also very strict, which is why we used it combined with MRR.

model	ARM64 mrr	ARM64 acc	AMD64 mrr	AMD64 acc	RISC-V mrr	RISC-V acc	I386 mrr	I386 acc	crosslingual mrr	crosslingual acc
Multilingual-ASMBERT	0.762	0.685	0.848	0.803	0.615	0.514	0.660	0.575	0.631	0.530