기술 요약: FTerViT: 완전 3 진수 비전 트랜스포머

문제 정의

비전 트랜스포머 (ViT) 는 강력한 이미지 분류기이지만, 상당한 메모리 발자국으로 인해 마이크로컨트롤러급 장치에는 일반적으로 적합하지 않습니다. DeiT-Small 과 같은 표준 컴팩트 ViT 는 FP32 정밀도에서 약 88.3 MB 를 필요로 하는 반면, 일반적인 마이크로컨트롤러 유닛 (MCU) 은 외부 RAM 을 몇 메가바이트만 제공합니다. 이러한 불일치는 클라우드 오프로딩 없이 엣지 장치에서 항상 켜져 있는 저전력 비전 애플리케이션의 배포를 방해합니다.

3 진수 양자화 (가중치를 $\{-1, 0, +1\}$로 제한) 는 FP32 대비 이론적으로 16 배의 압축을 제공하지만, 기존 최첨단 방법은 트랜스포머 인코더 레이어만 3 진수화합니다. 패치 임베딩, LayerNorm 파라미터, 분류기 헤드는 FP32 또는 낮은 비트 정수 (INT8) 의 전체 정밀도로 유지됩니다. 컴팩트 모델에서 이러한 나머지 전체 정밀도 구성 요소는 주요 병목 현상이 되어, 파라미터의 작은 비율을 차지함에도 불구하고 모델 크기의 최대 38% 를 소비합니다. 또한, 이러한 민감한 구성 요소를 완전히 3 진수화하려는 이전 시도 (예: TerViT) 는 첫 번째 및 마지막 레이어의 양자화에 대한 알려진 민감성으로 인해 상당한 정확도 하락 (최대 22.4 퍼센트 포인트) 을 초래했습니다.

방법론

저자들은 모든 가중치 행렬과 정규화 파라미터가 $\{-1, 0, +1\}$로 제한된 최초의 완전 3 진수 비전 트랜스포머인 FTerViT를 제안합니다. 이 접근 방식은 두 가지 주요 기둥으로 구성됩니다: 새로운 3 진수 연산자와 특수화된 학습 프로토콜.

1. 3 진수 원시 연산자

FP32 병목 현상을 제거하기 위해 저자들은 표준 구성 요소를 전용 3 진수 모듈로 대체합니다:

• TernaryBitLinear: 완전 연결 레이어에 사용됩니다. 이는 텐서별 절대 평균 ($absmean$) 스케일링 및 클램핑 반올림을 사용하여 가중치를${-1, 0, +1}$로 양자화합니다. 활성화 값은 토큰별 절대 최대 스케일링을 사용하여 8 비트 정수로 양자화됩니다.
• TernaryBitConv2d: 패치 임베딩 레이어를 위해 특별히 설계되었습니다. 이 방법은 이질적인 필터 크기를 처리하기 위해 채널별 스케일링을 적용하며, 이전 연구에서 가장 중요한 구성 요소로 식별된 패치 임베딩 레이어의 민감성을 해결합니다.
• TernaryLayerNorm: 동일한 $absmean$ 방식을 사용하여 LayerNorm 의 학습 가능한 아핀 파라미터 ($\gamma$ 및 $\beta$) 를 3 진수화합니다. 토큰별 통계 (평균 및 분산) 와 편향은 파라미터 수가 무시할 수 있을 정도로 적기 때문에 FP32 로 유지되지만, 이를 3 진수화하면 위치 정보가 파괴됩니다.

2. 학습 프로토콜

모델은 2 단계 지식 증류 (KD) 전략을 사용하여 학습됩니다:

• 1 단계 (학습): 3 진수 학생 모델은 사전 학습된 FP32 교사 모델에서 초기화되며, 학습률 $10^{-4}$ (코사인 감쇠) 로 250 에포크 동안 양자화 인식 증류 (QAD) 를 사용하여 학습됩니다. 손실 함수는 저비트 네트워크에서 KL 과 충돌할 수 있는 교차 엔트로피를 피하고, 학생과 교사 로짓 간의 순방향 KL 발산 항만 사용합니다.
• 2 단계 (복구): 학습률을 $10^{-5}$ (코사인 감쇠) 로 재시작하여 10 에포크 동안 수행하는 경량 미세 조정 단계입니다. 이 단계는 1 단계의 초기 포화 동안 손실된 스케일 재보정 및 정확도 회복을 가능하게 합니다.

주요 기여

1. 최초의 완전 3 진수 ViT: 저자들은 지식 증류로 학습될 때 ViT 의 가장 취약한 구성 요소 (패치 임베딩, LayerNorm, 분류기 헤드) 가 성공적으로 $\{-1, 0, +1\}$로 3 진수화될 수 있음을 입증하여, 이전에 양자화 가능하다고 여겨지던 것의 한계를 확장했습니다.
2. 새로운 연산자: 비인코더 레이어의 특정 민감성을 처리하기 위해 채널별 스케일링을 갖춘 TernaryBitConv2d 와 TernaryLayerNorm을 도입했습니다.
3. 학습 전략: 이전 3 진수 ViT 접근 방식보다 더 빠른 수렴과 훨씬 높은 최종 정확도를 가능하게 하는 간단하지만 효과적인 2 단계 학습 전략 (고학습률 학습 followed by 저학습률 복구) 입니다.
4. 장치 내 구현: 리소스 제약이 있는 듀얼 코어 ESP32-S3 마이크로컨트롤러 (8 MB PSRAM) 에서의 최초의 독립형 3 진수 비전 트랜스포머 C 구현입니다.

실험 결과

이 방법은 ImageNet-1K 및 기타 벤치마크에서 평가되었습니다:

• 정확도 및 압축:

  • FTerViT-DeiT-III-S (384$\times$384): FP32 교사 모델 대비 2.42 퍼센트 포인트만 손실한 6.09 MB에서 **82.43%**의 top-1 정확도를 달성하여 14.6 배 압축을 실현했습니다. 이는 TerViT 와 같은 이전 3 진수 ViT 방법보다 최대 8 퍼센트 포인트 더 높은 성능입니다.
  • FTerViT-DeiT-III-S (224$\times$224): 5.81 MB에서 **79.64%**의 top-1 정확도를 달성했습니다.
  • FTerViT-DeiT-S (224$\times$224): 5.81 MB에서 **77.47%**의 top-1 정확도를 달성하여 TerViT-DeiT-S 를 3.27 퍼센트 포인트 상회했습니다.
  • CIFAR 벤치마크: DeiT-Tiny 변형은 1.53 MB 의 저장 공간으로 CIFAR-10 에서 97.43%, CIFAR-100 에서 86.01% 를 달성했습니다.

• 구성 요소 민감성 분석:

  • 테일러 1 차 및 헤시안-trace 분석은 LayerNorm 과 패치 임베딩이 파라미터의 0.2% 미만을 차지함에도 불구하고 모델 중요도 (각각 34~39% 및 4.6%) 를 지배함을 확인했습니다. 이는 이전 방법들이 이를 더 높은 정밀도로 남겼던 이유와 그 성공적인 3 진수화가 중요한 이유를 설명합니다.

• 장치 내 배포 (ESP32-S3):

  • 5.81 MB 모델은 ESP32-S3 의 8 MB PSRAM 에 완전히 들어맞아 카메라 및 LCD 버퍼를 위해 약 4.5 MB 를 남깁니다.
  • 성능: 순방향 전달은 21.06 초가 소요됩니다.
  • 효율성: FP32 대비 패킹된 2 비트 3 진수 추론은 1.71 배 더 빠릅니다 (특정 레이어의 경우 804 ms 대 1376 ms) 그리고 유휴 상태 이상의 에너지 소비를 55% 감소시킵니다 (59 mJ 대 130 mJ).

의의

이 논문은 FTerViT 가 표준 ViT 정확도와 마이크로컨트롤러 규모 배포의 엄격한 메모리 제약 사이의 격차를 크게 좁힌다고 주장합니다. 이전에는 "손대기 힘든" 패치 임베딩 및 정규화 레이어를 포함한 전체 아키텍처를 성공적으로 3 진수화함으로써, 저자들은 저비용 ($10) 상용 MCU 에서 고품질 비전 모델의 배포를 가능하게 합니다. 이 작업은 올바른 학습 전략 (동일 아키텍처 KD 및 2 단계 복구) 을 사용하면 이러한 구성 요소의 민감성을 관리할 수 있으며, 완전 3 진수 ViT 가 엣지 AI 를 위한 실용적인 현실이 될 수 있음을 검증합니다. 저자들은 FP32 대비 정확도 격차가 모델 용량에 반비례한다고 언급하며, 더 큰 모델 (DeiT-III-S384) 이 더 작은 모델보다 상대적으로 적은 정확도 손실을 겪는다고 설명합니다.