기술 요약: MIRO (다중 보상 조건부 사전 학습)

1. 문제 제기

현재 텍스트-이미지 (T2I) 생성 파이프라인은 일반적으로 (1) 노이즈가 포함된 웹 데이터를 통한 사전 학습으로 이미지 구조를 학습하고, (2) 선별된 하위 집합에 대한 지도 미세 조정으로 분포를 좁히며, (3) 인간 피드백을 통한 강화 학습 (RLHF) 으로 사용자 선호도에 정렬하는 3 단계 패러다임을 따릅니다. 저자들은 이 패러다임이 다음과 같은 세 가지 치명적인 비효율성을 겪고 있다고 주장합니다.

• 데이터 폐기: 미세 조정 과정에서 '저품질' 데이터가 폐기되어 전체 데이터 분포에 존재하는 구조적 신호가 손실됩니다.
• 모드 붕괴 및 충실도 손실: RLHF 는 단일 스칼라 보상을 최적화하여 종종 분포를 붕괴시키고 (다양성 감소) 의미적 충실도를 해칩니다.
• 비유연성: 보상 간 트레이드오프는 학습 시점에 고정되어 추론 시 사용자가 선호도를 조정할 수 없습니다.
• 비효율성: 다단계 프로세스는 각 단계마다 상당한 계산 자원과 하이퍼파라미터 튜닝을 요구합니다.

본 논문은 사후 단계가 아닌 사전 학습 단계 자체에서 여러 보상을 정렬함으로써 전체 데이터 스펙트럼을 보존하면서도 제어 가능한 생성을 가능하게 할 수 있다고 제시합니다.

2. 방법론: MIRO

MIRO(다중 보상 조건부 사전 학습, MultI-Reward cOnditioned pretraining) 는 생성 모델을 보상 점수 벡터에 조건부로 설정함으로써 정렬을 사전 학습 단계에 직접 통합하는 프레임워크입니다.

2.1. 데이터셋 증강

데이터를 필터링하는 대신, MIRO 는 모든 학습 이미지 $(x, c)$에 $N$개의 보상 점수 벡터 $s = [s_1, s_2, \dots, s_N]$를 주석으로 달아줍니다. 저자들은 다음과 같은 7 가지 상호 보완적인 보상 모델을 활용합니다.

1. 미적 점수 (Aesthetic Score): 시각적 매력도.
2. HPSv2 및 PickScore: 인간 선호도 정렬.
3. ImageReward: 지시사항 준수 및 사용자 선호도.
4. VQAScore: 시각적 이해도.
5. JINA CLIP Score: 텍스트-이미지 정렬 (긴 캡션에 최적화됨).
6. SciScore: 과학적 정확성.

원시 점수는 정규화되어 다양한 보상 모델 간의 다양한 스케일을 처리하고 품질 수준 전반에 걸친 균형을 잡기 위해 $B$개의 이산 구간 (bins) 으로 분류됩니다.

2.2. 다중 보상 조건부 흐름 매칭

핵심 학습 목표는 표준 흐름 매칭 (Flow Matching) 수식을 수정합니다. 모델 $v_\theta$는 노이즈가 있는 샘플 $x_t$, 텍스트 프롬프트 $c$, 그리고 보상 벡터 $s$에 조건부로 속도장을 예측하도록 학습합니다:
$$\mathcal{L} = \mathbb{E}_{(x,c,s) \sim \tilde{D}, \epsilon \sim \mathcal{N}(0,I), t \sim U(0,1)} \left[ \| v_\theta(x_t, c, \hat{s}) - (\epsilon - x) \|_2^2 \right]$$
보상 수준 전체 스펙트럼에 대해 학습함으로써 모델은 데이터를 폐기하지 않고 특정 보상 수준과 특정 시각적 특성을 연관시키는 조건부 분포 $p(x|c, s)$를 학습합니다.

2.3. 보상 유도 추론

추론 시 MIRO 는 **다중 보상 분류기 없는 가이드 (Multi-Reward Classifier-Free Guidance, CFG)**를 통해 세밀한 제어를 가능하게 합니다. 가이드 벡터는 양의 보상 목표 $\hat{s}^+$(예: 최대 보상) 와 음의 목표 $\hat{s}^-$(예: 최소 보상) 를 대비하여 계산됩니다:
$$\hat{v}_\theta(x_t, c) = (1 + \omega)v_\theta(x_t, c, \hat{s}^+) - \omega v_\theta(x_t, c, \hat{s}^-)$$
이론적으로 이는 보상 기울어진 분포 $p_\omega(x|c) \propto p(x|c, s^+) \left( \frac{p(s^+|x, c)}{p(s^-|x, c)} \right)^\omega$에서 샘플링하는 것에 해당합니다. 이를 통해 사용자는 다음을 수행할 수 있습니다.

• 모든 보상을 동시에 최대화합니다.
• $\hat{s}^+$의 특정 구성 요소를 조정하여 (예: 텍스트 정렬보다 미적 요소 우선) 트레이드오프를 탐색합니다.
• 여러 목표를 균형 있게 맞추어 보상 해킹 (reward hacking) 을 방지합니다.

3. 주요 기여

1. MIRO 프레임워크: 보상 점수 벡터에 조건부로 설정하는 사전 학습 방법으로, 별도의 필터링, 미세 조정, 또는 사후 RLHF 단계를 제거합니다.
2. 효율성 및 성능: 0.36B 파라미터 MIRO 모델은 GenEval 및 사용자 선호도 벤치마크 (PickScore, ImageReward, HPSv2) 에서 최첨단 결과를 달성하여, 370 배 적은 학습 계산 자원을 요구하면서도 12B 파라미터 FLUX-dev 모델을 능가합니다.
3. 수렴 속도: 미적 및 선호도 지표에서 MIRO 는 기준 사전 학습보다 최대 19 배 빠르게 수렴합니다.
4. 추론 효율성: 샘플 기반 스케일링 하에서 MIRO 는 3 배 적은 추론 계산 자원으로 기준 모델과 동등하거나 더 나은 성능을 달성합니다.
5. 이론적 보장: 본 논문은 MIRO 가 전체 데이터 분포를 보존 (엔트로피 보존) 하고 RL 기반 정렬 (DDPO) 에 내재된 모드 붕괴 및 그래디언트 충돌을 방지한다는 증명을 제공합니다.

4. 실험 결과

• GenEval (구성적 벤치마크): MIRO 는 전체 점수 57(기준 52 대비) 을 달성했으며, Color Attribution(+31%) 및 Two Objects(+24%) 와 같은 까다로운 작업에서 큰 향상을 보였습니다. 합성 캡션과 결합될 때 MIRO 는 68에 도달하여 FLUX-dev(67) 를 능가합니다.
• 사용자 선호도: MIRO 는 ImageReward, HPSv2, PickScore 벤치마크에서 1 위를 차지했습니다.
• 학습 역학: 정성적 분석에 따르면 MIRO 는 50k 스텝까지 인식 가능하고 고품질의 이미지를 생성하는 반면, 기준 모델은 유사한 품질을 위해 200k~400k 스텝이 필요합니다.
• 보상 해킹: 단일 보상 모델 (예: 미적 점수만 최적화) 은 의미적 지표에서 심각한 저하를 보입니다. MIRO 의 다중 보상 조건 설정은 이를 방지하여 모든 지표에서 균형 잡힌 성능을 유지합니다.
• 테스트 시간 스케일링: MIRO 는 우수한 샘플 효율성을 보여줍니다. PickScore 와 같은 지표의 경우, MIRO 는 4 개의 샘플로 기준 모델이 128 개의 샘플로 달성하는 성능을 달성합니다 (32 배 효율성 향상).

5. 중요성 및 주장

본 논문은 MIRO 가 "보상을 최적화하는 것"에서 "보상 분포를 학습하는 것"으로의 패러다임 전환을 나타낸다고 주장합니다. 사전 학습 중 보상 점수를 제어 가능한 입력으로 취급함으로써 MIRO 는 다음과 같은 이점을 제공합니다.

• 내용과 품질의 분리: 모델은 재학습 없이 원하는 품질 수준에서 생성할 수 있습니다.
• "필터링 대 신호" 트레이드오프 제거: 선별된 데이터셋에서 종종 폐기되는 자연 이미지의 구조적 신호를 보존하면서 전체 데이터 스펙트럼을 활용합니다.
• 원칙적 제어 제공: 사용자는 추론 시 경쟁 목표 간 (예: 미적 요소 대 의미적 정렬) 의 트레이드오프를 명시적으로 조절할 수 있으며, 이는 표준 RLHF 파이프라인에서는 불가능한 기능입니다.
• 계산 비용 절감: 이 방법은 훨씬 더 작은 모델과 적은 계산 자원으로 최첨단 결과를 달성하여 T2I 모델 확장 시의 효율성 및 환경적 우려를 해결합니다.

저자들은 이 접근법이 보상 인식형 사전 학습에 대한 새로운 연구 방향을 열어주며, 현재 다단계 정렬 파이프라인에 대한 안정적이고 효율적이며 제어 가능한 대안을 제공한다고 결론지었습니다.