PaperlayerView exact revision on arXiv

Source-linked AI summary

LPASS: Linear Probes as Stepping Stones for vulnerability detection using compressed LLMs

Luis Ibanez-Lissen, Lorena Gonzalez-Manzano, Jose Maria de Fuentes, Nicolas Anciaux

arXiv:2505.24451v1cs.CRcs.AI

TL;DR

LPASS addresses the computational cost of fine-tuning and compressing LLMs for fine-grained vulnerability detection by using linear probes before those operations. It uses probes to select pruning cut-offs and estimate later effectiveness, and its compressed BERT and Gemma models outperform the state of the art while improving efficiency.

  • Problem

    Fine-grained vulnerability detection with LLMs requires costly fine-tuning and compression, while their effects on performance are unknown beforehand.

  • Method

    LPASS uses linear probes over LLM activations to estimate post-fine-tuning and post-compression effectiveness and select a layer-pruning cut-off.

  • Results

    LPASS-based compressed LLMs outperform the state of the art in vulnerability detection, with Gemma compressed versions improving over original models in F1-score while reducing training, inference, and model size costs.

  • Takeaways & Limitations

    LPASS provides an early basis for deciding whether fine-tuning and compression are worthwhile and helps build efficient vulnerability-detection LLMs.

  • Takeaways & Limitations

    The study does not establish that LPs transfer directly to other domains, and selecting suitable probes and features is not immediate.

Abstract

from arXiv · show

Large Language Models (LLMs) are being extensively used for cybersecurity purposes. One of them is the detection of vulnerable codes. For the sake of efficiency and effectiveness, compression and fine-tuning techniques are being developed, respectively. However, they involve spending substantial computational efforts. In this vein, we analyse how Linear Probes (LPs) can be used to provide an estimation on the performance of a compressed LLM at an early phase -- before fine-tuning. We also show their suitability to set the cut-off point when applying layer pruning compression. Our approach, dubbed $LPASS$, is applied in BERT and Gemma for the detection of 12 of MITRE's Top 25 most dangerous vulnerabilities on 480k C/C++ samples. LPs can be computed in 142.97 s. and provide key findings: (1) 33.3 \% and 72.2\% of layers can be removed, respectively, with no precision loss; (2) they provide an early estimate of the post-fine-tuning and post-compression model effectiveness, with 3\% and 8.68\% as the lowest and average precision errors, respectively. $LPASS$-based LLMs outperform the state of the art, reaching 86.9\% of accuracy in multi-class vulnerability detection. Interestingly, $LPASS$-based compressed versions of Gemma outperform the original ones by 1.6\% of F1-score at a maximum while saving 29.4 \% and 23.8\% of training and inference time and 42.98\% of model size.

1. Introduction

Vulnerability detection increasingly targets identifying specific vulnerability classes, but LLM fine-tuning and compression are computationally costly. LPASS uses linear probes to estimate these effects early and guide compression decisions, evaluating BERT and Gemma across vulnerability datasets.

  • Identifying which vulnerability is present is more challenging than binary vulnerability detection because it requires fine-grained knowledge of vulnerability classes.
  • LLMs can manage complex vulnerability-detection tasks, but their millions or billions of parameters make tuning and execution costly.
  • The study aims to save resources because vulnerability volume and severity are increasing, while identifying the vulnerability type supports appropriately scaled responses.
  • LPASS uses linear probes to estimate fine-tuning and compression effects early, while also providing an initial step toward explainable compression.
  • LPASS applies probes to BERT and Gemma, using layer pruning and quantization across three datasets and 12 MITRE Top 25 vulnerabilities.

2. Preliminaries

The paper situates LPASS among model-reduction methods and introduces linear probes as interpretable classifiers over internal LLM representations. It also frames CWE as the taxonomy used to organize software vulnerabilities.

  • Pruning reduces model size by removing architecture components or individual weights and activations, with structured pruning changing the architecture.
  • Quantization reduces model size and computational requirements by mapping high-precision weights to lower-precision data types.
  • Linear probes are simple classifiers placed over hidden states to assess whether specific information is represented internally at different LLM layers.
  • CWE provides a taxonomy of software and hardware vulnerabilities through identifiers, related classes, and potential damage scores.
  • MITRE Top 25 highlights common and impactful software weaknesses that are often easy for attackers to exploit.

3. LPASS foundations

LPASS is designed for static-analysis vulnerability detection, using inexpensive code features and layerwise probes to guide pruning and estimate later model effectiveness. Its goals include early assessment, detection, and efficiency in time and memory.

  • 3.1. Overview & Use case: LPASS targets static analysis within a broader vulnerability-detection pipeline that may also include dynamic analysis.
  • 3.1. Overview & Use case: LPASS requires users to extract simple code features and compute probes, while the remaining decision values are supplied by the approach.
  • 3.1. Overview & Use case: Code features act as easily extracted proxies for vulnerability presence or absence and can include routinely used representations such as lines of code or control-flow graphs.
  • 3.1. Overview & Use case: LPASS trains probes on internal activations to predict code features, uses probe accuracy to select a pruning cut-off, and estimates vulnerability-detection effectiveness without fine-tuning or executing the LLM.
  • 3.2. Goals: The approach seeks early assessment, vulnerability identification, reduced training and inference time, and lower memory requirements.
  • 3.2. Goals: Time and memory restrictions are aligned with achieving energy savings.

4. Description of LPASS

LPASS uses layerwise probes first to choose a pruning cut-off and then to estimate post-fine-tuning or post-compression effectiveness. The estimates are calibrated with a cross-dataset correction and evaluated through leave-one-out validation.

  • 4.1. Layer pruning leveraging Linear Probes: At each LLM layer, internal activations are supplied to a probe that predicts selected code features.
  • 4.1. Layer pruning leveraging Linear Probes: Layer pruning selects the cut-off layer by comparing average probe accuracy across layers and removing upper layers beyond the chosen point.The selected cut-off minimizes the aggregate absolute loss across datasets and code features.
  • 4.2. Estimating performance for vulnerability detection: LPASS estimates precision, recall, and F1-score by combining probe accuracy on a dataset with a correction value β.The estimated metrics are post-normalized to the [0, 1] range.
  • 4.2. Estimating performance for vulnerability detection: The correction value β is derived from the difference between real LLM performance and probe accuracy on a knowledge-base dataset.The authors report that users need not repeat this operation because the resulting values generalize across datasets.
  • 4.2. Estimating performance for vulnerability detection: Estimation quality is assessed with an error measure and leave-one-out cross-validation to test generalization across datasets.

5. Evaluation

The evaluation uses BERT and Gemma with three C/C++ datasets to study multi-class detection across selected MITRE Top 25 vulnerabilities.

  • Models and datasets: BERT and Gemma are evaluated on DiverseVul, Big-Vul, and PrimeVul, using models with 24 and 18 layers respectively.The selected configurations contain 334 million and 2 billion parameters, with Gemma limited to 1,024 tokens by resource constraints.
  • Vulnerabilities and samples: The study selects 12 CWEs from MITRE’s Top 25, choosing the most represented vulnerabilities across datasets and retaining 480k samples after token-length filtering.Ten CWEs are selected per dataset, with overlap producing 12 CWEs overall; samples beyond the applicable token limits are removed.
  • Vulnerabilities and samples: Table 1 summarizes the samples available for the ten most represented CWEs within MITRE’s Top 25 across the three datasets.The datasets support multi-class classification of vulnerability types.

5.3. Code features F

The evaluation uses Cyclomatic Complexity and Halstead Difficulty as code features for probing structural information in vulnerability-detection inputs.

  • Feature selection: Cyclomatic Complexity (CC) and Halstead Difficulty (HD) are selected because they capture structural code information that is harder to predict than simple counts.CC measures the number of linearly independent paths, while HD measures operand diversity.
  • Training settings: Training uses an 80%-20% random split, limits each CWE to 5,000 training samples, and oversamples underrepresented CWEs to balance classes.The experiments use two NVIDIA consumer GPUs with PyTorch and Hugging Face; Gemma uses Galore low-rank adaptation because of resource constraints.

5.5. Metrics

The study evaluates multi-class vulnerability recognition, binary vulnerable-versus-non-vulnerable discrimination, and computational efficiency.

  • Effectiveness metrics: Accuracy, F1 Score, Precision, and Recall measure multi-class capability and binary discrimination between CWE-bearing and non-vulnerable code.For the imbalanced binary setting, F1 is used as the computed metric.
  • Efficiency metrics: Training time, per-sample inference time, GPU memory, effective parameters, and model size quantify computational and memory efficiency.Memory measurements use a batch size of 1 for training and inference.

5.6. Results

LP-based layer pruning identifies cut-off points before compression, while vulnerability-detection results show near-baseline BERT performance and generally improved Gemma performance under selected configurations.

  • Compression comparisons: Tables 2 and 3 compare baseline models with layer-pruned and quantized variants across effectiveness, time, memory, and model size.Negative effectiveness differences indicate improvements over baseline, while positive time and model-resource differences indicate improvements.
  • 5.6.1. Layer pruning: Figure 3 shows loss decreasing until the cut-off layer and increasing afterward, indicating that later hidden representations add no extra information to the embeddings.The figure reports loss for datasets D and features CC and HD.
  • 5.6.1. Layer pruning: LP losses identify kcut=15 for BERT and kcut=5 for Gemma, with BERT’s CC probe showing the clearest trend across layers.BERT loss is generally minimal around k=15, while Gemma becomes relatively constant or decreases around k=5–6.
  • 5.6.2. Vulnerability detection: BERT remains near baseline after pruning and quantization, with binary F1 degradation below 1% and only limited recall degradation in selected datasets.For DiverseVul, 8-bit quantization degrades performance by 0.32%, while the pruned version degrades by 0.59%.
  • 5.6.2. Vulnerability detection: Gemma generally improves over its baseline after compression, whereas 4-bit quantization without pruning can cause substantial degradation, including 22.38% lower accuracy on PrimeVul.Pruned Gemma configurations with rank 1,024 perform better than baseline in some DiverseVul and Big-Vul settings.

Analysis per CWE.

CWE-119 and CWE-20 are the most difficult classes, while more specific vulnerabilities such as CWE-787 and CWE-125 are easier to detect. Additional pruning degrades all metrics, supporting the selected cut-off.

  • CWE-level performance: 76%/78% precision in BERT and 76.5%/74.5% in Gemma make CWE-119 and CWE-20 the weakest common classes.These abstract classes are frequently confused with each other and with the negative class.
  • CWE-level performance: 89.2%/86.1% precision in BERT and 87%/82.5% in Gemma make CWE-787 and CWE-125 easier to detect than broader classes.Their more fine-grained, concrete operations provide clearer distinctions for the models.
  • Additional pruning: 15 and 16 additional layers removed after halving the cut-off make all model metrics worse, corroborating the selected pruning boundary.The tests use half the original cut-off: 5 for Gemma and 15 for BERT, producing new cut-offs of 2 and 7.

Corroborating kcut appropriateness.

The selected LPASS cut-off preserves or improves task performance while reducing computational requirements, and it outperforms random layer removal. LPs also estimate downstream performance cheaply, although the evidence remains bounded by the tested models, languages, and linear estimation design.

  • Cut-off validation: 20% accuracy loss in Gemma and roughly 10% to 2.5% loss in BERT under random pruning support LPASS-based layer selection.Randomly removing the same number of layers was repeated three times; Figure 4 compares those results with LPASS pruning.
  • Efficiency: 71.36% lower Gemma training time and 78% lower inference time in reported settings demonstrate substantial gains from pruning.Gemma inference memory drops by 49%, while effective parameters decrease by 52.72%.
  • Efficiency: 32% lower inference time and about 30% to over 50% lower training time accompany roughly 30% metric improvements across datasets and models.These overall efficiency results summarize the compressed configurations evaluated in Tables 2 and 3.
  • Performance estimation: 4.6% average CC error for BERT and 13.1% average CC error for Gemma show that LPs estimate post-training and post-compression performance with model-dependent accuracy.Error remains similar across LLM configurations, while Table 8 reports leave-one-out estimates across datasets.
  • Limitations: The conclusions are limited to C/C++ and the tested BERT and Gemma models, while nonlinear relationships and other compression methods remain open directions.The authors also state that direct transfer of LPASS to other domains is not yet confirmed.
  • LP cost: 0.05 seconds per BERT layer and 1.49 seconds per Gemma layer make LP computation far cheaper than fine-tuning, with 99.11% and 98.26% average time reductions.A full LP round takes up to 98.18 seconds for BERT and 142.97 seconds for Gemma on DiverseVul.

6. Related work

Prior vulnerability-detection research spans general prompting, specialized fine-tuned models, and compression, but multi-class CWE identification remains challenging. LPASS compares favorably with existing approaches while avoiding additional pre-training.

  • Prior vulnerability detection: Existing prompt-based models struggle to identify the correct CWE, motivating specialized vulnerability-detection models.The literature covers binary and multi-class settings across several languages and prompting strategies.
  • Specialized models: Prior specialized approaches include broad fine-tuned model families, line-level prediction, and CWE grouping, but binary formulations cannot distinguish vulnerability types.Some specialized models require substantial pre-training resources, including 96 hours and 80 GB of VRAM.
  • Compression approaches: Shi et al. combine genetic search and knowledge distillation, but obtain 59% binary accuracy and require extensive search, pre-training, and fine-tuning.LPASS requires no additional pre-training, addressing the resource burden described for that approach.
  • Comparative results: 81% average multi-class accuracy for BERT and 76.6% for Gemma outperform other reported approaches across datasets.On Big-Vul, the models reach 64% and 65%, while the best LPASS result reaches 82%.
  • Comparative results: 96% accuracy on Big-Vul binary classification surpasses prior reported results ranging from 65% to 94%.The cited prior maximum is 93.83%.
  • LPASS positioning: LPASS evaluates C/C++ multi-class detection across datasets, uses 1,024-token inputs for Gemma, and reports compressed models that can improve performance rather than incur only losses.The reported improvement contrasts with the accuracy-loss characterization in Shi et al. [17, 18].

7. Conclusion

LPASS uses linear probes to guide layer-pruning decisions and estimate post-fine-tuning and post-compression performance. Its two compressed LLMs outperform both the state of the art and their non-compressed versions, while broader applicability remains open.

  • Conclusion: LPASS uses linear classifier probes to select a layer-pruning cut-off and estimate post-fine-tuning and post-compression performance.The approach is presented as a way to support informed decisions about costly model compression and fine-tuning.
  • Conclusion: LPASS-based versions of two LLMs outperform both the state of the art and non-compressed versions in vulnerability detection.This result supports the suitability of probes within the evaluated setting.
  • Future work: The suitability of linear classifier probes for other models and application domains remains to be established.The authors identify this as a future research direction because their use of probes had not previously been proven for model compression.
  • Future work: For vulnerability detection, the suitability of LPASS for languages beyond the evaluated setting remains an open issue.The authors specifically identify other programming languages as an unresolved direction.
Loading 2505.24451v1…