Predict to Skip: Linear Multistep Feature Forecasting for Efficient Diffusion Transformers

이 논문은 확산 트랜스포머의 반복적 탈노이즈 과정에 선형 다단계 예측 기법을 도입하여 계산 비용을 획기적으로 줄이면서도 생성 품질을 유지하는 'PrediT'라는 훈련 없는 가속화 프레임워크를 제안합니다.

핵심

이 논문은 **'PrediT(프레딧)'**이라는 새로운 기술을 소개합니다. 이 기술은 최근 화제가 되는 '이미지나 영상을 만드는 AI(확산 모델)'가 훨씬 더 빠르게 작동하도록 도와줍니다.

기존의 AI는 그림을 그리거나 영상을 만들 때, 마치 안개 낀 산길을 한 걸음 한 걸음 천천히 걸어가며 목적지에 도달하는 것처럼, 수백 번의 반복 작업을 거쳐야 했습니다. 이 과정은 시간이 너무 오래 걸리고 전기도 많이 먹습니다.

이 논문은 그 과정을 어떻게 스마트하게 단축했는지 설명합니다. 핵심 아이디어를 일상적인 비유로 풀어보겠습니다.

---

1. 문제: "이미지"를 그리는 AI의 지루한 여정

AI 가 그림을 그릴 때는 처음엔 완전히 잡히지 않은 안개 같은 상태 (노이즈) 에서 시작해, 점점 선명하게 다듬어 나갑니다.

• 기존 방식: AI 는 매번 "지금 이 단계에서 무엇을 그려야 하지?"라고 계산하며 한 걸음씩 나아갔습니다. (매번 정밀한 측정을 하는 것)
• 기존 가속화 방법 (단점): "아, 지난번에 계산한 게 비슷하니까 그냥 그걸 다시 쓰자!"라고 생각해서 계산을 건너뛰는 방법들이 있었습니다. 하지만 이는 오래된 지도를 들고 가는 것과 같아서, 갈수록 길이 틀어지고 (Latent drift), 그림이 흐릿해지거나 이상한 모양이 생기는 문제가 있었습니다.

2. 해결책: "PrediT" - 미래를 예측하는 똑똑한 나침반

저자들은 AI 가 그리는 과정이 매우 부드럽게 변한다는 사실을 발견했습니다. 급격하게 튀는 게 아니라, 일정한 흐름을 타고 변한다는 거죠. 이를 이용해 과거의 흐름을 보고 미래를 예측하는 방식을 도입했습니다.

🚗 비유 1: "운전사"와 "내비게이션"

• 기존 AI: 매 1 초마다 정차해서 "지금 속도가 얼마야? 방향은 어디야?"를 확인하고 다시 출발합니다. (매우 안전하지만 느림)
• 단순 건너뛰기 (기존 방법): "아까 봤을 때 직진이니까 그냥 10 초 동안 직진할게!"라고 말하며 계산을 멈춥니다. 하지만 갑자기 커브가 나오면 그대로 벽에 부딪힙니다.
• PrediT (새로운 방법):
  1. 예측 (Predictor): "지난 3 초 동안의 속도와 방향을 봤더니, 앞으로 2 초는 거의 직진할 것 같아!"라고 수학적으로 미래를 예측합니다. (아담스 - 바스포스 방법)
  2. 수정 (Corrector): 만약 갑자기 커브가 나오거나 길이 험해지면 (동적인 변화가 크면), "잠깐! 예측이 틀릴 수 있으니 다시 한번 정밀하게 확인하자!"라고 수정을 해줍니다. (아담스 - 몰턴 방법)
  3. 적응 (Dynamic Modulation): 길이 평탄할 때는 예측만 믿고 빠르게 가고, 길이 험할 때는 신중하게 갑니다. 상황에 따라 속도를 조절하는 똑똑한 운전사입니다.

3. 결과: 얼마나 빨라졌나요?

이 방법을 적용한 결과, AI 는 그림을 그리는 속도가 최대 5.5 배 빨라졌습니다.

• 이미지 생성 (FLUX 모델): 약 4.3 배 빨라졌습니다.
• 영상 생성 (HunyuanVideo): 약 3.3 배 빨라졌습니다.
• 품질: 속도가 빨라졌지만, 그림의 선명도나 아름다움은 원래와 거의 똑같이 유지됩니다. (오히려 더 좋아진 경우도 있음)

4. 왜 중요한가요?

• 전기세 절약: 같은 작업을 하더라도 훨씬 짧은 시간에 끝나므로 에너지 소비가 줄어듭니다.
• 접근성: 고사양 그래픽 카드가 없어도, 일반 컴퓨터에서 고화질 영상이나 그림을 빠르게 만들 수 있게 됩니다. (기존 방법들은 메모리가 부족해서 멈추는 경우가 많았는데, PrediT 는 메모리 사용량을 거의 늘리지 않습니다.)

📝 한 줄 요약

"PrediT 는 AI 가 그림을 그릴 때, 매번 정밀하게 계산하지 않고 '흐름'을 읽어서 미래를 예측하고, 필요할 때만 다시 확인하는 '스마트한 가속 기술'입니다. 덕분에 AI 는 훨씬 빠르게, 하지만 화질은 그대로 유지하며 작품을 만들어냅니다."

기술 요약

1. 문제 정의 (Problem)

Diffusion Transformer (DiT) 는 현재 고품질 이미지 및 비디오 생성의 핵심 백본으로 자리 잡았으나, 다음과 같은 치명적인 한계가 존재합니다.

• 높은 계산 비용: 이차적인 (quadratic) 어텐션 비용과 반복적인 디노이징 (iterative denoising) 과정으로 인해 추론 지연 시간 (latency) 이 매우 길어 실시간 생성이 어렵습니다.
• 기존 가속화 방법의 한계:
  • 학습 기반 방법 (Distillation, Quantization): 학습 데이터와 컴퓨팅 자원이 많이 필요하며, 분포 변화 (distribution shift) 로 인해 생성 품질이 저하될 수 있습니다.
  • 학습 없는 캐싱 방법 (Training-free Caching): 기존 방법들 (DeepCache, FORA 등) 은 시간적 중복성을 활용하여 이전 단계의 특징 (feature) 을 재사용합니다. 그러나 이는 '특징이 단계 간에 변하지 않는다'는 가정에 기반합니다.
  • 잠재 공간 드리프트 (Latent Drift): 디노이징 궤적의 동적 변화가 큰 구간에서 단순 재사용을 반복하면 오차가 누적되어 시각적 왜곡 (artifacts) 이 발생하거나 품질이 급격히 떨어집니다. 이를 방지하기 위해 가속 비율을 낮게 설정해야 하므로 속도 향상 효과가 제한적입니다.

2. 방법론 (Methodology)

저자들은 디노이징 과정의 특징 궤적이 국소적으로 매끄럽게 (locally smooth) 변화한다는 관찰을 바탕으로, 단순 재사용이 아닌 원칙적인 예측 (principled prediction) 을 통해 가속화를 달성하는 프레임워크 PrediT 를 제안합니다.

핵심 구성 요소:

1. 선형 다단계 예측 (Linear Multistep Feature Forecasting):

   • 특징 추정을 미분방정식 (ODE) 의 수치 해법 문제로 재정의합니다.
   • Adams-Bashforth (AB) 예측기: 과거의 함수 값 (이전 단계의 모델 출력) 을 선형 결합하여 미래의 출력을 외삽합니다. 미분 추정 없이 직접 과거 값을 활용하여 수치적 안정성을 확보하며, 오차 (Local Truncation Error) 를 $O(\Delta t^{k+1})$ 수준으로 줄입니다.
   • Adams-Moulton (AM) 수정기 (Corrector): 고동적 (high-dynamics) 구간에서 예측 오차가 누적되는 것을 방지하기 위해 미래 값을 포함하는 암시적 (implicit) 수정기를 사용합니다.

2. 예측기 - 수정기 결합 (Predictor-Corrector Scheme):

   • AB (예측기): 변화가 완만한 구간에서 효율적으로 사용되어 많은 단계를 건너뜁니다.
   • ABM (예측기 + 수정기): 변화가 급격한 구간에서 정확도를 높이기 위해 사용됩니다. AB 로 예측한 후 실제 모델을 한 번 호출하여 수정합니다.

3. 동적 단계 조절 (Dynamic Step Modulation, DSM):

   • 디노이징 궤적의 동적 변화는 일정하지 않습니다 (초기/후기 구간은 급변, 중간 구간은 완만).
   • 상대적 변화율 ($\delta_n$) 을 모니터링하여 얼마나 많은 단계를 건너뛸 수 있는지 적응적으로 결정합니다.
   • 동작 로직:
     • 변화율이 임계값 ($\tau$) 이상이면: ABM 사용 (거의 건너뛰지 않거나 적게 건너뜀).
     • 변화율이 낮으면: AB 사용 (많은 단계를 건너뜀).
   • 이를 통해 고동적 구간에서는 오차 누적을 방지하고, 저동적 구간에서는 최대한의 가속을 달성합니다.

3. 주요 기여 (Key Contributions)

1. 새로운 관점 제시: 단순한 특징 재사용 (Naive Reuse) 이 잠재 공간 드리프트를 유발함을 분석하고, 확산 궤적의 국소적 매끄러움을 활용한 고차 다항식 예측의 필요성을 증명했습니다.
2. PrediT 프레임워크 제안: 학습이 필요 없는 (training-free) 프레임워크로, Adams-Bashforth 예측기와 Adams-Moulton 수정기를 결합하고 동적 단계 조절 메커니즘을 도입하여 안정적이고 적응적인 특징 예측을 가능하게 했습니다.
3. 성능 입증: 다양한 DiT 기반 모델 (이미지 및 비디오) 에서 기존 캐싱 및 예측 기반 방법들을 압도하는 성능을 보여주었습니다.

4. 실험 결과 (Results)

PrediT 는 FLUX.1, HunyuanVideo, DiT-XL/2 등 다양한 모델에서 광범위한 실험을 수행했습니다.

• 이미지 생성 (FLUX.1):
  • 속도 향상: 최대 4.28 배 (50 스텝 기준) 의 속도 향상 달성.
  • 품질: ImageReward, CLIP Score, Aesthetic Score 등 모든 지표에서 원본 모델 (50 스텝) 과 유사하거나 더 우수한 성능을 보임. 기존 방법들 (TeaCache, FORA 등) 은 가속 시 품질 저하가 명확했으나 PrediT 는 유지함.
• 비디오 생성 (HunyuanVideo):
  • 속도 향상: 최대 3.28 배 의 속도 향상.
  • 메모리 효율성: TaylorSeer 나 ProfilingDiT 와 같은 다른 예측 기반 방법들이 고해상도/고프레임 설정에서 OOM (Out-of-Memory) 오류를 일으킨 반면, PrediT 는 최소한의 메모리 오버헤드 (1-2% 증가) 만으로 실행 가능.
  • 화질: 프레임 간 모션 일관성과 공간적 디테일을 유지하며 블러 (blur) 나 아티팩트가 없음.
• 클래스 조건부 이미지 생성 (DiT-XL/2):
  • FID 점수 향상: 50 스텝 설정에서 FID 를 2.28 에서 2.24 로 개선 (더 낮은 값이 좋음). 이는 고차 예측 기법이 이산화 오차 (discretization error) 를 줄여 오히려 생성 품질을 높일 수 있음을 시사합니다.

5. 의의 및 결론 (Significance)

• 효율성과 품질의 균형: 기존 가속화 방법들이 '속도 vs 품질'의 트레이드오프를 강요했다면, PrediT 는 원칙적인 수치 예측 기법을 도입하여 두 마리 토끼를 모두 잡았습니다.
• 실용성: 추가 학습이 불필요하고 메모리 오버헤드가 적어, 소비자 등급 GPU 에서도 고해상도 비디오 생성이 가능해졌습니다.
• 환경적 영향: 추론 지연 시간 단축을 통해 대규모 확산 모델의 에너지 소비와 탄소 발자국을 줄이는 데 기여합니다.
• 미래 전망: 이 프레임워크는 Mixture-of-Experts (MoE) 아키텍처나 상호작용형 생성 시나리오 등 차세대 모델에도 확장 적용될 수 있는 잠재력을 가집니다.

요약하자면, PrediT 는 Diffusion Transformer 의 반복적인 디노이징 과정을 단순한 '재사용'이 아닌 '수치적 예측'으로 접근함으로써, 기존 방법들의 한계를 극복하고 고품질 생성을 유지하면서 획기적인 추론 속도 향상을 달성한 혁신적인 프레임워크입니다.