MPDiT: Multi-Patch Global-to-Local Transformer Architecture For Efficient Flow Matching and Diffusion Model

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 최근 확산 모델 (Diffusion Models) 과 흐름 매칭 (Flow Matching) 모델에서 컨볼루션 기반 UNet 보다 성능이 뛰어난 Diffusion Transformer (DiT) 아키텍처가 표준으로 자리 잡았습니다.
• 문제점: 기존 DiT 는 모든 블록에서 동일한 수의 패치 (patch) 토큰을 처리하는 등방성 (isotropic) 설계를 따릅니다. 이로 인해 초기 블록에서도 불필요하게 많은 토큰에 대한 어텐션 연산을 수행하게 되어, 훈련 및 추론 시 계산 비용 (GFLOPs) 이 매우 높고 메모리 사용량이 큽니다.
• 기존 대안의 한계:
  • 선형 어텐션 (Linear Attention): 계산 복잡도를 줄이지만, 장기 의존성 (long-range dependencies) 을 포착하는 데 한계가 있어 성능이 저하됩니다.
  • 마스크 모델링 (Masked Modeling): 토큰을 마스킹하여 효율을 높이려 하지만, 마스킹 비율이 높을 경우 글로벌 및 로컬 정보 모델링이 저하되어 성능이 크게 떨어집니다.
  • 상태 공간 모델 (Mamba 등): 긴 토큰 시퀀스에서는 유리하지만, 잠재 공간 (latent space) 확산 모델은 토큰 수가 적어 (1,000 개 미만) 큰 효율 향상을 얻지 못합니다.

2. 제안 방법론 (Methodology)

저자들은 MPDiT (Multi-Patch Diffusion Transformer) 를 제안하며, 이는 전체적 (Global) 에서 지역적 (Local) 으로 점진적으로 해상도를 높이는 계층적 아키텍처를 채택합니다.

A. 멀티-패치 글로벌-로컬 아키텍처 (Multi-Patch Global-to-Local Design)

• 초기 블록 (Global Context): 모델의 앞부분 ($N-k$ 개 블록) 에서는 큰 패치 크기 (예: $p=4$) 를 사용하여 입력 토큰 수를 대폭 줄입니다 (예: 256 개 $\to$ 64 개). 이를 통해 적은 계산량으로 이미지의 전체적인 구조와 글로벌 컨텍스트를 효율적으로 학습합니다.
• 업샘플링 모듈 (Upsample Block): 큰 패치 토큰을 작은 패치 토큰으로 변환하는 모듈을 도입합니다. 기존 DiT 의 패치 임베딩 (Patch Embedding) 에서 생성된 원본 토큰 (스킵 연결) 과 결합하여 세부 정보 (Fine-grained details) 를 보존합니다.
• 후기 블록 (Local Refinement): 업샘플링된 작은 패치 토큰 (예: $p=2$) 만을 모델의 마지막 $k$ 개 블록 (예: 4~6 개) 에서 처리하여 고해상도 세부 묘사를 완성합니다.
• 효율성: 대부분의 블록이 적은 수의 토큰으로 작동하므로, 어텐션의 2 차 복잡도 특성상 GFLOPs 를 최대 50% 까지 절감하면서도 성능을 유지합니다.

B. 조건부 임베딩 개선 (Improved Conditioning)

• 시간 임베딩 (Time Embedding): 기존의 단순한 사인파 (Sinusoidal) + MLP 방식 대신, 푸리에 신경 연산자 (Fourier Neural Operator, FNO) 기반 임베딩을 도입했습니다. 이는 확산 과정의 연속적인 흐름 (SDE/ODE) 을 더 잘 포착하여 시간적 의존성을 풍부하게 표현하고 훈련 수렴 속도를 높입니다.
• 클래스 임베딩 (Class Embedding): 단일 토큰 대신 여러 개의 학습 가능한 토큰 (Multi-token) 을 사용하여 클래스 정보를 표현합니다. 이는 클래스 의미론을 더 풍부하게 인코딩하여 훈련 수렴을 가속화하고 생성 품질을 향상시킵니다.

3. 주요 기여 (Key Contributions)

1. 전체적 - 지역적 (Global-to-Local) 계층적 아키텍처: 어텐션 레이어 내부가 아닌 네트워크 구조 수준에서 글로벌/로컬 처리를 분리하여, 계산 효율성과 생성 품질 간의 균형을 최적화했습니다.
2. 임베딩 모듈 재설계: FNO 기반 시간 임베딩과 멀티-토큰 클래스 임베딩을 도입하여 기존 DiT 대비 약 10 점 이상의 FID 개선 효과를 입증했습니다.
3. 광범위한 실험 및 검증: ImageNet 256 및 512 데이터셋에서 기존 DiT, SiT, Mamba 기반 모델 등 다양한 베이스라인과 비교하여, 계산 비용 절감 (GFLOPs 50% 감소) 과 동시에 더 빠른 수렴 및 우수한 생성 품질을 달성함을 증명했습니다.

4. 실험 결과 (Results)

• ImageNet 256:
  • MPDiT-XL은 240 에포크 훈련 후 FID 7.36 (cfg 사용 시 2.05) 을 달성했습니다.
  • 기존 SiT-XL 이 1,400 에포크에 9.35 FID 를 달성한 것과 비교해 약 6 배 빠른 수렴 속도를 보였습니다.
  • GFLOPs는 118.66 (SiT) 에서 59.30으로 약 50% 감소했습니다.
• ImageNet 512:
  • 120 에포크 훈련 후 FID 2.47을 기록하여 모든 베이스라인 (DiT, SiT, DiM 등) 을 능가했습니다.
  • 총 훈련 계산량 (Total Training Compute) 이 기존 모델 대비 약 8.7% 수준으로 줄어들어, 단일 A100 노드에서도 대용량 배치 훈련이 가능해졌습니다.
• 추론 속도: 동일한 GPU 환경에서 DiT 대비 2 배 이상 빠른 샘플링 속도를 달성했습니다.

5. 의의 및 결론 (Significance)

이 논문은 확산 모델의 효율성 문제를 해결하기 위해 토큰 처리 전략의 근본적인 재설계를 제시했습니다.

• 효율성과 성능의 동시 달성: 기존에는 효율성을 높이면 성능이 떨어지는 트레이드오프 관계가 있었으나, MPDiT 는 계층적 패치 설계를 통해 이를 극복했습니다.
• 확장성: 제안된 아키텍처는 고해상도 이미지 생성뿐만 아니라, 향후 텍스트 - 이미지, 텍스트 - 비디오 모델 등 대규모 생성 모델의 훈련 비용과 메모리 요구 사항을 획기적으로 낮출 수 있는 잠재력을 가집니다.
• 실용성: 제한된 하드웨어 자원 (단일 노드 GPU) 으로도 대규모 모델을 효율적으로 훈련하고 배포할 수 있게 함으로써, 생성형 AI 의 대중화와 실용화를 가속화할 것으로 기대됩니다.

요약하자면, MPDiT는 "처음에는 큰 패치로 전체를 보고, 나중에는 작은 패치로 디테일을 다듬는다"는 직관적인 아이디어를 아키텍처 수준에 적용하여, 계산 비용을 절반으로 줄이면서도 더 빠르고高质量인 이미지를 생성하는 획기적인 확산 모델 설계입니다.