🤖 machine learning

Self-Supervised Representations for Binary Program Clustering: From Empirical Study to Retrieval-Augmented Learning

This paper presents the first systematic investigation of self-supervised and tabular representation learning for binary program clustering, identifying VIME as a new state-of-the-art method and proposing VIME-R, a retrieval-augmented variant that significantly improves malware clustering performance.

Original authors: Martin Mocko, Daniela Chudá

Published 2026-08-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Martin Mocko, Daniela Chudá

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine the internet as a giant, bustling city where millions of new "digital residents" arrive every single day. Most are friendly neighbors, but some are sneaky intruders trying to steal, break, or cause chaos. These intruders are called malware. In 2024 alone, over 80 million new ones showed up. To keep the city safe, security guards (cybersecurity experts) need to sort these millions of files into groups. If they can spot that a new file belongs to the same "family" as a known bad actor, they can stop the threat instantly. This sorting process is called clustering.

Usually, sorting is easy if you have a list of names (labels) telling you who is who. But in the real world, guards often have to sort a massive pile of mystery files without any names at all. This is where Self-Supervised Learning (SSL) and Tabular Representation Learning (TRL) come in. Think of SSL as a student who learns by looking at two slightly different photos of the same object and guessing they are the same thing, without a teacher telling them the answer. TRL is the specific art of teaching computers to understand "spreadsheet data" (rows and columns of numbers) rather than pictures or text. The big question researchers have been asking is: Can we teach computers to sort these mystery malware files into perfect families just by looking at their numbers, without any help from a human teacher?

The Mystery of the Mystery Files

In this study, researchers Martin Mocko and Daniela Chudá decided to play detective with a massive pile of digital files. They wanted to see if modern AI tricks, usually used for recognizing cats in photos or understanding sentences, could be repurposed to group malware files together. They used two huge, public datasets of Windows files called Ember and Bodmas. These datasets are like giant libraries containing hundreds of thousands of files, each described by a long list of numbers (features) that act like a fingerprint.

The researchers split their investigation into two phases, like a two-part mystery novel.

Phase 1: The "What If" Test
First, they wanted to know the absolute best possible score. Imagine a teacher giving the AI the answer key. They took famous AI models designed for images (like BYOL, SimSiam, Barlow Twins, and VICReg) and forced them to learn by comparing files that they knew belonged to the same family. This was a "supervised" test to see how high the ceiling could go.

The results were a mix of heroes and flops. BYOL and SimSiam turned out to be the superheroes. When given the answer key, they performed almost as well as a fully supervised model (one that knows every single family name), achieving a "Homogeneity" score (a measure of how pure the groups are) of nearly 99% for general categories and over 84% for specific families. However, Barlow Twins and VICReg stumbled badly, performing worse than even the simplest baseline methods. It turns out, not all "smart" AI models are equally smart when you switch from pictures to spreadsheets.

Phase 2: The Real Challenge
Next, the researchers removed the answer key. This is the real-world scenario: the AI has to sort the files without knowing any family names. They tested the best models from Phase 1 alongside newer methods designed specifically for spreadsheet data, like VIME, SCARF, and SwitchTab.

Here, the old image-based models (BYOL and SimSiam) struggled when forced to use the "corruption" tricks meant for spreadsheets. They couldn't beat the simple baselines. But VIME (Value Imputation and Mask Estimation) stepped up. It worked by randomly hiding parts of a file's data and trying to guess what was missing, learning the patterns in the process. VIME proved to be the new champion, beating the strong traditional baselines (like PCA and Autoencoders) and setting a new state-of-the-art record.

The Twist: VIME-R
The researchers weren't done yet. They noticed that VIME was guessing missing data by looking at random files from the whole library. But what if it only looked at files that were already very similar to the one it was studying? They invented VIME-R (Retrieval-Augmented). Instead of guessing from the whole crowd, VIME-R asks, "Who are the 100 neighbors most like this file?" and uses their data to fill in the blanks.

This simple change was a game-changer. VIME-R didn't just beat the other unsupervised methods; on the Ember dataset, it actually surpassed the "supervised ceiling" established by the best-performing supervised model in that specific test, SimSiam. While a fully supervised MLP classifier scored higher on the Bodmas dataset, on Ember, VIME-R achieved a family Homogeneity of 77.48%, beating the SimSiam score of 76.53%. This means that by using a clever "neighborhood" strategy, the unsupervised AI learned to group the malware families better than the supervised model that was explicitly told the family names in that specific context!

What This Means for the Future

The paper suggests that while some popular AI models (like Barlow Twins and VICReg) might not be the right tools for sorting malware spreadsheets, others (like BYOL and SimSiam) have huge potential if we can figure out how to create the right "pairs" of files for them to learn from.

Most importantly, the study proves that retrieval-augmented learning is a powerful new direction. By teaching the AI to learn from its closest neighbors rather than the whole world, we can get incredibly accurate results without needing expensive human labels. The researchers found that VIME-R improved the clustering quality by 2.7% to 5.8% over the previous best method.

While the study is limited to static features (like looking at a file's blueprint rather than watching it run) and specific datasets, the results are a strong signal. They show that with the right tricks, computers can learn to organize the chaotic world of malware on their own, potentially helping security guards spot new threats faster than ever before. The code for these methods will be made available, inviting others to build on this discovery.

Technical Summary: Self-Supervised Representations for Binary Program Clustering

Problem Statement

The paper addresses the critical challenge of binary program clustering, defined as the task of clustering all incoming executable files (both malware and benign) without prior knowledge of their labels. This is distinct from standard malware clustering, which often focuses only on malicious samples. The goal is to improve the quality of these clusters to aid in threat discovery, family analysis, and the creation of compressed representative datasets.

The core difficulty lies in the domain of tabular data (specifically static features extracted from Windows PE binaries). While Self-Supervised Learning (SSL) and Tabular Representation Learning (TRL) have achieved breakthroughs in computer vision and NLP, their application to binary program clustering remains largely unexplored. A primary barrier is the lack of effective data augmentation strategies for tabular data. Unlike images or text, tabular data lacks inherent sequential correlations or geometric structures that allow for simple perturbations to create "positive pairs" (semantically similar samples) required by contrastive and non-contrastive SSL methods.

Methodology

The research is structured into two distinct experimental phases using the public Ember and Bodmas datasets. The primary evaluation metric is Homogeneity, measured on both class labels (malware vs. benign) and malware family labels.

Phase 1: Establishing the Supervised Ceiling

The authors investigate the theoretical performance ceiling of popular SSL methods when they are provided with supervised positive pairs (pairs generated based on known malware family labels).

Models Tested: BYOL, SimSiam, Barlow Twins, and VICReg.
Adaptation: Vision-based SSL architectures were adapted for tabular data by replacing CNN backbones with fully connected (FC) layers.
Pair Generation: Positive pairs were created deterministically based on sample labels (samples with the same family label were paired).
Baseline: A Multi-Layer Perceptron (MLP) trained on family labels served as a fully supervised ceiling benchmark.

Phase 2: Purely Unsupervised Learning

This phase evaluates methods in a realistic scenario where no labels are available during training.

TRL Models: VIME, SCARF, and SwitchTab. These methods rely on data corruption (masking and imputing features) rather than explicit positive pair generation.
- VIME: Corrupts features by sampling from the global empirical marginal distribution.
- SCARF: Uses a similar corruption strategy but optimizes with InfoNCE contrastive loss.
- SwitchTab: Uses zero-corruption and an asymmetric autoencoder.
SSL Adaptation: The best-performing SSL models from Phase 1 (BYOL and SimSiam) were re-evaluated using VIME's corruption strategy as a proxy for data augmentation.
Baselines: Strong unsupervised baselines including PCA, Autoencoders, and UMAP.

Proposed Innovation: VIME-R

Informed by Phase 2 results, the authors propose VIME-R (Value Imputation and Mask Estimation with Retrieval).

Motivation: Standard VIME corrupts features using values from the global marginal distribution, which may introduce noise by mixing semantically distant samples.
Mechanism: VIME-R restricts the corruption pool. Instead of sampling from the global distribution, corrupted feature values are sampled from the local empirical marginal distribution of the $k$ -nearest neighbors (where $k=100$ ) of the target sample.
Goal: To generate more informative training pairs that preserve the local manifold structure of the data without requiring labels.

Key Results

Phase 1: Supervised Ceiling

Performance Gap: BYOL and SimSiam achieved performance comparable to fully supervised MLPs when trained on supervised pairs.
- On Bodmas: BYOL achieved 99.42% class Homogeneity and 84.92% family Homogeneity.
- On Ember: BYOL achieved 91.36% class Homogeneity and 75.87% family Homogeneity.
Underperformers: Barlow Twins and VICReg significantly underperformed, failing to beat strong baselines (PCA/Autoencoder) on the Bodmas dataset and showing high variance.
Ceiling Established: The study establishes a performance ceiling of ~99% class Homogeneity and ~85% family Homogeneity for supervised pair learning on these datasets.

Phase 2: Unsupervised Performance

TRL Dominance: VIME established a new state-of-the-art (SOTA) for unsupervised binary program clustering, outperforming PCA, Autoencoders, and UMAP.
- On Bodmas: VIME achieved 87.85% (max) class Homogeneity and 80.87% (mean) family Homogeneity.
- On Ember: VIME achieved 66.80% class Homogeneity and 74.36% family Homogeneity.
SSL Failure in Unsupervised Mode: When BYOL and SimSiam were forced to use VIME's corruption strategy (as a substitute for positive pairs), they failed to match the performance of unsupervised baselines like PCA, confirming that these SSL methods require specific, high-quality positive pairs to function effectively.
Gap to Supervised: VIME (unsupervised) came remarkably close to the supervised ceiling. On Ember, VIME's family Homogeneity (74.36%) was only 2.2% below the supervised SimSiam ceiling (76.53%).

VIME-R Results

SOTA Achievement: VIME-R further improved upon VIME, achieving new SOTA results.
- Bodmas: Improved family Homogeneity by 2.96% over VIME (83.83% vs 80.87%).
- Ember: Improved family Homogeneity by 3.9% over VIME (77.48% vs 74.36%).
Surpassing Supervision: Notably, on the Ember dataset, VIME-R's family Homogeneity (77.48%) surpassed the supervised ceiling established by SimSiam in Phase 1 (76.53%).

Contributions and Significance

The paper claims the following contributions:

First Systematic Investigation: It is the first study to systematically evaluate SSL and TRL methods specifically for the task of binary program clustering.
Performance Ceiling: It establishes a quantitative performance ceiling for binary program clustering using supervised pair learning, providing a benchmark for future augmentation strategies.
New SOTA: It demonstrates that unsupervised TRL methods (specifically VIME) can outperform strong traditional baselines (PCA, Autoencoders, UMAP) and approach supervised performance levels.
Retrieval-Augmented Learning: The proposal of VIME-R introduces a novel retrieval-augmented corruption strategy. The paper argues that this method bridges the gap between domain-agnostic corruption and supervised learning, proving that locality-aware augmentation can yield representations superior to those derived from explicit family-label supervision in certain contexts.

Conclusion

The study concludes that while standard vision-based SSL methods struggle without tailored data augmentations, Tabular Representation Learning (TRL) offers a viable path forward for malware analysis. Specifically, the retrieval-augmented corruption strategy (VIME-R) represents a promising direction, enabling unsupervised models to achieve clustering quality that rivals or even exceeds fully supervised baselines on family-level discrimination. The authors emphasize that this approach is particularly valuable for automated malware analysis where labeled data is scarce.