Interpreting and Steering for Safe and Correct Code Generation
Hao Yan, Ziyu Yao
EMNLP 2026 Main Conference
cs.AI
2026-08-31
DuoSteer steers safety and correctness heads, cutting Llama-3.1-8B vulnerability 26.9% on average and CWE-502 from 64.7% to 0%.
LLMs emit code with CWE bugs. Prior fixes sit at the I/O boundary: extra safety hints in the prompt, or SFT on secure code. Which layers and attention heads push a continuation toward safe versus vulnerable code was mostly unknown, so there was nothing to edit at inference.
Safety and correctness also tangle. Push only on safety and the vulnerability rate falls while functional correctness collapses. The need is a training-free, inference-time edit that keeps the code running.
CodeSec-Pairs is built first. Tasks come from Emergent Misalignment, SecurityEval, CodeLMSec, and CyberSecEval. Llama-3.1-8B-Instruct is sampled 10 times on a benign prompt and 10 times on a prompt that asks to skip validation. CodeQL labels the outputs. The result is 9,342 Python safe/vulnerable pairs across CWE-022 path traversal, 079 XSS, 094 code injection, 295 cert validation, and 502 unsafe deserialization. Intra-prompt pairs train probes and steering vectors; cross-prompt pairs are for causal localization only. The correctness vector uses a second pair set: CodeQL-safe outputs matched as safe-and-correct versus safe-but-incorrect, about 400 pairs per CWE, so correctness is not confounded with safety.
Two localizers. A linear probe asks whether the residual stream or a head's output linearly separates safe from vulnerable. Causal knockout zeros one head at all response positions and watches the log-likelihood margin of the safe continuation versus the vulnerable one. Layer probes reach 71% to 87% accuracy, but probe rank versus causal effect has Spearman |ρ| ≤ 0.072, and the top causal head per CWE sits at probe rank 65 to 250. Probes find where the distinction is encoded; knockout finds what drives generation.
Steering uses a mean-difference vector, divided by that head's activation standard deviation, scaled by strength α. DuoSteer adds the safety direction at safety-causal heads and the correctness direction at correctness-causal heads; heads in both sets get the sum. Overlap at the best configs is small, 4/32 to 25/64. A single joint vector of (safe ∧ correct) versus (vuln ∧ incorrect) collapses correctness at moderate α on three CWEs.
Evaluation is 274 SecCodePLT questions. Vulnerability V is CodeQL; correctness C is GPT-4.1 (Cohen's κ ≈ 0.87 versus a human). Table entries are the best joint score C(1−V) per method. Across five CWEs the paper reports −26.9% vulnerability and +7.5% correctness on average.
| CWE | Baseline V / C | DuoSteer V / C |
| 022 path traversal | 32.9 / 72.9 | 30.0 / 61.4, joint score below baseline |
| 079 XSS | 54.9 / 90.2 | 3.9 / 82.3 |
| 094 injection | 51.0 / 25.5 | 37.2 / 27.5 |
| 295 certs | 5.9 / 41.2 | 7.8 / 60.8, V barely moves, C +19.6 |
| 502 deserialization | 64.7 / 58.8 | 0.0 / 94.1 |
Causal-head mean-diff suppresses V more than probe-ranked heads, often at C's expense; safety-only steering loses the joint score on 022 and 094. Hint prompting cuts V harder on 022 and 094 and drops C by 7.9 to 43.1 points. SFT raises V on 079 and 295. Qwen-2.5-Coder-7B with another 2,500 pairs lifts the joint score on all five CWEs; 079 goes from 27.5% vulnerable to 2.0%.
Probe accuracy is not causal control, the usual failure when interpretability tries to intervene. DuoSteer treats safety and correctness as two vectors on two causal head sets and adds them at decode, with no finetune. When the fix pattern is concentrated, as in CWE-502's pickle-style API swap, the mean-difference vector is sharp. When fixes scatter, as in CWE-022, head-level edits still fail. Deployment needs white-box access and a per-model causal ranking. It is not an API switch.
Only 7B–8B open instruct models. 70B, 405B, and closed APIs have no per-head activations. Python only, five CWEs, no repo-level context. Single-head knockout is not a circuit; head interactions are unmeasured. Correctness is mainly GPT-4.1; execution tests exist only for 079/094/502. Hyperparameters were swept on the eval set; a 100-pair held-out protocol still beats baseline on four of five CWEs. Steering is sign-reversible; the paper releases it for defense. CWE-022 on Llama is an honest miss.