기술 요약: AEGPO – 확산 모델을 위한 적응형 엔트로피 유도 정책 최적화 (Adaptive Entropy-Guided Policy Optimization)

1. 문제 정의

인간 피드백 기반 강화학습(RLHF), 특히 그룹 상대 정책 최적화(GRPO)는 확산(Diffusion) 및 플로우 매칭(Flow-matching) 모델을 인간의 선호도에 맞게 정렬하는 표준 패러다임이 되었습니다. 그러나 기존의 GRPO 변형 방식들은 정적 휴리스틱(Static Heuristics)에 의존함으로 인해 다음과 같은 두 가지 근본적인 비효율성을 겪습니다:

1. 비효율적인 샘플 할당: 표준 방식은 모든 프롬프트에 균등한 수의 롤아웃(Rollout)을 할당하며, 샘플 간의 "학습 가치(Learning Value)" 차이를 무시합니다. 어떤 프롬프트는 정책 변화를 거의 일으키지 않는 반면, 어떤 프롬프트는 상당한 개선을 이끌어냅니다.
2. 정적 타임스텝 탐색: 탐색 전략(예: 트리 기반 롤아웃에서의 브랜칭)은 일반적으로 고정된 사전 정의 스케줄을 따릅니다. 이는 디노이징 과정의 동적인 특성을 반영하지 못하며, 프롬프트와 훈련 단계에 따라 결정적인 순간이 달라진다는 점을 간과합니다.

이러한 한계는 저가치 샘플에 계산 자원을 낭비하고, 비결정적 타임스텝에서 최적의 탐색을 방해하여, 수렴 속도를 늦추고 최종 정렬 성능을 제한하는 결과를 초낳습니다.

2. 방법론: AEGPO

저자들은 내재적 어텐션 신호를 활용하여 샘플 선택과 탐색을 동적으로 유도하는 플러그 앤 플레이(Plug-and-play) 프레임워크인 **적응형 엔트로피 유도 정책 최적화(AEGPO)**를 제안합니다. 핵심 통찰은 **어텐션 엔트로피(Attention Entropy)**가 강력한 이중 신호 대리 지표(Dual-signal proxy) 역할을 한다는 것입니다.

2.1 핵심 신호

이 방법은 모델의 어텐션 맵에서 파생된 두 가지 내재적 지표를 정의합니다:

• $\Delta$Entropy (샘플 수준 신호): 현재 정책과 베이스 정책 사이의 어텐션 분산(Dispersion)의 상대적 변화량입니다.
  • 정의: $\Delta\text{Entropy} = \frac{1}{T} \sum_{t=1}^{T} |\text{Entropy}_{\theta}(t) - \text{Entropy}_{\theta_{base}}(t)|$.
  • 기능: 이는 샘플 학습 가치의 견고한 대리 지표 역할을 합니다. 높은 $\Delta$Entropy는 프롬프트가 모델로 하여금 어텐션 전략을 크게 재구성하도록 강제함을 의미하며, 이는 높은 학습 잠재력을 시사합니다.
• Entropy($t$) (타임스텝 수준 신호): 특정 디노이징 단계 $t$에서의 절대적 어텐션 분산입니다.
  • 정의: 단계 $t$에서 이미지 특징에 대한 텍스트 토큰 간 어텐션 가중치의 평균 샤논 엔트로피(Shannon entropy)입니다.
  • 기능: Entropy($t$)의 피크(Peak)는 결정적인 탐색 타임스텝을 식별합니다. 이러한 피크는 높은 불확실성이나 광범위한 의미론적 통합이 발생하는 순간(주로 초기 구조 형성 단계 및 후기 정밀화 단계에서 발생)과 일치하며, 이 시기에 탐색을 수행할 때 보상 분산과 다양성이 가장 높게 나타납니다.

2.2 이중 수준 적응 전략

AEGPO는 총 롤아웃 예산을 늘리지 않으면서도 이러한 신호들을 사용하여 계산을 재할당합니다:

1. 글로벌 적응형 할당 (Global Adaptive Allocation):

   • 프롬프트 배치를 중앙값 $\Delta$Entropy를 기준으로 "고가치" 및 "저가치" 티어로 나눕니다.
   • 고가치 프롬프트(높은 $\Delta$Entropy)에는 더 큰 롤아웃 예산($r_{high}$)을 할당하고, 저가치 프롬프트에는 더 적은 롤아웃을 할당합니다.
   • 제약 조건: $r_{high} + r_{low} = 2r_{avg}$를 유지하여 총 계산량이 일정하게 유지되도록 합니다.

2. 로컬 적응형 탐색 (Local Adaptive Exploration):

   • 고정된 브랜칭 스케줄 대신, 각 롤아웃에서 가장 높은 Entropy($t$) 피크를 보이는 상위-$K$ 타임스텝을 동적으로 식별합니다.

• 브랜칭(탐색)은 오직 이 식별된 고분산 타임스텝에서만 수행됩니다.
  • 이를 통해 탐색이 가장 정보량이 많은 디노이징 궤적에 집중되도록 보장합니다.

3. 주요 기여

• 이중 신호 발견: 본 논문은 어텐션 엔트로피를 이중 신호 내재적 대리 지표로 확립하였으며, 여기서 상대적 변화($\Delta$Entropy)는 샘플 학습 가치를 나타내고, 절대적 피크(Entropy($t$))는 결정적인 탐색 타임스텝을 식별합니다.
• AEGPO 프레임워크: 경량화된 엔트로피 유도 메커니즘을 사용하여 샘플 선택과 타임스텝 탐색을 개선하는 이중 수준 적응형 프레임워크를 도입했습니다.
• 일반화 가능성: AEGPO가 모델 불가지론적(Model-agnostic)이며 프레임워크 불가지론적(Framework-agnostic)임을 입증하여, 다양한 베이스 모델(FLUX.1-dev, SD3.5-M)과 다양한 정책 최적화 방법(DanceGRPO, BranchGRPO, FlowGRPO, DiffusionNFT)을 성공적으로 향상시켰음을 보여주었습니다.

4. 실험 결과

FLUX.1-dev와 SD3.5-M을 베이스 모델로 사용하여 HPDv2.1 및 Pick-a-Pic 데이터셋에 대해 실험을 수행했습니다.

• 수렴 속도: AEGPO는 정책 최적화를 크게 가속화합니다. 표준 변형 모델 대비 DanceGRPO에서는 2배 빠른 수렴을, DiffusionNFT에서는 5배 빠른 수렴을 달로하며, 더 적은 단계 내에 우수한 최종 보상에 도달합니다.
• 성능 지표:
  • FLUX.1-dev에서 BranchGRPO + AEGPO 조합은 새로운 SOTA HPS-v2.1 점수인 0.374(기존 0.363에서 상승)와 Pick Score 0.232(기존 0.229에서 상승)를 달성했습니다.
  • SD3.5-M에서 DiffusionNFT + AEGPO는 HPS-v2.1을 0.343으로, ImageReward를 1.57로 개선했습니다.
• Reward-KL 트레이드오프: AEGPO 모델은 일관되게 베이스라인보다 Reward-KL 파레토 프런티어(Pareto frontier)에서 우위를 점하며, 주어진 정책 편차(KL divergence) 수준에서 훨씬 높은 정렬 보상을 달성합니다.
• 절제 연구 (Ablation Studies): 글로벌 할당(G)과 로컬 탐색(L)은 각각 독립적인 이점을 제공하며, 결합된 접근 방식(G+L)이 가장 높은 성능을 냅니다.
• 계산 오버헤드: 중간 어텐션 맵 추출로 인해 약 11.1%의 단계 시간 증가와 미미한 VRAM 증가라는 적은 오버헤드가 발생하지만, 이는 최고 성능에 도달하기 위해 필요한 전체 훈련 단계의 극적인 감소를 고려할 때 충분히 정당화됩니다.

5. 의의 및 주장

본 논문은 AEGPO가 확산 모델 정렬의 핵심적인 비효율성인 정적 샘플링 문제를 해결한다고 주장합니다. 휴리스틱 기반 할당에서 내재적 신호 기반 적응으로 전환함으로써, 이 방법은 계산 자원을 가장 정보가 많은 샘플과 생성 과정의 가장 결정적인 순간에 집중시킵니다.

저자들은 AEGPO가 새로운 손실 함수나 복잡한 보상 모델링을 요구하지 않는 일반적이고 플러그 앤 플레이가 가능한 모듈임을 강조합니다. 대신, 모델 자체의 내부 어텐션 역학을 활용하여 최적화를 가이드함으로써, 다양한 확산 아키텍처와 RLHF 프레임워크 전반에서 고품질 정렬을 위한 더 효율적인 경로를 제공합니다. 이 연구는 어텐션 엔트로피와 같은 내재적 신호가 생성 모델의 정책 최적화 효율을 높이는 데 있어 아직 충분히 활용되지 못한 자원임을 시사합니다.