JLT: Latent Diffusion Transformer에서의 Clean-Latent 예측

고정된 FLUX.2 VAE 잠재 표현에서 타깃 예측의 기하학적 차이가 학습 난이도와 생성 품질에 직접 영향을 준다. x, ϵ, v 간의 선형적 관계가 존재하나 finite-capacity 모델은 타깃의 기하학에 따라 회귀 문제의 난이도와 샘플링 효율이 다르게 나타난다. 이 연구는 latent diffusion에서도 타깃 지향이 단순한 algebraic 재표현이 아님을 보이며, 동일 표현 공간에서의 예측 타깃 설계가 성능에 큰 차이를 만든다는 점을 제시한다.

출발점: 픽셀 공간이 아니라 FIXED FLUX.2 VAE latent 공간에서의 타깃(x, ϵ, v) 차이가 모델 학습에 어떤 영향을 주는지 관찰한다. x, ϵ, v는 t에 따라 서로 선형적으로 변환되지만, finite-capacity 네트워크는 각 타깃의 분포가 다르게 다뤄진다. 이는 타깃의 기하학이 회귀 문제의 난이도와 일반화에 직결될 수 있음을 시사한다. 해결 원리: Cov(yx)=Σ, Cov(yε)=I, Cov(yv)=Σ+I로 인해 velocity 타깃은 모든 방향에 단위 분산의 바닥을 추가하고, 고변량 방향에서의 차이도 증가시키며, clean 타깃은 저변량 방향의 타깃 분산을 축소한다. 이는 x와 v가 선형적으로 변환 가능하더라도 학습 시점의 회귀 문제는 서로 다르게 구성된다는 것을 뜻한다. 달라지는 점: 고정된 표현 공간에서의 타깃 선택은 단순한 수식적 재정의가 아니라, 실제로 학습 커브와 샘플 효율, 생성 품질에 차이를 만들며, 이로써 latent diffusion에서도 타깃 기하학이 설계의 중요한 요소임을 확인한다.

1. 입력 정의: x ~ N(0, Σ)로 정해진 clean latent와 ε ~ N(0, I)로 정의된 잡음을 사용하고, zt = t x + (1 − t) ε를 따른다. 2) 목표 타깃: yx = x, yε = ε, yv = x − ε로 정의한다. 3) 타깃 간 변환 규칙: 예측된 타깃 ŷθ가 주어지면 Affine readout으로 다른 엔드포인트를 복원한다. 예를 들어 ϵ̂(x)θ = (zt − t x̂θ)/(1 − t), v̂(x)θ = (x̂θ − zt)/(1 − t)이다. 4) 수학적 관계의 재해석: 입력과 예측 타깃이 주어졌을 때의 오차는 [직관적 읽기 규칙]에 따라 재가중되며, x, ϵ, v 간의 변환은 예측 오차의 재가중치를 통해 학습에 다른 영향을 준다. 5) 로컬 가우시안 분석: Σ의 고유값 λi에 대해 Di = t^2 λi + (1 − t)^2로 정의하고, Var(xi|zi), Var(vi|zi) 등을 구해 타깃 간 조건부 불확실성 차이를 도출한다. 6) 구현 세부: JiT-B/16에 근접한 아키텍처로, 12 Transformer 블록, hidden 768, 12 heads, 128 패치 임베딩, 130M 파라미터, FLUX.2 VAE latent를 사용하며, 250K 스텝(200 에폭)으로 학습한다. 보조 구성요소로 class-token concatenation 및 ImageNet 보조 손실은 제외한다.

주요 벤치마크는 ImageNet 256 × 256에서 FID-50K로 측정한다. 표 1의 매칭 타깃ablation에서 /1의 경우 clean-latent(x)으로 예측하는 JLT-B/1이 DiT-B/1의 6.56에서 2.56으로 감소했고, /2의 경우 28.71에서 14.81로 감소했다. 최종적으로 JLT-B/1은 x w/ CFG에서 2.50, 232.51 IS를 달성했고 x w/o CFG에서는 14.00을 보였다. 추가로 4.2의 대표 baselines와의 비교에서 JLT-B/1의 최종 결과는 FID-50K 2.50, IS 232.51로 제시되며, JiT-L/16의 2.79, LDM의 3.60, JiT-B/16의 4.37 등과 비교된다. 학습 곡선은 Fig. 2에 요약되어 있으며, 100K 스텝 부근에서 aligned된 low-FID 영역으로 진입하고 최종 체크포인트에서 velocity 대비 확실한 여유를 유지한다. qualitative 샘플은 Figure 1의 teaser에 제시된다.

단락 1: 아키텍처 구조 — JLT는 12개의 Transformer 블록, hidden dimension 768, 12 heads, 128 차원의 패치 임베딩을 채택한 130M 파라미터의 Base-scale latent Transformer이다. 입력 표현은 fixed FLUX.2 VAE latent이며, /1과 /2의 VAE latent grid를 사용한다. 단, 125-epoch 수렴 시점의 평가를 위해 JiT-B/16과 같은 구성으로 비교가 가능하도록 설계되었다. 단락 2: 핵심 메커니즘의 수학적 기반 — 타깃(yx, yε, yv) 간의 관계는 zt = t x + (1 − t) ε에서 도출되며, x̂θ, ϵ̂θ, v̂θ 등의 예측은 Affine readout을 통해 다른 엔드포인트로 변환된다. 이는 수식적으로 ŷ를 예측한 후, 1−t, t 계수에 따라 남은 엔드포인트를 연산으로 구하는 형태이다. 또한 Cov(yx)=Σ, Cov(yε)=I, Cov(yv)=Σ+I로 정의되는 로컬 가우시안 근사를 통해 타깃의 분산 구조 차이가 학습에 미치는 영향을 설명한다. 단락 3: Prior work 대비 차별점 — latent space에서의 타깃 선택은 단순한 표현 압축의 영향이 아니라, 학습 시의 회귀 문제를 구성하는 핵심 요소임을 실험적으로 보인다. x와 v는 선형적으로 변환 가능하나 finite-capacity 모델의 학습 다이나믹스에서 차이를 보인다. 단락 4: 구현·학습 세부사항 및 이론적 분석 — FLUX.2 VAE latent 공간 고정, 250K 스텝(200 에폭) 학습, 130M 파라미터, 12블록 Transformer, 768 차원, patch-embedding 128, classifier-free guidance를 포함한 샘플링 설정, 보조 손실 제외 등의 설정이 기술된다. Appendix의 수식 변환 및 공분산 도출은 타깃 간 차이가 학습에 미치는 영향을 수학적으로 뒷받침한다.

본 연구는 ImageNet 256 × 256 및 130M 파라미터의 JLT-B/1 구성에 한정된다. 다른 데이터셋, 다른 토크나이저, 다른 샘플링 전략에서의 일반화는 확인되지 않는다. Appendix D에 제시된 추가 기하학 진단은 추가 검증 도구로 제시되지만 본 연구의 핵심 주장에는 포함되지 않았다.