Just-in-Time: 확산 트랜스포머를 위한 학습이 필요 없는 공간 가속화 기법

확산 모델은 이미지 생성 시 모든 영역을 동일한 연산량으로 처리하여 비효율적이다. 이 논문은 이미지의 전역적 구조가 세부 묘사보다 먼저 형성된다는 점에 착안하여, 초기 단계에서 일부 핵심 토큰만 계산함으로써 품질 저하 없이 추론 속도를 획기적으로 높였다.

Diffusion Transformer(DiT)는 고품질 이미지를 생성하지만, 입력 토큰 수의 제곱에 비례하는 연산 비용이 발생한다. 특히 이미지 생성 초기 단계에서는 대략적인 윤곽(저주파 성분)이 먼저 형성되는데, 기존 방식은 모든 픽셀 영역에 동일한 연산량을 쏟아부어 비효율적이다. JiT는 'Just-in-Time'이라는 이름처럼 필요한 시점에 필요한 영역만 계산하는 전략을 취한다.

초기에는 듬성듬성한 앵커 토큰들만 계산하고, 이들의 정보를 바탕으로 나머지 영역의 변화(Velocity Field)를 보간(Interpolation)하여 전체 이미지를 업데이트한다. 이는 마치 건물의 뼈대를 먼저 세우고 나중에 벽돌을 채우는 것과 유사한 원리다. 생성 과정이 진행됨에 따라 속도장의 변화가 큰 영역, 즉 세부 묘사가 필요한 부분을 동적으로 감지하여 계산할 토큰을 늘려나간다.

이 접근 방식은 FLUX.1과 같은 최신 모델에서 품질 손실을 최소화하면서도 연산량을 70% 이상 절감하는 성과를 거두었다. 특히 학습이 전혀 필요하지 않아 기존에 공개된 다양한 모델에 즉시 적용할 수 있다는 점이 강력한 차별점이다.

SAG-ODE(Spatially Approximated Generative ODE)는 전체 토큰 세트 중 일부인 앵커 토큰 집합에서만 Transformer 연산을 수행한다. [활성화된 앵커 토큰 y_k를 입력으로] → [Transformer 네트워크 u_theta를 통과시켜 앵커 토큰의 속도를 계산하고] → [증강 리프터 연산자 Π_k를 통해 전체 공간의 속도장 v_t를 생성하여] → [전체 잠재 상태 y(t)를 업데이트하는 방식]을 사용한다.

증강 리프터 Π_k는 앵커 토큰의 정확한 속도값과 비활성 토큰을 위한 보간된 속도값을 결합한다. 보간 연산자 I_k는 최근접 이웃 보간과 마스크된 가우시안 블러를 사용하여 [앵커 토큰의 속도 정보를 주변으로 확산시켜] → [비활성 영역의 속도를 추정하고] → [연속적인 속도장을 형성하여] → [모델 재학습 없이도 전체 구조의 일관성을 유지]하도록 설계되었다.

DMF(Deterministic Micro-Flow)는 새로운 토큰이 추가되는 전이 시점에서 상태 불연속성을 해결한다. [현재 앵커 토큰의 정보와 노이즈를 결합한 타겟 상태 y_k*를 설정하고] → [매우 짧은 시간 간격 동안 유한 시간 ODE를 풀어] → [새로 활성화된 토큰의 상태를 타겟으로 수렴시켜] → [시각적 아티팩트 없는 매끄러운 확장을 보장]한다. 이 과정에서 시간 변화율 (T_k - t)^-1을 사용하여 상태가 정확히 타겟으로 수렴하도록 제어한다.

FLUX.1-dev 모델을 대상으로 한 실험에서 JiT는 50단계 기본 모델 대비 약 7.07배의 속도 향상을 달성했다. CLIP-IQA(0.5397), ImageReward(0.9746), GenEval(0.6457) 등 주요 지표에서 기존 가속화 기법인 RALU나 TeaCache보다 우수한 성능을 보이며 고품질 생성을 유지했다. 특히 텍스트 렌더링과 같은 정교한 작업에서도 아티팩트 없는 결과물을 생성했다.

Ablation Study 결과, SAG-ODE의 공간 근사가 없을 경우 구조적 일관성이 무너져 노이즈가 발생함을 확인했다. 또한 중요도 기반 토큰 활성화(ITA)를 사용하지 않고 고정된 그리드 방식을 사용할 경우, 복잡한 고주파 세부 묘사 영역에서 선명도가 급격히 저하되는 현상이 나타나 동적 자원 할당의 중요성을 입증했다.

JiT는 Flow Matching 기반의 DiT 모델에서 발생하는 공간적 중복성을 활용한다. 전체 토큰 시퀀스 N에 대해 단계별로 확장되는 앵커 토큰 집합의 계층 구조를 정의하여 Coarse-to-Fine 생성을 구현한다. 이는 초기 단계에서 저주파 정보를 먼저 형성하고 나중에 고주파 세부 사항을 채우는 확산 모델의 특성을 수학적으로 공식화한 것이다.

SAG-ODE의 핵심은 일관성 속성(Consistency Property)이다. 선택 행렬의 직교성을 이용하여 앵커 토큰의 동역학이 Transformer의 원래 출력과 정확히 일치하도록 설계함으로써, 근사 과정에서도 학습된 원래의 데이터 분포를 최대한 보존한다. 보간 시에는 가우시안 블러의 커널 크기를 토큰 밀도에 따라 동적으로 조절하여 오버 스무딩을 방지한다.

ITA(Importance-guided Token Activation)는 속도장의 국소 분산을 측정하여 정보 밀도를 정량화한다. 분산이 큰 영역은 생성 프로세스가 활발한 곳으로 간주되어 다음 단계에서 우선적으로 활성화되며, 이는 정적인 샘플링 방식보다 훨씬 효율적인 자원 배분을 가능케 한다. 또한 Qwen-image 및 HunyuanVideo 모델로의 확장을 통해 모델 범용성을 검증했다.