Diagnosing Generalization Boundaries in AI-Based Vulnerability Detection via Calibrated Embedding Distances

arXiv:2610.14296 · cs.SE, cs.LG · Submitted 2026-10-09 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Diagnosing Generalization Boundaries in AI-Based Vulnerability Detection via Calibrated Embedding Distances".

Jane: The gist: calibrated embedding distances expose hidden generalization boundaries on syntactically similar vulnerable and patched code pairs.

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So we’re looking at the paper titled "Diagnosing Generalization Boundaries in AI-Based Vulnerability Detection via Calibrated Embedding Distances." Basically, it shows we can actually see where an AI model starts failing when it sees code outside of what it was trained on, just by measuring how far apart those pieces of code are in a mathematical space.

Jane: That’s what I was trying to understand. So instead of just looking at a final "vulnerable" or "not vulnerable" label, we can look at the actual distance between the vectors representing that code.

Tom: Exactly. They found that raw pretrained CodeBERT at layer three, when you apply a whitening transformation and use a k equals one nearest-neighbour scoring method, gives us this mean out-of-distribution AUROC of zero point seven three nine <ref:2610.14296#pg1>. That’s the baseline metric for how well the geometry helps separate things.

Jane: And then there’s this proximity effect they found in research question three, where patched false positives are systematically closer to those vulnerable training neighborhoods than true negatives <ref:2610.14296#pg3>.

Tom: They also found that cross-project shifts cause the sharpest degradation in accuracy, even though temporal shifts drift geometrically farthest <ref:2610.14296#pg1>.

Lu: The whole idea is that an embedding space should organize code fragments by security semantics, and in practice it doesn't, which is what this study investigates <ref:2610.14296#pg2>.

Meng: So if the geometry doesn't match the meaning, the model isn't learning what it needs to learn, right?

The paper's summary: Tom: So we’re looking at the detailed summary of "Diagnosing Generalization Boundaries in AI-Based Vulnerability Detection via Calibrated Embedding Distances." They are showing how this specific dataset and method let us diagnose those generalization boundaries by creating pairs of vulnerable and patched code that are strictly paired across different patterns, projects, and time periods <ref:2610.14296#pg3>.

Jane: It sounds like they’ve built a very controlled testbed. They construct this paired hard-negative ID/OOD benchmark from datasets like PrimeVul and TitanVul <ref:2610.14296#pg3>.

Tom: Right, and they also cleaned and normalized that dataset quite thoroughly, including character-level safety thresholds of one thousand five hundred characters and a multi-hunk reduction step <ref:2610.14296#pg3>.

Jane: And then the core of the work is how they train CodeBERT by fine-tuning it with a joint cross entropy and supervised contrastive learning objective to reshape that latent geometry <ref:2610.14296#pg3>.

Tom: After training, they extract embeddings from the CLS token at every transformer layer, and then apply a whitening transformation to decorrelate the dominant variance directions <ref:2610.14296#pg3>.

Jane: And then they score out-of-distribution samples using a non-parametric k nearest neighbour cosine distance, where the score is defined as OODk (x) equals one minus one over k summed over the cosine of the embedding E(x) and E(ni) <ref:2610.14296#pg1>.

Tom: That scoring method gives us that distance-based OOD score, which is how we measure deviation from what the model has seen during training.

Jane: And they analyze vulnerable-neighbourhood proximity by contrasting patched false positives against true negatives using metrics like the Vuln share at k equals ten <ref:2610.14296#pg3>.

Tom: That metric gave them a Vuln share score of zero point six three six for that specific measure <ref:2610.14296#pg3>.

The paper's improvements: Tom: Now they move into the suggested improvements, which basically shows how we can actually make these diagnostic tools better by adding more layers to how we check for generalization failures <ref:2610.14296#pg3>.

Jane: It sounds like they are suggesting a dual architecture again, but with an even clearer separation between the main prediction engine and this auxiliary representation check <ref:2610.14296#pg3>.

Lu: Yeah, the main thing is that instead of just one classifier, you use a high-capacity classifier alongside that frozen intermediate encoder to act as an independent guardrail <ref:2610.14296#pg3>. That way it’s not totally reliant on one part of the system for its safety checks.

Meng: So what does that mean for us shipping models? It means we can use the final output from the main model for quick triage, but if that score looks suspicious, we immediately run a check against that intermediate embedding distance to see if it’s truly out-of-distribution <ref:2610.14296#pg3>.

Tom: Exactly. The paper implies this dual architecture allows you to get an external measure of distributional proximity that is explicitly useful when the main model's confidence alone starts to fail on high-distance inputs <ref:2610.14296#pg3>.

Jane: That sounds like a practical way for an engineer to retrofit existing scanners without having to rebuild the whole system from scratch, right?

Lu: Precisely. It lets practitioners take those out-of-distribution scores derived from intermediate layers and apply post-hoc whitening, which is a low-compute way to introduce that geometric check into production pipelines <ref:2610.14296#pg3>.

Meng: That’s good because it keeps the overhead manageable; we aren't adding a massive new model to run every time we scan code. But it does require setting up that auxiliary encoder first.

Tom: The paper also suggests using those calibrated embedding distances for runtime risk stratification, where you group the distance bins into specific zones like Safe, Boundary, and Danger <ref:2610.14296#pg3>. That lets you make selective predictions instead of just a binary yes or no.

Jane: So instead of treating every piece of code equally based on its distance score, we can use that grouping to decide what action to take—like flagging something as dangerous for immediate review.

Lu: Yes, they showed that by doing this stratification, you can get a retention accuracy of zero point seven one six on the low-distance safe zone while only getting zero point five three four accuracy in the high-distance danger zone <ref:2610.14296#pg3>.

Meng: That specific number is interesting because it shows exactly what trade-off we’re making: sacrificing some recall in the danger zone to ensure high precision where it matters most.

Tom: And they also pointed out that vulnerability semantics matter, meaning you can dynamically adjust risk assessment based on the CWE level itself, not just the distance or shift type <ref:2610.14296#pg3>.

Jane: So even if two code snippets have a similar distance score, if one is a known high-risk class like CWE-four hundred sixteen and the other isn't, we prioritize that specific failure mode.

Lu: That’s the point: vulnerability semantics modulate the boundary independently of embedding distance or shift type, meaning you can guide your engineering efforts toward the most critical classes first <ref:2610.14296#pg3>.

Meng: So we get to prioritize which types of failures we spend our time fixing based on this geometric and semantic information.

Tom: It’s about moving from a system that just says "vulnerable" to a system that says "this code is potentially dangerous, and here is the geometric reason why."

Conclusion: Tom: So we’re wrapping up on "Diagnosing Generalization Boundaries in AI-Based Vulnerability Detection via Calibrated Embedding Distances." The main point is that we can empirically observe where these AI models start failing when they see code outside of their training distribution using calibrated embedding distances.

Jane: That means we have a measurable geometric signal that tells us exactly where the model’s knowledge breaks down, which is super powerful for building more reliable security pipelines.

Lu: The implication is that we need to adopt a dual architecture—a high-capacity classifier paired with an auxiliary frozen intermediate encoder—to use those distance scores as an independent guardrail instead of just relying on one final output number <ref:2610.14296#pg3>.

Meng: It changes how we budget our compute; instead of running every piece of code through the full heavy model, we can use that intermediate layer check to dynamically stratify risk into zones like Safe, Boundary, and Danger <ref:2610.14296#pg3>.

Lalam: I’m sure about this because the results show that proximity to vulnerable neighborhoods is a geometric contributor to false positives, which is an actionable signal we can use before the alert even gets to a developer <ref:2610.14296#pg2>.

Tom: It gives us a way to introduce dynamic risk stratification into our CI/CD pipelines with negligible compute overhead, which sounds like exactly what we need for real-world deployment.

Jane: It’s about moving from just looking at the final prediction score to understanding the underlying geometry that leads to that prediction.

Lu: And future work should look beyond just text representations and try applying this geometric framework to graph and data-flow based models, which might actually capture things like lifetime dependencies in flaws <ref:2610.14296#pg3>.

Meng: That sounds ambitious, but if it helps us catch use-after-free bugs better, we’re interested in seeing how that scales without the model becoming too slow to run.

Tom: We also need to explore explicit fine-tuning objectives tailored specifically to out-of-distribution density estimation rather than just relying on cross entropy classification for better boundary definition <ref:2610.14296#pg3>.

Jane: It really seems like the goal is making the AI model explicitly aware of its own uncertainty about what it sees, instead of just trying to minimize its classification error.

Lu: Exactly, and we have to remember that this proximity effect interacts with local classifier thresholds and input truncation, so the geometric signal isn't perfect by itself <ref:2610.14296#pg3>.

Meng: So we don't get a magic fix here; we still have to tune our operational thresholds carefully based on the geometry.

Tom: Overall, the paper shows how calibrated embedding distances make the generalization boundaries of vulnerability detection models empirically observable under hard-negative conditions for "Diagnosing Generalization Boundaries in AI-Based Vulnerability Detection via Calibrated Embedding Distances."

Jane: It’s a lot to take in, but it really sets up a clear path forward for building more robust security pipelines.

KONSTANTIN IEVTUSHENKO, RALF LÄMMEL, JAN JÜRJENS

University of Koblenz

cs.SE, cs.LG

Submitted: 2026-10-09

Updated: 2026-10-09

Comments: Submitted for review

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 87/100

The gist: The gist: calibrated embedding distances expose hidden generalization boundaries on syntactically similar vulnerable and patched code pairs.

Key concepts

Generalization Boundaries
These are the limits of a model's reliable performance when encountering data it hasn't seen during training. The study used embedding distances to map these boundaries, showing exactly where the model transitions from performing well on known code to failing on new, unseen code.
Hard-Negative Samples
These are pairs of code snippets that are semantically similar but have different outcomes—one vulnerable and one patched. These pairs were used as test cases to probe the model's ability to distinguish between true vulnerabilities and correctly fixed code.
Vulnerability-Neighbourhood Proximity
This concept measures how close a patched false positive is geometrically in the embedding space to the training data representing vulnerable code. The findings show that these false positives are systematically closer to vulnerable training examples than true negatives, indicating a geometric shortcut in model learning.

Terminology

Summary

The gist: calibrated embedding distances expose hidden generalization boundaries on syntactically similar vulnerable and patched code pairs.

How it works

The study investigates generalization boundaries by operationalizing hard-negative samples using vulnerable functions and their corresponding patched counterparts The researchers evaluate two model families, CodeBERT and Qwen2.5-Coder, across three controlled shift scenarios guided by three research questions

The methodology involves several key steps to diagnose these boundaries

  1. A paired hard-negative ID/OOD benchmark is constructed from PrimeVul, TitanVul, and VulGate datasets, restricted to semantically coherent CWE classes.

  2. The dataset undergoes cleaning and normalization, including character-level safety thresholds of 1,500 characters and pair-aligned multi-hunk reduction.

  3. Model training involves fine-tuning the CodeBERT encoder with a joint Cross-Entropy (CE) and Supervised Contrastive Learning (SCL) objective to reshape latent geometry.

  4. Embeddings are extracted from the [CLS] token at every transformer layer, and a whitening transformation is applied to decorrelate dominant variance directions.

  5. Distance-based OOD scoring is computed using a non-parametric k-nearest-neighbour cosine distance to the training distribution, where the score is defined as OODk (x) = 1 - (1/k) sum(cos E(x), E(ni)).

  6. Vulnerable-Neighbourhood Proximity is analyzed by contrasting patched false positives against true negatives using metrics like the Vuln. share (k=10).

How it works

The analysis focuses on three main research questions to map representation geometry directly to model degradation

RQ1 (Embedding Space Calibration): Which representation choices improve the separability of in-distribution and out-of-distribution samples in the embedding space?

The strongest overall configuration uses raw pretrained CodeBERT at layer 3 with whitening and k = 1, achieving a mean OOD-AUROC of 0.739. This result indicates that whitened variants reach 0.739, confirming that decorrelating dominant variance directions restores metric utility.

RQ2 (Generalization Boundaries): To what extent do embedding-based distance metrics reveal generalization boundaries, and how do these boundaries behave under specific distribution shifts?

The analysis shows that accuracy drops from 0.834 to 0.479 as distance increases for CodeBERT CE+SCL when evaluated over its own embedding distance. Furthermore, cross-project shift causes the sharpest accuracy degradation, even though temporal shifts drift geometrically farthest.

RQ3 (Vulnerable-Neighbourhood Proximity): To what extent are false positives on patched code characterized by geometric proximity to vulnerable training neighbourhoods?

Patched false positives reside systematically closer to vulnerable training neighbourhoods than true negatives, with the Vuln. share (k=10) AUROC of 0.636. This proximity effect is confirmed by showing that cross-project patched samples are geometrically closer to the vulnerable centroid than to the patched centroid.

How it works

The results establish a diagnostic signal that maps representation geometry directly to model degradation without requiring auxiliary models

Classification Performance and Representation Geometry:

The decoupling between classification accuracy and OOD separability suggests that fine-tuning reshapes the embedding space toward the classification boundary at the expense of broader distributional structure needed for reliable OOD ranking. This implies that production security tools should adopt a dual-representation architecture, using a high-capacity classifier alongside an auxiliary, frozen intermediate encoder representation to act as an independent guardrail.

Distance-Based Boundaries and Model Dependence:

The analysis of CodeBERT distance and Qwen2.5-Coder confidence shows that while the LLM confidence remains comparatively high even as accuracy declines with distance, the CodeBERT-derived distance proves informative where model confidence alone fails.

Shift-Specific and Geometric Failure Modes:

The geometric signal reveals that cross-project shift is harder than its distance alone would predict because patched samples are geometrically closer to the vulnerable class centroid. This confirms that representation geometry grounds shift-induced failures directly in this geometry, showing measurable shortcut learning.

Vulnerability-Class-Dependent Failure Modes:

CWE-level results show that vulnerability semantics modulate the generalization boundary independently of embedding distance, with CWE-416 being the weakest class despite substantial training support.

How it works

The practical implications suggest a system can selectively abstain from fully automated decisions on distant code by grouping distance bins into empirical zones—Safe (Bins 1–3), Boundary (Bins 4–6), and Danger (Bins 7–10)

Practical Implications:

This approach enables selective prediction, allowing automated tools to flag vulnerabilities in familiar code regions while routing geometrically distant or entangled functions to security analysts. The system can also indicate which operational trade-off to budget for: auditing unfamiliar repositories requires reviewer capacity for inflated false alarms, whereas inspecting temporally shifted code necessitates recall-focused monitoring against silently missed defects.

The study concludes that calibrated embedding distances make the generalization boundaries of vulnerability detection models empirically observable under hard-negative conditions. The findings suggest that a dual architecture combining a high-capacity classifier with an auxiliary, frozen intermediate encoder representation is necessary for robust security pipelines. This framework allows practitioners to retrofit existing scanners directly with out-of-distribution scores derived from intermediate layers via post-hoc whitening. The ability to introduce dynamic risk stratification into CI/CD pipelines with negligible compute overhead is a key practical outcome.

The study also identifies that future work should extend this geometric framework beyond textual representations to graph- and data-flow-based models, which may better capture lifetime-dependent vulnerabilities such as Use-After-Free flaws. Additionally, future research should explore explicit fine-tuning objectives tailored specifically to OOD density estimation rather than relying solely on cross-entropy classification. Finally, the authors note that proximity to vulnerable neighbourhoods is a geometric contributor to patched false positives, but this effect interacts with local classifier thresholds and input truncation.

This study demonstrates that calibrated embedding distances make the generalization boundaries of vulnerability detection models empirically observable under hard-negative conditions. The findings suggest that a dual architecture combining a high-capacity classifier with an auxiliary, frozen intermediate encoder representation is necessary for robust security pipelines. This framework allows practitioners to retrofit existing scanners directly with out-of-distribution scores derived from intermediate layers via post-hoc whitening. The ability to introduce dynamic risk stratification into CI/CD pipelines with negligible compute overhead is a key practical outcome. This study also identifies that future work should extend this geometric framework beyond textual representations to graph- and data-flow-based models, which may better capture lifetime-dependent vulnerabilities such as Use-After-Free flaws. Additionally, future research should explore explicit fine-tuning objectives tailored specifically to OOD density estimation rather than relying solely on cross-entropy classification. Finally, the authors note that proximity to vulnerable neighbourhoods is a geometric contributor to patched false positives, but this effect interacts with local classifier thresholds and input truncation. The study concludes that calibrated embedding distances make the generalization boundaries of vulnerability detection models empirically observable under hard-negative conditions.

Improvements for AI systems

  1. Improved AI system can incorporate an auxiliary, frozen intermediate encoder representation to serve as an independent guardrail to flag out-of-distribution inputs. This dual architecture allows a high-capacity classifier (like a fine-tuned LLM) to perform primary triage while the CodeBERT layer 3 embedding provides an external measure of distributional proximity that is explicitly stated as being useful when model confidence alone would fail to flag exactly the high-distance region.

  2. Improved AI system can utilize calibrated embedding distances for runtime risk stratification by grouping distance bins into empirical zones—Safe, Boundary, and Danger—to enable selective prediction. This allows the system to selectively abstain from fully automated decisions on distant code in the highest-distance Danger Zone, yielding a 71.6% accuracy on the retained low-distance code (Safe Zone) compared to 53.4% in the rejected region (Danger Zone).

  3. Improved AI system can implement shift-specific failure mode diagnostics by analyzing group centroids and margins in the calibrated embedding space to understand why certain distribution shifts cause specific degradation. This enables developers to distinguish between temporal samples [which] push samples furthest from the training distribution and cross-project shift[s] causes the sharpest accuracy degradation, leading to targeted monitoring strategies.

  4. Improved AI system can leverage geometric proximity diagnostics for patched code by quantifying the relationship between false positives and vulnerable neighborhoods, as patched false positives reside systematically closer to vulnerable training neighbourhoods than true negatives. This provides a diagnostic signal where proximity to vulnerable training clusters provides an actionable signal to flag uncertain alerts before reaching developers.

  5. Improved AI system can dynamically adjust its risk assessment based on vulnerability semantics by analyzing CWE-level performance, recognizing that vulnerability semantics modulate the generalization boundary independently of embedding distance and shift type. This allows the system to prioritize specific failure modes, such as identifying that CWE-416 remains the weakest class, guiding where to focus engineering efforts.

Abstract

High benchmark scores in learning-based vulnerability detection mask severe degradation under distribution shifts. This paper demonstrates that calibrated embedding geometry exposes these hidden generalization boundaries on syntactically similar vulnerable and patched code pairs. An evaluation of CodeBERT and Qwen2.5-Coder across pattern, project, and temporal shifts on a strictly paired dataset reveals three main findings. First, whitened intermediate-layer embeddings with local nearest-neighbour scoring separate unfamiliar samples with a mean out-of-distribution AUROC of 0.739 (RQ1). Second, distance scores reveal a shift-specific reliability gradient where accuracy drops by over 35 percentage points and cross-project shifts cause the sharpest degradation, even though temporal shifts drift geometrically farthest (RQ2). Third, false positives on patched code systematically correlate with geometric proximity to vulnerable training neighbourhoods, increasing from 9.4% in the lowest-risk decile to 58.8% in the highest (RQ3). These calibrated distances establish a diagnostic signal that maps representation geometry directly to model degradation without requiring auxiliary models.

Sources

Related papers