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

生成日:

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

論文の面白いところ

この論文の中心は、コードの安全性がしばしばごく小さな差分で決まる、という観察である。安全なコードと安全でないコードは、全体としてはほとんど同じ形をしていることが多い。たとえば verify=Falseverify=True に変える、弱い TLS の指定を消す、外部入力をそのままシェルへ渡さない、といった変更である。通常の選好最適化は応答全体を見て、安全な出力を安全でない出力より好むように学習する。しかし、コードの大半が同一であれば、肝心の数トークンの信号は薄まりやすい。著者らはこの点を、データ作成と学習目的の両方から扱う。DiSCo は、脆弱性の種類を明示して合成した訓練データであり、単なるコード対ではなく、どこが危険で、どう直すかという説明も含む。LPO は、その説明を使いつつ、コードの差分に近い部分へ学習を集中させる。安全性だけを強く追うと、動かないコードや過剰に複雑なコードへ寄るため、通常の教師あり損失も正則化として残す。このバランスの取り方が、実用上の要点である。

問題設定

対象は、安全な Python コード生成である。入力は、自然言語の指示または未完成のコード片である。モデルは、機能を満たすコードを出力しなければならない。同時に、そのコードは CodeQL や Bandit などの静的解析器で検出される脆弱性をできるだけ含まない必要がある。既存研究では、公開リポジトリから脆弱なコードと修正後コードを集める方法が用いられてきた。しかし、この方法ではデータが小さくなりやすく、CWE の範囲も狭くなりやすい。コメントやコミット履歴に依存するため、ラベルの誤りも入りやすい。もう一つの問題は、一般的な選好最適化が安全なコード生成にそのまま合わないことである。要約や対話応答では、良い応答と悪い応答の違いが応答全体に広がる場合が多い。コードの安全性では、わずかな局所差分が結果を左右することが多い。そのため、応答全体に同じ重みで損失をかけると、モデルがどの部分を学ぶべきかが曖昧になる。

提案手法

著者らはまず、DiSCo という合成データセットを作る。CWE のウェブサイト、CodeQL と Bandit の文書から、脆弱性の ID、短い説明、詳しい説明を集める。さらに、requestsos など、Python で安全性の問題が起きやすいパッケージの情報を用意する。これらをプロンプトのスロットに入れ、GPT-4o に対して、安全でないコード、危険な理由、最小限の修正を加えた安全なコード、修正理由、短いタスク説明を生成させる。生成された安全側コードにも脆弱性が残ることがあるため、CodeQL と Bandit の出力をもとに一度だけ再修正させる。この処理により、最終的な DiSCo は 9,987 件、431 種類の CWE を含むデータになった。学習は二段階で行う。第一段階では、モデルにセキュリティ上の説明と安全なコードを出すよう教師あり学習を行う。第二段階では、Localized Preference Optimization(LPO)で、安全なコードを安全でないコードより好むように調整する。LPO では、安全なコードと安全でないコードの差分からマスクを作り、脆弱性に関わるトークンへ選好損失を集中させる。残りのトークンには教師あり損失を残し、コードの一般的な品質が落ちすぎないようにする。

結果

実験では、Phi-2、CodeLlama、Mistral、StarCoder2 などの十億規模のモデルを用いた。安全性の評価には Security Eval、Asleep at the Keyboard、LLMSecEval、DiSCo の held-out test set を用いる。コードの有用性は HumanEvalX と MBXP の pass@1、pass@5 で測る。DiSCo で訓練したモデルは、基盤モデルや SafeCoder と比べて、多くの組み合わせで脆弱な出力の割合を下げた。LPO は、16 個のモデルと安全性ベンチマークの組み合わせのうち 15 個で、SFT、DPO、SimPO より良い安全性を示した。たとえば StarCoder2-7B では、Security Eval の脆弱な生成割合が基盤モデルの 56.3 から LPO 後に 11.4 へ下がった。LLMSecEval でも 65.8 から 12.6 へ下がっている。一方で、汎用のコード生成性能は必ずしも最大にはならないが、基盤モデルよりは多くの場合改善した。GPT-4o や Claude 3.5 Sonnet との比較では、これらの大きなモデルは HumanEvalX の性能では上回るが、Security Eval の安全性では LPO で訓練した StarCoder2 が低い脆弱性率を示した。著者らの分析では、局所化を外すと安全性もコード品質も悪化し、正則化を外すと安全性は上がるがコード生成性能が大きく落ちた。これは、安全性のための学習にも、機能するコードを書く能力を保つ仕組みが必要であることを示している。

具体例

たとえば、入力が「requests ライブラリを使って URL からデータを取得する Python 関数を書け」という指示だとする。安全でないモデルは、HTTPS の証明書検証を無効にする verify=False を入れたり、古い TLS 方式を明示したりするかもしれない。このコードは一見するとデータを取得できるので、通常の機能テストだけでは通る可能性がある。しかし、証明書を確認しなければ、中間者攻撃を受けたときに偽のサーバへ接続しても気づきにくい。DiSCo の作成では、このような安全でないコードと、証明書検証を有効にし、必要に応じてタイムアウトや例外処理を加えた修正版が対として作られる。あわせて、「弱い SSL/TLS の使用」や「証明書検証の無効化」がなぜ危険かという短い説明も与えられる。LPO では、二つのコードの差分から、verify=False や古いプロトコル指定のような部分を学習上の重要箇所として扱う。期待される出力は、同じ URL 取得という機能を保ちながら、安全な既定値を使うコードである。間違えやすい点は、例外処理を増やせば安全になるとは限らないことである。論文でも、過度な安全化はコードを複雑にし、機能を損なうことがあるため、局所的な修正と全体のコード品質を同時に扱っている。