효율적인 Flow Matching 및 Diffusion 모델을 위한 멀티 패치 Global-to-Local Transformer 아키텍처

기존 Diffusion Transformer는 모든 층에서 동일한 해상도로 계산하여 비효율적이었으나, 이 논문은 초반에는 큰 그림을 보고 후반에 세부 묘사를 하는 계층 구조로 연산량을 절반으로 줄였습니다. 이를 통해 고해상도 이미지 생성을 더 적은 자원으로 빠르게 수행할 수 있는 길을 열었습니다.

Transformer 기반 Diffusion 모델(DiT)은 이미지를 작은 패치로 나누어 토큰화하는데, 모든 층에서 동일한 개수의 토큰을 처리하는 등방성(Isotropic) 구조를 가집니다. 이는 고해상도 이미지일수록 토큰 수가 급격히 늘어나 연산량이 시퀀스 길이의 제곱에 비례해 폭증하는 원인이 됩니다.

MPDiT는 사람이 그림을 그릴 때 먼저 큰 붓으로 구도를 잡고 나중에 작은 붓으로 묘사하는 원리를 아키텍처에 도입했습니다. 네트워크 초반부에는 패치 크기를 키워 토큰 수를 줄임으로써 전역적 맥락을 빠르게 파악하고, 후반부의 소수 블록에서만 패치를 작게 쪼개어 세부적인 질감을 다듬는 방식입니다.

이러한 계층적 접근은 연산의 대부분을 차지하는 초반부의 토큰 수를 1/4 이하로 줄여 전체 연산량을 50%까지 절감합니다. 또한, 단순한 선형 임베딩 대신 푸리에 변환 기반의 시간 임베딩을 도입해 확산 과정의 연속적인 흐름을 더 정확하게 학습하도록 설계하여 효율성과 품질을 동시에 잡았습니다.

전체 구조는 N개의 Transformer 블록으로 구성되며, 처음 N-k개 블록은 큰 패치(p=4 또는 8)를 입력받아 적은 수의 토큰으로 연산합니다. 이후 Upsample Block을 통해 토큰 시퀀스를 확장하고, 마지막 k개 블록에서 작은 패치(p=2) 수준의 고해상도 정보를 정교화합니다.

Upsample Block은 64개 토큰을 256개로 확장하기 위해 Linear Projection과 Pixel-unshuffle 연산을 수행합니다. [입력 토큰 시퀀스 → 선형 투영 및 재배열 → 4배 확장된 토큰 시퀀스] 과정을 거치며, 이때 원본 이미지의 고주파 정보를 보존하기 위해 초기 패치 임베딩 결과와의 Skip Connection을 결합하여 정보 손실을 방지합니다.

FNO Time Embedding은 스칼라 값인 timestep t를 1D 그리드와 결합하여 고차원 신호로 변환합니다. [t값 입력 → 1D 그리드 가산 → Spectral Convolution 연산 → 시간적 특징 추출] 순으로 계산되며, 이는 확산 모델의 역과정이 가지는 연속적인 궤적을 주파수 영역에서 더 정교하게 모사하도록 돕습니다.

ImageNet 256x256 실험에서 MPDiT-XL 모델은 59.3 GFLOPs만 사용하여 기존 DiT-XL(118.6 GFLOPs)의 절반 수준 연산량으로 더 낮은 FID(2.05)를 달성했습니다. 이는 연산 효율성뿐만 아니라 생성 품질 면에서도 계층적 구조가 유리함을 입증한 결과입니다.

학습 수렴 속도 측면에서 MPDiT는 기존 모델보다 약 11.5배 빠르게 수렴하는 것으로 나타났습니다. 특히 제안된 FNO 타임 임베딩과 멀티 토큰 클래스 임베딩을 적용했을 때 FID가 약 10포인트 가량 개선되는 효과를 확인했으며, 샘플링 처리량(throughput)은 DiT-XL/2 대비 2배 이상 향상되었습니다.

MPDiT는 DiT의 등방성 설계를 탈피하여 Global-to-Local 계층 구조를 제안합니다. 초반 블록은 큰 패치 사이즈를 사용하여 토큰 수를 256개에서 64개로 줄여 Self-Attention의 이차 복잡도 문제를 완화합니다. 마지막 4~6개의 refinement 블록만 고해상도 토큰을 처리함으로써 전체 연산 효율을 극대화합니다.

Upsample Block은 단순 보간이 아닌 학습 가능한 파라미터를 포함하며, LayerNorm과 GELU 활성화 함수를 거쳐 세부 정보를 복원합니다. 또한 원본 이미지 특징을 Skip Connection으로 주입하여 계층 간 해상도 불일치로 인한 아티팩트를 최소화합니다.

FNO(Fourier Neural Operator) 기반 시간 임베딩은 1D Spectral Convolution을 사용하여 주파수 영역에서 시간 정보를 처리합니다. 이는 확산 모델의 노이즈 제거 과정이 연속적인 함수 형태를 띤다는 점에 착안하여, 기존 MLP 방식보다 시간적 의존성을 더 정교하게 포착하며 ODE/SDE 궤적 학습에 최적화되어 있습니다.

ImageNet 데이터셋에서는 강력한 성능을 보였으나, SDv3나 Flux와 같은 초대형 텍스트-이미지 모델이나 Sora와 같은 비디오 생성 모델로의 확장은 여전히 탐구가 필요한 영역으로 남아 있습니다.