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

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
Downloads last month
-
Safetensors
Model size
90.4M params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Paper for NetherlandsForensicInstitute/Multilingual-ASMBERT