Probing the Prefill: Detecting Code Vulnerabilities via Latent Activations

Alizishaan Khatri

Source-linked AI summary

Probing the Prefill: Detecting Code Vulnerabilities via Latent Activations

Alizishaan Khatri

arXiv:2608.16970v1Aug 17, 2026cs.CRcs.AIcs.LG

TL;DR

Post-hoc defenses screen completed vulnerable code without using the generating LLM’s internal state. This paper probes frozen coding-model activations across four benchmarks, finding a recoverable vulnerability signal that averages 41.7% F1 and matches published Devign SOTA while trailing on harder benchmarks.

Problem
Existing vulnerability defenses mainly screen completed code post hoc, leaving open whether an LLM’s internal state already encodes a recoverable vulnerability signal.
Method
The study extends activation probing to code vulnerability by training MLP probes on frozen prefill activations from four LLMs evaluated across four C/C++ benchmarks.
Results
Across four benchmarks, probes average 41.7% F1, roughly match published Devign SOTA, and trail SOTA substantially on harder, more imbalanced benchmarks.
Takeaways & Limitations
The findings provide early, preliminary evidence that coding-LLM activations contain information predictive of vulnerability status, while dataset quality may constrain harder benchmarks.
Takeaways & Limitations
Evaluation uses existing corpus functions rather than model-generated completions, and in-distribution results may not predict worst-case or obfuscated performance.

Abstract

from arXiv · show

LLM-based code generation is now embedded in mission-critical pipelines, but defenses against vulnerable output remain post-hoc -- static analyzers, fine-tuned classifiers, or an LLM judge that screen completed code, ignoring the generating model's own internal state. We test a narrower, directly measurable question: when an LLM reads a piece of C/C++ code as context, do its hidden activations already carry a signal about that code's vulnerability status? We extract last prefill token activations from four LLMs (Granite-4.1-8B, Qwen3.5-9B, Qwen3.6-27B, Gemma-4-12B) across three model families and train MLP probes on these activations. We evaluate them on four function-level C/C++ benchmarks (Devign, Big-Vul, Draper VDISC, PrimeVul). Our probes achieve 41.7\% average F1 using 13.4--16.0M-parameter probes -- under 0.2\% of base-model size. On Devign, the best probe (Qwen3.5-9B, 68.8\% F1) matches the published fine-tuned-classifier SOTA (67.9\%) despite reading only a frozen, general-purpose LLM's activations; on the harder, more imbalanced benchmarks (Big-Vul, Draper VDISC, PrimeVul) probes trail SOTA substantially. This is early evidence that a coding LLM's own representation of arbitrary code is informative about that code's vulnerability status, motivating further work toward lightweight, model-native vulnerability screening.

I. INTRODUCTION

This preliminary study asks whether an LLM’s internal state contains a recoverable signal about code vulnerability status when reading code as context. It extends activation probing to function-level C/C++ vulnerability detection and examines model coverage, CWE patterns, and generation-time gating implications.

Motivation and scope: The study tests whether an LLM’s internal state encodes a recoverable vulnerability signal consistently enough to motivate generation-time gating.It uses existing human/CVE-labeled corpus functions rather than model-generated completions, so it does not test whether a model can flag its own forthcoming vulnerable output.
Novelty: The work extends activation probing from LLM safety and prompt-harm classification to code vulnerability detection.Code previously appeared only incidentally as a minority of Aegis dataset samples.
Method: Last-prefill-token activations from four LLMs across three model families are used to train MLP probes evaluated on Devign, Big-Vul, Draper VDISC, and PrimeVul.The models are Granite-4.1-8B, Qwen3.5-9B, Qwen3.6-27B, and Gemma-4-12B.
Evaluation and limitations: The paper reports precision, recall, and F1 for every model–dataset pair and discusses probe-based generation-time gating against post-hoc baselines.It also identifies known generalization failure modes of probing-based detectors.
Results: Big-Vul accuracy is uniformly high and tightly clustered across its five most frequent CWE types for all four probes.CWE-119 is the single hardest category for every model.

II. RELATED WORK · A. Post-Hoc Vulnerability Detection · B. Probing LLM Internals for Vulnerability Detection

Existing vulnerability screening is largely post-hoc, operating on emitted code or separately trained detectors rather than the generating model’s internal computation. This work instead probes hidden states during the forward pass, extending prior internal-probing approaches to candidate code read by dedicated LLMs.

A. Post-Hoc Vulnerability Detection: Production screening commonly applies static or dynamic analyzers to generated code or fine-tunes classifiers on labeled code.These approaches inspect completed outputs rather than the generating model’s computation.
A. Post-Hoc Vulnerability Detection: Post-hoc detectors are black-box to the generating model’s computation and can be evaded by surface transformations preserving the vulnerability.The passage contrasts output-level screening with access to internal representations.
A. Post-Hoc Vulnerability Detection: Detection models already encode bug-relevant semantics, although those representations are not necessarily linearly separable.This motivates probing internal state directly rather than analyzing emitted code after generation.
A. Post-Hoc Vulnerability Detection: The proposed method probes the generating model’s internal state through its forward pass instead of screening emitted code after the fact.It directly targets the model computation used while reading the code.
B. Probing LLM Internals for Vulnerability Detection: LPASS uses layerwise linear probes on a compressed LLM to guide pruning, rather than to screen code vulnerabilities.Its purpose is model compression, distinguishing it from vulnerability detection.
B. Probing LLM Internals for Vulnerability Detection: Sternfeld et al. train logistic-regression/MLP probes on prompt-end hidden states to predict whether an ungenerated completion will be secure.They also show that minimal prompt perturbations can flip the predicted completion security outcome.
B. Probing LLM Internals for Vulnerability Detection: This work differs by using deeper 6-layer MLP probes and activations from candidate code itself across four dedicated corpora.The comparison is with prompt-paired generation settings in prior work.

C. Latent-State Probing Beyond Code … A. Overview

The paper situates latent-state probing within prior work on jailbreak and malware detection, then defines vulnerability prediction from a frozen LLM’s last-prefill-token activation. Its overview uses a two-stage pipeline that extracts labeled activations and trains only a class-weighted MLP probe.

C. Latent-State Probing Beyond Code: Prior latent-state probing has addressed malicious prompts, jailbreaks, and malware, but reported probes can generalize poorly out of distribution.Chia et al. separate safe and jailbroken states, Fomin uses leave-one-dataset-out evaluation, Wang et al. report poor out-of-distribution generalization, and Ajayi et al. study malware outside LLMs.
III. PROBLEM FORMULATION: The vulnerability oracle labels a completion positive when dataset ground truth identifies one or more target vulnerability classes, such as a CWE.The oracle is defined over completions produced by an LLM from coding prompts.
III. PROBLEM FORMULATION: The probe input is h(x), the final-layer hidden state at the last prompt token immediately before generation begins.Formally, h(x) = Mhidden(x)[−1] ∈ R^d.
A. Overview: The method separates activation extraction from probe training, following a two-stage design.Each coding prompt is first run through the frozen LLM, after which an MLP is trained on the stored activations.
A. Overview: In extraction, each function passes through the frozen LLM, and its last-prefill-token activation is cached with its vulnerability label.The pipeline denotes frozen model weights and stores paired activation-label examples for subsequent training.
A. Overview: In training, an MLP probe learns from cached activation-label pairs using class-weighted binary cross-entropy, while the base LLM remains frozen.Only the probe parameters θ are updated.

B. Models

The study evaluates vulnerability probes across four coding-capable LLMs from three publishers, including two Qwen models for within-family comparison. Each probe uses one d-dimensional activation vector from the final layer’s last prefill token.

Model selection: Four LLMs spanning three publishers test whether vulnerability-relevant representations generalize beyond a single coding model.The models include Granite-4.1-8B, Qwen3.5-9B, Qwen3.6-27B, and Gemma-4-12B.
Model selection: The two Qwen models probe within-family generalization across a scale and generation jump.Qwen3.5-9B and Qwen3.6-27B are from the same family.
Activation extraction: Each example is represented by a single d-dimensional vector from the final hidden layer’s last prefill token.Activations are extracted using last-token pooling rather than a layerwise sweep.

C. Datasets · D. Probe Architecture

The study evaluates frozen LLM activations on four binary-labeled, function-level C/C++ vulnerability datasets spanning different labeling methodologies and label-quality regimes. It uses model-specific MLP probes with a shared progressively shrinking architecture and a single class-weighted binary logit.

C. Datasets: Four widely used datasets provide binary vulnerable/benign labels at C/C++ function granularity.The datasets are selected to cover varied vulnerability-detection settings.
C. Datasets: The datasets span manual expert review, CVE/patch mining, and static-analyzer-derived labeling methodologies.This variation reflects differing known label-quality regimes, including active concerns about label noise.
C. Datasets: Devign contains 27.4K rows from Linux kernel, QEMU, Wireshark, and FFmpeg, with 21.9K/2.7K/2.7K train/validation/test splits and approximately 46% vulnerable functions.Devign uses manual labels and has a reported 67.9% F1 SOTA from SNOPY.
C. Datasets: Each corpus function c replaces x, and pooled h(c) from the same contextual forward pass serves as the probe input rather than a generated completion.This operationalization treats the LLM as a fixed feature extractor and avoids post-hoc completion labeling; emitted completions remain unevaluated.
D. Probe Architecture: The probe uses progressively decreasing hidden dimensions, GELU activations, and dropout, following the referenced architecture.The architecture is shared across models except for the input dimension d.
D. Probe Architecture: A single raw logit with class-weighted binary cross-entropy replaces the referenced two-class softmax head for the vulnerable/benign task.The change adapts the output layer from a multi-class-style prompt-harm setting to binary classification.
D. Probe Architecture: Each model receives one probe with its own input width d while all other architecture choices remain identical.This design makes cross-model comparisons target representation differences rather than probe capacity.

E. Training Procedure · F. Evaluation Metrics

The probes are trained with class-weighted binary cross-entropy and checkpoint-based threshold selection, using AdamW for up to 50 epochs. Evaluation reports standard classification metrics on held-out tests, plus per-CWE accuracy where dataset schemas permit it.

E. Training Procedure: Training caches each (h_i, y_i) activation-label pair before optimizing the probe.The procedure adapts prior work from prompt-level harm prediction to completion-level vulnerability prediction.
E. Training Procedure: Class-weighted binary cross-entropy is used as the probe-training objective.
E. Training Procedure: AdamW training uses η = 2.5 × 10−4, weight decay 10−2, β = (0.9, 0.9999), and ϵ = 10−8.
E. Training Procedure: Training uses batch size 1024 for up to 50 epochs, with ReduceLROnPlateau halving η after two unimproved validation epochs.The required validation improvement is ≥10−2.
E. Training Procedure: F1 is reported at each checkpoint’s tuned threshold when available, otherwise at the untuned 0.5 cutoff.Table IV marks results using the untuned cutoff.
F. Evaluation Metrics: Precision, recall, F1, and AUC are reported for every model–dataset pair on held-out test splits.
F. Evaluation Metrics: Per-CWE accuracy is additionally reported for Draper VDISC and Big-Vul, whose schemas expose CWE labels.Draper VDISC provides per-CWE boolean columns, while Big-Vul uses regex extraction from a free-text field; Devign and PrimeVul lack schema-specific extractors.
F. Evaluation Metrics: The study does not yet report a cross-model/cross-dataset generalization matrix.The paper identifies this as a direction motivated by §VI-C and [13].

G. Experimental Setup · V. RESULTS · A. Per-Vulnerability-Class Breakdown

Using native dataset splits and GPU-only activation extraction, the study evaluates probes across 16 model–dataset pairs, averaging 41.7% F1 overall. Performance varies mainly by dataset, while Big-Vul’s class breakdown shows uniformly high accuracy that requires cautious interpretation.

G. Experimental Setup: Native train/validation/test splits were used without additional label- or CWE-based stratification.Activation extraction used 2× NVIDIA L4 and 2× NVIDIA A100 GPUs; probe training and visualization ran on CPUs.
V. RESULTS: 41.7% F1 averaged across all 16 model–dataset pairs.The probes used four models across four function-level datasets.
V. RESULTS: 62.1–68.8% F1 on Devign made it the easiest dataset and roughly matched its published 67.9% SOTA for three of four models.Big-Vul reached 44.6–59.5% F1 but remained below its reported 96.5% SOTA.
V. RESULTS: 30.5–30.9% F1 on Draper VDISC and 17.6–19.5% F1 on PrimeVul indicated difficult, model-invariant performance.These ranges persisted despite differences in base-model size.
A. Per-Vulnerability-Class Breakdown: CWE-119, with n = 4478, was hardest, whereas CWE-416 and CWE-399 were consistently easiest.The authors caution that near-ceiling accuracy may reflect CWE-independent signal or a general benign-versus-vulnerable distinction, given Big-Vul’s benign skew and high tuned thresholds of 0.919–0.945.

VI. DISCUSSION · A. Why Probe at Generation Time: The Resource-Savings Case

Generation-time probes are positioned as a resource-saving alternative to external vulnerability screening: they add little parameter cost, reuse the primary model’s activations, and avoid separate sequential services. Their comparisons to external detectors are primarily parameter-count comparisons, not experiments rerun by the authors.

A. Why Probe at Generation Time: The Resource-Savings Case: 0.16% of an 8B-parameter prompt-harm model, ’s probe matched 7B guard models at 82.7% F1 versus WildGuard’s 84.4% on BeaverTails.The cited comparison also trailed GPT-4-scale judges by only a few points.
A. Why Probe at Generation Time: The Resource-Savings Case: 13.41–16.03M parameters make the probes only 0.17%–0.06% of their respective base models.Absolute probe size barely grows with activation dimension, so the fraction shrinks as the base model scales.
A. Why Probe at Generation Time: The Resource-Savings Case: Concurrent activation probing bounds total latency by the slower of T_llm and T_probe, whereas external screening adds sequential filtering and generation stages.Because the probe is much smaller than the primary model, T_probe ≪ T_llm in practice; measured under 1 ms against 50–500 ms generation.
A. Why Probe at Generation Time: The Resource-Savings Case: External detector sizes are literature-reported rather than results from benchmarks rerun by the authors, so the comparison is parameter-count only.The representative external systems include CodeBERT, GraphCodeBERT, and instruction-tuned guard models such as Guard 3, Granite Guardian, and WildGuard.
A. Why Probe at Generation Time: The Resource-Savings Case: 8–9× smaller than CodeBERT/GraphCodeBERT and 437–597× smaller than 7–8B guard models, the probes use a lightweight MLP head over frozen activations.The head reads activations from a model the pipeline already needs for generation, unlike end-to-end fine-tuned encoder models.
A. Why Probe at Generation Time: The Resource-Savings Case: Five orders of magnitude smaller than a GPT-4-scale judge, the probes avoid deploying a second model for screening.This supports the resource argument through parameter-count reduction and reuse of the generation model.

B. Relation to Prior Probing Work

This work differs from prior probing studies in both design and evaluation metric. On Devign, probes on unmodified activations achieve 68.8% F1, matching fine-tuned-classifier SOTA and providing the first reported comparable quantitative evidence.

Design and comparability: The design differs from prior work by [7] and [11].§II specifies these differences.
Design and comparability: LPASS targets pruning, while Sternfeld et al. report completion-flip rates under prompt perturbation rather than held-out-corpus F1.Their metrics are not directly comparable to this study’s F1 evaluation.
Quantitative comparison: 68.8% F1 on Devign matches fine-tuned-classifier SOTA and is presented as the first comparable quantitative evidence from a probe on unmodified activations.The result indicates that unmodified activation probes recover a signal comparable to the published classifier baseline.

C. Threats to Generalization · VII. LIMITATIONS •

Generalization remains uncertain because probes may overfit dataset artifacts, while evaluation is limited in model, representation, dataset, baseline, architecture, adversarial, and split coverage. A shared hardest-CWE pattern across models is encouraging but remains single-dataset evidence.

C. Threats to Generalization: Probing detectors can overfit in-distribution artifacts, although all four models agree on the hardest CWE type in a single-dataset breakdown.Prior work reports that in-distribution probe accuracy can fail under distribution shift, while standard splits can overstate generalization.
VII. LIMITATIONS •: Coverage includes 4 models across 3 families, but excludes other families and closed-weight models whose activations are inaccessible.Qwen3.5-9B and Qwen3.6-27B share a family.
VII. LIMITATIONS •: The detector uses only final-layer, last-prefill-token activations and binary labels; per-CWE results are post-hoc diagnostics from Big-Vul, not risk scores.No layerwise sweep is reported.
VII. LIMITATIONS •: Dataset limitations include Big-Vul spurious-correlation risk, Draper VDISC’s imperfect static-analyzer ground truth, and PrimeVul’s severe imbalance of ∼3% vulnerable functions.Evaluation uses existing corpus functions rather than model-generated completion slices, and in-distribution results may not bound worst-case or obfuscated performance.
VII. LIMITATIONS •: No controlled head-to-head baseline study or wall-clock latency measurement was conducted; Tables II/VI rely on literature-reported comparator numbers.The omitted baselines include static analyzers, fine-tuned classifiers, and an LLM judge.
VII. LIMITATIONS •: The inherited 6-layer MLP lacks architecture ablations, while logistic-regression and non-LLM TF-IDF/AST probes were not tested.These omissions leave dataset quality versus model choice unresolved as a performance constraint.
VII. LIMITATIONS •: The study reports no adversarial or evasion analysis despite motivating the work partly by the evadability of post-hoc detectors.This leaves robustness against deliberate attacks unmeasured.
VII. LIMITATIONS •: Datasets retain native rather than shared re-stratified splits, and several Table IV cells use an untuned threshold, limiting interpretation of within-column rankings.The threshold caveat is marked by † in Table IV.

VIII. CONCLUSION

The study finds that LLM activations over code contain a predictive signal for vulnerability status. Lightweight MLP probes achieve 41.7% average F1 and match published SOTA on Devign.

VIII. CONCLUSION: The study trained MLP probes on activations from four LLMs across four vulnerability datasets.The models were Granite-4.1-8B, Qwen3.5-9B, Qwen3.6-27B, and Gemma-4-12B.
VIII. CONCLUSION: 41.7% average F1 was achieved with 13.41–16.03M-parameter probes, representing 0.06–0.17% of base-model size.These results support the feasibility of lightweight activation-based vulnerability screening.
VIII. CONCLUSION: 68.8% vs. 67.9% F1 on Devign shows that the probes match published SOTA on that benchmark.The comparison is reported for the best probe result on Devign.