Learning to Focus: Context Extraction for Efficient Code Vulnerability Detection with Language Models

TL;DR

FocusVul enhances language models for code vulnerability detection by learning to select and extract semantically rich context, significantly improving accuracy and efficiency on real-world benchmarks.

Contribution

Proposes a novel framework that learns to identify vulnerability-relevant regions and extract context, improving LM-based detection without relying on commit annotations during inference.

Findings

164.04% improvement in classification performance

19.12% reduction in FLOPs

Outperforms heuristic and full-function fine-tuning methods

Abstract

Language models (LMs) show promise for vulnerability detection but struggle with long, real-world code due to sparse and uncertain vulnerability locations. These issues, exacerbated by token limits, often cause models to miss vulnerability-related signals, thereby impairing effective learning. A key intuition is to enhance LMs with concise, information-rich context. Commit-based annotations offer precise, CWE-agnostic supervision, but are unavailable during inference, as they depend on historical code changes. Moreover, their extreme sparsity, often covering only a few lines, makes it difficult for LMs to process directly. In this paper, we propose FocusVul, a model-agnostic framework that improves LM-based vulnerability detection by learning to select sensitive context. FocusVul learns commit-based annotation patterns through hierarchical semantic modeling and generalizes them to identify line-level vulnerability-relevant regions during inference. It then extracts LM-oriented context via both dependency and execution flows surrounding selected regions, yielding semantically rich inputs for effective vulnerability detection. Experiments on real-world benchmarks show that FocusVul consistently outperforms heuristic-based and full-function fine-tuning approaches, improving classification performance by 164.04% and reducing FLOPs by 19.12% on average.