README.md

---
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 (ARM64, AMD64, 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](https://github.com/NetherlandsForensicInstitute/asmtransformers)

# Version
TODO: calver of new release

# Usage
TODO: do we refer to the inference file of our repo? 

# 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

**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"
    }
  ]
```

**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
  }
```

Note that during pretraining, the --mlm-prob parameter has been set to 0.4. 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, in accordance to [Wettig et al., 2023](https://aclanthology.org/2023.eacl-main.217.pdf).

**Estimated time it has cost to train:**

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

# 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](https://wiki.debian.org/DebianRepository). 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](https://github.com/NetherlandsForensicInstitute/asmtransformers/tree/main/asmtransformers/asmtransformers/preprocessors).

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|