[논문 리뷰] Temporal Dynamic Quantization for Diffusion Models

Introduction

• Diffusion models는 GAN과 같은 기존의 모델들에 비해 훨씬 더 우수한 성능을 보였지만, 높은 비용을 수반함
  ⇒ 리소스가 제한적인 모바일 환경에서는 적용 어렵고, 고성능의 서버에만 적용 가능
  ⇒ 따라서 성능을 보전하면서도 계산 비용과 메모리를 줄일 수 있는 방안에 대한 연구 ↑
  ⇒ 전에 비해서는 적은 sampling steps를 이용하는 등 개선되긴 했지만, 여전히 높은 계산 비용과 방대한 메모리가 수반
  ⇒ performance enhancements & model compression의 필요성
• 기존의 연구들: reducing the # of sampling steps 에 초점
  but 개별적인 denoising steps를 가볍게 하는 것도 중요함.
  • Single denoising step ≈ a conventional deep learning model inference
    ⇒ 다양한 압축 방법이 사용될 수 있음. (ex. Quantization, Knowledge Distillation 등)
  • 그 중에서 자주 쓰이는 기법인 Quantization의 문제점 = significant performance degradation
  • QAT & PTQ: 일반적으로 고정된 데이터 분포를 가진 네트워크에서 잘 작동 but 확산 모델 같은 경우에는 시간에 따라 변하는 활성화 분포를 가져서 적용하기 어려움.
  • cf) 확산 모델에서 양자화 시 성능이 하락하는 주요한 원인은 뉴럴 네트워크의 각 레이어의 활성화 값의 분포가 시간에 따라 크게 차이가 나기 때문 ⇒ 이러한 분포차이에 대응하기 위해, 확산 모델의 시간 정보를 입력으로 하여 각 시간의 양자화 간격을 학습할수 있는 네트워크를 두는 방법 제안

• Temporal Dynamic Quantization (TDQ) 모듈 제안
  • generates a time-dependent optimal quantization configuration that minimizes activation quantization errors
  • 기존의 QAT와 PTQ 알고리즘의 통합 및 확장 → activation quantization error 최소화하는 시간 의존적인 최적의 양자화 구성 생성
  • 추론 중 계산 오버헤드 발생시키지 않도록 설계 ⇒ 기존의 가속 프레임워크와 호환 & 수정 없이 사용 가능
  • 각 time step에서 최적의 양자화 매개변수 생성

---

2. Backgrounds and Related Works

2.1. Diffusion Model

• 확산 모델은 고품질 이미지를 생성할 수 있지만, 반복적인 denoising 과정 때문에 실제 환경에서 사용하기 어려움
• DDPM
• DDIM
• Latent Diffusion Model(LDM): 이미지 대신 잠재 변수 사용 대규모 확산 모델(ex. Stable Diffusion)은 LDM을 활용하고 LAION-Dataset과 같은 대규모 데이터셋에서 학습을 통해 테스트 입력에 따른 고화질, 고해상도 이미지 생성 가능하게 함

2.2. Quantization

• Quantization: a prominent neural network optimization technique that reduces storage requirements with low-precision representation and performance improvement when the corresponding low- precision acceleration is available
• 저장 요구량을 줄이고 해당 저정밀 표현을 사용할 때 성능을 개선하는 신경망 최적화 기술
• 양자화된 모델 표현의 자유도 제한: $b$ 비트 정밀도 사용 시 오직 $2^b$ 의 양자화 레벨만 가능하기 때문
• 최종 품질 ← the tuning of quantization hyperparameters (such as the quantization interval or zero offset)

2.3. Quantization-aware Training and Post-training Quantization

• Quantization-aware Training (QAT)
  • 양자화 연산자를 도입 → 양자화의 효과를 고려하여 최종 손실 값을 최소화하도록 네트워크 매개변수 업데이트 허용 → 추가 훈련(additional training)
• Post-training Quantization (PTQ)
  • 양자화 이후 end-to-end forward/backward propagation 적용X
  • 대신 양자화에 의해 발생하는 block-wise reconstruction errors를 줄이는 것에 집중
• QAT와 PTQ: static quantization에 집중 → 최근 연구: the benefits of input-dependent dynamic quantization을 중요시함.
• Dynamic Quantization
  • input-dependent distribution of activations 의 변화에 따라 양자화 간격을 조정 ⇒ 양자화 오류↓
  • 하지만 activations로부터 통계학적 정보를 추출하는 과정에서 추가 비용 발생 가능
• 연구 목적 및 의의:
  추가적인 추론 비용을 발생시키지 않으면서, 시간 단계 정보에 기반하여 적절한 양자화 구간의 생성을 가능하게 함으로써 활성화 양자화 오류 최소화

3. Temporal Dynamic Quantization

3.1. Quantization Methods

• $b$ 비트 선형 양자화에 중점, 여기서 가능한 $2^b$ 개의 양자화 수준은 균등하게 배치
• 선형 양자화는 두 가지 주요 하이퍼파라미터인 양자화 간격 $s$와 제로 오프셋 $z$를 포함
• 전체 정밀도 데이터 $x$ 가 주어지면, 양자화된 데이터 $\hat{x}$ 는 다음과 같이 계산

• n(p): 가장 작은(큰) 양자화 인덱스
• clip( ): 클리핑 함수
• [ ]: 반올림 함수

• conv_in.0:
  • 활성화 값이 비교적 작은 범위 내에서 변동 (-4 ~ +4)
  • 활성화 값은 시간이 지남에 따라 점차적으로 감소하는 경향
• up.1.upsample.conv.0
  • 훨씬 넓은 범위의 활성화 값 (-100 ~ +200).
  • 시간이 지남에 따라 활성화 값이 감소하는 경향을 보이나, 변동 범위가 훨씬 큼

정적 양자화의 한계: 확산 모델에서 활성화 분포는 시간 단계가 진행됨에 따라 변화하는데, 정적 양자화는 이러한 동적 변화를 적절히 처리하지 못하는 문제

⇒ 활성화 값의 변화가 큰 확산 모델에서는 각 시간 단계에 따라 양자화 간격을 조정하는 동적 양자화 방식이 더 적합할 수 있음

 

 

3.2. Challenges of Diffusion Model Quantization

양자화 오류를 최소화하는 활성화를 위한 최적의 양자화 매개변수(s와 z)를 찾는 것

• 여전히 정적 매개변수에 의존 → 입력 활성화 분포의 동적 특성을 고려하여 양자화 매개변수를 업데이트할 수 있도록 함
• 입력 데이터에 따라 활성화 분포가 변경될 수 있지만 전체적인 추세는 유사하게 유지되므로 시간 단계의 차이를 기반으로 한 최적의 간격 결정을 통해 안정적이고 강력한 양자화 결과를 얻을 수 있음.
• 각 시각 단계에 대해 스칼라 양자화 간격을 사용하면 되겠다 라고 생각할 수 있으나.. 이는 시간 간 관계를 포착하는 데에 부족하고 훈련과 추론 사이에 서로 다른 경우에는 사용할 수 없음. 따라서 시간 단계 정보를 통해 양자화 간격을 예측하는 학습 가능한 소규모 네트워크 사용
• TDQ 모듈에서 동적 양자화 간격 간격 $\tilde{s}$ 는 제공된 시간 단계 $t$ 에 기반하여 생성됨:
  $$I=enc(t), \tilde{s}=f(I)$$
  
  • enc(⋅): 시간 단계의 인코딩 함수
  • 𝐼: 인코딩된 특징
  • 𝑓(⋅): 생성기 모듈 함수

 

• 각 양자화 연산자에 부착되고 주어진 시간 단계에 기초하여 독립적인 동적 양자화 간격이 생성됨
• 생성기는 선형 층 여러 개를 단순하게 쌓고 softplus 함수로 값을 제한하여 구현하며, 생성기의 모든 구성 요소는 미분 가능함. 따라서 PTQ 또는 QAT 과정 동안, 간격은 양자화 오류를 최소화하는 방향으로 그래디언트 강하를 사용하여 업데이트될 수 있음
  예를 들어, LSQ와 같은 잘 알려진 QAT 알고리즘을 사용할 때, 정적 간격이 TDQ 모듈의 출력으로 대체되고 양자화 함수는 다음과 같이 수정됨.:

• 직통 추정기(a straight-through estimator)를 사용하면 gradient는 학습 가능한 생성기의 매개변수를 통해 전파될 수 있으며, 이 매개변수들은 최종 작업 손실을 최소화하는 목표를 가지고 반복적인 그래디언트 강하를 통해 업데이트됨.
• 동일한 pipeline은 (재구성 오류를 최소화하는 로컬 그래디언트를 활용하는) PTQ 알고리즘에도 적용 가능.
  ⇒ TDQ 모듈은 QAT 및 PTQ 체계에 쉽게 적용 가능함.

• TDQ 모듈의 이점
  • (input-dependent dynamic quantization과는 달리) 비용이 들지 않는 추론이 가능함.
  • PTQ나 QAT 과정 이후에 time-dependent interval이 오프라인에서 사전 계산될 수 있고, 추론 중에 사전 계산된 값을 사용할 수 있음.
  • 기존 프레임워크에 수정 없이 원활하게 통합될 수 있음.

3.4. Engineering Details

• Frequency Encoding of Time Step:
  • time step을 직접 생성기에 공급 시 열등한 수렴 품질 발생 (저주파 유도 편향 때문, time step이 생성기에 직접 입력되는 경우 time step에 관계없이 거의 변하지 않는 간격 생성하는 경향)
    cf) 이미지에서 frequency가 높다 = 픽셀값의 변화량이 잘 변한다는 것을 의미함. 수렴하려면 픽셀값이 어느 정도 유동적이어야 하는데, 저주파 유도편향 때문에 픽셀값이 잘 변하지 안게되니 열등한 수렴 품질이 발생한다는 것!
  • 저주파 유도 편향을 완화하기 위하여 → geometric Fourier encoding 사용
  • t: the current time step, d: the dimension of encoding vector, I: the encoded vector

• 이를 통해 TDQ 모듈이 time step의 high-frequency dynamics(고주파 동적)을 수용할 수 있게 됨

• Initialization of TDQ Module:
  • TDQ 모듈의 출력을 원하는 값으로 초기화해야 함. (기존의 양자화 기술은 정적 단계 값만 초기화하면 됨.)
  • 이를 위해 가중치에 대해 He 초기화를 사용하고 MLP의 마지막 선형 층의 편향을 원하는 값으로 설정
    MLP에 입력되는 것이 평균이 0인 무작위 변수로 간주될 수 있기 때문에, He 초기화된 MLP의 출력도 평균이 0이 됨.
    따라서 편향 조정을 통해 MLP 출력의 평균을 원하는 값으로 제어 가능.

4. Experimental Setup

• DDIM(a pixel space diffusion model) ← CIFAR-10 dataset
• LDM(a latent space diffusion model) ← LSUN Churches dataset

⇒ 저해상도 및 고해상도 이미지 생성 시나리오에서 제안된 방법의 효과를 보여줄 수 있음.

 

• 두 모델 모두에 PTQ와 QAT를 적용
• 잠재 공간 모델이 VAE와 확산 모델을 포함하는데 본 연구에서는 확산 모델 부분에만 양자화 적용함. (VAE에는 적용X)
• TDQ 모듈의 출력으로 정적 간격을 대체하는 것

5. Results

5.1. Quality Analysis after QAT and PTQ

기존의 정적 양자화 방법과 TDQ 모듈 비교

• 활성화 비트 수가 감소함에 따라 모든 정적 양자화 방법은 품질이 떨어짐
• 반면에 TDQ 모듈을 적용한 법은 일관된 품질을 보여주며, 8비트에서도 상당한 품질 개선을 보여줌. 또한, 전체 정밀도 출력에 대한 품질 저하가 거의 없음.
  ⇒ 제한된 활성화 양자화 레벨을 효율적으로 할당함으로써 이러한 이점을 달성함.
• NIPQ는 straight-through estimator의 불안정성을 해결하기 위해 인공 노이즈를 기반으로 한 가짜 양자화(pseudo quantization)를 도입함. 그러나 NIPQ의 노이즈는 입력 노이즈와 구분이 안 되며, 확산 모델의 수렴을 방해함.

 

기존의 정적 PTQ 체계와 TDQ 모듈의 PTQ 비교를 제시

• Min-Max 방법: target tensor의 최소 및 최대 값에 의해 범위가 결정되는 naive linear quantization approach (단순한 선형 양자화 접근법).
• 대체로 활성화 비트가 감소하면 성능 저하가 발생(FID ↑)하는데, TDQ는 FID 저하가 약간만 발생함.

 

기존의 QAT와 PTQ는 흐릿하고 인식할 수 없는 이미지를 생성하는 반면, TDQ가 적용된 방법은 현실적인 이미지를 생성

⇒ 활성화 양자화에서 시간 정보를 통합하는 것은 출력의 인식적 품질을 유지하는 데 매우 효과적임을 보여줌.

 

 

5.2 Generalization Performance of TDQ Module

빠른 추론을 위한 TDQ 모듈의 성능에 대한 실험 결과

• LSQ: time step이 감소함에 따라 FID↑ = 성능이 크게 저하
• TDQ module: LSQ에 비해 비교적 적은 성능 저하
  ⇒ 샘플링 시간 단계가 변화해도 TDQ 모듈이 효과적으로 기능함

 

5.3 Ablation Studies of TDQ module

TDQ 모듈의 출력 동역학을 조사

동적 간격의 업데이트를 시각화하여 시간 단계와의 관계 보여줌

 ✅ 그래프 해석하기 ✅ 

1. 그래프의 구성
- 파란색 선 (Step Size): TDQ 모듈에서 생성된 양자화 간격의 크기
- 빨간색 선 (Activation Variance): 해당 레이어의 입력 활성화의 분산 (활성화 값의 변동성을 의미, 네트워크가 다양한 입력에 얼마나 민감하게 반응하는지)

2. 그래프 해석
- 일반적으로, 활성화 분산이 크게 변할 때 양자화 간격도 유사한 패턴으로 변화하는 것을 볼 수 있음. 
= TDQ 모듈이 입력 활성화의 변화에 따라 동적으로 간격을 조정하고 있음을 의미
- 몇몇 그래프에서는 활성화 분산의 변화가 간격의 변화보다 더 뚜렷하게 나타나며, TDQ 모듈이 이러한 변화를 감지하고 적절히 반응하고 있음

• DDIM에서 W4A4 LSQ를 사용하여 훈련된 간격은 활성화 변화와 일치하여 변화하는 경향을 보여주지만, 이 패턴은 모든 계층에서 일관되게 관찰되지는 않음.

 

• TDQ는 계층의 분포 정보를 직접 사용하지 않고도 동적 양자화에 비해 우수한 성능
• TDQ 모듈 내의 계층 수가 성능에 미치는 영향: 계층 수가 4개를 초과할 때 출력 품질이 상당히 일관되게 된다는 것 ⇒ 성능과 복잡성의 균형을 고려하여 MLP의 4 계층을 경험적으로 선택

 

각 시간 단계마다 양자화 간격을 직접 학습하는 구성과 TDQ의 성능 비교

• 각 시간 단계마다 양자화 간격을 직접 학습하고 여러 시간 단계에 걸쳐 공유 양자화 간격을 사용하는 개념은 기준선(S1)에 비해 출력 품질을 향상시킬 수 있지만, TDQ는 여전히 우수한 출력 품질을 보이며 더 큰 다양성을 제공.

 

6. Discussion

• 많은 계층에서 강한 시간적 의존성을 보이지만, 약 30%의 계층에서는 이러한 상관관계를 보이지 않음. 이러한 계층들은 주로 U-Net의 중간 블록에 위치하며, 그들의 인스턴스별 의미 정보에 크게 영향을 받음.
• LDM 모델은 시간적 의존성이 적어 성능 향상이 덜 나타남. 하지만, 이러한 경우에도 TDQ 모듈이 수렴을 보장하여 시간 단계에 걸쳐 일관된 출력을 결과로 낳는다는 점을 강조 → 시간적 의존성이 적은 시나리오에서도 출력 품질이 기존 PTQ/QAT 알고리즘과 비교할 수 있도록 보장함.

 

7. Conclusion

• 확산 모델에서의 활성화 양자화 문제, 특히 시간 단계에 걸친 활성화의 동역학 탐구
• 기존의 정적 양자화 방법이 이 문제를 효과적으로 해결하지 못함.
• TDQ 모듈을 소개: QAT와 PTQ 모두에 원활하게 통합되면서 동적 양자화 간격 생성을 기반으로 출력 품질을 크게 향상시킴. 추론 중에 어떠한 오버헤드도 발생하지 않는 접근 방식을 구현할 수 있으므로 모바일 및 엣지 디바이스를 위한 저비트 확산 모델의 성능 개선과 품질 유지에 도움이 될 것

---

cf) 기존 양자화 기법과 TDQ의 주요 차이점

기존의 정적 양자화 기법과 Temporal Dynamic Quantization (TDQ)의 주요 차이점:

1. 동적 양자화 간격 생성:
   • 기존 정적 양자화: 양자화 간격은 모델 훈련 전에 설정되며, 학습 및 추론 과정에서 변경되지 않음. 이 방식은 입력 데이터의 변화나 시간에 따른 활성화의 분포 변화를 반영하지 못하는 단점이 있음.
   • TDQ: 시간 단계 정보에 기반하여 양자화 간격을 동적으로 조정함. 이는 각 시간 단계에서 최적의 양자화 설정을 생성하여 활성화 양자화 오류를 최소화하는 데 도움.
2. 계산 오버헤드:
   • 기존 방식: 일부 동적 양자화 방법은 실행 시 추가적인 계산 비용을 발생시킴
     (→ 모바일이나 엣지 디바이스에서 문제가 될 수 있음.)
   • TDQ: 추론 중 추가적인 계산 오버헤드 없이 구현됨. 
3. 통합과 확장성:
   • 기존 방식: 입력 데이터의 변화에 따라 유연하게 조정되지 않으며 대부분 정적임.
   • TDQ: QAT와 PTQ에 쉽게 통합될 수 있음 → 기존 양자화 알고리즘을 확장하여 시간에 따라 최적의 양자화 구성 생성 가능.
4. 출력 품질 및 일관성:
   • 기존 방식: 낮은 비트 정밀도에서 성능 저하가 발생 가능.
   • TDQ: 낮은 비트 정밀도에서도 출력 품질을 상당히 향상시키는 동시에 시간 단계에 걸쳐 일관된 출력을 보장함.
     (이는 특히 시간적 의존성이 적은 시나리오에서도 PTQ/QAT 알고리즘과 비교할 수 있는 출력 품질을 유지함.)