Back to Blog
Probing the Prefill: Can an LLM's Own Activations Tell You Its Code Is Vulnerable?

Probing the Prefill: Can an LLM's Own Activations Tell You Its Code Is Vulnerable?

Wrynx Team
10 min read
AI Safety

Coding assistants and autonomous coding agents are quietly becoming part of the software supply chain for things that actually matter — defense systems, autonomous vehicles, robot swarms. Every one of those pipelines needs some way to catch vulnerable code before it ships. Today, that almost always means a post-hoc check: a static analyzer, a fine-tuned classifier, or a second LLM acting as a judge, all of which read the code only after it has already been generated.

At Wrynx, we've been asking a more basic question: when a coding LLM reads a function, does it already know, internally, whether that function looks vulnerable — before it writes a single token of a completion? This post walks through what we did and what we found.

The idea: read the model's mind instead of its output

The intuition draws on a growing body of interpretability work showing that LLMs encode surprisingly rich, linearly-recoverable information in their hidden states — including whether an incoming prompt is harmful. That's territory we've worked in before (our DSN-W paper on latent-state harm detection, and a follow-up reproducibility study across model families). If a model's activations already carry a "this prompt is harmful" signal that a lightweight probe can extract, we wanted to know whether a coding model's activations carry a "this code is vulnerable" signal too.

The formal target is: given a coding prompt x, an LLM M produces a completion o = M(x). Does the model's last hidden state over x predict whether o will turn out vulnerable, before o is even generated?

That's a hard thing to test directly — it requires generating completions and having reliable vulnerability ground truth for them. So we did the necessary first step instead: we swapped in existing, human/CVE-labeled functions from established benchmarks as the "context" c, and asked whether the model's representation of c alone predicts c's known label. We want to flag this distinction clearly up front, because it matters for how to read everything below: this is evidence that vulnerability status is recoverable from activations in general, not yet evidence that a model can flag its own forthcoming vulnerable output. That stronger claim is future work.

The pipeline

The method is a clean two-stage setup, adapted from our own prior activation-probing work on prompt-harm detection and repurposed for code:

Two-stage extraction and training pipeline

Stage 1 — Extraction. Each candidate function is passed through a frozen target LLM, and the last-prefill-token activation — the final hidden state at the last token of the prompt, right before generation would begin — is pulled out and cached alongside its ground-truth label.

Stage 2 — Training. A small MLP probe (6 layers, tapering from the model's hidden dimension down to a single logit) is trained on those cached, frozen activations with a class-weighted binary cross-entropy loss. The base LLM never gets touched — only the probe's ~13–16M parameters are updated.

This is a deliberately narrow instrument: one vector per function, from the final layer only, no layer sweep, and a binary label only.

Four models, four benchmarks

To check whether any signal we found was a property of "coding LLMs" broadly rather than a quirk of one model, we spanned three publishers and, via two Qwen checkpoints, one within-family generation jump:

Model Params Hidden dim Layers
Granite-4.1-8B (IBM) 8B 4096 40
Qwen3.5-9B (Alibaba) 9B 4096 32
Qwen3.6-27B (Alibaba) 27B 5120 64
Gemma-4-12B (Google DeepMind) 12B 3840 48

And four function-level C/C++ vulnerability benchmarks, which we chose deliberately to span very different labeling methodologies and difficulty regimes:

Dataset Size Vulnerable % Labeling method Published SOTA F1
Devign 27K 46% Manual expert review 67.9%
Big-Vul 217K 5.8% CVE/patch mining 96.5%*
Draper VDISC 1.27M ~9.6% Static-analyzer ensemble 60.9%
PrimeVul 236K ~3% Auto-labeling (bias-corrected) 24.5%*

(Both starred SOTA numbers come with a documented asterisk in the underlying literature — Big-Vul's 96.5% is reachable by a bag-of-words baseline with no code structure at all, and PrimeVul's SOTA collapses to 3.1% when a checkpoint trained on Big-Vul is evaluated on it, underscoring how gameable some of these benchmarks are.)

What we found: signal is real, but capped by dataset quality, not model choice

Averaged across all 16 (model, dataset) pairs, our probes hit 41.7% F1 using probes that are 13.4–16.0M parameters — well under 0.2% of the size of the base model whose activations they're reading.

Probe F1 vs published SOTA across datasets and models

A few things jump out of this chart:

  • On Devign, our probes get close to — and in one case beat — the published fine-tuned-classifier SOTA. Qwen3.5-9B's probe hits 68.8% F1 against a reported SOTA of 67.9%, despite the probe reading nothing but a frozen, general-purpose model's activations rather than being trained end-to-end for the task.
  • On Big-Vul, probes land in the 44.6–59.5% range — well short of the (likely inflated) 96.5% SOTA, but respectable given that SOTA number's own caveats.
  • On Draper VDISC and PrimeVul, something more interesting happens: performance is nearly identical across all four models, despite those models differing in size by more than 3×. Draper VDISC probes cluster tightly at 30.5–30.9% F1; PrimeVul probes cluster at 17.6–19.5% F1.

That tight clustering, independent of which model's activations are being read, is the result we find most interesting here. It's consistent with a fairly intuitive explanation: on datasets with known label noise and severe class imbalance, the ceiling on achievable performance is being set by the data, not by which LLM's internal representation the probe is trained on. We want to be upfront, though, that confirming this really is a data ceiling — rather than our probes simply underperforming an achievable target — would require a non-LLM baseline (e.g., TF-IDF or AST features) that we haven't run yet. It's on our list.

Vulnerability class breakdown

For Big-Vul, where CWE labels are available, we broke classification accuracy down by the five most frequent vulnerability types:

CWE Type (n) Granite Qwen3.5 Qwen3.6 Gemma
Benign / No CWE (6424) 98.8 98.6 98.9 98.3
CWE-119 — buffer bounds (4478) 97.9 97.5 97.7 97.7
CWE-20 (3388) 98.1 97.8 98.0 97.8
CWE-399 (2889) 98.8 98.6 98.7 98.6
CWE-264 (2287) 98.1 98.3 98.3 98.6
CWE-416 (1559) 99.0 98.8 98.8 98.3

Accuracy is uniformly high (97.5–99.0%) and tightly clustered — but this is on a heavily benign-skewed dataset with a high tuned decision threshold, so near-ceiling accuracy is consistent both with genuine CWE-independent vulnerability signal and with a general shortcut correlated with vulnerability. We can't fully distinguish those from accuracy alone. One notable, model-agnostic pattern did stand out though: CWE-119 (buffer-bounds violations) is the single hardest category for all four models, while CWE-416 and CWE-399 are consistently the easiest — the same ranking shows up regardless of which LLM is doing the representing.

Why we think probing is worth pursuing over just running a bigger judge model

The case for this approach is fundamentally about resource cost. A probe adds a vanishingly small parameter overhead relative to deploying a second model as an external judge:

Probe parameter count vs external detectors

Our probes are roughly 8–9× smaller than fine-tuned CodeBERT/GraphCodeBERT classifiers, 400–600× smaller than 7–8B guard models, and about five orders of magnitude smaller than a GPT-4-class judge (to be clear: those comparison figures are literature-reported parameter counts — we haven't benchmarked them head-to-head ourselves).

Beyond raw parameter count, there's a latency and infrastructure argument: an external judge is additive — a separate model, often a separate GPU service, that runs sequentially after generation and adds a network round-trip. A probe reading the primary model's own activations is concurrent with generation and needs no separate infrastructure. In our own prior work on this same probing technique for prompt-harm classification, we measured probe inference at under 1ms against 50–500ms generation times — effectively free by comparison. And because the probe reads prefill-time activations, it can in principle flag risk before a vulnerable completion is even fully written, unlike a filter that has to wait for the whole response.

There's also a deployment angle we're excited about for tactical or air-gapped settings — embedded toolchains in autonomous vehicles or robot swarms, for instance — where routing generated (possibly sensitive) code out to an external guard model over a network is itself a liability, or simply infeasible under field compute/connectivity constraints. A same-process probe sidesteps that entirely, though we haven't tested in such an environment yet.

What we think this study is — and isn't — evidence for

We tried to be explicit and disciplined about scope while writing this up, and we want to carry that same honesty into this post rather than oversell it:

  • This tests recoverability from existing labeled code, not self-flagging of generated output. Every dataset here pairs a corpus function with its ground-truth label — none of them pairs a prompt with a model-generated completion. So this is direct evidence that vulnerability status is linearly-ish recoverable from a coding LLM's representation of arbitrary code. It's necessary, but not sufficient, groundwork toward the more ambitious claim we're actually after: that a model could flag its own about-to-be-generated vulnerable code.
  • No non-LLM baseline yet. Our hypothesis that dataset quality (not model choice) caps performance on the harder benchmarks is plausible and consistent with the tight cross-model clustering we saw, but confirming it needs a baseline like TF-IDF or AST features trained the same way. That's next on our list.
  • No architecture ablation. The 6-layer MLP probe design is inherited wholesale from our earlier prompt-harm probing work, not independently justified for code. We haven't tested whether a simple linear (logistic regression) probe does just as well — which matters, because "recoverable by an MLP" is a different and weaker claim than "linearly separable."
  • No adversarial or obfuscation robustness testing, despite our own motivation citing the evadability of post-hoc detectors as the reason to look at activations in the first place. We know this is a gap.
  • No cross-dataset generalization matrix. Prior work on probing-based detectors for malicious prompts has shown they can look accurate in-distribution and fail hard under real distribution shift. The one encouraging sign we found — all four models agree on CWE-119 being the hardest category — is single-dataset agreement, not the cross-dataset test that would actually address this concern.

This is a workshop-scale study, and we've tried to frame it that way rather than as a finished system. In the paper itself we lay out five concrete follow-ups — a linear-probe ablation, a non-LLM baseline, a model-generated-completion pilot, an adversarial-robustness check, and a cross-model/cross-dataset generalization matrix — as the controls we still need before we'd stand behind the stronger claim.

Where we land

The headline result — a probe roughly 1/600th the size of a typical external guard model, reading only a frozen coding LLM's own last-prefill activation, matching published fine-tuned-classifier SOTA on Devign — is one we find genuinely exciting as an existence proof. It adds code vulnerability detection to the growing list of properties (alongside prompt harm, jailbreak status, and others) that appear to leave a recoverable trace in an LLM's hidden states without any task-specific fine-tuning of the base model.

Whether that trace is strong enough, general enough, and robust enough to actually gate generation in a production pipeline is still an open question for us. We're actively working through the follow-up list above, and we'll be writing about what we find as we go.

---

Paper: Alizishaan Khatri, "Probing the Prefill: Detecting Code Vulnerabilities via Latent Activations," arXiv:2608.16970 [cs.CR], August 2026. Questions or collaboration ideas? Reach us at research@wrynx.com. Read Full Paper · Read the paper on our Research page →

Citation

To cite this work, please use:

@article{khatri2026probing,
  title   = {Probing the Prefill: Detecting Code Vulnerabilities via Latent Activations},
  author  = {Khatri, Alizishaan},
  journal = {arXiv preprint arXiv:2608.16970},
  year    = {2026},
  eprint  = {2608.16970},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CR},
  url     = {https://arxiv.org/abs/2608.16970}
}