Conceptos Básicos
Diffusion Transformer 모델의 추론 속도를 높이기 위해 레이어 캐싱 메커니즘을 활용하는 Learning-to-Cache(L2C) 기법을 소개합니다. 이 기법은 레이어 간의 중복 계산을 활용하여 모델 파라미터 업데이트 없이도 상당 부분의 계산을 제거할 수 있습니다.
Resumen
Diffusion Transformer 모델 가속화를 위한 새로운 접근 방식: Learning-to-Cache
본 연구 논문에서는 Diffusion Transformer 모델의 추론 속도를 향상시키기 위한 혁신적인 방법인 Learning-to-Cache(L2C)를 제안합니다. Diffusion Transformer는 이미지 생성, 비디오 생성, 음성 합성, 3D 생성 등 다양한 분야에서 뛰어난 생성 능력을 보여주었지만, 높은 계산 비용과 느린 추론 속도라는 단점을 가지고 있습니다.
본 논문에서는 Diffusion Transformer의 계층적 구조와 Diffusion 모델의 순차적 특성을 활용하여, 시간 단계 간의 중복 계산을 제거하는 캐싱 메커니즘을 제안합니다. 특히, Transformer의 각 레이어를 캐싱의 기본 단위로 취급하여 시간 단계 간의 중복 계산을 효과적으로 활용합니다.
차별화 가능한 최적화 목표: 캐싱할 레이어를 식별하기 위한 기존의 exhaustive search 방식은 모델의 깊이에 따라 검색 공간이 기하급수적으로 증가한다는 문제점을 가지고 있습니다. 본 논문에서는 이러한 문제를 해결하기 위해 레이어 선택 문제를 차별화 가능한 최적화 문제로 변환하여 효율적인 학습을 가능하게 합니다.
입력 불변, 시간 단계 가변 라우터: 두 인접한 시간 단계 사이의 예측을 보간하여 빠른 구성(모든 레이어 캐싱)과 느린 구성(모든 레이어 유지) 사이의 최적의 캐싱 방식을 찾습니다. 이를 위해 시간에 따라 변하지만 입력과는 무관하게 동작하는 특수 라우터를 최적화하여 정적 계산 그래프를 생성합니다.
모델 파라미터 업데이트 불필요: 라우터는 모델 파라미터를 업데이트하지 않고 최적화되므로 계산 비용이 적고 최적화가 용이합니다.
ImageNet 데이터셋을 사용한 실험 결과, L2C는 DDIM, DPM-Solver와 같은 기존의 빠른 샘플러뿐만 아니라 기존의 캐시 기반 방법보다도 뛰어난 성능을 보여주었습니다. 특히 U-ViT-H/2 모델의 경우, 캐시 단계에서 최대 93.68%의 레이어를 캐싱할 수 있었으며, FID의 감소는 0.01 미만으로 나타났습니다.