캐시된 적응형 토큰 병합: 확산 모델에서의 동적 토큰 감소 및 중복 계산 제거

기본 정보

• 논문 ID: 2501.00946
• 제목: Cached Adaptive Token Merging: Dynamic Token Reduction and Redundant Computation Elimination in Diffusion Model
• 저자: Omid Saghatchian, Atiyeh Gh. Moghadam, Ahmad Nickabadi (Amirkabir University of Technology)
• 분류: cs.CV (컴퓨터 비전)
• 발표 시간: 2025년 1월 1일 (arXiv 사전 인쇄본)
• 논문 링크: https://arxiv.org/abs/2501.00946
• 코드 링크: https://github.com/omidiu/ca_tome

초록

확산 모델은 고품질의 고차원 이미지 생성을 위한 유망한 방법이 되었습니다. 그러나 이러한 모델은 높은 계산 비용과 느린 추론 속도로 인한 장애물에 직면해 있으며, 이는 부분적으로 자기주의 메커니즘의 입력 크기에 대한 이차 계산 복잡도 때문입니다. 본 논문은 토큰 간의 유사성을 계산하고 유사도가 임계값 매개변수 t보다 큰 토큰을 병합하여 이 문제를 해결하는 캐시된 적응형 토큰 병합(CA-ToMe) 방법을 제안합니다. 인접한 단계에서 관찰되는 반복 패턴과 유사성 빈도의 변화로 인해, 이 방법은 적응형 임계값을 구현하고 캐시 메커니즘을 추가하여 토큰 병합 방법을 강화합니다. 실험 결과는 이 방법이 훈련 없는 가속 방법으로서 제거 과정에서 1.24배의 가속을 달성하면서 기존 방법과 동일한 FID 점수를 유지함을 보여줍니다.

연구 배경 및 동기

문제 정의

확산 모델은 이미지 생성 작업에서 우수한 성능을 보이지만 심각한 계산 효율 문제에 직면해 있습니다:

1. 높은 계산 비용: 자기주의 메커니즘의 이차 복잡도로 인한 느린 추론 속도
2. 순차적 제거 과정: 병렬화 불가능, 각 제거 단계에서 반복 계산 필요
3. 중복 계산: 인접한 시간 단계 간의 대량의 반복 계산

문제의 중요성

• 확산 모델의 높은 지연 시간은 빠른 추론이 필요한 응용 프로그램에서의 사용을 제한합니다
• 높은 계산 비용으로 인해 모델 배포가 어려우며, 특히 리소스가 제한된 환경에서 그렇습니다
• 기존 가속 방법은 재훈련이 필요하거나 품질에 상당한 손실을 초래합니다

기존 방법의 한계

1. 샘플링 단계 감소 방법은 일반적으로 재훈련이나 복잡한 수치 해석기를 필요로 합니다
2. 토큰 가지치기 방법은 정보 손실과 성능 저하를 초래합니다
3. **기존 토큰 병합(ToMe)**은 고정 병합률을 사용하여 다양한 시간 단계와 계층의 유사성 분포 변화에 적응할 수 없습니다

연구 동기

관찰된 두 가지 핵심 현상에 기반합니다:

1. 다양한 시간 단계와 계층에서 토큰 유사성 분포의 상당한 변화
2. 인접한 추론 단계 간의 토큰 쌍의 높은 유사성

핵심 기여

1. 적응형 임계값 메커니즘 제안: 토큰 유사성 분포에 따라 병합 전략을 동적으로 조정하여 고정 병합률을 대체
2. 캐시 메커니즘 설계: 인접한 단계 간의 유사성을 활용하여 토큰 쌍을 캐시하여 반복 계산 감소
3. 훈련 없는 가속 구현: 방법을 사전 훈련된 모델에 직접 적용 가능하며 재훈련 불필요
4. 더 나은 품질-속도 균형 달성: 기본 ToMe 방법과 비교하여 이미지 품질을 유지하면서 더 빠른 추론 속도 실현

방법 상세 설명

작업 정의

입력: 확산 모델 제거 과정의 토큰 시퀀스 출력: 적응형 병합 및 캐시 최적화를 통한 가속 추론 과정 제약: 생성된 이미지 품질의 현저한 저하 없음

모델 아키텍처

1. 적응형 토큰 병합 메커니즘

기존 ToMe 방법은 고정 비율 r을 사용하여 토큰을 병합하는 반면, CA-ToMe는 유사성 임계값 t를 도입합니다:

핵심 개념:

• 이미지를 sx × sy 크기의 스트라이드 영역으로 분할
• 각 스트라이드 영역의 왼쪽 상단 토큰을 대상 토큰으로 선택
• 소스 토큰과 대상 토큰 간의 코사인 유사도 계산
• 임계값 t를 초과하는 유사도를 가진 토큰 쌍만 병합

장점 분석:

• 시나리오 A: 대부분의 토큰 유사도가 낮을 때, 고정 병합률은 유사하지 않은 토큰을 강제로 병합하여 정보 손실을 초래합니다. 적응형 임계값은 높은 유사도 토큰만 병합되도록 보장합니다
• 시나리오 B: 대부분의 토큰이 높은 유사도를 가질 때(제거 초기 등), 고정 병합률은 병합 수량을 제한합니다. 적응형 임계값은 더 많은 토큰 병합을 허용하여 효율성을 높입니다

2. 캐시 메커니즘 설계

Jaccard 거리 분석을 기반으로 인접한 단계 간 토큰 쌍의 높은 유사성 발견:

$JaccardDistance(A_n, A_{n+1}) = 1 - \frac{|A_n \cap A_{n+1}|}{|A_n \cup A_{n+1}|}$

여기서 An은 n번째 단계의 모든 소스-대상 토큰 쌍 집합을 나타냅니다.

구현 전략:

• 체크포인트(checkpoints) 설정, 특정 시간 단계에서만 유사성 행렬 계산
• 비체크포인트 단계에서 이전에 계산된 토큰 쌍 재사용
• 유사성 행렬의 반복 계산 오버헤드 현저히 감소

기술 혁신 포인트

1. 동적 적응성: 유사성 분포에 따라 병합 전략을 자동으로 조정하여 고정 매개변수의 한계 극복
2. 시간 차원 최적화: 시간 단계 간의 중복성을 활용하여 캐시를 통한 계산량 감소
3. 계층 선택적 적용: 계산 집약적인 U-Net 상단 계층(D1 및 U1)에 특별히 최적화 적용
4. 재훈련 불필요: 즉시 사용 가능한 가속 방법으로 기존 모델에 직접 적용 가능

실험 설정

데이터셋

• ImageNet-1k 데이터셋: 512×512 해상도의 2000개 이미지 생성(클래스당 2개, 총 1000개 클래스)
• 검증 집합: ImageNet-1k 검증 이미지 5000개를 사용하여 FID 점수 계산
• 프롬프트 템플릿: "A high-quality photograph of a classname."

평가 지표

1. FID (Fréchet Inception Distance): 생성된 이미지 품질을 측정하는 주요 지표
2. 추론 시간: 2000개 이미지 생성의 평균 시간
3. PSNR: 피크 신호 대 잡음 비, 픽셀 수준 재구성 품질 측정
4. SSIM: 구조 유사성 지수, 공간 및 구조 일관성 평가

비교 방법

• 기본선(Baseline): 원본 Stable Diffusion v1.5
• ToMe: 기존 토큰 병합 방법(r=50%)

구현 세부사항

• 하드웨어: Tesla V100S GPU
• 확산 단계: 50단계 PLMS 샘플링
• CFG 스케일: 7.5
• 스트라이드 크기: 고정 2×2
• 적용 계층: U-Net의 D1 및 U1 계층에만 적용

실험 결과

주요 결과

주요 발견:

• CA-ToMe는 가장 빠른 추론 속도(6.09초) 달성
• FID 점수(34.05)는 ToMe(34.16)보다 우수하고 기본선(33.66)에 근접
• 속도와 품질 간의 최적 균형 달성

임계값 매개변수 분석

관찰 결과:

• 임계값 0.4-0.6 범위에서 변화가 작음, 대부분의 토큰 유사도 ≥0.6이기 때문
• 임계값 0.7은 최적의 품질-속도 균형 제공
• 더 높은 임계값은 품질을 향상시키지만 속도를 감소시킵니다

캐시 구성 비교

구성 3이 최고 성능을 보이며, Jaccard 거리 분석과 일치하여 8, 11, 13단계 및 최종 단계에서 더 많은 체크포인트를 설정합니다.

절제 실험

다양한 구성 요소의 기여도 비교:

1. 적응형 임계값만: 고정 병합률과 비교하여 이미지 품질 향상
2. 캐시 메커니즘만: 계산 시간 현저히 감소
3. 완전한 CA-ToMe: 두 기술의 결합으로 최고 성능 달성

관련 연구

확산 모델 가속 방법

1. 샘플링 단계 감소:
   • 지식 증류 방법26,51,28
   • 암시적 샘플링32
   • 고급 미분 방정식 해석기52,33
   • 대부분 재훈련 필요
2. 단계별 계산 감소:
   • 양자화 방법31,36
   • 토큰 감소21,40,41,43,44
   • 캐시 기술24,37,38,39
   • 즉시 사용 가능, 재훈련 불필요

토큰 감소 기술

1. 토큰 가지치기: 중요하지 않은 토큰 직접 삭제, 정보 손실 가능성
2. 토큰 병합: 유사한 토큰 병합, 정보 완전성 유지
   • ToMe21: 고정 병합률 사용
   • 본 논문 CA-ToMe: 적응형 임계값 + 캐시 메커니즘

캐시 기술

기존 캐시 방법은 다양한 구성 요소를 대상으로 합니다:

• 교차 주의 캐시38
• U-Net 인코더 캐시39
• 고급 특성 캐시24

본 논문은 토큰 병합의 유사성 계산에 캐시를 처음 적용합니다.

결론 및 토론

주요 결론

1. 적응형 임계값은 고정 병합률의 한계를 효과적으로 해결하여 유사성 분포에 따라 병합 전략을 동적으로 조정합니다
2. 캐시 메커니즘은 시간 단계 간의 중복성을 활용하여 반복 계산을 현저히 감소시킵니다
3. CA-ToMe 방법은 1.24배 가속을 달성하면서 이미지 품질을 유지하거나 약간 향상시킵니다
4. 훈련 없는 특성은 방법에 우수한 실용성과 확장성을 제공합니다

한계

1. 임계값 매개변수 조정: 다양한 모델과 작업에 대해 최적 임계값을 조정해야 합니다
2. 적용 범위 제한: 주로 U-Net 아키텍처의 확산 모델을 대상으로 합니다
3. 캐시 오버헤드: 캐시된 토큰 쌍 정보 저장을 위한 추가 메모리 필요
4. 계층 제한: 상단 계층에만 적용하여 다른 계층의 최적화 기회를 놓칠 수 있습니다

향후 방향

1. 자동 임계값 학습: 최적 임계값을 자동으로 결정하는 방법 개발
2. 다른 아키텍처로 확장: DiT 등 새로운 확산 모델 아키텍처에 적응
3. 더 정교한 캐시 전략: 콘텐츠 기반 적응형 캐시 메커니즘
4. 하드웨어 최적화: 특정 하드웨어를 위한 최적화 구현

심층 평가

장점

1. 높은 혁신성: 적응형 개념을 토큰 병합에 도입하고 캐시 메커니즘과 결합하여 완전한 솔루션 형성
2. 높은 실용 가치: 훈련 없는, 즉시 사용 가능한 특성으로 배포가 용이합니다
3. 충분한 실험: 포괄적인 절제 실험과 매개변수 분석이 방법의 유효성을 뒷받침합니다
4. 견고한 이론적 기초: Jaccard 거리 기반 유사성 분석이 캐시 메커니즘에 이론적 지원을 제공합니다

부족한 점

1. 불충분한 이론 분석: 적응형 임계값 선택에 대한 이론적 지침 부족
2. 제한된 실험 범위: ImageNet에서만 검증, 다른 데이터셋과 작업의 평가 부족
3. 적은 비교 방법: 주로 ToMe와 비교, 다른 가속 방법과의 비교 부족
4. 단일 품질 평가: 주로 FID 지표에 의존, 인간 평가 및 다른 품질 지표 부족

영향력

1. 학술적 기여: 확산 모델 가속을 위한 새로운 아이디어와 방법 제공
2. 실용적 가치: 기존 확산 모델에 직접 적용 가능하며 광범위한 응용 전망
3. 재현성: 완전한 코드 구현 제공으로 재현 및 확장 용이
4. 영감: 적응형 및 캐시 개념이 더 많은 관련 연구에 영감을 줄 수 있습니다

적용 시나리오

1. 리소스 제한 환경: 모바일 장치, 엣지 컴퓨팅 등의 시나리오
2. 실시간 응용: 빠른 이미지 생성이 필요한 대화형 응용 프로그램
3. 대규모 배포: 서버 계산 비용 및 지연 시간 감소
4. 연구 프로토타입: 다른 가속 기술을 위한 기초 구성 요소

참고문헌

본 논문은 54개의 관련 문헌을 인용하며, 주로 다음을 포함합니다:

• 확산 모델 기초 이론1,2,3
• 이미지 생성 응용4,5,18,19,20
• 가속 기술24,25,26,27,28
• 토큰 처리 방법21,40,41,43,44
• 캐시 기술24,37,38,39

종합 평가: 이는 확산 모델 가속 분야에서 실용적 가치를 가진 연구입니다. 적응형 임계값과 캐시 메커니즘의 교묘한 결합을 통해 이미지 품질을 유지하면서 현저한 속도 향상을 달성합니다. 이론 분석과 실험 범위에서 개선의 여지가 있지만, 훈련 없는 특성과 우수한 실험 결과는 높은 실용 가치와 영향력을 제공합니다.