통합 토크나이저와 잠재 디노이징을 위한 단일 단계 학습

초록

UNITE는 이미지 토크나이저와 잠재 확산 디노이저를 하나의 Generative Encoder 에 공유함으로써, 토크나이징과 생성 과정을 단일 단계에서 동시에 학습한다. 재구성 손실과 흐름‑매칭 디노이징 손실만으로 사전 학습된 인코더 없이도 ImageNet‑256에서 FID 2.12(베이스)·1.73(라지) 수준의 성능을 달성한다.

상세 분석

UNITE의 핵심 아이디어는 토크나이징과 생성이 “조건부 잠재 추론”이라는 동일한 문제의 두 변형이라는 점이다. 토크나이징은 완전 관측된 이미지 x 로부터 잠재 z₀ 를 거의 단일점 분포로 추정하고, 생성은 잡음 z_t 과 텍스트·클래스 조건을 이용해 동일한 z₀ 를 복원한다는 관점에서 두 과정은 동일한 네트워크 GE_θ 에 의해 구현될 수 있다. 논문은 이를 실현하기 위해 GE_θ 를 두 모드로 동작하도록 설계했으며, 첫 번째 순전파에서는 x → z₀ 를, 두 번째 순전파에서는 z_t → \hat{z₀} 를 수행한다. 두 순전파 모두 동일한 파라미터를 공유함으로써 재구성 손실과 흐름‑매칭 디노이징 손실이 동시에 파라미터를 업데이트한다.

이 구조는 기존 LDM 파이프라인이 겪는 “토크나이저 고정 → 디퓨전 모델 학습”이라는 단계적 학습의 비효율성을 근본적으로 제거한다. 또한, 별도의 adversarial loss나 DINO·MAE와 같은 사전 학습된 비전 인코더를 필요로 하지 않는다. 실험에서는 ImageNet‑256에서 베이스 모델(FID 2.12)과 라지 모델(FID 1.73)을 달성했으며, 분자 생성에서도 경쟁력 있는 결과를 보였다.

분석 파트에서는 GE_θ 의 레이어별 표현 정렬을 CKA 지표로 측정했는데, 파라미터를 완전히 분리한 경우에도 높은 정렬을 보이며 토크나이징과 디노이징이 본질적으로 호환되는 작업임을 확인했다. 그러나 파라미터 공유가 rFID·gFID 트레이드오프에서 가장 우수한 결과를 제공한다는 점에서, 공유가 “공통 잠재 언어”를 형성하는 데 기여함을 실증했다.

또한, 흐름‑매칭(Flow‑Matching) 손실을 사용해 잡음 단계 t 에 따라 z_t 를 생성하고, 이를 다시 GE_θ 에 입력해 z₀ 를 예측한다. 이는 전통적인 DDPM의 스코어 매칭보다 학습 안정성이 높고, 고해상도 이미지에서도 효율적인 역전파를 가능하게 한다. 전체 파이프라인은 Encoder → Decoder 구조만을 유지하면서도, 토크나이저와 생성기를 동시에 최적화하는 단일 학습 루프를 제공한다.

결과적으로 UNITE는 “하나의 네트워크가 토크나이저와 디노이저 역할을 동시에 수행한다”는 새로운 패러다임을 제시하며, 파이프라인 복잡성을 크게 낮추고, 학습 비용과 저장 공간을 절감한다. 이는 향후 멀티모달·대규모 생성 모델 설계에 중요한 설계 원칙이 될 것으로 기대된다.