Pruning Attention Heads Based on Semantic and Code Structure for Smart Contract Vulnerability Detection

Abstract

Along with the sustained occurrence of black swan events in the decentralized application ecosystem, smart contract security is a growing concern. Traditional solutions mainly rely on predefined rules, while highly accurate, require intensive manual code analysis. Machine learning methods (mainly based on BERT) leverage semantic and contextual information, but overlook crucial code structure features, which are critical for identifying vulnerability. Furthermore, the presence of useless or harmful attention heads in the BERT model leads to less robust predictions and slows down processing speeds. We propose a novel method named Pruning attention Heads based on Semantic and Code Structure (PHSCS) for Smart Contract Vulnerability Detection. Specifically, we introduce a new structure-aware pre-training programming language task, Variable Edge Prediction, which bypasses the use of data flow nodes as input and directly predicts data flow edges between variables, aiming to efficiently learn code structure while ensuring the ability to process extensive code. Additionally, we present a pruning strategy to optimize BERT, tailored to the semantic and structural peculiarities of code. By employing Taylor Expansion for evaluating attention heads' significance and guiding their pruning, iteratively refined the BERT model. Experiment results on 8 vulnerability types illustrate that the PHSCS method surpasses state-of-the-art methods.