TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward

이 논문은 비미분 가능한 보상 신호를 활용하여 소수 단계 확산 모델의 학습을 가능하게 하는 새로운 강화 학습 패러다임인 TDM-R1 을 제안하고, 텍스트 렌더링 및 시각적 품질 등 다양한 작업에서 최첨단 성능을 입증합니다.

핵심

TDM-R1: "한 번에 완성하는 그림 그리기"를 위한 혁신적인 학습법

이 논문은 인공지능이 **매우 적은 단계 (4 단계)**로 고품질 이미지를 만들어내는 기술을 더 똑똑하게 만드는 방법을 소개합니다. 마치 "한 번에 그림을 완성하는 마법"을 더 정교하게 다듬는 과정이라고 생각하시면 됩니다.

이 복잡한 기술을 일상적인 비유로 쉽게 설명해 드릴게요.

---

1. 문제: "빠르지만, 가끔 엉망이 되는 그림"

지금까지 AI 는 그림을 그릴 때 수백 번의 작은 수정 (단계) 을 거쳐서 완성했습니다. 하지만 최근 기술 발전으로 4 번의 큰 수정만으로도 아주 빠른 속도로 그림을 그릴 수 있게 되었습니다.

하지만 여기서 문제가 생겼습니다.

• 기존의 학습법: AI 에게 "이 그림이 마음에 드니?"라고 물어볼 때, AI 가 그 답을 이해하려면 수학적으로 계산 가능한 (미분 가능한) 신호만 받아들일 수 있었습니다.
• 현실의 문제: 하지만 우리가 실제로 원하는 건 "사람이 보기에 예쁜가?", "문자가 정확하게 적혔는가?", "개 3 마리가 그려졌는가?" 같은 계산으로 바로 풀 수 없는 (비미분 가능한) 기준들입니다.
• 결과: 기존 방법으로는 AI 가 "빠르지만, 글자는 엉망이고, 개는 5 마리나 그려진" 엉뚱한 그림을 그릴 수밖에 없었습니다.

2. 해결책: TDM-R1 (새로운 학습 시스템)

저자들은 이 문제를 해결하기 위해 TDM-R1이라는 새로운 학습 방식을 개발했습니다. 이 방식은 두 가지 핵심 아이디어를 사용합니다.

비유 1: "확실한 길 (Deterministic Trajectory)"을 이용한 지도

기존의 AI 학습은 길을 찾을 때 "아마도 이쪽일 거야"라고 **추측 (확률적)**하며 나아가서, 중간에 어디를 지났는지 정확히 알 수 없었습니다. 그래서 "어디서 실수했는지"를 알기 어려웠습니다.

하지만 TDM-R1 은 **확실한 길 (Deterministic)**을 사용합니다.

• 비유: 마치 GPS 가 "이 길로 가세요"라고 정확한 경로를 보여주는 것처럼, AI 가 그림을 그리는 중간 단계 (노이즈에서 선명한 그림으로 변하는 과정) 를 정확히 추적할 수 있습니다.
• 효과: "어디서 실수했는지"를 중간중간 정확히 파악할 수 있게 되어, "글자가 틀렸네? 여기 고쳐야지!"라고 정확한 피드백을 줄 수 있게 됩니다.

비유 2: "가상의 코치 (Surrogate Reward)"와 "학생 (Generator)"의 팀워크

가장 어려운 점은 "사람이 예쁘다고 한 그림"을 AI 가 직접 이해하고 수정하는 것이 어렵다는 것입니다.

• 학생 (Generator): 그림을 그리는 AI (4 단계로 그림을 그리는 학생).
• 가상의 코치 (Surrogate Reward): 학생이 그린 그림을 보고 "이 부분은 좋아, 저 부분은 고쳐야 해"라고 수학적으로 계산 가능한 점수를 매겨주는 AI 코치.

TDM-R1 의 마법:

1. 학생이 그림을 그립니다.
2. 실제 인간이나 평가 도구가 "이 그림은 글자가 잘 써져서 점수 100 점!"이라고 점수를 줍니다 (이건 계산할 수 없는 점수).
3. 가상의 코치는 이 점수를 보고 "아, 이 학생이 이런 그림을 그렸을 때 점수가 높았구나. 다음엔 비슷한 경로를 따라가면 될 것 같아"라고 학습합니다.
4. 이제 학생은 가상의 코치가 알려준 "수학적인 길"을 따라 그림을 고쳐갑니다.

이렇게 학생과 코치가 서로를 가르치며 동시에 성장하는 방식 (GAN 과 유사한 구조) 을 통해, AI 는 복잡한 지시사항도 정확히 따르는 능력을 갖게 됩니다.

---

3. 놀라운 성과: "4 단계로 80 단계를 이기다"

이 새로운 방식을 적용한 결과, 놀라운 일이 일어났습니다.

• 기존의 80 단계 모델: 그림을 그리는 데 80 번의 수정을 거친 고가의 모델.
• TDM-R1 (4 단계 모델): 그림을 그리는 데 단 4 번의 수정만 거친 모델.

결과:

• GenEval (명령어 이해도 테스트): 기존 80 단계 모델이 63% 점수를 받았는데, TDM-R1 은 **92%**라는 압도적인 점수를 받았습니다. (비유하자면, 80 번 수정한 화가보다 4 번 수정한 화가가 지시사항을 더 정확히 이해한 셈입니다.)
• 실제 적용: "마스크에 'TDM-R1'이라고 적힌 우주인", "반짝이는 와이어로 만든 늑대" 같은 복잡한 요청도 정확하게 그려냈습니다.
• 확장성: 최신 대형 모델 (Z-Image) 에 적용해도 100 단계 모델보다 더 좋은 결과를 냈습니다.

4. 요약: 왜 이것이 중요한가?

이 논문은 **"빠른 AI 도구를 더 똑똑하게 만드는 방법"**을 제시했습니다.

• 과거: "빠르게 그리려면 정확한 지시를 따르기 힘들어."
• 지금 (TDM-R1): "빠르게 그리면서도, 사람이 원하는 복잡한 지시 (글자, 개수, 위치) 를 완벽하게 따를 수 있어!"

이 기술은 앞으로 AI 가 만드는 이미지와 동영상이 더 빠르고, 더 정확하며, 더 인간적인 취향을 반영할 수 있는 길을 열었습니다. 마치 초고속으로 그림을 그리는 화가가, 동시에 미술평론가의 눈까지 갖게 된 것과 같습니다.

기술 요약

TDM-R1: 비가환적 (Non-Differentiable) 보상을 활용한 Few-Step 확산 모델 강화

이 논문은 텍스트 - 이미지 생성 분야에서 Few-Step (소수 단계) 확산 모델의 성능을 향상시키기 위해, 기존 RL(강화학습) 방법론의 한계를 극복하고 비가환적 (Non-Differentiable) 보상을 효과적으로 활용하는 새로운 강화학습 패러다임인 TDM-R1을 제안합니다.

1. 문제 정의 (Problem)

• Few-Step 모델의 한계: 최근 확산 모델의 증류 (Distillation) 기술을 통해 4~10 단계 이내의 매우 빠른 이미지 생성이 가능해졌으나, 이러한 모델들은 정밀한 지시 따르기, 복잡한 텍스트 렌더링, 객체 배치 등에서 여전히 어려움을 겪습니다.
• 기존 RL 의 제약: Few-Step 모델에 적용된 기존 강화학습 방법들은 **보상 신호가 미분 가능 (Differentiable)**해야 한다는 전제에 의존합니다. 이는 역전파 (Back-propagation) 를 통해 보상을 모델에 전달해야 하기 때문입니다.
• 실제적 요구사항의 부재: 실제 세계의 중요한 보상 신호 (예: 인간의 이진 선호도, 객체 수 세기, OCR 을 통한 텍스트 정확도 등) 는 대부분 **비가환적 (Non-Differentiable)**입니다. 기존 방법론은 이러한 중요한 신호를 활용하지 못해 Few-Step 모델의 잠재력을 제한하고 있습니다.

2. 방법론 (Methodology)

TDM-R1 은 Trajectory Distribution Matching (TDM) 기반의 Few-Step 모델을 베이스로 하여, 학습 과정을 **대리 보상 학습 (Surrogate Reward Learning)**과 **생성기 최적화 (Generator Optimization)**로 분리하여 문제를 해결합니다.

• 결정론적 궤적 (Deterministic Trajectories) 의 활용:

  • TDM-R1 은 TDM 의 결정론적 샘플링 경로 (ODE 기반) 를 활용합니다.
  • 기존 확률적 (Stochastic) 경로는 중간 단계의 보상을 추정할 때 편향과 분산이 크지만, 결정론적 경로를 사용하면 중간 노이즈 단계 ($x_t$) 에서 최종 이미지 ($x_0$) 로 가는 경로가 명확하므로, 중간 단계에 대한 편향 없는 보상 추정이 가능합니다.
  • 이를 통해 전체 생성 경로의 각 단계에 대해 정확한 보상 신호를 할당할 수 있습니다.

• 대리 보상 (Surrogate Reward) 학습:

  • 비가환적 보상 ($r$) 을 직접 역전파할 수 없으므로, 확산 모델 파라미터 ($p_\phi$) 로 매개변수화된 미분 가능한 대리 보상을 학습합니다.
  • 그룹 기반 선호도 최적화 (Group-based Preference Optimization): Bradley-Terry (BT) 모델을 기반으로, 양 (+) 과 음 (-) 의 샘플 그룹 간의 선호 관계를 학습합니다. 각 그룹 내 샘플의 이점 (Advantage) 에 따라 가중치를 부여하여 세밀한 학습 신호를 제공합니다.
  • 동적 참조 모델 (Dynamic Reference Model): 고정된 참조 모델 대신, EMA(Exponential Moving Average) 를 적용한 동적 참조 모델을 사용하여 보상이 학습되는 과정에서 발생하는 과적합 (Overfitting) 및 불안정성을 방지합니다.

• Few-Step 생성기 학습:

  • 생성기 ($p_\theta$) 는 학습된 대리 보상을 최대화하면서, 사전 학습된 기본 확산 모델과의 역 KL 발산 (Reverse KL Divergence) 을 최소화하도록 학습됩니다.
  • 이는 생성된 샘플이 기본 분포와 너무 멀어지는 것을 방지하면서도, 비가환적 보상에 맞춰 성능을 향상시키는 균형을 맞춥니다.

3. 주요 기여 (Key Contributions)

1. 비가환적 보상을 활용한 첫 번째 Few-Step RL 패러다임: 인간의 선호도, 객체 수, 텍스트 정확도 등 미분 불가능한 다양한 보상 신호를 Few-Step 확산 모델에 성공적으로 통합했습니다.
2. 대리 보상 메커니즘: 비가환적 보상을 미분 가능한 형태로 변환하여 단계별 (Per-step) 학습을 가능하게 하는 새로운 아키텍처를 제안했습니다.
3. 결정론적 경로의 이점 증명: Few-Step 모델의 결정론적 특성을 활용하여 중간 단계 보상 추정의 분산을 줄이고 수렴 속도를 높였음을 이론적 및 실험적으로 입증했습니다.

4. 실험 결과 (Results)

TDM-R1 은 SD3.5-M 및 6B 파라미터의 Z-Image 모델 등 다양한 베이스 모델에서 광범위한 실험을 수행했습니다.

• GenEval (구성적 이미지 생성) 성능:
  • 4 NFE(No. of Function Evaluations) 기준 TDM-R1 은 **92%**의 점수를 기록했습니다.
  • 이는 80 NFE 를 사용하는 기존 베이스 모델 (63%) 과 상용 SOTA 모델인 GPT-4o(84%) 를 모두 압도하는 결과입니다.
  • 특히 객체 수 세기 (Counting), 위치 (Position), 속성 바인딩 (Attribute Binding) 등 복잡한 지시 따르기에서 뛰어난 성능을 보였습니다.
• 시각적 텍스트 렌더링 (Visual Text Rendering):
  • OCR 정확도 평가에서 기존 Few-Step 모델 대비 획기적인 개선을 보여주었습니다.
• 아웃-오브-도메인 (Out-of-Domain) 평가:
  • 학습에 사용되지 않은 이미지 품질 지표 (Aesthetic Score, ImageReward 등) 에서도 성능이 저하되지 않고 오히려 향상됨을 확인하여, 보상 해킹 (Reward Hacking) 이 발생하지 않음을 입증했습니다.
• Z-Image 모델 확장성:
  • 강력한 6B 파라미터 모델인 Z-Image 에 적용 시, 100 NFE 버전과 Few-Step Turbo 버전을 모두 4 NFE 로 압도하는 성능을 보였습니다.

5. 의의 및 결론 (Significance)

• 산업적 적용 가능성: Few-Step 모델은 이미 산업계에서 대규모 서비스의 표준으로 자리 잡고 있습니다. TDM-R1 은 이러한 모델에 인간 선호도나 복잡한 규칙 기반의 보상을 적용하여 성능을 극대화할 수 있는 길을 열었습니다.
• 효율성과 성능의 동시 달성: 기존에는 고품질 보상을 얻기 위해 많은 단계 (80~100 NFE) 가 필요했으나, TDM-R1 은 **4 단계 (4 NFE)**만으로 80 단계 모델보다 뛰어난 성능을 달성하여 생성 효율성을 극대화했습니다.
• RL 의 새로운 방향: 확산 모델 분야에서 비가환적 보상을 다루는 RL 방법론의 표준을 제시하며, LLM(대형 언어 모델) 에서의 성공적인 RLHF 경험을 이미지 생성 영역으로 성공적으로 확장했습니다.

요약하자면, TDM-R1은 Few-Step 확산 모델이 비가환적 보상 신호를 통해 인간 수준의 지시 따르기 능력을 갖추도록 하며, 적은 계산 비용으로 최첨단 (SOTA) 성능을 달성할 수 있게 하는 혁신적인 강화학습 프레임워크입니다.