TDM-R1: 비미분 보상을 활용한 소수 단계 확산 모델의 강화학습

기존의 빠른 이미지 생성 모델(Few-step)은 복잡한 글자나 사물 개수를 맞추는 데 한계가 있었고, 이를 개선하기 위한 강화학습은 수학적으로 계산 가능한 특수한 점수 모델만 사용할 수 있었다. TDM-R1은 사람이 직접 매긴 점수나 정답 여부 같은 '비미분' 점수도 학습에 활용할 수 있게 하여, 생성 속도는 유지하면서도 품질을 비약적으로 높였다.

기존 확산 모델은 수십 번의 노이즈 제거 단계를 거쳐 이미지를 만드는데, 이를 1~4번으로 줄인 것이 소수 단계(Few-step) 모델이다. 하지만 단계가 너무 적으면 복잡한 지시사항을 놓치기 쉽다. 이를 강화학습으로 교정하려면 '이미지가 얼마나 좋은지'를 나타내는 보상 함수의 기울기(Gradient)가 필요한데, 글자가 맞았는지 틀렸는지 같은 이진 보상은 기울기를 구할 수 없어 학습이 불가능했다.

TDM-R1은 이 문제를 해결하기 위해 '대리 보상 모델'을 도입한다. 먼저 생성된 이미지들에 대해 사람이 매긴 점수나 OCR 결과 같은 실제 보상을 수집한다. 그 다음, 이 데이터를 바탕으로 '어떤 이미지가 더 좋은지'를 판단할 수 있는 미분 가능한 별도의 딥러닝 모델(대리 보상 모델)을 먼저 학습시킨다.

특히 TDM 모델의 특징인 '결정론적 경로'를 이용한다. 노이즈에서 이미지로 가는 길이 정해져 있으므로, 최종 결과물이 좋다면 그 과정에 있는 중간 단계 이미지들도 좋은 점수를 받을 자격이 있다고 보고 단계별로 세밀하게 학습을 가이드한다. 결과적으로 단 4번의 계산만으로도 수십 번 계산한 모델보다 더 정확한 이미지를 생성하게 된다.

결정론적 샘플링 기반 중간 보상 할당은 TDM의 ODE 샘플링 경로를 활용한다. [노이즈 x_T에서 깨끗한 이미지 x_0까지의 결정론적 경로 입력] → [중간 단계 x_t에 대해 최종 보상 r(x_0)를 기반으로 한 기댓값 계산] → [각 단계별 보상 신호 생성] → [중간 단계에서도 정확한 가이드 제공].

대리 보상 모델(Surrogate Reward) 학습은 확산 모델 파라미터를 공유하는 보상 모델을 구축한다. [이미지 쌍(Positive/Negative) 입력] → [Bradley-Terry 모델 기반의 그룹 선호도 최적화 수행] → [미분 가능한 보상 값 출력] → [생성기 학습을 위한 안정적인 기울기 제공].

생성기 최적화는 보상 최대화와 역 KL 발산(Reverse KL) 최소화를 동시에 수행한다. [대리 보상 모델의 점수 입력] → [보상을 높이는 방향으로 가중치 갱신] → [동시에 베이스 모델의 분포에서 너무 벗어나지 않도록 제약] → [이미지 품질 유지 및 보상 해킹 방지].

GenEval 벤치마크에서 SD3.5-M 기반 TDM-R1은 92%의 정확도를 기록했다. 이는 80단계를 사용하는 베이스 모델(63%)과 상용 모델인 GPT-4o(84%)를 압도하는 수치다. 특히 사물 개수(Counting)와 위치 관계(Position)에서 큰 폭의 개선을 보였다.

시각적 텍스트 렌더링(OCR Accuracy) 실험에서 TDM-R1은 95%의 정확도를 달성하여, 강화학습 전(55%) 대비 비약적인 성능 향상을 입증했다. 이는 복잡한 텍스트가 포함된 이미지 생성 능력이 실질적으로 강화되었음을 의미한다.

60억 파라미터 규모의 Z-Image 모델에 적용했을 때도 4단계(4 NFE)만으로 100단계를 사용하는 원본 모델보다 높은 성능을 기록하며 확장성을 증명했다.

TDM-R1은 Trajectory Distribution Matching(TDM) 아키텍처를 기반으로 하며, 생성기(Generator)와 대리 보상 모델(Surrogate Reward Model)을 분리하여 교차 최적화하는 구조를 취한다.

Surrogate Reward는 확산 모델의 중간 단계 노이즈 이미지 x_t에 대해 조건부 확률 p(c|x_t)를 추정하도록 설계되었으며, 이는 Bradley-Terry 모델을 통해 그룹 단위 선호도 최적화(Group-based Preference Optimization)로 학습된다.

학습 안정성을 위해 EMA(Exponential Moving Average)를 적용한 동적 참조 모델(Dynamic Reference Model)을 사용하여 과도한 정규화를 방지하고 보상 모델의 오버피팅을 억제한다.

생성기 학습 시에는 인스턴스 레벨의 제약 대신 분포 레벨의 역 KL 발산 최소화를 사용하여, 소수 단계 증류(Distillation) 과정에서 발생하는 흐릿함(Blurriness) 문제를 해결하고 선명한 결과물을 얻는다.

논문에서는 구체적인 한계점을 명시적으로 언급하지 않았으나, 실험 결과에서 특정 벤치마크의 일부 지표가 베이스 모델 대비 소폭 하락하는 경우가 관찰되었다.