Vulnerability Detection via Topological Analysis of Attention Maps

TL;DR

This work addresses vulnerability detection in source code by leveraging the topology of attention maps from a BERT-based model. It constructs attention graphs from transformer attention matrices, applies persistent homology to obtain topological features across a filtration, and uses simple classifiers to detect vulnerabilities. On the Devign dataset, topology-based models surpass the pre-trained CodeBERTa baseline but do not reach the performance of a fine-tuned CodeBERTa, suggesting that attention-head topology encodes meaningful vulnerability signals yet benefits from task-specific fine-tuning. The results demonstrate that semantic information captured in attention heads is extractable via topological data analysis, motivating future work that blends semantic features with structural/topological cues and explores advanced PH techniques.

Abstract

Recently, deep learning (DL) approaches to vulnerability detection have gained significant traction. These methods demonstrate promising results, often surpassing traditional static code analysis tools in effectiveness. In this study, we explore a novel approach to vulnerability detection utilizing the tools from topological data analysis (TDA) on the attention matrices of the BERT model. Our findings reveal that traditional machine learning (ML) techniques, when trained on the topological features extracted from these attention matrices, can perform competitively with pre-trained language models (LLMs) such as CodeBERTa. This suggests that TDA tools, including persistent homology, are capable of effectively capturing semantic information critical for identifying vulnerabilities.