DINOv3에 대한 내용

참고하면 좋은 유튜브

EMA 1 EMA 2 EMA 3

DINOv2, DINOv3

Abstract

1. 좋은 데이터 준비 + 설계 디자인 ⇒ 데이터 셋, 모델 scale up
2. 학습 길어짐 ⇒ 세밀한 특징에 대한 인지 저하 (너무 러프하게 이미지를 봄) ⇒ Gram anchoring
3. 모델 유연성 ↑

1. Introduction

SSL: 메타데이터(라벨링) 없는 데이터로 학습

⇒ 입력 분포 변화에 견고, 강한 전역/지역 특징 제공 ⇒ 이미지 내 패턴의 자연스러운 공통 발생 활용

모델 scale up 미해결 과제

1. 데이터 수집 어케하지? ⇒
2. 코사인 스케줄 사용 X ⇒
3. 학습 계속 ⇒ 국소 특징 질 저하 ⇒ Gram anchoring

Strong & Versatile Foundational Models

1. fine tuning 없이도, 기존 finetuned 모델보다 좋음
2. 과학적 응용 good

Superior Feature Maps Through Gram Anchoring

SSL ⇒ data ↑ ⇒ 전역적 표현 학습 ⇒ 조밀한 특징 질 저하

The DINOv3 Family of Models

Gram anchoring ⇒ 모델 scale up 가능 ⇒ 7B 모델 (DINO) 성공

+

distillation ⇒ 더 작은 모델 ㄱㄴ

2. Related Work

SSL

SSL (자기 지도 학습)은 어노테이션 없이 인위적인 학습 과제(Pre-text tasks)를 통해 데이터의 표현을 학습하며, 신호가 연속적인 컴퓨터 비전 분야는 언어 분야보다 이러한 과제 설계의 난도가 높습니다.

이미지의 일부를 복원하는 인페인팅(Inpainting) 기반 방식은 MAE(Masked Auto-Encoder)와 JEPA(Joint-Embedding Predictive Architecture) 등으로 발전하며 고수준의 시각 특징을 추출하는 데 성공했습니다.

잠재 공간을 예측하는 JEPA 의 방식은 iBOT 에 이르러 이미지 간 차이를 식별하는 판별적 손실 함수 와 결합되며 더욱 정교하고 강력한 시각 표현 학습으로 발전했습니다.

1. Summarization

Pixel 기반 방식이 이미지의 외형적인 "복사"에 집중한다면, Latent Space 예측 방식은 이미지 속에 담긴 객체의 본질적인 "의미" (Semantic Representation) 를 추론하는 것이 목적입니다.

2. Example

이 차이를 "고양이의 귀" 부분이 마스킹 된 상황으로 비유하여 구체적인 수치와 함께 설명하겠습니다.

• Pixel-based (MAE): 마스킹 된 영역의 픽셀 값 개 ( RGB) 를 하나하나 맞추려 노력합니다. 만약 조명이 어둡거나 노이즈가 섞여 있다면, 모델은 본질적인 "고양이 귀" 의 형태보다는 픽셀의 어두운 색상 값을 복원하는 데 더 많은 계산 자원을 소모하게 됩니다.
• Latent-based (JEPA / DINOv3): Target Encoder 가 해당 영역을 보고 생성한 차원의 고차원 벡터를 예측합니다. 이 벡터 안에는 "이것은 동물의 귀이며, 뾰족한 형태를 가졌고, 털 질감을 가졌다" 라는 High-level semantic 정보가 압축되어 있습니다. 모델은 픽셀의 미세한 색상 차이보다는, 그 위치에 어떤 특징 (Feature) 이 와야 하는지 그 "맥락" 을 맞추는 데 집중합니다.

3. Conclusion

결론적으로, 픽셀 자체를 복구하는 것은 조명이나 각도 변화 같은 불필요한 정보 (Redundancy) 까지 학습하게 만들 수 있습니다. 반면, Latent Space 를 예측 목표로 삼으면 모델은 이미지의 지엽적인 노이즈를 무시하고 객체의 형태, 구조, 종류와 같은 핵심적인 Visual representation 을 더 효과적으로 학습하게 됩니다. 이것이 DINOv3 가 다양한 Downstream tasks 에서 강력한 성능을 내는 근본적인 이유입니다.

Vision Foundataion Models

생략

Dense Transformer Features

생략

3. Training at Scale Without Supervision

데이터 셋 확장 ⇒ 모델 확장 + SSL ⇒ 특정 작업 국한 X, 범용 vision 분야 가능

3.1. Data Preparation

단순 양 증가 보단, 좋은 데이터 양 증가

Data Collection and Curation

1. 계층적 클러스터링 기반 큐레이션: Vo et al. (2024) 의 계층적 k-means 기반 자동 큐레이션 방식을 적용했습니다. DINOv2 를 이미지 임베딩으로 사용하고, 5단계의 클러스터링을 적용했습니다 (클러스터 수는 최하위부터 최상위까지 각각 2억 개, 800만 개, 80만 개, 10만 개, 2.5만 개). 계층 구조를 구축한 후 균형 잡힌 샘플링 알고리즘을 적용하여, 웹에 등장하는 모든 시각적 개념을 균형 있게 포괄하는 16억 8,900만 장의 큐레이션된 하위 집합 (LVD-1689M) 을 구축했습니다.

2. 검색 기반 큐레이션: Oquab et al. (2024) 이 제안한 절차와 유사한 검색 기반 시스템을 채택했습니다. 선택된 시드 (seed) 데이터셋과 유사한 이미지를 데이터 풀에서 검색하여, 다운스트림 작업과 관련된 시각적 개념을 다루는 데이터셋을 생성했습니다.

3. 공개 데이터셋 활용: ImageNet1k (Deng et al., 2009), ImageNet22k (Russakovsky et al., 2015), 그리고 Mapillary Street-level Sequences (Warburg et al., 2020) 를 포함한 공개 컴퓨터 비전 데이터셋을 사용했습니다. 이 마지막 부분은 Oquab et al. (2024) 을 따라 모델의 성능을 최적화하는 역할을 합니다.

Q1. 계층적 클러스터링을 통한 자동 정제에서 임베딩 추출 부분에서 DINOv2를 사용해서 image to vector로 변환했다고 하는데, DINOv2는 이미지를 벡터로 만들어줘?

A1. DINOv2와 임베딩 (Embedding) 추출 DINOv2 는 비전 트랜스포머 (Vision Transformer) 기반의 자가 지도 학습 모델로, 이미지를 입력받아 해당 이미지의 시각적 특징을 함축하고 있는 고차원 벡터로 변환하는 역할을 합니다. 이를 임베딩 (Embedding) 이라 부릅니다. 즉, 사진이라는 데이터를 컴퓨터가 계산할 수 있는 숫자 나열(벡터)로 바꾸어 주는 것이 맞습니다. 이렇게 벡터화가 되어야만 이미지 사이의 유사도를 계산하거나 클러스터링을 할 수 있습니다.

Q2. 계층적 클러스터링을 통한 자동 정제에서 계층적 K-means 알고리즘으로 비슷한 특징을 갖는 데이터를 묶었다고 하는데, 그럼 원시 데이터 170억장을 1단계에서는 2억개의 그룹으로 분류해서 각 그룹에 170억/2억 개의 이미지가 들어있다는 거야?

A2. 계층적 K-means 알고리즘의 구조 네, 이해하신 개념이 맞습니다. 다만 계층 구조를 이해할 필요가 있습니다.

• 1단계 (Lowest Level): $170$ 억 장의 원시 데이터를 $2$ 억 개의 매우 세밀한 그룹으로 나눕니다. 산술적으로는 평균 $170$ 억 $/$ $2$ 억 $=$ 약 $85$ 장의 이미지가 각 그룹(클러스터)에 들어가는 셈입니다.
• 상위 단계로 이동: 이후 이 $2$ 억 개의 그룹들을 다시 비슷한 것끼리 묶어 $800$ 만 개, $80$ 만 개, $10$ 만 개, 최종적으로는 $2$ 만 $5$ 천 개의 큰 그룹으로 계층화합니다.
• 결과적으로 거대한 이미지 더미가 "대분류 - 중분류 - 소분류 - 아주 세밀한 분류"의 트리 구조로 정리되는 것입니다.
  Q3. 계층적 클러스터링을 통한 자동 정제에서 각 계층 구조에서 골고루 뽑았다는 게 뭐지?
  A3. 계층 구조에서의 균형 샘플링 (Balanced Sampling) 균형 샘플링 이란 특정 카테고리의 데이터만 너무 많이 뽑히지 않도록 조절하는 것입니다. 인터넷 데이터에는 셀카, 풍경, 음식 사진은 엄청나게 많지만, 희귀한 동물이나 특수 장비 사진은 매우 적습니다. 만약 무작위로 뽑는다면 모델은 셀카만 공부하게 될 것입니다. 하지만 앞서 만든 계층 구조(클러스터)를 이용해 모든 그룹에서 골고루 데이터를 추출하면, 흔한 이미지의 비중은 줄이고 희귀한 시각적 개념의 비중을 높여 전 세계의 다양한 시각적 정보를 빠짐없이 학습할 수 있게 됩니다.
  Q4. 검색 기반 큐레이션에서 결국 170억개 데이터 중에 일부를 추출했다는 거지?
  A4. 검색 기반 큐레이션의 대상맞습니다. 검색 기반 큐레이션 역시 처음에 확보한 $170$ 억 장 의 거대한 데이터 풀에서 필요한 것만 골라내는 과정입니다. 전체 데이터 중 우리가 관심 있는 특정 분야와 닮은 것들만 "검색"해서 추출한 것입니다.
  Q5. 검색 기반 큐레이션에서 일부 데이터 추출이 어떻게 다운스트림 태스크와 관련된 시각적 개념을 집중시키지?
  A5. 다운스트림 태스크와 시각적 개념의 집중 다운스트림 태스크 (Downstream Task) 란 모델이 최종적으로 해결해야 할 구체적인 문제(예: 의료 영상 판독, 자율 주행 도로 인식 등)를 의미합니다.
• 시드 데이터셋 (Seed Dataset): 우리가 목표로 하는 태스크와 밀접한 관련이 있는 검증된 데이터입니다.
• 집중 과정: 이 시드 데이터와 유사한 이미지를 $170$ 억 장 중에서 찾아내면, 모델은 해당 분야의 시각적 특징(질감, 형태, 구도 등)을 더 집중적으로 학습하게 됩니다. 결과적으로 실제 현장에 적용했을 때 더 높은 성능을 내는 최적화된 데이터셋 이 구축됩니다.
  Q6. 그래서 결과적으로 총 데이터 수가 뭐야?
  A6. 최종 데이터 수 텍스트에 명시된 부분과 수치 관계를 정리하면 다음과 같습니다.
• 첫 번째 파트 (LVD-1689M): 계층적 클러스터링의 결과물인 16억 8,900만 장 ($1,689 \text{ M}$) 입니다.
• 두 번째 파트 (검색 기반): $170$ 억 장 중 시드 데이터와 유사한 이미지들입니다.
• 세 번째 파트 (공개 데이터셋): ImageNet 등을 포함한 약 $1,500$ 만 장 내외입니다.
• 전체 학습 데이터셋: 실제 총합은 $16$ 억 $8,900$ 만 장 + $\alpha$ (Part 2) + 약 $1,500$ 만 장 (Part 3) 입니다. 하지만 Part 1 의 비중이 $99$ % 이상으로 압도적이기 때문에, 논문에서는 최종 결과물을 메인 데이터셋 명칭인 LVD-1689M 으로 통칭합니다.

Data Sampling

ImageNet1k로만 구성된 균일 배치를 전체 훈련의 10% + 나머지 90%는 다른 모든 구성 데이터셋을 혼합한 불균일 배치

Data Ablation

3.2. Large-Scale Training with Self-Supervision

모델 크기를 키우면서, 전역 및 지역적 특성을 모두 개선하자!

Learning Objective

$$L_{\text{Pre}} = L_{\text{DINO}} + L_{\text{iBOT}} + 0.1 \times L_{\text{Koleo}}$$

• $L_{\text{DINO}}$: 이미지 전체의 전역적 특징을 학습하는 이미지 레벨 손실 함수입니다.

• $L_{\text{iBOT}}$: 이미지 내부 패치의 국소적 특징을 학습하는 패치 레벨 손실 함수입니다.

• $0.1 \times L_{\text{Koleo}}$: 배치 내 특징들을 공간상에 균일하게 분산시켜 특징 붕괴를 방지하는 규제 항입니다.

• $L_{\text{DINO}}$ (Image-level Objective): 이미지 전체를 대상으로 하는 손실 함수입니다. 이미지의 글로벌 특징(Global Features)을 학습하며, 서로 다른 뷰(View)의 이미지가 유사한 특징을 갖도록 유도하는 이미지 단위 의 목표를 가집니다.
  
  

• $L_{\text{iBOT}}$ (Patch-level Objective): 이미지 내의 특정 패치(Patch)들을 대상으로 하는 손실 함수입니다. 이미지의 일부를 가리고(Masking) 이를 복원하는 잠재적 재구성(Latent Reconstruction) 과정을 통해 국소적인 특징 (Local Features)을 학습하도록 돕습니다.
  
  아래 그림에는 student에서 local crop을 사용하는 것 처럼 표현되었는데, 아닌 것 같음...
  
  이때, 동일한 아키텍처, 다른 가중치를 사용함. 교사($\text{Teacher}$)와 학생($\text{Student}$) 모두 동일한 백본(예: $\text{ViT-L}$)을 사용합니다. 하지만 실시간으로 학습되는 학생 모델과 달리, 교사 모델은 학생 모델의 과거 기록을 바탕으로 만들어진 '복제본' 에 가깝습니다.

• $L_{\text{Koleo}}$ (Koleo Regularizer): 배치(Batch) 내의 특징들이 공간상에 골고루(Uniformly) 퍼지도록 장려하는 규제 항입니다. 특정 특징들이 한곳으로 뭉치는 붕괴(Collapse) 현상을 방지하며, 특징 공간을 효율적으로 활용하게 만듭니다.

Updated Model Architecture

Optimization

• 복잡한 스케줄링 대신 상수 값을 써서 학습 기간을 유연하게 늘릴 수 있도록 설계

• 초기 안정성을 위해 선형 웜업을 사용

• 데이터 이해도를 높이기 위해 글로벌/로컬 크롭을 섞어 씀

• AdamW

• 10만 회의 웜업

• 0.0004의 일정한 학습률

• 0.04의 가중치 감쇠

• 레이어당 0.98의 학습률 감쇠 계수

• 0.4의 레이어 드롭아웃

• 교사 모델에 대해 0.999의 EMA 계수

우리는 모든 학생 지역 크롭과 두 교사 전역 크롭의 클래스 토큰, 그리고 두 모델의 서로 다른 전역 크롭 쌍 사이에 $L_{DINO}$ 손실을 적용합니다.

학생이 보는 전역 크롭 패치 토큰 중 $[0.1, 0.5]$ 사이의 무작위 비율을 50% 의 확률로 마스킹하며, 이들과 교사 EMA 가 보는 가시적 토큰 사이에 $L_{iBOT}$ 손실을 적용합니다.

학생이 보는 첫 번째 전역 크롭의 클래스 토큰 16개로 구성된 소규모 배치에는 $L_{Koleo}$ 손실을 적용합니다.

4. Gram Anchoring: A Regularization for Dense Features

Gram anchoring을 통한 지역적 작업 성능 향상

4.1. Loss of Patch-Level Consistency Over Training

60만 회 이후부터는 유사도 지도가 상당히 모호해짐 ⇒ 참조 패치와 무관한 패치들이 높은 유사도를 보임

CLS 토큰과 패치 간 유사도 증가가 국소성(Locality) 감소를 의미하는 이유?

CLS 토큰 은 이미지 전체의 정보를 요약하는 전역적(Global) 특징을 담습니다. 이와 패치들 사이의 유사도가 높아진다는 것은 다음과 같은 상황을 의미합니다.

• 정보의 획일화: 원래 각 패치(Patch) 는 이미지의 특정 부분(국소적 특징)만 선명하게 보여줘야 합니다. 그런데 전역 정보를 담은 CLS 토큰 과 너무 비슷해지면, 각 패치가 자기만의 고유한 지역 정보를 잃어버리고 전체적인 배경 정보만 똑같이 공유 하게 됩니다.

• 해상도 저하: 결과적으로 모든 패치가 비슷비슷한 전역 정보만 내뱉게 되므로, 사물의 경계를 정밀하게 구분해야 하는 세밀한 작업(Dense tasks, 예: 세그멘테이션) 에서 성능이 떨어지게 되는 것입니다.

4.2. Gram Anchoring Objective

학습 초기의 우수한 조밀 특성을 보이는 교사의 패치간 유사도(Gram matrix)를 학생의 gram이 따라가게함.

그러니까 학습 초기에는 우수한 조밀 특성을 보임 ⇒ 그 모델의 패치간 유사도를 학생이 따라감

$$L_{Gram} = ∥X_{S} · X^⊤_{S} − X_G · X^⊤_G∥^2_F$$

$X_S$: 학생의 패치간 유사도 행렬 $X_G$: 선생의 패치간 유사도 행렬

⇒ $$L_{Ref}=w_DL_{DINO}+L_{iBOT}+w_{DK}L_{DKoleo}+w_{Gram}L_{Gram}$$

• 100만 회 반복 학습 이후에 refinement 시작
• 1만 회 반복 학습마다 Gram 교사를 업데이트하여, 메인 EMA 교사와 동일하게 만듦

1. 첫 번째 훈련 phase (100만 회): $$L_{\text{Pre}} = L_{\text{DINO}} + L_{\text{iBOT}} + 0.1 \times L_{\text{Koleo}}$$

2. 두 번째 훈련 phase(refinement) (1만 회 마다 Gram 교사를 메인 EMA 교사와 동일하게 바꿈): $$L_{Ref}=w_DL_{DINO}+L_{iBOT}+w_{DK}L_{DKoleo}+w_{Gram}L_{Gram}$$

4.3. Leveraging Higher-Resolution Features

Gram 교사를 더 좋게 하기 위해,

Gram 교사에 일반 해상도의 두배인 이미지를 입력 ⇒ 얻은 특징 지도를 bicubic 보간법으로 2⨉ 다운샘플링하여 학생 모델의 출력 크기와 일치 시킴

⇒ 매끄러운 고해상도 특징의 개선된 패치 일관성을 학생 모델로 효과적으로 증류!!!

5. Post-Training

5.1. Resolution Scaling

1. 메인 훈련: 효율성 위주의 기초 학습

   • 설정: $256 \times 256$ 해상도와 $16$ 패치 크기로 훈련하여 속도와 효율성을 확보합니다.

2. 고해상도 적응 단계: 가변 해상도 대응

   • 필요성: 현대 비전 응용 분야에서 요구되는 $512 \times 512$ 이상의 고해상도 및 가변적인 추론 환경에 대응하기 위해 도입되었습니다.
   • 방법: 혼합 해상도(Mixed Resolution) 방식을 사용하며, 전역 크롭(${512, 768}$)과 지역 크롭(${112, 168, 224, 336}$)을 섞어 1만 회 추가 학습을 진행합니다.

고해상도 적응 단계는 Gram anchoring에 필수적임.

1. 첫 번째 훈련 phase (100만 회): $$L_{\text{Pre}} = L_{\text{DINO}} + L_{\text{iBOT}} + 0.1 \times L_{\text{Koleo}}$$

2. 두 번째 훈련 phase(refinement) (1만 회 마다 Gram 교사를 메인 EMA 교사와 동일하게 바꿈): $$L_{Ref}=w_DL_{DINO}+L_{iBOT}+w_{DK}L_{DKoleo}+w_{Gram}L_{Gram}$$

3. High resolution adaptation (1만 회): $$L_{Ref}=w_DL_{DINO}+L_{iBOT}+w_{DK}L_{DKoleo}+w_{Gram}L_{Gram}$$

1. 첫 번째 훈련 단계 (Initial Training Phase)
   • 기간: 약 $100$ 만 회 ($1\text{M iterations}$)
   • 손실 함수: $L_{\text{Pre}} = L_{\text{DINO}} + L_{\text{iBOT}} + 0.1 \times L_{\text{DKoleo}}$
   • 특징: $256 \times 256$ 단일 해상도에서 모델의 기초적인 전역 및 국소 특징을 학습합니다.

2. 두 번째 훈련 단계 (Refinement Step)
   • 특징: $100$ 만 회 이후 시작되며, 망가진 패치 수준 일관성을 '수리'하는 단계입니다.
   • 손실 함수: $L_{\text{Ref}} = w_D L_{\text{DINO}} + L_{\text{iBOT}} + w_{DK} L_{\text{DKoleo}} + w_{\text{Gram}} L_{\text{Gram}}$
   • 메커니즘: $1$ 만 회마다 메인 EMA 교사 의 상태를 Gram 교사 로 복사하여 패치 간 유사도 구조를 고정(Anchoring)합니다.

3. 고해상도 적응 단계 (High-resolution Adaptation)
   • 기간: 약 $1$ 만 회 ($10\text{k iterations}$) 추가 학습
   • 손실 함수: $L_{\text{Pre}}$ 가 아니라 $L_{\text{HRef}}$ (또는 $L_{\text{Ref}}$ 와 유사한 구성)를 사용합니다.
   • 논문에서는 고해상도 적응 시 $\text{Gram anchoring}$ 이 필수적이라고 명시하고 있으며, 고해상도 특징을 다운샘플링하여 계산하는 $L_{\text{HRef}}$ 를 사용하여 학습합니다.
   • 특징: ${512, 768}$ 등 다양한 크기의 혼합 해상도 를 입력하여 가변 해상도 대응력을 높입니다.

5.2. Model Distillation

A Family of Models for Multiple Use-Cases

ViT-7B를 ViT-S,B,L로 증류함. 증류할 때, $L_{\text{Pre}} = L_{\text{DINO}} + L_{\text{iBOT}} + 0.1 \times L_{\text{DKoleo}}$로 100만 번 학습 후, 코사인 스케줄에 따라 25만회 학습률 쿨다운 + high-resolution adaptation But, 패치 수준 일관성 문제가 없어, Gram anchoring X.

교사를 EMA가 아닌 7B 모델을 직접 사용하여, 고정함.

• 패치 수준 일관성(Patch-level consistency) 문제: 학습이 진행됨에 따라 개별 패치 특징이 고유의 지역적 정보를 잃고 전역 정보($\text{CLS}$ 토큰)와 지나치게 유사해지거나 서로 획일화되어 세밀한 표현력이 떨어지는 현상을 의미합니다.

Efficient Multi-Student Distillation

생략

5.3. Aligning DINOv3 with Text

생략

6. Results

6.1. DINOv3 provides Exceptional Dense Features

6.2. DINOv3 has Robust Versatile Global Image Descriptors

6.3. DINOv3 is a Foundation for Complex Computer Vision Systems

MoR에 대한 내용

Abstract

• Parameter sharing (Recursive)
• Adaptive computation (EE)

EE ⇒

• non-RT, RT, RRT: 모든 토큰이 동일한 횟수 만큼 재귀 함 ⇒ $O(N^2)$ EERT(MoR): 각 토큰별로 recursive 깊이가 다름 ⇒ $O(N_{active}^2)$
• Active 토큰의 KV만 caching

1. Introduction

1) Parameter sharing ⇒ param ↓ 2) EE ⇒ FLOPs ↓ 3) 재귀별 KV caching ⇒ memory footprint ↓

2. Method

2.1. Preliminary

Parameter-sharing strategies

2.2. MoR

2.2.1. Routing Strategies: Expert-choice vs. Token-choice

Expert-choice routing

재귀마다 전문가를 다르게 둠.

ex) $N_r=3$ 일 때, 전문가는 3명.

$$H_t^{r+1} = \begin{cases} g_t^r f(H_t^r, \Phi') + H_t^r, & \text{if } g_t^r > P_\beta(G^r) \ H_t^r, & \text{otherwise} \end{cases} \quad (2.1)$$

⇒ 에서 Stepwise + Cycle로 예를들면.

$N_r=2$임. 즉, 전문가 2명.

1, 3, 6 지난 후에, 전문가 1번이 평가해서, $g_t^r = G(\theta_r^\top H_t^r)$ 계산!

$H_t^r$: 전문가 1이 평가할 hidden state $\theta_r^\top$: 전문가 1의 채점표(parameter)

전문가 1의 점수가 $g_t^r$!

Token-choice routing

매번 평가하는 것이 아니라, loop에 본격적으로 들어가기 전에 모든 token마다 얼마나 loop를 돌릴지 정하고 가는 형태.

$$H_t^{r+1} = \begin{cases} g_t^r f(H_t^r, \Phi') + H_t^1, & \text{if } r = i \ g_t^r f(H_t^r, \Phi'), & \text{otherwise} \end{cases} \quad (2.2)$$

pros & cons

Router가 미래의 정보까지 보고 학습됨 ⇒ cheating임 ⇒ Auxiliary router or a regularization loss ⇒ 미래의 토큰 정보에 접근하지 않고 추론 시 상위 k 토큰을 정확하게 탐지

2.2.2. KV Caching Strategies: Recursion-wise Caching vs. Recursive Sharing

예시 상황 요약

입력: 토큰 5개 ($T_1, T_2, T_3, T_4, T_5$), Hidden Dimension: 10 라우팅 결과: $T_1, T_4$(1회 종료), $T_3$(2회 종료), $T_2, T_5$(3회 완주)

기법 요약

• Recursion-wise KV caching: 각 재귀 단계에 생존한 토큰의 KV(크기 $N \times 10$)만 해당 단계 전용 캐시에 저장함
• Recursive KV sharing: 첫 번째 재귀 단계의 전체 KV(크기 $5 \times 10$)만 저장하고, 이후 모든 깊은 단계에서 이를 재사용함

예시를 통한 차이점 설명 (Loop 3 시점에서 $T_6$ 생성 시)

$T_6$가 생성되면서 3번째 루프를 돌며 어텐션(Attention)을 수행할 때:

• Recursion-wise:Loop 3 캐시를 참조하는데, 여기엔 생존한 $T_2, T_5$의 정보(크기 $2 \times 10$)만 저장되어 있음3.탈락한 $T_1, T_3, T_4$의 정보는 재계산하지 않고 무시하며, 오직 $T_2, T_5$에만 집중(Attention)함 (문맥 손실 발생)
• Sharing:Loop 1 캐시를 참조하는데, 여기엔 $T_1 \sim T_5$가 모두 저장(크기 $5 \times 10$)되어 있음 5.탈락한 토큰들도 포함한 전체 문맥을 참조할 수 있음 (단, Query는 Loop 3 수준이나 Key/Value는 Loop 1 수준이라 깊이가 얕음)

pros & cons

3. Experiments

Relaxed RT에 대한 내용

Abstract

Contribution

1. 기존 weight unshared model로 RT를 initialized.
2. LoRA를 통한 accuracy 회복.
3. Continuous Depth-wise Batching을 통한 inference speedup.

1. Introduction

큰 골조

$1^{st}$: 기존 잘 훈련된 non-RT 모델로 shared weight를 잘 초기화 $2^{nd}$: 약간의 finetune

Relaxed란? 모든 레이어가 동일한 가중치를 공유하는 대신, 각 레이어 LoRA 모듈을 둠 ⇒ non-RT와 RT의 중간의 의미로 Relaxed

LoRA

기존 non-RT가 12개 layer이고, 총 4개의 layer로 압축해서 3번 반복하는 RT를 구성하고 싶음 ⇒ non-RT의 3개의 layer를 1개로 압축하고 싶은거니 ⇒ 3개의 LoRA module을 둠

이때, LoRA 행렬의 rank가 relaxnity가 됨.

LoRA 행렬의 초기화

non-RT의 3개 layer를 1개의 layer로 압축했으니 총 4개의 가중치가 나옴 non-RT에서 3개, RT에서 1개

⇒ $W_{non-RT,1번 레이어} - W_{RT,1번 레이어}$로 truncated SVD 적용해서, LoRA 행렬 1을 초기화 $W_{non-RT,2번 레이어} - W_{RT,1번 레이어}$로 truncated SVD 적용해서, LoRA 행렬 2을 초기화 $W_{non-RT,3번 레이어} - W_{RT,1번 레이어}$로 truncated SVD 적용해서, LoRA 행렬 3을 초기화

Continuous Depth-wise Batching

Continuous Sequence-wise Batching

이 기법은 토큰의 위치(Position) 에 상관없이 모델 파라미터가 동일하다는 점을 이용합니다.

1. 예시 상황

• Batch Size: $2$ (동시에 2개의 문장 생성 가능)
• 요청 A: "안녕" (3개 토큰 생성 후 종료 예상)
• 요청 B: "오늘 날씨에 대해 길게 설명해줘" (100개 토큰 생성 예상)

2. 작동 원리
   1. Step 1~3: 요청 A와 B가 동시에 토큰을 하나씩 생성합니다.
   2. Step 4: 요청 A가 작별 토큰(<EOS>)을 내뱉고 생성을 마칩니다.
   3. 핵심: 기존 방식은 요청 B가 100개를 다 만들 때까지 기다려야 하지만, 이 기법은 요청 A가 비운 슬롯에 즉시 새로운 요청 C 를 투입합니다.
   4. 이유: 모델의 가중치는 문장의 시작 부분($t=1$) 을 계산할 때나 중간 부분($t=50$) 을 계산할 때나 똑같기 때문입니다.

Continuous Depth-wise Batching

이 기법은 Recursive Transformer (RT ) 에서만 가능하며, 가중치가 레이어의 깊이(Depth) 에 상관없이 동일하다는 점을 이용합니다. 이는 조기 종료(Early Exiting) 기법과 결합될 때 극대화됩니다.

1. 예시 상황

• 모델 구조: 1개의 가중치 블록을 최대 10번 반복(loop) 하는 RT
• 요청 A (쉬운 질문): "1+1은?" (연산량이 적어 4번의 반복만으로 정답 도출 가능)
• 요청 B (어려운 질문): "양자역학을 설명해줘" (연산량이 많아 10번의 반복이 모두 필요)

2. 작동 원리
   1. Loop 14: 요청 A와 B가 동일한 가중치 블록에서 동시에 14회차 반복 연산을 수행합니다.
   2. Loop 5: 요청 A는 이미 답을 찾았으므로 4회차에서 Early Exit 하여 계산을 마칩니다.
   3. 핵심: 이때 요청 B는 여전히 5회차 연산이 필요합니다. RT 는 가중치가 모든 루프에서 동일하므로, 요청 A가 빠진 자리에 새로운 요청 C의 1회차 연산 을 요청 B의 5회차 연산과 동시에 배치(jointly computing) 로 묶어서 처리할 수 있습니다.
   4. 결과: 하드웨어는 요청 B의 남은 계산을 처리하는 동시에 새 요청의 시작 계산을 병렬로 수행하게 되어 쉬는 구간이 없어집니다.

사용자님의 질문은 매우 날카로운 지적입니다. "어차피 하나의 토큰을 완성하려면 여러 번( 번) 돌려야 하는 것은 똑같지 않나?"라는 의문이 생길 수 있습니다. 하지만 RT (Recursive Transformer)와 non-RT 의 결정적인 차이는 '조기 종료(Early Exiting)' 시 하드웨어 가용성을 어떻게 활용하느냐에 있습니다.

설정하신 규칙에 따라 강조 표시 와 수식 주변에 공백을 적용하여 설명해 드립니다.

1. 가중치의 특수성 vs. 범용성

가장 큰 차이는 물리적인 가중치() 가 메모리에 로드되었을 때, 어떤 작업을 처리할 수 있는지에 대한 범용성 입니다.

• Non-RT (Weight Specialization): 각 레이어가 서로 다른 가중치() 를 가집니다. 번 레이어 연산을 위해 을 메모리에 올렸다면, 그 시점에는 오직 '3층 연산' 만 할 수 있습니다.
• RT (Weight Universality): 모든 루프가 동일한 가중치() 를 공유합니다. 이 가중치가 메모리에 올라와 있다면, 그것이 누군가에게는 '1회차 루프' 일 수도 있고, 다른 누군가에게는 '10회차 루프' 일 수도 있습니다.

  2. Early Exiting 상황에서의 처리량(Throughput) 비교

  배치(Batch) 내에 쉬운 질문(A)과 어려운 질문(B)이 섞여 있다고 가정해 봅시다.

  Non-RT의 경우 (Sequence-wise Batching)

1. Step 1: [A의 1층, B의 1층] 연산
2. Step 2: [A의 2층, B의 2층] 연산 (A가 여기서 정답을 찾아 조기 종료함)
3. Step 3: [None, B의 3층] 연산

• 문제: 하드웨어는 을 로드했지만, 새로 들어온 요청 C는 이 필요하므로 지금 당장 처리할 수 없습니다. A의 빈자리는 요청 B가 문장을 완전히 끝낼 때까지 비어 있게 됩니다.

  RT의 경우 (Continuous Depth-wise Batching)

1. Step 1: [A의 1회차, B의 1회차] 연산
2. Step 2: [A의 2회차, B의 2회차] 연산 (A가 여기서 조기 종료함)
3. Step 3: [C의 1회차, B의 3회차] 연산

• 해결: RT 는 모든 루프가 같은 가중치를 쓰므로, A가 빠진 자리에 즉시 새로운 요청 C의 시작 부분 을 끼워 넣을 수 있습니다. 하드웨어 입장에서는 똑같은 를 쓰기 때문입니다.

  3. 결론: 왜 RT가 더 빠른가?

  사용자님의 말씀대로 토큰 하나당 필요한 총 연산량은 비슷할 수 있습니다. 하지만 RT 는 가중치의 공유 특성 덕분에 '깊이(Depth)' 라는 차원에서 동적인 그룹화() 가 가능해집니다.

• Non-RT: 문장이 완전히 끝나야 빈자리를 채울 수 있음.
• RT: 어떤 요청이 특정 깊이에서 일찍 끝나면, 그 즉시 다음 요청의 첫 번째 루프를 시작 하여 빈자리를 채울 수 있음.
  이러한 유연성 덕분에 하드웨어 유휴 시간이 줄어들고, 이론적으로 의 추론 처리량 향상을 달성할 수 있는 것입니다.

2. Effective Model Compression with Recursive Patterns

2.1. 트랜스포머

$h_{t}^{l} = f(h_{t}^{l-1}; \Phi_{l}), l \in [1,L],$

이전 레이어 $f(;\Phi_{l-1})$의 출력, $h_{t}^{l-1}$를 입력으로 받아, 현재 레이어 $f(;\Phi_{l})$를 지나, $h_{t}^{l}$를 출력함.

이때, $t$는 타임스텝인데, 그냥 $n$번째 토큰으로 이해해도 됨.

결국 $t$가 증가하려면, (즉, 다음 토큰으로 넘어가려면) 모든 레이어를 지나, $y_{t}$를 출력해서, 그걸 다시 $f(;\Phi_{1})$에 넣는것 임.

2.2. 기존 non-RT로 RT를 initialized

Looping 전략

Gemma 2B: 18 layers 를 CYCLE 전략으로 2배 압축 ⇒ $L=18, B=2$

⇒ $B$는 loop 횟수

$$h_{t}^{l} = f(h_{t}^{l-1} ; \Phi'_{( (l-1) \text{ mod } L/B) + 1}), l \in [1,L],$$

$l = 1$: $((1-1) \text{ mod } 9) + 1 = 1 \implies \Phi'{1}$ 사용 $l = 2$: $((2-1) \text{ mod } 9) + 1 = 2 \implies \Phi'{2}$ 사용 $l = 3$: $((3-1) \text{ mod } 9) + 1 = 3 \implies \Phi'{3}$ 사용 ... $l = 8$: $((8-1) \text{ mod } 9) + 1 = 8 \implies \Phi'{8}$ 사용 $l = 9$: $((9-1) \text{ mod } 9) + 1 = 9 \implies \Phi'{9}$ 사용 $l = 10$: $((10-1) \text{ mod } 9) + 1 = 3 \implies \Phi'{1}$ 사용 (다시 처음으로)

⇒ 즉, loop 도는 주기에 관한 얘기!

초기화 전략

• Stepwise , Average , Lower 초기화 방식 중 Recursive Transformers 에는 레이어를 특정 간격으로 추출하는 Stepwise 방식이 가장 적합했습니다.

• 반면 Relaxed Recursive Transformers 의 경우에는 가중치 행렬들을 평균 내어 초기화하는 Average 방식이 가장 높은 성능 회복력을 보였습니다.

2.3. relaxed via LoRA

$$h^{l}{t} = f (h^{l-1}{t} ; \Phi'{ ( (l-1) \text{ mod } L/B ) + 1}, \Delta \Phi'{l}), \quad l \in [1, L]$$

$\Delta \Phi'_{l}$: LoRA module 파라미터

non-RT의 2개 레이어를 RT에서 1개로 공유해서 쓸 때, RRT는 해당 1개의 shared layer에서 2개의 LoRA 모듈이 필요. 그 2개 모듈은 각각의 loop에서 쓰임.

즉, 첫 번째 loop에서, 1번 LoRA 모듈, 두 번째 loop에서, 2번 LoRA 모듈.

첫 번째 루프 (Loop 1): 공유 가중치($W'$) 에 첫 번째 LoRA 모듈을 더해 연산합니다. $h=W'x + \Delta W'_1x = W'x + B_1A_1x$ ($\Delta W'_1 = W_1-W'$)

두 번째 루프 (Loop 2): 동일한 공유 가중치($W'$) 에 두 번째 LoRA 모듈을 더해 연산합니다. $h=W'x + \Delta W'_2x = W'x + B_2A_2x$ ($\Delta W'_2 = W_2-W'$)

LoRA rank

r ↑ ⇒ non-RT 스러움 r ↓ ⇒ RT 스러움

2.4. Continuous Depth-wise Batching으로 infernece speedup

Depth-wise batching은 RT의 특정으로 인해, 더 세밀히 batch control이 가능.

3. Experiments

3.1. Experimental Setup

model

• Gemma 2B
• TinyLlama 1.1B
• Pythia 1B

를2배 압축.

後 SlimPajama로 uptraining.

3.2. Non-Recursive Model Baselines

Full-size model

Pretrain dataset과 Uptraining dataset이 다르거나, Pretrain dataset 품질 > Uptraining dataset 품질

⇒ 성능 저하

⇒ Gemma, Pythia는 dataset 품질이 SlimPajama(Uptraining dataset) 품질보다 좋음

⇒ non-RT를 SlimPajama로 같은 토큰 만큼 uptraining vs. RRT를 SlimPajama로 같은 토큰 만큼 uptraining 으로 비교함.

3.3. 결과

3.4. RT의 초기화 방법

Stepwise가 짱임.

3.5. RRT의 초기화 방법

• Stepwise: Residual ($W_{orig} - W'_{shared}$)이 너무 큼 ⇒ LoRA로 매꾸기 힘듦
• Average: Residual이 상대적으로 작음 ⇒ LoRA로 매꾸기 good

Average가 짱임.

3.6. Uptraining & KD

생략

3.7 EE 가능성

가능

3.8. Continuous Depth-wise Batching

굿.

EE에 대한 내용

Abstract

EE 장점

1. 추론 가속화
2. 기울기 소실 문제 완화
3. overfitting, overthinking 감소

1 Introduction

생략

2 Motivation and Venefits of Early-exit DNNs

1. 추론 가속화
2. vanishing gradient problem 해결
3. overfitting problem 해결
4. 병렬 구현의 어려움 해결

2.1 Speed Up the DNN Inference

쉬운 입력 ⇒ 초기 layer에서 추출된 특징을 사용 ⇒ 초기 EE

어려운 입력 ⇒ 출력 layer까지 처리 필요

2.2 An Alternate Solution to the Overfitting Challenge of DNNs

• EE의 요소 및 방법: 모든 측면 분기(Side Branches)의 가중 손실(Weighted Losses)을 단일 최적화 문제로 결합하여 공동으로 학습시킵니다.
• 해결 과제 및 효과: 각 분기가 상호 정규화 요소(Regularizer)로 작용함으로써 과적합(Overfitting) 문제를 해결하고 모델의 일반화 성능(Genralization) 및 강건성(Robustness)을 향상시킵니다.

2.3 Rectify Overthinking Issue of DNNs

단순 입력 ⇒ EE ⇒ overthinking X

2.4 Mitigate the Vanishing Gradients Problem of DNNs

EEDNN은 네트워크의 심층부까지 가기 전, 중간 종료 지점을 통해 역전파 경로를 단축시키고, 하위 레이어에 즉각적인 판별 목표를 부여함으로써 기울기가 사라지는 것을 방지하고 학습의 효율성을 극대화합니다.

2.5 Provide a Multi-tiered Platform for Training and Deploying DNNs

생략

3. Related Works

생략

4 Early-Exit DNN-Architecture

4.1 Conventional DNN Architecture

생략

4.2 Early-exit DNN Architecture

EE branch: Conv layer, FC layer 또는 이들의 조합

4.3 Design Constraints of Early-exit DNNs

4.3.1 분기 구조

• 분기마다 동일한 구조 (대다수)
  • 단일 FC layer (대다수)
  • Pooling layer
  • Max-pooling layer
  • Mixed max-average pooling layer
  • 단일 conv layer + fc layer
  • 여러 conv layer
  • 조합
• 분기마다 다른 구조
  • 보통 backbone의 얕은 conv layer는 고차원 특징(이미지 전체에 대한 특징)을 출력하는데, 그래서 더 복잡한 분류기 필요

4.3.2 분기 위치 및 개수

• 균등 분기 배치
• 불균등 분기 배치

4.3.3 조기종료 정책

• 규칙 기반 정적 종료 정책
• 학습 기반 동적 종료 정책

5. 학습

• Joint : Backbone + side branch 포함 전체를 하나로 통합해 학습
• Branch-wise: 각 side branch를 backbone 레이어들과 함께 반복적으로 학습
• Separate: side branch를 독립적인 분류기로 취급하여 각각 별도로 학습
• 2-stage: 먼저 backbone 학습 후 freeze하고, side branch 별도 학습
• KD: side branch가 student, backbone의 최종 출력이 teacher

5.1 Joint

5.2 Branch-wise

1단계 (초기 학습): $L_1, L_2$ (백본 레이어)와 $E_1$ (첫 번째 측면 분기)을 함께 학습시킵니다. 학습이 완료되면 $L_1, L_2$의 가중치($\theta_1, \theta_2$)를 고정(Freeze)합니다. 이제 이 값들은 더 이상 변하지 않습니다.

2단계 (확장 학습): 이제 $L_3, L_4$와 $E_2$ (두 번째 측면 분기)를 학습시킬 차례입니다.이때 $L_1, L_2$는 업데이트하지 않고, 여기서 나온 출력값(Feature map)을 $L_3$의 입력으로 그대로 사용합니다. 즉, 오차 역전파($Backpropagation$)가 $E_2 \rightarrow L_4 \rightarrow L_3$까지만 일어나고, $L_2$와 $L_1$으로는 전달되지 않습니다.

5.3 Separate

Branch-wise와 달리, freeze 안함

5.4 2-stage

그림 7에 나타난 바와 같이, 백본 레이어($L_1$에서 $L_n$ 및 $F_E$)가 먼저 학습됩니다. 그 후 해당 파라미터들은 각 분기를 독립적으로 학습시키는 데 사용됩니다. 예를 들어, 측면 종료 분기 $E_1$은 고정된 $L_1$과 $L_2$의 파라미터를 사용하여 학습됩니다. 마찬가지로, 측면 종료 분기 $E_2$는 고정된 $L_1, L_2, L_3, L_4$의 파라미터를 사용하여 학습되며, 이 과정이 모든 분기에 대해 반복됩니다.

backbone이 이미 학습되어 있을 때, ㅈㄴ 유용.

5.5 KD

최종 출력 레이어가 높은 acc를 달성했어야, 의미가 있음.

5.6 Hybrid

조합해서 사용.

6 Inference

• 규칙 기반 정적 종료 정책
• 학습 가능 동적 종료 정책

6.1 규칙 기반 정적 종료 정책

• Entropy
• Maximum softmax probability
• 사용자 정의 점수 함수

를 통해 confidence score 측정

Threshold vector $T= [T_1, T_2, ..., T_i, ..., T_{n-1}]$에서 $T_i$는 $i$번째 side branch의 threshold를 나타냄.

즉, 특정 시점의 threshold 값 보다 confidence scored와 비교해서, EE!

Entropy

인데, 값의 분포가 skewd 되어 있을수록, 낮은 엔트로피(불확실성)를 갖게됨. ⇒ $H(p) < Threshold$ 이면 EE!

Softmax 기반

각 분기의 fc layer 출력 벡터에 softmax 적용한 것 ⇒ $Softmax(v) < Threshold$ 이면 EE!

사용자 정의 점수 함수

6.2 학습 가능 동적 종료 정책

6.3 비교

• 엔트로피 비교 및 소프트맥스 확률 비교와 같은 정적 방법(Static methods)은 미리 정의된 임계값에 의존하기 때문에 학습 복잡도와 추론 지연 시간이 낮습니다. 이는 구현이 간단하고 실행 속도가 빠르다는 장점이 있습니다. 그러나 적응력이 부족하기 때문에 다양한 입력 특성에 대한 강건성은 낮게 나타납니다.

• 반대로 종료 선택 컨트롤러 및 강화 학습 기술과 같은 동적 방법(Dynamic methods)은 정교한 학습 구성 요소로 인해 학습 복잡도가 높습니다. 하지만 최적의 종료 지점을 학습하고 변화하는 입력과 조건에 적응함으로써 강건성 면에서 탁월합니다. 이러한 동적 방법은 계산 집약적이지만, 더 큰 유연성과 맥락에 기반한 의사 결정을 제공하여 다양한 시나리오에서 최적의 성능을 보장합니다.

• 맥락 기반 추론 이력 및 목적 함수 기반 추론 방법은 이력 데이터와 학습된 목적을 활용함으로써 중간 정도의 추론 지연 시간과 높은 강건성 사이의 균형을 맞춥니다.

대부분의 방법은 높은 확장성을 유지하여 다양한 규모에 적응 및 배포가 가능하지만, 동적 접근 방식은 계산 자원 수요가 더 높을 수 있습니다.

1. H module과 L module이 같나?

   • 같으면 TRM

     • 세팅: .sh 에서, arch=trm // .sh 에서, arch=trm으로 하더라도, arch.H_layers=T로 하면, HRM과 같이 됨.
     • 결과: all_config.yaml 에서, H_layers: 0

   • 다르면 HRM

     • 세팅: .sh 에서, arch=hrm
     • 결과: all_config.yaml 에서, H_layers: 4

2. 1-step gradient 유무?

   • 없으면, TRM

     • 세팅:

   • 있으면, HRM

     • 세팅:

3. ACT에서 continue_loss가 있나?

   • 없으면, TRM
     • 세팅: .sh 에서, arch=trm
     • 결과: all_config.yaml 에서, no_ACT_continue: True
   • 있으면, HRM
     • 세팅: .sh 에서, arch=hrm
     • 결과: all_config.yaml 에서, no_ACT_continue 항목 X

4. Recurrent 횟수

   • T = 3, n = 6 TRM
     • 세팅: .sh 에서, *arch.H_cycles=3 arch.L_cycles=6 *
     • 결과: all_config.yaml 에서, H_layers: 3, H_layers: 6
   • T = 2, n = 2 HRM
     • 세팅: .sh 에서, *arch.H_cycles=2 arch.L_cycles=2 * // TRM 논문에서는 T=3, n=3이 optimal이라고 함.
     • 결과: all_config.yaml 에서, H_layers: 2, H_layers: 2

5. Supervision 횟수

   • $N_{sup} = 16$ 인가?
     • 세팅: .sh 에서, 별다른 세팅 X
     • 결과: all_config.yaml 에서, halt_max_steps: 16

6. Transformer block 수

   • 2, TRM
     • 세팅: .sh 에서, arch.L_layers=2
     • 결과: all_config.yaml 에서, H_layers: 0, L_layers: 2
   • 4, HRM
     • 세팅: .sh 에서, arch.H_layers=4, arch.L_layers=4
     • 결과: all_config.yaml 에서, H_layers: 4, L_layers: 4

7. Attention-free?

   • O
     • 세팅: .sh 에서, 별다른 세팅 X
     • 결과: all_config.yaml 에서, mlp_t: True
   • X
     • 세팅: .sh 에서, 별다른 세팅 X
     • 결과: all_config.yaml 에서, mlp_t: False

8. EMA?

   • O

     • 세팅: .sh 에서, ema=True
     • 결과: all_config.yaml 에서, ema: true

   • X

     • 세팅: .sh 에서, 별다른 세팅 X
     • 결과: all_config.yaml 에서, ema: false

9. halt_exploration_prob?

   ⇒ ACT에서 최소 세그먼트 수($M_{min}$)(min_halt_steps)는 확률적으로 결정하는데, 그 확률이 halt_exploration_prob임

   // sudoku4x4에서, halt_exploration_prob이 0인데, // 그 뜻은 min_halt_steps가 0이라는 뜻이고, // halt는 is_last_step 이거나, q_halt_logits > 0 일 때만 한다는 뜻.

Model size

Model size에 관여하는 부분은 3가지가 있다.

1. H_layers
2. L_layers
3. Transformer Block(layer)의 크기

Transformer Block 하나의 크기는 로 총 3,407,872개

• HRM은 $H_layers = 4, L_layers = 4$ ⇒ 대략, $3,407,872 ⨉ 8 = 27,262,976 ≈ 27M$

• TRM은 $H_layers = 0, L_layers = 2$ ⇒ 대략, $3,407,872 ⨉ 2 = 6,815,744 ≈ 7M$

Emulated Depth

Default로

• HRM은 $T=2, N=2, layer=4$ ⇒ $Emulated Depth=24$

• TRM 논문에서 HRM $T=3, N=3, layer=4$이 optimal이라고함 ⇒ $Emulated Depth=48$

• TRM은 $T=3, N=6, layer=2$ ⇒ $Emulated Depth=42$

실험

1. H module, L module의 size를 다르게 한다면, 어떤 것의 size가 큰 것이 좋을까?

⇒ H module, L module의 size를 다르게 한다는 뜻은, model 구조가 HRM을 따라간다는 뜻

즉, $H_layers ≠ 0$.

HRM_module_with_TRM_training 실험

$T=3, N=3, layer=4$로 $Emulated Depth = 48$ 실험임.

HRM_module_with_TRM_training.sh

arch로 사용하는 hrm_module_with_trm_training은 다음과 같이 작성함 (config/arch/hrm_module_with_trm_training.yaml) 위 name 에서 사용하는 HrmModuleWithTrmTrainingModel_ACTV1은 다음과 같이 작성함(models/recursive_reasoning/hrm_module_with_trm_training.py)

그 결과 다음과 같은 all_config.yaml이 생성 됨

결과 모델:

HRM_heavyH_module_with_TRM_training 실험

HRM_heavyL_module_with_TRM_training 실험

2. H module, L module 중 어느 module이 quantization에 더 강건할까?

L module quantization

L_module이 forward pass에서 대다수를 차지한다.

$T=3, N=3, layer=4$로 $Emulated Depth = 48$ 로 예시를 들면,

에서 H module 횟수 : L module 횟수 = 1 : 3 임

⇒ L module을 quantization 할 때, H module을 quantization 할 때보다 3배 더 빠를 것임

3. EE

EE에 대해 참고할 것

고려사항

우선 제일 먼저, H_module 끝 마다 출력되는 hidden state를 뽑아서, 정답과 얼마나 유사한지 가시화 + 유사도 검사 즉, overthinking 여부 검사

⇒ 정답과 거의 유사하거나, 유사도가 높으면, EE 가능성 ↑

+

ACT에서 조기 훈련 종료 통계 조사

1. 전체 $N_{sup}=16$ 동안의 inference 시간 대비 중간 분류기(?)로 인한 overhead 비율 조사
2. 임계값 설정
3. 훈련법 설정
4. 신뢰도 정의
5. 일반화 가능성 일반적으로 임계값은 dataset-dependent함 하지만 내 실험에서는 크게 신경 안써도 ㄱㅊ을 것 같은 이유가 애초에 sudoku가 너무 domain-specific 함 만약 다른 일반적인 task에 대해서 모델을 확장한다면,
   1. 임계값 기반 적응: 미지의 도메인에 따라 임계값을 조정함
   2. 특징 기반 적응: 모든 계층에서 애초에 도메인 불변 특징을 학습하는 것

TRM에 대한 내용

Abstract

TRM은 HRM보다 훨씬 더 단순한 재귀적 추론 접근 방식으로, 단 2개의 레이어만 가진 하나의 초소형 네트워크를 사용하면서도 HRM보다 훨씬 더 높은 일반화 성능을 달성.

HRM: 27M TRM: 7M

1. Introduction

HRM

1) recursive hierarchical reasoning 2) deep supervision

Recursive hierarchical reasoning

$f_L$(고빈도)와 $f_H$(저빈도) 두 네트워크가 서로의 잠재 특징($z_L, z_H$)을 입력받아 재귀적으로 작동하며, 이는 뇌의 계층적 처리를 모방

Deep supervision

잠재 특징의 계산 그래프를 분리(detach)하여 다음 단계의 입력으로 재사용하고 단계마다 지도(학습)함으로써, 메모리 폭발 없이 매우 깊은 신경망을 흉내 냄

ARC 벤치마크 분석 결과, HRM 모델의 성능 향상은 모델 내부의 재귀적 순환(H/L 모듈) 때문 X

⇒ 모델 전체를 여러 번 반복하며 단계마다 지도(학습)하는 '심층 지도(Deep Supervision)' 방식이 거의 모든 성능을 이끌어낸 핵심 요인

2. Background

2.1. Structure and goal

• HRM은 입력 임베딩($f_I$), 하위($f_L$) 및 상위($f_H$) 순환 네트워크, 출력 헤드($f_O$)라는 4개의 학습 가능한 구성요소를 가짐

• 하위($f_L$) 및 상위($f_H$) 순환 네트워크는 RMSNorm, 편향 없음, 회전식 임베딩(RoPE), SwiGLU 활성화 함수를 갖춘 4계층 트랜스포머 아키텍처를 기반으로 함

2.2. Recursion at two different frequencies

$x \leftarrow f_I(\tilde{x})$ $z_L \leftarrow f_L(z_L + z_H + x)$ # 기울기 없음 $z_L \leftarrow f_L(z_L + z_H + x)$ # 기울기 없음 $z_H \leftarrow f_H(z_L + z_H)$ # 기울기 없음 $z_L \leftarrow f_L(z_L + z_H + x)$ # 기울기 없음 $z_L \leftarrow z_L.\text{detach()}$ $z_H \leftarrow z_H.\text{detach()}$ $z_L \leftarrow f_L(z_L + z_H + x)$ # 기울기 있음 $z_H \leftarrow f_H(z_L + z_H)$ # 기울기 있음 $\hat{y} \leftarrow \text{argmax}(f_O(z_H))$

2.3. Fixed-point recursion with 1-step gradient approximation

마지막 $f_L$ 및 $f_H$ 단계만 역전파하여 기울기를 근사하는 데 사용.

6개 중 마지막 두 단계의 기울기만 추적하는 것을 정당화하는 데 사용되며, 이는 메모리 요구량을 크게 줄여줌.

2.4. Deep supervision

유효 깊이(effective depth)를 향상시키기 위해 사용.

⇒ 이전의 잠재 특징($z_H$ 및 $z_L$)을 다음 순전파를 위한 초기화로 재사용

2.6. Deep supervision and 1-step gradient approximations replaces BPTT

• 심층 지도 + 1단계 기울기 근사 X: 384계층 전체에 대해 역전파(BPTT)를 한 번에 수행해야 하므로, 메모리 폭발이 일어나고 계산 비용이 엄청나게 비쌈

• 심층 지도 + 1단계 기울기 근사 O: 전체 깊이 대신 L_net과 H_net의 마지막 계산 블록에 대해서만 역전파를 N_sup번 반복하므로, 메모리 문제를 해결하고 계산 비용을 획기적으로 줄일 수 있음.

cf) $384 = (#{TransformerBlocks})(n + 1)TN{sup} = 4 \times (2 + 1) \times 2 \times 16 = 384$

2.7. Summary of HRM

생략

3. Target for improvements in Hierarchical Reasoning Models

3.1. Implicit Function Theorem (IFT) with 1-step gradient approximation

HRM은 6번의 재귀(recursions) 중 마지막 2번을 통해서만 역전파를 수행함.

저자들은 '암시적 함수 정리(IFT)'와 '1단계 근사(one-step approximation)'를 활용하여 이를 정당화함.

이 정리는 순환 함수가 고정점(fixed point)에 수렴할 때, 그 평형점(equilibrium point)에서 단일 단계로 역전파가 적용될 수 있음을 의미함.

But, "고정점(fixed-point)에 도달했다"라고 주장하며 수학적 "꼼수"(1단계 경사도 근사)를 사용하는데, 실제로는 고정점에 도달할 때까지 반복 계산을 하는 것이 아니라, "단순히 몇 번 순전파"를 실행하고 멈춤

3.2. Twice the forward passes with Adaptive computational time (ACT)

Q-러닝 목표는 정지 손실(halting loss)과 계속 손실(continue loss)에 의존한다.

계속 손실은 HRM을 통한 추가적인 순전파(extra forward pass)를 필요로 함.

⇒ 이는 ACT가 샘플당 시간을 더 효율적으로 최적화하는 반면, 최적화 단계(optimization step)당 2번의 순전파가 필요함을 의미

3.3. Hierarchical interpretation based on complex biological arguments

TBW

4. Tiny Recursion Models

4.1. No fixed-point theorem required

HRM: 1단계 기울기 근사를 활용하기 위해, $z_L$, $z_H$ 가 고정점에 수렴할 것이라고 가정.

But, 수렴 X

Let, (full recursion process) $=$ $z_L \leftarrow f_L(z_L + z_H + x)$ $...$ $z_L \leftarrow f_L(z_L + z_H + x)$ $z_H \leftarrow f_H(z_L + z_H)$ 일 때,

Full recursion process $T$번 중 마지막 $T$번째의 과정만 기울기를 계산해서 backward!!!

HRM vs. TRM

• HRM (1단계 근사 방식): $N_sup$ 루프의 각 단계에서, $T$번의 재귀 중 마지막 1번 재귀의 "마지막 2개 함수($f_L, f_H$)"에 대해서만 경사도를 계산.

• TRM (전체 역전파 방식): $N_sup$ 루프의 각 단계에서, $T$번의 재귀 중 마지막 1번 재귀의 "프로세스 전체($n+1$번의 net 호출)"에 대해 경사도를 계산.

4.2. Simpler reinterpretation of zH and zL

• HRM: 2개의 latent features 사용(H module, L module)

• TRM: 1개의 latent features 사용

그럼 H module의 생성 $z_H$는? ⇒ $y = z_H ← f_H(x, z_L, z_H= y)$

즉, 계층 구조는 필요 X 단순히 입력 $x$, 제안된 해답 $y$ (이전의 $z_H$), 그리고 잠재 추론 특징 $z$ (이전의 $z_L$)가 있을 뿐

입력 질문 $x$, 현재 해답 $y$, 현재 잠재 추론 $z$가 주어지면, 모델은 재귀적으로 잠재 $z$를 개선 ⇒ $z_i = f(x, y, z_{i-1})$

後 현재의 잠재 $z$와 이전 해답 $y$가 주어지면, 모델은 새로운 해답 $y$를 제안 ⇒ $y = f(y, z_i)$

코드를 보면, n loop(L_cycles)에서 net(z, y+z) 임을 알 수 있음.

다양한 feature에 대해 실험 해봤는데, 2개 feature가 최적임을 찾음 ※ TRM은 1개의 latent feature를 사용한다의 'feature'와 다른 feature임

1개의 laytent feature라는 것은 H module과 L module이 같은 가중치를 쓴다는 거고, 2개의 feature가 최적이라는 것은 z의 개수 1(y 1개, z 1개 ⇒ 2개 latent feature)이 최적이라는 뜻

4.3. Single network

위와 동일한 내용임.

4.4. Less is more

Model 확장을 위해, layer(Transformer block) 수를 늘려 capacity를 늘리려함

⇒ overfitting

⇒ 재귀 횟수(n)을 비례적으로 늘리면서 block 수를 줄였을 때, 일반화 성능 극대화

⇒ $#{TransformerBlock} \leftarrow \frac{#{TransformerBlock}}{2}$

⇒ $n \leftarrow n ⨉ 2$

4.5. attention-free architecture for tasks with small fixed context length

Attention의 경우 $L≫D$ 일 때, $[D, 3D]$ 행렬만으로 전체 시퀀스를 설명할 수 있음.

But, $L ≤ D$ 일 때, linear layer는 $[L, L]$ 행렬만 필요로 하므로 더 저렴함. (MLP-Mixer)

⇒ 작고, 고정된 context lenght에서 잘 작동

But,

긴 context suboptimal

4.6. No additional forward pass needed with ACT

HRM의 ACT는 "계속 손실"을 계산하기 위해 비효율적인 2차 순전파가 필요하여 훈련 속도가 느려짐

⇒ TRM은 "계속 손실"을 제거하고 "중단 손실"만 학습함으로써, 비용이 드는 2차 순전파 과정을 없앰

⇒

4.7. Exponential Moving Average (EMA)

생략

4.8. Optimal the number of recursion

HRM: $T=3, n=3$이 최적 TRM: $T=3, n=6$이 최적

// full recursion process의 gradient를 backward하기에 n이 크면 OOM

5. Results

Setups

Dataset

• Sudoku-Extreme Train: 1K Test: 423K

• Maze-Hard Train: 1K Test: 1K

• ARC-AGI1

• ARC-AGI-2

Augmentation

Shuffling, dihedral transformations, color permutation, dihedral-group, translations, flips, reflection 등 많이 함

Results

HRM에 대한 내용

HRM 관련 유용스

Abstract

기존 CoT는 방대한 data 요규량, 높은 latency ⇒ 인간의 hierarchical하고, multi-timescale한 처리 과정에서 영감을 받아, HRM

HRM은 한 번의 forward pass로 interdependent한 recurrent module 2개를 씀.

1. High-level module: 느리고, 추상적인 계획
2. Low-level module: 빠르고, 상세한 실행

1. Introduction

Reasoning 이 ㅣ필요한 작업에서, model depth를 증가시키면, 성능 ↑ But, 매우 깊더라도, 성능 최적 X 아래 Figure 2에서 볼 수 있듯, Scaling Depth를 늘려도, 109M 정도 부터 saturate됨.

⇒ CoT prompting

CoT prompting: 복잡한 문제를 간단한 중간 언어 단계로 분해하여 순차적으로 해결.

문제점: 단 한 번의 실수로도 전체 추론이 실패할 수 있고, 많은 토큰을 생성하여 느리고 비효율적.

⇒ latent reasoning Latent reasoning은 언어가 인간 의사소통을 위한 도구일 뿐, 생각 그 자체의 기질(substrate)은 아니라는 이해와 일치. 즉, 마지막에 출력만 언어로 하면 되지, 중간 언어 단계(CoT prompting)는 필요없다는 뜻.

⇒ 단순히 층을 쌓는 것
⇒ vanishing gradient 문제

⇒ Recurrent architectures
⇒ early convergence ⇒ 후속 계산 단계를 비활성(inert) 상태로 만듦 + BPTT

⇒ 층 多 + without BPTT ⇒ HRM

• 두 개의 coupled된 recurrent module(H, L module)
• One-step gradient approximation (BPTT 제거)

2. Hierarchical Reasoning Model

Let, input vector $x$, output prediction vector $\hat{y}$

Learnable component, input network $f_I(\cdot; \theta_I)$에 의해, input vectoer $x$는 $\tilde{x}$가 됨.

$\tilde{x} = f_I(x; \theta_I)$

Let, time step $i = 1, \ldots, N \times T$. ( N: $f_L(\cdot; \theta_L)$의 recurrent 횟수, T: $f_H(\cdot; \theta_H)$의 순환 주기 )

H module은 해당 사이클(N번의 $f_L(\cdot; \theta_L)$ recurrent)이 끝날 때의 L module의 최종 상태를 사용하여, 사이클당 한 번만 H의 hidden state $z_H$를 update.

$$z_L^i = f_L(z_L^{i-1}, z_H^{i-1}, \tilde{x}; \theta_L) ,$$

$$z_H^i = \begin{cases} f_H(z_H^{i-1}, z_L^i; \theta_H) & \text{if } i \equiv 0 \pmod{T}, \ z_H^{i-1} & \text{otherwise}. \end{cases}$$

전체 $N$번의 사이클($N \times T$의 time step)이 끝난 후, H module의 hidden state $z_H^{N \times T}$을 output network $f_O(\cdot; \theta_O)$에 넣어, $\hat{y}$ 얻음.

$\hat{y}=f_O(z^{NT}_H; \theta_O)$

이 과정 한 번이 forward 한 번임.

Hierarchical convergence

표준 RNN의 문제점: Hidden state가 너무 빨리 고정점(fixed point)으로 수령하여, 후속 계산이 중단되고 NN의 유효 깊이가 제한.

⇒ 계산 능력을 보존하기 위해 hidden state가 매우 느리게 수렴하도록 NN를 설계.

But, 수렴 속도를 억지로 늦추면, 시스템 전체가 불안정해질 위험 有.

⇒

• H module: 하위 계산 결과($z_L$)를 통합하여, 전반적인 전략을 지시하고, L module이 새로운 계산을 시작하도록 컨텍스트($z_H$)를 업데이트.

• L module: H module이 설정한 컨텍스트($z_H$) 내에서 '국소 평형'에 도달할 때까지 집중적인 탐색 또는 정제 계산을 수행.

⇒ Early convergence를 방지 + 표준 RNN($T$ 단계)보다 훨씬 깊은 $NT$ 단계의 향상된 유효 깊이를 안정적으로 확보.

⇒ HRM이 (표준 RNN처럼 활동성이 급격히 감소하는 것과는 대조적으로) 많은 단계에 걸쳐 높은 계산 활동성(forward residual)을 유지하면서도 + 안정적인 수렴을 누릴 수 있음을 보여줌.

Approximate gradient

RNN은 BPTT를 사용. BPTT는 forward pass에서 얻은 모든 hidden state를 저장했다가, backward pass 중에 gradient와 결합해야함.

⇒ memory 부담 ↑ BPTT 참고자료

만약, RNN이 fixed point로 convergence된다면, 그 equilibrium point에서 단일 단계로 backward해서, BPTT 피할 수 있음(왜?...)

⇒ 즉, 아래 그림의 빨간 포인트들에서 단일 단계 backward할 수 있다는 뜻.

⇒ one-step approximation ⇒ 각 module의 마지막 hidden state의 gradient만 사용하고, 다른 hidden state는 상수 취급

Deep supervision

Let, 한 번의 forward pass를 segment. 총 segment 수를 M이라고 한다면, 각 segment $m \in {1, ..., M}$.

Let, segment $m$에서의 최종 H/L module의 hidden state를 $z^m=(z^{mNT}_H, z^{mNT}_L)$.

⇒ $$(z^m, \hat{y}^m) \leftarrow \text{HRM}(z^{m-1}, x; \theta)$$

$$L_m \leftarrow \text{LOSS}(\hat{y}^m, y)$$

$$\theta \leftarrow \text{OPTIMIZERSTEP}(\theta, \nabla_{\theta}L_m)$$

+

Hidden state $z^m$을 다음 segment로 넘길 때 계산 그래프에서 '분리(detach)'하여, gradient가 이전 segment로 backpropagation되지 않도록 차단!!!

⇒ H module에 더 빈번한 피드백을 제공 + 메모리 부담 X

'Approximate gradient' vs. 'Deep supervision'

• 근사 기울기 (Approximate Gradient): detach로 인해 과거의 계산 기록이 의도적으로 잘린 상태에서, 각 지도 단계(step)마다 계산되는 기울기 값.

• 심층 지도 (Deep Supervision): detach로 계산 그래프를 분리하며 (순전파 → 손실 계산 → 역전파) 과정을 n_supervision번 반복하는 전체 훈련 기법.

Adaptive computational time (ACT)

뇌는 과제 복잡성과 잠재적 보상에 따라 학습 시간을 동적으로 조절함.

⇒ Training 中에만 adaptive halting strategy를 HRM에 통합.

⇒ 학습 中 segment 수를 동적으로 조절.

Q-head를 통해, H module의 최종 상태($z_H^{mNT}$)를 사용하여, halt, continue 행동의 Q-value를 예측.

$Q-value = \hat{Q}^m = \sigma(\theta_Q^T z_H^{mNT}) = (\hat{Q}^m_{halt}, \hat{Q}^m_{continue})$

if ((m ≥ $M_{max}$) or (($\hat{Q}{halt}$ > $\hat{Q}{continue}$) and ($m ≥ M_{min}$))) ⇒ halt

cf) $M_{max}$: 고정 하이퍼 파라미터, $M_{min}$: 확률적으로 랜덤하게 1 ~ $M_{max}$에서 선택.

⇒ Q-value를 통해, halt 여부 결정

Q-head update 과정

$\hat{G}^m_{halt}$: "halt" action을 선택했을 때의 보상이고, $\hat{G}^m_{continue}$: "continue" action을 선택했을 때의 미래의 보상임

$BinaryCrossEntropy(\hat{Q}^m, \hat{G}^m)$: "action"에 대한 loss

⇒ 예측($\hat{Q}^m$)이 실제 목표($\hat{G}^m$)와 비슷해 지도록 함!!!

전체 $L_m^{ACT} = \text{LOSS}(\hat{y}^m, y)+BinaryCrossEntropy(\hat{Q}^m, \hat{G}^m)$

cf) 이때, halt action을 했다고 해서, 배치의 모든 샘플이 새로운 샘플로 대체 X ⇒ 해당 샘플만 대체 O

3. Results

3.1 Benchmarks

ARC-AGI Challenge

Sudoku-Extreme

Maze-Hard

3.2 Evaluation Details

3.3 Visualization of intermediate timesteps

4. Brain Correspondence

5. Related Work

Reasoning and algorithm learning

Brain-inspired reasoning architectures

Hierarchical memory

6. Discussions

Turing-completeness of HRM

Reinforcement learning with chain-of-thought

Linear attention

7. Conclusion

RepQ-ViT에 대한 내용

Abstract

P1. post-LayerNorm activation의 inter-channel variation S1. Channel-wise Q @ quantization step. Layer-wise Q @ inference step.

P2. Power-law distribution을 가진 post-Softmax activation S2. $log_{\sqrt{2}}$ Q @ quantization step. $log_2$ Q @ inference step.

1. Introduction

전통적인 'Q-infernece 종속' paradigm은 유일한 선택일까? ⇒ Q, infernece step을 decoupling

P1. post-LayerNorm activation의 inter-channel variation S1. Channel-wise Q @ quantization step. Layer-wise Q @ inference step. ⇒ Q step에서 channel-wise Q하고, I step에서 layer-wise로 reparameterization하여, LayerNorm의 affine factoer, 다음 layer의 weight를 조정해서 acc ↑

P2. Power-law distribution을 가진 post-Softmax activation S2. $log_{\sqrt{2}}$ Q @ quantization step. $log_2$ Q @ inference step.

2. Related Works

생략

3. Methodology

3.1 Preliminaries

Hardware-friendly quantizers

Uniform Q

$log_2$ Q

3.2 Scale Reparam for LayerNorm Activations

⇒ 단순히 통일된 Q SF를 적용하는 layer-wise Q는 위와 같은 큰 inter-channel variation에서 acc ↓

⇒ Q step: channe-wise Q ⇒ $s \in \R^D$, $z \in \Z^D$를 얻음

I step: $s, z$를 $\tilde{s} \in \R^1, \tilde{z} \in \Z^1$로 reparam.

이때, $\tilde{s}, \tilde{z}$k는 당연히 미리 계산됨(PTQ니까...)

$\tilde{s} = E[s], \tilde{z} = E[z]$

라고 하면,

이고, 이러한 연산은 Layernorm의 affine factor를 다음과 같이 조정함으로써 달성 ㄱㄴ.

위에처럼 하면, $s, z$를 reparam하지만, activation 분포를 이렇게 이동시킴.

⇒

이런 식으로 뒤에 layer의 weight에 오차를 보상하게 함.

3.3 Scale Reparam for Softmax Activations

FQ-ViT의 $log_2$ Q는 중요한 소수의 attention 값이 큰 부분에 적은 bin을 할당.

⇒ acc ↓

⇒ $log_{\sqrt{2}}$ Q

But HW frendly X (∵ $log_2$ Q처럼 bit-shifting 연산 X)

⇒ Q step: $log_{\sqrt{2}}$ Q I step: $log_2$ Q

⇒

s와 비교할 때, reparam된 $\tilde{s}$는 I step에서 dequant 절차 또한 효율적인 bit-shift로 할 수 있게함! // 물론 $A^{(\Z)}$가 홀수일 때, $\tilde{s}=s \cdot [\sqrt{2}]$ 여서, // $\hat{A}=2$를 $\lfloor -\frac{A^{(\Z)}}{2}\rfloor$ 만큼 shift하고, $\tilde{s}$를 곱하긴 함.

// 짝수일 떄도, $s$가 FP32이면, fp연산을 하긴 함...

4. Experiments

4.2. Quantization Results on ImageNet Dataset

4.3. Quantization Results on COCO Dataset

4.4 Ablation Studies

4.5 Efficiency Analysis

ARQ-ViT에 대한 내용

Abstract

P1. 기존 SF 지표(calibration metric)가 low-bit에 대한 Q 영향을 측정하는데 부정확 S1. Block-wise로 Q perturbation을 감지 + 중대한 오차에 대한 집중도를 높이기 위해, 사소한 오차를 무시 ⇒ Bottom-elimination Blockwise Calibration

P2. 기존 Q 패러다임이 Softmax의 power-law-distribution에 친화적이지 X S2. softmax의 power-law-distribution을 유지하기 위해, Matthew-effect Preserving Quantization 설계

1 Introduction

S1. Blockwise Bottom-elimination Calibration, BBC는 block-wise하게 Q error를 감지 + 사소한 error에 해당하는 2차 그래디언트를 생략하여, 중대한 오차에 집중

S2. Mattew-effect Preserving Quantization, MPQ는 softmax의 목적에 맞게, 큰 값(중요한 값)에 더 많은 bin 할당

2 Related Work

생략

3 Method

3.1 Preliminaries

PTQ4ViT는 최적의 SF를 결정하기 위해, Hessian guided metric을 사용.

위 Q로 인한 작업 손실의 기댓값을 최소화하는 SF를 찾아야함.

⇒ $min_\Delta(E[L(\hat{x})]-E[L(x)])$

⇒ ≈ $min_\Delta((\hat{O}-O)^T\bar{H}^{(O)}(\hat{O}-O))$

이후 구간을 $n$등분한 $\Delta_x$의 search space를 순회하며, ≈ $min_\Delta((\hat{O}-O)^T\bar{H}^{(O)}(\hat{O}-O))$를 만족하는 $\Delta_w, \Delta_a$를 교대 탐색!!!

3.2 Blockwise Bottom-elimination Calibration

PTQ4ViT의 hessian guided metric은 low-bit에서, layer-wise한 최적화가 block level에서의 Q를 감지할 수 없어, 부정확 + 낮은 비트에서 오차(hessian matrix)가 클 수 밖에 없는데, dense한 hessian matrix는 사소한 오차'까지 모두 보려다가, '중대한 오차'에 대한 집중도를 잃음

⇒

4bit에서 인접한 후보 SF들끼리 loss 차이가 큼 + 너무 뾰족 뾰족해서, 최적 찾기 어려움

⇒ Blockwise Bottom-elimination Calibration

1. Blockwise하다?

어떤 b번째 block의 input/output을 $a^b, O^b$라고 할 때, $O^b={w_L^b}^T{w_{L-1}^b}^T...{w_1^b}^Ta^b$ 이다. 또한 $l$번째 layer가 보정될 때, $L$번째부터 $l$번째 layer까지의 복합 레이어로 간주될 수 있고, 그 weight와 activation은 다음과 같음

weight 보정을 예로들면, $b$번째 block의 $l$번째 layer에 대한 2차 항은 다음과 같음

이때,

⇒ (7)을 보면, 결국, 어떤 block의 어떤 layer든 해당 block의 2차 정보로 최적의 SF를 결정할 수 있음 ⇒ Blockwise!!!

2. Bottom-elimination?

bit-width ↓ ⇒ Q error ↑ ⇒ '사소한 error' ↑ ⇒ '중대한 error' 관심 ↓

⇒ 최종 출력에 perturbation을 일으키는 큰 error에 더 주목하기 위해, '최하위 제거'메커니즘을 적용한, sparse hessian matrix!!!

⇒

⇒

3.3 Matthew-effect Preserving Quantization

다른 연산들과 달리 softmax의 output distribution은 비대칭적이고 극도로 불균형한 power-law probability distribution을 따름.

이전 많은 연구들 PTQ4ViT(twin-Q), FQ-ViT(log방법)는 모두 낮은 softmax 값(대다수의)에 많은 bin을 부여하고 싶어함.

But 그렇게 되면, softmax의 원래 기능(작은 값은 더 작게, 큰 값은 더 크게(∵exponential function 때문에))을 상실...

⇒ 중요한 큰 값들의 정보 손실...

⇒ Matthew-effect Preserving Quantization, MPQ

이는 양자화 전후의 상호 정보량 최대화를 순수하게 추구하는 대신, 양자화 과정 동안 소프트맥스 출력의 마태 효과를 유지

⇒ 즁요한 값은 더 fine하게 Q (그림 5(d))

3.4 Framework of APQ-ViT

4. Experiment

4.1 Settings

생략

4.2 Ablation Study

4.3 Comparison on Classification Task

4.4 Comparison on Object Detection Task

5. Conclusion

생략

FQ-ViT에 대한 내용

Abstract

P1. LayerNorm input의 심각한 inter-channel variation ⇒ performance ↓ S1. Power-of-Two Factor

P2. Attention map에서의 extreme skwed distribution ⇒ performance ↓ S2. Log-Int-Softmax

1 Introduction

Post-Training Quantization for Vision Transformer에서 LayerNorm과 Softmax를 Q하는 것은 심각한 acc ↓를 가져옴을 보임. ⇒ Model이 fully Q되지 않아, HW에서 FP 연산 장치를 유지 + resource 소모 ↑ + inference speed ↓

⇒ P1. LayerNorm input에서 significant한 inter-channel variation 발견. 일부 channel의 범위는 median의 40배를 초과 ⇒ acc ↓ S1. LayerNorm의 input을 Q하기 위한, PTF

P2. Attention-map 값이 extremely skwed distribution ⇒ 대부분 값 0 ~ 0.01, 소수 높은 attention 값만 1에 가깝 ⇒ acc ↓ S2. 대부분의 작은 값에 더 높은 Q 해상도를 제공하기위해, LIS

2 Related Work

생략

3 Proposed Method

3.1 Preliminary

Uniform Quantization

⇒ $s, zp$는 $X$의 하한($l$)과 상한($u$)에 의해 결정.

Log2 Quantization

3.2 Power-of-Two Factor for LayerNorm Quantization

LayerNorm은 forward step에서 $μ_X, σ_X$를 계산하여, input $X$를 normalized함. 이후, affine param $\gamma, \beta$가 nomalized된 input을 재조정.

P1. ViT에서 LayerNorm input의 inter-channel variation이 큼.

⇒ 모든 channel에 동일한 Q param을 적용하는 layer-wise Q는 acc ↓

⇒ Group-wise, Channel-wise Q

But

여전히 FP 영역에서 평균과 분산 계산...

⇒ Channel-wise하게, factor $\alpha$를 도입하는, PTF

⇒ Channel별 고유한 factor $\alpha$를 갖고, layer-wise 고유한 $s, zp$를 갖게됨!!!

⇒ Infernce시 $s, zp$는 미리 추출될 수 있음 + LayerNorm input이 Q되서, 평균, 분산이 int 영역에서 계산 됨!!!

3.3 Log-Int-Softmax for Softmax Quantization

P2. Attention map의 상당히 많은 값이 상당히 작은 값에 분포. 1에 가까운 큰 값은 이상치로 소수에 불과.

이렇게 많은 값들($10^{-4}$ ~ $10^{-2}$의 값들)에 하나의 구간만 할당하는 uniform Q는 않좋음.

⇒ $log_2$ 방식은 이 값들을 위해, 더 많은 구간을 할당 可能 + 지수 함수 (Softmax의 exp($X$))의 다항식 근사치인 i-exp

⇒ integer-only + fast + low-cost

4 Experiments

4.1 Implementation Details

• Calibration: ImageNet, COCO에서 1000개 training images
• Weight Q: symmetric channel-wise MinMax
• Activation Q: asymmetric layer-wise MinMax
• PTF의 K = 3

4.2 Comparison with State-of-the-art Methods

4.3 Ablation Studies

⇒ Softmax를 낮은 비트로 Q 하려면, LIS 必

4.4 Visualization of Quantized Attention Map

⇒ Softmax를 낮은 비트로 Q 하려면, LIS 必

5 Conclusions

생략

PTQ4ViT에 대한 내용

Abstract.

Softmax, GeLU 이후 activation의 distribution ≠ gaussian distribution + MSE, cosine distance를 통한, optimal SF 결정은 정확 X

⇒ Twin uniform Q + Hessian guided metric을 통한, SF optimizing

1 Introduction

생소했던 개념들

• Self-attention module은 global information을 포착
• PTQ는 label 없는 calibraion images를 사용

문제1. 특수 분포 activation

1. Softmax: Softmax를 거친 값들은 [0, 1] 범위에서 매우 unbalanced된 distribution을 갖음. 대부분 값이 0에 매우 가까움. + 큰 값의 수는 매우 적지만, 이 값들은 두 patch 간 높은 연관성(attention)을 의미. ⇒ 매우 중요한 값 ⇒ 이 큰 값을 Q 범위에 포함시키기 위해, 큰 SF 要 ⇒ 大多數의 작은 값들이 0으로 Q됨 ⇒ Q Error ↑
2. GeLU: GeLU를 거친 값들은 asymmetrical한 distribution 갖음. 양수 값들은 분포가 넓음. + 음수 값들은 skewed됨.

⇒ 이런 특수 distribution을 해결하기 위해, Twin uniform Q 이는 값들은 두 범위로 분리하여 각각 Q.

문제2. SF를 결정하는 지표 정확 X

MSE, Cosine distance, Pearson correlation coefficient등이 있었음. ⇒ local information만 사용 ⇒ suboptimal ⇒ Hessian guided metric

2 Background and Related Work

2.1 Vision Transformer

생략

2.2 QuantizationFang, J., Shafiee, A., Abdel-Aziz, H., Thorsley, D., Georgiadis, G., Hassoun, J.: Post-training piecewise linear quantization for deep neural networks. In: ECCV (2020)

3 Method

3.1 Base PTQ for Vision Transformer

생략

3.2 Twin Uniform Quantization

• Softmax 후 활성화 분포: 활성화 분포가 매우 불균형하며(unbalanced), 대부분의 값이 0에 매우 가깝고 소수의 값만이 1에 가깝습니다.
• GELU 후 활성화 분포: GELU 함수를 거친 값들은 높은 비대칭 분포(highly asymmetric distribution)를 가지며, 무한정한 양수 값들은 범위가 큰 반면 음수 값들은 분포 범위가 매우 작습니다.

⇒ Twin uniform Q

Twin uniform Q는 2개의 Q 범위 $R_1, R_2$를 갖음. 각각은 다른 SF $\Delta_{R1}$과 $\Delta_{R2}$에 의해 제어됨. ⇒

• Softmax 後 값에 대한 범위 설정 작은 $\Delta_{R1}^{softmax}$를 사용하여, $R_1=[0, 2^{k-1}\Delta{R1}^{softmax})$를 잘게 Q. + 상대적으로 큰 $\Delta_{R2}^{softmax}=\frac{1}{2^{k-1}}$로 고정하여, $R_2=[0, 1]$을 크게 크게 Q. ∵ calibration 영향을 피하기 위해.

• GeLU 後 값에 대한 범위 설정 작은 $\Delta_{R1}^{GeLU}=\frac{0.170}{2^{k-1}}$로 고정하여, $R_1=[-2^{k-1}\Delta_{R1}^{GeLU}, 0]$를 잘게 Q. ∵ 음수 전체 범위 $[GeLU_{min}, 0]$를 포괄하려고. + 상대적으로 큰 $\Delta_{R2}^g$를 사용하여, $R_2=[0, 2^{k-1}\Delta_{R2}^{GeLU}]$를 크게 크게 Q.

⇒ 이제 optimal $\Delta_{R1}^{softmax}, \Delta_{R2}^{GeLU}$만 찾으면 됨.

⇒ ![] (https://velog.velcdn.com/images/sangjun_moon/post/ecdb6670-ce8c-4f84-8d80-f5bfdd8f274a/image.png)

그림4에서, optimal 빨간색 범위만 찾으면 됨!!!

3.3 Hessian Guided Metric

Optimal SF를 찾기위해, 여러 metric(MSE, Cosine, Pearson, ...)으로 layer-wise + greedy하게 했었음

물론 Q 前/後 마지막 layer output 간의 distance ($L = CE(\hat{y}, y)$)를 사용하는 것이 PTQ에서 가장 정확. But, forward 너무 많이 해야됨.

⇒ Hessian Guided Metric

weight $W$를 변수로 취급할 때, Loss의 expectation은 $E[L(W)]$

Q는 weight에 작은 perturbation $\epsilon$을 갖게 함. ⇒ $\hat{W}=W+\epsilon$

⇒ Taylor series expansion을 통해, Q가 Loss에 미치는영향을 분석 可能

⇒

optimal SF를 찾아야 함!!! ⇒ $min_{\Delta}(E[L(\hat{W})]-E[L(W)])$을 layer-wise하게 구성해야함.

⇒

∵ $\bar{g}(W)$는 잘 training된 model에서 0을 갖고, $\epsilon = (\hat{O_l}-O_l)$이어서.

3.4 PTQ4ViT Framework

1단계: Output($y$) 및 gradient 수집

1 ~ 6 line

2단계: Optimal SF찾기

7 ~ 14 line

문제점

10 ~ 13 line에서 특정 layer를 Q하면, $\epsilon = (\hat{O}^l-O^l)$이 바뀌는데, 그럼 $\bar{H}(W) = diag((\frac{∂L}{∂O_1^l})^2, ..., (\frac{∂L}{∂O^l_{|O_l|}})^2)$ 도 매번 새로 계산해야 되는데, 그렇게 안함... ∵ GPU 메모리 줄이려고...

4 Experiments

4.1 Experiment Settings

• Softmax 後: $\Delta^{softmax}_{R1}$의 탐색 공간은 $[\frac{1}{2^k}, \frac{1}{2^{k+1}}, ..., \frac{1}{2^{k+10}} ]$
• $\alpha, beta$, # Round
  • BasePTQ
    • $\alpha = 0.5, beta = 1.2$, # Round = 1
  • PTQ4ViT
    • $\alpha = 0, beta = 1.2$, # Round = 3
• ImageNet
• Calibration 32개

4.2 Results on ImageNet Classification Task

4.3 Ablation Study

Conclusion

생략

PTQ for ViT에 대한 내용

Abstract

• Attention mechanism 기능을 보존하기 위해, Q된 attention map의 relative order를 유지하게하는 ranking loss를 최종 Q 목적 함수에 도입
• MHA의 Attention map과 MLP의 output feature 간 nuclear norm을 활용하여, MixQ

1 Introduction

이 논문 이전의 PTQ 방법들은 CNN이나 RNN을 위해 설계됨 ⇒ transformer의 특징 고려 X

⇒

• Attention map의 relative order를 유지하기 위해, ranking loss 도입
• Attention map과 MLP의 output feature의 sensitivity 측정을 위해, nuclear norm을 사용해서, MixQ
• Bias correction을 통한, Q error 보정

2 Related Works

생략

3 Methodology

• Linear layer를 위한 similarity-aware quantization + ranking-aware quantization
• Accuracy ↑을 위한 bias correction
• Nuclear norm을 통한 MixQ

3.1 Preliminaries

에서 $X$의 가로 한 줄이 patch 하나에 대한 정보

에서 $X_lW_l^Q$의 결과 $Q_l$ 또한 가로 한 줄이 patch 하나에 대한 정보

3.2 Ranking-Aware Post-Training Quantization

는 Q 함수이다.

는 Q된 출력이다.

scale factor ($\Delta_W^l, \Delta_X^l$)가 Q 결과에 큰 영향을 미침을 알 수 있다. ∵ 사실상 조절 가능한 것이 $Y$와 $\Delta$이기 때문에...

⇒ Calibration Dataset에서 생성된 weight와 activation에 대한 optimal scale factor를 찾자!!!

Self-attention layer는 전역적 관련성을 계산하는 CNN과 차별화를 갖는 구성 요소임.

Q 後 attention map의 relative order가 변경되는 것을 empirically 관찰 ⇒ 이는 심각한 performance degradation 발생

Relative order란? Attention map의 한 행 내에서, 임의의 두 값 $A_{kp}, A_{kq}$의 대소 관계를 의미

Q 後 attention map의 relative order가 변경됬다는 것은 Q 前에는 한 row에 ..., 3.61, 5.234, ..., 4.24, ... 이었는데, Q 後에는 한 row에 ..., 4, 5, ..., 4, ... 가 되면, $3.61 < 4.24$의 대소 관계가 사라짐

⇒ Ranking Loss 도입

⇒ Ranking Loss function을 통해, Q Attention 값의 쌍별 순서가 원본 attention 값의 순서를 유지하도록 함.

전체 최적화 목표

전체 Loss Runction은 다음과 같음

$\Gamma$는 pearson correlation coefficient로 Q 전/후 텐서의 유사도 측정 지표로 사용됨

Alternative searching method

위의 식(11)을 해결하기 위해, weight scale factor 최적화 後 input scale factor 최적화

Bias Correcition

Q로 인해 발생하는 오차를 출이기 위해, 최적의 scale factor를 찾고, biased 된 error를 보정

Q errors는 다음과 같아,

Q된 한 layer의 출력 $\hat{O}$의 error 기댓값은 다음과 같음

∵

⇒ $\hat{O}-\mathbb{E}[\hat{O}]$ 를 통해, 보정!!!

3.3 Nuclear Norm Based Mixed-Precision Quantization

MHA, MLP를 MPQ 대상으로 보고, 두 모듈의 sensitivity를 계산하기 위해, nuclear norm 사용

Nuclear norm이란 SVD의 $\Sigma$ matrix의 tr($\Sigma$)임

즉, 임

위 nuclear norm을 특정 MPQ 설정에서 model 전체에 얼마나 큰 perturbation을 일으키는지 측정하는 종합 sensitivity indicator $\Omega$에 적용

$tr()$: nuclear norm $||\hat{X}-X||^2_2$ : Q error의 곱

⇒

4 Exprimental results

이태까지 Vit의 PTQ 연구가 없어서, 기존 CNN을 위한 PTQ 기법들과 비교함

4.1 Implementation details

생략

4.2 Results and Analysis

4.3 Ablation study

5 Conclusion

생략

PACT에 대한 내용

ABSTRACT

CNN의 FC layer에서 Activation이 weight보다 크기가 큼 But, weight Q 논문만 다수 존재

PACT = learnable activation clipping param $\alpha$를 사용

1 INRODUCTION

PACT = learnable activation clipping param $\alpha$를 사용

+ 더 빠른 수렴을 위해, loss function에 $\alpha$에 대한 regularization 적용

2 RELATED WORK

생략

3 CHALLENGES IN ACTIVATION QUANTIZAITON

Weight Q는 loss function을 weight 변수에 대해 이산화하는 것과 같음 ⇒ Training 中 Q error를 보상 가능

But, Activation Q는 learnable param X ⇒ Training 中 Q error를 보상 불가능

Activation Q는 특히 ReLU를 쓸 때(그림 1의 relu with quantization) 더 어려워짐 ∵ ReLU는 upperbound가 없어, 더 많은 bit-width요구

그림 1을 보면, relu with quantization에서 error가 relu에 비해 엄청 큼

⇒ Globally uppder-bound 설정하는 (그림 1의 clippping with quantization) clipping activation function

But, Model가 차이로 인해, globally optimal clipping value를 결정 hard

⇒ Globally optimal clipping value를 결정하기 위해, learnable $\alpha$ 도입

4 PACT: PARAMETERIZED CLIPPING ACTIVATION FUNCTION

$$ y = PACT(x) = 0.5(|x| - |x - \alpha| + \alpha) = \begin{cases} 0, & x \in (-\infty, 0) \ x, & x \in [0, \alpha] \ \alpha, & x \in [\alpha, +\infty) \end{cases} \quad (1) $$

⇒ activation의 범위를 $[0, \alpha]$로 제한

이후, clipping된 activation $y$는 dot-product를 위해, $k$ bitfh linearly quantized됨 $$ y_q = round(y \cdot \frac{2^k - 1}{\alpha}) \cdot \frac{\alpha}{2^k - 1} \quad (2) $$

⇒ $\frac{\alpha}{2^k-1}$을 scale factor로도 씀 ⇒ 단순히 clipping range만 학습하는 것이 아니라, 양자화 구간(interval)과 라운딩 함수(mapping function)의 형태를 함께 학습

4.1 UNDERSTANDING HOW PARAMETERIZED CLIPPING WORKS

$\alpha$의 범위를 sweep해가며, Q를 적용했을 때의 training loss를 기록

Pretrained된 SVHN에 대해, $\alpha$의 범위에 따른, cross-entropy및 training loss를 측정

한 layer씩, ReLU대신 PACT로 했을 때의 cross-entropy가 그림 3(b)에 있음

• 그림 3(a)는

• FP forward-pass*를 통해 계산된 cross-entropy인데, Q가 적용되지 않았을 때, 즉, 식 (1)만 됬을 때, $\alpha$가 증가함에 따라(ReLU에 가까워짐에 따라) 많은 layer에서 cross-entropy가 작은 값으로 수렴 ⇒ ReLU가 좋은 activation function임을 보임 But, Q가 없는 현재 상황에서도, 특정 layer에 대해서는 잘 훈련된 $\alpha$를 쓰는게 cross-entropy를 줄임을 알 수 있음 (act0 layer를 $\alpha = 1.25$정도로 했을 때, 가장 낮은 cross-entropy, act6 layer를 $\alpha = 6.0$정도로 했을 때, 가장 낮은 cross-entropy를 갖음) ⇒ ReLU가 항상 optimal은 아님

• 그림 3(b)는

• Q forward-pass*를 통해 계산된 cross-entropy인데, $\alpha$가 증가함에 따라(ReLU에 가까워짐에 따라) 대부분의 layer에서 cross-entropy가 커짐 ⇒ ReLU가 좋은 activation function이 아님을 보임 + 최적의 $\alpha$가 layer마다 다른 범위를 가진다는 것을 관찰, 이는 training을 통해 양자화 스케일을 '학습'할 필요성을 뒷받침

4.2 EXPLORATION OF HYPER-PARAMETERS

생략

5 EXPERIMENTS

5.1 ACTIVATION QUANTIZATION PERFOMANCE

• Bit-width가 높을수록 train/test error가 FP(baseline)에 더 가까워짐
• 3bit보다 높은 bit-width를 사용하여 훈련하면 FP(baseline)과 거의 동일하게 수렴하는 것을 볼 수 있음
• 그림 4(f)를 보면, ResNet18과 ResNet50 모두에서 PACT는 다른 Q 기법들과 비교하여 일관되게 더 낮은 acc degradation을 달성했으며, 이는 이전의 Q 접근 방식들에 비해 PACT의 강건성(robustness)을 입증

  5.2 PACT PERFORMANCE FOR QUANTIZED CNNs

  PACT가 activation Q에 초점을 맞추고 있음에도 불구하고, weight Q와 함께 사용할 수 있음을 입증하기 위해, activation PACT Q, weight DoReFa Q 해봄

표 1은 다양한 model에서, activation, weight 모두 DoReFa Q 했을 때와, activation PACT Q, weight DoReFa Q 했을 때를 비교하여 보여줌

• 가중치와 활성화 모두 4비트 정밀도일 때, PACT는 테스트된 모든 네트워크에서 전체 정밀도(full-precision) 정확도를 일관되게 달성

6 SYSTEM-LEVEL PERFORMANCE GAIN

생략

7 CONCLUSION

생략

MPQ via Learned Layer-wise Importance에 대한 내용

Abstract.

Scale factor를 특정 bit-width에서 최종 accuracy에 대한 해당 layer의 기여도를 반영하는 importance indicator로 사용.

Scale factor는 QAT 중에 수치적 변환을 자연스럽게 인지해서, 해당 layer의 quatization sensitivity indicator로서 효과적.

Joint training scheme을 통해 다양한 bit-width에 대한 지표를 한 번에 얻을 수 있음.

1 Introduction

Uniform-precision Q는 suboptimal임 ⇒ Mixed-precision Q

MPQ에서 bit-width를 더 finer-grained ⇒ search space가 exponentially ↑

layer 수: $L$, 한 layer에서 activation에 대해 선택 가능한 bit-widht: $n$, 한 layer에서 weight에 대해 선택 가능한 bit-widht: $n$, ⇒ search scape $= L ⨉ (n+n)$

search-based MPQ

HAQ[25], AutoQ[21]: bit-width 결정 문제를 Markov Decision Process로 modeling해서 Deep Reinforced Learing 활용. ⇒ 시간 ↑

DNAS[26], SPOS[14]: differentiable search process를 달성하기 위해, NAS algorithm을 적용. ⇒ search space를 수동으로 제한

criterion-base MPQ

HAWQ[12], HAWQ-v2[11]: layer의 sensitivity를 측정하기 위해, 2차 정보(hessian eigenvalue, trace)를 사용. MPQCO[6]: Hessian matrix를 효율적으로 계산 후, MCKP를 공식화. ⇒ Biased approximation: Q되지 않은 network의 hessian 정보를 사용 ⇒ Q의 존재 인지 X ⇒ Limited search space: MPQCO의 근원적인 문제로 인해...

⇒ End-to-end importance criterion 사용! ⇒ Q에 있는 learnable scale factor를 사용!

Scale factor는 end-to-end로 학습되므로, Biased approximation X + limited search space X

+ joint training scheme을 통해, cirterion을 얻는 시간 ↓

2 Related Work

2.1 Neural Network Quantization

Fixed-Precision Quantization

• Lq-nets[28]: learnable quantizer 도입
• Pact [7]: activation을 위한 learnable upper bound 도입
• D.S. Learned step size quantization[13], Learning to quantize deep networks by optimizing quantization intervals with task loss.[18]: learnable scale factor 도입

Mixed-Precision Quantization

위에서 충분히 언급함

2.2 Indicator-Based Model Compression

위에서 충분히 언급함

3 Method

3.1 Quantization Preliminary

Scale factor ↑ ⇒ 더 많은 서로 다른 연속적인 값들이 동일한 Q된 값으로 mapping 됨

즉, Scale factor는 mapping 범위임

아래 그림 보면 Layer 1이 더 넓은 연속적인 수를 하나로 뭉뚱그림

3.2 From Accuracy to Layer-wise Importance

$$ S^{*} = \argmax_{S \sim \Gamma(\mathcal{A})} ACC_{val}(f(x; S, \mathcal{W}_S), y) \tag{3} $$

$$ \text{s.t.} \quad \mathcal{W}S = \argmin{\mathcal{W}} \mathcal{L}_{train}(f(x; S, \mathcal{W}), y), \tag{3a} $$

$$ \quad BitOps(S) \leq C \tag{3b} $$

위 3a, 3b 제약 조건 下 $S^*$를 찾아함.

식의 의미: $S$: 전체 network의 구성 bit-width임. (e.g., $L_1 = (3bit, 2bit), L_1 = (5bit, 6bit) ...$)

즉, 특정 Q 정책 $S$를 써서, train한 loss의 최소값을 갖는 parameter $W$를 $W_S$로 하고 (3a), 특정 $BitOps(S) ≤ C$를 지키면서,

valid ACC가 최대화 되는 S를 제한된 search space $\Gamma(\mathcal{A})$에서 찾겠다의 의미!!!

근데 search space $\mathcal(A)$가 너무 크니까 시간 오래걸리니까

⇒ joint training scheme

3.3 Learned Layer-wise Importance Indicators

같은 bit-width로 잘 훈련된 $s$와 weight를 가진 2개의 예시 layer를 Q해 봤을 때,

그림 1에서 볼 수 있듯이, layer 1의 분포가 layer 0보다 훨씬 넓다.

layer 1에서는 더 많은 서로 다른 연속적인 값들이 동일한 Q된 값으로 mapping ⇒ $s$ ↑ ⇒ 원래 연속적인 값들의 내재된 차이를 더 많이 소멸 시켜 expressiveness ↓ ⇒ bit-width 늘려야 함

$s$ ↑ ⇒ 중요함 ⇒ bit-width ↑

Feasibility Verification

MobileNet에서 Q sensitivity: depth-wise conv > point-wise conv 라고 알려짐.

MobileNetv1에 있는 5개의 DW-PW 쌍 각각에 대해, 개별적으로 Q해가며, scale factor와 accuracy를 측정.

총 10개의 실험을 함.

DW-PW 0을 2, 4 bit Q, DW-PW 1을 2, 4 bit Q, DW-PW 2를 2, 4 bit Q, DW-PW 3을 2, 4 bit Q, DW-PW 4를 2, 4 bit Q.

그림 2를 보면,

• 4bit에서 2bit로 감소할 때, DW-convs의 accuracy degradation이 더 큼 ⇒ DW-convs가 매우 민감하다는 사전 지식과 일치
• 동일 bit-width에서 DW-convs의 scale factor가 더 큼 ⇒ scale factor가 해당 layer의 Q sensitivity를 반영함

3.4 One-time Training for Importance Derivation

중요도 지표를 뽑는 cost ↑

⇒ 한 번의 훈련으로 모든 layer의 n개 bit-width에 해당하는 중요도 지표를 얻기 위한 joint training scheme

각 훈련 step $t$에서, n개의 bit-width 옵션들에 대해, forward, backward를 n번 수행, + one-shot NAS에서 영감을 받아, 각 layer에 대해 무작위 bit-width 할당 과정을 한 번 도입하여, 다른 layer의 서로 다른 bit-width들이 서로 소통할 수 있게 함.

⇒ 한 step에서, gradient 계산 $n+1$번 수행.

그림 4는 이 방법으로 단일 훈련 세션(물론 gradient 계산 $n+1$번)에서 얻은 모든 layer의 importance indicator를 나타냄.

3.5 Mixed-Precision Quantization Search Through Layer-wise Importance

그림 2에서 보았듯이, DW-convs가 높은 중요도 점수를 갖음 ⇒ bit-width ↑

제약사항을 지키면서, 각 bit-width 별 중요도 점수의 총합이 가장 작은 조합을 찾으면 됨!

4 Experiments

Dataset: ImageNet Model: ResNet18, ResNet50, MobileNetv1

• UniQ Baseline들: PACT[7] (PACT), PROFIT[23], LQ-Net[28]
  • MixQ Baseline들: HAQ [25], AutoQ [21], SPOS [14], DNAS [26], BP-NAS [27], MPDNN [24], HAWQ [12], HAWQv2 [11], DiffQ [9], MPQCO [6]

4.1 Mixed-Precision Quantization Performance Effectiveness

4.2 Mixed-Precision Quantization Policy Search Efficiency

4.3 Ablation Study

생략

5 Conclusion

생략

Deep Taylor Decomposition에 대한 내용

Abstract

DNN은 성능은 good, but, lack of transparency == limiting the interpretability

==> DNN은 multilayer nonlinear structure 때문에, blackbox 처럼 작동

==> Deep Taylor Decomposition을 통해 분류에 쓰인 입력 요소들의 기여도(ex) 사진의 어느 부분이 사진을 cat으로 분류하는 데 어느 정도의 기여를 함)를 분해하여, 네트워크의 작동 방식을 설명

1. INTRODUCTION

Blackbox처럼 작동하는 DNN은 왜 특정 결정을 내렸는지 설명 못하는 transparent 문제 ==> interpretable classifier로 model이 input을 기반으로 nonlinear classification desicion을 어떻게 내리는지 설명

DTD는 interpretable classifier를 만드는 데 사용되는 핵심 방법론

궁극적으로 입력 픽셀 공간에 분류 결정의 기여도를 보여주는 relevance map (heat map)(관련성 지도)를 생성

2. PIXEL-WISE DECOMPOSITION OF A FUNCTION

Heat map은 함수 값(즉, 신경망 출력)을 입력 변수에 redistribution 함으로써 생성

Redistribution 되는 양은 해당 입력 변수가 함수 값에 기여하는 정도!!!

이때, 함수(신경망)는 $f : \R^d \rarr \R+$ 이다. 즉, 이미지 $x \in \R^{H \times W}$이면, $x \in \R^{d = H \times W}$이고, 이 이미지의 prediction 결과는 고양이일 확률이기에, 0이상의 실수인 것이다.

이미지 $x$의 각 픽셀 $p$에 관련성 점수 $R_p(x)$를 mapping 해야한다. 이는 이미지 $x$에 대해 픽셀 $p$가 분류 결정 $f(x)$를 설명하는 데 어느 정도 기여하는지를 나타낸다. (그러니까 이 픽셀이 ai의 결정에 어느정도 영향을 미쳤는 지를 나타냄)

우리는 Heatmapping이 아래에서 정의하는 특정 속성을 충족하길 바란다.

Definition 1. conservative(보존성)

$$ \forall x:\ f(x) = \sum\limits_{p} R_p(x). $$

히트맵에 있는 모든 픽셀의 관련성 점수를 합친 값이 신경망 모델이 최종적으로 출력한 값, 즉 모델이 감지한 총 관련성($f(x)$)과 정확히 같아야 한다.

당연하다. $R_p(x)$는 $f(x)$를 분배한 것이기에,,,

Definition 2. positive(양수성)

$$ \forall x, p:\ R_p(x) \geq 0 $$

히트맵의 모든 픽셀 관련성 점수가 0보다 크거나 같아야 한다. 음수 값이 없어야 한다는 뜻이다.

특정 픽셀이 모델의 결정에 '반대되는 증거'를 제공해서는 안 된다는 것을 의미. 예를 들어, 고양이 사진에서 어떤 픽셀은 '고양이일 가능성을 높이고', 다른 픽셀은 '고양이가 아닐 가능성을 높인다'는 식의 모순된 기여는 허용하지 않는다. 모든 픽셀은 감지된 객체의 존재를 지지하거나(양수 값), 아무런 영향도 미치지 않아야(0 값) 한다.

당연하다. $f(X) \geq 0$ 이고, 이를 분배한 것이 $R_p(x)$이고, $f(x)$는 고양이일 확률이지 고양이가 아닐 확률이 아니기 때문이다.

Definition 3. consistent(일관성)

$$ (f(x) = 0) \Rightarrow (R(x) = 0) $$

히트맵이 보존성(정의 1)과 양수성(정의 2)을 모두 만족할 때 '일관적'이라고 합니다.

특히, 모델이 이미지에서 아무것도 감지하지 못했을 경우($f(x)=0$), 히트맵의 모든 픽셀 관련성 점수도 0이어야 합니다. 이는 객체가 없을 때 히트맵이 완전히 비어있어야 하며, 긍정적 기여와 부정적 기여가 상쇄되어 0이 되는 경우가 아니어야 함을 강조. 이 '일관성'은 논문에서 제안하는 히트맵 기법의 정확성을 평가하는 기준으로 사용.

즉, 이제 Definition 3. consistent를 만족하며, 의미 있는 heatmapping 기법을 정해야한다.

A. Taylor Decomposition

논문에서는 적절한 heatmapping 기법으로 $f(\tilde{x})=0$인 $\tilde{x}$에서의 $f$의 taylor expansion을 선택했다.

이는 아래와 같다.

$$ \left( \left. \frac{\partial f}{\partial x} \right|_{x = \tilde{x}} \right)^\top$$

에서 전치가 있는 이유는 $x$가 벡터이기 때문이다.

또한 $R(x)={R_p(x)}$ 이므로, $$\sum_p \left. \frac{\partial f}{\partial x_p} \right|_{x = \tilde{x}} \cdot (x_p - \tilde{x}_p)$$는 $R(x)$가 된다.

이때, $x$는 벡터이므로, $f(x)=R(x)+\epsilon = \left. \frac{\partial f}{\partial x} \right|_{x = \tilde{x}} \odot (x - \tilde{x}) + \epsilon$ 이다.

여기서 $\epsilon$은 2차 이상의 항들인데, 이 항들은 동시에 여러 픽셀의 정보를 포함해서 재분배하기 어렵기에 무시한다.

(자세한 이유는 $x$가 벡터이므로, 2차 미분일 경우 이런 식으로 되서 그럼)

$R(x)=\left. \frac{\partial f}{\partial x} \right|_{x = \tilde{x}} \odot (x - \tilde{x})$ 의 각 항을 설명하면, 아래와 같다.

• $\tilde{x}$ : $f(\tilde{x})=0$ 인 이미지 $\tilde{x}$ 이므로, 물체를 판별하는 데 쓸모 없는 부분임 ex) 고양이가 아닌 부분

• $(x-\tilde{x})$ : 물체를 판별하는데 쓸모 있는 부분

• $\left. \frac{\partial f}{\partial x} \right|_{x = \tilde{x}} \odot (x - \tilde{x})$ : $\tilde{x}$의 각 픽셀이 $f(x)$에 영향을 주는 정도

이제 남은 일은 최적의 $\tilde{x}$를 찾는 것이다

좋은 root point는 일부 픽셀에서 정보를 선택적으로 제거하는 동시에 주변 환경을 변경하지 않아야 한다. 즉 고양이 빼고는 정확한 이미지여야한다.

그렇기에, 다음의 $\xi$를 찾는 것이 목표이다. $$ \min_{\xi} \ |\xi - x|^2 \quad \text{subject to} \quad f(\xi) = 0 \quad \text{and} \quad \xi \in \mathcal{X} $$

이때, $\mathcal{X}$는 input domain 이다.

하지만, 위 방법은 시간이 오래 걸리고, non-convexity 문제 발생

B. Sensitivity Analysis

위의 문제(적절한 root point를 찾는 것이 어렵다.)를 해결하는 가장 간단한 방법은 픽셀에 relevance를 부여할 때, squared derivatives of the classifier에 비례하게 부여하는 것이다.

$$ R(x) \propto \left( \frac{\partial f}{\partial x} \right)^2 $$

이때, 제곱은 아다마르 곱셈이다.

그럼 이때의 root point는 무엇인가? ==> $\xi = x - \delta \cdot \frac{\partial f}{\partial x}$

여기서 점 $ξ$는 실제 점 $𝑥$로부터 함수 $f$가 가장 급격하게 감소하는 방향으로 무한히 작은 거리만큼 떨어져 있다.

따라서 $\xi$에서의 1차 taylor expansion은 다음과 같다

$$ f(x) = f(\xi)+ \left( \left. \frac{\partial f}{\partial x} \right|{x=\xi} \right)^\top \cdot \left( x - \left( x - \delta \frac{\partial f}{\partial x} \right) \right) + 0 \ = f(\xi) + \delta \left( \frac{\partial f}{\partial x} \right)^\top \frac{\partial f}{\partial x} + 0 \ = f(\xi) + \sum_p \underbrace{ \delta \left( \frac{\partial f}{\partial x_p} \right)^2 }{R_p} + 0, $$

여기서 2차 이상의 항들은 1차 taylor expansion이기에 0이 된다.

하지만 여기서 문제가 있는데, $\xi$를 $x$에서 아주 살짝 떨어진 이미지로 했기에, $f(\xi)$가 $f(x)$와 비슷할 수 밖에 없어, $\sum_p \underbrace{ \delta \left( \frac{\partial f}{\partial x_p} \right)^2 }_{R_p}$는 작을 수 밖에 없다.

즉, 모든 관련성이 재분배 되지 않는 0차 항($f(\xi)$)에 흡수되기 때문에 보존적이지 않게 된다.(Definition 1 위반)

C. Deep Taylor Decomposition

Deep taylor decomposition은 divide-and-conquer에서 영감을 받았다. DNN이 학습한 함수가 인접한 layer 간의 값을 연결하는 더 간단한 하위 함수들의 집합으로 구조적으로 분해될 수 있다는 특성을 이용. 즉, input -> output이 아니라, 그 사이에 여러 layer를 하위의 함수로 취급하여, divide-and-conquer하는 것임

==> 전체 신경망 함수 $f$를 직접적으로 고려하는 대신, 특정 layer의 뉴런 집합 ${x_i}$을 다음 layer의 뉴런 $x_j$에 할당된 relevance $R_j$로 mapping하는 국소 문제를 푸는 것임

// forward 방향: i -> j임 그럼 backward로 relevance backpropagation하면 j -> i 인데, 왜 뉴런 집합 ${x_i}$을 $x_j$에 할당된 $R_j$로 mapping 하지???

==> 그니까 taylor decomposition은 함수를 분해하는 거임 여기서 그 함수는 ${x_i}$을 통해 얻은 $x_j$의 relevance $R_j$를 분해하는 것이기에 그렇게 표현한 것

즉, 이 두 대상이 특정 함수 $R_j({x_i})$에 의해 함수적으로 연결되어 있다고 가정하면, 이 국소적(local) 함수에 Taylor decomposition를 적용하여 상위 층의 관련성 $R_j$를 하위 층의 관련성 점수 집합 ${R_i}$로 재분배하는 것!

이렇게 하면 함수가 더 작아져서 taylor decomposition이 더 쉬워지며, 특히, root point를 찾는 것이 용이

이제 중간 hidden layer에서의 taylor decomposition을 해보자.

$$ f(x) = f(\tilde{x}) + \left( \left. \frac{\partial f}{\partial x} \right|{x=\tilde{x}} \right)^\top \cdot (x - \tilde{x}) + \varepsilon\ = 0 + \sum_p \underbrace{ \left. \frac{\partial f}{\partial x_p} \right|{x=\tilde{x}} \cdot (x_p - \tilde{x}p) }{R_p(x)} + \varepsilon \tag{3} $$

앞에 식은 뉴런 하나($x_j$)에서 뉴런 여러개(${x_i}$)로 relevance를 분배하는 것이다.

반면 아래 식은 여러 뉴런(${x_j}$)에서 여러 뉴런(${x_i}$)로 relevance를 분배하는 것이다.

$$ \sum_j R_j = \left( \left. \frac{\partial \left( \sum_j R_j \right)}{\partial {x_i}} \right|{{\tilde{x}_i}} \right)^\top \cdot ({x_i} - {\tilde{x}_i}) + \varepsilon \ = \sum_i \underbrace{\sum_j \left. \frac{\partial R_j}{\partial x_i} \right|{{\tilde{x}i}} \cdot (x_i - \tilde{x}_i) }{R_i} + \varepsilon, \tag{4} $$

어렵게 보이지만 간단하다. 그냥 위에 식 (3)을 여러번(hidden 뉴런 개수번 ($num({x_j})$) 적용하여, 누산 한 것이다.

3. APPLICATION TO ONE-LAYER NETWORKS

$$ x_j=\max\left(0, \sum_i x_i w_{ij} + b_j \right) \tag{5} $$

$$ x_k = \sum_j x_j \tag{6} $$

인 간단한 nn을 가정하자.

이때, 모든 j에 대해 $b_j\leq0$이도록 하자.

==> 함수 $g$의 root point ${\tilde{x_i}}$의 존재를 보장.

==> root point가 반드시 필요한 taylor decomposition 적용 可

예측된 출력값을 역방향으로 전파되어야 하는 총 관련성과 동일하다고 설정하면, $R_k = x_k$ 가 된다.

이때, $x_k=\sum_j x_j$ 이기에, $R_k = x_k=\sum_j x_j$ 이다.

이 $R_k$를 각 뉴런 ${x_j}$로 분배하자. ==> taylor decomposition

Step 1. $x_k$에서 ${x_j}$로의 taylor decomposition

$$ R_j = \left. \frac{\partial R_k}{\partial x_j} \right|_{{\tilde{x}_j}} \cdot (x_j - \tilde{x}_j) \tag{8} $$

당연히, root point $\tilde{x_j}$는 다음을 만족해야 한다. $\sum_j \tilde{x_j} = 0$ 왜냐하면, $R_k$를 $0$이 되게 하는 ${x_j}$는 $\sum_j x_j =0$ 이기 때문이다.

또한 root point가 역할을 제대로 하려면 반드시 허용 가능한(admissible) 값이어야 하는데,

그 값은 ${x_j}\ge0$ 이다.($ReLU$ 때문에)

위 두 조건($\sum_j \tilde{x_j} = 0$) , (${\tilde{x_j}}\ge0$)에 의해, ${\tilde{x_j}}=0$ 이다.

또한 $R_k=\sum_j x_j$ 를 $x_j$로 미분하면, 1이기에, $$ R_j = \left. \frac{\partial R_k}{\partial x_j} \right|_{{\tilde{x}_j}} \cdot (x_j - \tilde{x}_j) \tag{8} =x_j $$ 이다.

Step 2. ${x_j}$에서 ${x_i}$로의 taylor decomposition

$R_j=max(0,\sum_i x_iw_{ij}+b_j)$ 이다.

이를 통해 다시 ${R_i}$로의 재분배를 하자. ==> deep taylor decomposition

식 (4)를 적용하면, 다음과 같다.

$$ R_i = \sum_j \left. \frac{\partial R_j}{\partial x_i} \right|_{{\tilde{x}_i}^{(j)}} \cdot \left( x_i - \tilde{x}_i^{(j)} \right). \tag{11} $$

각 관련성 $R_j$를 분해할 때마다 서로 다른 root point를 선택한다는 점을 강조하기 위해, root point(${\tilde{x_i}}$)에 위첨자 ($j$)를 추가.

cf) 괜찮은 참고자료 (https://angeloyeo.github.io/2019/08/17/Layerwise_Relevance_Propagation.html)

A. Unconstrained Input Space and the $w^2$-Rule

Input을 이미지가 아닌 일반적인 상황일 때를 먼저 고려해보자.

그럼 $Input = \mathcal{X} = \R^d$ 이다.

최적의 root point ($\tilde{x}$)를 찾자! $$ \min_{\tilde{x}} \ |\tilde{x} - x|^2 \quad \text{subject to} \quad f(\tilde{x}) = 0 \quad \text{and} \quad \tilde{x} \in \mathcal{X} $$

의 $\tilde{x}$를 찾으면 된다.

이는

• plane equation $\sum_i \tilde{x}i^{(j)} w{ij} + b_j = 0$ (식 (10)에서 정의된 $R_j$ 평면이다)에 수직이며, (이 plane equation의 법선 벡터는 $w^{(j)}=(w_{1j},w_{2j},...,w_{ij})$이다.)
• input $x$를 지나는 직선

과 plane equation의 교점이다!

이 교점은 다음과 같다. $$ {\tilde{x}i}^{(j)} = {x_i - \frac{w{ij}}{\sum_i w_{ij}^2} \left( \sum_i x_i w_{ij} + b_j \right)} $$

이 교점을 식(11)에 대입하면, 뉴런 $i$로 재분배되는 relevance를 구할 수 있고, 이는 다음과 같다. $$ R_i = \sum_j \frac{w_{ij}^2}{\sum_{i'} w_{{i'}j}^2} R_j \tag{12} $$

즉, $w^2-rule$은 input space에 제약이 없는 경우 relevance를 weight 제곱에 따라 안정적으로 분배하는 일관된 방법이다.

하지만, 만약 input space가 이미지 처럼 $R_+^d$ 이면, $w^2-rule$을 적용하지 못할 수 있다. 왜냐하면, 아래 그림과 같이, 구한 ${x_i}^{(j)}$가 $R_+^d$가 아닐 수 있기 때문이다.

그래서, B의 $z-rule$을 사용하는 것이다.

B. Constrained Input Space and the z-Rules

C. Experiment

TBW

4. APPLICATION TO DEEP NETWORKS

A. Min-Max relevance Model

B. Training-Free Relevance Model

C. Experiment on MNIST

5. CONCLUSION

생략

Mix-QViT에 대한 내용

Abstract

Quantization은 model size를 줄이는 좋은 기법이지만,

전체 network에서 uniform bit-width quantization하는 것은 suboptimal performance임

⇒ MPQ(Mixed-Precision Quantization)

하지만,

1) Exhaustive search algorithms에 의존 2) 각 layer의 특정 역할과 quantization sensitivity를 고려하지 않고, arbitrarily different bit-width로 quantization 함

⇒ Mix-QViT (explainability-driven MPQ framework)

아래 2가지 criteria로 각 layer별로 bit-width를 부여

1) LRP(Layer-wise Relevance Propagation)으로 평가된 layer importance (각 layer가 final classification에 얼마나 영향을 미치는 지 나타냄) 2) Quantization sensitivity (다른 layer들을 baseline precision으로 유지하면서, 각 layer를 다양한 precision으로 quantization 했을 때, 성능에 미치는 영향을 평가)

+ 추가적으로, PTQ에서 post-LayerNorm activation에서 발생하는 extreme outliers를 줄이기 위한, clipped channel-wise quantization 제안

Q1. Exhaustive search algorithms 란?

A1. 예시를 들어 설명하겠다.

상황 설정 및 필요한 가정:

• 모델: 비전 트랜스포머(ViT) 모델이라고 가정합니다.
• 레이어 수: 간단하게 이 ViT 모델이 총 3개의 레이어(Layer 0, Layer 1, Layer 2)로 구성되어 있다고 가정합시다.
• *가능한 비트 너비 집합 B: * 모델의 가중치(weights)와 활성화(activations)에 적용할 수 있는 비트 너비 선택지가 {2비트, 4비트, 8비트}라고 가정합니다. 즉, $B = {2, 4, 8}$ 입니다. 여기서 비트 수가 낮을수록 모델 크기는 작아지고 빨라지지만 정확도 손실 가능성이 커집니다.
• *각 레이어의 비트 너비 선택: * 각 레이어 $l$마다 가중치($b_l^w$)와 활성화($b_l^a$)에 대해 B에서 비트 너비를 선택할 수 있습니다. 예를 들어, Layer 0의 가중치를 4비트로, 활성화를 8비트로 설정할 수 있습니다. (즉, $(b_0^w, b_0^a) = (4, 8)$)
• 전체 아키텍처 비트 너비 구성 S: 전체 모델의 비트 너비 구성은 각 레이어의 가중치-활성화 비트 너비 쌍의 집합입니다. 예를 들어, $S = {(b_0^w, b_0^a), (b_1^w, b_1^a), (b_2^w, b_2^a)}$ 와 같이 표현됩니다.
• 탐색 공간 A: 가능한 모든 $S$의 조합을 포함하는 공간입니다.
• 목표: 자원 제약 조건 $C$를 만족하면서 검증 성능(validation performance)을 최대화하는 최적의 비트 너비 구성 $S^*$를 찾는 것입니다.

Exhaustive Search Algorithm (전수 탐색 알고리즘)의 예시:

전수 탐색 알고리즘은 가능한 모든 $S$ 조합을 만들고, 각 조합에 대해 모델을 양자화하고, 보정(calibration)하고, 성능을 평가하는 과정을 거칩니다.

1. 가능한 비트 너비 조합 계산:
   • 각 레이어는 가중치와 활성화에 대해 각각 $|B|=3$가지 비트 너비 선택(2, 4, 8비트)을 가집니다.
   • 따라서 한 레이어에서 가능한 $(b_l^w, b_l^a)$ 조합의 수는 $3 \times 3 = 9$가지 입니다. (예: (2,2), (2,4), (2,8), (4,2), (4,4), (4,8), (8,2), (8,4), (8,8))
   • 총 3개의 레이어가 있으므로, 전체 아키텍처에 대한 가능한 비트 너비 구성 $S$의 총 수는 $9 \times 9 \times 9 = 9^3 = 729$가지 입니다.

2. 전수 탐색 과정: 전수 탐색 알고리즘은 이 729가지의 모든 $S$ 조합에 대해 다음 과정을 수행합니다:
   • 조합 1: $S_1 = {(2,2), (2,2), (2,2)}$ (모든 레이어의 가중치와 활성화를 2비트로)
     • 모델을 이 구성으로 양자화합니다.
     • 양자화된 모델을 보정합니다.
     • 검증 데이터셋으로 성능(예: 정확도)을 평가합니다.
     • 자원 제약 $C$ (예: 모델 크기 5MB 이하, 추론 시간 10ms 이하)를 만족하는지 확인합니다.

   • 조합 2: $S_2 = {(2,2), (2,2), (2,4)}$
     • 위와 동일한 과정을 반복합니다.

   • ... (중략) ...

   • 조합 729: $S_{729} = {(8,8), (8,8), (8,8)}$ (모든 레이어의 가중치와 활성화를 8비트로, 즉 Full-Precision에 가까움)
     • 위와 동일한 과정을 반복합니다.

3. 최적 $S^*$ 선택:
   • 729가지의 모든 $S$ 조합에 대한 평가가 끝난 후, 자원 제약 $C$를 만족하면서 가장 높은 검증 성능을 보인 $S$를 $S^*$로 선택합니다.

전수 탐색의 문제점:

이 예시에서는 레이어가 3개에 불과했지만, 실제 ViT 모델은 수십, 수백 개의 레이어를 가집니다. 예를 들어, 레이어가 12개라면 가능한 조합은 $9^{12}$가지가 되어 상상할 수 없을 정도로 많은 계산이 필요하게 됩니다. 이러한 이유로 논문에서는 전수 탐색 방법이 "계산적으로 매우 집약적(computationally intensive)"이며 "높은 비용(high cost)"을 요구한다고 지적하는 것입니다.

Mix-QViT는 바로 이 지점에서 각 레이어의 중요도와 민감도를 체계적으로 분석하여, 비효율적인 전수 탐색 없이도 최적에 가까운 비트 너비 구성을 찾으려는 목표를 가지고 있습니다.

** Abstract 요약** Mix-QViT가 제안한 2가지 기법

1) LRP를 이용한 Layer Importance, Quantization Sensitivity. 2가지 criteria를 사용한 Explainability-driven MPQ framework

2) Post-LayerNorm activation에서 발생하는 extreme outliers의 영향을 줄이기 위한, severe inter-channel variations 제거 기법 Clipped channel-wise quantization

Q2. Layer importance와 Quantization sensitivity가 뭔 차이지?

A2.

bit-width 줄일 순서 : C -> A -> B -> D

1. INTRODUCTION

문제점 1.

PTQ에서 LayerNorm, Softmax, GELU가 quantization에 sensitive해서, performance degradation에 큰 영향을 준다.

=>

• APQ-ViT Matthew-effect-preserving quantization을 사용하는 block-wise calibration
• FQ-ViT LayerNorm의 channel 간 변화를 처리하기 위해 Power-of-Two Factor를 도입하고, Softmax layer quantization을 위해, Log-Int-Softmax 사용
• RepQ-ViT 및 AdaLog calibration 단계에서는 LayerNorm에 channel-wise quantization을 사용하고 , Softmax layer의 경우 RepQ-ViT는 밑이 $\sqrt{2}$인 log quantization을, AdaLog는 적응형 밑 log quantization을 사용하여 표현 능력을 향상.
  inference 단계에서는 두 프레임워크 모두 LayerNorm에 layer-wise quantization을, Softmax layer에 log2 quantization를 채택하여 하드웨어 효율적인 형식으로 전환.

=> 여전히 low-bit(e.g. 4-bit)에서는 FP model과 큰 차이 존재

문제점 2.

Uniform bit precision은 suboptimal

=> MPQ

• MHSA, MLP module의 sensitivity를 측정하기 위해, nuclear norm metirc 사용.

• PMQ는 각 layer를 제거했을 때, 유도되는 오류를 평가

• Hessian based 방법을 통해 중요도 점수를 측정

=> high computational overhead + fail to explain

해결법

2가지 complementary score를 활용한 mixed-precision quantization strategy

1. Explainability-based layer importance score
2. Quantization sensitivity score

1. Explainability-based layer importance score

이는 LRP(Layer-wise Relevance Propagation)를 사용하여 도출되며, 각 layer가 최종 예측에 얼마나 기여하는지에 대한 interpretable metirc 제공

2. Quantization sensitivity score

이는 각 layer의 성능이 baseline 대비 다른 bit 구성에서 어떻게 변화하는지 포착

=> 이 2가지를 통해, 주어진 제약 조건 下에서 각 layer의 optimal bit-width를 할당하는 Integer Quadratic Problem(IQP)을 공식화

2. METHODS

A. Vision Transformer Architecture

생략

B. Model Quantization

TBW

C. Clipped Channel for LayerNorm Activations

생략

D. Layer Importance Score using LRP

LRP는 gradient와 relevance scores를 사용하여, 각 feature가 model의 output에 미치는 영향을 측정한다.

gradient 구하기

$x^l$이 forward에서 앞에 layer input이고, $x^{l+1}$이 뒤에 layer input.

$x^1$이 network input, $x^L$이 output에 해당할 때, classifier's output $y_c$에 대한 $l$번 째 layer의 gradient는 다음과 같다.

$$ \nabla x_j^{l} = \frac{\partial y_c}{\partial x_j^{l}} = \sum_i \frac{\partial y_c}{\partial x_i^{l+1}} \cdot \frac{\partial x_i^{l+1}}{\partial x_j^{l}}, \tag{18} $$

당연히, index $j$는 $x^l$의 뉴련이고, $i$는 $x^{l+1}$의 뉴런이다.

relevance scores 구하기

$R^l$이 forward에서 앞에 layer relevance score이고, $R^{l+1}$이 뒤에 layer relevance score.

LRP의 relevance propagation은 DTD에 기반한다. $l$번 째 layer가 input feature map $X^l$과 weight $W^l$에 대해 수행하는 연산을 $M_i^l(X^l,W^l)$라고 하면, relevance 는 다음과 같다.

$$ \begin{aligned} R_j^{l} &= \mathcal{G}(X^l, W^l, R^{l+1}) \ &= \sum_i x_j^l \frac{\partial M_i^l(X^l, W^l)}{\partial x_j^l} \frac{R_i^{l+1}}{M_i^l(X^l, W^l)}. \end{aligned} \tag{19} $$

Transformer block의 activation function이 GeLU와 같이 양수, 음수 모두 만들 경우, 이런 혼합된 값들로 인해 발생할 수 있는 Relevance score의 부정확성을 해결하기 위해, 음수 요소를 제거한, subset $p ={ (i,j) \mid x_j w_{ji} \geq 0 }$ 을 사용하여, relevance score를 정의한다.

$$ \begin{aligned} R_j^l &= \mathcal{G}p(x, w, p, R^{l+1}) \ &= \sum{{i \mid (i, j) \in p}} \frac{x_j w_{ji}} {\left( \sum_{{j' \mid (j', i) \in p}} x_{j'} w_{j'i} \right)} R_i^{l+1}. \end{aligned} \tag{20} $$

Q3. Transformer Interpretability Beyond Attention Visualization의 LRP는 MHSA module 내 attention layer에 대해, [CLS] token의 relevance와 그 gradient 간의 hadamard 곱을 수행하여, relevance score map을 계산한다. 하지만 이 방식은 [CLS] token을 사용하는 architecture에만 적용 가능하다는 제한이 있다.

본 연구에서는 [CLS] token 대신 image tokens를 활용함으로써, transformer의 모든 layer에 대한 relevance mapping을 일반화하고, 이를 모든 transformer architecture로 확장한다.

이게 무슨 말이지?

A3.

🔵 1. CLS 기반 relevance 계산이란?

Transformer 구조에서는 보통 [CLS] 토큰이 전체 입력 시퀀스를 대표하는 요약 정보로 사용되어 최종 분류에 활용됩니다. 따라서 많은 LRP 기반 해석 방법들은 다음처럼 relevance를 계산합니다:

✔️ 계산 방식

• 모델의 출력 $y \in \mathbb{R}^C$ 중에서 **관심 있는 클래스 $t$**에 해당하는 값을 선택
• 그 출력에 대한 **relevance $R^{(0)}$**를 초기화:
  $R^{(0)} = \mathbf{1}_t$
• relevance를 역전파할 때, 주로 [CLS] 토큰 위치에 집중해서 relevance를 할당하고 해석

  ✔️ 한계

• [CLS] 토큰만 해석하므로 **입력 이미지의 공간적인 정보(패치 별 중요도)**를 파악하기 어려움
• 모든 정보를 [CLS]가 대표한다고 가정하는데, 이는 일부 모델에서 적절하지 않음 (예: patch-based 모델들)

🟢 2. Image token 기반 relevance 계산이란?

이 논문에서는 [CLS] 토큰 대신, 이미지 패치에 해당하는 토큰들(image tokens)을 기반으로 relevance를 계산합니다.

✔️ 계산 방식

• 각 attention 레이어의 image token들 전체에 대해 relevance $R$를 계산
• 예를 들어 $s$개의 image token에 대해 $R \in \mathbb{R}^{s}$ 형태로 relevance를 유지
• 이 relevance를 통해 어떤 패치(영역)가 분류에 기여했는지를 파악 가능

  ✔️ 장점

• relevance가 공간 분포를 가지므로, 시각적으로 어떤 패치가 중요했는지 해석 가능
• [CLS] 없이도 모든 transformer 구조에 적용 가능 (예: Swin Transformer처럼 CLS가 없는 구조)

앞서 Relevance score를 구하는 방법을 설명했다.

이제 결과적으로 Mixed-Precision Bit Allocation Strategy에서 쓰일 layer importance score, $\Omega^L_u$에 대해 알아보자.

Layer Importance Score 구하기

$\Omega^L_u$는 $L$ 번째 block 내의 layer $u$가 model 중요도에서 차지하는 비중을 의미한다.

그 식은 다음과 같다.

$$ \Omega_a^L = \frac{C_a^L}{\sum_{l=1}^{L} \sum_{u \in U_l} C_u^l}. \tag{23} $$

$C_a^L$ : 블록 $L$의 특정 레이어 $a$의 contribution score $U_l$ : 블록 $l$에 포함된 모든 양자화된 레이어 집합

그니까 $U_l$은 QKV, MatMul1, MatMul2, Projection, FC1, FC2 이다.

즉, 모든 블록의 모든 레이어의 contribution score의 총합에 대한 현재 블록의 현재 레이어의 contribution score가 Layer Importance Score 이다.

그럼 이제 contribution score에 대해 알아보자.

Contribution Score 구하기

Contribution score $C_a^L$는 ImageNet1k validation dataset에서 $T=256$개 샘플을 임의로 뽑아, 하나의 샘플에 대한 relevance score map $S_{a,t}^L$를 구하고, 그 relevance score map의 원소 전체 평균 $\mathbb{E}(S_{a,t}^L)$구하고, 그걸 전체 샘플에 대해 반복하여, 평균을 얻은 것이다.

그 식은 다음과 같다.

$$ C_a^L = \frac{1}{T} \sum_{t=1}^{T} \mathbb{E}\left(S_{a,t}^L\right). \tag{22} $$

즉, Contribution Score는 하나의 레이어가 여러 샘플에서 얼마나 일관되게 모델 출력에 영향을 미치는지를 나타내는 값으로, 각 샘플에 대해 생성된 relevance score map의 평균 기여도를 다시 평균내어 계산한 값이다.

그럼 이제 relevance score map에 대해 알아보자.

Relevance Score Map 구하기

Relevance Score Map $S^L_{a,t}$는 하나의 샘플 t에 대해, 블록 $L$의 특정 레이어 $a$에서의 gradient $∇^L_{a,t}$와 relevance $R_{a,t}^L$의 hadamard 곱을 하여, attention head 차원에 대한 평균을 내면된다.

그 식은 다음과 같다.

$$ S_{a,t}^L = \mathbb{E}h \left( \nabla{a,t}^L \odot R_{a,t}^L \right)^+. \tag{21} $$

각 요소의 차원에 대하여...

$R_{a,t}^L ∈ \R^{h × N × N}$ $\nabla_{a,t}^L ∈ \R^{h × N × N}$ $S_{a,t}^L ∈ \R^{N × N}$ $\mathbb{E}(S_{a,t}^L) ∈ \R$ $C_a^L ∈ \R$ $\Omega_a^L∈ \R$

아래 그림은 DeiT-Small의 각 레이어별 Layer Importance Score를 나타낸 것이다. 즉, 출력 분류에 기여하는 정도를 표현한 것이다.

E. Quantization Sensitivity Analysis

PTQ를 사용하여 ViT 내 개별 transformer layer들의 quantization sensitivity를 분석

모든 layer를 동일한 기준 bit(예: 4/4비트)로 설정한 뒤, 특정 layer 하나의 bit만 다양하게 조정(예: 2/2, 3/3, 5/5, 6/6비트)하며 성능 변화를 측정 ➡ layer별 quantization sensitivity 측정

그 식은 다음과 같다.

$$ \Delta \mathcal{L}{T_b} = \mathcal{L}{(U \setminus T){b'}, T_b} - \mathcal{L}{U_{b'}} , \tag{24} $$

$$ \Delta \mathcal{L}{T_b}^{+} = \Delta \mathcal{L}{T_b} + \left| \min_{u \in U,, b \in \mathcal{B}} \Delta \mathcal{L}_{ub} \right| , \tag{25} $$

$$ \Lambda_{T_b} = \frac{\Delta \mathcal{L}{T_b}^{+}}{\sum{b \in \mathcal{B}} \sum_{u \in U_b} \Delta \mathcal{L}_{ub}^{+}} . \tag{26} $$

• $L_{U_{b'}}$ : 모든 layer가 기준 비트($b'$)로 quantization되었을 때의 loss 값
• $L_{(U\setminus T)_{b'},T_b}$ : 전체는 기준 비트인데, 특정 레이어 T만 다른 비트 폭 $b$로 설정했을 때의 loss
• $\left| \min_{u \in U,, b \in \mathcal{B}} \Delta \mathcal{L}_{ub} \right|$ : 음수 값으로 인한 왜곡을 방지하기 위해, 전체 레이어 ($U$)에 대해 관찰된 최소 변화값의 절대값을 더함

아래는 ImageNet-1K validation dataset에서 무작위 256개 image로 DeiT-Small model를 추적한 결과이다.

보면,

• MatMul2 연산의 경우 : 정밀도를 5/5 또는 6/6으로 높였을 때는 성능에 거의 영향을 주지 않았지만, 이를 3/3 이하로 낮췄을 때는 성능이 크게 저하되었다.

• FC2 연산의 경우 : 정밀도를 높이면, 다른 layer 보다 성능 향상이 두드러졌다.

F. Mixed-Precision Bit Allocation Strategy

주어진 레이어 $l$에 대해, weight, activation에 모두 적용 가능한 bit-width 선택지 집합 $B$로부터 다음과 같은 조합이 가능하다.

$(b_l^w, b_l^a), \quad \text{where } b_l^w \in B,, b_l^a \in B$

전체 architecture bit-width 설정은 다음과 같다.

$S={(b_l^w,b_l^a)}^L_{l=0}$

이때, search space $A$는 가능한 모든 $S$의 조합으로 이루어지는데, MPQ의 목표는 자원 제약 조건 $C$하에서 검증 성능을 최대화하는 최적의 비트 폭 설정 $S^*∈A$를 찾는 것이다.

이 과정은 각 후보 S에 대해 quantization, calibration, evaluation이 필요하므로 계산 비용이 매우 크다.

➡ pre-computed indicator를 활용한 mixed-precision bit selection strategy

$$ \Phi = \max_{\alpha} \sum_{l=1}^{L} \sum_{u \in U} \left( \Omega_u^l \odot \left( \sum_{j=1}^{|\mathcal{B}|} \alpha_{u,j}^l b_j \right)

• \sum_{j=1}^{|\mathcal{B}|} \Lambda_{u,j} \alpha_{u,j}^l b_j \right), \tag{28a} $$

$$ \text{s.t.} \quad \alpha_{u,j}^l \in {0, 1}, \quad \sum_{j=1}^{|\mathcal{B}|} \alpha_{u,j}^l = 1, \quad \forall, l, u, \tag{28b} $$

$$ \sum_{l=1}^{L} \sum_{u \in U} \sum_{j=1}^{|\mathcal{B}|} \alpha_{u,j}^l |w_u^l |b_j \leq C_M, \tag{28c} $$

$$ \sum_{l=1}^{L} \sum_{u \in U} \sum_{j=1}^{|\mathcal{B}|} \mathrm{MAC}u^l \alpha{u,j}^l b_j^2 \leq C_{\mathrm{BitOps}}, \tag{28d} $$

$$ l \in {1, \dots, L},\quad u \in {u_1, \dots, u_{|U|} },\quad \mathcal{B} = { b_1, \dots, b_{|\mathcal{B}|} }. \tag{28e} $$

$\alpha_u^l ∈ {0,1}^{|B|}$ 이며, $\sum_{b=1}^{|\mathcal{B}|} \alpha_{u,b}^l = 1$ 는 각 layer의 bit-width 결정이다. 이게 무슨 말이냐면,

• $\alpha^l_u \in {0,1}^{|\mathcal{B}|}$는 → $\alpha^l_u$가 **길이 $|\mathcal{B}|$**인 0-1 벡터라는 의미입니다. → 즉, 각 원소 $\alpha^l_{u,b} \in {0,1}$

• 그런데 그 원소들의 합이 1이어야 하므로, → 전체 벡터 중 오직 하나의 위치만 1이고, 나머지는 전부 0인 벡터입니다.

즉, 주어진 조건을 만족하는 $\alpha^l_u$는 one-hot 벡터 (one-hot encoding)입니다.

예를들어, $\alpha^l_u$ 가 (1, 0, 0, 0) 이면 1 bit. $\alpha^l_u$ 가 (0, 0, 1, 0) 이면 4 bit.

3. EXPERIMENTS

A. Experimental Setup

B. Post Training Quantization

1) Results on ImageNet1K

2) Results on COCO

C. Quantization-Aware Training

D. Ablation Study

E. Efficiency Analysis

4. CONCLUSIONS

SpecEE에 대한 내용

Abstract

Early Exiting(EE)는 최근 하드웨어 연산량과 메모리 접근을 효과적으로 줄임으로써 LLM의 inference 속도를 높이는 유망한 기법으로 주목받고 있다.

Q1. EE란? A1. 모델이 모든 레이어를 거치지 않고, 중간에서 추론을 종료하는 방식

논문에서 제시하는 문제점 : LLM의 vocabulary가 EE predictor의 runtime search space로 작동하며, predictor workload에 큰 영향을 미친다. (e.g., Llama2에서 약 30,000 단어 크기의 vocabulary를 search 하는 데 전체 inferecne 지연 시간의 20% 차지)

구체적으로 다음 세 가지 문제 해결.

1. 높은 계산 복잡도의 predictor 기존 predictor의 경우, vocabulary size 전체가 search space이다. => 문제점
2. 계층별 예측기 배치의 비효율적 활용
3. speculative decoding에서의 예측기 매핑 복잡도의 기하급수적 증가

1. Introduction

LLM inference 속도를 높이고, deployment cost를 줄이기 위한 연구들이 다수 진행되어 왔다. 이들 중 일부 연구(fast decoding)는 결과의 일관성을 유지하지만, 다른 연구들(pruning, quantization)은 정확도 손실을 초래할 수 있어, Figure 1(a)와 같이 LLM inference 및 deployment에서 정확도와 속도 향상 간의 pareto frontier를 형성한다.

더 속도 향상을 하기 위해서 Early Exiting(EE)와 같은 기법이 있다. 이는 LLM inference에서 일련의 decoder layer를 전체 다 사용하는 것이 아니라, 동적으로 조정하여, 부분 사용하는 것이다.

이런 방법이 속도 향상을 이루는 이유는 Figure 1(b)에서 보듯이 decoder layer가 LLM inference에서 대부분의 시간을 쓰기 때문이다.

EE의 신뢰성은 Figure 1(c)에서 보듯이 token 생성 시 각 token마다 필요한 decoder layer가 다름에서 얻을 수 있다.

+ 또한 EE는 파라미터를 조정하지 않으므로, 일반성에서도 좋다.

EE 구조 물론입니다. 아래는 말씀하신 내용을 포함해 AdaInfer의 Early Exit 구조를 구체적이고 체계적으로 정리한 설명입니다. 각 단계는 연산 방식, 입력 차원, 분류기 학습 방식까지 포괄합니다.

🔍 AdaInfer의 Early Exit 구조 (차원・학습 방식 포함한 상세 설명)

1. Logits 기반 Feature 추출 및 차원 축소 각 디코더 레이어에서 마지막 토큰의 hidden state $h \in \mathbb{R}^{d}$ (예: $d = 4096$)를 LM Head $W \in \mathbb{R}^{d \times V}$와 곱해 logit $z \in \mathbb{R}^{V}$ 계산 후, softmax를 통해 확률 분포 $p \in \mathbb{R}^{V}$를 생성한다. 이 확률 분포에서 다음 두 가지 scalar 통계치를 추출하여 feature vector $x_d \in \mathbb{R}^2$ 를 구성:
   • gap: 상위 두 토큰 확률 차이 → $\text{gap} = p_1 - p_2$
   • top prob: 가장 높은 확률값 → $\text{top prob} = p_1$

2. 이진 분류기 학습을 위한 데이터 구성 각 레이어에서 생성한 feature vector $x_d \in \mathbb{R}^2$에 대해 정답 레이블 $y_c \in {0,1}$ 을 다음 기준으로 부여:
   • 해당 레이어에서 예측한 토큰 $\hat{y}$가 최종 레이어에서 예측한 토큰 $y$ 와 동일하면 → label $y_c = 1$
   • 다르면 → label $y_c = 0$ 수식:

     $$ y_c = \begin{cases} 1 & \text{if } \hat{y} = y \ 0 & \text{otherwise} \end{cases} $$

     즉, L개의 레이어가 있는 LLM에 대해 하나의 입력 $x$ 는 총 L개의 (feature, label) 쌍 $\langle x_d, y_c \rangle$ 을 생성하게 된다.

3. 통계적 이진 분류기 학습 (SVM, CRF)
   • SVM: 각 레이어에서 독립적으로 분류. 입력 $x_d \in \mathbb{R}^2$ 만으로 “Exit 여부”를 판단.
   • CRF: 디코더 레이어 전체를 시퀀스로 보고, 레이어 간 문맥 의존성(e.g., 점진적 확신 증가)을 학습하여 더 정교한 exit 타이밍 예측 가능.

4. 추론 시 Early Exit 실행 추론 중 각 디코더 레이어에서 feature vector $x_d$를 추출 → 학습된 분류기에 입력 → 출력이 1이면 해당 레이어에서 inference 중단 후 LM Head로 최종 토큰 예측 수행.
   • 이때 사용되는 logit 계산은 마지막 토큰 하나에 한정되므로 전체 시퀀스 길이에 비해 연산량은 매우 작음 (논문 기준: 전체 FLOPs의 0.03% 수준).

5. 효율성과 확장성
   • feature vector가 2차원이므로, 분류기의 파라미터 수 및 연산량은 매우 작고 빠름.
   • 분류기는 task-independent하게 학습되므로 다양한 LLM에 generalize 가능.
   • Llama2, OPT 등 다양한 LLM에서 평균 17.8%의 레이어 pruning, sentiment task에서는 최대 43% pruning을 성능 저하 없이 달성함.

📌 요약 포인트

• 입력 차원: $x_d \in \mathbb{R}^2$
• 레이블 조건: $\hat{y} = y \rightarrow y_c = 1$
• 분류기 종류: SVM (독립적), CRF (시퀀스 모델링)
• 학습 데이터 수: 입력 하나당 $L$개의 훈련 샘플 생성
• 효율성: logit 연산은 마지막 토큰만 사용하므로 미미한 오버헤드

🔁 기본 LLM 추론 흐름 LLM은 한 번에 토큰 하나만 생성합니다.

하나의 토큰을 생성하려면 입력 토큰 시퀀스를 받아서, 디코더 블록 전체(L개 레이어) 를 순차적으로 통과해야 합니다.

마지막 레이어까지 도달한 후 LM Head를 통해 확률 분포 → 다음 토큰 선택 → 이 과정을 반복하여 시퀀스를 생성합니다.

즉,

“하나의 토큰 생성을 위해 디코더 전체를 1회 통과”

🚀 AdaInfer는 이 중간을 “스마트하게 잘라내는” 구조 디코더 블록을 통과하면서 매 레이어에서 EE predictor가 판단:

“지금 레이어까지만 계산해도, 예측이 충분히 확신 있어 보여?”

그렇다면 해당 레이어에서 멈추고, 그 상태로 다음 토큰을 확정합니다.

그리고 다음 토큰 생성을 위해 다시 디코더 블록을 layer 1부터 시작합니다.

문제점 => LLM vocabulary가 EE predictor의 online search space(LLM에서 $hidden_{dim} ⨉ vocabulary_{size}$ 가중치를 가진 linear 연산, 즉 LM Head)로 작동하여, 연산량에 상당한 영향을 미침

해결책 => Speculative token 생성을 통해 predictor의 search space를 $10^4$배 축소

3가지 Challenge

Challenge-1: Time-consuming predictor with high design complexity

기존 LLM EE predictor는 예측 전에 full search space(완전한 LM Head와 곱해지는)을 순회해야 하며, 이후 고차원 원시 데이터(4,000차원 초과)를 특성 분석이나 추출 없이 그대로 예측 입력으로 사용한다. 이러한 고차원 입력을 처리하기 위해 예측기는 AdaInfer의 SVM과 같은 기본 모델을 사용하며, GPU 병렬성을 고려하지 않아 전체 연산의 약 30%, 전체 추론 지연의 약 15%를 차지하게 된다.

Challenge-2: Under-utilization of layer-wise predictor deployment

기존 LLM EE system에서 predictor를 모든 decoder layer뒤에 삽입한다. 하지만 통계적으로 예측 성공의 확률은 skewed distribution을 따른다. => 대부분의 토큰이 특정 소수의 레이어에서만 EE 된다. => 대부분의 predictor 계산이 불필요.

Challenge-3: Exponential mapping complexity of predictor in speculative decoding.

Speculative decoding은 트리 기반의 토큰 구조를 통해, 초안 생성과 토큰 검증을 수행하는 방식으로, autoregressive decoding의 낮은 throughput 문제를 해결하고자했다.

Speculative decoding? LLM inference throughput을 높이기 위한 기법으로,

1. Lightweight Draft Model이 main LLM 대신 먼저 K개의 토큰을 "추측"하여 생성(빠름)
2. Main LLM이 해당 K개의 토큰이 실제로 나올 수 있는지 검증

• 정확하면 -> 그대로 채택
• 틀리면 -> 일부 교체, 다시 생성
  위에서 트리 기반이라고 한것은 첫 speculative token을 K개를 생성 그 다음 speculative token은 $K^2$개, 그 다음은 $K^3$개 ... 여서.

하지만 speculative decoder에서의 EE mapping은 각 토큰을 문맥적 의미 고려 없이 predictor의 독립적인 탐색 공간으로 취급하여, 매핑 복잡도가 기하급수적으로 증가하고, 높은 throughput의 이점을 활용 X.

speculative decoder에서의 EE mapping은 다음을 의미

📌 목적 Speculative Decoding 중에 나온 여러 후보 토큰(트리 구조)을 어느 레이어에서 종료할 수 있을지(Predictor의 exit 판단) 매핑하는 과정.

⚠️ 기존 방식의 문제 각 후보 토큰을 독립적으로 매핑함 (e.g., Tree 내 각 노드를 개별적인 입력으로 처리)

토큰 간 문맥/의미적 연관성을 고려하지 않음

결과적으로 Token 수가 늘면 Predictor search space도 지수적으로 늘어남

3가지 해결방법 요약

(1) Speculation-based lightweight predictor design at the algorithm level.

speculative 토큰의 확률 변화가 정답 여부와 강하게 상관되어 있다. => 이를 예측 특징으로 활용할 수 있는 여러 의미 있는 지표들을 추출한다.

+ GPU의 병렬성을 최대한 활용하기 위해, 우리는 예측기로 경량 MLP를 채택

(2) Two-level heuristic predictor scheduling at the system level.

추론 중 예측기의 통합 및 계산을 휴리스틱하게 제어하기 위해 오프라인 및 온라인 스케줄링을 포함하는 2단계 휴리스틱 예측기 스케줄링 기법

오프라인 스케줄링은 광범위한 통계 분석을 통해 얻은 skewed distribution을 기반으로 예측기를 할당한다.

온라인 스케줄링은 실행 시간 중에 수행되며, exit 레이어 위치의 context similarity에 기반한다. 즉, 현재 토큰의 exit 레이어가 이전 다섯 개 토큰의 exit 레이어로부터 ±2 레이어 이내에 있을 확률은 70%를 초과한다.

(3) Context-aware merged mapping for predictors at the mapping level.

매핑 수준에서의 문맥 인식 기반 예측기 병합 매핑. 기법 (2)에서 언급한 종료 레이어 위치 간 문맥 유사성에 기반하여, 우리는 이러한 특성이 트리 기반 speculative decoding에도 적용된다는 점을 지적한다. 이는 입력 토큰 트리 간에도 문맥적 의존성이 존재함을 의미한다. 이에 따라 speculative decoding을 지원하는 효율적인 GPU 구현 기반의 문맥 인식 기반 예측기 병합 매핑을 제안하며, 트리 기반 토큰의 각 경로를 하나의 하이퍼 토큰으로 병합함으로써 기하급수적 매핑 복잡도를 선형 복잡도로 전환

전체 SpecEE 구조

1. 입력 프롬프트를 받은 후, 오프라인 및 온라인 스케줄링 메커니즘으로 구성된 휴리스틱 스케줄링 엔진이 활성화가 필요한 예측기를 식별한다.
2. speculative 모델이 호출되어 speculative 토큰을 생성한다.
3. 연속된 디코더 레이어 쌍 사이마다, 예측기를 활성화해야 하는 경우에는 특징(feature)을 추출하고, 예측기 모델을 활용하여 추론을 계속 진행할지 종료할지를 판단한다.

위 1, 2, 3을 시퀀스 끝까지 반복

2. Background

생략

3. Motivation

3.1 Key Challenges of Early Exiting

vocabulary size가 크면 클수록, predictor가 순회해야하는 search space가 크다. => 어휘 공간을 줄여, online token prediction을 효과적으로 수행.

3.2 Key Insight

컴퓨터 시스템 설계에서의 추론(speculation) 개념과 2.2절에서 다룬 speculative decoding에서 영감을 받아, 우리는 speculative decoding에서 DLM(Draft LLM)의 역할이 TLM(Original LLM)을 위한 speculative 토큰을 생성하는 것이라 본다. TLM의 관점에서 볼 때, DLM의 출력은 실제로 항상 정답이 아닐 수는 있으나, 토큰 선택 범위(즉, 탐색 공간)를 좁히는 잠재적인 수단을 제공한다. => DLM을 학습시키는 목적은 TLM의 결과가 가능한 한 이 speculative 토큰들과 잘 일치하도록 만드는 데 있다. => 즉, 충분히 성능이 좋은 DLM을 사용한다면, TLM의 결과를 speculative 토큰이 포함된 좁은 공간(그림 2(a)에서 언급된 유효한 소규모 공간)으로 완전히 제한할 수 있다.

=> DLM이 충분히 강력하면, TLM이 굳이 모든 디코더 레이어를 사용하지 않아도 되고, 결과는 speculative token 내부에서 결정되므로 search space가 작아진다

4 Speculation-based Lightweight Predictor

4.1 Motivation: Time-consuming Predictor

Speculative 모델이 줄여주는 건 “TLM의 토큰 후보군”, 하지만 EE Predictor가 쓰는 feature는 여전히 “디코더 내부의 고차원 hidden state” 이기 때문에, search space가 줄었다고 해서 predictor의 연산량이 자동으로 줄지는 않는다.

4.2 Insight: Probability Shift

• (아래의 위 그림) 최종 출력이 축소된 search space에 존재하는 경우, 해당 토큰의 확률은 특정 레이어에서 급격히 상승하고, 반면, 다른 토큰들의 확률은 낮은 수준에서 안정적으로 유지. => EE를 했어야 되는 상황에서 EE를 하게 해줌
• (아래의 아래 그림) 최종 출력이 축소된 search space에 없는 경우, 모든 토큰의 확률은 낮은 값에서 안정적으로 유지. => EE를 안했어야 되는 상황에서 EE를 안하게 해줌

4.3 Approach: Lightweight Design

위의 motivation, insight를 통해, speculation-based lightweight predictor 설계. 이는

• feature selection
• judgment mechanism
• correction algorithm

단계로 이루어짐.

4.3.1 Feature selection

Predictor가 사용할 feature는 다음과 같다. (1) Speculative token logits 각 레이어의 출력(hidden_states)과 speculative 토큰에 해당하는 lm_head의 컬럼 간의 행렬 곱(1 × hidden_dim × num_speculatives) 의 결과로, LLM이 해당 speculative token에 대해 얼마나 확신을 가지는 지 직관적으로 반영.

(2) Local probabilities speculative token logits에 softmax 한 것

(3) Probability variation 현재 레이어와 이전 레이어 간의 local probabilities의 차이로, 레이어를 지남에 따라 확률이 어떻게 달라지는지 포착 => 양수이면, 해당 토큰에 더 확신이 생기고 있는 거임

보면, (a)에서 team의 probability gap을 보면 0.12, yourself 또한 0.12이다. 그러나, local probabilites를 보면, yourself가 더 높다. 즉, yourself가 exiting 확률이 더 높아야한다. => probability gap과 local probabilities를 같이 사용

+ (b)에서 yourself와 your는 local probabilities는 0.58로 같지만, speculative token logits는 your가 9.80으로 더 크다. 이럴때, 같은 local probability일 때, logit 값이 더 작다 => 상대적 우세 => yourself가 exiting 확률이 높음

=> (b)에서 왼쪽 case의 후보들 (your, yourself, my)이 마지막 layer의 lm head로 가서 최종 토큰 생성!

4.3.2 Judment mechanism

앞서 언급된 특징들을 기반으로, speculative 모델은 매번 4개의 speculative 토큰을 생성하도록 구성되며, 이에 따라 총 특징 차원은 12개(4 × 3)가 된다.

GPU의 Tensor Core가 제공하는 높은 연산 성능을 최대한 활용하기 위해, 우리는 기존의 머신러닝 방식(SVM 등) 대신 은닉 차원이 512인 2-레이어 MLP를 예측기로 채택하였다. 이 예측기는 ReLU 활성 함수를 사용하고, 이진 분류 작업을 처리하기 위해 출력층에 Sigmoid 함수를 설정한다.

입력된 특징은 예측기로 전달되며, 예측기의 출력이 사전 정의된 임계값(예: 0.5)과 비교되어 exiting 여부가 결정된다.

4.3.3 Verification algorithm

위에 Judgment mechanism에 추가로 Verificator를 장착하여, 안전장치를 하나 추가했다.

verificator는 pred가 threshold 보다 클 경우, Global LM Head를 통해(Hidden state * vocalbulary) global token logits를 계산하고, 제일 높은 확률을 가진 global token이 speculative 목록에 있다면, 해당 토큰을 출력하고, 추론을 종료한다. // 여기서 말하는 speculative 목록은

Q. 근데 그러면, speculative 방식을 써서 LM Head의 search space를 극단적으로 줄이는 방법이 의미가 없지 않나? 결국 verificator에서 global logits을 계산하니까

A. ㄴㄴ. 애초에 MLP Predictor를 통해 나온 pred 값이 threshold 보다 높을 결우가 드물게 일어남. (당연하지, 보통 exit 직전에 일어나겠지...) 그래서 global logits을 계산하는 횟수 자체가 줄어듦. 이전 방식 (AdaInfer)는 모든 레이어에서 exit predictor를 수행하는데, 그때마다 full LM Head를 사용했었음.

4.4. Evaluation

생략

5. Two-level Heuristic Scheduling Engine

5.1 Motivation

7.5 Ablation study에서 Llama2-7B를 사용했을 때, end-to-end inference speedup이 15% 밖에 안됨. => 그럼에도 평균 실행 layer 수는 23개, 이는 이론적으로 약 33%(32/(23+1))이 가능함을 시사 => 어딘가 문제가 있음

문제의 원인을 찾기위해, speculative model의 overhead를 생각해보면, 대략 decoder layer 한 개의 실행 시간. => 전체적인 inference speed를 저하시키는 원인은 predictor overhead

predictor overhead는 $(단일 predictor 실행시간)⨉(레이어 수)$이다. 이때, 단일 predictor 실행 시간은 다음 실험을 통해, 최적으로 구성하였기에 => 레이어 수를 줄여야됨

또한 Figure 10(a)(c)에서 나타난 바와 같이, 종료 확률이 하위 50%에 해당하는 레이어들(가로 점선 아래)의 확률 합은 20%를 넘지 않으며, 이는 해당 레이어들에서의 예측이 대부분 불필요함을 시사한다. 종료 확률이 낮은데 예측을 하고 있는 상황

Figure 10(b)에서 32레이어 쓸 때 보다, 8레이어 쓸 때, average forward layers 가 3.1 증가한다고 한다. => predictor가 너무 적어, 적절한 시점에 종료 실패해서...

=> end-to-end inference acceleration을 위해서는 적절한 predictor 수와 위치를 제어해야 한다.

5.2 Insights and Analysis: Skewed Distribution and Context Similarity

LLM inference 과정에서 model 선택과, 문맥 입력이라는 두 변수에 따라 예측기의 수와 위치를 동적으로 조정해야 한다.

Skewed Distribution

Figure 10(a)(c)에서 보여주는 예측기 결과 분포를 보면, 모델에 따라 편향의 정도가 다름을 알 수 있다.

Context Similarity

현재 토큰의 종료 레이어와 상관관계가 있는 레이어가 있나?에서 출발하여, 직전 N개의 토큰의 종료 레이어와의 관계에 주목.

통계 결과에 따르면, 현재 토큰의 종료 레이어는 이전 5개의 토큰의 종료 레이어로부터 ±2 레이어 이내에 위치할 확률이 약 80%에 달한다.

그래프 예시

가정 N(x-axis)이 3, 현재 7번 째 token 종료 layer 10.

N이 3이므로, 이전 3개 layer를 보자

6번 째 token: 8에서 EE → {6, 7, 8, 9, 10} 5번 째 token: 8에서 EE → {6, 7, 8, 9, 10} 4번 째 token: 10에서 EE → {8, 9, 10, 11, 12}

→ 수집된 exit layer = {6, 7, 8, 9, 10, 11, 12} (중복 제거)

이 수집된 exit layer set에 현재 종료 layer(10)가 포함될 확률이 (y-axis)

Figure 11에 따르면, N이 5일 때, 수집된 exit layer set에 현재 종료 layer가 포함될 할 확률이 80%

만약 이런 수집된 exit 레이어와 현재 토큰의 종료 레이어가 상관관계가 없다면, 실제 현재 토큰의 종료 레이어가 수집된 exit 레이어에 들어갈 확률은 $\frac{(수집된 exit 레이어 수)}{전체 레이어 수} = \frac{7}{32} = 0.218$ 이다. // 위의 그래프 예시임

하지만 실제 현재 토큰의 종료 레이어가 수집된 exit 레이어에 들어갈 확률이 0.8이므로 이 둘은 상관관계가 있다고 볼 수 있다.

=> exit 레이어의 위치는 뚜렷한 context similarity가 존재한다.

Q. 왜 이전 5 토큰을 기준으로 exit layer를 수집하나? A. 참고하는 과거 토큰이 많으면 많을수록, 수집된 exit 레이어의 수 많아짐. => predictor 많아짐!!! => 연산 多 참고하는 과거 토큰이 적으면 적을수록, hit ratio가 확 떨어짐.

=> 연산이 적으면서도 hit ratio가 높은, trade-off 지점이 과거 5개의 토큰을 보는 거임

5.3 Approach: Two-level Heuristic Scheduling

위의 Figure 12에서 보이듯 offline scheduling과 online scheduling 두 부분으로 구성된다.

Offline Scheduling

서로 다른 LLM에 대한 데이터를 사전에 수집하여 사용. 수많은 프롬프트를 이용하여 모든 predictor가 완전히 통합된 상태로 LLM inference를 수행한뒤, 각 predictor에서 수집된 데이터를 기반으로 predictor의 사용 빈도를 순위화한다. 이 결과는 모델 구성 파라미터로 통합되며, 이는 LLM마다 달라지며 한 번만 오프라인으로 수행하면 된다.

Online Scheduling

앞서 언급했듯 context similarity에 기반하기 위해, 최근 N 토큰을 기준으로 exit 레이어를 수집했다. 이때 N이 circular queue의 길이가 된다.

또한, 전체 레이어 수 L에 해당하는 길이의 배열도 사용한다.

circular queue는 최근 N개의 토큰에 대한 exit layer 위치를 순차적으로 기록한다.

전체 레이어 수 L에 해당하는 길이의 배열의 i번째 요소는 해당 레이어가 큐에 기록된 최근 N개의 토큰의 exit 레이어 주변 (±2)에 위치한 횟수를 누산한다.

최종적으로 예측기의 수와 위치는 offline scheduling에서 선택된 결과의 하위 집합과, online scheduling에서 파생된 결과의 합집합에 의해 결정된다.

위 그림 설명

길이 5짜리 버퍼로, 최근 5개의 토큰이 어느 레이어에서 exit됐는지를 저장

최근 5개 토큰의 exit 레이어 = [16, 19, 20, 19, 20] = [16, 19, 20]

16레이어 주변 (±1)에 위치한 횟수를 누산 => 16 → [15, 16, 17] → array[15]++, array[16]++, array[17]++ 19 → [18, 19, 20] → array[18]++, array[19]++, array[20]++ 20 → [19, 20, 21]→ array[19]++, array[20]++, array[21]++

=> | Layer | Count | | ----- | ----- | | 15 | 1 | | 16 | 1 | | 17 | 1 | | 18 | 1 | | 19 | 2 | | 20 | 2 | | 21 | 1 |

Union 생성: 카운트가 1 이상인 레이어들의 모음 + offline에서 exit된 횟수가 상위 M개인 레이어들의 모음

6. Context-aware Merged Mapping for Predictor

6.1 Motivation

Speculative model은 autoregressive 방식으로 여러 개의 토큰을 병렬적으로 생성하여, token tree를 구성한다.

① 일반 Autoregressive 방식 "How are you" → "?" 생성

32개 레이어 × 1회 토큰

→ 한 번에 하나의 토큰 생성만 가능

② Speculative Decoding (DLM) "How are you" → Top-3 토큰 생성: ["I", "It", "Thank"]

3개의 후보 토큰 → 병렬

이어서:

"I" → [thank, am, can]

"It" → [is, go, s]

"Thank" → [you, ...]

→ 이렇게 해서 “Token Tree”가 만들어지고 → 이후 이 전체 후보 트리 구조를 한 번의 TLM 추론으로 검증

Q. “DLM이 트리 전체를 만들고 나서야 TLM이 검증을 시작한다면, 그 트리 생성 자체가 느리면 전체 속도 이득이 상쇄되는 것 아닌가?” A. DLM 자체가 엄청 빨라서 ㄱㅊ

검증 추론 과정에서 EE를 적용할 때, 현재의 predictor mapping 방식은 토큰 트리의 각 토큰을 문맥적 의미를 고려하지 않은 독립적인 탐색 공간으로 취급한다.

즉, 각 분기의 토큰들을 predictor에 mapping하는 것이다. 위 Figure 13에서는 ✅ "?" → Top-3 후보 ["I", "It", "Thank"] → 1개의 predictor: 이 3개를 묶어서 exit 판단 ✅ "I" → ["am", "can", "think"] ✅ "It" → ["is", "go", "do"] ✅ "Thank" → ["you", "him", "them"] → 각각 1개 predictor씩 할당

총 4개의 predictor이다. 이런 predictor는 서로 독립적으로 작동하기 때문에, 전체 mapping 복잡도는 각 예측기의 복잡도의 곱으로 계산되어, 지수적으로 증가!

6.2 Approach: Context-aware Merged Mapping Algorithm.

6.3 Extension: Support for Orthogonal Acceleration Techniques

7. Evaluation

7.1 Experimental Setup

7.2 Evaluation on Speedup and Throughput

7.3 Hardware Evaluation

7.4 Overhead Evaluation

7.5 Ablation Study

8. Conclusion

BinaryViT에 대한 내용

Abstract

Binarization은 weight와 activation이 binary 상태일 때, popcount 연산을 사용함으로써, ViT model의 크기와 computational cost를 크게 줄일 수 있다.

Vanilla ViT들은 CNN이 갖고 있는 핵심적인 구조적 특성들이 결여되어 있어, binary CNN보다 representational capability가 현저히 떨어진다.

1. Introduction

Activation, Weight 모두 binarize된 ViT의 accuracy를 full-precision model과 비슷하게 하기 위한 방법들

1. Distillation
2. Scaling factor를 learnable parameter로 설정
3. Bi-RealNet: 추가적인 residual connection을 통해, 정보를 보존
4. ReActNet: sign함수 앞에 learnable threshold를 추가하고, 각 residual connection 뒤에 RPReLU activation 함수를 추가하여, output distribution을 재구성
5. BiMLP: multibranch block을 제안하여 patch 혼합과, channel 혼합이 동시에 이루어지게 하여, FC layer의 제한된 표현력을 극복

2. Designing a full binarized ViT baseline

CNN과 BERT에 적용된 기존 binarization 기법을 활용하여, A1W1 Vanilla ViT 설계

2.1. Binarized fully-connected layer

$$X∈{\R}^{N⨉D_{in} }$$에 sign()를 적용하기 직전에, threshold vector $${\beta}X∈{\R}^{D{in}}$$를 적용한다.

$$W∈{\R}^{D_{in}⨉D_{out}}$$의 threshold는 $${\mu}(W)∈\R^D_{out}$$이다. $${\mu}(W)$$는 $$ W$$의 평균이다.

전체 행렬 곱셈식은 다음과 같다. $$Y(X) = \alpha_Wsign(X+\beta_X) ⊗ sign(W-\mu(W))$$

실제 구현은 조금 다르다($$\beta_X$$가 없고, popcount가 아니며, $$\alpha_W$$가 $$Sign(W - \mu(W))$$에 곱해진다. 아래 코드를 실행하면 그 과정을 알 수 있다.

PS) 아니다, $\beta_X$가 없는게 아니고, $Y(X)$를 호출하는 쪽에서 $X+\beta_X$처리를 하고, $X$로 넣기에 BinaryQuantizer에 없는 것이다.

import torch
import torch.nn as nn

# BinaryQuantizer: 이진 양자화를 위한 autograd.Function
class BinaryQuantizer(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        # forward 단계: 입력 텐서를 이진화 (부호 함수: -1, 0, 1 반환)

        ctx.save_for_backward(input)  # backward 단계에서 사용하기 위해 입력 저장
        out = torch.sign(input)
        return out

    @staticmethod
    def backward(ctx, grad_output):
        # backward 단계: 기울기를 계산함

        input = ctx.saved_tensors  # 저장한 입력값을 불러옴
        input = input[0]
        # 입력이 [-1, 0] 구간에 해당하는지 판별 (float 텐서로 변환)
        indicate_leftmid = ((input >= -1) & (input <= 0)).float()
        # 입력이 (0, 1] 구간에 해당하는지 판별
        indicate_rightmid = ((input > 0) & (input <= 1)).float()
        # 각 구간에 대해 선형 보정된 기울기를 계산한 후 grad_output과 곱함
        grad_input = (indicate_leftmid * (2 + 2 * input) + indicate_rightmid * (2 - 2 * input)) * grad_output.clone()
        return grad_input

    # QuantizeLinear: 양자화된 선형(fully-connected) 계층

class QuantizeLinear(nn.Linear):
    def __init__(self,  *kargs, bias=False, config=None):
        super(QuantizeLinear, self).__init__(*kargs, bias=bias)
        # 가중치와 입력의 양자화 비트 수를 config에서 가져옴
        self.weight_bits = config.weight_bits
        self.input_bits = config.input_bits

        # 가중치 양자화 방법 선택 (비트 수에 따라 다른 양자화 함수 사용)
        if self.weight_bits == 1:
            self.weight_quantizer = BinaryQuantizer
        elif self.weight_bits == 2:
            # self.weight_quantizer = TwnQuantizer
            # 가중치 클리핑 범위를 버퍼에 등록
            self.register_buffer('weight_clip_val', torch.tensor([-config.clip_val, config.clip_val]))
        elif self.weight_bits < 32:
            # self.weight_quantizer = SymQuantizer
            self.register_buffer('weight_clip_val', torch.tensor([-config.clip_val, config.clip_val]))

        # 입력(activation) 양자화 방법 선택
        if self.input_bits == 1:
            self.act_quantizer = BinaryQuantizer
        elif self.input_bits == 2:
            # self.act_quantizer = TwnQuantizer
            self.register_buffer('act_clip_val', torch.tensor([-config.clip_val, config.clip_val]))
        elif self.input_bits < 32:
            # self.act_quantizer = SymQuantizer
            self.register_buffer('act_clip_val', torch.tensor([-config.clip_val, config.clip_val]))


    def forward(self, input):

        # 가중치 양자화: 비트 수에 따라 다르게 처리
        if self.weight_bits == 1:
            print("=== Weight ===\n")
            print(f"W : {self.weight}\n")

            # 이진 가중치의 경우, 각 행별 평균 절댓값을 스케일링 팩터로 사용
            scaling_factor = torch.mean(abs(self.weight), dim=1, keepdim=True)
            print(f"a_W : {scaling_factor}\n")
            scaling_factor = scaling_factor.detach()

            real_weights = self.weight - torch.mean(self.weight, dim=-1, keepdim=True)
            print(f"M(W) : {torch.mean(self.weight, dim=-1, keepdim=True)}\n")
            print(f"W - M(W) : {real_weights}\n")

            binary_weights_no_grad = scaling_factor * torch.sign(real_weights)
            print(f"sign(W - M(W)) : {torch.sign(real_weights)}\n")

            cliped_weights = torch.clamp(real_weights, -1.0, 1.0)
            # 이진화된 가중치로 기울기 흐름은 유지하면서 클램핑된 값을 사용
            weight = binary_weights_no_grad.detach() - cliped_weights.detach() + cliped_weights
            print(f"a_W * sign(W - M(W)) : {weight}\n")
        elif self.weight_bits < 32:
            # 지정된 양자화 함수를 적용하여 가중치를 양자화
            weight = self.weight_quantizer.apply(self.weight, self.weight_clip_val, self.weight_bits, True)
        else:
            # 양자화를 사용하지 않는 경우 원래의 가중치를 사용
            weight = self.weight

        # 입력(activation) 양자화: 입력이 이진인 경우에만 적용
        if self.input_bits == 1:
            print("=== Activation ===\n")

            print(f"X : {input}\n")
            input = self.act_quantizer.apply(input)
            print(f"sign(X + b_X) : {input}\n")

        # 선형 연산 수행
        print("=== Y ===\n")
        out = nn.functional.linear(input, weight)
        print(f"(sign(X + b_X)) * (a_W * sign(W - M(W))) : {out}\n")

        # bias가 존재하면 bias 추가 (출력 텐서 크기에 맞게 확장)
        if not self.bias is None:
            out += self.bias.view(1, -1).expand_as(out) 

        return out


# ------------------------------------------------------------------------------
# 테스트 코드를 위한 DummyConfig 클래스 정의 (실제 config 대신 사용)
class DummyConfig:
    def __init__(self, weight_bits, input_bits, clip_val):
        self.weight_bits = weight_bits
        self.input_bits = input_bits
        self.clip_val = clip_val

# QuantizeLinear 테스트 함수
def test_quantize_linear():
    print("\n=== QuantizeLinear 테스트 ===\n")
    # Dummy config 생성: 가중치와 입력 모두 1비트, 클리핑 값 1.0
    config = DummyConfig(weight_bits=1, input_bits=1, clip_val=1.0)
    # 입력 차원 5, 출력 차원 3인 QuantizeLinear 계층 생성
    qlinear = QuantizeLinear(5, 3, bias=False, config=config)
    # print(f"초기 W : {qlinear.weight}\n")

    # 더미 입력 생성: 배치 크기 2, 입력 차원 5
    dummy_input = torch.randn(2, 5)
    # print(f"X : {dummy_input}\n")

    # forward pass 실행
    output = qlinear(dummy_input)
    print(f"Y : {output}\n")

if __name__ == "__main__":
    test_quantize_linear()

이렇게 행렬 곱셈 이후엔, Batch Normalization을 적용하고, residual connection을 연결하고, RPReLU activation function을 적용한다.

$$BiFC(X) = RPReLU[BN{\alpha_Wsign(X+\beta_X) ⊗ sign(W-\mu(W))} + R(X)]$$

$$ \mathbb{R}(\mathbf{X}) = \begin{cases} \mathbf{X}, & c_{in} = c_{out} \ \text{Cat}([\mathbf{X}, \text{for } i \text{ in range}(n)], \text{dim}=2), & nc_{in} = c_{out} \ \left( \frac{1}{n} \sum_{i=1}^{n} \mathbf{X}(:, :, \frac{i-1}{n} d_{in} : \frac{i}{n} d_{in}) \right), & c_{in} = n c_{out} \end{cases} $$

2.2. Binarized vision transformer

기존 Binary ViT의 과정

1. 처음 이미지 $$x∈\R^{H⨉W⨉C}$$는 embedding layer에서, patch $$x_p∈\R^{N⨉(P^2•C)}$$로 분할된다. (이때, patch의 수 $$N = HW/P^2$$이다.)
   
   
2. 각 patch에 대해 linear projection, $$E∈\R^{(P^2⨉C)⨉D}$$가 적용되어, $$x∈\R^{N⨉D}$$가 되고, class token embedding $$x_{class}∈\R^D$$가 추가되어 $$x∈\R^{(N+1)⨉D}$$가 되고, position embedding $$E_{pos}∈\R^{(N+1)⨉D}$$가 더해져,

$$H∈\R^{(N+1)⨉D}$$가 출력된다.
$$H = [\mathbf{x}{class}; \mathbf{x}_p^1E; \mathbf{x}_p^2E; ,,,; \mathbf{x}_p^NE] + E{pos}$$ (';'는 행방향으로 붙이는 것임)

3. Embedding layer의 output $$H∈\R^{(N+1)⨉D}$$가 첫 번째, transformer block(encoder block)에 입력으로 들어간다.

4. Transformer block에서 입력 $$H∈\R^{(N+1)⨉D}$$는 pre-batch-normalization layer를 거쳐, $$\hat{H} = BN_1(H) ∈ \R^{(N+1)⨉D}$$로 변환 된다.

5. $$N_H$$개의 atttention head를 가진 binarized된 MHA module에서 $$\hat{H}$$를 BiFC(Binarized FC)에 넣으면, $$Q_h,K_h,V_h$$가 다음과 같이 계산된다.
$$Q_h=BiFC_{Q_h}(\hat{H}h)∈\R^{(N+1)⨉D_h}$$ $$K_h=BiFC{K_h}(\hat{H}h)∈\R^{(N+1)⨉D_h}$$ $$V_h=BiFC{V_h}(\hat{H}h)∈\R^{(N+1)⨉D_h}$$
$$D_h=D/N_H$$
6. Attention score 구하기. Attention score은 다음과 같이 계산된다. $$A_h=Rsign(Q_h)•Rsign(K_h^T)∈\R^{(N+1)⨉(N+1)}$$
7. Attention score $$A_h$$를 softmax하고, 0 또는 1로 binarize하여, attention probability를 얻는다. 그 식은 다음과 같다.
$$\mathbf{P}_h = \alpha{\mathbf{P}} \left[ \sigma \left( \frac{1}{\alpha_{\mathbf{P}}} \mathrm{Softmax} \left( \frac{\mathbf{A}h}{\sqrt{D_h}} \right), 0, 1 \right) \right] ∈ \R^{(N+1)⨉(N+1)}$$
이때, $$\alpha_P ∈ \R$$는 learnable scaling factor이다.
8. Attention probability $$P_h ∈\R^{(N+1)⨉(N+1)}$$를 $$Rsign()$$으로 binarized된 $$\bar{V}_h=Rsign(V_h) ∈ R^{(N+1)⨉D}$$과 곱하여, 각 token에 대해 value 정보를 반영한다.
$$P_h•\bar{V}_H∈ R^{(N+1)⨉D_h}$$

9. Token에 대한 value 정보가 반영된 attention probability를 BatchNormalization 해주고, Q, K, V 정보 보존을 위해, $$Q_h, K_h, V_h$$를 residual connection으로 연결한다. 그 값을 RPReLU를 거치면, 하나의 head에 대한 output이 나온다. 그 식은 아래와 같다.
$$head_h=RPReLU(BN{at}(P_h\bar{V}H)+Q_h+K_h+V_h)∈ \R^{(N+1)⨉D_h}$$

10. 모든 head($$N_H=D/D_h$$)의 output($$head_h$$)는 서로 concatenate된 후, $$BiFC$$를 거치면, 전체 Bi-MHA의 output이 나온다. 그 식은 아래와 같다.
$$F=BiFC_O(Cat(head_1, ..., head{N_H}))∈ \R^{(N+1)⨉D}$$

11. Main residual connection이 Bi-MHA의 output에 적용된다. 그 식은 아래와 같다.
$$F=BiFC_O(Cat(head_1, ..., head_{N_H}))+H∈ \R^{(N+1)⨉D}$$
이때, $H$는 다음과 같고, 2번 과정에서의 $H$이다.
$$H = [\mathbf{x}{class}; \mathbf{x}_p^1E; \mathbf{x}_p^2E; ,,,; \mathbf{x}_p^NE] + E{pos}∈\R^{(N+1)⨉D}$$

12. Residual output $F∈\R^{(N+1)⨉D}$는 두 번째 pre-batch-normalization layer인 $BN_2$를 거쳐 정규화된다. 그 식은 아래와 같다. $\hat{F}=BN_2(F)∈\R^{(N+1)⨉D}$
    
    

13. 정규화된 $\hat{F}∈\R^{(N+1)⨉D}$는 두 개의 BiFC로 구성된 BiFFN을 통과하고, 마지막으로 BiFFN의 output에 두 번째 main residual connection이 적용되어, 다음과 같이 최종 출력 $R$을 얻는다. $R=BiFFN(\hat{F})+F∈\R^{(N+1)⨉D}$
    
    

Distillation

Binarized된 ViT의 성능을 향상시키기 위해, student model의 logit과 teacher model의 logit 간의 soft crossentropy loss를 최소화함으로써, full-precision model의 knowledge를 binarized model로 distill한다.

3. What else do binary CNNs have that binary transformers do not have?

위 2장에서 언급된 모든 기법들을 적용하더라도 (Table 1의 ReActNet), 정확도가 낮다. 이는 대부분의 SOTA binary CNN 정확도보다 훨씬 낮다. => CNN architecture의 세부 요소 및 특성을 분석하여, ViT에 적용 => Binary activation / weight로도 많은 수의 클래스를 가진 dataset에서도 accuracy ↑

그러기 위해, Binary ViT의 표현력을 증가시키는 방법 3가지

1. Global Average Pooling(GAP)를 classifier layer 전에 삽입

   

2. Multiple average pooling branch를 추가

   

3. CNN에서 피라미드 구조를 차용하기

추가적으로 ResNet, MobileNet의 설계 아이디어 차용

Main residual branch의 scale이 Bi-MHA, Bi-FFN 출력 같은 main branch의 scale을 압도하지 않도록 residual branch 앞에 affine 변환을 배치.

3.1. Global average pooling before classifier layer

Binary CNN: classifier 앞에 average pooling layer 有 Vanilla ViT: classifier 앞에 단일 cls-token pooling layer만 有

==> cls-token pooling을 Gloval average pooling(GAP)을 통해 모든 token 정보를 반영하자! + Embedding에서 cls-token embedding을 제거하자!

==> 위 2.2.의 과정 2의 $$H = [\mathbf{x}{class}; \mathbf{x}_p^1E; \mathbf{x}_p^2E; ,,,; \mathbf{x}_p^NE] + E{pos}$$ 이 $$H = [\mathbf{x}p^1E; \mathbf{x}_p^2E; ,,,; \mathbf{x}_p^NE] + E{pos}$$ 로 교체된다. 그렇게 되면 당연히, $H, E_{pos} ∈\R^{N⨉D}$ 가 된다.

// 연산량 소폭증가, 성능 대폭 증가

3.2. More branches

Binary Conv layer가 Binary FC layer보다 3배 정도 높은 표현력을 가진다.

==> Transformer block 내부의 Bi-FFN 옆에 4개의 average pooling branch를 추가.

각 average pooling layer는 다음과 같다.

• 1⨉5 Avg. Pool
• 1⨉3 Avg. Pool
• 5⨉1 Avg. Pool
• 3⨉1 Avg. Pool

// 연산량 소폭증가, 성능 대폭 증가

3.3. Scaling right before residual connection

사전 지식: Affine 변환(선형 변환으로 생각하면 됨)(ex, BN)은 model의 표현력에 영향을 미치지 않는다.

깊은 layer의 main branch에서 전달되는 정보가 residual branch에 의해 압도되어, 제 성능 발휘 못하는 문제가 생긴다.

Affine 변환이 표현력에 영향을 미치지 않지만, 여러가지 실험을 통해 residual connection 전에 affine 변환을 사용하는 것이 전혀 사용하지 않는 것보다 좋다는 것을 알게되었다.

==> Attention과 FFN의 각 main residual connection 마다 main branch(Bi-MHA, Bi-FFN의 output)에 affine 변환을 한다.

기존 방법

• $$F=BiFC_O(Cat(head_1, ..., head_{N_H}))+H∈ \R^{(N+1)⨉D}$$
• $R=BiFFN(\hat{F})+F∈\R^{(N+1)⨉D}$

Main branch에 Affine transformation 적용 방법

• $F=\alpha_1⊙BiFC_O(Cat(head_1, ..., head_{N_H}))+\beta_1+H∈ \R^{N⨉D}$
• $R=\alpha_2⊙BiFFN(\hat{F})+\beta_2+F∈ \R^{N⨉D}$
• $\alpha_1, \alpha_2, \beta_1, \beta_2 ∈ \R^{D}$

코드 비교

ReActNet

DATA_DIR=/path/to/dataset

torchrun --nproc_per_node=8 main.py \
    --num-workers=40 \
    --batch-size=64 \
    --epochs=300 \
    --dropout=0.0 \
    --drop-path=0.0 \
    --opt=adamw \
    --sched=cosine \
    --weight-decay=0.00 \
    --lr=5e-4 \
    --warmup-epochs=0 \
    --color-jitter=0.0 \
    --aa=noaug \
    --reprob=0.0 \
    --mixup=0.0 \
    --cutmix=0.0 \
    --data-path=${DATA_DIR} \
    --output-dir=logs/reactdeit-small-patch16-224 \
    --teacher-model-type=deit \
    --teacher-model=configs/deit-small-patch16-224 \
    --teacher-model-file=logs/deit-small-patch16-224/best.pth \
    --model=configs/deit-small-patch16-224 \
    --model-type=extra-res \
    --replace-ln-bn \
    --weight-bits=1 \
    --input-bits=1 \
    --enable-cls-token \
    --disable-layerscale \
    # --resume=logs/reactdeit-small-patch16-224/checkpoint.pth \
    # --current-best-model=logs/reactdeit-small-patch16-224/best.pth \

BinaryViT

DATA_DIR=/path/to/dataset

torchrun --nproc_per_node=8 main.py \
    --num-workers=32 \
    --batch-size=64 \
    --epochs=300 \
    --dropout=0.0 \
    --drop-path=0.0 \
    --opt=adamw \
    --sched=cosine \
    --weight-decay=0.00 \
    --lr=5e-4 \
    --warmup-epochs=0 \
    --color-jitter=0.0 \
    --aa=noaug \
    --reprob=0.0 \
    --mixup=0.0 \
    --cutmix=0.0 \
    --data-path=${DATA_DIR} \
    --output-dir=logs/binaryvit-small-patch4-224 \
    --teacher-model-type=deit \
    --teacher-model=configs/deit-small-patch16-224 \
    --teacher-model-file=logs/deit-small-patch16-224/best.pth \
    --model=configs/binaryvit-small-patch4-224 \
    --model-type=extra-res-pyramid \
    --replace-ln-bn \
    --weight-bits=1 \
    --input-bits=1 \
    --avg-res3 \
    --avg-res5 \
    # --resume=logs/binaryvit-small-patch4-224/checkpoint.pth \
    # --current-best-model=logs/binaryvit-small-patch4-224/best.pth \

ReActNet, BinaryViT 두 코드 모두에서의 Transformer block(class ViTLayer(nn.Module))에서Bi-MHA, Bi-FFN이 호출하는 class ViTOutput(nn.Module)

class ViTOutput(nn.Module):
    def __init__(self, config: ViTConfig, layer_num, drop_path=0.0) -> None:
        super().__init__()
        self.dense = QuantizeLinear(config.intermediate_size[config.stages[layer_num]], config.hidden_size[config.stages[layer_num]], config=config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

        self.move = nn.Parameter(torch.zeros(config.intermediate_size[config.stages[layer_num]]))
        self.norm = config.norm_layer(config.hidden_size[config.stages[layer_num]], eps=config.layer_norm_eps)
        self.rprelu = RPReLU(config.hidden_size[config.stages[layer_num]])
        self.pooling = nn.AvgPool1d(config.intermediate_size[config.stages[layer_num]] // config.hidden_size[config.stages[layer_num]])
        self.layerscale = LayerScale(config.hidden_size[config.stages[layer_num]]) if not config.disable_layerscale else nn.Identity()


    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        out = self.norm(self.dense(hidden_states + self.move)) + self.pooling(hidden_states)
        out = self.rprelu(out)
        out = self.dropout(out)

        out = self.layerscale(out)

        out = self.drop_path(out)

        return out

반면, BinayViT에서는 Affine 변환을 하고 있음을 알 수 있다.

3.4. Pyramid structure

최신 SOTA Binary CNN은 고해상도에서 저해상도로 점진적으로 feature map 크기를 줄이고, hidden dimension은 증가시키는 pyramid structure를 갖는다.

이런 pyramid structure는 binary nn의 표현력을 향상시킨다.

==> pyramid structure를 통해, 계산 복잡도를 증가시키지 않으면서도, 표현력 ↑

첫 번째와 두 번째 스테이지에서 sequence size가 3316과 784이므로, 이 크기에서 attention을 적용하는 것은 계산 비효율적이다.

==> key, value 행렬을 계산하기 직전에 입력에 downsampling을 한다.

이렇게 downsampling된 값들은 residual connection 이전에 upsampling된다.

그 전체 식은 아래와 같다.

$head_h = RPReLU(BN_{at}(P_h\bar{V}_h) + Q_h + N(K_h) + N(V_h)$ 여기서 $N()$은 upsampling으로, nearest-neighbor interpolation function이다.

$K_h=BiFC_{K_h}(BiFC_R(AvgPool(\hat{H}))h)$ $V_h=BiFC{V_h}(BiFC_R(AvgPool(\hat{H}))_h)$

$AvgPool()$은 kernel 크기와 stride가 R인 average pooling이다.

* * * * 피라미드 코드의 ViTEncoder 부분의 ViTLayer를 반복 호출하는 루프에서 PatchEmbed 클래스를 주기적으로 호출하는 부분이다. * * * *

Survey of Quantization에 대한 내용

Abstract

continuous한 real-value 숫자들을 discreted한 숫자 집합에 분배하는 방식을 어떻게 설계해야 최소한의 bit 수로 요구를 충족시키면서도 computation accuracy를 최대화할 수 있을까?'에 대한 내용이다.

이러한 quantization 문제는 memory / computation resource가 심각하게 제한된 경우 중요하다.

1. Introduction

NNs의 accuracy가 크게 개선되었다. 이러한 accuracy 향상은 종종 over-parameterized된 대규모 model을 통해 이루어졌다.

==> resource가 제한된 많은 응용 환경에서는 이를 배포할 수 없음

위 문제의 해결법

a) Designing efficient NN model architectures

NN model의 micro-architecture 또는 macro-architecture의 optimization에 중점을 둔 연구. 즉, model size, depth, width 등의 주어진 제약 조건 내에서 적합한 NN architecture를 자동으로 찾기위해 AutoML 및 방법을 설계하는 연구이다. 즉, NAS(Nureal Architecture Search).

b) Co-designing NN architecture and hardware together

특정 대상 hardware paltform에 NN architecture를 적응시키고 co-design하는 것. 이는 NN 구성 요소(latency, energy)에 따른 overhead는 hardware-dependent하기 때문이다. 예를들어, 전용 cache hierarchy를 가진 HW는 그런 hierarchy가 없는 HW에 비해 bandwidth 의존적인 연산을 훨씬 효율적으로 실행할 수 있다.

c) Pruning

saliency(중요도)가 낮은 뉴런을 제거하여 spare computational graph를 만드는 것. // saliency가 낮은 뉴런은 제거하더라도 model 출력이나 loss function에 최소한의 영향을 미치는 뉴런이다.

unstructured pruning

saliency가 낮은 뉴런을 위치에 관계없이 제거하는 것.

대부분의 뉴런을 aggressive하게 pruning할 수 있다. 그러나 이방식은 sparse matrix operation을 초래하며, 이러한 operation은 accelerate하기 어렵고, 일반적으로 memory-bounded하다.

Q1. 왜 sparse matrix operation은 accelerate 하기 어렵고, 일반저으로 memory-bounded한가? A1.

structured pruning

saliency가 낮은 뉴런을 위치에 관계없이 제거하는 것이 아닌, 특정 parameter group(ex) entire convolutional filters)을 제거하는 것. 이는 layer과 weight matrix의 input, output 형상을 변경하게 되어, 여전히 dense한 matrix operation을 허용한다.

d) Kenowledge distillation

큰 모델을 training 한 후 이를 teacher model로 사용하여 더 간결한(student) model을 training하는 과정을 포함한다. student model을 training할 때 "hard" class label 대신, teacher model이 생성한 "soft probability"를 활용한다. 이 확률은 input에 대한 더 많은 정보를 포함할 수 있다. distillation은 다른 quantization, pruning에 비해, 공격적인 compression시 꽤나 큰 accuracy loss를 겪는 경향이 있다.

e) Quantization

Quantization은 NN model의 training, inference에서 모두 뛰어나고 일관된 성공을 보여준 접근법이다.

특히, half-precision and mixed-precision training의 돌파구는 AI accelerator에서 한 차원 높은 troughput을 가능하게 하는 주요 동력으로 작용했다. 그러나 half-precision 이하로 내려가는 것은 상당한 tuning 없이는 매우 어렵다는 것이 증명되었으며, 최근 quantization 연구의 대부분은 inference에 초점을 맞추고 있다.

2. General History Of Quantization

Quantization은 큰 집합의 input 값을 작은 집합의 outpur 값으로 매핑하는 방법으로, 오랜 역사를 가지고 있다. 반올림과 truncation은 전형적인 예이다

A) Quantization in Neural Nets

NN은 quantization 문제에 독특한 도전과 기회를 제공한다.

1. NN의 inference와 trainin은 모두 computationally intensive하여, 수치 값을 효율적으로 표현하는 것이 특히 중요하다.
2. 현재의 NN은 over-parameterized 되어 있어, accuracy에 영향을 미치지 않고, bit precision을 줄일 수 있는 충분한 여지가 있다.
3. NN은 공격적인 quantization과 극단적인 discretization에 매우 robust하다. 이러한 강인성은 NN이 가진 parameter 수에 따른 새로운 자유도를 제공한다.
4. NN의 서로 다른 layer는 loss function에 대해 다른 영향을 미치며, 이는 quantization에 대한 mixed-precision 접근 방식을 고무한다.

3. Basic Concepts Of Quantization

A. Problem Setup and Notations

NN이 학습 가능한 매개변수 ${W_1, W_2, ..., W_L}$을 가지는 $L$개의 layer로 구성되어 있다고 가정하자. 이때, 모든 parameter의 조합을 $\theta$로 나타낸다. WLOG, supervised learning problem에 초점을 맞추며, 이 문제의 명목적인 목표는 다음의 empirical risk minimization 함수를 최적화하는 것이다.

$L(\theta) = \frac{1}{N} \sum_{i=1}^{N} l(x_i, y_i; \theta)$                $(1)$

$(x,y)$ : input data와 이에 해당하는 label $l(x,y;\theta)$ :loss function $N$ : 전체 데이터 포인트의 수 $h_i$ : $i$번째 layer의 input hidden activation $a_i$ : $i$번째 layer의 output hidden activation 학습된 model parameter $\theta$가 fp precision으로 저장되어 있음.

Quantization의 목표는 parameter $\theta$ 뿐만 아니라 intermediate activation map(즉, $h_i, a_i$)의 precision을 낮은 precision으로 줄이면서도 model의 일반화 능력 및 accuracy에 미치는 영향을 최소화 하는 것이다.

B. Uniform Quantization

NN의 weight와 activation을 유한한 값 집합으로 quantization 할 수 있는 함수를 정의해야한다. ==> real value 값을 input으로 받아, 이를 낮은 precision의 범위로 매핑한다. 아래 Figure 1에서 이를 시각적으로 설명하고 있다.

널리 사용되는 양자화 함수는 다음과 같이 정의된다.

$Q(r) = Int(\frac{r}{S}) - Z$                $(2)$

$Q$ : quantization operator $r$ : real value $S$ : 실수형 scaling factor $Z$ : 정수형 zero point $Int$ : 실수 값을 정수 값으로 매핑

Q1. 위 식(2)의 quantization은 위 Figure 1의 uniform quant인가? 아니면 non-uniform quant인가?

A1. 함수는 실수 r을 정수 값 집합으로 매핑하는 방법이다. 이러한 양자화 방식은 결과적으로 생성된 양자화 값(즉, 양자화 수준)이 균일하게 간격을 이루기 때문에 균일 양자화(uniform quantization)로도 알려져 있다(Figure 1, 왼쪽).

그렇다면 non-uniform quantization은? III-F에서 다룰 것.

quantized된 $Q(r)$에서 실수 r을 복원할 수 있는 연산을 dequantization이라고 하고 이는 다음과 같다.

$\tilde{r} = S \cdot (Q(r) + Z)$                $(3)$

Q2. 여기서 복원된 실수 값 $\tilde{r}$은 $r$과 왜 다를까?

A2. quantizing 과정에서의 precision loss로 인해서

C. Symmetric and Asymmetric Quantization

uniform quantization에서 중요한 요소 中 하나는 식(2)에서 사용되는 Scaling factor $S$이다. scaling factor는 주어진 실수 값 r의 범위를 여러 개의 구간으로 나누는 역할을 한다. (즉, 얼마나 잘게 나눌까?의 문제) scaling factor는 다음과 같이 정의된다.

$S=\frac{\beta-\alpha}{2^b-1}$                $(4)$

$[\alpha, \beta]$ : clipping range로 실수 값을 제한하는 유한한 범위를 나타낸다. $b$ : quantization bit width

따라서, scaling factor를 정의하려면, 먼저 clipping range를 결정해야한다. 이 과정을 calibration이라고 한다.

Clipping range 선택

• 기본 접근접 : input의 min,max를 clipping range로 선택하는 것($\alpha = r_{min}, \beta = r_{max}$). 이 접근법은 asymmetric quantization 방식으로 clipping range가 원점에 대해 대칭적이지 않을 수 있다. ($-\alpha ≠ \beta$). 이는 아래 Figure 2(오른쪽)에 설명되어 있다.

• symmetric quantization : symmetric한 clipping range를 사용하여, $\alpha = -\beta$로 설정한다. 일반적인 선택은 input의 min,max를 기반으로 $-\alpha = \beta = max(|r_{max}|,|r_{min}|)$으로 설정하는 것이다. asymmetric quant는 symmetric quant보다 더 좁은 clipping range를 제공하며, 이는 ReLU이후 activation 처럼 항상 양수인 데이터에 유리하다.

  Q3. asymmetric quant의 경우 왜 symmetric quant보다 더 좁은 clipping range를 제공하는가? 또한 ReLU 이후 activation 처럼 항상 양수인 데이터에 왜 유리한가?

  A3. input의 range가 [-2.435, 162.624]이라고 하자. symmetric quant의 clipping range : [-162.624, 162.624] asymmetric quant의 clipping range : [-2.435, 162.624] ==> asymmetric이 더 좁음 ==> 더 좁은 범위를 자르네? ==> 더 잘게 자른다는 거네? ==> quantizing 이후 값이 상대적으로 더 정확하다는 거네? ==> ReLU이후 activation은 결국 양수니까 양수의 범위에서 더 잘게 짜를 수 있는 asymmetric이 더 유리하네!

symmetric quant를 사용할 경우

영점을 식(2)에서 Z를 제거하여 식이 단순화된다.

$Q(r) = Int(\frac{r}{s})$                $(5)$

symmetric quant의 scaling factor 선택

• full range : $S=\frac{2\cdot max(|r|)}{2^n -1}$로 설정하여 INT8의 전체 범위([-128,127])를 사용한다.

• restricted range : $S=\frac{max(|r|)}{2^{n-1}-1}$로 설정하여 범위를 [-127,127]로 제한한다.

  • 정확도 : full range의 경우가 더 정확하다.

symmetric quant의 실제 적용

• weight quant : symmetric quant는 주로 weigt를 quantizing 하는 데 사용된다. 이는 zero-point를 0으로 설정하면 inference 중 computatino cost를 줄일 수 있고, 구현이 간단해지기 때문이다.

• activation quant : asymmetric한 activation으로 인해 발생하는 offset cross terms는 데이터에 독립적인 static term으로 간주되며, bias에 흡수되거나 accumulator를 초기화하는 데 사용할 수 있다.

이렇게 min,max를 symmetric, asymmetric quant 모두에서 사용하는 것은 널리 사용되는 방법이다. 그러나 이 접근법은 activation에서의 outlier에 민감하다. 이러한 outlier는 불필요하게 범위를 증가시켜 quantization resolution(해상도)를 낮춘다.

이를 해결하기 위한 한 가지 방법은 input의 min,max 대신 percentile을 사용하는 것이다. 즉, min,max 대신 i번째로 작은, 큰 값을 $\alpha, \beta$로 사용한다.

또 다른 방법은 , KL divergence (Kullback-Leibler divergence) 즉 실제 값과 quantized된 값 간의 정보 손실을 최소화하도록 $\alpha, \beta$를 선택하는 것이다.

Summary (Symmetric vs Asymmetric Quantization)

symmetric quant는 clipping range를 대칭적으로 분할한다. 이는 구현이 더 쉬워지는 장점이 있으며, 식 (2)에서 Z가 0이 되기 때문이다. 그러나 범위가 비대칭적으로 skewed된 경우에는 최적의 선택이 아니다. 이러한 경우에는 asymmetric quant가 선호된다.

D. Range Calibration Algorithms : Static vs Dynamic Quantization

Clipping range가 언제 결정되는가의 문제.

Weight : inference시 parameter는 대부분 고정되어 있어, static하게 결정됨

Activation : input sample에 따라 activation map이 달라서, activation을 quant하는데 dynamic, static 두가지 접근 방식이 쓰인다.

Dynamic에서는 clipping range를 run-time 동안 각 activation map에 대해 dynamic하게 계산한다. 이는 real-time으로 signal statistics (min, max, percentile, etc.)를 계산해야 하기에 overhead가 높다. 하지만 input 별로 signal 범위를 정확히 계산하므로 더 높은 accuracy를 갖는다.

Static에서는 clipping range를 pre-calculate 하고, inference 동안 고정한다. pre-calculate에 자주 사용되는 방법 중 하나는 일련의 calibration input을 넣어봐서 activation의 일반적인 범위를 계산하는 것이다.

그 계산법에는 다음과 같은 계산법들이 있다.

• unquantized된 weight의 분포와 해당 quantized된 값 간의 MSE(Mean Squared Error)를 최소화하는 방법
• entropy
• NN training 중에 clipping range를 training하거나 impose하는 방법(LQNets, PACT, LSQ, LSQ+, etc.)

Summary (Dynamic vs Static Quantization)

Dynamic quant는 각 activation의 clipping range를 dynamic하게 계산하며, 일반적으로 가장 높은 accuracy를 달성한다. 그러나 signal 범위를 동적으로 계산하는 것은 비용이 매우 높아 실무에서는 대부분 모든 입력에 대해 고정된 clipping range를 사용하는 static quant를 선택한다.

E. Quantization Granularity

CV 작업에서 한 layer의 activation input은 여러 다른 convolution filter와 함께 convolved된다(아래 Figure 3 참고). 이때 filter마다 값의 분포와 범위가 다를 수 있따. 따라서 quantization시 weight clipping range가 계산되는 범위(Granularity)가 중요하다. 이는 주로 Layerwise quant, Channelwise quant로 나뉘며, 각 방식은 quantization resolution, accuracy에 영향을 미친다.

a) Layerwise Quantization

한 layer의 filter 전체 weight를 고려하여 clipping range를 결정(Figure 3을 보면 빨간색 점선, 초록색 점선의 폭이 다름을 볼 수 있음). layer내 모든 parameter 통계(min, max, percentile, etc.)를 분석한 후, 동일한 clipping range를 layer의 모든 conv filter에 적용하는 것. But 각 conv filter 마다 범위가 크게 다를 수 있어, 최적의 accuracy를 달성하지 못함

그니까 쉽게 말해서, Layerwise는 Figure 3의 Filter 1 ~ Filter C 까지의 모든 filter에 대해 paramter 통계를 내고, 그 통계 하나를 가지고, 모든 conv filter에 동일한 clipping range를 적용하는 것.

예를 들어, 상대적으로 좁은 범위의 parameter를 가진 conv kernel(filter)은 동일 layer 내 넓은 범위를 가진 kernel 때문에 quantization resolution이 손실될 수 있다.

b) Groupwise Quantization

layer내 여러 channel을 그룹화하여 clipping range를 계산하는 것(activation이나 conv kernel 모두 가능). 이는 하나의 conv 또는 activation 내에서 parameter 분포가 크게 달라지는 경우에 유용하다. 예를들어, Transformer의 fully-connected attention layer를 quantize하는데 Q-BERT에서 이 방법이 유용함을 발견함. 그러나 당연히 이 방법은 서로 다른 scaling factor를 처리해야한다.

c) Channelwise Quantization

다른 channel과 독립적으로 각 conv filter에 고정된 값을 사용하는 방식. 즉, 각 channel마다 전용 scaling factor를 할당한다.

d) sub-channelwise Quantization

channelwise quant를 더욱 세분화하여 conv 또는 fc layer 내의 모든 parameter 그룹에 대해 clipping range를 결정하는 것.

Summary (Quantization Granularity)

channelwise quantiztaion은 현재 conv kernel을 quantize하는 표준 방법. 이는 각 개별 kernel에 대해 clipping range를 조정할 수 있도록 하면서도 무시할 만한 수준의 overhead만 발생.

F. Non-Uniform Quantization

quantization step, quantization level이 균등하게 분포되지 않는 방식. Non-Uniform Quant의 정의는 다음(식 6)과 같다.

$Q(r) = X_i, if r ∈ [\Delta_i, \Delta_{i+1})$                $(6)$

$X_i$ : discrete quantization level $\Delta_i$ : quantization step (threshold)

실수 r의 값이 quantzation step $\Delta_i$와 $\Delta_{i+1}$ 사이에 위치하면, quantizer Q는 quantization level $X_i$에 매핑된다. 이때, $X_i$와 $\Delta_i$는 모두 non-uniform하게 분포한다.

Non-Uniform Quant의 특징

• 고정된 bit-width에서 더 높은 accuracy를 달성한다. 왜냐하면 중요한 값 영역에 집중하거나 적절한 동적 범위를 찾음으로써 데이터 분포를 더 잘 포착할 수 있기 때문. 예를들어 bell-shaped distribution 또는 긴 꼬리를 가지는 weight, activation에 대해 설계된 non-uniform quant방법이 많이 있다.

• rule-based non-uniform quant는 logarithmic distribution을 사용하는 방법이 일반적이다. 이는 quantzation level과 step이 linera하게 증가하는 대신 exponent하게 증가하게된다.

• binary-code-based quant는 실수 vector $r ∈ R^n$을 binary vector로 quant하는 방식이다. 이는 다음과 같이 표현된다. $r \approx \sum_{i=1}^{m} \alpha_i b_i$ $\alpha_i ∈ R$ : scaling factor $b_i ∈ {-1, +1}^n$ : binary vector

이후 생략...

G. Fine-tuning Methods

Quantization 후에는 NN의 parameter를 조정해야한다. 이는 model을 retraining하는 QAT로 하거나, retraining 없이 수행하는 PTQ 방식으로 진행할 수 있다.

1) Quantization-Aware Training

trained model에 quant를 적용하면 parameter에 섭동이 발생하여 fp precision으로 training된 수렴 지점에서 벗어날 수 있다. 이를 해결하기 위해, quantized된 parameter로 model을 retraining하여 수렴 지점에서 벗어나는 정도를 회복할 수 있다.

이러한 방법들 중 QAT는 fp precision으로 forward pass, backward pass를 하고, 각 gradient update 후에 parameter를 quant하는 방식으로 진행된다. 이 과정에서 weight update는 반드시 fp precision으로 수행한 후에 quant를 적용해야한다. 그렇지 않을 경우 low-precision 환경에서는 0-gradient나 높은 error가 발생할 수 있다.

backpropagation에서 non-differentiable인 $Q(r)$ (식 2)을 어떻게 처리할까? 만약 별다른 approximate 없이 $Q(r)$의 gradient를 계산하면, $Q(r) = Int(\frac{r}{S}) - Z$                $(2)$ 의 Round 연산이(Int()에 있음) 조각별로 평탄한 연산자이기 때문에 거의 모든 구간에서 gradient가 0이 된다. ==> STE(Straight Through Estimator) STE는 round 연산을 무시하고, 이를 항등 함수로 근사하는 방식으로 작동하며, 이를 위 Figure 5에서 보여주고 있다.

STE는 coarse한 approximate임에도 불구하고, 대체로 잘 동작.

model parameter 조정 이외에도, 이전 연구들은 QAT동안 quantization parameter를 학습하는 것이 효과적임을 발견했다. PACT는 uniform quant에서 activation의 clipping range를 training하며, QIT는 non-uniform quant에서 확장하여 quantization level과 step을 학습.

Summary (QAT)

QAT는 STE의 거친 근사에도 불구하고 효과적으로 작동하는 것으로 나타났습니다. 그러나 QAT의 주요 단점은 NN model을 retraining하는 데 드는 높은 계산 비용이다. 특히 low-bit precision quant의 경우 accuracy를 회복하려면 수백 epoch 동안 retraining이 필요할 수 있습니다.

만약 quantized된 model이 오랜 기간 배포될 예정이고 효율성과 accuracy가 특히 중요하다면, 이러한 retraining에 대한 투자는 가치가 있을 가능성이 높습니다. 그러나 모든 경우에 해당하지는 않으며, 일부 model은 상대적으로 짧은 수명을 가질 수 있다. 다음으로, 이러한 오버헤드가 없는 대안 접근법을 논의합니다.

2) Post-Training Quantization

빠름 but accuracy ↓ ==> quant이후 weight의 mean과 variance에서 발생하는 고유한 bias를 보정 + 서로 다른 layer나 channel 간의 weight 범위를 균등화하여 quant error를 줄임

ACIQ는 PTQ를 위해 최적의 clipping range와 channel별 bit width 설정을 분석적으로 계산. But channel별 activation quant는 HW에서 효율적으로 구현하기 어렵 ==> OMSE activation에 대한 channel별 quant를 제거하고, quantized된 tensor와 fp tensor간의 L2 distance 최적화를 통해 PTQ + outlier로 인한 문제점을 완화하기 위해, OCS(Outlier Channel Splitting). OCS는 oulier를 포함하는 channel을 복제하고 절반으로 나누어 처리.

3) Zero-shot Quantization

위에 G. Fine-tuning Methods에서 살펴보았듯이 qauntization후 accuracy 저하를 최소화하기 위해서는 전체 또는 일부 training data를 이용해야한다.

1. 값을 clip하고, 적절한 scaling factor를 찾기위해(calibration)activation의 범위를 알아한다.
2. quantized model은 param을 조정하고 accuracy 저하를 회복하기 위해 fine-tuning이 필요하다.

But training dataset은 배포하기에 너무 크거나, 개인정보의 문제, 보안의 문제등으로 인해, quantization procedure 중에 original training data에 접근하는 것은 불가능하다.

==> ZSQ

ZSQ의 2가지 level

• Level 1 : No data and no finetuning (ZSQ + PTQ)
• Level 2 : No data but requires finetuning (ZSQ + QAT)

ZSQ의 인기있는 연구 분야 중 하나는 pre-trained model이 training된 실제 데이터와 유사한 synthetic data를 생성하는 것이다. Synthetic data는 calibrating 및 finetuning에 쓰인다.

Summary (ZSQ)

ZSQ는 training/validation data에 접근하지 않고, 전체 quantization 과정을 수행한다. 이는 특히 고객의 데이터셋에 접근하지 않고 workload 배포를 가속화하려는 MLaaS(Machine Learning as a Service) 제공없체에게 중요하다. 또한, 보안 또는 프라이버시 문제로 이냏 trianing data 접근이 제한되는 경우에도 중요한 방법이다.

H. Stochastic Quantization

inference 중에는 quantization scheme이 deterministic하게 이루어진다. 그러나 이것이 유일한 가능성은 아니다. 몇몇 연구에서는 QAT 및 reduced precision training을 위해 stochastic quantization을 연구했다.

deterministic한 quant는 작은 weight update가 weight 변화를 일으키지 않을 수 있다. 왜냐하면, rounding 연산은 항상 같은 값을 리턴하기 때문이다.

하지만 stocahstic quant는 NN에게 그런 rounding 연산의 단점을 escape할 수 있는 기회를 주어, parameter를 update할 기회를 제공할 수 있다.

Stochastic quant는 weight update에 따라 fp를 올림하거나 내림하는 확률로 mapping한다. 이게 무슨말이냐는 그냥 아래 식 (8)을 보라.

$ \text{Int}(x) = \begin{cases} \lfloor x \rfloor & \text{with probability } \lceil x \rceil - x, \ \lceil x \rceil & \text{with probability } x - \lfloor x \rfloor. \end{cases} $                $식(8)$

$x$가 3.7인경우 30%의 확률로 3으로 quantized 되고, 70%의 확률로 4로 quantized 된다. ==> 즉 deterministic X, Stochastic O

But binary quant에는 적용될 수 없다.

$$ \text{Binary}(x) = \begin{cases} -1 & \text{with probability } 1 - \sigma(x), \ +1 & \text{with probability } \sigma(x). \end{cases} $$

이것들 말고도 많음...

But 모든 weight update마다 random number를 생성하는 데 드는 overhead가 있어, 실무 채택 X.

4. Advanced Concepts : Quantization Below 8 bits

A. Simulated and Integer-only Quantizaiton

Quantized 된 NN model을 배포하는 데는 2가지 방법이 있다.

• simulated quant(fake quant) : model param이 low-precision으로 저장되지만, 연산은 fp 연산으로 수행. 그래서 dequantization 필요. 아래 Figure 6 중간 그림.

• integer-only quant : 모든 연산이 low-precision의 integer 연산으로 수행. 이른 아래 Figure 6 오른쪽 그림.

low-precision logic은 latency, power consumption, area efficiency 측면에서 full-precision 보다 여러 이점이 있다. 아래 Figure 7 왼쪽을 보면, 많은 HW processor들이 low-precision 산술의 빠른 처리를 위해 inference throughtput과 latency를 향상시켰다. 또한, Figure 7 오른쪽을 보면 45nm 기술을 기준으ㅗㄹ, low-precision logic은 power, energy면에서 훨씬 효율적임을 보여준다. 예를들어 INT8 덧셈은 FP32 덧셈보다 energy cost가 30배, area cost가 116배 뛰어나다.

B. Mixed-Precision Quantization

low-precision quant를 사용할수록 하드웨어 성능이 향상된다는 것은 명백하다. 그러나 모델을 균일하게 ultra low-precision으로 quant하면 accuracy가 크게 저하될 수 있다. 이를 해결하기 위해 mixed-precision quant가 제안되었다.이 방법에서는 각 계층을 서로 다른 bit precision으로 quant하며, 이는 아래 Figure 8에 도식적으로 나타나 있다.

mixed-precision searching problem

어느 layer를 어떤 precision을 할지 찾는 문제이다.

탐색 기반 해결법

1. RL을 활용한 자동 qauntizztion policy 결정 + HW accelerator의 feedback를 RL agent의 feedback에 반영.
2. mixed-precision configuration searching problem을 NAS문제로 공식화하고, Differentiable NAS방법을 사용하여 search space를 효율적으로 탐색.

이런한 탐색 기반 방법의 한 가지 단점은 대규모 계산 자원이 필요하며, hyperparam과 initializaion에 민감하다는 것이다.

주기적 함수 정규화 해결법

periodic function regularization을 사용하여 mixed-precision model을 training 하는 것. 이 방법은 자동으로 서로 다른 layer과 accuracy에 대한 중요도의 차이를 구별하며, 각 층의 bit width를 학습하는 방식으로 동작한다.

C. Hardware Aware Quantization

Quantization의 목표 중 하나는 Inference Latency를 향상시키는 것. 그러나 모든 하드웨어가 특정 layer 또는 연산이 quantized된 후 동일한 속도 향상을 제공하지는 않는다. 실제로 quantization의 이점은 하드웨어에 따라 다르며, on-chip memory, bandwidth, cache hierarchy 구조와 같은 다양한 요인이 quantization에 의한 속도 향상에 영향을 미친다. ==> 하드웨어 친화적 quantizaitno을 통해 최적의 이점을 얻는 것이 중요. 특히, RL agent를 사용하여, 각 layer의 bit width에 따른 Latency 정보를 포함한 look-up table을 기반으로 하드웨어 친화적 mixed-precision configuration을 결정 그러나 이 접근 방식은 하드웨어 latency를 시뮬레이션하여 계산.

이를 해결하기 위해 최근 연구에서는 quantized된 연산을 실제 하드웨어에 직접 배포하여, 각 layer에 대해 다양한 quantization bit precision에서의 실제 배포 latency를 측정하는 방법을 제안.

D. Distillation-Assisted Quantization

model distillation은 더 높은 accuracy를 가진 대형 모델(Teacher)을 활용하여 소형 모델(Student)의 학습을 돕는 방법.

학생 모델을 학습하는 동안, 단순히 실제 클래스 레이블(Ground-Truth Class Labels)만 사용하는 대신, distillation은 Teacher가 생성한 Soft Probabilities를 활용하는 것을 제안. 이 Soft Probabilities는 입력 데이터에 대한 추가 정보를 포함할 수 있다.

즉, 전체 loss function은 student loss와 distillation loss를 결합하여 다음과 같이 공식화 한다.

$L = \alpha H(y,σ(z_s)) + \beta H(σ(z_t,T),σ(z_s,T))$

• $\alpha, \beta$ : 학생과 distillation loss간의 가중치를 조정하는 계수
• $y$ : 실제 클래스 레이블
• $H$ : cross entropy loss function
• $z_s, z_t$ : 학생, 선생 모델이 생성한 logits
• $T$ : softmax temperature $$ p_i = \frac{\exp\left(\frac{z_i}{T}\right)}{\sum_j \exp\left(\frac{z_j}{T}\right)} $$

E. Extreme Quantization

F. Vector Quantization

5. Quantization And Hardware Processors

Quantization는 모델 크기를 줄이는 것뿐만 아니라 속도를 높이고 전력 소비를 줄이는 데도 유용하다. 특히 저정밀 로직을 지원하는 하드웨어에서는 그 효과가 두드러진다. 그러므로 IoT와 모바일 애플리케이션에서 엣지 배포(edge deployment) 시 중요한 역할을 한다.

엣지 디바이스는 보통 계산 능력, 메모리, 그리고 전력 예산(power budget) 같은 자원이 매우 제한적이다. 이러한 제약으로 인해 많은 딥 뉴럴 네트워크(Deep Neural Network, DNN) 모델을 실행하는 데 높은 비용이 들거나 실행이 어려울 수 있다.

게다가 많은 엣지 프로세서는 부동소수점 연산(Floating Point Operations)을 아예 지원하지 않으며, 특히 마이크로컨트롤러(micro-controllers)에서는 이러한 연산이 불가능한 경우도 많다. 이러한 환경에서는 양자화가 필수적인 기술이 된다.

양자화는 다양한 하드웨어 플랫폼에서 NN 모델을 최적화하고 배포하는 데 중요한 역할을 한다. ARM Cortex-M은 저비용, 저전력 임베디드 디바이스를 위해 설계된 32비트 RISC ARM 프로세서 코어군이다. 예를 들어, ARM Cortex-M 코어를 기반으로 한 STM32 마이크로컨트롤러는 엣지 NN 추론에 널리 사용된다. 일부 ARM Cortex-M 코어는 전용 부동소수점 연산 장치를 포함하지 않으므로, 모델을 배포하기 전에 반드시 양자화를 진행해야 한다. ARM의 CMSIS-NN 라이브러리는 이러한 코어에 NN 모델을 양자화하고 배포하는 과정을 돕는 도구로, 비트 쉬프팅(bit-shifting) 연산을 활용한 고효율 고정소수점 양자화와 복원 과정을 지원한다. GAP-8은 엣지 추론을 위한 CNN 가속기를 내장한 RISC-V SoC로, 정수 연산만 지원하는 또 다른 사례이다.

한편, Google Edge TPU와 같은 목적 기반 프로세서는 소형 저전력 디바이스를 위해 설계되었으며 8비트 연산만 지원한다. 따라서 TensorFlow 모델을 사용할 경우, 양자화 인식 학습(quantization-aware training) 또는 학습 후 양자화(post-training quantization)가 필요하다. 최근 몇 년간 엣지 프로세서의 계산 성능이 크게 향상되면서, 이전에는 서버에서만 실행 가능했던 복잡한 NN 모델도 엣지에서 배포하고 추론할 수 있게 되었다. Figure 9는 상용 엣지 프로세서들의 처리량 변화를 보여주며, 이는 양자화, 저정밀 로직, 그리고 전용 딥러닝 가속기의 발전이 크게 기여했음을 나타낸다.

양자화는 엣지 프로세서뿐만 아니라 비 엣지 하드웨어에서도 유용하며, 예를 들어 99번째 백분위 지연(latency)과 같은 서비스 수준 협약(SLA) 요구사항을 충족하는 데 큰 도움을 줄 수 있다. NVIDIA Turing GPU, 특히 T4 GPU는 이를 잘 보여주는 사례이다. Turing Tensor Cores는 저정밀 행렬 곱셈을 효율적으로 수행하도록 설계된 전용 실행 유닛으로, 현대 NN 모델 최적화에 있어 양자화의 폭넓은 활용 가능성을 보여준다.

6. Future Directions For Research In Quantization

생략

7. Summary And Conclusion

생략

SmoothQuant에 대한 내용

Abstract

LLM inference Quantization for accuracy + hardware efficiency ==> SmoothQuant

SmoothQuant의 3가지 목표

1. training-free
2. accuracy-preserving
3. general-purpose PTQ

위 목표는 8-bit weight, 8-bit activation (W8A8) quantization으로 이룰 수 있다.

SmoothQuant는 weight quant가 activation quant 보다 쉽다는 점에 기반하여, activation의 outlier(극단값)을 부드럽게 완화하여 quantization의 어려움을 activation에서 weight로 이전하는 수학적으로 동등한 변환을 수행

Q1. 왜 weight quantization이 activation quantization 보다 쉬운가?

A1. Weight Quantization은 학습된 파라미터(Weight)에만 적용되므로, 추가적으로 입력 데이터에 따른 동적인 변화에 신경 쓰지 않아도 된다. 반면 Activation Quantization은 입력 데이터와 모델의 동작에 따라 결과값이 달라질 수 있으므로, 훈련 데이터 분포를 잘 반영하지 못하면 성능 저하가 심각할 수 있기 때문에. + Activation의 경우 outlier가 있음 weight의 분포는 비교적 uniform하고 flat하여, quantizing이 용이 반면 activation의 distribution은 outlier로 인해 ununiform하고, unflat해서, quantization 어려움

==> 여러 LLM model에서의 모든 matrix multiplication에 대해 weight, activation을 INT8로 quantization

1. Introduction

Q2. GPT-3 model은 1750억 개의 parameters를 가진다. 이때 parameter는 무엇인가? 또한 만약 FP16이면 저장 및 실행하려면 최소 몇 Byte의 메모리가 필요한가?

A2. parameter는 matrix, bias에 들어있는 하나의 값의 총 개수를 의미한다. 16bit * 1750억 = 2조8000억 bit = 350GB

LLM은 크기가 크다 ==> computation 비쌈 ==> quantization으로 LLM cost 줄이자

이때, weight, activation을 INT8로 Quantized하자. ==> GPU 메모리 사용량 절반 + matrix multiplication throughput 2배

SmoothQuant

LLM에 대해 정확하고, 효율적인 PTQ 솔루션임.

activation은 outlier로 인해 weight보다 quantization이 훨씬 어렵지만, 서로 다른 token은 channel간 유사한 변동성을 나타낸다. ==> SmoothQuant는 quantization의 어려움을 activation에서 weight로 이전하는 offline migration을 한다.

즉, SmoothQuant는 channel별로 수학적으로 동등한 스케일 변환을 제안하여 channel 간 크기를 크게 평탄화(smooth)하여 model을 quantization에 적합하게 한다.

2. Preliminaries

Integer Uniform Quantization

$\overline{X}^{INT8} = ⌊\frac{X^{FP16}}{Δ}⌉ , Δ = \frac{max(|X|)}{2^{N-1}-1}$ (1)

• $X$는 FP tensor
• $\overline{X}$는 quantization된 tensor
• $Δ$는 quantization step size (scaling factor)
• $⌊⋅⌉$는 반올림 연산
• $N$은 bit 수 (이 경우 8)

만약, tensor가 asymmetric이라면(ex) ReLU 이후) zero-point를 추가하여 처리할 수 있다.

이러한 quantization 방식은 activation의 outlier를 보존하기 위해 $max(|X|)$를 사용한다. (for accuracy 유지)

또한 $Δ$는 activation의 일부 calibraion sample을 사용하여 offline으로 계산할 수 있다. 이를 static quantization이라고 한다.

quantization granularity levels

• per-tensor quantization 위 그림의 (a)이다. 이는 전체 matrix에 대해 single step size를 사용한다.
• per-token + per-channel quantization 위 그림의 (b)이다. 이는 각 토큰에 연관된 activation에 대해 서로 다른 quantization step size를 사용하는 per-token quantization이나, weight의 각 output channel에 대해 다른 quantization step size를 사용하는 per-channel quantization를 적용할 수 있다.
• group-wise quantization 이는 channel의 group 별로 서로 다른 quantization step size를 사용하는 방법이다.

3. Review of Quantization Difficulty

LLM에서 activation의 quantization이 어려운 이유 = outlier로 인해

1. Activations are harder to quantize than weights.

weight의 분포는 비교적 uniform하고 flat하여, quantizing이 용이 반면 activation의 distribution은 outlier로 인해 ununiform하고, unflat해서, quantization 어려움

2. Outliers make activation quantization difficult.

activation에서 나타나는 outlier의 크기는 대부분의 activation 보다 약 100배 이상 크다. ==> per-tensor quantization에서 얼마 안되는 outlier 값들로 인해

$\overline{X}^{INT8} = ⌊\frac{X^{FP16}}{Δ}⌉ , Δ = \frac{max(|X|)}{2^{N-1}-1}$ 에서 $max(|X|)$의 값이 너무 의미없이 커지게 된다.

즉, 얼마 안되는 outlier가 $Δ$를 지배하게된다. ==> non-outlier channel들이 무의미해 진다.

3. Outliers persist in fixed channels.

1. 한 channel에서 outlier가 발생하면, 특정 token에서만 outlier가 발생하는 것이 아닌, 모든 token에서 지속적으로 나타난다. (figure 4 참고)

2. channel 간 분산은 크고 (즉, 특정 token에 대해 일부 channel은 activation이 매우 크고, 나머지 channel은 작음. 즉, channel들 간 activation 크기 차이가 큼), channel 내 분산은 작다 (즉, 특정 channel에서의 activation은 여러 token 간 거의 일정하게 유지됨).

위 2가지 이유(outlier의 지속성, 각 channel 內의 작은 분산)로 인해, activation의 per-channel quantization을 수행할 수 있다면, per-tensor quantization과 비교했을 때 quantization error를 훨씬 줄일 수 있다. 이는 아래 Table 1에서 확인할 수 있다.

But,

per-channel quant는 hardware-accelerated GEMM kernel과 잘 맞지 않는다. 이 kernel은 높은 throughput을 요구하는 연산 sequence에 의존하며, 낮은 throughput을 가지는 instruction을 해당 sequence에 삽입하는 것을 허용하지 않는다.

==

==>

이전 연구들은 모두 linear layer에서의 per-token activation quantization을 사용하였으나, 이는 activation quantization의 어려움을 완전히 해결하지는 못함

4. SmoothQuant

위의 이유로인해 per-channel activation quantization 대신 acrivation을 smoothing factor $s ∈ R^{C_i}$로 나누어 smooth하고자 한다.

$Y = (X \cdot diag(s)^{-1}) \cdot (diag(s) \cdot W) =\hat{X} \cdot \hat{W}$ $(식 3)$

입력 $X$는 일반적으로 이전 linear operation으로 생성되므로, $s$는 이전 layer의 parameter에 offline으로 쉽게 병합할 수 있다.

==> 추가 scaling으로 인한 kernel overhead X

Migrate the quantization difficulty from activations to weights.

이 논문의 목표는 quantization의 어려움을 activation에서 weight로 이전 하는 것이다.

이를 위해 $s$를 선택하여 $\hat{X} = X \cdot diag(s)^{-1}$으로 quantization하기 쉬운 형태가 되도록 하고자 한다.

quantization error를 줄이기 위해, 모든 channel에서 effective quantization bits를 증가시켜야 하는데, effective quantization bits의 총합은 모든 channel이 동일한 최대 크기를 가질 때 가장 커진다.

// effective quantization bits : quantization 시 표현 가능한 데이터의 정확도.

==> 걍 outlier를 줄이려는 거임

==> $s_j = max(|X_j|), j=1,2,...,C_i$ 예를들어, 의 $s$는 아래와 같다.

activation의 범위는 dynamic하다. 입력 sample에 따라 달라질 수 있다. 이를 해결하기 위해, pre-training dataset에서 calivration sample을 사용하여 activation channel의 scale을 추정한다.

But

모든 quantization 어려움이 weight로 이전된다. ==> weight quantization error ↑

∴ 반대로 $s_j = 1/max(|W_j|)$를 선택하면, quantization 어려움을 activation으로 이전한다.

==> 적절하게 합쳐야한다

==> $\LARGE s_j = \frac{\max(|X_j|)^\alpha}{\max(|W_j|)^{1-\alpha}}$ $식(5)$

대부분의 model에서 $α = 0.5$일 때, quantization 어려움을 균등하게 분할함.

==> 해당 channel에서 weight와 activation이 유사한 최대 값을 가지도록 보장하여, 비슷한 quantization 어려움을 공유하게 한다.

// activation outlier가 더 두드러지는 model의 경우 $α = 0.75$가 더 낫다.

위 Figure 5는 $α = 0.5$일 때의 예시이다.

이때, $\hat{X} = X \cdot diag(s)^{-1}$은 다음과 같다.

Applying SmoothQuant to Transformer block

LLM model에서 linear layer는 대부분의 parameter와 computation을 차지한다.

==> self-attention 및 feed-forward layer의 input activation에 대해 scale smoothing을 수행하고, 모든 linear layer를 W8A8로 quantization한다. 또한 attention 계산에서 BMM operator도 quantization 한다.

5. Experiments

생략

6. Related Work

생략

Conclusion

생략