Mixture of Experts (MoE) for Malware Classification

This repository contains a full pipeline for evaluating a multimodal Mixture of Experts (MoE) architecture designed to classify malware and benign files. The MoE model intelligently fuses embeddings generated by three distinct "expert" models to achieve state-of-the-art retrieval and classification performance using FAISS.

🏗️ Architecture

🚀 Quickstart Installation

Because this repository relies on specific machine learning libraries, we highly recommend using Conda to set up a clean Python 3.10 environment.

1. Clone the repository and enter the directory:

   git clone <your-repo-url>
   cd "moe research"

2. Create and activate a fresh Conda environment:

   conda create -n moeresearch python=3.10 -y
   conda activate moeresearch

3. Install the required packages:

   python -m pip install -r requirements.txt

4. Download the Large Datasets: Because the dataset CSV files are too large for GitHub, we host them publicly on Hugging Face. To automatically download them into your root directory, run:

   python download_datasets.py

(Note: If you plan on extracting SAFE embeddings from raw binaries, you will also need to install Radare2 on your system: brew install radare2 on macOS, or apt install radare2 on Linux).

---

📂 Project Structure & Experts

This repository contains stripped-down, streamlined versions of the expert models. We have completely removed their heavy training code and bloated assets, leaving behind only the pure inference code needed to generate embeddings for our MoE pipeline.

1. safe_final/ (SAFE Expert)

SAFE (Self-Attentive Function Embeddings) analyzes the assembly structure of a binary to understand its control flow.

• What it does: Disassembles raw binaries using Radare2, extracts the assembly instructions, and runs them through an attention-based neural network to create a mathematical embedding.
• Why it's stripped down: We omitted the original repo's massive training pipeline and datasets. We only include the inference scripts (embed_text.py).
• Important Note on Weights: Because the pre-trained SAFE weights (safe.pb) and the vocabulary file (i2v.tar.bz2) are over 200MB, they exceed GitHub's strict 100MB file limit and are not included in this repository. To get them, simply navigate into safe_final and run ./download_model.sh. This script will automatically download the required weights into a data/ folder for you!
• Original Repository: gadiluna/SAFE

2. MalConv2/ (MalConv Expert)

MalConv looks at the raw bytes of a binary file without needing to decompile or disassemble it.

• What it does: Treats the binary as a continuous stream of bytes and processes it through a convolutional neural network (CNN) to identify malicious patterns.
• Why it's stripped down: Like SAFE, we only retained the core model architecture and inference script necessary to generate the raw byte embeddings for our MoE pipeline.
• Original Repository: FutureComputing4AI/MalConv2

3. node2vec_final/ (Node2Vec Expert)

Node2Vec analyzes the mathematical shape and control flow graph (CFG) of the binary.

• What it does: Disassembles raw binaries using Radare2, extracts the directed Control Flow Graph, and runs hundreds of random walks across the graph nodes to generate a comprehensive structural embedding.
• Why it's stripped down: We merged the original separate malware and benign pipelines into a single, clean inference script (embed_graph.py) that handles the Radare2 graph extraction and Node2Vec embedding generation automatically.

---

🧠 The Final MoE Brain

At the root of the repository, you will find the scripts that act as the "brain" of the operation. These files take the individual outputs from SAFE, MalConv, and Node2Vec, and fuse them together.

• best_moe_test.pt: This is the prized possession! It contains the final, pre-trained weights for the MoE Neural Network.
• MoE_Training_dataset_form.py: The training pipeline used to train the MoE router. It aligns the embeddings from all experts and trains the Gating Network to learn which expert to trust for which specific type of malware.
• evaluate_experts_moe.py: The ultimate evaluation script. It loads the best_moe_test.pt weights, builds a separate FAISS index for each individual expert, and then builds a final FAISS index for the combined MoE model. It then runs tests to prove that the fused MoE embeddings outperform the individual experts!

🧪 Testing the Vector Database

If you want to test how the embeddings are stored and searched, check out the build_and_test_faiss/ folder! It contains scripts to simulate building a standalone vector database from an embeddings .csv file and querying it for nearest neighbors.

To run the test:

cd build_and_test_faiss

# 1. Build the FAISS database index
python build_faiss_index.py

# 2. Query the index for the nearest neighbors
python test_search.py