기술 요약: 이산 확산을 활용한 제약 조건부 코드 생성

문제 정의

코드 생성 모델, 특히 이산 확산 모델은 부분적으로 손상된 토큰 시퀀스를 반복적으로 정제하여 프로그램을 구축하는 유망한 패러다임으로 부상했습니다. 그러나 중요한 한계가 존재합니다. 이러한 생성 모델의 학습 목적은 목표 분포를 맞추는 것이지만, 기능적 정확성, 보안 속성, 또는 엄격한 구문 유효성과 같은 프로그램 수준의 제약 조건을 본질적으로 포착하지는 못합니다. 생성된 프로그램이 기본 모델 하에서 매우 그럴듯할지라도, 여전히 잘못된 동작을 보이거나 악용 가능한 취약점을 포함할 수 있습니다.

기존의 안내 메커니즘 (프롬프팅, 재순위화, 거부 샘플링, 사후 수리 등) 은 일반적으로 후보 프로그램이 완전히 생성된 후에만 개입합니다. 이 접근 방식은 검색 공간이 매우 결합적이며, 의미론적 또는 보안 제약 조건을 만족하는 프로그램의 부분집합이 작은 실현 가능 영역을 구성할 수 있기 때문에 코드 생성에 있어 비효율적이고 신뢰할 수 없습니다. 사후 접근 방식은 종종 전역적 구조적 또는 의미론적 위반을 수정하기 위해 큰 접미사 전체 또는 전체 프로그램을 재생성해야 합니다.

이산 확산 모델은 이를 해결할 고유한 기회를 제공합니다. 역방향 탈노이즈 과정은 모든 중간 단계에서 전역 프로그램 상태를 노출합니다. 이 구조는 최종 코드가 확정되기 전에 제약 조건을 생성 과정에 직접 통합할 수 있는 자연스러운 개입 지점을 제공합니다.

방법론: 코드에 대한 제약 확산 (CDC)

본 논문은 프로그램 수준의 제약 조건을 이산 확산 모델의 역방향 탈노이즈 과정에 직접 통합하는 훈련 불필요 신경 심볼릭 추론 프레임워크인 **코드에 대한 제약 확산 (Constrained Diffusion for Code, CDC)**을 소개합니다. CDC 는 확산 샘플링 궤적을 편집 가능한 프로그램 표현으로 취급합니다. 각 탈노이즈 단계에서 프레임워크는 네 가지 주요 작업을 수행합니다.

1. 클린 상태 제안: 기본 탈노이저는 프로그램 시퀀스 전체에 대한 클린 토큰 분포 ($\hat{x}^{(t)}_0$) 를 예측합니다.
2. 제약 조건 인식 로컬라이제이션: 제약 조건 평가기 ($E$) 는 단위 테스트, 보안 분석기 등의 목표 명세에 대해 후보 프로그램을 평가합니다. 이는 위반 점수와 구조화된 피드백을 반환합니다. 로컬라이제이션 맵 ($M$) 은 이 피드백을 위반에 책임 있는 특정 토큰 영역 ($S_t$) 으로 변환합니다.
3. 제약 조건 인식 수정: 제약 조건 인식 연산자 ($P_C$) 는 로컬라이즈된 영역 내에서 제안 분포를 수정합니다. 목표는 기본 탈노이저의 분포 (Kullback–Leibler 발산을 통해 고정됨) 에 근접하면서 제약 위반을 줄이는 것입니다.
4. 제약 조건부 역방향 업데이트: 수정된 제안은 다음 시간 단계로 역방향 체인을 진행하는 데 사용됩니다.

CDC 는 서로 다른 제약 조건 유형에 맞춰 조정된 두 가지 특정 연산자로 이 프레임워크를 구체화합니다.

• GradGuide (기능적 정확성): 이 연산자는 실행 기반 레이블 (예: 테스트 통과 결과) 로 훈련된 학습된 대리 모델 ($g_\phi$) 을 사용하여 소프트 제약 조건을 처리합니다.
  • 토큰 분포를 소프트 임베딩으로 변환하여 만족도 점수를 예측합니다.
  • 대리 위반에 대한 토큰 임베딩의 기울기를 계산하여 오류를 로컬라이즈하고, 고-주요성 영역을 식별합니다.
  • KL 고정 프로젝션 (확률 분포를 실현 가능한 토큰 방향으로 이동) 과 제약 조건 트리거 로컬 리마스크 (중간 프로그램이 여전히 제약 조건을 위반하는 경우 확정된 토큰을 탈노이즈를 위해 다시 개방) 를 통해 수정을 적용합니다.
• MDFI (보안): 이 연산자는 정적 프로그램 분석을 사용하여 하드 제약 조건을 처리합니다.
  • 정적 분석기는 중간 상태에서 부분 프로그램 그래프 (AST 와 데이터 플로우 결합) 를 구성합니다.
  • 취약점 증인 (예: 안전하지 않은 데이터 플로우 경로 또는 구조적 패턴) 을 식별하고 이를 토큰 스팬으로 로컬라이즈합니다.
  • 이산적 개입을 적용합니다: 대체 (안전하지 않은 구성 요소 리마스크) 및 삽입 (누락된 가드 또는 확인을 위한 새 마스킹 토큰 주입).
  • 나머지 탈노이즈 단계를 안내하기 위해 조건부 컨텍스트에 직접 수정 힌트를 주입합니다.

주요 기여

본 논문은 세 가지 주요 기여를 합니다.

1. 제약 조건 인식 로컬라이제이션: 탈노이즈 중 수정을 위해 대리 모델 또는 정적 분석기로부터의 프로그램 수준 피드백을 특정 토큰 영역으로 매핑하는 일반 메커니즘.
2. 제약 조건 인식 수정: 기본 탈노이저에 근접하면서 로컬라이즈된 영역을 수정하는 훈련 불필요 연산자. 이들은 미분 가능한 안내 (기울기 및 KL 고정 업데이트를 통해) 와 비미분 가능한 심볼릭 피드백 (표적 리마스크 및 삽입을 통해) 을 지원합니다.
3. 실증적 평가: 탈노이즈 중 제약 조건 통합이 기능적 정확성, 구문, 보안 벤치마크 전반에 걸쳐 코드 생성을 개선하며, 더 적은 편집 토큰으로 이산 확산 및 자기회귀 베이스라인 모두를 능가한다는 증거.

실험 결과

저자들은 Dream-Coder 7B, DiffuCoder 7B, Dream-7B 를 백본으로 사용하여 HumanEval-X C++, MBPP-C++, CWEval, LLMSecEval+ 에서 CDC 를 평가했습니다.

• 기능적 정확성: CDC 는 pass@1 및 pass@10 비율을 크게 향상시켰습니다. HumanEval-X C++ 에서 CDC 는 Dream-Coder 7B 의 pass@1 을 34.1% 에서 65.2% 로, 컴파일 성공률을 67.1% 에서 79.2% 로 향상시켰습니다. MBPP-C++ 에서 pass@1 은 27.7% 에서 59.2% 로 향상되었습니다. 이러한 결과는 여러 지표에서 최첨단 자기회귀 모델 (예: DeepSeek-Coder 6.7B) 을 능가했습니다.
• 보안: CDC 는 기능 - 보안 공동 성공률을 크게 향상시켰습니다. CWEval 에서 공동 pass@1 (fs@1) 은 12.04% 에서 34.26% 로 증가했습니다. LLMSecEval+ 에서 보안 pass@1 (sec@1) 은 일반 (vanilla) 상태의 54.67% 에서 80.67% 로 상승하여 자기회귀 재프롬프팅을 능가했습니다.
• 효율성 및 로컬라이제이션: CDC 는 자기회귀 재프롬프팅에 비해 우수한 효율성을 보여주었습니다. 자기회귀 방법은 종종 전체 프로그램을 재생성하는 반면 (예: CWEval 에서 중앙값 113 개 토큰 편집), CDC 는 매우 로컬라이즈된 편집을 수행했습니다 (CWEval 에서 중앙값 12 개 토큰). 이는 CDC 가 광범위한 재생성에 의존하는 대신 제약 관련 영역에 수정을 집중한다는 것을 확인시켜 줍니다.

중요성 및 주장

본 논문은 이산 확산 모델이 제약 조건부 코드 생성에 대한 자연스러운 기질 역할을 한다고 주장합니다. 샘플링 중에 노출되는 전역적이고 편집 가능한 상태를 활용함으로써, CDC 는 사후 필터링이나 전체 프로그램 재생성에 의존하는 대신 생성 중 에 집중된 제약 수정을 가능하게 합니다.

저자들은 이 접근 방식이 기본 모델을 재훈련하지 않고도 대리 모델 기반의 소프트 제약과 심볼릭 분석기 기반의 하드 제약 모두를 통합할 수 있음을 강조합니다. 결과는 탈노이즈 시간 제약 통합이 확산 기반 코드 모델을 강력한 자기회귀 생성기와 경쟁 가능하게 만들 수 있음을 시사하며, 안전하고 기능적으로 정확한 코드를 생성하는 더 효율적이고 표적화된 경로를 제공합니다. 이 연구는 최종 출력이 확정되기 전에 모델을 실현 가능한 솔루션으로 유도함으로써 소프트웨어 공학을 위한 더 신뢰할 수 있고 안전한 생성형 AI 로 나아가는 한 걸음으로 자리매김합니다.