손상된 잠재에서 의미와 세부를 예측해 확산 모델이 배우기 쉬운 토크나이저를 만든다

무엇인가

이 논문은 잠재 확산 모델이 이미지를 생성할 때 쓰는 시각 토크나이저의 잠재 공간을 다룬다. 토크나이저는 재구성을 위해 많은 정보를 보존해야 하지만, 그 정보가 확산 모델이 효율적으로 학습할 수 있는 형태로 조직되어야 한다. 논문은 이 두 목표가 충돌할 수 있다는 전제에서 출발해, REPA-E의 개선이 공동 확산 학습 때문인지 손상된 잠재로부터의 의미 예측 때문인지 분리해 묻는다. REPA-E에서 흐름 매칭 목적과 확산 전용 생성 블록을 제거하고 손상된 잠재에서 사전학습 특징을 예측하는 목적만 남겨도, 동일 조건의 다운스트림 생성 성능이 비슷했다는 관찰이 핵심이다. 저자들은 이를 일반 원리인 견고한 예측으로 정리한다. 즉 생성에 중요한 정보는 손상된 잠재에서도 복구 가능해야 한다.

어떻게 동작하나

방법은 두 보조 목적을 표준 VAE 목적에 더한다. RobustAlign은 보조 의미 예측기 Fη가 손상된 잠재 z_t(φ)에서 동결된 DINOv2-B 특징 T(x)를 예측하도록 학습한다. 여기서 z_t(φ)는 정규화된 인코더 잠재에 (1-t)를 곱하고 tε를 더한 것이며 t는 0과 1 사이에서 균일하게 샘플링된다. 손실은 패치별 코사인 유사도를 1에서 뺀 값이고, 이 손실은 손상된 잠재를 통해 인코더 파라미터까지 역전파된다. RobustAlign은 객체 수준 구조와 전경-배경 분리를 강조하지만 미세한 공간·외관 변화는 덜 담을 수 있다. RobustRecon은 이를 보완하기 위해 별도 보조 디코더 Gξ가 손상된 잠재에서 원본 이미지를 직접 복원하도록 MAE와 LPIPS 손실로 학습한다. Gξ는 1단계 VAE 디코더에서 초기화되지만 독립 파라미터를 쓰고 토크나이저 학습 후 버려진다. 최종 목적은 L_VAE에 λA L_RobustAlign과 λR L_RobustRecon을 더한 형태이며, 두 보조 목적 모두 손상된 잠재를 통해 인코더를 조정하고 표준 VAE 목적은 깨끗한 잠재에서의 재구성을 유지한다.

무엇과 다른가

논문은 견고한 예측이 왜 확산 학습성을 높이는지에 대한 분석도 제시한다. 등방성 가우시안 손상 아래에서 정규화된 잠재의 i번째 고유방향 분산을 ν_i라 하면, 그 방향의 신호 대 잡음비는 (1-t)^2 ν_i / t^2가 된다. 따라서 같은 손상 수준에서 분산이 큰 잠재 방향은 상대적으로 더 많은 신호를 유지한다. 저자들은 견고한 예측 목적이 생성 관련 정보를 이런 고분산 주성분 방향에 더 접근 가능하게 만들고, 선행 연구에 따르면 확산 모델이 고분산 공분산 모드를 더 일찍 학습하는 경향이 있다고 설명한다. 잠재 PCA 시각화에서 RobustAlign은 객체 수준 조직과 전경-배경 분리를, RobustRecon은 더 풍부한 국소 변화를 드러내며, 둘을 결합하면 두 구조가 함께 나타난다고 보고한다.

어떻게 쓰나

실험은 ImageNet 256×256에서 토크나이저를 학습해 재구성과 다운스트림 텍스트-이미지 생성을 평가한다. 토크나이저는 8배 공간 다운샘플링, 32개 잠재 채널, 약 84M 파라미터의 합성곱 구조다. 200k 스텝 VAE 워밍업 후 λA=1.5, λR=0.5로 250k 스텝 RoAR 정제를 수행한다. 다운스트림 평가에서는 토크나이저를 동결하고, 24층·히든 1024·약 0.8B 파라미터의 MMDiT와 동결된 Qwen3-0.6B 텍스트 인코더를 필터링·재캡션된 LAION 256×256에서 처음부터 학습한다. 비교 대상은 SD3-VAE, FLUX.1-VAE, VA-VAE, VTP, PAE, REPA-E-FLUX이며, FLUX.2-VAE는 토크나이저 학습 데이터와 레시피가 공개되지 않아 참고로만 제시된다. 재구성은 ImageNet 검증셋의 rFID, PSNR, LPIPS, SSIM과 TokBench의 텍스트·얼굴 세부 보존으로, 생성은 GenEval, GenEval2, DPG-Bench에서 네 시드 평균으로 평가한다.

전제와 한계

고정된 32k 스텝 다운스트림 학습 예산에서 RoAR-VAE는 PSNR 32.88 dB와 LPIPS 0.0349를 기록했다. 재구성 중심의 FLUX.1-VAE와 비교해 PSNR은 1.80 dB 높고 LPIPS는 0.0510에서 0.0349로 낮아졌다. 동시에 학습성 강화 토크나이저 중 가장 강한 기준보다 평균 다운스트림 생성 점수가 2.19점 높았고, FLUX.2-VAE의 재구성·생성 성능에 근접한다고 보고한다. 텍스트와 얼굴 세부 재구성에서는 모든 지표에서 가장 좋았고, 작은 텍스트 재구성 정확도는 FLUX.1-VAE보다 16.93%p 높았으며 작은 얼굴 유사도는 0.53에서 0.65로 올랐다. 절제 실험에서 Vanilla-KL의 평균 점수는 50.50이었고 RobustRecon만으로 59.38, RobustAlign만으로 62.49, 둘을 결합해 63.86이었다. 손상된 잠재를 제거하면 평균이 63.86에서 58.47로 떨어져 잠재 손상 자체가 핵심 요소임을 보였다. RobustAlign의 의미 예측기를 Transformer 대신 MLP로 바꾸면 성능이 크게 떨어지고, Gξ와 VAE 디코더의 파라미터를 공유하면 약간 나빠진다는 설계 실험도 제시된다.

개발자 관점에서 이 논문은 토크나이저를 단순 재구성 압축기로만 고르면 다운스트림 확산 모델의 학습 효율을 놓칠 수 있음을 보여준다. 특히 전체 확산 모델을 공동 학습하지 않고도 손상된 잠재에서 깨끗한 목표를 예측하는 보조 목적을 붙여 잠재 공간을 다듬을 수 있고, 보조 예측기는 학습 후 버릴 수 있다는 점이 실무적이다. 다만 이득은 동일한 다운스트림 학습 조건에서 측정된 것이므로, 실제 파이프라인에서는 손상 스케줄, λA와 λR, 예측기 용량, 그리고 재구성 지표만이 아니라 GenEval류 생성 지표와 텍스트·얼굴 세부 보존을 함께 확인해야 한다. 논문이 제안하는 것은 특정 토크나이저 구조에 붙일 수 있는 학습 원리이지, 기존 VAE를 그대로 두고 성능이 보장되는 플러그인은 아니다.

저자들이 밝힌 범위와 전제도 분명하다. 토크나이저는 통제 비교를 위해 ImageNet 256×256에서만 학습되었고, 다운스트림 생성 모델은 필터링·재캡션된 LAION에서 학습된다. FLUX.2-VAE는 산업용 참고로만 쓰였고 학습 데이터와 레시피가 비공개라 기준 비교에서 제외되었다. 또한 이 토크나이저는 범용 정보 압축기이며 인간 중심 데이터로 특별히 학습되거나 인간 관련 속성에 최적화된 것이 아니므로, 이를 기반으로 한 생성 시스템에는 워터마킹과 모니터링 같은 보호 장치를 권고한다. 논문은 코드와 학습된 모델을 공개할 예정이라고 밝힌다.