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

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

왜 중요한가

핵심 기여

학습이 필요 없는 공간 가속화 프레임워크 JiT 제안

추가적인 모델 학습이나 증류 과정 없이 기존 확산 트랜스포머(DiT) 모델에 즉시 적용 가능한 공간 도메인 가속화 기법을 개발했다.

공간 근사 생성 상미분 방정식(SAG-ODE) 설계

소수의 앵커 토큰 연산만으로 전체 잠재 상태의 속도장(velocity field)을 근사하여 연산량을 동적으로 줄이는 수식을 정립했다.

결정론적 마이크로 플로우(DMF) 메커니즘 도입

새로운 토큰이 활성화될 때 발생할 수 있는 시각적 아티팩트를 방지하기 위해 통계적 정합성을 유지하며 상태를 전이시키는 유한 시간 ODE를 제안했다.

중요도 기반 동적 토큰 활성화(ITA) 전략

속도장의 국소 분산을 측정하여 정보 밀도가 높은 영역을 우선적으로 계산 자원에 할당하는 동적 선택 방식을 구현했다.

핵심 아이디어 이해하기

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 모델로의 확장을 통해 모델 범용성을 검증했다.

실무 활용

별도의 모델 학습이나 미세 조정 없이 기존 DiT 기반 파이프라인에 즉시 통합하여 추론 속도를 획기적으로 높일 수 있는 실용적인 프레임워크이다.

고해상도 이미지 생성 서비스의 서버 비용 절감 및 사용자 응답 시간 단축
제한된 GPU 자원을 가진 모바일 또는 소비자용 기기에서의 실시간 이미지 생성
긴 시퀀스 처리가 필요한 비디오 생성 모델의 연산 효율화 및 메모리 최적화

코드 공개 여부: 비공개

키워드

DiT(확산 트랜스포머)Spatial Acceleration(공간 가속화)Flow Matching(플로우 매칭)ODE(상미분 방정식)Training-free(학습 불필요)

용어 해설

확산 트랜스포머(Diffusion Transformer): — Transformer 아키텍처를 확산 모델의 백본으로 사용하는 구조이다. 이미지나 비디오 생성 시 토큰 간의 장거리 의존성을 효과적으로 모델링하여 고품질의 결과물을 생성하지만, 토큰 수 증가에 따른 연산 비용이 크다는 특징이 있다.
플로우 매칭(Flow Matching): — 노이즈 분포와 데이터 분포 사이의 결정론적 이동 경로를 정의하는 속도장(Velocity Field)을 직접 학습하는 프레임워크이다. 기존 확산 모델보다 샘플링 효율이 높고 수학적으로 간결하여 최신 생성 모델에서 널리 채택된다.
공간적 중복성(Spatial Redundancy): — 이미지 내의 인접한 픽셀들이나 토큰들이 서로 유사한 정보를 공유하고 있어, 모든 부분을 독립적으로 연산하지 않아도 전체적인 구조를 파악할 수 있는 특성이다. 이를 활용하면 연산량을 획기적으로 줄일 수 있다.
상미분 방정식(Ordinary Differential Equation): — 하나의 독립 변수(주로 시간)에 대한 함수와 그 도함수들 사이의 관계를 나타내는 방정식이다. 생성 AI에서는 노이즈가 데이터로 변해가는 궤적을 수학적으로 정의하고 추적하는 데 사용된다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

관련 토론

댓글

관련 피드