2025-11-23T17:13:17.428108

Token Pruning for Caching Better: 9 Times Acceleration on Stable Diffusion for Free

Zhang, Xiao, Tang et al.
Stable Diffusion has achieved remarkable success in the field of text-to-image generation, with its powerful generative capabilities and diverse generation results making a lasting impact. However, its iterative denoising introduces high computational costs and slows generation speed, limiting broader adoption. The community has made numerous efforts to reduce this computational burden, with methods like feature caching attracting attention due to their effectiveness and simplicity. Nonetheless, simply reusing features computed at previous timesteps causes the features across adjacent timesteps to become similar, reducing the dynamics of features over time and ultimately compromising the quality of generated images. In this paper, we introduce a dynamics-aware token pruning (DaTo) approach that addresses the limitations of feature caching. DaTo selectively prunes tokens with lower dynamics, allowing only high-dynamic tokens to participate in self-attention layers, thereby extending feature dynamics across timesteps. DaTo combines feature caching with token pruning in a training-free manner, achieving both temporal and token-wise information reuse. Applied to Stable Diffusion on the ImageNet, our approach delivered a 9$\times$ speedup while reducing FID by 0.33, indicating enhanced image quality. On the COCO-30k, we observed a 7$\times$ acceleration coupled with a notable FID reduction of 2.17.
academic

토큰 프루닝을 통한 캐싱 개선: Stable Diffusion에서 9배 가속화 무료 달성

기본 정보

  • 논문 ID: 2501.00375
  • 제목: Token Pruning for Caching Better: 9× Acceleration on Stable Diffusion for Free
  • 저자: Evelyn Zhang, Bang Xiao, Jiayi Tang, Qianli Ma, Chang Zou, Xuefei Ning, Xuming Hu, Linfeng Zhang
  • 분류: cs.CV (컴퓨터 비전), cs.LG (머신러닝)
  • 발표 시간: 2024년 12월 31일
  • 논문 링크: https://arxiv.org/abs/2501.00375
  • 코드 링크: github.com/EvelynZhang-epiclab/DaTo

초록

Stable Diffusion은 텍스트-이미지 생성 분야에서 현저한 성공을 거두었으나, 반복적 노이즈 제거 메커니즘으로 인한 높은 계산 비용과 느린 생성 속도가 문제이다. 특징 캐싱 같은 방법들이 효과성과 단순성으로 주목받고 있지만, 이전 시간 단계에서 계산된 특징을 단순히 재사용하면 인접한 시간 단계의 특징이 유사해져 시간에 따른 특징의 동적성이 감소하고, 결국 생성된 이미지의 품질에 영향을 미친다. 본 논문은 동적 인식 토큰 프루닝(DaTo) 방법을 제안하여 특징 캐싱의 한계를 해결한다. DaTo는 동적성이 낮은 토큰을 선택적으로 제거하여 높은 동적성을 가진 토큰만 자기 주의 층에 참여하도록 함으로써 시간 단계 간 특징 동적성을 확장한다. ImageNet에 Stable Diffusion을 적용할 때 이 방법은 9배 가속화를 달성하면서 FID는 0.33 감소했고, COCO-30k에서는 7배 가속화와 FID 2.17의 현저한 감소를 관찰했다.

연구 배경 및 동기

문제 배경

확산 모델은 생성 모델링 분야에서 중대한 진전을 이루었으며, 텍스트-이미지 생성, 비디오 생성 등의 작업에 광범위하게 적용되고 있다. 그러나 확산 모델의 반복적 노이즈 제거 메커니즘은 막대한 계산 비용과 느린 생성 속도를 초래하여 더 광범위한 응용을 제한한다.

기존 방법의 한계

현재 확산 모델을 가속화하는 방법은 주로 다음을 포함한다:

  1. 샘플링 단계 감소: DDIM 등의 빠른 샘플러
  2. 각 단계의 계산 비용 감소: 지식 증류, 구조 프루닝, 양자화, 토큰 프루닝 및 특징 캐싱

이 중 특징 캐싱은 효과성과 단순성으로 널리 인기를 얻고 있으며, 이전 시간 단계에서 계산된 특징을 저장하고 후속 시간 단계에서 재사용한다. 그러나 특징 재사용은 서로 다른 시간 단계의 특징이 유사한 값을 갖도록 강제하여 시간 단계에 따른 특징의 동적성을 감소시키고, 원래의 확산 과정을 손상시켜 생성 품질을 저하시킨다.

연구 동기

논문은 실험 관찰을 통해 원래 Stable Diffusion과 비교할 때, 특징 캐싱을 사용하는 모델에서 인접한 시간 단계 간의 특징 차이가 현저히 감소함을 발견했다. 이는 핵심 질문을 제기한다: 특징 캐싱을 수행하면서도 여전히 올바른 특징 동적성을 유지할 수 있는가?

핵심 기여

  1. 동적 인식 토큰 프루닝(DaTo) 방법 제안: 서로 다른 시간 단계에서 특징 캐싱으로 인해 동적성이 감소한 토큰을 제거하고 높은 동적성을 가진 토큰으로 복구하여 특징 캐싱으로 인한 생성 품질 저하를 방지한다.
  2. 진화 탐색 전략 설계: 진화 방법을 통해 최적의 특징 캐싱 및 토큰 프루닝 전략을 탐색하여 DaTo의 잠재력을 충분히 발휘한다.
  3. 현저한 성능 향상 달성: Stable Diffusion 및 SDXL에 대한 광범위한 실험은 훈련 및 추가 데이터 없이 Stable Diffusion에서 최대 9배 가속화와 생성 품질 손실 없음을 보여준다.

방법 상세 설명

작업 정의

본 논문의 작업은 이미지 생성 품질을 유지하면서 Stable Diffusion 모델의 추론 과정을 현저히 가속화하는 것이다. 입력은 텍스트 프롬프트이고, 출력은 해당하는 고품질 이미지이며, 제약 조건은 모델을 재훈련할 필요가 없다는 것이다.

모델 아키텍처

1. 동적 인식 토큰 프루닝(DaTo)

기본 토큰 선택:

  • 시간 노이즈 차이 점수: t번째 시간 단계에 대해 이전 두 인접 시간 단계 출력의 절댓값 차이를 계산한다: 
    DiffScore = (1/C) * Σ|f_up_0(x_{t+2})_c - f_up_0(x_{t+1})_c|
  • 패치 기반 토큰 선택: 이미지를 s×s의 겹치지 않는 패치로 나누고, 각 패치에서 가장 높은 DiffScore를 가진 토큰을 기본 토큰으로 선택한다.

CFG 정렬: 분류기 자유 유도(CFG)를 처리하기 위해 조건부 생성의 기본 토큰 위치를 무조건부 생성에 복사한다:

X_base,i,j[k] = X_base,i,j[k - B/2], k ∈ {B/2, B/2+1, ..., B-1}

프루닝 토큰 선택: 코사인 유사도를 기반으로 기본 토큰과 가장 유사한 K개의 토큰을 프루닝 대상으로 선택한다:

X_prune = arg topK max Cosine_Similarity(X_i, X_j)

프루닝 토큰 복구: 프루닝된 토큰을 가장 유사한 기본 토큰을 직접 복사하여 복구한다.

2. 시간 단계 인식 특징 캐싱

탐색 공간 프루닝:

  • 캐싱 깊이 d는 {0, 1, 1/2}로 제한
  • 프루닝 비율 r은 {0.3, 0.4, 0.5, 0.6, 0.7}로 제한

진화 탐색 알고리즘: NSGA-II 다목적 최적화 알고리즘을 사용하여 다음을 최적화한다:

  • 추론 지연
  • 생성 품질(FID)

탐색 과정은 선택, 교배, 변이 등의 표준 진화 연산을 포함하며, 최종적으로 최적의 단계 인식 전략 F(t)를 얻는다.

기술 혁신점

  1. 동적성 복구 메커니즘: 낮은 동적성을 가진 토큰을 선택적으로 제거하고 높은 동적성을 가진 토큰으로 복구하여 특징 캐싱으로 인해 손상된 특징 동적성 분포를 성공적으로 복구한다.
  2. 통합된 캐싱-프루닝 프레임워크: 특징 캐싱과 토큰 프루닝을 훈련 무관 프레임워크에 결합하여 시간 및 토큰 수준의 정보 재사용을 실현한다.
  3. 적응형 전략 탐색: 서로 다른 시간 단계의 서로 다른 중복 특성에 대해 최적의 캐싱 깊이 및 프루닝 비율을 자동으로 탐색하는 방법을 제안한다.

실험 설정

데이터셋

  • ImageNet-1k: 2000개의 512×512 이미지 생성(클래스당 2개)
  • COCO-30k: 30000개의 이미지 생성(캡션당 1개)
  • MS COCO 검증 세트: SDXL 평가에 사용, 5k개의 1024×1024 이미지 생성

평가 지표

  • FID (Fréchet Inception Distance): 생성 품질 측정
  • CLIP Score: 텍스트-이미지 정렬 평가
  • Inception Score: 이미지 품질 평가
  • 지연 및 가속 비율: 효율성 평가

비교 방법

  • DDIM/DPM: 빠른 샘플러
  • ToMeSD: 토큰 병합 방법
  • DeepCache: 특징 캐싱 방법
  • DeepCache & ToMeSD: 단순 결합 방법

구현 세부사항

  • NSGA-II 진화 알고리즘 사용, 모집단 크기 20, 100세대 실행
  • CFG scale: 7.5 (SD v1.5), 9.0 (SD v2), 7.0 (SDXL)
  • 샘플링 단계: 50단계 PLMS
  • 단일 4090 GPU에서 테스트

실험 결과

주요 결과

Stable Diffusion v1.5 (ImageNet):

  • 구성 e1: 9.01배 가속화, FID 27.64에서 27.31로 감소
  • 모든 구성에서 비교 방법을 능가

Stable Diffusion v2 (ImageNet):

  • 구성 e2: 7.25배 가속화, FID 28.20
  • 원래 모델 대비 FID 29.8에서 28.20으로 감소

COCO-30k 데이터셋:

  • SD v1.5: 7배 가속화, FID 12.15에서 9.98로 감소(2.17 감소)
  • SD v2: 7.25배 가속화, FID 13.68에서 13.88로

SDXL (MS COCO):

  • 2.32배 가속화, FID 24.25에서 23.10으로 감소
  • DeepCache(1.75배) 및 DeepCache&ToMeSD(1.78배)보다 현저히 우수

소거 실험

DiffScore의 효과성: 서로 다른 캐싱 설정 및 프루닝 비율에서 DiffScore를 사용하면 일관되게 FID 점수를 향상시켜 시간 노이즈 차이 점수의 효과성을 입증한다.

CFG 정렬의 영향: 프루닝 비율이 증가함에 따라 CFG 정렬 구성으로 인한 이득이 점진적으로 증가하며, 높은 프루닝 비율(0.7)에서 FID 개선은 13포인트에서 30포인트까지 다양하다.

사례 분석

시각적 비교 결과는 DaTo가 여러 측면에서 우수한 성능을 보여준다:

  • 내용 충실도: 원본 이미지 내용과 높은 유사성
  • 세부 사항 보존: 고세부 장면에서 미세한 질감 유지
  • 스타일 적응: 이미지-이미지 작업에서 내용 보존과 스타일 정확성 균형
  • 프롬프트 정렬: 복잡한 텍스트 프롬프트의 모든 요소 정확하게 생성

실험 발견

  1. 특징 동적성 복구: DaTo는 특징 차이 분포를 원래 Stable Diffusion 수준에 가깝게 성공적으로 복구
  2. 희소 인코딩 효과: 적절한 토큰 프루닝 및 특징 캐싱은 핵심 특징에 집중함으로써 모델 성능을 개선할 수 있음
  3. 전략 일반화성: SD v1.5에서 탐색한 전략이 SDXL 및 다른 데이터셋에서 좋은 성능을 발휘

관련 연구

효율적인 확산 모델

  • 샘플링 단계 감소: DDIM, 일관성 모델 등
  • 네트워크 압축: 양자화, 프루닝, 증류 등
  • 아키텍처 최적화: U-Net 개선, Transformer 최적화 등

토큰 감소 전략

  • 학습 방법: DynamicViT, A-ViT 등 보조 모델을 사용한 정렬 및 프루닝
  • 휴리스틱 방법: Token Pooling, Token Merging 등 훈련 무관 방법
  • 확산 모델 응용: ToMeSD, AT-EDM 등 생성 작업에 맞춘 적응

캐싱 메커니즘

  • U-Net 캐싱: DeepCache는 시간 중복성을 활용한 특징 캐싱
  • DiT 캐싱: Δ-DiT는 Diffusion Transformer에 대한 캐싱 전략
  • 최적화 과제: 효율성 향상과 생성 품질 유지 간의 균형

결론 및 논의

주요 결론

  1. DaTo는 특징 캐싱으로 인한 특징 동적성 손실 문제를 성공적으로 해결
  2. 진화 탐색으로 얻은 적응형 전략이 고정 구성을 현저히 능가
  3. 이 방법은 여러 모델 및 데이터셋에서 현저한 가속화 및 품질 향상을 달성

한계

  1. 탐색 비용: ≤20 GPU 시간이 수용 가능하지만 여전히 추가 계산 자원 필요
  2. 하드웨어 의존성: 성능 향상은 하드웨어 구성에 따라 달라질 수 있음
  3. 극단적 설정에서의 한계: 과도한 프루닝 비율 또는 과도하게 낮은 캐시 업데이트 빈도는 성능을 손상시킴

향후 방향

  1. 적응형 전략 학습: 더 지능적인 적응형 캐싱 및 프루닝 전략 개발
  2. 다른 아키텍처 적응: 더 많은 확산 모델 아키텍처로 확장
  3. 이론적 분석: 희소 인코딩 원리가 확산 모델에서 작용하는 방식에 대한 깊이 있는 이해

심층 평가

장점

  1. 높은 혁신성: 특징 캐싱의 동적성 손실 문제를 체계적으로 해결한 첫 시도
  2. 실용적 방법: 훈련 무관, 배포 및 통합이 용이
  3. 충분한 실험: 여러 모델, 데이터셋에 대한 포괄적 평가
  4. 이론적 지원: 희소 인코딩에 대한 이론적 설명 제공
  5. 오픈소스 친화적: 완전한 코드 구현 제공

부족한 점

  1. 이론적 분석 부족: 이 방법이 FID를 개선하는 이유에 대한 이론적 설명이 상대적으로 단순
  2. 탐색 알고리즘 의존성: 최적 전략을 찾기 위해 진화 탐색이 필요하여 사용 복잡도 증가
  3. 평가 지표 단일성: 주로 FID에 의존하며 더 다양한 품질 지표 부족
  4. 사용자 연구 부재: 생성 품질을 검증하기 위한 인간 평가 없음

영향력

  1. 학술적 가치: 확산 모델 가속화에 새로운 사고와 방법 제공
  2. 실용적 가치: 기존 Stable Diffusion 모델에 직접 적용 가능
  3. 재현성: 상세한 구현 세부사항 및 오픈소스 코드 제공
  4. 영감: 생성 모델에서 토큰 수준 최적화 응용에 대한 사례 제시

적용 시나리오

  1. 자원 제한 환경: 모바일 기기, 엣지 컴퓨팅 시나리오
  2. 실시간 응용: 빠른 생성이 필요한 대화형 응용
  3. 대량 생성: 대규모 이미지 생성 작업
  4. 연구 프로토타입: 빠른 반복이 필요한 연구 프로젝트

참고문헌

논문은 46개의 관련 문헌을 인용하며, 확산 모델, 토큰 감소, 캐싱 메커니즘 등 여러 관련 분야의 중요한 연구를 포함하여 본 연구에 견고한 이론적 기초와 비교 기준을 제공한다.

종합 평가: 이는 확산 모델 가속화라는 중요한 문제에 대해 혁신적인 해결책을 제시한 고품질의 컴퓨터 비전 논문이다. 방법 설계가 정교하고, 실험 평가가 포괄적이며, 실용적 가치가 뛰어나다. 이론적 분석 깊이에서 다소 부족하지만, 그 실제 기여와 영향력은 인정할 만하다.