Distillation과 호환되는 학습 가능한 특징 캐싱으로 비디오 디퓨전 트랜스포머 가속화

초록

본 논문은 기존의 학습‑무료 특징 캐싱과 단계‑증류 방식을 결합한 새로운 가속 프레임워크인 DisCa를 제안한다. 경량의 학습 가능한 신경 예측기를 도입해 고차원 특징의 시간적 변화를 정밀히 예측하고, 대규모 비디오 디퓨전 모델에 적용 가능한 제한된 MeanFlow 기법으로 증류 안정성을 확보한다. 실험 결과 11.8배의 속도 향상을 달성하면서도 영상 품질 손실을 거의 보이지 않게 유지한다.

상세 분석

DisCa는 비디오 디퓨전 트랜스포머(ViDT)의 두 가지 주요 병목인 “샘플링 단계 수 감소”와 “디노이징 네트워크 연산량”을 동시에 완화한다. 기존의 특징 캐싱은 인접 타임스텝 간 특징이 거의 동일하다는 가정 하에 단순 재사용 혹은 선형/고차 보간을 적용했지만, 단계‑증류가 적용된 모델에서는 타임스텝 간 차이가 크게 증가해 캐시 오류가 급격히 누적된다. 이를 해결하기 위해 저자들은 “경량 신경 예측기(Predictor)”를 설계하였다. Predictor는 DiT의 중간 특징 맵을 입력으로 받아, 다중 이전 타임스텝(보통 2~3개)의 캐시와 타임스텝 임베딩을 결합해 다음 타임스텝의 특징을 직접 예측한다. 네트워크는 2개의 DiT 블록으로 구성되어 전체 DiT 파라미터 대비 4% 이하의 규모이며, 학습 과정에서는 기존 DiT와 예측된 특징을 모두 디스크리미네이터에 전달해 적대적 손실과 L2 손실을 동시에 최소화한다. 이렇게 하면 전통적인 수식 기반 예측이 포착하지 못하는 비선형·고차원 변화를 학습적으로 보정할 수 있다.

또한, 단계‑증류 자체의 불안정성을 완화하기 위해 “Restricted MeanFlow”를 도입한다. 기존 MeanFlow는 평균 속도 필드를 학습해 몇 단계만에 전체 샘플링 경로를 근사했지만, 대규모 비디오 모델에 적용하면 JVP 연산에서 발생하는 수치 오차와 과도한 압축 비율 때문에 훈련이 발산한다. Restricted MeanFlow는 (1) 평균 속도 계산 시 시간 구간을 제한하여 급격한 변화가 포함되지 않도록 하고, (2) 압축 비율이 높은 구간을 사전 pruning하여 학습 샘플을 정제한다. 결과적으로 평균 속도 추정이 더 보수적으로 이루어져, 증류 과정에서 발생할 수 있는 아티팩트와 불안정을 크게 감소시킨다.

실험에서는 텐센트 혼원·샤오동대 비디오 데이터셋을 사용해 기존 최첨단 방법(FORA, TaylorSeer, MeanFlow, Consistency 등)과 비교하였다. 16→2, 16→1 단계 압축 시 PSNR/SSIM 차이는 0.1dB 이하이며, 사용자 연구에서도 시각적 품질 차이가 거의 감지되지 않았다. 또한, 전체 추론 시간은 GPU A100 기준 11.8배 가속되었으며, 메모리 사용량도 약 30% 절감되었다. Ablation study에서는 Predictor의 깊이와 입력 타임스텝 수, Restricted MeanFlow의 구간 제한 비율이 성능에 미치는 영향을 정량적으로 분석해, 제안된 설계가 최적임을 입증한다.

한계점으로는 Predictor 학습을 위해 별도의 라벨이 필요하고, 매우 극단적인 압축(예: 1→0.5 단계)에서는 여전히 품질 저하가 관찰된다. 또한, 현재 구현은 DiT 기반 모델에 특화돼 있어 CNN 기반 디퓨전 모델에는 직접 적용하기 어렵다. 향후 연구에서는 멀티모달(텍스트·오디오) 디퓨전 모델에 대한 일반화와, 예측기와 MeanFlow를 공동 최적화하는 end‑to‑end 프레임워크 개발이 기대된다.