AEGPO 확산모델에 정책최적화를 섞는 아이디어

 

AEGPO 확산모델에 정책최적화를 섞는 아이디어

GRPO 계열로 확산/플로우 모델을 정렬할 때, 병목은 “무엇을 더 많이 굴릴지(샘플)”와 “언제 분기할지(타임스텝)”를 정적으로 정해 생기는 낭비입니다. AEGPO는 이 문제를 내부 신호 하나(Attention Entropy)로 샘플 수준과 타임스텝 수준으로 분해해 해결합니다. ∆Entropy로 프롬프트별 롤아웃을 재배분하고, Entropy(t) 피크로 분기 시점을 고르는 방식이라 reward attribution 없이도 동작한다는 점이 실용적입니다. 다만 ∆Entropy=학습가치라는 인과, median split의 조잡함, 피크=좋은 탐색이라는 해석, 그리고 “경량” 대비 오버헤드·RM 오버피팅 우려는 더 엄밀한 보강이 필요합니다.

∆Entropy: ‘학습가치’ 프록시로서 강점과 인과의 빈틈

논문은 Attention Entropy를 “이중 신호(dial-signal)”로 정의합니다. 먼저 각 타임스텝 t에서 텍스트 토큰에 대한 이미지 쿼리의 어텐션 분포로 Shannon entropy를 만들고(식 1~4), 현재 정책 θ와 베이스 정책 θ_base의 엔트로피 차이를 ∆Entropy(t)=|Entropyθ(t)−Entropy_base(t)|로 정의합니다(식 5). 그리고 이를 전체 디노이징 구간에 평균내어 샘플(프롬프트) 가치로 ∆Entropy를 만듭니다(식 6). 이 문제 분해가 깔끔한 이유는 “보상 신호를 어느 토큰/어느 스텝에 귀속시키지 않고도” 정책 변화의 크기를 내부에서 측정할 수 있기 때문입니다. 논문은 Figure 2에서 ∆Entropy가 낮은 프롬프트는 학습이 진행돼도 이미지 변화가 작고 엔트로피 궤적이 서로 뭉쳐 있으며, ∆Entropy가 높은 프롬프트는 시각적 개선이 크고 엔트로피 궤적이 크게 벌어진다는 정성 예시를 제공합니다.

정량 근거도 있습니다. Figure 3(Left)에서 평균 ∆Reward와 평균 ∆Entropy가 함께 상승/정체하는 동조 관계를 보여주며, Figure 3(Right)에서는 high-∆Entropy 데이터만으로 학습한 모델이 low-∆Entropy 데이터만 학습한 모델보다 더 빠르게 수렴하고 최종 reward도 높다고 제시합니다.
이 흐름은 “∆Entropy가 큰 샘플이 정책 업데이트를 더 강하게 유도한다”는 사실을 꽤 설득력 있게 보여줍니다.

하지만 사용자 비평대로, 여기서 “학습가치”라는 단어는 인과 주장까지 포함해 위험해집니다. ∆Entropy는 본질적으로 “현재 정책이 베이스 대비 얼마나 흔들리나”에 가깝습니다. 그러면 다음 대안 설명을 강하게 배제해야 합니다.

난이도/불확실성 가설: ∆Entropy가 큰 샘플은 단지 현재 모델이 불확실해서 엔트로피가 흔들리는 ‘어려운 샘플’일 수 있고, 그게 장기적으로 좋은 일반화를 만든다는 보장은 없습니다.

드리프트 편향 가설: ∆Entropy가 큰 샘플을 더 자주 학습하면, 정책이 베이스에서 크게 벗어나는(드리프트가 큰) 주제/스타일로 과도하게 쏠릴 수 있습니다. 논문은 reward-KL 파레토에서 AEGPO가 같은 KL에서 더 높은 reward를 만든다고 주장하지만(Figure 7), 샘플 단위로 “어떤 종류의 드리프트가 강화됐는지”까지는 충분히 분해되지 않습니다.

논문이 스스로도 힌트를 주는 지점이 하나 있습니다. Discussion에서 KL divergence는 “프롬프트별 신호”가 아니라 전역 드리프트라 샘플 배분에 부적합하다고 말하고, 대신 엔트로피가 prompt-resolved라 유리하다고 주장합니다.
이 말은 곧 “∆Entropy가 KL과는 다른 정보를 준다”는 뜻이기도 합니다. 따라서 더 강한 주장으로 가려면, KL(또는 드리프트 통제) 하에서도 ∆Entropy가 추가 예측력을 갖는지가 실험으로 보여져야 합니다. 예를 들면, 샘플별 KL을 맞춘 상태에서 ∆Entropy 상/하가 다음 스텝의 reward gain이나 held-out 일반화에 여전히 영향을 주는지 같은 분석입니다.

정리하면, 논문이 이미 보여준 것은 “∆Entropy가 단기 학습 속도를 올린다”이고, 독자가 더 보고 싶은 것은 “∆Entropy가 장기 일반화/안정성에도 좋은 샘플을 고른다”입니다. 이 간극이 메워지면 ‘학습가치’라는 표현이 훨씬 안전해집니다.

논문 주장	독자가 요구하는 추가 근거
∆Entropy가 샘플 학습가치 프록시	난이도/불확실성 vs 장기 가치 분리, 드리프트(KL) 통제 후에도 예측력이 남는지, held-out 프롬프트·스타일 일반화에서 이득 유지 여부
high-∆Entropy 학습이 더 빠르고 더 높게 수렴	초기/후기 학습 단계별로도 같은 결론인지, 상위 극소수(long-tail) 집중이 더 좋은지(배분 전략 비교)

Entropy(t): 피크 분기가 ‘탐색 가치’라는 직관은 좋지만 정의가 더 필요함

논문이 제시한 두 번째 축은 타임스텝 가치입니다. Entropy(t)는 t에서 텍스트 토큰에 대한 어텐션이 얼마나 분산됐는지(주의가 넓게 퍼졌는지)를 나타내고, 논문은 Entropy(t)의 상위 피크가 “고가치 탐색 타임스텝”이라고 주장합니다. Figure 4에서 Top-K 엔트로피 피크가 디노이징 초반(t≈1)과 후반(t≈13~15)에 몰리는 U-shape 분포를 보이며, 이를 “초반은 거친 구조 결정, 후반은 세부 정련”이라는 해석으로 연결합니다. 그리고 Local Adaptive Exploration은 각 롤아웃에서 Entropy(t) 전체 궤적을 계산한 뒤, TopK(Entropy(t))에서만 branching을 수행합니다(K=4로 고정, 페이지 6).

정량 근거로는 Table 1이 있습니다. HPDv2.1 1000 prompts 평균에서 고정 분기 스케줄들(예: (0,2,4,8) 등)보다 Entropy-Guided가 Reward Std, LPIPS MPD, TCE가 모두 조금씩 높다고 보고합니다.
즉 “피크에서 분기하면 결과 다양성과 보상 변동성이 커진다”는 사실은 데이터로 제시합니다.

하지만 사용자 비평이 말한 것처럼, Reward Std가 커지는 것이 항상 좋은 탐색은 아닙니다. 불안정한 탐색이나 모드 붕괴도 분산을 키울 수 있기 때문입니다. 논문은 LPIPS MPD/TCE 같은 다양성 지표를 같이 올려 “좋은 다양성”을 시사하지만, 결국 핵심은 “정렬 품질(선호도, 조건 충족)”과의 트레이드오프를 함께 보여주는 것입니다.
특히 이 논문은 RLHF 맥락이고 최종 평가는 RM score 중심으로 이루어지므로, “피크 분기가 품질도 같이 유지/개선한다”는 것을 더 명확히 보여줘야 합니다. 예를 들어 분기 전략별로 reward 평균, 하위 10% 꼬리 성능, 그리고 인간 선호(가능하다면)까지 함께 제시하면 “분산 증가=탐색 가치” 해석이 훨씬 단단해집니다.

또 하나의 보강 포인트는 “왜 피크가 아닌 곳에서 분기하면 덜 유용한가”를 반증 실험으로 보여주는 것입니다. 논문은 고정 스케줄과 비교하지만, 피크 주변을 일부러 피하는 스케줄(anti-peak) 같은 대조군이 있으면 “피크의 의미”가 더 선명해집니다.

배분: 플러그인 장점은 확실하지만 median split은 거칠고 비용 논쟁이 남음

AEGPO의 실전 매력은 플러그인 성격입니다. 논문은 AEGPO가 새로운 loss를 제안하는 것이 아니라, 기존 GRPO 변형들(DanceGRPO, BranchGRPO, Flow-GRPO, DiffusionNFT)에 “샘플 배분 + 분기 시점 선택”만 얹는 구조라고 강조합니다. 실제로 Global Adaptive Allocation은 간단합니다. 배치 내 각 샘플 i에 대해 vi=∆Entropy_i를 계산하고, 배치 median으로 high/low 2티어로 나눈 뒤 rollout budget을 다르게 줍니다. 평균 12 롤아웃이면 (rlow,rhigh)=(8,16)처럼 배정하되 총 예산은 유지합니다(페이지 5). 이 단순함 덕분에 기존 파이프라인 변경 범위가 제한되고, “reward attribution 없이도” 자동으로 돌아갑니다.

또한 효율/수렴 가속을 전면 지표로 내세운 점은 설득 포인트입니다. Figure 1에서 DanceGRPO에서는 2×, DiffusionNFT에서는 5× 빠르게 baseline의 최종 reward에 도달한다고 시각화합니다.
Table 2에서는 FLUX.1-dev와 SD3.5-M에서 여러 RM(HPSv2.1, PickScore, ImageReward, GenEval)로 AEGPO가 일관 개선을 보인다고 보고합니다(예: FLUX.1-dev에서 BranchGRPO HPS-v2.1 0.363→0.374).
Ablation(Table 3)에서도 Global(G)만, Local(L)만 켜도 개선되고 둘을 합치면 더 좋아진다고 정리합니다.

다만 사용자 비평대로 median split은 너무 거칩니다. 분포가 long-tail이면 상위 극소수에 더 몰아야 할 수도 있고, 초기에는 거의 모두가 high-∆Entropy로 뭉쳐 분리가 무의미해질 수도 있습니다. 논문은 20-step warmup 동안은 uniform allocation을 쓰고(페이지 6), 그 후에야 ∆Entropy가 “정보성을 갖는다”고 말해 초기 붕괴를 완화하지만, 여전히 “2티어 고정”의 한계는 남습니다.
따라서 다음 비교가 있으면 논문이 더 강해집니다.

median split(2티어) vs top-p quantile(상위 p% 집중) vs 연속 배분(softmax/temperature)

학습 단계별(early vs late) 최적 배분이 달라지는지

비용 논쟁도 남습니다. 논문은 attention map을 추출하기 위해 특정 레이어(FLUX.1-dev/SD3.5-M에서 5,10,15)를 사용하고, 오버헤드 분석에서 FLUX.1-dev 기준 step당 52초 추가(469→521초, +11.1%)와 VRAM 33.5GB→34.5GB(+1GB)를 “무시 가능”이라 말합니다(페이지 8).
그런데 인프라가 빡빡한 환경(더 큰 해상도, 긴 rollout, 멀티노드)에서는 11%가 병목이 될 수 있습니다. 논문이 더 실전적으로 설득되려면 “wall-clock to target reward”와 “GPU-hours”로 비교해야 합니다. 논문도 “수렴이 빨라져서 이득”을 주장하니, 그 주장을 바로 비용 단위로 환산해 제시하면 논쟁이 줄어듭니다.
또한 “어떤 레이어/헤드만 보면 충분한가” 같은 더 적극적 경량화(부분 샘플링, 낮은 해상도 주의맵, head subset) 분석이 붙으면, 플러그인 가치가 더 커집니다.

마지막으로 평가가 RM score 중심이라는 점에서 RM 오버피팅 우려는 남습니다. 논문은 여러 RM으로 평가해 편향을 줄이려 하지만(Table 2), 근본적으로는 RM을 더 잘 맞추는 방향으로 갈 가능성은 항상 있습니다.
따라서 사람이 직접 보는 선호 평가나 RM-agnostic 지표(예: 특정 프롬프트 조건 충족률의 수동 검증, 혹은 규칙 기반 체크)를 소량이라도 추가하면, “AEGPO는 RM 오버피팅을 더 효율화하는 도구”라는 반론을 약화시킬 수 있습니다.

AEGPO는 Attention Entropy를 샘플(∆Entropy)과 타임스텝(Entropy(t) 피크)으로 분해해, 롤아웃 배분과 분기 시점을 동시에 적응화한 점이 강점입니다. 다만 ∆Entropy=학습가치 인과, median split의 거칠음, 피크 분기의 품질-다양성 트레이드오프, 오버헤드·RM 오버피팅은 추가 실험(드리프트 통제, 연속 배분, wall-clock/GPU-hours, 인간 평가)로 보강돼야 합니다.

자주 묻는 질문 (FAQ)

Q. ∆Entropy는 정확히 무엇을 의미하나요? A. 각 timestep의 평균 Attention Entropy를 현재 정책 θ와 베이스 θ_base에서 계산한 뒤, 그 차이의 절댓값을 ∆Entropy(t)=|Entropyθ(t)−Entropy_base(t)|로 두고, 이를 전체 디노이징 구간에 평균낸 값이 ∆Entropy입니다. 논문은 이를 “샘플이 정책의 어텐션 전략을 얼마나 바꾸게 하는가”의 프록시로 사용합니다.

Q. Entropy(t) 피크에서 분기하면 왜 좋은가요?
A. 논문은 피크가 디노이징 초반과 후반에 몰리는 U-shape 분포를 보이며(Figure 4), 피크 시점이 “주의가 넓게 퍼져 민감도가 큰 구간”이라 해석합니다. 실제로 피크 기반 분기 전략은 고정 스케줄보다 Reward Std와 다양성(LPIPS MPD, TCE)이 조금 더 높게 보고됩니다(Table 1). 다만 분산 증가는 불안정성과도 연결될 수 있어 품질 지표와의 동시 보고가 더 필요합니다.

Q. AEGPO의 오버헤드는 정말 ‘경량’인가요?
A. 논문은 FLUX.1-dev에서 attention map 추출로 step당 시간이 469초에서 52초 늘어 +11.1%가 되고, VRAM은 33.5GB→34.5GB로 +1GB 증가한다고 보고합니다. 대신 수렴이 2×~5× 빨라져 전체 효율이 좋아진다는 논리입니다. 실무에서는 wall-clock과 GPU-hours로 목표 reward까지의 총 비용을 함께 비교하는 것이 가장 안전합니다.

[출처]
https://arxiv.org/html/2602.06825v1