[Paper] End-to-End 학습을 통한 통합 토크나이제이션 및 잠재 디노이징

Source: arXiv - 2603.22283v1

번역할 텍스트를 제공해 주시면, 요청하신 대로 한국어로 번역해 드리겠습니다.

개요

이 논문은 이미지 토큰화와 잠재 확산을 하나의 공유 모델로 통합하는 UNITE라는 단일‑단계 오토인코더를 소개합니다. 토큰화(이미지 인코딩)와 생성(노이즈에서 디코딩)을 동일한 추론 문제의 양면으로 간주함으로써, UNITE는 현재 잠재 확산 모델(LDM)들이 의존하는 번거로운 다‑단계 파이프라인을 제거합니다.

Key Contributions

• 통합 아키텍처: 가중치 공유를 통해 토크나이저와 잠재 생성기 역할을 동시에 수행하는 생성 인코더.
• 단일‑단계 학습: 배치당 두 번의 순전파만으로 토크나이제이션과 디퓨전을 공동 최적화하여 별도로 사전 학습된 인코더가 필요 없게 함.
• 공통 잠재 언어: 공유 파라미터가 재구성 및 조건부 생성 모두를 위한 잠재 공간을 정렬하여 표현 품질을 향상시킴.
• 최첨단 결과, 적대적 기법 없이: ImageNet 256×256에서 FID 2.12 (Base)와 1.73 (Large)을 달성하여, 대규모 사전 학습 백본에 의존하는 기존 LDM들을 능가하거나 동등한 성능을 보임.
• 교차‑모달 검증: 동일한 프레임워크가 분자 그래프 생성에도 적용 가능함을 보여, 이 접근법이 자연 이미지에만 국한되지 않음을 입증함.

방법론

1. Generative Encoder (GE) – 트랜스포머‑스타일 인코더로, 다음 중 하나를 입력받습니다:

   • 이미지 (전체 관측) → 잠재 토큰 시퀀스를 생성합니다 (토큰화).
   • 무작위 노이즈 + 조건부 입력 (예: 텍스트 프롬프트) → 이미지로 디코딩될 잠재 시퀀스를 예측합니다 (생성).

2. 미니‑배치당 두 단계 훈련 루프:

   • 1단계 (Tokenization) – 실제 이미지를 GE에 입력하여 잠재값을 얻고, 가벼운 디코더로 이미지를 복원합니다. 복원 손실 (L2 + 퍼셉추얼)을 계산합니다.
   • 2단계 (Diffusion) – 가우시안 노이즈를 샘플링하고 동일한 조건부 입력을 연결한 뒤 GE에 통과시켜 잠재 시퀀스를 예측하고, 이를 디코딩합니다. 디퓨전 손실 (denoising score matching)과 복원된 이미지에 대한 동일한 복원 손실을 함께 계산합니다.

3. 가중치 공유는 두 작업에서 오는 그래디언트가 동일한 잠재 공간을 형성하도록 하여, 인코더가 압축 가능하고 (좋은 토큰화) 생성하기 쉬운 (샘플링이 용이한) 표현을 학습하도록 장려합니다.

4. 훈련 세부 사항 – 외부 사전 학습된 인코더 (예: DINO)나 적대적 목표는 사용되지 않습니다. 모델은 ImageNet (및 분자 데이터셋)에서 스크래치로 훈련되며, 표준 AdamW 옵티마이저와 코사인 학습률 스케줄을 사용합니다.

결과 및 발견

• 표현 분석: 토크나이제이션 및 diffusion에서 얻은 잠재 벡터가 높은 코사인 유사도(↑)를 보이며 “공통 잠재 언어”가 존재함을 확인.
• 압축: 학습된 토큰 어휘는 최소한의 지각 손실로 약 8배 압축을 달성하여 전용 자동인코딩 방식과 경쟁.
• 소거 실험: 가중치 공유를 제거하거나 두 작업을 별도로 학습하면 FID가 약 0.5–0.8 정도 악화되어 공동 최적화의 이점을 강조.

Practical Implications

• Simplified pipelines: 개발자는 이제 별도의 토크나이저 단계 없이 확산 모델을 훈련할 수 있어 엔지니어링 오버헤드와 저장소(동결된 인코더 체크포인트 필요 없음)를 줄일 수 있습니다.
• Faster prototyping: 단일 훈련 스크립트가 인코딩과 생성 모두를 처리하므로 새로운 조건 신호(예: 맞춤 태그, 멀티모달 입력)를 실험하기가 쉬워집니다.
• Resource efficiency: 처음부터 훈련함으로써 대규모 사전학습 비전 모델이 필요 없어 GPU 메모리와 연산 요구량을 낮출 수 있어, 대규모 사전학습 예산을 감당할 수 없는 조직에 유리합니다.
• Cross‑domain generation: 동일한 아키텍처를 이미지가 아닌 데이터(예: 분자 그래프)에도 재활용할 수 있어, 여러 데이터 유형을 단일 코드베이스로 처리하는 통합 생성 AI 플랫폼의 문을 엽니다.
• Potential for on‑device inference: 토크나이징과 디코딩이 가중치를 공유하므로, 컴팩트한 런타임이 엣지 디바이스에서 “토큰으로 인코딩” 및 “토큰에서 샘플링” 기능을 모두 제공할 수 있습니다.

제한 사항 및 향후 연구

• 스케일링 동작이 완전히 탐색되지 않음: 실험은 256×256 해상도에서 멈추며, 이 접근법이 초고해상도 합성(예: 1024×1024)으로 어떻게 확장되는지는 아직 불분명합니다.
• 조건부 유연성: 논문은 클래스 수준 조건부에 초점을 맞추고 있으며, 보다 풍부한 텍스트 프롬프트나 멀티모달 신호로 확장하려면 아키텍처 조정이 필요할 수 있습니다.
• 학습 안정성: 적대적 손실을 피하면서도, 재구성 손실과 확산 손실 사이의 가중치에 민감할 수 있어 신중한 하이퍼파라미터 튜닝이 요구됩니다.
• 벤치마크 범위: 평가가 ImageNet과 단일 분자 데이터셋에만 제한되어 있으며, 의료 영상이나 비디오와 같은 더 넓은 도메인 테스트가 모달리티‑불가지론 주장에 힘을 실어줄 것입니다.

향후 연구 방향으로는 UNITE를 더 큰 해상도로 확장하고, 보다 표현력 있는 조건부(예: 전체 문장 캡션)를 통합하며, 학습 중 토크나이제이션에서 확산으로 점진적으로 강조점을 이동시키는 커리큘럼 전략을 조사하는 것이 포함됩니다.

저자

• Shivam Duggal
• Xingjian Bai
• Zongze Wu
• Richard Zhang
• Eli Shechtman
• Antonio Torralba
• Phillip Isola
• William T. Freeman

논문 정보

• arXiv ID: 2603.22283v1
• 카테고리: cs.CV, cs.AI, cs.GR, cs.LG
• 발행일: 2026년 3월 23일
• PDF: PDF 다운로드