확산 모델에 친화적인 잠재 매니폴드를 위한 조건은 무엇인가? 잠재 확산을 위한 사전 정렬 오토인코더

기존의 토크나이저들은 단순히 이미지를 잘 복원하는 데만 집중하여 확산 모델이 학습하기 어려운 복잡한 잠재 공간을 만들었습니다. 이 논문은 확산 모델이 학습하기 쉬운 공간의 특성을 정의하고 이를 강제하는 PAE를 제안하여 학습 효율과 생성 품질을 동시에 획기적으로 개선했습니다.

Latent Diffusion Model은 이미지를 압축된 잠재 공간(Latent Space)으로 변환한 뒤 그 안에서 데이터를 생성하는 법을 배운다. 기존에는 원본 이미지를 얼마나 똑같이 복원하느냐(Reconstruction Fidelity)를 기준으로 토크나이저를 만들었으나, 복원이 잘 된다고 해서 확산 모델이 그 공간의 규칙을 쉽게 배울 수 있는 것은 아니다. 마치 복잡하게 얽힌 실타래는 풀기 어렵지만, 색깔별로 잘 정리된 실타래는 다루기 쉬운 것과 같은 이치다.

PAE는 이 실타래를 정리하기 위해 시각 기초 모델(VFM)이 이미 알고 있는 '이미지의 구조적 특징'을 가이드로 삼는다. 잠재 공간 내의 인접한 점들이 시각적으로도 유사한 이미지를 생성하도록 연속성을 부여하고, 비슷한 의미를 가진 데이터들이 근처에 모이도록 배치한다. 이는 확산 모델이 노이즈를 제거하는 과정에서 예측해야 할 목표를 훨씬 단순하고 매끄럽게 만들어준다.

결과적으로 확산 모델은 복잡한 공간 구조를 해석하는 데 에너지를 쓰는 대신, 실제 이미지의 패턴을 배우는 데 집중할 수 있게 된다. 이는 적은 학습 횟수로도 더 고품질의 이미지를 생성할 수 있게 하며, 특히 적은 단계의 샘플링(Few-step sampling)에서도 안정적인 성능을 보장하는 결과로 이어진다.

PAE는 고정된 시각 기초 모델(VFM)을 참조하여 잠재 공간을 정형화하는 세 가지 정규화 기법을 핵심으로 한다. 첫째, Spatial Structure Regularization (SSR)은 잠재 토큰 간의 관계를 나타내는 Gram Matrix를 계산하고 이를 VFM의 특징 맵 구조와 일치시킨다. [잠재 토큰 간 유사도 행렬 입력 → VFM 특징 간 유사도 행렬과 비교 → 차이 최소화 → 개별 인스턴스의 공간적 일관성 확보]

둘째, Manifold Continuity Regularization (MCR)은 잠재 벡터에 미세한 섭동(Perturbation)을 가한 뒤 디코딩된 결과물들이 시각적으로 연속성을 유지하도록 강제한다. [잠재 벡터 z와 섭동된 z' 입력 → 디코더 통과 후 이미지 생성 → LPIPS 거리 계산 → 국부적 매니폴드의 부드러움 확보] 이때 계층적 섭동 방식을 사용하여 복원 품질 손상을 최소화한다.

셋째, Semantic Consistency Regularization (SCR)은 잠재 공간의 전역적 의미 구조를 VFM의 임베딩 공간과 정렬한다. [전역 풀링된 잠재 벡터 입력 → VFM의 의미론적 특징과 코사인 유사도 계산 → 손실 반영 → 클래스별 군집화 유도] 또한, Detail-aware Modulator (DAM)를 도입하여 VFM의 강력한 의미 정보는 유지하면서도 세부적인 픽셀 디테일을 보완하여 높은 복원 성능을 동시에 달성한다.

ImageNet 256x256 데이터셋에서 PAE는 압도적인 성능을 보였다. 80 에포크 학습만으로도 기존의 강력한 베이스라인인 RAE가 800 에포크 이상 학습했을 때와 유사한 성능에 도달하며 약 13배의 수렴 가속화를 입증했다. 최종적으로 800 에포크 학습 시 gFID 1.03을 기록하며 기존 최고 기록을 경신했다.

Ablation Study 결과, SSR, MCR, SCR 세 가지 정규화 요소가 각각 공간 구조(SSC), 연속성(LPC), 의미론(GSQ) 지표를 독립적으로 개선함이 확인되었다. 특히 LPC 지표가 개선될수록 확산 모델의 학습 안정성과 생성 품질이 비례하여 상승하는 강한 상관관계가 나타났다.

Few-step Sampling 실험에서도 PAE는 뛰어난 효율성을 보였다. 단 15단계의 샘플링만으로도 기존 FAE 모델의 250단계 샘플링 결과와 대등한 gFID 수치를 기록했으며, 45단계에서는 gFID 1.05를 달성하여 적은 연산으로도 고품질 이미지 생성이 가능함을 증명했다.

PAE 아키텍처는 Frozen VFM(DINOv2 등)을 인코더의 백본으로 활용하며, 여기에 학습 가능한 Detail-aware Modulator (DAM)를 결합한다. DAM은 입력 이미지를 패치화하여 Transformer 블록을 통해 VFM 특징에 픽셀 레벨의 세부 정보를 주입한다. 이후 RMS 정규화를 거쳐 구형(Sphere) 잠재 공간을 형성함으로써 확산 모델의 추론 효율을 높이고 섭동 기반 정규화의 안정성을 확보한다.

학습 과정은 두 단계로 나뉜다. 먼저 VFM의 고차원 특징을 토크나이저의 병목 차원에 맞게 압축하고 공간적 노이즈를 제거하는 VFM Refine 단계를 거친다. 이후 정제된 VFM Prior를 타겟으로 삼아 토크나이저의 본 학습을 진행한다. 이러한 설계는 VFM의 강력한 사전 지식을 직접 상속받는 대신, 확산 모델이 작동하기에 최적화된 '매니폴드 구조'를 형성하는 데 초점을 맞춘다.

현재 연구는 ImageNet 256x256 해상도에 집중되어 있어, 더 높은 해상도나 다양한 도메인(비디오 등)으로의 일반화 가능성에 대해서는 추가 검증이 필요하다. 또한 가변 해상도나 동적 토큰 할당과 같은 복잡한 설정에서의 동작 방식은 아직 다루지 않았다.