통합 토큰화 및 잠재 노이즈 제거를 위한 엔드 투 엔드 훈련

End-to-End Training for Unified Tokenization and Latent Denoising

잠재 확산 모델(LDMs)은 학습된 잠재 공간에서 작동하여 고품질의 결과를 생성합니다. 그러나 최첨단 LDM을 훈련하려면 확산 모델을 훈련하기 전에 먼저 토크나이저를 훈련해야 하는 복잡한 단계를 거쳐야 합니다. 본 연구에서는 통합 토큰화 및 잠재 확산을 위한 오토인코더 아키텍처인 UNITE를 제안합니다. UNITE는 가중치 공유를 통해 이미지 토크나이저와 잠재 생성기를 모두 수행하는 생성 인코더로 구성됩니다. 우리의 핵심 아이디어는 토큰화와 생성이라는 두 가지 작업이 서로 다른 조건 하에서 동일한 잠재 추론 문제로 볼 수 있다는 것입니다. 토큰화는 완전하게 관찰된 이미지로부터 잠재 변수를 추론하는 반면, 생성은 노이즈와 함께 텍스트 또는 클래스 조건 정보를 사용하여 잠재 변수를 추론합니다. 이러한 점에 착안하여, 동일한 생성 인코더를 통해 두 번의 순방향 과정을 거치면서 두 작업을 동시에 최적화하는 단일 단계 훈련 절차를 도입했습니다. 공유된 파라미터는 잠재 공간을 공동으로 형성하도록 하여 "공통 잠재 언어"를 장려합니다. 이미지 및 분자 모달리티에서 UNITE는 적대적 손실이나 사전 훈련된 인코더(예: DINO) 없이 최첨단 성능에 근접한 결과를 달성하며, ImageNet 256 x 256 데이터셋에서 Base 모델과 Large 모델의 경우 각각 FID 점수가 2.12와 1.73을 기록했습니다. 또한, 생성 인코더를 표현 정렬 및 압축의 관점에서 분석했습니다. 이러한 결과는 처음부터 토큰화 및 생성 작업을 단일 단계로 공동 훈련하는 것이 가능하다는 것을 보여줍니다.

Latent diffusion models (LDMs) enable high-fidelity synthesis by operating in learned latent spaces. However, training state-of-the-art LDMs requires complex staging: a tokenizer must be trained first, before the diffusion model can be trained in the frozen latent space. We propose UNITE - an autoencoder architecture for unified tokenization and latent diffusion. UNITE consists of a Generative Encoder that serves as both image tokenizer and latent generator via weight sharing. Our key insight is that tokenization and generation can be viewed as the same latent inference problem under different conditioning regimes: tokenization infers latents from fully observed images, whereas generation infers them from noise together with text or class conditioning. Motivated by this, we introduce a single-stage training procedure that jointly optimizes both tasks via two forward passes through the same Generative Encoder. The shared parameters enable gradients to jointly shape the latent space, encouraging a "common latent language". Across image and molecule modalities, UNITE achieves near state of the art performance without adversarial losses or pretrained encoders (e.g., DINO), reaching FID 2.12 and 1.73 for Base and Large models on ImageNet 256 x 256. We further analyze the Generative Encoder through the lenses of representation alignment and compression. These results show that single stage joint training of tokenization & generation from scratch is feasible.