WelkIR: Flow-Sensitive Pre-trained Embeddings from Compiler IR for Vulnerability Detection
摘要
While most vulnerability detection methods rely on source code features, such as Abstract Syntax Tree (AST), some have adopted compiler intermediate representations (IRs) for their efficient code representations. However, IR-based methods struggle with modeling long-range dependencies, exhibit limitations in semantic extraction, and suffer from inadequate model evaluation. To address these challenges, we propose WelkIR, a two-tier IR-based Transformer model for vulnerability detection, which is fine-tuned from a flow-sensitive pre-trained model on task-specific data. During pre-training, we introduce three novel pre-training tasks to capture the program semantics: Control Flow Prediction (CFP), Data Def-Use Prediction (DDP), and Data Reachability Prediction (DRP). The pre-trained model is then fine-tuned on task-specific data to refine its understanding of vulnerable code patterns, resulting in a specialized vulnerability detection model. We evaluate WelkIR through comprehensive experiments. On the Juliet benchmark, WelkIR achieves a 7.01%–23.26% improvement in macro F1-score compared to state-of-the-art baselines. We additionally construct ARVul, which, to the best of our knowledge, is the first large-scale LLVM IR-based vulnerability dataset, comprising 3,732 vulnerable functions from 217 real-world projects. On ARVul, WelkIR outperforms baselines by 9.01%–32.74% in F1-score. Ablation studies confirm that the proposed pre-training tasks contribute to WelkIR’s performance in vulnerability detection.