CMT: Mid-Training for Efficient Learning of Consistency, Mean Flow, and Flow Map Models

이 논문은 사전 학습된 확산 모델과 최종 흐름 지도 모델 학습 사이에 경량 중간 단계를 도입하여 학습 안정성과 효율성을 극대화하고, 기존 방법 대비 훨씬 적은 데이터와 계산 자원으로 최첨단 생성 성능을 달성하는 'CMT(Consistency Mid-Training)' 프레임워크를 제안합니다.

핵심

🎨 배경: 왜 AI 그림 그리기는 느릴까요?

기존의 '확산 모델 (Diffusion Model)'은 그림을 그릴 때, 완전히 검은 화면 (소음) 에서 시작해 한 방울 한 방울씩 물감을 섞어가며 그림을 완성합니다.

• 비유: 마치 안개 낀 산을 한 걸음 한 걸음 천천히 걸어서 정상에 도달하는 것과 같습니다.
• 문제: 이 과정이 너무 오래 걸립니다. 그림 하나를 그리기 위해 100 번 이상의 단계 (걸음) 가 필요할 수 있어, 시간이 많이 들고 비용이 비쌉니다.

🚀 해결책: "직행 버스" (Flow Map)

연구자들은 "왜 한 걸음 한 걸음 걸어가나요? 처음부터 끝까지 직행 버스를 타면 안 될까요?"라고 생각했습니다.

• Flow Map (흐름 지도): 소음 상태에서 바로 완성된 그림으로 '점프'하는 기술을 말합니다.
• 기존 방식의 문제: 이 직행 버스를 가르치는 것은 매우 어렵습니다.
  • 무작위 출발: 아무것도 모르는 상태에서 시작하면 (랜덤 초기화), 버스가 길을 잃고 엉뚱한 곳에 도착합니다.
  • 기존 지도자 (Pre-trained Diffusion) 의 한계: 기존에 잘 훈련된 '걸음걸이 지도자'를 빌려와도, 그는 '한 걸음'만 가르칠 줄 알지, '직행'을 가르치는 법을 모릅니다. 그래서 학생 (새 모델) 이 지도자를 따라 하다가 혼란을 겪고 훈련이 불안정해집니다.

✨ CMT 의 등장: "중간 훈련 (Mid-Training)"

이 논문은 **'중간 훈련 (Mid-Training)'**이라는 새로운 단계를 제안합니다. 마치 **유아기 (Pre-training) 와 성인이 되기 전 (Post-training) 사이에 있는 '청소년기'**를 거치는 것과 같습니다.

🚌 CMT 의 작동 원리 (비유)

1. 유아기 (기존 훈련): AI 는 소음에서 그림을 그리는 법을 배웁니다. (걸음걸이 학습)
2. 청소년기 (CMT, 중간 훈련):
   • 여기가 핵심입니다! AI 는 **완벽한 가이드 (선생님)**가 그린 '직행 버스 경로'를 따라가며 연습합니다.
   • 상황: 선생님이 "소음 (A) 에서 출발해서 16 걸음 만에 그림 (B) 에 도달하는 경로"를 보여줍니다.
   • 학습: AI 는 이 경로상의 중간 지점을 보고, "아, 여기서 바로 그림 (B) 으로 점프하면 되구나!"라고 배웁니다.
   • 효과: AI 는 이제 '걸음걸이'만 아는 게 아니라, '어디서 시작해서 어디로 점프할지'를 정확히 아는 상태가 됩니다.
3. 성인기 (최종 훈련): 이제 AI 는 이 '청소년기'에 배운 지식을 바탕으로, 최종적인 '직행 버스'를 완성합니다.

🌟 CMT 가 가져온 놀라운 변화

이 '중간 훈련'을 거친 결과, 다음과 같은 기적이 일어났습니다.

• ⏱️ 속도: 그림을 그리는 시간이 최대 98% 단축되었습니다. (예: 100 시간 걸리던 것을 2 시간으로 줄임)
• 💰 비용: 필요한 데이터 양과 GPU(컴퓨터) 사용 시간이 획기적으로 줄었습니다.
• 🎯 품질: 2 단계 (2 걸음) 만에 그리는 그림의 품질이 기존 최고 수준 (SOTA) 을 넘어섰습니다.
  • 예: CIFAR-10 데이터셋에서 1.97, ImageNet 512x512 에서 1.84 의 매우 낮은 FID 점수 (낮을수록 좋음).
• 🛡️ 안정성: 훈련이 자주 망가지거나 (수렴 실패) 불안정해지는 일이 거의 없어졌습니다.

💡 핵심 요약

이 논문은 **"AI 가 그림을 그릴 때, 무작정 처음부터 가르치지 말고, '중간 단계'에서 완벽한 경로를 먼저 익히게 하라"**고 말합니다.

• 기존 방식: "걸음걸이부터 배워라" → "그럼 이제 점프해봐!" (혼란스럽고 느림)
• CMT 방식: "걸음걸이 배웠지? 이제 '걸음걸이'를 이용해 '직행 경로'를 먼저 연습해보자" → "자, 이제 점프해봐!" (매우 빠르고 정확함)

이 기술은 AI 가 고화질 이미지를 훨씬 더 저렴하고 빠르게 만들어낼 수 있는 길을 열었습니다. 마치 비행기를 타고 여행하듯, AI 가 그림을 그리는 시대가 온 것입니다.

기술 요약

1. 문제 정의 (Problem)

확산 모델 (Diffusion Models) 은 이미지 생성 분야에서 뛰어난 성능을 보이지만, 추론 시 확률 흐름 상미분 방정식 (PF-ODE) 을 수치적으로 풀어야 하므로 많은 단계 (NFE, Number of Function Evaluations) 가 필요하여 계산 비용이 높고 지연 시간이 길다는 한계가 있습니다. 이를 해결하기 위해 Flow Map Models(예: Consistency Models, Mean Flow) 가 제안되었습니다. 이 모델들은 ODE 의 해를 직접 학습하여 적은 단계 (few-step) 로 고품질 이미지를 생성할 수 있게 합니다.

그러나 Flow Map 모델의 학습에는 다음과 같은 심각한 문제가 존재합니다:

• 불안정성 (Instability): 학습 목표가 Stop-gradient 를 사용하는 가상의 타겟 (pseudo-target) 에 의존하기 때문에 최적화 신호가 불안정하고 수렴이 느립니다.
• 하이퍼파라미터 민감도: 학습 시간 가중치 (time weighting) 나 샘플링 스케줄 등 많은 휴리스틱 (heuristic) 에 의존해야 합니다.
• 높은 비용: 기존 방법론들은 사전 학습된 확산 모델 (Pre-trained Diffusion Model) 로 초기화하더라도, 확산 모델이 포착하는 '무한소 이동 (infinitesimal movement)'과 Flow Map 이 학습해야 하는 '장거리 점프 (long-jump)' 사이의 불일치로 인해 여전히 학습이 어렵고 비용이 많이 듭니다.

2. 제안 방법 (Methodology: CMT)

저자들은 **Consistency Mid-Training (CMT)**을 제안합니다. 이는 사전 학습 (Pre-training) 과 최종 Flow Map 학습 (Post-training) 사이에 삽입되는 경량화된 중간 학습 단계입니다.

핵심 아이디어

CMT 는 사전 학습된 모델 (Teacher) 이 생성한 ODE 궤적 (Trajectory) 을 활용하여, 모델이 궤적 상의 임의의 점 (noisy state) 을 해당 궤적의 깨끗한 끝점 (clean sample) 으로 직접 매핑하도록 학습시킵니다.

학습 파이프라인

1. Stage 1: Pre-training (Teacher Sampler)
   • 기존 확산 모델 (EDM 등) 이나 작은 Flow Map 모델을 사전 학습하여 결정론적 ODE 솔버 (Teacher) 를 준비합니다.
2. Stage 2: Mid-training (CMT)
   • 목표: Teacher 가 생성한 궤적 $\{ \hat{x}_{t_i} \}$ 상의 임의의 점 $\hat{x}_{t_i}$를 해당 궤적의 시작점인 깨끗한 데이터 $\hat{x}_{t_0}$ (또는 Flow Map 의 경우 $\hat{x}_{t_i} \to \hat{x}_{t_j}$) 로 매핑하는 모델을 학습합니다.
   • 손실 함수 (Loss):
     • Consistency Models (CM) 의 경우: $L_{CMT-CM} = \sum d(f_\theta(\hat{x}_{t_i}, t_i), \hat{x}_{t_0})$
     • Mean Flow (MF) 의 경우: 평균 드리프트를 학습하도록 $L_{CMT-MF} = \sum \| h_\theta(\hat{x}_{t_i}, t_i, t_j) - \frac{\hat{x}_{t_i} - \hat{x}_{t_j}}{t_i - t_j} \|^2$
   • 특징: Stop-gradient 가 필요 없으며, 고정된 명확한 회귀 타겟 (Explicit Regression Target) 을 사용합니다. 이는 학습을 매우 안정적으로 만듭니다.
3. Stage 3: Post-training (Flow Map Learning)
   • CMT 로 학습된 가중치를 초기값으로 사용하여 최종 Few-step Flow Map 모델 (ECT, ECD, MF 등) 을 학습합니다.
   • CMT 초기화 덕분에 Post-training 은 더 빠르게 수렴하며, 복잡한 휴리스틱 (시간 가중치 조정 등) 없이도 높은 성능을 달성합니다.

3. 주요 기여 (Key Contributions)

1. 새로운 학습 패러다임 (Mid-training): Few-step Flow Map 모델 학습을 위한 체계적인 '중간 학습' 개념을 처음 도입했습니다. 이는 LLM 의 Mid-training 에서 영감을 받았으나, 생성 모델의 ODE 궤적 특성에 맞춰 재해석되었습니다.
2. 이론적 근거: CMT 초기화가 Oracle Loss 와 실제 Loss 간의 **기울기 편차 (Gradient Bias)**를 크게 줄인다는 것을 이론적으로 증명했습니다. 이는 무작위 초기화나 단순 확산 모델 초기화보다 훨씬 강력한 초기 상태를 제공함을 의미합니다.
3. 범용성: Consistency Models (CM) 와 Mean Flow (MF) 모두에 적용 가능하며, 픽셀 공간 (Pixel-space) 과 잠재 공간 (Latent-space) 고해상도 모델 모두에서 유효함을 입증했습니다.
4. 간소화된 학습 프로세스: Stop-gradient, 복잡한 시간 샘플링, 수동 가중치 조정 등 기존 Flow Map 학습의 복잡한 공학적 트릭 (Tricks) 을 제거하여 학습을 단순화하고 안정화했습니다.

4. 실험 결과 (Results)

CMT 는 다양한 벤치마크에서 SOTA(State-of-the-Art) 성능을 달성하면서 학습 비용 (데이터 양 및 GPU 시간) 을 획기적으로 줄였습니다.

• 성능 (FID):
  • CIFAR-10: 2-step FID 1.97 (기존 CM 들보다 우수).
  • ImageNet 64x64: 2-step FID 1.32.
  • ImageNet 512x512: 2-step FID 1.84 (기존 ECD 의 3.38 대비 압도적).
  • ImageNet 256x256: 1-step FID 3.34 (Scratch MF 의 3.43 보다 우수).
  • MS-COCO (T2I): 2-step FID 5.01.
• 효율성 (비용 절감):
  • 학습 데이터: 기존 방법 대비 최대 98% 감소 (예: ImageNet 64x64 에서 sCT 대비 98% 절감).
  • GPU 시간: ImageNet 512x512 기준, 기존 ECD 학습 시간 대비 91.4% 단축 (400 H100 GPU 시간 vs 4643 시간).
  • 수렴 속도: ImageNet 256x256 에서 MF 를 처음부터 학습하는 것보다 약 50% 빠른 수렴.

5. 의의 및 결론 (Significance)

이 논문은 CMT를 통해 Few-step Flow Map 모델의 학습 장벽을 효과적으로 해소했습니다.

• 안정성 확보: 불안정했던 Flow Map 학습을 궤적 일관성 (Trajectory-consistent) 을 가진 초기화를 통해 안정화시켰습니다.
• 비용 효율성: 고품질 Few-step 생성 모델을 학습하는 데 필요한 막대한 컴퓨팅 자산을 획기적으로 줄여, 실제 적용 가능성을 높였습니다.
• 일반성: 확산 모델 초기화의 한계를 극복하고, 다양한 아키텍처와 모델 유형 (CM, MF) 에 적용 가능한 범용적인 프레임워크를 제시했습니다.

결론적으로, CMT 는 Few-step 생성 모델의 학습을 더 효율적이고, 안정적이며, 접근하기 쉽게 만드는 중요한 이정표가 되는 연구입니다.