There is No VAE: End-to-End Pixel-Space Generative Modeling via Self-Supervised Pre-training

이 논문은 사전 훈련된 VAE 없이 픽셀 공간에서 직접 학습하는 새로운 2 단계 프레임워크를 제안하여 이미지넷에서 VAE 기반 모델 및 DiT 를 능가하는 최첨단 생성 품질과 훈련 효율성을 달성했습니다.

핵심

"VAE 는 없다: 픽셀 공간에서 직접 그림을 그리는 새로운 방법"

이 논문은 인공지능이 그림을 그리는 방식에 대한 혁신적인 변화를 제안합니다. 기존의 방식이 겪던 문제를 해결하고, 더 빠르고 더 잘 그리는 새로운 방법을 소개합니다.

🎨 기존 방식의 문제점: "번거로운 번역가"

기존의 AI 그림 그리기 기술 (확산 모델) 은 보통 두 단계로 나뉩니다.

1. 번역 (VAE): 먼저 고해상도의 원본 그림을 압축해서 '요약본 (잠재 공간, Latent Space)'으로 바꿉니다.
2. 그림 그리기: AI 는 이 요약본을 보고 상상력을 발휘해 새로운 그림을 그립니다.
3. 해석: 다시 그 요약본을 원본 그림으로 되돌려서 보여줍니다.

비유하자면:
이것은 고급 요리사가 요리할 때, 재료를 먼저 '분말'로 갈아서 섞고, 다시 '수프'로 끓인 뒤, 마지막에 다시 '원재료'로 되돌리는 과정과 같습니다.

• 단점: 이 '분말'로 바꾸는 과정 (VAE) 자체가 매우 어렵고, 원래의 맛 (세부적인 디테일) 이 조금씩 사라집니다. 또한, AI 가 그리는 동안 이 '번역가' (VAE) 가 실수를 하면 최종 결과물도 망가집니다.

💡 이 논문의 해결책: "직접 그리는 천재 화가"

이 논문은 **"번역가 (VAE) 는 필요 없다!"**라고 선언합니다. 대신 AI 가 원래의 고해상도 그림 (픽셀) 을 직접 보고, 직접 배우고, 직접 그리는 방법을 제시합니다.

🚀 새로운 훈련 방법: "두 단계 교육 시스템"

이 새로운 방법 (EPG) 은 두 단계로 이루어진 특별한 교육 과정을 거칩니다.

1 단계: "눈을 뜨는 훈련" (프리트레이닝)

AI 에게는 아직 그림을 그리는 능력 (디코더) 이 없습니다. 먼저 **그림의 '의미'를 파악하는 능력 (인코더)**만 집중적으로 훈련시킵니다.

• 비유: 마치 미술관 큐레이터가 됩니다. 흐릿하게 흐려진 그림 (노이즈가 섞인 이미지) 을 보더라도, 이것이 '고양이'인지 '자동차'인지 그 본질적인 의미를 알아내는 훈련입니다.
• 핵심: 단순히 흐릿한 그림을 보는 게 아니라, 같은 그림이 시간이 지남에 따라 어떻게 깨끗해지는지 (확산 경로를 따라) 그 흐름을 이해하도록 가르칩니다.

2 단계: "손을 움직이는 훈련" (파인튜닝)

이제 의미 파악 능력을 갖춘 AI 에게 **그림을 그리는 도구 (디코더)**를 붙여줍니다.

• 비유: 이제 큐레이터가 화가가 됩니다. 그가 파악한 '의미'를 바탕으로, 처음부터 끝까지 원본 그림의 디테일 그대로를 픽셀 단위로 직접 그려냅니다.
• 결과: 번역 과정이 없으니, 원래의 맛 (디테일) 이 그대로 살아납니다.

🏆 놀라운 성과: "왜 이것이 중요한가?"

이 새로운 방식은 기존 방식보다 훨씬 훌륭합니다.

1. 더 잘 그립니다 (성능):

   • ImageNet(유명한 그림 데이터셋) 에서 FID 점수 1.58이라는 역대 최고 기록을 세웠습니다. (숫자가 낮을수록 더 잘 그린다는 뜻입니다.)
   • 기존에 '번역가 (VAE)'를 쓰던 최고의 모델들보다 훨씬 더 선명하고 자연스러운 그림을 그립니다.

2. 더 빠르고 저렴합니다 (효율):

   • 기존 방식보다 학습에 필요한 계산 비용이 약 30% 적게 듭니다.
   • 그림을 한 장 뽑아내는 속도도 매우 빠릅니다.

3. 한 번에 그리는 consistency 모델의 성공:

   • 보통 AI 는 그림을 그릴 때 여러 번 반복해서 수정합니다 (예: 75 번). 하지만 이 모델은 **한 번의 시도 (One-step)**로도 매우 훌륭한 그림을 그릴 수 있습니다.
   • 이는 고해상도 그림을 VAE 없이 직접 학습시킨 세계 최초의 성과입니다.

🌟 요약: "왜 이 논문이 획기적인가?"

이 논문은 **"복잡한 번역 과정 (VAE) 을 버리고, AI 가 원본을 직접 이해하고 그리는 직관적인 방법"**을 증명했습니다.

• 기존: "재료를 분말로 → 섞어서 → 다시 수프로 → 원재료로" (비효율적이고 맛 떨어짐)
• 새로운 방법 (EPG): "재료를 바로 보고 → 바로 요리해서 → 바로 상에 올림" (맛있고 빠름)

이 기술은 AI 가 그림을 그릴 때 더 이상 '번역가'에 의존하지 않아도 된다는 것을 보여주며, 앞으로 더 고화질이고 효율적인 AI 예술의 시대를 열 것이라고 기대됩니다.

기술 요약

1. 문제 정의 (Problem)

최근 생성 모델 (Diffusion, Consistency Models) 의 고해상도 이미지 생성 성능은 대부분 사전 학습된 VAE(Variational Autoencoder) 의 잠재 공간 (Latent Space) 에서 훈련되는 방식에 의존하고 있습니다. 그러나 이 VAE 기반 접근법은 다음과 같은 근본적인 한계를 가집니다.

• VAE 훈련의 어려움: 압축과 고충실도 재구성의 균형을 맞추는 것이 매우 어렵습니다.
• 성능 병목 현상: 훈련된 VAE 는 훈련 분포에서 벗어난 잠재 변수 (Latents) 에 대해 불완전한 재구성을 수행하며, 생성 모델의 적응 능력을 VAE 의 고정된 용량으로 제한합니다.
• 픽셀 공간 (Pixel-space) 훈련의 비효율성: VAE 없이 원본 픽셀 (Raw Pixels) 에서 직접 훈련하는 시도는 많았으나, 높은 계산 비용과 느린 수렴 속도로 인해 잠재 공간 기반 방법보다 성능과 효율성이 현저히 떨어졌습니다.

이 논문은 VAE 없이 픽셀 공간에서 직접 훈련하되, 잠재 공간 기반 방법과 견줄 만한 성능과 효율성을 달성하는 것을 목표로 합니다.

2. 방법론 (Methodology)

저자들은 생성 모델을 **자기지도 학습 (Self-Supervised Learning, SSL)**의 관점에서 재해석하여, 인코더 (Encoder) 와 디코더 (Decoder) 의 역할을 분리하고 **2 단계 훈련 프레임워크 (EPG)**를 제안합니다.

2.1 핵심 아이디어: 인코더 - 디코더 역할 분리

• 인코더: 노이즈가 포함된 입력에서 고수준의 시각적 의미 (Visual Semantics) 를 학습합니다.
• 디코더: 인코더의 표현을 기반으로 저수준의 픽셀을 생성합니다.
• 이 구조는 SSL 에서 인코더가 일반 시각 학습을, 디코더가 태스크 특화 예측 헤드를 수행하는 방식과 유사합니다.

2.2 1 단계: 자기지도 사전 학습 (Pre-training)

전통적인 SSL 은 깨끗한 이미지의 증강 (Augmentation) 을 사용하지만, 이 논문은 확산 (Diffusion) 과정의 노이즈를 활용합니다.

• 목표: 깨끗한 이미지에서 의미 있는 시각적 특징을 추출하되, 이를 **동일한 결정론적 샘플링 궤적 (Deterministic Sampling Trajectory, ODE)**상의 점들과 정렬 (Align) 합니다.
• 손실 함수:
  1. 대비 손실 (Contrastive Loss): 데이터 증강된 뷰 쌍을 양의 쌍 (Positive pair) 으로, 배치 내 다른 샘플을 음의 쌍 (Negative pair) 으로 사용하여 의미 학습을 촉진합니다.
  2. 표현 일관성 손실 (Representation Consistency Loss): 동일한 ODE 궤적 상의 인접한 시간 단계 (예: $x_{t_n}$과 $x_{t_{n-1}}$) 의 특징을 정렬합니다. 이는 노이즈가 큰 이미지에서도 의미 있는 표현을 학습하도록 돕습니다.
• 안정화 기법: 초기 훈련 불안정성을 해결하기 위해 온도 파라미터 ($\tau$) 를 선형적으로 조정하는 스케줄을 도입하여, 큰 노이즈 단계에서는 느슨한 정렬을, 작은 노이즈 단계에서는 강한 정렬을 유도합니다.

2.3 2 단계: 엔드 - 투 - 엔드 파인튜닝 (Fine-tuning)

• 사전 학습된 인코더와 무작위 초기화된 디코더를 결합합니다.
• 확산 모델 (Diffusion Model): 기존 확산 손실 (Denoising objective) 로 엔드 - 투 - 엔드 파인튜닝합니다.
• 일관성 모델 (Consistency Model): ODE 궤적 상의 일관성 손실과 함께, 모델 출력과 깨끗한 원본 이미지 ($x_0$) 간의 보조 손실 (Auxiliary Loss) 을 추가하여 학습을 안정화하고 수렴 속도를 높입니다.
• 아키텍처: Vision Transformer (ViT) 를 백본으로 사용하며, 이미지 해상도가 증가함에 따라 패치 크기 (Patch Size) 를 비례하여 조정 (예: 256x256 은 16x16, 512x512 는 32x32) 하여 입력 토큰 길이를 고정함으로써 계산 효율성을 극대화합니다.

3. 주요 기여 (Key Contributions)

1. 새로운 훈련 프레임워크 (EPG): 고해상도 픽셀 공간 생성을 위한 효율적이고 확장 가능한 2 단계 훈련 프레임워크를 제안했습니다. 이는 확산 모델 훈련을 이미지 분류와 유사한 자기지도 학습 문제로 재정의합니다.
2. SOTA 성능 달성: ImageNet 데이터셋에서 VAE 기반 방법 (DiT, SiT 등) 을 능가하는 성능을 기록했습니다.
   • 확산 모델: ImageNet-256 에서 FID 1.58, ImageNet-512 에서 FID 2.35 (75 NFE 기준) 를 달성하여 기존 픽셀 공간 방법 및 VAE 기반 방법보다 월등히 뛰어났습니다.
   • 일관성 모델: 사전 학습된 VAE 나 확산 모델을 전혀 사용하지 않고, 픽셀 공간에서 직접 훈련된 일관성 모델로서 ImageNet-256 에서 FID 8.82 (1 단계 생성) 를 기록했습니다. 이는 해당 분야에서 최초의 성과입니다.
3. 효율성: DiT(XL/2) 모델보다 약 30% 적은 훈련 계산 자원으로 더 나은 성능을 달성했습니다. 또한, 다양한 해상도에서 패치 크기를 조정하여 GFLOPs 와 훈련 시간을 최적화했습니다.

4. 실험 결과 (Results)

• ImageNet-256: EPG-XXL/16 모델은 FID 1.87 (75 NFE) 을 기록하며, DiT-XL/2 (FID 2.27) 보다 우수한 성능을 보였습니다. 더 큰 모델 (EPG-G/16) 을 사용하면 FID 1.58 까지 개선되었습니다.
• ImageNet-512: EPG-L/32 모델은 FID 2.35 를 기록하며 고해상도 생성에서도 강력한 확장성을 입증했습니다.
• 계산 효율성: Table 5 에 따르면, EPG-XXL/16 은 DiT-XL/2 보다 160 시간 (훈련 시간) 대비 160 시간 (EPG) 으로 유사하거나 더 적은 비용으로 더 좋은 성능 (FID 1.87 vs 2.27) 을 냈습니다. (VAE 훈련 비용 제외 시 비교).
• 일관성 모델: 1 단계 생성 (One-step generation) 에서 FID 8.82 를 기록하며, 잠재 공간 기반의 iCT-XL/2 (FID 34.24) 를 압도적으로 능가했습니다.

5. 의의 및 결론 (Significance)

이 논문은 **"VAE 는 더 이상 필요하다"**는 메시지를 전달하며, 생성 모델 분야에서 VAE 의존성을 탈피할 수 있는 새로운 패러다임을 제시합니다.

• 비즈니스/연구적 가치: VAE 훈련의 복잡성과 성능 병목 현상을 제거하여, 새로운 데이터나 도메인에 생성 모델을 더 쉽게 적응시킬 수 있게 합니다.
• 기술적 혁신: 픽셀 공간에서 직접 훈련하는 것이 반드시 비효율적이라는 편견을 깨뜨렸으며, 자기지도 학습의 원리를 확산/일관성 모델에 성공적으로 적용하여 SOTA 성능을 달성했습니다.
• 미래 전망: 이 프레임워크는 고해상도 및 멀티모달 이미지 생성의 확장 가능한 솔루션으로, 향후 더 강력한 백본을 활용하여 잠재 공간 기반 방법과의 성능 격차를 완전히 해소할 가능성을 보여줍니다.

요약하자면, 이 논문은 자기지도 사전 학습을 통해 인코더의 의미 표현 능력을 극대화하고, 이를 픽셀 공간 생성 모델에 적용함으로써 VAE 없이도 최상급의 생성 품질과 효율성을 달성한 획기적인 연구입니다.