디퓨전 강화학습 RAM — 사전학습식 회귀로 Flow-GRPO 대비 학습 스텝 50배 절감

핵심 메시지

• 디퓨전과 플로우 매칭 모델의 사전학습이 잘 확장되는 이유는 그것이 결국 닫힌 형태의 목표를 향한 지도학습식 회귀이기 때문인데, 보상을 쓰는 사후학습은 그 구조를 버리면서 비싸고 불안정해졌다.
• RAM(Reinforce Adjoint Matching)은 KL 규제 아래 보상을 최대화하면 최적 생성 과정이 노이즈를 섞는 법칙은 그대로 둔 채 깨끗한 최종 샘플의 분포만 높은 보상 쪽으로 기울인다는 성질에서 출발한다.
• 학습 한 스텝은 현재 모델로 최종 샘플을 뽑고, 보상을 매기고, 사전학습과 똑같은 방식으로 노이즈를 섞은 뒤 회귀하는 것이 전부다. 확률미분방정식 롤아웃도, 역방향 adjoint 계산도, 보상 함수의 미분도 필요 없다.
• 보상이 미분 가능할 필요가 없다는 점이 실용적으로 크다. 이미지 속 글자를 읽는 광학문자인식 모델이나 사람 선호 모델처럼 미분이 정의되지 않는 채점기도 그대로 보상으로 쓸 수 있다.
• Stable Diffusion 3.5M에서 구성력, 글자 렌더링, 사람 선호 세 축 모두 가장 높은 보상을 기록했고, 비교 대상인 Flow-GRPO의 최고 보상에 최대 50배 적은 학습 스텝으로 도달했다.

쉽게 이해하기

발표자는 뮌헨 공대 박사과정의 안드레아스 베르크마이스터로, 디퓨전 모델을 강화학습으로 사후학습하는 연구를 소개했다. 문제의식은 단순하다. 디퓨전 사전학습은 깨끗한 샘플에 해석적으로 노이즈를 섞고 닫힌 형태의 목표에 회귀하는 것이라 놀랄 만큼 잘 확장되는데, 보상으로 모델을 정렬하는 사후학습은 그 회귀 구조를 포기하면서 훨씬 복잡해졌다는 것이다. 연구진은 사전학습을 확장 가능하게 만든 본질이 무엇인지로 되돌아가, 같은 원리에서 사후학습 목적함수를 다시 유도했다.

설정은 사전학습된 모델과 보상 함수가 주어진 상황이다. 실험에서는 텍스트-이미지 모델과 구성력 벤치마크를 썼다. 예컨대 '빨간 얼룩말'을 요구했을 때 사전학습 모델은 얼룩말은 그리지만 빨갛게 만들지 못하고, 비전 모델이 프롬프트와 내용을 대조해 부분 점수를 준다. 목표는 기대 보상을 최대화하되 사전학습 모델에서 너무 멀어지지 않도록 KL로 붙잡아 두는 것이고, 이 문제의 최적해는 참조 분포를 보상의 지수로 기울인 분포라는 잘 알려진 결과가 있다.

핵심 유도는 최적 속도장과 참조 속도장의 차이가 가치함수의 그래디언트, 즉 adjoint와 같다는 데서 시작한다. 기존 adjoint matching은 이를 얻기 위해 확률적 롤아웃을 시뮬레이션하고 그 전체를 거슬러 미분해야 했고, 이 때문에 보상이 미분 가능해야 하며 비싸고 수치적으로 불안정했다. RAM은 가치함수를 보상 항과 경로 비용 항으로 쪼갠 뒤, 보상 항에는 REINFORCE 항등식을 적용하고, 거기 남는 '브리지 스코어'를 베이즈 정리로 조건을 뒤집어 모델이 이미 갖고 있는 속도장으로 해석적으로 계산한다.

남은 경로 비용 항은 아예 버리는 근사를 택했다. 고차원에서 이 항이 골칫거리이기 때문이다. Stable Diffusion 3.5의 출력은 6만 5천 차원 수준이라 모델이 조금만 움직여도 참조 모델과의 차이가 커지고, 이를 스칼라로 뭉개면 분산이 보상 신호를 덮어버려 좋은 이미지라서 점수가 높은 건지 경로 비용이 커서 그런 건지 구분할 수 없게 된다. 발표자는 이것이 KL 규제를 통째로 버리는 것과는 다르다고 강조했다. 애초에 이 고정점 조건 자체가 KL 규제된 최적해에서 나온 것이고, 회귀 목표 안에 참조 모델이 앵커로 남아 있기 때문이다.

그 결과 남는 손실은 전체 경로가 아니라 세 점, 즉 모델이 뽑은 깨끗한 샘플과 가우스 노이즈, 그리고 둘을 섞은 노이즈 샘플에만 의존한다. 구현은 분류기 없는 가이던스로 20스텝 샘플링을 하고, 프롬프트 하나당 24번 굴려 평균 보상을 기준선으로 빼 분산을 줄인 뒤, 최종 샘플 하나에서 노이즈 샘플 여덟 개를 만들어 학습 타깃을 재사용한다. 벤치마크는 구성력, 문자 인식 기반 글자 렌더링, 사람 선호 세 가지였다.

주요 인사이트

• 속도 향상의 근원은 학습 신호의 상관관계에 있다. 정책 그래디언트 계열은 샘플링 경로 위의 지점에서만 손실을 계산하는데 그 지점들은 서로의 노이즈만 다른 버전이라 그래디언트가 강하게 상관된다. RAM은 최종 샘플에서 독립적인 노이즈 벡터로 되돌아가므로 훨씬 덜 상관된 학습 지점을 원하는 만큼 만들 수 있다.
• 비싼 것은 최종 샘플 생성과 보상 계산뿐이고, 거기서 파생되는 학습 타깃은 노이즈를 섞고 참조 모델을 한 번 평가하는 정도라 사전학습과 비슷한 비용이다. 그래서 한 번 만든 최종 샘플을 여러 타깃으로 재사용하는 것이 이득이 된다.
• 플로우 매칭 모델은 대응하는 확률적 스케줄의 계수가 노이즈 끝단에서 발산해 확률적 샘플링 자체가 까다롭고 결과가 흐려지는데, RAM은 결정론적 상미분방정식 샘플링만으로 충분해 이 문제를 비껴간다.
• 브리지 스코어에는 두 가지 추정기가 있다. 베이즈 정리로 얻은 쪽은 최적점에서 편향이 있지만 분산이 없고, 이론적으로 무편향인 쪽은 분산이 붙는다. 실험에서는 저차원과 고차원 모두에서 편향된 쪽이 문제를 일으키지 않아 실용적인 선택이 됐다.
• 보상만 좇으면 글자는 정확해지지만 이미지가 현실감을 잃는 보상 해킹이 생긴다. 발표자는 회귀 타깃에 참조 모델이 남아 있는 것이 RAM에서 이 현상이 덜한 이유라고 설명하며, 참조 모델 앵커가 없는 비교 대상과 대비했다.

자주 묻는 질문

관련 AI 소식