🤖 AI

LESA: Learnable Stage-Aware Predictors for Diffusion Model Acceleration

본 논문은 확산 모델의 계산 비용을 줄이면서도 고품질 생성을 유지하기 위해, 확산 과정의 단계별 특성을 학습하는 KAN 기반의 'LESA' 예측 프레임워크를 제안하고 텍스트 - 이미지 및 텍스트 - 비디오 생성 작업에서 기존 최첨단 방법보다 우수한 가속화와 품질 개선을 입증합니다.

원저자: Peiliang Cai, Jiacheng Liu, Haowen Xu, Xinyu Wang, Chang Zou, Linfeng Zhang

게시일 2026-03-17

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Peiliang Cai, Jiacheng Liu, Haowen Xu, Xinyu Wang, Chang Zou, Linfeng Zhang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

🎨 1. 문제: AI 화가는 왜 이렇게 느릴까요?

지금까지의 AI 그림 생성 기술 (확산 모델) 은 그림을 그릴 때 매우 천천히, 한 번에 한 번씩 작업을 반복합니다.

비유: 마치 거대한 캔버스에 그림을 그릴 때, 화가가 "이제 이 부분을 칠하고, 다음엔 저 부분을 칠하고..."라고 매번 모든 물감을 다시 섞고 붓을 담가서 한 번씩 칠하는 것과 같습니다.
결과: 그림이 아주 예쁘기는 하지만, 완성까지 시간이 너무 오래 걸려서 실생활에 쓰기 어렵습니다.

🚀 2. 기존 해결책의 한계: "그냥 복사" vs "예측"

속도를 높이기 위해 연구자들은 두 가지 방법을 썼습니다.

그냥 복사 (Reuse): "아까 그렸던 부분과 비슷하니까 그냥 복사해서 쓰자." -> 하지만 그림이 변할 때는 이 방법이 실패해서 그림이 뭉개집니다.
예측 (Forecasting): "앞으로 어떻게 변할지 수학 공식으로 예측해서 미리 그려보자." -> 하지만 AI 가 그리는 과정은 처음엔 거친 스케치, 중간엔 형태 잡기, 마지막엔 디테일 다듬기 등 단계마다 성격이 완전히 다릅니다. 하나의 공식으로 모든 걸 예측하려니 엉뚱한 그림이 나옵니다.

✨ 3. LESA 의 등장: "상황을 아는 똑똑한 예언자"

LESA 는 이 문제를 해결하기 위해 두 가지 핵심 아이디어를 도입했습니다.

① "단계별 전문가" (Stage-Aware Multi-Expert)

LESA 는 그림을 그리는 과정을 세 단계로 나누고, 각 단계에 맞는 전문가를 배치합니다.

초기 단계 (거친 스케치): 변화가 매우 빠르고 거칠 때 → 빠르게 반응하는 전문가가 담당.
중간 단계 (형태 잡기): 변화가 안정적이고 부드럽게 이어질 때 → 안정적인 전문가가 담당.
마지막 단계 (디테일 다듬기): 아주 미세한 수정이 필요할 때 → 정교한 전문가가 담당.
비유: 마치 건축 현장 같습니다. 기초를 다질 때는 중장비 크레인이 필요하고, 벽을 쌓을 때는 벽돌공이, 인테리어를 할 때는 세련된 디자이너가 필요합니다. LESA 는 이 상황에 맞춰 가장 적합한 전문가를 그 순간마다 투입합니다.

② "KAN"이라는 새로운 두뇌

이 전문가들이 사용하는 두뇌는 기존 AI 와 다릅니다. **KAN (콜모고로프 - 아르논드 네트워크)**이라는 기술을 썼는데, 이는 유연한 점토와 같습니다.

기존 AI (MLP): 미리 정해진 모양의 블록 (레고) 만으로만 쌓을 수 있어 복잡한 곡선을 표현하기 어렵습니다.
LESA 의 KAN: 점토처럼 스스로 모양을 변형하며 데이터에 맞춰 최적의 예측을 만들어냅니다. 그래서 복잡한 AI 의 움직임도 정확하게 따라잡을 수 있습니다.

🛠️ 4. 어떻게 훈련시킬까? (두 단계 학습)

LESA 를 가르치는 방법도 특별합니다.

1 단계 (정답 학습): 화가가 그린 '정답'을 보며 "이런 상황에서는 이렇게 변한다"는 것을 외웁니다.
2 단계 (실전 연습): 아까 배운 대로 스스로 예측을 하고, 그 예측을 바탕으로 다음 그림을 그리는 실전 연습을 합니다. 이렇게 하면 예측이 조금 틀려도 (오차가 생기더라도) 그 오차를 바로잡아 나가는 튼튼한 능력을 기릅니다.

🏆 5. 결과는 어떨까요?

LESA 는 기존 기술보다 압도적으로 빠르면서도 화질은 거의 유지합니다.

FLUX.1-dev: 약 5 배 빨라졌는데, 그림 품질은 1% 도 떨어지지 않았습니다.
Qwen-Image: 약 6.25 배 빨라졌고, 오히려 기존 최고 기술보다 화질이 20% 이상 좋아졌습니다.
HunyuanVideo (영상): 영상 생성도 5 배 빨라졌으며, 영상의 흔들림이나 왜곡이 훨씬 적습니다.

💡 요약

LESA 는 **"AI 가 그림을 그릴 때, 상황에 따라 가장 적합한 전문가를 투입하고, 유연한 두뇌로 미래를 예측하게 만든 기술"**입니다.

이전에는 "무조건 빨리 그리면 화질이 나빠진다"는 것이 상식이었는데, LESA 는 **"빠르면서도 예쁘게 그릴 수 있다"**는 새로운 가능성을 열었습니다. 이제 AI 가 그림이나 영상을 만드는 데 걸리는 시간이 획기적으로 줄어들어, 우리가 더 쉽게 고화질 콘텐츠를 즐길 수 있게 될 것입니다.

1. 연구 배경 및 문제점 (Problem)

배경: 확산 모델 (Diffusion Models), 특히 Diffusion Transformers (DiTs) 는 이미지 및 비디오 생성 분야에서 뛰어난 성능을 보이지만, 반복적인 노이즈 제거 (denoising) 과정과 깊은 트랜스포머 아키텍처로 인해 엄청난 계산 비용이 발생하여 실제 배포에 제약이 있습니다.
기존 방법의 한계:
- Feature Caching (특성 캐싱): 인접한 시간 단계 (timesteps) 간의 특성 (feature) 유사성을 활용하여 계산을 생략하거나 예측하는 방식입니다.
- Cache-Then-Reuse: 이전 단계의 특성을 그대로 재사용합니다. (단순 반복)
- Cache-Then-Forecast (예측 기반): Taylor 급수 확장 등 훈련 없는 (training-free) 기법을 사용하여 미래 특성을 예측합니다.
핵심 문제: 기존 예측 방법들은 확산 과정의 특성이 시간에 따라 매끄럽게 (smoothly) 변한다는 가정을 기반으로 합니다. 그러나 실제 분석 (Fig. 2) 에 따르면, 확산 과정은 단계별 (stage-dependent) 동역학을 가집니다.
- 초기 (고노이즈): 급격하고 불안정한 특성 변화.
- 중간: 안정적이고 연속적인 변화.
- 후기 (저노이즈): 미세한 디테일 재구성.
- 이러한 비균일한 (non-uniform) 동역학을 단일 고정된 예측 규칙으로 처리하면 생성 품질이 저하되고 시간적 일관성이 깨지는 문제가 발생합니다.

2. 제안 방법: LESA (Methodology)

저자들은 확산 과정의 단계별 특성을 학습 가능한 예측기로 모델링하는 LESA (Learnable Stage-Aware) 프레임워크를 제안합니다.

가. 핵심 아키텍처

Kolmogorov-Arnold Network (KAN) 기반 예측기:
- 기존 MLP 와 달리, KAN 은 각 뉴런에서 학습 가능한 스플라인 (spline) 기반의 일변수 함수를 사용하여 복잡한 비선형 시간적 동역학을 더 정밀하고 파라미터 효율적으로 학습합니다.
- 캐시된 과거 특성들과 시간 단계 (timestep) 정보를 입력받아 미래 특성을 예측합니다.
단계 인식 다중 전문가 구조 (Stage-Aware Multi-Expert Architecture):
- 확산 과정을 노이즈 수준에 따라 3 개의 단계로 나누고, 각 단계에 전용 예측기 (Expert) 를 할당합니다.
  - 고노이즈 전문가: 초기 급격한 변화 처리.
  - 중노이즈 전문가: 중간 안정적 변화 처리.
  - 저노이즈 전문가: 후기 미세 디테일 정제 처리.
- 각 단계마다 다른 시간 창 (Temporal Window, $K$ ) 을 적용하여 (예: 초기에는 짧은 창, 후기에는 긴 창) 단계별 특성에 최적화합니다.

나. 2 단계 학습 전략 (Two-Stage Training Strategy)

Ground-Truth Guided Training (GT-Guided Training):
- 실제 모델의 정답 (Ground Truth) 특성을 사용하여 예측기가 denoising 동역학을 학습하도록 지도 학습을 수행합니다.
Closed-Loop Autoregressive Training (CL-AR Training):
- 실제 추론 환경과 유사하게, 예측기의 이전 예측값을 입력으로 사용하여 다음 단계를 예측하는 순환 학습을 수행합니다.
- 이를 통해 빠른 추론 과정에서 발생할 수 있는 예측 드리프트 (prediction drift) 에 대한 강건성 (robustness) 을 확보합니다.

3. 주요 기여 (Key Contributions)

학습 가능한 시간 예측 프레임워크: 훈련 없는 (training-free) 방법의 한계를 극복하고, KAN 을 기반으로 확산 모델의 활성화 (activation) 에서 직접 시간적 동역학을 학습하는 효율적인 예측기를 도입했습니다.
단계 인식 다중 전문가 아키텍처: 단일 예측기가 아닌, 노이즈 수준별 단계에 특화된 전문가 예측기를 배치하여 각 단계의 고유한 특성 변화를 정밀하게 모델링했습니다.
2 단계 학습 전략: 정답 기반 학습과 폐루프 자기회귀 학습을 결합하여, 가속화 추론 중 발생하는 오차 누적 문제를 해결하고 안정적인 성능을 보장합니다.

4. 실험 결과 (Results)

LESA 는 이미지 (Text-to-Image) 및 비디오 (Text-to-Video) 생성 작업에서 다양한 모델 (FLUX.1, Qwen-Image, HunyuanVideo) 에서 SOTA(최고 성능) 를 달성했습니다.

FLUX.1-dev:
- 5.00 배 가속화 달성.
- 품질 저하 최소화 (ImageReward 1.0% 감소, PSNR 30.17 유지).
Qwen-Image:
- 6.25 배 가속화 달성.
- 이전 SOTA(TaylorSeer) 대비 20.2% 품질 향상 (ImageReward 기준).
HunyuanVideo (비디오):
- 5.00 배 가속화 달성.
- TaylorSeer 대비 24.7% PSNR 향상 및 시간적 일관성 유지.
기타: Stable Diffusion XL 등에서도 DeepCache 등 기존 방법 대비 우수한 품질과 속도 trade-off 를 보였습니다.

5. 의의 및 결론 (Significance)

확산 모델 가속화의 패러다임 전환: 단순한 재사용이나 고정된 수학적 근사 (Taylor 등) 를 넘어, 데이터 기반 학습 (Learnable) 을 통해 확산 과정의 복잡한 동역학을 모델링하는 새로운 접근법을 제시했습니다.
실용성: 고해상도 이미지 및 긴 비디오 생성에서도 높은 가속화 비율을 유지하면서 품질을 보존하여, 실제 애플리케이션에서의 확산 모델 배포 장벽을 낮춥니다.
일반화 능력: 다양한 아키텍처 (DiT 기반 모델) 와 작업 (이미지/비디오) 에 걸쳐 일관된 성능 향상을 보여주어, 제안된 프레임워크의 강력한 일반화 능력을 입증했습니다.

요약하자면, LESA는 확산 모델의 단계별 특성을 학습 가능한 KAN 기반 예측기로 정밀하게 포착함으로써, 기존 방법들의 한계를 극복하고 고속かつ고품질의 생성을 가능하게 한 획기적인 연구입니다.