17일 전
MDTv2: 마스킹된 확산 트랜스포머는 강력한 이미지 합성기이다.
Shanghua Gao, Pan Zhou, Ming-Ming Cheng, Shuicheng Yan

초록
이미지 합성 분야에서 성공을 거두고 있음에도 불구하고, 확산 확률 모델(DPMs)은 이미지 내 객체 부분 간의 관계를 학습하는 데 있어 맥락적 추론 능력이 부족함을 관찰할 수 있으며, 이는 학습 속도 저하로 이어진다. 이러한 문제를 해결하기 위해, 우리는 이미지 내 객체의 의미적 부분 간 맥락적 관계 학습 능력을 명시적으로 향상시키기 위해 마스크 잠재 모델링 기법을 도입한 마스크 확산 트랜스포머(Masked Diffusion Transformer, MDT)를 제안한다. 훈련 과정에서 MDT는 잠재 공간에서 특정 토큰을 마스크 처리하고, 이를 통해 비대칭 확산 트랜스포머를 설계하여 마스크 처리되지 않은 토큰들로부터 마스크된 토큰을 예측하면서도 확산 생성 과정을 유지한다. 이로 인해 MDT는 부분적인 맥락 입력으로부터 이미지의 전체 정보를 재구성할 수 있으며, 이는 이미지 토큰 간의 상관 관계를 효과적으로 학습할 수 있도록 한다. 또한, 더 효율적인 매크로 네트워크 구조와 훈련 전략을 도입하여 MDTv2라는 개선된 버전을 제안한다. 실험 결과, MDTv2는 이미지 합성 성능에서 뛰어난 성과를 보였으며, ImageNet 데이터셋에서 새로운 SOTA(FID: 1.58)를 달성하였고, 이전 SOTA인 DiT보다 10배 이상 빠른 학습 속도를 기록하였다. 소스 코드는 https://github.com/sail-sg/MDT 에 공개되어 있다.