Interpreting and Steering for Safe and Correct Code Generation

Large language models (LLMs) frequently generate code with security vulnerabilities. While existing mitigations like prompting or fine-tuning operate at the input/output boundary, they often fail to address the internal mechanisms that drive these vulnerabilities. This paper investigates where safe-versus-vulnerable code generation is represented within LLMs and whether this internal understanding can be used to steer models toward safer outputs without retraining.

The authors introduce CodeSec-Pairs, a dataset of 9,342 contrastive Python code pairs (safe vs. vulnerable) sampled from Llama-3.1-8B-Instruct. They employ two interpretability techniques: linear probing to identify where safety information is decodable, and causal head knockout to identify which attention heads actually drive the model's preference for safe or vulnerable code. They find that these two sets of components are largely disjoint. Based on this, they propose DuoSteer, a double-steering approach that injects a safety direction and a correctness direction at their respective causally identified attention heads.

DuoSteer significantly outperforms existing baselines, including prompting and supervised fine-tuning. Across five vulnerability types, it achieves an average 26.9% reduction in vulnerability rates while simultaneously improving functional correctness by 7.5%. The authors demonstrate that the safety and correctness directions are near-orthogonal in the model's representation space, which allows them to be combined additively without the performance collapse seen in single-vector steering methods. These results were successfully replicated on Qwen-2.5-Coder-7B-Instruct.

This work provides a principled, training-free method for improving the security of LLM-generated code. By moving beyond black-box interventions and targeting the specific attention heads responsible for safety and correctness, the authors offer a more precise way to mitigate vulnerabilities. This approach highlights the importance of distinguishing between where a model stores information and where it actually uses that information to make decisions.