[NYU] 표현 오토인코더를 활용한 디퓨전 트랜스포머 (Diffusion Transformers with Representation Autoencoders)

배경 및 개요

이미지 생성 모델의 표준으로 자리 잡은 잠재 확산 모델(Latent Diffusion Model, LDM)과 디퓨전 트랜스포머(DiT)는 사전 훈련된 오토인코더를 사용하여 고차원의 픽셀 공간을 저차원의 잠재 공간(latent space)으로 매핑하고, 이 잠재 공간 내에서 디퓨전 프로세스를 수행합니다. 이 접근법은 계산 효율성과 생성 품질을 크게 향상시켰지만, 핵심 요소인 오토인코더 컴포넌트는 거의 발전 없이 기존의 VAE(Variational Autoencoder) 구조, 특히 Stable Diffusion의 SD-VAE에 머물러 있었습니다.

뉴욕 대학교(NYU) 연구팀은 이 VAE가 가진 근본적인 한계점을 지적합니다.

구식 아키텍처: 무거운 컨볼루션(Convolutional) 기반 백본을 사용하여 계산적으로 비효율적입니다.

낮은 차원의 잠재 공간: 정보를 과도하게 압축하여(예: 256x256 이미지를 32x32x4로) 원본 이미지의 세부 정보를 손실하고 정보 용량을 제한합니다.

빈약한 표현력: 순수하게 이미지 복원(reconstruction)만을 목표로 훈련되어, 생성된 잠재 벡터가 이미지의 의미론적(semantic) 구조를 제대로 담아내지 못하여 최종 생성물의 품질을 저해합니다.

이 논문은 이러한 한계를 극복하기 위해 표현 오토인코더(Representation Autoencoder, RAE)라는 새로운 패러다임을 제안합니다. RAE는 기존 VAE를 완전히 대체하는 개념으로, 다음과 같이 구성됩니다.

인코더: DINO, SigLIP, MAE와 같이 자기 지도 학습(self-supervised learning)이나 다중모드 학습을 통해 사전 훈련된 표현 인코더(representation encoder)를 그대로 사용(frozen)합니다. 이 인코더들은 이미지의 풍부한 의미론적 정보를 담고 있는 고품질의 표현(representation)을 추출하도록 이미 최적화되어 있습니다.

디코더: 고정된 표현 인코더가 추출한 고차원의 잠재 벡터를 다시 픽셀 공간으로 복원하는 경량의 디코더를 새로 훈련시킵니다.

RAE는 VAE와 달리 정보를 압축하지 않으며(예: 256x256 이미지 -> 256개 토큰), 의미적으로 매우 풍부하고 구조화된 고차원의 잠재 공간을 제공합니다. 그러나 이러한 고차원 잠재 공간에서 디퓨전 트랜스포머를 효과적으로 작동시키는 것은 새로운 도전 과제였습니다. 본 연구는 이 문제의 원인을 분석하고, 이론적/경험적으로 검증된 해결책을 제시하여 RAE 기반의 디퓨전 모델이 기존 VAE 기반 모델보다 훨씬 빠른 수렴 속도와 압도적인 생성 품질을 달성할 수 있음을 입증합니다.

모델별 테스트 결과

RAE와 새롭게 제안된 디퓨전 트랜스포머 아키텍처(DiTDH)를 결합한 모델은 ImageNet 데이터셋에서 기존의 모든 이미지 생성 모델을 뛰어넘는 SOTA(State-of-the-Art) 성능을 달성했습니다.

ImageNet 256x256 해상도 (Class-conditional generation):

가이던스 미적용 (without guidance): FID 1.51을 기록하여, 이전의 모든 공개 모델들을 큰 차이로 능가했습니다. 이는 REPA-E(1.70), VA-VAE(2.17) 등 기존 SOTA 모델보다 월등한 성능입니다.

가이던스 적용 (with AutoGuidance): FID 1.13이라는 경이적인 점수를 달성하며 새로운 기록을 세웠습니다.

ImageNet 512x512 해상도 (Class-conditional generation):

가이던스 적용 (with AutoGuidance): 256x256과 동일하게 FID 1.13을 달성했습니다. 이는 이전 SOTA 모델인 EDM-2(1.25), DDT(1.28) 등을 넘어서는 최고 성능입니다.

수렴 속도 및 효율성:

제안된 모델(DiT-XL on RAE)은 동일한 FID 점수에 도달하기까지 기존 VAE 기반의 SiT-XL 모델보다 47배, 표현 정렬(representation alignment) 기법을 사용한 REPA-XL 모델보다 16배 적은 훈련 연산량(GFLOPs)을 요구하여 압도적인 훈련 효율성을 보였습니다.

테스트 데이터셋

ImageNet (ILSVRC 2012): 1,000개의 클래스와 약 128만 개의 이미지로 구성된 대규모 이미지 데이터셋으로, 모델의 클래스 조건부 이미지 생성 성능을 평가하는 표준 벤치마크로 사용되었습니다.

평가 지표: 생성된 이미지의 품질과 다양성을 평가하기 위해 FID(Fréchet Inception Distance) 점수가 핵심 지표로 사용되었습니다. FID는 점수가 낮을수록 실제 이미지 분포와 유사함을 의미합니다.

연구 방법론 및 결론

본 연구는 RAE 기반의 디퓨전 트랜스포머를 성공시키기 위해 다음과 같은 체계적인 방법론을 제시합니다.

RAE의 우수성 입증: 먼저, 제안된 RAE가 기존 SD-VAE보다 우수함을 보였습니다. DINOv2 인코더와 ViT-XL 디코더를 사용한 RAE는 SD-VAE보다 훨씬 낮은 복원 FID(rFID 0.49 vs 0.62)와 훨씬 높은 표현력(선형 탐색 정확도 84.5% vs 8.0%)을 달성하면서도, 연산량(GFLOPS)은 1/3 수준에 불과함을 입증했습니다.

고차원 잠재 공간을 위한 디퓨전 트랜스포머 길들이기: 표준 DiT는 고차원의 RAE 잠재 공간에서 제대로 작동하지 않았습니다. 연구팀은 그 원인을 세 가지로 진단하고 해결책을 제시했습니다.

문제 1: 부적합한 트랜스포머 설계: DiT 모델의 내부 히든 차원(d)이 RAE 잠재 토큰의 차원(n)보다 작을 경우, 모델이 데이터 분포를 제대로 학습하지 못하는 ‘차원 병목 현상’이 발생함을 이론적(Theorem 1) 및 실험적으로 규명했습니다. 해결책: DiT의 히든 차원 d를 잠재 토큰의 차원 n 이상으로 설정해야 합니다.

문제 2: 부적합한 노이즈 스케줄: 기존 노이즈 스케줄은 채널 수가 적은 픽셀/VAE 공간에 맞춰져 있어, 채널 수가 많은(고차원) RAE 잠재 공간에 적용 시 노이즈의 영향력이 상대적으로 감소하는 문제가 있었습니다. 해결책: 잠재 공간의 전체 차원(토큰 수 × 채널 수)을 고려한 ‘차원 의존적 노이즈 스케줄 시프트’를 적용하여 노이즈 레벨을 보정합니다.

문제 3: 디코더의 강건성 부족: VAE 디코더는 본질적으로 노이즈가 낀 잠재 벡터로 훈련되지만, RAE 디코더는 깨끗한 잠재 벡터로만 훈련됩니다. 이로 인해 디퓨전 과정에서 발생하는 미세한 노이즈나 분포 이탈에 취약할 수 있습니다. 해결책: RAE 디코더 훈련 시, 잠재 벡터에 의도적으로 약간의 가우시안 노이즈를 추가하는 ‘노이즈 증강 디코딩’을 통해 디코더의 강건성을 높입니다.

DiTDH: 확장성을 위한 새로운 아키텍처: ‘차원 병목 현상’을 해결하기 위해 전체 DiT 모델의 너비를 키우는 것은 계산 비용이 막대합니다. 이를 해결하기 위해 DDT(Decoupled Diffusion Transformer)에서 영감을 받아 DiTDH (DiT with wide DDT Head) 아키텍처를 제안했습니다. 이는 기존 DiT 백본은 그대로 두고, 마지막 단에 얕지만 매우 넓은(shallow yet wide) 트랜스포머 헤드를 추가하는 방식입니다. 이 헤드가 최종적으로 고차원 잠재 공간의 디노이징을 담당함으로써, 전체 모델의 계산량을 크게 늘리지 않으면서도 차원 병목 문제를 효과적으로 해결합니다.

결론적으로, 이 논문은 의미론적으로 풍부한 표현 인코더가 고품질 이미지 생성에 매우 효과적일 수 있음을 증명하며, 기존의 VAE 기반 잠재 확산 모델의 패러다임을 전환했습니다. RAE라는 새로운 오토인코더 개념과, 고차원 잠재 공간에 최적화된 훈련 기법 및 DiTDH 아키텍처를 통해, 이미지 생성 분야에서 새로운 SOTA를 달성하고 향후 디퓨전 트랜스포머 연구의 새로운 표준을 제시했습니다.

시사점

표현 학습과 생성 모델링의 통합: DINO와 같은 자기 지도 학습 모델과 이미지 생성 모델은 오랫동안 별개의 영역으로 간주되었습니다. 이 연구는 두 분야를 성공적으로 결합하여, 표현 학습의 발전이 곧바로 생성 모델의 성능 향상으로 이어질 수 있는 길을 열었습니다.

차세대 잠재 확산 모델의 청사진: RAE는 기존 VAE의 단점을 완벽하게 보완하며, 더 가볍고, 더 빠르고, 더 높은 품질의 생성을 가능하게 합니다. 이는 향후 개발될 LDM, DiT 기반 모델들의 기본 오토인코더로 채택될 가능성이 매우 높습니다.

고차원 데이터 처리의 새로운 접근법: DiTDH 아키텍처는 고차원 토큰 시퀀스를 효율적으로 처리하는 방법을 제시함으로써, 이미지 생성을 넘어 비디오, 3D 등 다른 고차원 데이터 생성 모델에도 적용될 수 있는 잠재력을 가집니다.

생성 AI의 효율성 증대: 더 적은 연산량으로 더 빠르게 수렴하고 더 좋은 결과를 낸다는 것은 AI 모델 개발 및 배포 비용을 크게 절감할 수 있음을 의미하며, 이는 생성 AI 기술의 대중화와 발전을 더욱 가속화할 것입니다.

출처: https://www.arxiv.org/pdf/2510.11690

DiffusionTransformer #Autoencoder #GenerativeAI #SOTA #ComputerVision #DeepLearning

디퓨전트랜스포머 #오토인코더 #생성형AI #컴퓨터비전 #딥러닝 #표현학습