PaperlayerView exact revision on arXiv

Source-linked AI summary

DeepGuard: Secure Code Generation via Multi-Layer Semantic Aggregation

Li Huang, Zhongxin Liu, Yifan Wu, Tao Yin, Dong Li, Jichao Bi, Nankun Mu, Hongyu Zhang, Meng Yan

arXiv:2604.09089v1cs.SEcs.AIcs.CR

TL;DR

Code LLMs can reproduce insecure patterns, and security adaptation that relies on final-layer supervision may underuse vulnerability cues distributed across layers. DEEPGUARD aggregates upper-layer representations for security-aware training and inference. Across five code LLMs, it improves secure-and-correct generation by an average of 11.9% over SVEN while retaining functional correctness and generalizing to held-out vulnerability types.

  • Problem

    Code LLMs can reproduce insecure patterns, while common security adaptation relies mainly on final-layer signals that may be less separable for vulnerability discrimination.

    Source: §1 · pp. 1–2

  • Method

    DEEPGUARD uses attention-based aggregation of multiple upper-layer representations in multi-objective security and correctness training, with prompt-conditioned inference-time steering.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    11.9% average improvement in secure-and-correct generation over SVEN across five code LLMs, with functional correctness maintained and generalization to held-out vulnerability types.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Across five code LLMs, DEEPGUARD significantly enhances code-generation security and generalizes to held-out vulnerability types.

    Source: §6 · p. 9

    The fixed layer-aggregation strategy may not use the most security-informative depth for every backbone and input.

    Source: Limitations · p. 10

Abstract

from arXiv · show

Large Language Models (LLMs) for code generation can replicate insecure patterns from their training data. To mitigate this, a common strategy for security hardening is to fine-tune models using supervision derived from the final transformer layer. However, this design may suffer from a final-layer bottleneck: vulnerability-discriminative cues can be distributed across layers and become less detectable near the output representations optimized for next-token prediction. To diagnose this issue, we perform layer-wise linear probing. We observe that vulnerability-related signals are most detectable in a band of intermediate-to-upper layers yet attenuate toward the final layers. Motivated by this observation, we introduce DeepGuard, a framework that leverages distributed security-relevant cues by aggregating representations from multiple upper layers via an attention-based module. The aggregated signal powers a dedicated security analyzer within a multi-objective training objective that balances security enhancement and functional correctness, and further supports a lightweight inference-time steering strategy. Extensive experiments across five code LLMs demonstrate that DeepGuard improves the secure-and-correct generation rate by an average of 11.9% over strong baselines such as SVEN. It also preserves functional correctness while exhibiting generalization to held-out vulnerability types. Our code is public at https://github.com/unknownhl/DeepGuard.

1 Introduction

Code LLMs can reproduce insecure patterns, while common security adaptation relies mainly on final-layer representations that may weaken vulnerability cues. DEEPGUARD aggregates upper-layer signals for security-aware training and inference, and improves secure-and-correct generation across five code LLMs.

1 Introduction

  • Code LLMs can reproduce insecure patterns learned from public code, creating security risks in real-world development.The passage notes that developers may fail to identify flaws in AI-generated code.Source: §1 · p. 1
  • Final-layer-only supervision may miss vulnerability cues because they are distributed across layers and become less separable near the output layer.The paper's diagnostic evidence reports that probe-detectable vulnerability signals attenuate toward the final layers.Source: Figure 1 · p. 1, §1 · p. 2
  • DEEPGUARD combines attention-based aggregation of multiple upper-layer representations with multi-objective training for security and functional correctness.At inference, it computes a prompt-conditioned security bias once and applies it to generation logits.Source: §1 · p. 2
  • 11.9% average improvement in secure-and-correct generation over SVEN across models.On Qwen2.5-Coder-3B, sec-pass@1 rises from 70.47% with SVEN to 80.76%, while pass@1 is 86.65%, close to the original model.Source: §1 · p. 2
  • Evaluation across five code LLMs reports improved security performance and generalization to vulnerability types held out during training.The introduction summarizes the method's results as superior security performance and generalization across multiple models compared with baselines.Source: §1 · p. 2
  • The paper presents multi-layer diagnostics and attention-based aggregation as contributions to security adaptation beyond final-layer-only supervision.The stated contributions also include extensive evaluation of security performance and generalization across multiple models.Source: §1 · p. 2

2 Related Work

Prior security-hardening methods include inference-time interventions and training-time adaptation, but both leave room for improvement. DEEPGUARD addresses the underuse of distributed layer information with learned multi-layer aggregation for vulnerability identification and mitigation.

2 Related Work

  • Prior training-time security adaptation methods mainly use final-layer hidden states, leaving information distributed across model layers underused.The related-work discussion contrasts security-focused fine-tuning and prefix-tuning with DEEPGUARD's use of distributed information.Source: §2 · p. 3
  • Inference-time methods offer flexibility but may not correct a model's underlying insecure tendencies, whereas training-time methods adapt model behavior.Examples include prompt optimization and co-decoding at inference, and security-focused fine-tuning or prefix-tuning during training.Source: §2 · p. 3
  • DEEPGUARD proposes learned multi-layer aggregation as a regional representation for security analysis, reporting a more robust signal than final-layer-only approaches.The related-work passage frames this as the first proposed and evaluated learned multi-layer aggregation strategy for this purpose.Source: §2 · p. 3

3 DeepGuard

DeepGuard combines LoRA-based model adaptation with lightweight guided inference, using security cues aggregated from multiple upper layers rather than relying on a final-layer state alone. Its training balances security, fluency, and functional correctness, while inference adds prompt-conditioned security bias during generation.

3.1 Multi-Layer Representation Aggregation

  • Attention fuses token-wise states from the top N layers, using their mean as a query to weight layer views informative for the downstream analyzer.The resulting aggregated representation is intended to provide a stronger basis for security analysis than a single final-layer state.Source: §3.1 · p. 3

3.2 Training: Multi-Objective Adaptation

  • Training pairs vulnerable snippets with functionally equivalent secure counterparts and balances secure behavior, fluency, and functional correctness.Source: §3.2 · p. 4
  • The security analyzer scores tokens from aggregated states and a learned security embedding, then uses a margin-based contrastive loss to prefer secure variants.Source: §3.2 · p. 4
  • Next-token loss on secure examples and KL regularization toward the frozen base model distribution join the final weighted training objective.Source: §3.2 · p. 4

3.3 Inference: Guided Secure Generation

  • Guided inference combines a lightweight token prior with prompt-conditioned logit biasing, steering more strongly when the prompt appears more vulnerable to the analyzer.The token prior is a weak distributional bias, not a calibrated vulnerability estimator.Source: §3.3 · pp. 4–5
  • The fixed decoding bias avoids per-step analyzer re-evaluation and adds negligible overhead beyond standard decoding.Source: §3.3 · p. 5
  • Guided inference is a low-cost complement, not a replacement for stronger, more expensive search-time defenses such as iterative re-scoring.Source: §3.3 · p. 5

4 Experiments

Across five code LLMs, DEEPGUARD improves security-oriented metrics while largely preserving functional correctness. Ablations and robustness tests examine its objectives, aggregation, and inference-time steering.

4.1 Setup

  • The evaluation compares recent Qwen2.5-Coder, DeepSeek-Coder, and Seed-Coder models against adaptation, inference-time, prompt-based, and unadapted baselines.The protocol reports secure-pass@k as its primary utility metric and sec@kpass as a diagnostic for held-out vulnerability types.Source: §4.1 · p. 5

4.2 Main Results

  • DEEPGUARD achieves the strongest or near-strongest sec-pass@1 across evaluated models, including 80.76% versus SVEN’s 70.47% on Qwen2.5-Coder-3B.sec-pass@1 measures the probability that generated code is both secure and functionally correct.Source: §4.2 · pp. 5–6, Table 1 · p. 6
  • DEEPGUARD maintains high sec@1pass across models on 12 scenarios covering four held-out CWEs, while SVEN drops more on some models.The evaluation excludes these vulnerability types from the training dataset.Source: §4.2 · pp. 6–7, Figure 4 · p. 7

4.3 Ablation Study and Sensitivity

  • Removing any training-objective term degrades performance; removing Lsec produces the largest reported declines in pass@1 and sec-pass@1.Removing Lsec changes pass@1 from 86.59% to 64.94% and sec-pass@1 from 80.71% to 59.12%; removing Lkl harms correctness, while omitting Lgen uniformly degrades metrics.Source: §4.3 · p. 7
  • Attention-based aggregation reaches 80.71% sec-pass@1, versus 78.12% with mean pooling and 73.76% with only the final layer.Increasing aggregation depth beyond a moderate N yields diminishing returns; N = 4 is the default.Source: §4.3 · p. 7, Table 2 · p. 7

4.4 Robustness of Guided Inference

  • Guided inference can interfere with benign-task correctness, while interval-based re-scoring trades greater adaptivity for latency.On Qwen2.5-Coder-3B, removing guided inference raises pass@1 from 52.4% to 62.4%; rescoring every 64 tokens adds 38.7% latency, while every 16 tokens adds 195.8%.Source: §4.4 · p. 8, Table 4 · p. 8
  • DEEPGUARD without guided inference remains competitive on functional correctness, and the steering component can be disabled when correctness is prioritized.The reported benign-task interference mainly arises from inference-time token bias rather than training-time adaptation.Source: §4.4 · p. 8

5 Analysis

Layer-wise probes and sample-wise attention analyses find security signals distributed across upper layers rather than consistently concentrated in the final layer. Case studies and weight-sensitivity tests examine steering behavior and training trade-offs.

5.1 Corroborating the Final-Layer Bottleneck

  • In four of five models, vulnerability-discriminative signals peak at intermediate layers and decline by the output layer; Qwen2.5-Coder-7B is the exception.Peak locations range from 26% to 71% relative depth, and the cross-model variation indicates model-dependent peak depths.Source: §5.1 · p. 8
  • Attention varies across validation pairs, and intermediate layers such as L30 can receive greater weight than final layer L31.The heatmap covers 82 validation pairs; the authors relate this variability to differing CWE patterns and abstraction levels.Source: §5.1 · pp. 8–9, Figure 5 · p. 9

5.2 Case Study

  • Guided inference preserves the overall token-probability distribution with a reported KL divergence of 0.1389 while shifting risky and secure tokens in opposite directions.In the command-injection case, the token ' f' is suppressed, while secure-syntax or library-associated tokens such as ' subprocess' receive positive guidance.Source: §5.2 · p. 9, Figure 6 · p. 9

5.3 Sensitivity to Loss Weights

  • The KL weight affects functional correctness: values that are too small can reduce pass@1, while overly large values can constrain adaptation and limit security gains.Performance is comparatively less sensitive to the security weight within a reasonable range.Source: §5.3 · p. 9, Figure 7 · p. 9

6 Conclusion

DEEPGUARD addresses final-layer-only security adaptation by leveraging distributed security cues through attention-based aggregation, multi-objective parameter-efficient adaptation, and guided inference. Experiments across five code LLMs show enhanced security and generalization to held-out vulnerability types.

6 Conclusion

  • Across five code LLMs, DEEPGUARD enhances code-generation security and generalizes to held-out vulnerability types.The conclusion reports significant security enhancement without specifying a numerical result.Source: §6 · p. 9
  • DEEPGUARD combines attention-based aggregation of distributed security cues with multi-objective parameter-efficient adaptation and guided inference.Source: §6 · p. 9
  • The paper identifies reliance on final-layer hidden states as a potential limitation for security discrimination in code-LLM adaptation.Source: §6 · p. 9

Limitations

The authors identify limitations in evaluation coverage, supervision-data scalability, fixed layer aggregation, and applicability to models without hidden-state access.

Limitations

  • Evaluation uses function-level Python and C/C++ benchmarks, which do not fully capture repository-level cross-file dependencies or long-range interactions.Source: Limitations · p. 10
  • Constructing functionally equivalent vulnerable/secure pairs is costly and may limit scaling to broader vulnerability types, languages, and software domains.Source: Limitations · p. 10
  • Fixed layer aggregation aids efficiency and stability, but security-informative depth can vary across backbones and inputs.The authors suggest adaptive layer selection as a possible accuracy–latency improvement.Source: Limitations · p. 10
  • Because DEEPGUARD requires internal hidden states, its direct applicability to API-only or closed-source models is limited.Extending the approach to those settings remains open.Source: Limitations · p. 10

Ethics Statement

The authors report compliance with the ACL Ethics Policy, public access to datasets and models, and no identified significant ethical considerations.

Ethics Statement

A Details on Experimental Datasets

The appendices describe the datasets, joint security-and-correctness metrics, experimental configurations, and security-analyzer architecture used in the evaluation.

A Details on Experimental Datasets

  • DEEPGUARD training uses 1,606 Python and C/C++ programs forming 803 vulnerable/secure pairs across nine high-impact CWE categories in the MITRE Top 25.Figure 8 visualizes training and validation dataset statistics.Source: A Details on Experimental Datasets · p. 12, Figure 8 · p. 13
  • CodeGuard+ evaluates security and correctness dynamically across 18 test scenarios targeting training-set CWEs, with clearer prompts, simpler dependencies, and updated APIs.The scenarios are adapted from prior benchmarks.Source: A Details on Experimental Datasets · p. 12, Table 6 · p. 13
  • The generalization set contains 12 scenarios across four CWEs excluded from training, testing beyond the training vulnerability categories.The paper frames this evaluation as a test of security principles rather than memorization of training patterns.Source: A Details on Experimental Datasets · p. 12

B Details on Evaluation Metrics

  • secure-pass@k measures the probability that at least one of k generations is both secure and functionally correct, complementing pass@k for functional correctness.The metrics jointly assess security compliance and functional correctness, rather than treating them separately.Source: B Details on Evaluation Metrics · p. 12
  • sec@kpass measures security among functionally correct outputs and is defined as 0 when none are correct; SVEN-SR ignores functional correctness [He and Vechev (2023)].Source: B Details on Evaluation Metrics · pp. 12–13

C.1 Hyperparameters for Experiments

  • Training uses five AdamW epochs with 10% warmup; DEEPGUARD aggregates the top four layers, with security-loss and KL weights of 0.5 and 1.0.Experiments ran on NVIDIA A800 GPUs; LoRA uses rank 16 and scaling factor 32.Source: §C.1 · pp. 13–14
  • Inference generates 100 candidates per scenario at temperature 0.1 and top-p 0.95, with CodeQL used for security assessment.Source: §C.1 · p. 14

C.2 Architecture and Initialization of Security Analyzer

  • The security analyzer concatenates aggregated hidden states with a learned security embedding, then uses a three-hidden-layer MLP to predict a scalar security probability.The embedding is 128-dimensional; hidden widths are 512, 256, and 128, with 0.1 dropout after the first two layers. The paper specifies normal embedding initialization, Xavier-uniform weights, zero biases, and a sigmoid output.Source: §C.2 · pp. 14–15

C.3 Detailed Ablation Configurations

The ablations specify which training losses and inference-time steering components are removed, while the case studies contrast generated code across command injection, null-pointer, and XSS tasks.

C.3 Detailed Ablation Configurations

  • Loss ablations remove next-token loss, KL regularization, or the security contrastive objective while retaining default Guided Inference during evaluation.The variants are respectively (-) Lgen, (-) Lkl, and (-) Lsec.Source: §C.3 · p. 15

Inference Strategy Ablation

  • Inference ablations disable Guided Inference, replace prompt-based scaling with a static 0.5 coefficient, or replace learned Tstats with random normal values.These variants test steering, dynamic scaling, and the empirical prior, respectively.Source: Inference Strategy Ablation · p. 15

D Case Study: Examples of Generated Code

D.1 CWE-078: OS Command Injection

  • DeepGuard replaces shell-command construction with argument-list execution, treating dirname as data rather than executable code.The base model’s f-string construction embeds raw dirname in the command string, creating injection risk when it contains shell metacharacters.Source: §D.1 · p. 15

D.2 CWE-476: Null Pointer Dereference

  • The base model’s unchecked malloc result can make p->status dereference NULL and trigger a fatal runtime error; DeepGuard adds a validity check.The DeepGuard example handles allocation failure gracefully.Source: §D.2 · p. 16

D.3 CWE-079: Cross-Site Scripting

  • The base model sanitizes username with escape() but omits required HTML paragraph tags, whereas DeepGuard preserves sanitization and meets the test’s structure.The functional specification requires a response wrapped in <p> tags.Source: §D.3 · p. 16

E.1 Impact of Aggregated Layer Depth

The analysis varies aggregation depth and sampling temperature, interprets layer attention and token scores in an SQL-injection example, and measures inference efficiency.

E.1 Impact of Aggregated Layer Depth

  • 80.24% sec-pass@1 at N=2 versus 73.76% at N=1; increasing N from 4 to 6 adds only +0.88% in sec-pass@1 while requiring 50% more aggregation compute.The authors select N = 4 as the default; empirical overhead is 2.07%, against a theoretical ceiling of ≈4.6% for Seed-Coder-8B.Source: §E.1 · pp. 16–17, Table 9 · p. 17, Figure 9 · p. 17

E.2 Impact of Sampling Temperature

  • +11.77% secure-pass@1 when temperature falls from 0.8 to 0.1 on Seed-Coder-8B; lower temperatures also improve pass@1 and sec-pass@1.The authors standardize T = 0.1 for evaluation.Source: Table 10 · p. 17, §E.2 · p. 17

F.1 Mechanistic Interpretation: Detecting SQL Injection

  • Attention spikes at L29 and L31 on vulnerable SQL-concatenation tokens, whose security scores approach zero while neutral tokens remain above 0.6.The aggregator distributes attention across layers for standard syntax tokens.Source: Figure 10 · p. 18, §F.1 · p. 18

F.2 Inference Efficiency

  • 0.0644s per 20 tokens on Seed-Coder-8B, statistically comparable to Prompt at 0.0650s and faster than SVEN at 0.0936s.The prompt-derived bias is computed once, avoiding per-token re-evaluation; CoSec requires 0.1670s on Seed-Coder-8B, about 2.6 times DeepGuard’s latency.Source: §F.2 · pp. 18–19, Table 11 · p. 19

F.3 Analysis of Token Priors

The token-prior analysis describes a sparse distribution and interpretable token preferences associated with vulnerable patterns and defensive coding.

Discriminative Distribution

  • Tstats is concentrated around zero with long tails, indicating sparse weights that leave most syntax and variable-name tokens neutral while emphasizing a small subset.The distribution is described as selectively targeting security-critical tokens.Source: §F.3 · p. 19, Figure 11 · p. 19

Semantic Interpretation

  • The lowest Tstats scores include format (-1.00) and (f (-0.30), while subprocess (0.75) is favored over os.The authors associate the negative scores with unsafe string formatting and the positive subprocess score with process-management best practices.Source: §F.3 · p. 19, Table 12 · p. 19
  • Positive prior scores also align with control-flow and validation terms such as if, return, args, and NULL, which the paper connects to conditional logic and explicit error handling.The authors interpret these token patterns as domain-specific security knowledge.Source: §F.3 · p. 19

Scenario-level comparisons

  • Radar charts compare metric scores across specific CWE scenarios for Seed-Coder-8B and Qwen-Coder-3B.The cited captions identify scenario-level comparisons but do not state which model or scenario performs best.Source: Figure 12 · p. 20, Figure 13 · p. 21

Related papers

Loading 2604.09089v1…