1. 문제 정의 (Problem)

확산 모델 (Diffusion Models), 특히 최신 Diffusion Transformer (DiT) 아키텍처는 고품질 이미지 및 비디오 생성에서 뛰어난 성능을 보이지만, 다단계 추론 (multi-step inference) 으로 인해 막대한 계산 비용이 소요됩니다. 이를 해결하기 위해 기존 연구들은 레이어 캐시나 토큰 캐시 (Token Caching) 기술을 도입하여 중간 특징을 재사용함으로써 추론 속도를 높였습니다.

그러나 기존 토큰 캐시 기반 방법론들은 다음과 같은 한계를 가지고 있습니다:

• 비효율적인 특징 캐싱 전략: 수동으로 설계된 희소성 (Sparsity) 할당 방식과 고정된 스케줄을 사용하여 최적의 가속 성능을 내지 못함.
• 불완전한 전방 계산 (Incomplete Forward Computation): 기존 방법들은 품질 저하를 방지하기 위해 특정 단계에서는 캐시를 사용하지 않고 완전한 전방 계산 (Full Forward Pass) 을 수행하도록 강제로 설계되어 있음. 이는 토큰 수준 캐싱의 잠재적인 가속 효과를 제한함.
• 수동 튜닝의 필요성: 각 레이어와 시간 단계별 재사용 비율을 수동으로 조정해야 하며, 이는 확장성과 성능에 부정적인 영향을 미침.

2. 제안 방법론: DiffSparse (Methodology)

이 논문은 DiffSparse를 제안합니다. 이는 확산 트랜스포머 모델의 레이어별 토큰 희소성을 학습 가능한 (Learnable) 방식으로 최적화하는 차분 가능 (Differentiable) 프레임워크입니다. 핵심 구성 요소는 다음과 같습니다.

가. 학습 가능한 희소성 비용 예측기 (Learnable Sparsity Cost Predictor)

• 모델은 모든 레이어와 탈노이즈 (Denoising) 단계에 대해 다양한 희소성 비율 (Sparsity Rates) 을 적용했을 때 발생할 '비용 (Quality Degradation)'을 예측하는 행렬을 학습합니다.
• 이 예측기는 토큰 시퀀스 길이 ($N$) 에 의존하지 않으며, 시간 단계 ($T$), 레이어 수 ($L$), 그리고 후보 희소성 집합 ($|S|$) 의 크기에만 의존하므로 고해상도 모델로의 확장성이 뛰어납니다.

나. 동적 프로그래밍 솔버 (Dynamic Programming Solver)

• 전체적인 희소성 제약 조건 (Target Sparsity Ratio) 을 만족하면서 총 비용을 최소화하는 최적의 레이어별 희소성 구성을 찾기 위해 동적 프로그래밍 (Dynamic Programming) 을 사용합니다.
• 이 과정을 통해 각 단계와 레이어에서 얼마나 많은 토큰을 재사용 (Reuse) 하고 얼마나 많은 토큰을 다시 계산 (Recompute) 할지 결정하는 이진 마스크 (Binary Mask) 를 생성합니다.

다. 토큰 선택기 (Token Selector)

• 학습된 희소성 비율에 따라 각 레이어에서 중요도가 높은 토큰을 동적으로 선택합니다.
• 토큰 중요도 점수는 자기 주의 (Self-attention), 교차 주의 (Cross-attention), 캐시 재사용 빈도, 공간적 분포 등을 기반으로 계산된 복합 점수를 사용하여 결정됩니다.

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

기존 방법들이 필요로 하는 '수정된 완전 전방 계산 단계'를 제거하고 효율성을 극대화하기 위해 도입된 전략입니다.

1. 1 단계: 완전 전방 계산이 필요한 단계 (Full-step positions) 와 레이어별 희소성 비용을 독립적으로 학습합니다.
2. 2 단계: 학습된 완전 전방 계산 비용을 레이어 희소성 비용에 통합 (Warm-start) 하여, 전체 샘플링 단계에 걸쳐 FLOPs 를 동적으로 재분배하도록 미세 조정합니다. 이 과정을 통해 완전 전방 계산 단계에 대한 의존성을 제거하고 순수한 캐시 기반 가속을 가능하게 합니다.

마. 학습 손실 함수

• 생성 품질 저하를 최소화하기 위해 LPIPS (Learned Perceptual Image Patch Similarity) 손실 함수를 사용하여 원본 모델 (Teacher) 과 희소화 모델 (Student) 간의 지각적 유사성을 최대화하도록 학습합니다.

3. 주요 기여 (Key Contributions)

1. DiffSparse 프레임워크: 수동 휴리스틱 없이 차분 가능한 희소성 비용 예측기, 동적 프로그래밍 솔버, 적응형 토큰 선택기를 통합하여 레이어별 토큰 희소성을 자동 최적화하는 방법을 제안했습니다.
2. 2 단계 학습 전략: 기존 방법들의 한계인 '수정된 완전 전방 계산' 필요성을 제거하여 토큰 수준 특징 캐싱의 가속 잠재력을 완전히 개방했습니다.
3. SOTA 성능 달성: 다양한 기반 모델 (DiT-XL/2, PixArt-α, FLUX, Wan2.1) 에 대한 광범위한 실험을 통해 기존 최첨단 (SOTA) 방법들 (ToCa, TaylorSeer 등) 을 크게 능가하는 효율성과 생성 품질을 입증했습니다.

4. 실험 결과 (Results)

• PixArt-α (텍스트 - 이미지 생성): 20 단계 샘플링 시, 54% 의 토큰을 제거하여 계산 비용을 54% 절감하면서도 FID 점수를 27.79로 유지했습니다. 이는 원본 모델 (FID 28.20) 보다 오히려 품질이 더 좋았으며, 기존 방법인 ToCa (28.35) 나 TaylorSeer (29.08) 보다 월등히 우수한 성능을 보였습니다. 속도 향상은 1.91 배를 달성했습니다.
• DiT-XL/2 (이미지넷 클래스 조건 생성): 2.07 배 가속 시 FID 를 3.05 (ToCa) 에서 2.81로 개선하여 세부 사항 보존 및 이미지 충실도를 증명했습니다.
• Wan2.1 (텍스트 - 비디오 생성): VBench 벤치마크에서 DuCa 대비 더 높은 가속비 (2.05 배) 를 유지하면서 최고 VBench 점수를 기록했습니다.
• 고해상도 일반화: 256x256 해상도에서 학습된 희소성 예측기가 512x512 해상도에서도 재학습 없이 효과적으로 작동하여 메모리 및 학습 비용을 절감할 수 있음을 보였습니다.

5. 의의 및 결론 (Significance)

DiffSparse 는 확산 모델의 추론 비용을 획기적으로 줄이면서도 생성 품질을 유지하거나 오히려 향상시킬 수 있는 새로운 패러다임을 제시합니다.

• 효율성: 수동 튜닝 없이 모델이 스스로 최적의 계산 자원을 할당하는 방식을 도입하여 실제 배포 (Edge/Mobile 등) 에 적합한 경량화 솔루션을 제공합니다.
• 품질 보장: 단순히 계산량을 줄이는 것을 넘어, 지각적 손실 (Perceptual Loss) 을 통해 품질 저하를 방지하고, 오히려 학습된 희소성 스케줄을 통해 생성 분포의 수렴을 가속화하여 품질을 개선하는 효과를 입증했습니다.
• 범용성: 다양한 DiT 아키텍처와 생성 작업 (이미지, 비디오) 에 적용 가능하며, 기존 캐싱 방법론들의 구조적 한계를 극복했습니다.

결론적으로, 이 연구는 확산 트랜스포머의 실용적 배포를 위한 핵심 기술인 '지능형 토큰 희소성 최적화'의 새로운 기준을 제시했습니다.