PAVR: A Pre-Training Approach with Self-attention for Vulnerability Repair

As vulnerability detection tools evolve, more and more vulnerabilities are found. Due to the high labor cost required to fix vulnerabilities manually, the AI-based approaches are introduced. Recently, Neural Machine Translation (NMT) has been introduced for application in the field of vulnerability repairing. However, However, most of the existing NMT approaches focus on bug fixing tasks in Automated Programmed Repair (APR) tasks and lack specialization in vulnerability fixing tasks. In this paper, we present PAVR, an approach based on deep learning models combined with vulnerability-related pre-training tasks and attention enhancement. We propose three customized pre-training tasks to enable our deep learning model to acquire code generation capabilities related to vulnerability repair during the pre-training process. At the same time, we impose certain restrictions on the output based on the Decoder part, which makes the model’s attention to vulnerability types enhanced. We demonstrate that our pre-training tasks and augmented attention mechanism are effective, and we implement PAVR.