Stable-GFlowNet: 대조적 궤적 균형을 통한 다양하고 강력한 LLM 레드티밍

기존 LLM 레드티밍 도구들은 특정 공격 방식에만 매몰되거나 학습이 불안정하여 다양한 보안 취약점을 찾아내는 데 한계가 있었다. 이 논문은 수학적으로 불안정한 계산 과정을 제거한 Stable-GFN을 통해 기존 대비 7배 더 다양한 공격 프롬프트를 생성하며 LLM의 안전성을 획기적으로 강화한다.

기존의 GFlowNet은 전체 보상의 합인 분배 함수(Z)를 학습 과정에서 계속 추정해야 했다. 이는 마치 전 세계 모든 맛집의 점수 합계를 알아야 특정 식당의 순위를 매길 수 있는 것과 같아서, 데이터가 방대한 LLM 환경에서는 계산 오차가 발생하기 쉽고 학습이 금방 망가지는 원인이 되었다.

Stable-GFN은 이 문제를 해결하기 위해 '상대적 비교' 개념을 도입했다. 전체 합계를 구하는 대신, 두 개의 공격 프롬프트를 직접 비교하여 보상이 더 높은 쪽이 더 자주 생성되도록 확률 비율을 조정한다. 이는 절대적인 점수 대신 'A가 B보다 2배 더 맛있다'는 상대적 정보만으로도 전체적인 선호도 분포를 정확히 학습할 수 있다는 원리를 이용한 것이다.

결과적으로 모델은 분배 함수 추정이라는 불안정한 단계를 건너뛰면서도, 보상이 높은 다양한 공격 경로를 골고루 탐색할 수 있게 된다. 이를 통해 특정 공격 패턴에만 갇히지 않고 LLM이 가진 수많은 잠재적 취약점을 넓은 범위에서 찾아낼 수 있다.

Stable-GFN은 세 가지 핵심 모듈로 구성된다. 첫째, Contrastive Trajectory Balance (CTB)는 두 독립 샘플 y1, y2의 정책 확률 비율 [pi(y1)/pi(y2)]과 보상 비율 [R(y1)/R(y2)]의 로그 차이를 제곱하여 손실을 계산한다. 이 과정에서 공통 분모인 Z가 상쇄되어 제거되므로 수치적 안정성이 극대화된다.

둘째, Noisy Gradient Pruning (NGP)은 두 샘플의 로그 보상 차이가 특정 임계값 sigma보다 작을 경우 해당 쌍의 손실 계산을 생략한다. [|log R(y1) - log R(y2)| > sigma] 조건이 충족될 때만 학습을 진행함으로써, 독성 분류기의 미세한 노이즈가 모델 가중치에 악영향을 주는 것을 방지한다.

셋째, Min-K Fluency Stabilizer (MKS)는 참조 모델을 활용해 생성된 문장에서 가장 확률이 낮은 k개 토큰의 평균 로그 확률을 계산한다. 이 값이 임계값 미만이면 해당 샘플에 강한 페널티 보상을 부여하여, 모델이 의미 없는 문자열(gibberish)을 생성해 보상을 가로채는 리워드 해킹을 차단한다.

실험 결과, Stable-GFN은 기존 GFN 베이스라인 대비 약 7배 더 많은 고유한 공격 프롬프트(Unique Attacks)를 생성하는 데 성공했다. 구체적으로 고유 공격 수는 17개에서 134개로 증가했으며, 공격 성공률(ASR)은 92.55%의 높은 수준을 유지했다.

Ablation Study에서는 CTB와 NGP를 결합했을 때 가장 높은 UA(121개)와 ASR(92.2%)을 기록함을 확인했다. 특히 분배 함수 Z를 직접 학습하는 기존 방식(GFN-TB)이 노이즈가 많은 환경에서 학습에 실패하거나 특정 모드로 급격히 수렴하는 것과 달리, 제안 기법은 일관된 성능 향상을 보였다.

전이 공격(Transfer Attack) 테스트에서도 Stable-GFN으로 생성된 공격 세트는 학습에 사용되지 않은 다른 LLM(Gemma, Llama 3.2 등)들에 대해서도 기존 기법들보다 압도적으로 높은 공격 성공률을 기록하며 범용적인 취약점 발견 능력을 입증했다.

Stable-GFN의 핵심 이론적 기여는 CTB 목적 함수가 기존 Trajectory Balance(TB)와 동일한 최적 정책을 공유함을 수학적으로 증명한 것이다. CTB의 그래디언트 분석 결과, 이는 Reinforce Leave-One-Out (RLOO)과 유사한 형태의 내재적 분산 감소 메커니즘을 포함하고 있음이 밝혀졌다.

NGP의 경우, 샘플 간의 연결성을 나타내는 saliency graph가 연결 그래프(connected graph) 형태를 유지한다면 최적 정책으로의 수렴이 보장됨을 이론적으로 제시했다. 이는 고차원 공간에서도 리플레이 버퍼를 활용한 앵커링을 통해 실무적으로 달성 가능하다.

MKS는 기존의 KL-divergence 정규화가 참조 모델의 분포에 과도하게 고착되어 탐색 성능을 저하시키는 문제를 해결한다. 전체 분포가 아닌 하위 k개 토큰의 유창성만을 제약함으로써, 모델이 언어적 형식을 파괴하지 않으면서도 보상 지형을 자유롭게 탐색할 수 있는 유연성을 제공한다.