DeepGuard: 다층 의미 집계를 통한 안전한 코드 생성

DeepGuard: Secure Code Generation via Multi-Layer Semantic Aggregation

코드 생성에 사용되는 대규모 언어 모델(LLM)은 학습 데이터에서 불안전한 패턴을 모방할 수 있습니다. 이를 완화하기 위한 일반적인 전략은 최종 트랜스포머 레이어에서 파생된 감독 신호를 사용하여 모델을 미세 조정하는 것입니다. 그러나 이러한 설계는 최종 레이어 병목 현상을 겪을 수 있습니다. 즉, 취약점 식별에 중요한 신호가 여러 레이어에 분산되어 있으며, 다음 토큰 예측을 위해 최적화된 출력 표현 근처에서는 감지하기 어려워질 수 있습니다. 이 문제를 진단하기 위해, 우리는 레이어별 선형 프로빙을 수행했습니다. 분석 결과, 취약점과 관련된 신호는 중간에서 상위 레이어에서 가장 잘 감지되지만, 최종 레이어에 가까워질수록 약화되는 것을 확인했습니다. 이러한 관찰을 바탕으로, 우리는 DeepGuard라는 프레임워크를 제안합니다. DeepGuard는 어텐션 기반 모듈을 통해 여러 상위 레이어의 표현을 집계하여 분산된 보안 관련 신호를 활용합니다. 집계된 신호는 보안 향상과 기능적 정확성을 균형 있게 유지하는 다중 목표 학습 목표 내에서 전용 보안 분석기를 구동하며, 또한 경량화된 추론 시간 제어 전략을 지원합니다. 다섯 가지 코드 LLM에 대한 광범위한 실험 결과, DeepGuard는 SVEN과 같은 강력한 기본 모델보다 평균 11.9% 더 높은 안전하고 정확한 생성률을 보였습니다. 또한, DeepGuard는 기능적 정확성을 유지하면서도 새로운 유형의 취약점에 대한 일반화 능력을 보여줍니다. 저희의 코드는 다음 링크에서 공개적으로 이용 가능합니다: https://github.com/unknownhl/DeepGuard.

Large Language Models (LLMs) for code generation can replicate insecure patterns from their training data. To mitigate this, a common strategy for security hardening is to fine-tune models using supervision derived from the final transformer layer. However, this design may suffer from a final-layer bottleneck: vulnerability-discriminative cues can be distributed across layers and become less detectable near the output representations optimized for next-token prediction. To diagnose this issue, we perform layer-wise linear probing. We observe that vulnerability-related signals are most detectable in a band of intermediate-to-upper layers yet attenuate toward the final layers. Motivated by this observation, we introduce DeepGuard, a framework that leverages distributed security-relevant cues by aggregating representations from multiple upper layers via an attention-based module. The aggregated signal powers a dedicated security analyzer within a multi-objective training objective that balances security enhancement and functional correctness, and further supports a lightweight inference-time steering strategy. Extensive experiments across five code LLMs demonstrate that DeepGuard improves the secure-and-correct generation rate by an average of 11.9% over strong baselines such as SVEN. It also preserves functional correctness while exhibiting generalization to held-out vulnerability types. Our code is public at https://github.com/unknownhl/DeepGuard.