PaperlayerView exact revision on arXiv

Source-linked AI summary

Interpreting and Steering for Safe and Correct Code Generation

Hao Yan, Ziyu Yao

arXiv:2608.30025v1cs.AI

TL;DR

LLMs often generate vulnerable code, while the internal mechanisms distinguishing safe from vulnerable continuations remain comparatively understudied. The paper builds contrastive code-pair data, localizes causally relevant components, and introduces training-free DUOSTEER, which combines safety and correctness steering. Across five CWEs, DUOSTEER reduces vulnerability by an average of 26.9% and improves functional correctness by 7.5%, outperforming prompting and supervised fine-tuning baselines and replicating on Qwen-2.5-Coder-7B-Instruct.

  • Problem

    The internal computation distinguishing safe from vulnerable code generation has received comparatively little attention, motivating causal analysis of its components and inference-time vulnerability mitigation.

  • Method

    The paper constructs 9,342-pair CODESEC-PAIRS, localizes safety-related layers and heads with probes and causal patching, and applies separate safety and correctness steering directions at causal heads.

  • Results

    26.9% average vulnerability rate reduction and 7.5% functional correctness improvement are achieved by DUOSTEER across five vulnerability types, outperforming prompting and supervised fine-tuning baselines.

  • Takeaways & Limitations

    Causally localized, two-direction steering outperforms single-vector steering, prompting, and supervised fine-tuning on both Llama-3.1-8B-Instruct and Qwen-2.5-Coder-7B-Instruct.

  • Takeaways & Limitations

    The study covers two 7–8B open-weight instruct models, five Python CWE classes, and isolated function stubs; broader scales, languages, CWEs, and contexts require further re-collection and causal identification.

Abstract

from arXiv · show

Large language models (LLMs) frequently generate source code containing vulnerabilities, yet little work studies the internal mechanisms that distinguish safe from vulnerable generation in them. In this work, we systematically perform a mechanistic interpretation of LLMs, aiming at both understanding how code safety-vs-vulnerability is represented or driven by components in an LM and turning the insights into actionable steering strategies to encourage safer code generation. To this end, we introduce CodeSec-Pairs, a dataset of 9,342 Python safe-and-vulnerable contrastive code pairs, sampled from Llama-3.1-8B-Instruct. Utilizing the dataset, we explore approaches to localize layers and attention heads that relate to code safety, and further experiment with different steering strategies for inference-time vulnerability reduction. In particular, we propose DuoSteer, a double-steering approach that simultaneously applies safety and code-correctness steering to attention heads. In experiments over five vulnerability types, DuoSteer leads to an average of -26.9% vulnerability rate reduction and +7.5% functional correctness improvement, which outperforms not only other steering variants but also prompting and supervised fine-tuning baselines. The advantage also replicates on Qwen-2.5-Coder-7B-Instruct with another 2,500 contrastive pairs sampled from that model.

1 Introduction

This work examines the internal components associated with safe versus vulnerable code generation and turns that analysis into training-free steering interventions. It introduces CodeSec-Pairs, compares localization and steering approaches, and proposes DUOSTEER to jointly improve safety and correctness.

  • LLMs frequently generate code containing CWE-categorized security vulnerabilities, while prior mitigations mainly operate at the input/output boundary.
  • The study constructs CODESEC-PAIRS with 9,342 safe-and-vulnerable Python code pairs sampled from Llama3.1-8B-IT.Responses are sampled under benign and vulnerability-eliciting prompts and paired within and across prompts.
  • Linear probes localize decodable safety information, whereas causal patching identifies attention heads that drive the safe-vulnerable distinction.At the head level, probe and causal rankings diverge sharply, with |ρ|≤0.072 across every studied CWE.
  • All evaluated steering variants exhibit a safety-correctness trade-off, with stronger vulnerability reduction costing functional correctness.Probe-targeted head steering underperforms layer-level steering by 3.4% and causal-targeted head steering by 14.6% in average vulnerability reduction.
  • DUOSTEER applies separate safety and correctness directions at their respective causal heads, improving over single-vector steering and other baselines.The two-direction intervention is training-free and is reported to alleviate the safety-correctness trade-off without retraining.

2 Related Work

Related work spans security mitigations, mechanistic interpretability, and training-free probing and steering. This paper extends these lines by causally targeting code vulnerability and composing separately localized safety and correctness directions.

  • Prior work measures vulnerable code generation across CWE classes and develops fine-tuning, prompt-level, and retrieval-augmented mitigations.
  • Mechanistic interpretability explains transformer behaviors through circuits, features, causal mediation, activation patching, sparse autoencoders, and attribution graphs.
  • Existing probing and steering methods modify model behavior without fine-tuning, often targeting tonal axes whose targets share surface structure across inputs.
  • This work targets code vulnerability without shared surface tokens and additively composes a separately localized correctness direction with safety.The paper identifies this composition as new to the steering literature.

3 Methodology

The methodology builds contrastive safe/vulnerable code pairs, localizes safety signals with probes and causal head knockout, and tests inference-time steering, including DuoSteer’s joint safety-correctness intervention.

  • CODESEC-PAIRS: CODESEC-PAIRS contains matched safe and vulnerable Python generations for studying code-security representations and interventions.Intra-prompt pairs train probes and steering vectors, while cross-prompt pairs augment causal patching.
  • CODESEC-PAIRS: Five CWEs cover path traversal, code injection, resource exhaustion, unsafe deserialization, and security-check bypass patterns.Four CWEs come from the Top-25 list, with CWE-295 added for vulnerability diversity.
  • Component localization: Linear probes measure where safe-versus-vulnerable information is decodable, while causal head knockout tests which heads drive the model’s preference.Knockout zeroes a head across response positions and evaluates changes in safe-versus-vulnerable continuation preference.
  • Representation steering: Mean-difference steering uses safe-minus-vulnerable representation vectors, normalized by activation variability and scaled by strength α.The study compares layer-level and probe- or causally selected head locations across multiple head counts and strengths.
  • DuoSteer: DuoSteer simultaneously injects safety and correctness directions at independently selected causal headsets to reduce vulnerability while retaining functional correctness.The correctness signal is built from an expanded pool because safe-only pairs provide too few reliable safe-and-correct examples.
  • DuoSteer: Safety and correctness headsets overlap only partially, with best configurations ranging from 2/32 to 25/64 shared heads.At overlapping heads, the two steering vectors are added; their mean cosine is approximately −0.16.

4 Experimental Setup

Experiments primarily use Llama-3.1-8B-Instruct and held-out SecCodePLT questions across five CWEs, with Qwen-2.5-Coder-7B-Instruct used for replication. Evaluation measures detectable vulnerability, functional correctness, and their joint score.

  • Models and data: Llama-3.1-8B-Instruct is the primary model, evaluated on held-out SecCodePLT questions covering five CWEs.Qwen-2.5-Coder-7B-Instruct is additionally tested for replication.
  • Metrics: Vulnerability rate V is the fraction of generations flagged by CodeQL for the target CWE, while correctness rate C is judged by GPT-4.1.V is interpreted as detectable-vulnerability rate because static analysis may have false negatives.
  • Benchmark scope: SecCodePLT was later renamed SeCodePLT and expanded to additional languages, but the five evaluated CWE questions remained unchanged.The study retains the original benchmark name because that was the release used in its experiments.
  • Metrics: The joint safety-correctness score is C(1−V), combining higher functional correctness with lower detectable vulnerability.The study also validates correctness with execution-based unit tests on three CWEs that provide them.

5 Where Vulnerabilities Are Linearly Encoded vs. Where They Are Driven

Safety information is broadly linearly accessible in residual-stream layers, but head-level encoding varies sharply by CWE. Causal influence is largely dissociated from probe ranking, so probe localization alone is insufficient for steering.

  • Linear encoding: 71–87% layer-level probe accuracy across every CWE shows that safe-versus-vulnerable information is linearly accessible in the residual stream.The best layer is typically ℓ∈[9, 11], except CWE-295, which peaks at the final layer.
  • Linear encoding: Head-level encoding varies from zero CWE-022 heads above 0.65 accuracy to 960 such CWE-295 heads.CWE-079, CWE-094, and CWE-502 fall between these extremes.
  • Causal localization: |ρ|≤0.072 between probe and causal effects on every CWE indicates that probe rank does not predict causal influence.The most causally important head appears at probe rank 64–249 among the top 256 candidates.
  • Causal localization: Causal safe- and vulnerability-promoting heads are scattered across the head grid rather than concentrated near the best probe layer.Figure 2 colors the top-256 probe-ranked heads by their causal Δ(ℓ,j).
  • Steering implications: Probe-ranked steering underperforms causal-ranked steering because it can perturb vulnerability-promoting components instead of suppressing them.Using the same mean-difference direction at probe-ranked versus causal-ranked heads produces different behavioral effects.

6 Steering for Secure Code Generation

The section compares causal-head, probe-ranked, layer-level, prompting, and fine-tuning interventions, showing that DUOSTEER best manages vulnerability reduction and correctness preservation across most CWEs.

  • DUOSTEER: DUOSTEER matches or exceeds CAUSALMD on the joint safety-correctness score across all five CWEs.It restores correctness above baseline on CWE-094 while reducing vulnerability, and gains on CWE-295 come entirely from correctness.
  • DUOSTEER: DUOSTEER outperforms a single joint direction because separate safety and correctness directions can target their respective causal head sets with independent strengths.The single direction collapses correctness at moderate steering strength on anti-aligned CWEs such as CWE-079.
  • Steering variants: Causal-head steering reduces vulnerability more effectively than probe-ranked or layer-level steering, but single safety steering often sacrifices functional correctness.CAUSALMD outperforms PROBEMD and LAYERMD on vulnerability reduction, while safety steering can cause correctness collapse.
  • Alternative defenses: DUOSTEER improves the joint score over hint prompting on every CWE, while prompting lowers correctness on every CWE and can increase vulnerability.Prompting raises vulnerability on CWE-295 and leaves the joint score below baseline on three CWEs.
  • Alternative defenses: DUOSTEER also avoids the asymmetric brittleness of SFT, which raises vulnerability on CWE-079 and CWE-295 and degrades correctness.SFT improves the joint trade-off only on CWE-502, where DUOSTEER still performs better.
  • Robustness and generalization: DUOSTEER’s gains remain positive across all five CWEs on Qwen-2.5-Coder-7B-Instruct and under sampling decoding, with paired-bootstrap support.On Qwen, CWE-022 improves through both vulnerability reduction and correctness gains, unlike the Llama boundary case.

7 Post-Hoc Accounts of Per-CWE Steerability

The section explains per-CWE variation in DUOSTEER through causal traction, vector geometry, and contrastive-pair structure, while identifying boundary cases where steering is weak or costly.

  • Motivation: DUOSTEER’s benefit over single safety steering varies sharply across CWEs, ranging from a +50.0 joint-score gain on CWE-502 to a tie on CWE-295.The analysis therefore examines both steering-vector geometry and CODESEC-PAIRS data patterns.
  • Vector geometry: Safety and correctness vectors are near-orthogonal across the representation space, with causal-head cosine means ranging from −0.16 to +0.08.CWE-022 and CWE-079 are anti-aligned, CWE-094 and CWE-295 co-aligned, and CWE-502 is near-zero.
  • Vector geometry: When vectors are anti-aligned or orthogonal, correctness steering can add safety-oriented components and compound vulnerability reduction; co-alignment adds little.Four of five CWEs follow this heuristic, with CWE-502 and CWE-079 showing the largest safety gains over CAUSALMD.
  • Pair structure: More concentrated contrastive-pair fix patterns correlate with larger steering effects, including a Refactor fraction correlation of ρ = −0.90 with steered vulnerability.Refactor pairs can capture abstract restructuring signals that transfer across surface forms, whereas heterogeneous API substitutions transfer less sharply.
  • Boundary cases: CWE-022 is a boundary case because the safety direction has low causal traction, while CWE-079 shows a modest execution-based pass-rate decline despite improved joint score.No CWE-022 head-level configuration beats the unsteered baseline, and DUOSTEER reduces pass rate on CWE-079 by 7.9 Llama points and 3.9 Qwen points.

8 Conclusion

The conclusion identifies a gap between probe-ranked and causally driven heads and presents DUOSTEER as a training-free, dual-direction intervention. It reports stronger safety-correctness outcomes than competing defenses on both evaluated models.

  • Conclusion: Probe-ranked and causally-ranked vulnerability heads are largely disjoint, motivating interventions at causally identified heads.The conclusion treats causal localization as the basis for selecting steering sites.
  • Conclusion: DUOSTEER injects safety and correctness directions at their respective causal head sets without training.It outperforms prompting, SFT, and single-vector steering on both Llama-3.1-8B-Instruct and Qwen-2.5-Coder-7B-Instruct.

Limitations

The study’s evidence is bounded by model scale and access, language and CWE coverage, contextual setting, head-localization scope, and correctness evaluation limitations.

  • Model scale: The experiments cover only two open-weight 7–8B instruct models, and whether the causal-head structure generalizes to larger or closed models remains open.The method requires white-box access to per-head activations during decoding, which closed APIs do not expose.
  • Language and CWE coverage: The analysis uses Python and five CWE classes, while extending it to other languages or vulnerabilities requires new contrastive pairs and causal-head identification.The evaluation also uses isolated function stubs rather than the larger repository and multi-turn contexts common in production.
  • Mechanistic scope: Single-head knockout ranks causally relevant heads but does not explain a full mechanistic circuit or measure interactions among heads.Counterfactual activation patching, path patching, and automated circuit discovery are identified as complementary directions.
  • Correctness evaluation: Correctness evaluation relies solely on the GPT-4.1 judge for CWE-022 and CWE-295, while available tests remain narrow even for other CWEs.The judge agrees with a human annotator at Cohen’s κ ≈0.87, and execution-based tests agree in direction on the three CWEs that provide them.

Ethics Statement

The paper frames DuoSteer as a defensive intervention and documents public, privacy-preserving research materials alongside its annotation and evaluation procedures.

  • Ethics Statement: Activation steering is sign-reversible, so the authors release DuoSteer strictly for defensive deployment and recommend against reversing its safety direction.Negating the safety direction could in principle steer generation toward more vulnerable code.
  • Ethics Statement: All models, benchmarks, and analysis tools are open-source or publicly available, and CodeSec-Pairs contains only model-generated code without personally identifying information.No human subjects were involved beyond author annotation.
  • Annotation protocol: Each contrastive pair is independently annotated by GPT-4.1 for structural distance and fix mechanism using the safe and vulnerable code side-by-side.Structural distance distinguishes Minimal, Refactor, and Divergent changes; fix mechanism distinguishes Substitution, Guard-Addition, Deletion, and Unclear.
  • Annotation validation: Automatic annotations show substantial agreement for fix mechanism and almost-perfect agreement for structural distance on 200 hand-labeled pairs.Cohen’s κ is 0.73 for fix mechanism and 0.81 for structural distance.
  • Dataset composition: The dominant fix mechanism varies sharply by CWE, motivating the paper’s observations about per-CWE steering behavior.The full annotation distribution is reported in Table 5, with case studies in Appendix J.
  • Prompting: The study uses distinct benign and vulnerability-eliciting prompt templates for constructing safe and vulnerable responses.The vulnerability-eliciting template injects a CWE class and description, while the prompt-level hint baseline uses a separate hint-augmented template.
  • Interpretability protocol: Linear probes are trained per CWE at both layer and attention-head representations to predict safe versus vulnerable code.Per-layer probes use mean-pooled residual streams, while per-head probes use value-projection outputs before the head output projection.
  • Evaluation protocol: Security failures are excluded from the correctness label so vulnerability and functional correctness measure disjoint failure modes.Correctness is assigned by GPT-4.1 using the CodeJudge template and a binary decision procedure.

C Joint-Vector and Direction Ablations

The ablations show that separately steering safety and correctness is more effective and stable than combining both objectives into one vector, while head and direction selection matter.

  • Joint vector vs. DuoSteer: The joint vector underperforms DuoSteer on C(1−V ) for every CWE and raises vulnerability above baseline on CWE-094 and CWE-295.Table 6 compares the best joint-vector configuration with the unsteered baseline and best DuoSteer configuration.
  • Steering stability: At α > 3, the joint vector collapses on all five CWEs, producing degenerate or syntactically broken outputs that cannot be reliably scored.Observed failures include empty strings, repeated tokens, and non-code text.
  • Random controls: Random directions and random head selections do not reduce vulnerability, supporting the informativeness of learned directions and causally selected heads.Probe-ranked head steering is the informative mismatch: its vulnerability reduction collapses to ∆V ∈[−5.9, +1.4] on CWE-022/079/094.
  • Direction geometry: Safety and correctness directions can be anti-aligned, co-aligned, or transition toward zero depending on the CWE and head budget.CWE-022 and CWE-079 remain anti-aligned, CWE-094 and CWE-295 remain co-aligned, and CWE-502 transitions from anti-aligned to near-zero.
  • Supervised fine-tuning baseline: The SFT baseline is trained on 5,000 balanced safe-and-correct examples and evaluated with the same prompts, detector, and correctness judge as steering conditions.The corpus contains 1,000 examples per CWE, with 4,750 training and 250 validation examples.

F Hint-vs-Steering Geometric Analysis

The analysis compares hint-induced activation shifts with DUOSTEER’s safety direction and finds them essentially perpendicular, indicating complementary interventions. Combining hints with DUOSTEER helps most where steering alone has limited leverage, while unretuned stacking can reduce correctness elsewhere.

  • Geometric comparison: Successful hints shift safety-causal head activations substantially, but their direction is essentially perpendicular to the learned safety vector.Across all five CWEs, the shift norm is approximately 0.13 versus a prompt-noise floor of approximately 0.02, while cosine similarity is below 0.02.
  • Geometric comparison: Hint prompting and DUOSTEER are therefore non-redundant and complementary defenses in representation space.The paper identifies the combination as most useful on CWEs where steering alone has little leverage.
  • Combined intervention: On three of five CWEs, adding hints to DUOSTEER improves the joint score relative to hints alone, with the largest gain on CWE-502.For CWE-502, vulnerability falls from 11.8% to 2.0% and the joint score rises from 13.8 to 52.6.
  • Combined intervention: Against DUOSTEER alone, hint stacking improves only CWE-022 and CWE-094, where steering has the weakest causal traction.The joint score increases from 43.0 to 47.4 on CWE-022 and from 17.2 to 18.9 on CWE-094.
  • Robustness and selection: DUOSTEER improves the joint score over baseline under sampling decoding on all five CWEs, with significant gains on three.The sampling regime uses temperature 0.7 and top-p 0.95; CWE-094 remains at the reported p = 0.052 boundary.
  • Robustness and selection: Held-out validation selects configurations that improve the joint score over baseline on four of five CWEs, while CWE-022 remains a boundary case.Validation recovers the same head budget on every CWE except CWE-094 and exactly matches the test-selected configuration on CWE-502.
Loading 2608.30025v1…