Teaching an Old LLM Secure Coding: Localized Preference Optimization on Distilled Preferences

TL;DR

This paper introduces a systematic approach to improve secure code generation in large language models by distilling a preference dataset and developing a localized preference optimization algorithm that focuses on security-related code regions.

Contribution

The paper presents a novel method for distilling security-focused preferences from LLMs and a localized preference optimization algorithm that enhances secure code generation.

Findings

Significant reduction in insecure code generation.

Improved overall code quality.

Effective handling of localized security issues.

Abstract

LLM generated code often contains security issues. We address two key challenges in improving secure code generation. First, obtaining high quality training data covering a broad set of security issues is critical. To address this, we introduce a method for distilling a preference dataset of insecure and secure code pairs from frontier LLMs, along with a security reasoning that explains the issues and the fix. The key idea here is to make use of security knowledge sources to devise a systematic prompting strategy that ensures broad coverage. Second, aligning models to secure code requires focusing on localized regions of code. Direct preference optimization methods, like SimPO, are not designed to handle these localized differences and turn out to be ineffective. We address this with a new localized preference optimization algorithm that masks the security related tokens in both the winning (secure) and losing (insecure) responses. To prevent loss in code quality, we also add a regularizer. Evaluations show that both training on our dataset, DiSCo, and the new preference optimization algorithm, LPO, yield substantial reductions in code insecurity while also improving overall code quality. Code and dataset are available at https://github.com/StonyBrookNLP/disco-lpo.