TLDR: WAVES is a novel approach that uses weakly supervised learning to detect software vulnerabilities at the function level and precisely locate them at the statement level, without needing expensive, detailed statement-level labels for training. It leverages Multiple Instance Learning and a Transformer-based model with dual pooling channels to achieve state-of-the-art localization performance, significantly aiding developers in identifying and fixing code flaws.
Software vulnerabilities are a persistent and critical challenge in the world of technology. These flaws in code can be exploited by malicious actors, leading to data breaches, system compromises, and significant financial and reputational damage. While many tools exist to detect vulnerabilities, most operate at a broad level, identifying entire functions or files as problematic. This leaves developers with the arduous and time-consuming task of manually sifting through large blocks of code to pinpoint the exact vulnerable statements that need fixing.
The problem is compounded by the difficulty and high cost of creating detailed, statement-level labels for training advanced detection models. Such labeling requires expert knowledge and is a major bottleneck for developing more precise tools.
Introducing WAVES: A Smarter Way to Find Vulnerabilities
To address this critical gap, researchers have introduced a novel approach called WAVES, which stands for “WeAkly supervised Vulnerability Localization via multiplE inStance learning.” This innovative system aims to not only detect if a piece of code is vulnerable but also to precisely locate the specific lines of code responsible for the flaw, all without needing expensive, manual statement-level labels during its training phase. You can read the full research paper here.
WAVES draws inspiration from a concept in artificial intelligence known as Multiple Instance Learning (MIL). Imagine you have a collection of bags, and each bag contains several items. You know if a bag is “good” or “bad,” but you don’t know which specific item inside a “bad” bag makes it bad. MIL helps a system learn to identify the problematic items even with only bag-level labels.
In the context of WAVES, a software function is treated as a “bag,” and each statement within that function is an “item” or “instance.” If a function is labeled as vulnerable, WAVES understands that at least one statement within it must be the culprit. Conversely, if a function is safe, all its statements are considered safe. This clever reframing allows WAVES to convert the readily available function-level labels into “pseudo-labels” for individual statements, guiding its learning process.
How WAVES Works: A Dual-Perspective Approach
The WAVES system employs a sophisticated architecture, starting with a Transformer-based encoder. This component processes the raw code, breaking it down into smaller units (tokens) and understanding their relationships. Crucially, WAVES then generates two different types of “representation vectors” for each statement, offering a dual perspective on potential vulnerabilities:
- Max Pooling Channel: This channel focuses on capturing “local” suspicious information. Think of it as looking for specific keywords, variable misuses, or API calls that are commonly associated with vulnerabilities.
- Mean Pooling Channel: This channel takes a broader, “statement-wide” view, looking for patterns that might indicate issues like integer overflows or other contextual flaws that aren’t tied to a single keyword.
By combining the insights from both these channels, WAVES creates a more comprehensive understanding of each statement’s vulnerability status. These insights are then fed into classifiers, and their scores are combined to make a final prediction for each statement.
Learning Without Direct Labels
The core innovation lies in WAVES’s training strategy. Since it doesn’t have explicit statement-level labels, it uses the function-level labels to create these “pseudo-labels.” For non-vulnerable functions, it’s straightforward: all statements are pseudo-labeled as non-vulnerable. For vulnerable functions, WAVES sorts statements based on their predicted vulnerability scores and assigns the function’s “vulnerable” label to the top few (a configurable ‘k’ number) statements. This approach helps the model learn to identify the most suspicious statements within a vulnerable function, even if the initial pseudo-labels aren’t perfectly accurate.
Impressive Results in Detection and Localization
Extensive experiments on several benchmark datasets have shown WAVES to be highly effective. In terms of overall vulnerability detection at the function level, WAVES performs comparably to existing state-of-the-art methods. However, where WAVES truly shines is in its ability to localize vulnerabilities at the statement level. It achieved state-of-the-art performance across various metrics, significantly improving the precision with which vulnerable lines of code can be identified.
The research also explored how different design choices impact WAVES. For instance, combining both the max pooling and mean pooling channels generally led to the best performance, highlighting the benefit of a dual perspective. The amount of training data also played a significant role; as more data was used, WAVES became even better at both detecting and localizing vulnerabilities, even without explicit statement-level annotations.
Understanding Different Vulnerability Types
WAVES demonstrated a consistent ability to detect various types of vulnerabilities, with a true positive rate of around 56% for most Common Weakness Enumeration (CWE) types. It showed particular strength in identifying vulnerabilities like Integer Overflow or Wraparound (CWE-190), Double Free (CWE-415), and Use After Free (CWE-416), which often have more fixed and recognizable patterns. However, it found it more challenging to detect vulnerabilities related to Exposure of Sensitive Information to an Unauthorized Actor (CWE-200) and Resource Management Errors (CWE-399).
Also Read:
- AI Agents Fortify LLM-Generated Firmware for Embedded Systems
- SmartCoder-R1: Generating Secure and Transparent Smart Contracts with AI
The Human Factor: Developers Need Localization
A study involving software developers confirmed the practical value of WAVES. Most developers admitted to having limited confidence in finding vulnerabilities on their own, even when told a function was vulnerable. They expressed a strong need for automated tools that not only detect but also pinpoint the exact location of vulnerabilities. Over 70% of participants agreed that the localization information provided by WAVES helped them identify and fix vulnerabilities more quickly, underscoring the real-world impact of this research.
While WAVES represents a significant leap forward, the researchers acknowledge areas for future improvement. These include exploring dynamic ways to assign pseudo-labels and expanding datasets to handle even longer code segments, further enhancing the tool’s accuracy and applicability in complex real-world scenarios.
