FewStepRL: 소수 단계 확산 모델을 위한 강화학습 미세 조정

확산 모델의 후처리 학습에서는 생성 결과를 보상으로 평가하고, 그 신호를 이용해 모델을 업데이트하는 강화학습이 널리 활용됩니다. 하지만 현재 정책을 매 반복마다 롤아웃하는 방식에서는 샘플링 단계가 많을수록 롤아웃 비용이 학습 비용의 큰 부분을 차지합니다. 소수 단계 확산 모델은 이 비용을 줄일 수 있지만, 단순히 기존 강화학습 목표에서 샘플링 단계만 줄이는 것으로는 충분하지 않습니다.

가장 중요한 차이는 속도의 정의입니다. 다단계 확산 모델은 수송 과정의 특정 시점에서의 순간 속도를 중심으로 설명되는 경우가 많습니다. 반면 소수 단계 모델은 일정한 시간 구간에 걸친 평균 속도를 학습합니다. 이 차이를 고려하지 않고 기존 목표를 적용하면, 목적함수가 가정하는 모델의 동작과 실제 모델의 표현 사이에 불일치가 생길 수 있습니다.

FewStepRL은 새로운 단일 보상함수를 제안하기보다, 여러 강화학습 목표를 소수 단계 모델에 맞게 변환하는 공통 프레임워크를 제시합니다. 핵심 절차는 세 부분으로 구성됩니다.

• x 공간으로의 변환: 서로 다른 목표를 공통 정책 예측의 형태로 표현합니다. 기본적으로는 깨끗한 샘플의 끝점에서 다루고, 전이 또는 플로우 맵 목표에서는 해당 중간 수준을 명시적으로 유지합니다. 이를 통해 밀도, 속도, 보상을 기반으로 한 목표를 같은 표현에서 비교할 수 있습니다.
• 그래디언트 분해: 각 목표의 그래디언트를 기준값을 향한 자기 회귀 형태의 회귀 항과 보상으로 인한 증분 항의 합으로 나눕니다. 따라서 방법 간 차이는 그래디언트 관점에서 앵커와 가중치를 어떻게 선택하는지로 정리됩니다.
• 앵커 적응: 소수 단계 모델의 수송 종점이나 접선 방향을 이용해 앵커를 계산합니다. 순간 속도가 필요하면 모델 자체에서 복원하거나 보조 네트워크를 사용할 수 있습니다.

이 구조를 바탕으로 FlowGRPO, AWM, DiffusionNFT, DiffusionOPSD, ReFL의 소수 단계 버전이 하나의 모듈형 코드베이스에 구현되었습니다. 연구자는 알고리즘마다 별도의 학습 파이프라인을 새로 만들기보다, 앵커와 가중치, 속도 추정 방식을 바꾸며 여러 목표를 비교할 수 있습니다.

실험 결과는 보상의 성질에 따라 적합한 방법이 달라진다는 점을 보여줍니다. 보상이 미분 가능하고 그 그래디언트를 계산하는 비용이 감당 가능한 경우, 보상 그래디언트 기반 방법이 가장 효과적이었습니다. 반대로 보상이 미분 불가능하거나 그래디언트 계산 비용이 큰 경우에는 DiffusionNFT 계열이 비교적 적은 반복만으로 높은 보상에 도달하는 경향을 보였습니다. 다만 이 방식은 반복 한 번당 비용이 더 높을 수 있습니다.

학습에 사용한 단계보다 추론 단계를 늘리는 방식은 전체적으로 일관된 이점을 제공하지 않았습니다. 그러나 일부 출력은 학습 시점의 두 배에서 네 배 정도로 단계를 늘렸을 때 눈에 띄게 개선되었습니다. 따라서 추가 단계는 특정 샘플의 품질을 높일 수 있지만, 전체 성능 향상을 보장하지는 않습니다.

FewStepRL의 주요 의미는 소수 단계 확산 모델의 강화학습을 개별적인 구현 문제에서 재사용 가능한 설계 문제로 바꿔 바라보게 한다는 데 있습니다. 빠른 샘플링을 제대로 활용하려면 단계 수뿐 아니라 속도의 정의, 정책 예측 위치, 보상 신호, 회귀 기준점이 서로 일관되어야 합니다. 이는 추론 비용과 보상 최적화 비용을 함께 줄이려는 확산 모델 시스템에 유용한 분석 관점을 제공합니다.

출처: Hugging Face Daily Papers