NormGuard: Reward-Preserving Norm Constraints in Flow-Matching Reinforcement Learning

기술 요약: NormGuard: Flow-Matching 강화 학습에서의 보상 보존형 노름 제약 (Reward-Preserving Norm Constraints)

문제 정의
강화 학습(RL) 사후 학습(post-training)은 플로우 기반 생성 모델을 인간의 선호도에 정렬하기 위한 표준적인 접근 방식이 되었습니다. 그러나 이 과정은 빈번하게 '보상 과최적화(reward over-optimization)'를 초래합니다. 이는 모델이 보상 프록시(reward proxy)의 불완전함을 악용하여, 보상 모델이 패널티를 부여하지 못하는 지점에서 지각적 품질 저하(예: 과도한 샤프닝, 색상 변화, 부자연스러운 조명)를 일으키는 현상입니다. 초기 중단(early stopping)이나 KL 정규화(KL regularization)와 같은 기존의 완화 전략은 사후 학습으로 인한 드리프트(drift)를 하나의 집합적인 양으로 취급합니다. 이들은 참조 모델(reference)로부터의 속도 편차 총량을 제한하지만, 보상 신호를 전달하는 성분과 시각적 아티팩트를 유발하는 성분을 구분하지 못합니다.

핵심 관찰 및 진단
저자들은 이러한 드리프트의 구체적인 구조적 징후를 식ر히 찾아냈습니다: RL 미세 조정은 다양한 방법론(NFT, AWM, DPO)과 모델(SD3.5, FLUX)에 걸쳐, 사전 학습된 참조 모델($v_{ref}$) 대비 스텝당 속도 노름($\|v_\theta\|$)을 5%에서 15%까지 일관되게 팽창시킵니다.

본 논문은 새로운 해결책을 제안하기에 앞서 두 가지 잠재적 솔루션을 조사했습니다:

1. 추론 시 재정규화(Inference-Time Renormalization): 추론 시 속도를 재조정하여 아티팩트를 완화하는 Classifier-Free Guidance(CFG)에서 영감을 얻어, 저자들은 RL로 미세 조정된 속도를 추론 시점에 참조 노름으로 다시 스케일링하는 실험을 진행했습니다. 그 결과, 이는 RL 학습된 모델에는 효과가 없음을 발견했습니다. CFG에서는 추론 시 속도 헤드(velocity heads)의 명시적인 결합으로 인해 팽창이 발생하지만, RL에 의한 팽창은 모델 가중치 내에 공동 적응(co-adapted)되어 있습니다. 따라서 추론 시 재스케일링은 보상은 유지할 수 있으나, 시각적 아티팩트(과도한 샤프닝, 부자연스러운 조명)를 해결하는 데는 실패했습니다.
2. 아드요인트 민감도 분석(Adjoint Sensitivity Analysis): 저자들은 훈련 중에 속도 크기를 억제하는 것이 일관된 보상 신호를 제거하는지 분석했습니다. 6,400개의 샘플에 대해 아드요인트 민감도 분석을 수행한 결과, 곱셈적 속도 스케일링(multiplicative velocity scaling) 하에서의 1차 보상 변화는 프롬프트 전반에 걸쳐 높은 노이즈 대 신호 비율(3배에서 100배)과 부호의 이질성을 보였습니다. 배치 수준의 평균은 0에 가깝게 유지되었는데, 이는 노름 팽창을 억제하는 것이 보상 이득을 체계적으로 제거할 가능성이 낮음을 시사합니다.

방법론: NormGuard
추론 시의 수정이 효과가 없고 노름 억제가 일관된 보상 비용을 수반하지 않는다는 발견을 바탕으로, 저자들은 훈련 시 정규화 도구인 NormGuard를 제안합니다.

• 메커니즘: NormGuard는 현재 속도 노름이 참조 모델의 노름을 초과할 때만 활성화되는 일방향 힌지 패널티(one-sided hinge penalty)입니다. 이는 모든 속도 국소적 기본 손실(NFT, AWM, DPO 등에 사용됨)에 가산적으로 더해집니다.
• 수식: 샘플링된 상태 $(x_t, t, c)$에 대하여, 정규화 항은 다음과 같이 정의됩니다:
  $$R_{norm}(\theta) = \lambda \cdot \max\left(0, \frac{\|v_\theta(x_t, t, c)\|^2 - \|v_{ref}(x_t, t, c)\|^2}{\|v_{ref}(x_t, t, c)\|^2}\right)$$
  여기서 $\lambda > 0$은 강도를 조절합니다. 이 패널티는 $\|v_\theta\| \leq \|v_{ref}\|$인 동안에는 업데이트를 제한하지 않으며, 오직 "초과된" 노름만을 처벌합니다.
• 호환성: 이 패널티는 기존의 파라미터 공간 제약 없이, 기본 손실과 동일한 속도 출력 공간에서 작동하므로 기존의 사후 학습 목적 함수와 자연스럽게 결합됩니다.

실험 결과
저자들은 두 가지 베이스 모델(SD3.5-Medium, FLUX.2-4B), 세 가지 사후 학습 방법(NFT, AWM, DPO), 그리고 두 가지 보상 모델(PickScore, HPSv2)에 대해 NormGuard를 평가했습니다.

• 지각적 품질: 텍스처 충실도 및 아티팩트 탐지와 같은 축에서 이미지 품질을 평가하기 위해 MLLM 심판(Qwen3.5 및 GPT-4.1)을 사용했을 때, NormGuard는 모든 7가지 구성에서 비정규화된 베이스라인보다 일관되게 우수한 성능을 보였습니다. 승률(win rates)은 유의미하게 높았습니다 (예: SD3.5-M + NFT의 경우 64% vs 35%).
• 포렌식 리얼리즘: NormGuard는 7개 설정 중 6개에서 포렌식 리얼리즘 점수(RealScore)를 개선하였으며, 부자연스러운 조명이나 과도한 샤프닝과 같은 아티팩트를 줄였습니다.
• 보상 보존: 이 방법은 베이스라인의 보상 이득을 대부분 보존했습니다. PickScore와 HPSv2의 차이는 미미하였으며(-0.004 ~ +0.011), 이는 방사형 팽창(radial inflation)을 제어하는 것이 보상 정렬을 담당하는 방향성 변화를 방해하지 않음을 나타냅니디.
• 적은 단계의 추론: NormGuard의 이점은 추론 단계 수가 감소함에 따라 증폭되었습니다. MLLM 승률의 격차는 28단계에서의 9%에서 4단계에서의 20%로 넓어졌으며, 이는 정규화 도구가 팽창된 속도로 인해 발생하는 이산화 오차(discretization errors)로부터 ODE 적분 과정을 안정화시킨다는 것을 시사합니다.
• 어블레이션(Ablations): 개선 사항은 초기 중단(early stopping)에 의한 것이 아니었습니다 (정규화된 모델이 베이스라인의 더 이른 체크포인트보다 우수한 성능을 보임). 또한, 이는 KL 정규화와 상보적이었으며, 서로 다른 실패 모드를 다루고 있었습니다.

의의 및 주장
본 논문은 표준적인 KL 정규화에 의해 패널티를 받는 집합적 드리프트가 유익하게 분해될 수 있다고 주장합니다: 노름 팽창은 지각적 아티팩트를 유발하지만 보상 신호는 거의 운반하지 않는 반면, 방향성 재정렬(directional realignment)이 보상 이득의 대부분을 차지한다는 것입니다.

NormGuard는 플로우 매칭 모델의 보상 과최적화를 해결하기 위한 정밀한 진단 및 해결책을 제공합니다. 훈련 중 과도한 속도 노름 성장을 억제함으로써, 보상 정렬을 희생하지 않고도 사진적 사실성을 개선하고 아티팩트를 줄입니다. 저자들은 이를 KL 정규화의 전역적 제약이나 CFG의 효과 없는 추론 시 수정과는 구별되는, RL 드리프트의 특정 구조적 징후를 다루는 단순하고 효과적인 개입으로 규정합니다.

한계점
저자들은 방사형 팽창의 동적인 기원이 여전히 미해결 과제로 남아있다고 언급했습니다. 보상 중립성에 대한 그들의 주장은 배치 수준에서의 1차 신호에 근거합니다. 개별 샘플의 방사형 섭동(radial perturbations)은 보상에 영향을 미치지만, 이러한 효과는 프롬프트 전반에 걸쳐 평균적으로 0이 됩니다. 또한, 현재의 분석은 속도 국소적 목적 함수(NFT, AWM, DPO)에 적용되며, 이를 Flow-GRPO와 같은 궤적 수준의 목적 함수로 확장하는 것은 향후 과제로 남겨두었습니다.