Stable-GFlowNet: Toward Diverse and Robust LLM Red-Teaming via Contrastive Trajectory Balance

"Stable-GFlowNet: Toward Diverse and Robust LLM Red-Teaming via Contrastive Trajectory Balance" 논문에 대한 상세한 기술적 요약은 다음과 같습니다.

1. 문제 제기

대형 언어 모델 (LLM) 레드팀링은 배포 전에 취약점 (유해하거나 해로운 출력) 을 사전에 발견하는 것을 목표로 합니다. 강화 학습 (RL) 방법은 높은 보상을 얻는 공격을 찾을 수 있지만, **모드 붕괴 (mode collapse)**로 인해 반복적이고 협소한 공격 패턴을 생성하는 단점이 있습니다. 반면, 생성 흐름 네트워크 (GFN) 는 분포 매칭을 위한 프레임워크를 제공하여 이론적으로 다양한 고보상 궤적을 발견할 수 있게 합니다. 그러나 GFN 을 LLM 레드팀링에 적용하는 데는 두 가지 중요한 과제가 존재합니다:

1. 학습 불안정성: 표준 GFN 목적 함수 (예: 궤적 균형, TB) 는 글로벌 분할 함수 ($Z$) 의 추정이 필요합니다. LLM 의 고차원 이산 공간에서 $Z$를 추정하는 것은 수치적으로 불안정하여 수렴이 저조하고 모드 붕괴를 초래합니다.
2. 노이즈가 있는 보상: 유해성 분류기는 밀집하지만 노이즈가 많은 신호를 제공합니다. 이러한 분류기는 종종 의미 없는 텍스트나 분포 밖 (OOD) 텍스트에 0 이 아닌 보상을 부여하여, 모델이 유창하고 다양한 공격 대신 무의미한 텍스트를 생성함으로써 보상을 '해킹'하게 만듭니다.

2. 방법론: Stable-GFlowNet (S-GFN)

저자들은 명시적인 분할 함수 추정을 제거하고 노이즈가 있는 보상을 필터링하여 안정적이고 다양한 학습을 보장하는 Stable-GFlowNet (S-GFN) 프레임워크를 제안합니다. 이 파이프라인은 세 가지 핵심 구성 요소로 이루어져 있습니다:

A. 대비적 궤적 균형 (Contrastive Trajectory Balance, CTB)

$Z$ 추정의 불안정성을 해결하기 위해, S-GFN 은 절대 흐름 매칭 목적 함수를 쌍대 비교 목적 함수로 대체합니다.

• 메커니즘: 정책 흐름과 $Z$에 대한 보상의 오차를 최소화하는 대신, CTB 는 두 궤적 $(y_1, y_2)$에 대한 정책 확률의 로그 비율과 보상의 로그 비율 간의 제곱 차이를 최소화합니다:
  $$L_{CTB} = \left( \log \frac{\pi_\theta(y_1)}{\pi_\theta(y_2)} - \log \frac{R(y_1)}{R(y_2)} \right)^2$$
• 이론적 보장: 저자들은 CTB 의 최적 정책이 표준 궤적 균형 (TB) 의 최적 정책과 동일함을 증명했습니다 ($\pi_\theta(y) \propto R(y)$).
• 분산 감소: Reinforce Leave-One-Out 과 유사하게 한 샘플을 다른 샘플의 기준선으로 사용함으로써, CTB 는 외부 기준선 네트워크나 $Z$ 추정 없이도 기울기 분산을 암묵적으로 감소시킵니다.

B. 노이즈 기울기 가지치기 (Noisy Gradient Pruning, NGP)

쌍대 비교에 내재된 보상 노이즈 증폭 위험을 완화하기 위해, S-GFN 은 중요도 기반 필터를 도입합니다.

• 메커니즘: 기울기는 로그 보상의 절대 차이가 임계값 $\sigma$를 초과하는 궤적 쌍에 대해서만 계산됩니다:
  $$L_{NGP} = \mathbb{I}(|\log R(y_1) - \log R(y_2)| > \sigma) \cdot L_{CTB}$$
• 효과: 이는 최적화기가 피해 LLM 이나 유해성 분류기의 확률적 특성으로 인한 미미한 보상 차이에 기반하여 업데이트하는 것을 방지하여, 노이즈가 있는 기울기를 효과적으로 가지치기합니다.

C. Min-K 유창성 안정화기 (Min-K Fluency Stabilizer, MKS)

모델이 종종 가짜 유해성 점수를 받는 의미 없는 텍스트의 지역 최소값으로 수렴하는 것을 방지하기 위해, S-GFN 은 유창성 제약을 적용합니다.

• 메커니즘: 참조 모델에 대한 과적합을 유발할 수 있는 전역 KL-발산 패널티 대신, MKS 는 생성된 시퀀스에서 참조 모델 하에서 가장 확률이 낮은 $k$개의 토큰에 대한 평균 로그 확률인 Min-K 통계량을 계산합니다.
• 강제 패널티: Min-K 점수가 임계값 $T_{mks}$ 미만으로 떨어지면 보상은 0 (또는 큰 음수 값) 으로 잘립니다.
  $$R_{MKS}(y) = \mathbb{I}(M_k(y) \ge T_{mks}) \cdot R(y)$$
• 장점: 이는 유창하고 언어적으로 일관된 프롬프트만이 고보상으로 간주되도록 보장하여, 모델이 무의미한 텍스트를 생성하지 않고도 유효한 공격의 더 넓은 공간을 탐색할 수 있게 합니다.

3. 주요 기여

1. Stable-GFN 프레임워크: 분할 함수 추정 ($Z$) 을 제거하는 새로운 레드팀링 프레임워크로, LLM 을 위한 GFN 의 학습 불안정성 주요 원인을 해결합니다.
2. 대비적 궤적 균형 (CTB): TB 와 동일한 최적 분포를 달성하지만 분산을 현저히 낮추고 학습 안정성을 높이는 새로운 목적 함수입니다.
3. 노이즈에 대한 강건성: NGP 와 MKS 의 통합은 노이즈가 있는 분류기와 의미 없는 텍스트 보상이라는 레드팀링의 고유한 과제를 해결하여 모드 붕괴와 보상 해킹을 방지합니다.
4. 보편성: 이 논문은 CTB 와 NGP 가 레드팀링뿐만 아니라 분자 생성 및 하이퍼그리드 탐색과 같은 다른 분포 매칭 작업에서도 효과적임을 입증했습니다.

4. 실험 결과

저자들은 다양한 피해 모델 (Qwen, Llama, Gemma) 과 방어 설정에서 PPO, Jailbreak R1, Rainbow Teaming, 표준 GFN (TB) 을 포함한 베이스라인과 S-GFN 을 비교 평가했습니다.

• 다양성 (고유 공격): S-GFN 은 134 개의 고유한 공격 프롬프트를 생성하여 표준 GFN 베이스라인 (17 개) 과 PPO(3 개) 대비 획기적인 개선을 보였습니다. 이는 GFN 베이스라인보다 약 7 배 더 많은 고유 공격을 생성한 것입니다.
• 공격 성공률 (ASR): S-GFN 은 **약 92.5%**의 높은 ASR 을 유지하여 최상위 성능의 베이스라인 (GFN: 93.7%, PPO: 91.7%) 과 비교할 수 있었으며, 다양성이 효과성을 희생한 것이 아님을 입증했습니다.
• 교차 공격 전이성: S-GFN 의 공격을 사용하여 피해 모델을 안전 미세 조정했을 때, 다른 방법의 공격에 대해 훨씬 더 강건해졌습니다 (GFN 공격의 ASR 은 거의 0% 로 하락). 반대로, S-GFN 의 공격은 다른 방법으로 미세 조정된 모델에 대해서도 유효하게 유지되어 우수한 일반화 능력을 보여주었습니다.
• 전이 공격: S-GFN 은 보지 못한 피해 모델 (예: Gemma3, Llama3.2) 로 공격을 성공적으로 전이하여, 전이 시나리오에서 ASR 과 다양성 모두 다른 방법들을 능가했습니다.
• 효율성: $O(N^2)$의 쌍대 비교에도 불구하고, 이러한 연산이 스칼라 연산이므로 계산 비용은 무시할 수준입니다. 실제 역전파 비용은 $O(N)$으로 유지되어 S-GFN 은 표준 GFN 과 거의 동일한 효율성을 가집니다.

5. 의의

이 논문은 LLM 안전성 연구에서 중요한 진전을 의미합니다. GFN 의 학습 불안정성을 해결함으로써, 모드 붕괴로 인해 기존 RL 방법이 놓치는 다양하고 영향력 있는 취약점의 체계적 발견을 가능하게 합니다.

• 안전성 측면: 배포 전에 더 넓은 범위의 공격 벡터를 식별할 수 있도록 레드팀링을 위한 더 강건한 도구를 제공합니다.
• 방어 측면: S-GFN 이 생성한 다양한 공격은 더 넓은 범위의 미래 공격에 대해 더 탄력적인 안전 미세 조정 모델을 이끕니다.
• 기계 학습 측면: 제안된 대비적 궤적 균형은 고차원 노이즈 환경에서의 분포 매칭을 위한 일반화 가능한 솔루션을 제공하여, GFN 의 적용 범위를 레드팀링을 넘어 신약 개발 및 재료 과학과 같은 분야로 확장합니다.

요약하자면, Stable-GFlowNet은 GFN 의 이론적 다양성과 실제 LLM 레드팀링에 필요한 실용적 안정성 사이의 간극을 성공적으로 연결하여, 다양하고 효과적인 적대적 프롬프트 생성을 위한 새로운 최첨단 (state-of-the-art) 을 설정합니다.