주파수별 칼만 필터로 디퓨전 트랜스포머 추론을 6배 가속한다

무엇인가

디퓨전 트랜스포머(DiT)는 이미지·영상 생성 품질을 크게 끌어올렸지만, 수십 번의 디노이징 스텝마다 대형 트랜스포머를 다시 통과시켜야 해서 추론 지연이 크다. 기존의 feature caching은 인접 스텝의 특징이 비슷하다는 점을 이용해 이전 스텝의 특징을 재사용하지만, 스텝 간격이 벌어질수록 유사도가 무너져 가속비를 조금만 올려도 품질이 급락한다. 최근의 forecasting 계열(TaylorSeer, Spectrum 등)은 재사용 대신 과거 관측으로 다음 스텝의 특징을 외삽해 더 긴 스킵을 노린다. 이 논문은 그 forecasting 방식이 두 가지를 균일하게 가정한다는 점을 문제로 잡는다. 첫째, 예측 대상인 모델 출력의 모든 성분에 같은 규칙을 적용하는데 실제로는 저주파 성분이 초반 몇 스텝에서 급격히 변하고 끝으로 갈수록 정체되는 반면, 고주파 성분은 초반에 잠잠하다가 후반에 빠르게 움직인다. 둘째, 샘플러의 스텝 인덱스나 원시 타임스텝을 시간 좌표로 쓰는데, 스케줄러에 따라 스텝당 실제 디노이징량이 균일하지 않아 같은 간격이 전혀 다른 변화량을 의미하게 된다.

어떻게 동작하나

SpecKCast의 핵심은 예측 대상을 중간 활성값이 아니라 샘플러가 최종적으로 소비하는 velocity(모델 출력)로 두고, 이를 특징 공간이 아닌 공간-주파수 영역에서 다루는 것이다. 각 스텝의 출력 Y_j에 채널별로 직교 2D DCT를 적용해 Z_j를 얻고, 주파수 평면을 겹치지 않는 방사형 대역 L=7개로 나눈다. 각 계수는 독립된 시계열로 취급되어 자기 값을 따로 추적하지만, 시간 좌표·공분산·프로세스 노이즈 통계는 같은 대역 안에서 공유한다. 이 덕분에 계수 수만 개를 추적하면서도 공분산 행렬은 대역 수만큼(7개)만 유지하면 된다. 스킵 스텝에서는 예측된 계수들을 모아 역 DCT로 공간 영역에 되돌린 뒤 샘플러에 그대로 넘긴다.

무엇과 다른가

시간 좌표는 로그 신호대잡음비 λ_j = 2log((1-σ_j)/σ_j)에서 출발하되, 대역마다 따로 워핑한다. t_ℓ(j) = sigmoid((λ_j + log p_ℓ)/T_w)로 각 대역의 활성 구간을 잡고 이를 [-1,1]로 재스케일해 모든 대역이 같은 범위를 쓰도록 맞춘다. p_ℓ는 그 대역의 보정된 신호 전력, T_w는 활성 창의 폭이다. 결과적으로 대역이 실제로 많이 변하는 구간에서는 좌표가 빠르게 진행하고 안정된 구간에서는 느리게 진행해, 국소 외삽의 정확도가 올라간다. 이 변환은 예측 좌표만 바꿀 뿐 디퓨전 과정 자체는 건드리지 않는다.

어떻게 쓰나

각 계수는 칼만 필터로 추적되며 상태는 (값 m, 시간 미분 ν) 두 개다. 매 샘플러 스텝마다 m은 Δ·ν만큼 전진하고 ν는 ρ = exp(-Δ/τ)로 감쇠한다(τ=0.1). 이 감쇠가 오래된 변화율의 영향을 줄여 긴 스킵에서의 오버슈트를 막는다. 전이행렬 A와 프로세스 노이즈 Q는 이 등속+감쇠 동역학에 맞춰 구성되고, 관측이 있는 계산 스텝에서만 칼만 업데이트가 일어난다. 스킵이 길어지면 공분산이 커져 예측 불확실도가 자연히 증가한다.

전제와 한계

스킵 스케줄은 고정하지 않는다. 대역별 공분산의 좌상단 원소(계수 추정 오차의 분산)에 대역 계수 수 d_ℓ를 곱해 합산하고, 후반부 위험을 눌러주는 가중치 ω(j)=exp(-β·j/(N-1))를 곱해 스칼라 위험도 R(j)를 만든다. 계산 스텝 이후 공분산을 앞으로 전파해 미래 스텝의 위험도를 예측하고, R(k)가 임계값 ε을 처음 넘는 스텝에서 트랜스포머를 다시 평가한다. ε이 작으면 자주 계산하고, 크면 더 길게 건너뛴다.

실험은 이미지 모델 FLUX.1-dev(1024×1024)와 Qwen-Image(1328×1328), 영상 모델 Wan2.1-14B와 HunyuanVideo에서 수행했고, 50스텝 완전 계산을 기준으로 삼았다. 비교 대상은 FORA, TeaCache, HiCache, SeaCache, TaylorSeer, Spectrum이다. FLUX.1-dev에서 5.54배 가속 시 가장 강한 베이스라인보다 PSNR이 1.1dB 높았고, Qwen-Image에서는 6.00배 가속에서 2.0dB 높았다. 같은 속도로 스텝만 줄인 Base 대비로는 PSNR이 5.5~9.0dB 개선됐는데, 이는 단순히 스텝을 줄이는 것보다 50스텝 궤적을 훨씬 잘 따라간다는 뜻이다. Qwen-Image 8 NFE에서 Spectrum의 LPIPS는 0.527까지 올라간 반면 SpecKCast는 0.267을 유지했다. 메모리 오버헤드는 FLUX 기준 최대 0.01GiB로, TaylorSeer·HiCache·Spectrum이 쓰는 0.73~18.97GiB와 대비된다. 영상 쪽에서는 HunyuanVideo 14 NFE에서 SeaCache보다 PSNR 2.1dB 높으면서 2.57배가 아닌 3.33배로 돌았고, Wan2.1-14B에서는 모든 예산에서 VBench Quality가 가장 좋았으며 14 NFE에서 기준 모델과 0.25점 이내에 머물렀다. Spectrum이 Wan2.1-14B 14·10 NFE에서 LPIPS를 근소하게 앞서는 구간이 있지만 최대 13.12GiB의 추가 메모리를 요구한다.

FLUX.1-dev 14 NFE에서의 절제 실험도 방법의 각 요소가 실제로 기여함을 보여준다. 고정 스케줄 아래에서 칼만과 테일러 예측기는 비슷했고(칼만이 약간 빠르고 PSNR 조금 높음), 체비쇼프는 더 느리고 부정확했다. 흥미롭게도 주파수 분해와 시간 워핑만 갖추면 테일러 예측기라도 TaylorSeer보다 PSNR이 4.2dB 높고 지연이 4.59초에서 3.46초로 줄었다. 즉 성능 향상의 상당 부분은 예측기 종류가 아니라 무엇을 어떤 좌표에서 예측하는가에서 나온다. 불확실도 기반 스케줄링은 고정 스케줄 대비 2.2dB, 대역별 시간 워핑은 원시 log-SNR 대비 3.3dB, 스텝 인덱스 대비 4.7dB를 더했고, 주파수 분해를 제거하면 3.1dB가 떨어졌다. 모든 변형이 체비쇼프를 제외하면 0.1초 이내 차이라 이득은 사실상 공짜에 가깝다.

개발자 관점에서 이 방법의 매력은 학습이 필요 없고 추가 메모리가 수 MiB 수준이라는 점이다. 모델 가중치를 건드리지 않고 샘플러 루프 바깥에서 출력을 가로채 DCT·역 DCT와 소형 칼만 필터 몇 개를 돌리는 구조라, FLUX나 Wan 계열 체크포인트를 그대로 두고 래퍼만 끼워 지연을 줄이고 싶은 경우가 1차 적용처다. 다만 도입 전에 확인할 것이 있다. 임계값 ε을 목표 NFE 예산에 맞춰 캘리브레이션해야 하고, τ·T_w·β·대역 수 L 같은 하이퍼파라미터가 모델별로 조정되어야 한다(논문은 L=7, τ=0.1, T_w=2.0, β=0.6을 기본값으로 쓴다). 또 평가가 50스텝 기준과의 재현도(PSNR/SSIM/LPIPS) 중심이므로, 실제 서비스에서 문제가 되는 프롬프트 다양성이나 극단적 프롬프트에서의 실패 모드는 별도로 확인해야 한다.

저자들이 명시한 한계는 명확하다. SpecKCast는 디노이징 스텝 사이의 시간적 중복만 활용하며, 계산하는 스텝에서는 여전히 전체 토큰 그리드를 처리하므로 트랜스포머 한 번의 평가 안에 존재하는 공간적 중복은 건드리지 못한다. 시간적·공간적 중복을 함께 활용하는 것이 향후 과제로 남아 있다는 것이 논문의 마지막 문장이다.