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

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

왜 중요한가

핵심 기여

Controlled latent target ablation으로 타깃 기하학의 효과 증명

같은 표현 공간과 학습 설정에서 clean-latent 예측(x)과 velocity 예측(v) 간 직접 비교를 수행해, FID-50K가 /1에서 6.56에서 2.50으로, /2에서 28.71에서 14.81로 큰 차이를 보임을 확인했다. 동일한 토큰화와 구성, 학습 일정 하에서 타깃 선택이 결과에 직접 기여함을 보여준다.

Latent 공간에서의 예측 타깃 설계가 성능에 미치는 영향의 메커니즘 제시

로컬 가우시안 분석을 통해 velocity(v=x−ϵ) 타깃은 등방성 Covariance Floor를 추가하고 저변량 방향을 증폭하는 반면, clean-latent(yx=x) 타깃은 저변량 방향을 억제한다는 기전을 수학적으로 제시했다.

JLT: FLUX.2 VAE 잠재 공간에서의 Base-scale latent Transformer 구성

JLT는 130M 파라미터의 Base-scale latent Transformer로, 고정된 FLUX.2 VAE latent space에서 타깃 변경만을 실험 변수로 삼는다. 12 Transformer 블록, hidden 768, 12 heads, 128 패치 임베딩, 250K 스텝(200 에폭)으로 학습한다. JiT-B/16과의 비교를 의도적으로 통제한다.

코드 공개 의의

JLT의 구현 및 실험은 https://github.com/akatsuki-neo/JLT에서 공개될 수 있음을 명시한다(문서에 명시된 코드 저장소 URL).

핵심 아이디어 이해하기

출발점: 픽셀 공간이 아니라 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에서도 타깃 기하학이 설계의 중요한 요소임을 확인한다.

방법론

입력 정의: 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을 통해 다른 엔드포인트로 변환된다. 이는 수식적으로 ŷ를 예측한 후, 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에 제시된 추가 기하학 진단은 추가 검증 도구로 제시되지만 본 연구의 핵심 주장에는 포함되지 않았다.

실무 활용

고정된 latent 표현에서의 타깃 지향 설계가 latent diffusion의 학습 및 샘플링 품질에 실질적 영향을 준다는 점을 시사한다. latent 공간의 타깃 선택이 모델의 표현력과 샘플 품질 간의 trade-off를 좌우할 수 있음을 보여준다.

latent-space 기반 고해상도 이미지 합성에서 타깃 예측 방식의 조합으로 샘플 품질 향상
FIXED latent 공간에서의 flow matching 및 diffusion 모델 실험 설계 시 타깃 geometry를 고려한 ablation 연구
라벨-조건 없는(class-conditional 없이) latent diffusion 모델의 샘플링 품질 향상을 위한 타깃 설계 연구

코드 공개 여부: 공개

코드 저장소 보기

키워드

flow matchingclean-data predictionlatent-spacelatent diffusion TransformerFLUX.2 VAEclassifier-free guidanceFID-50K

용어 해설

타깃 기하학(Target Geometry): — 잠재 공간에서의 타깃 기하학은 신경망이 학습하는 목표 분포의 모양을 의미한다. x(clean latent), ϵ(noise), v(x−ϵ) 간의 선형적 관계가 존재하더라도 finite-capacity 모델은 각 타깃에 따라 오차 분포와 학습 난이도가 달라진다. 이로 인해 동일한 표현 공간에서도 학습 효율과 샘플링 속도가 달라지며, 본 연구는 latent 공간에서의 타깃 선택이 성능 차이에 직접 기여함을 보인다.
로컬 가우시안 근사(Local Gaussian Approximation): — 데이터 분포를 로컬한 선형-가우시안 근사로 가정하고 공분산 스펙트럼을 해석한다. 이 근사는 x의 공분산 Σ를 통해 고변량 방향과 저변량 방향을 구분하고, 타깃(yx, yε, yv)의 조건부 분산 차이를 도출하는 기초가 된다.
등방성 공분산 바닥(Isotropic Covariance Floor): — _velocity 타깃은 모든 방향에 대해 단위 분산의 등방성 바닥을 추가하는 효과를 준다. 고유 방향의 분산이 작을수록 clean 타깃은 이를 억제하고, velocity 타깃은 저분산 방향까지 영향력을 확장한다는 점이 분석된다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

GenFirst로 잠재 붕괴 없는 생성 학습

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

GenFirst로 잠재 붕괴 없는 생성 학습