SensiMix: Sensitivity-Aware 8-bit Index & 1-bit Value Mixed Precision Quantization for BERT Compression

Posted by 날짜:

 본 문서에서는 2022년 PLOS ONE 저널에 발표된 "SensiMix: Sensitivity-Aware 8-bit Index & 1-bit Value Mixed Precision Quantization for BERT Compression" 논문을 소개합니다. 논문에 대한 상세한 정보는 다음과 같습니다.

  • Title: SensiMix: Sensitivity-Aware 8-bit Index & 1-bit Value Mixed Precision Quantization for BERT Compression
  • Authors: Tairen Piao, Ikhyun Cho, and U Kang
  • Journal: PLOS ONE (2022)

BERT Compression

어떻게 하면 기학습된 BERT 모델을 효과적으로 압축할 수 있을까요? BERT, RoBERTa, 그리고 ERNIE와 같이 기학습된 언어 모델(pre-trained language model)들은 문법 적합성 판정, 문장 감정 분석, 문장 유사도 분석 등 다양한 자연어처리(Natural Language Processing, NLP) 분야에서 모델의 성능을 크게 향상시켰고, 그 중심에는 BERT가 있습니다. 하지만, 이러한 기학습된 언어 모델들은 크기가 매우 크고 연산량(FLOPs)이 매우 많다는 단점이 있기 때문에, 이들을 압축하여 크기를 줄이고, 연산량을 줄이는 연구들이 활발하게 이루어지고 있습니다. 본 문서에서 소개하는 SensiMix 또한 대표적인 압축 기법 중 하나인 양자화(quantization) 기법을 바탕으로 기학습 언어 모델인 BERT를 정확도 손실을 최소화하며 효율적으로 압축하는 것을 목표로 합니다.

SensiMix (Overview)

SensiMix는 양자화 기법을 이용하여 작고, 빠르면서도 정확한 BERT 압축 모델을 만드는 것을 목표로 합니다. 이러한 목표를 달성하기 위해 본 논문에서 제안하는 주요 아이디어들은 아래와 같습니다.
  • Sensivity-aware mixed precision quantization
    • 양자화로 인한 정보손실을 최소화하기 위하여 모델 내의 민감한 부분과 상대적으로 덜 민감한 부분을 구분하여 양자화 정도를 선택합니다.
  • Training methods for 1-bit quantization
    • 1비트 양자화로 인한 정확도 손실을 막기 위해 학습 과정에서 ABWR, PT, ILF등의 새로운 학습 기법들을 도입하였습니다.
  • Fast matrix multiplication for inference
    • 양자화 후의 추론 속도 향상을 극대화하기 위하여 FP16(half precision) GEMM(general matrix multiplication), XNOR-Count GEMM등을 사용합니다.
지금부터는 위 아이디어들에 대해 자세하게 설명해드리도록 하겠습니다.


Sensivity-aware mixed precision quantization

양자화란 32-bit로 표현되어 있는 기학습된 모델의 파라미터를 16-bit, 8-bit, 혹은 1-bit의 낮은 정밀도(precision)의 숫자들로 표현하여 모델의 크기를 줄이고 연산 속도를 향상시키는 기법을 의미합니다. 이때, 낮은 정밀도를 사용할수록 모델의 크기는 작게 만들 수 있지만, 이에 따라 모델의 성능이 낮아진다는 단점이 있기 때문에 어떤 정밀도를 사용할지는 신중하게 선택해야 합니다. 본 논문에서는 다음과 같은 두 가지 발견을 바탕으로 파라미터들의 정밀도를 선택하였습니다. (1) 모듈별로 정밀도에 대한 민감도(sensivity)가 다르며, BERT 내부에서는 self-attention 모듈이 이후에 이어지는 FFN(feedforward network) 모듈들에 비해 정밀도에 더욱 민감합니다. (2) BERT에서는 입력층에 가까운 레이어일수록 정밀도에 민감합니다. 그러므로, SensiMix에서는 self-attention 모듈에는 모두 8-bit 인덱스 양자화를 사용하였고, 입력층에 가까운 층들의 FFN에 대해서는 8-bit 인덱스 양자화를, 출력층에 가까운 층들의 FFN에 대해서는 1-bit 양자화를 사용하였습니다.  또한, 단어 임배딩(embedding)들의 경우에도 정밀도에 민감하기 때문에 8-bit 인덱스 양자화를 사용하였습니다.
 
그림 1. BERT(왼쪽)와 SensiMix(오른쪽)의 구조

그림 1은 기학습 모델인 BERT(왼쪽)와 양자화를 진행한 SensiMix(오른쪽)를 나타내고 있습니다. 그림을 보면 BERT에서는 모든 모듈들이 32-bit로 나타나있지만, SensiMix에서는 모듈에 따라 8-bit 및 1-bit 양자화가 되어 있는 것을 확인할 수 있습니다. 본 논문에서는 self-attention과 FFN에 모두 8-bit 인덱스 양자화 적용한 층을 8-bit encoder라고 명명하고, self-attention에는 8-bit를 , FFN에는 1-bit 양자화를 적용한 층을 MP encoder(mixed precision encoder)라 명명하였습니다. 앞서 설명한 바와 같이 입력층과 가까운 저층부에는 8-bit encoder들이 배치되어 있고, 출력과 가까운 고층부에는 MP encoder들이 배치되어 있는 것을 확인할 수 있습니다. 또한, SensiMix에서는 기존 BERT 모델의 12개의 층을 모두 사용하지 않고, 일부만 사용하며, 전체 층의 개수 및 MP encoder 층의 개수를 조절함으로써 압축 모델의 크기를 조절할 수 있습니다.

Training methods for 1-bit quantization

양자화를 진행할 때, 낮은 정밀도를 사용할수록 모델의 크기를 효과적으로 줄일 수 있지만, 모델의 성능이 낮아지는 특징이 있습니다. 그러므로 1-bit 양자화를 진행할 경우 모델의 크기를 기존의 1/32까지 줄일 수 있지만, 이에 따라 모델의 성능이 줄어들 수 있기 때문에 이를 방지해주는 것이 매우 중요합니다. SensiMix 또한 MP encoder에서 FFN 모듈에 1-bit 양자화를 사용하기 때문에 Absolute Binary Weight Regularization(ABWR), Prioritized Training (PT), 그리고 Inverse Layer-wise Fine-tuning(ILF)라는 세 가지 새로운 기법을 도입하여 성능 하락을 방지하였습니다.

Absolute Binary Weight Regularization(ABWR)

1-bit 양자화를 진행하게 되면 모델의 파라미터들은 -1 또는 1의 값만을 가지게 됩니다. 하지만, 양자화를 진행하기 전의 학습된 값들은 -1 및 1이 아닌 다양한 값들을 가지고 있으며, 이 때문에 양자화를 진행하게 되면 성능 하락이 발생하게 됩니다. 그러므로 SensiMix에서는 기존의 파라미터들이 -1 또는 1의 값에 가까운 값을 가지게 해줄 수 있는 기법인 ABWR 기법을 제안하였습니다. 해당 기법은 원본 파라미터 값의 절댓값이 1과 가깝도록 만들어주는 항을 목적함수에 추가하는 방식으로 동작합니다.
그림 2. ABWR을 적용하기 전(왼쪽)과 후(오른쪽)의 파라미터 분포

그림 2는 ABWR을 적용하기 전(왼쪽)과 후(오른쪽)의 파라미터 분포를 나타내고 있습니다. ABWR을 적용하기 전에는 파라미터들이 0을 중심으로 분포하며, 1 및 -1과의 거리가 상당하여 1-bit 양자화를 진행했을 경우 성능 하락이 심할 것으로 보입니다. 하지만, ABWR을 적용한 오른쪽 그림에서는 원래에 비해 상당히 많은 파라미터들이 1 또는 -1에 가깝게 분포하게 되어 1-bit 양자화로 인한 성능 하락이 상당히 줄어들 것으로 예상할 수 있습니다.

Prioritized Training (PT)

통상적으로 1-bit 양자화된 모델을 학습할 때, 입력 및 파라미터를 모두 학습 시작때부터 양자화한 후에 학습을 진행하게 됩니다. 하지만, 이렇게 학습을 진행할 경우 모델의 파라미터들은 1-bit 양자화된 입력에 맞추어 학습되기도 전에 자신들 또한 양자화가 되기 때문에 성능 하락이 심해질 수 있습니다. 그러므로, 본 논문에서는 학습 초기 단계에서 입력에 대해서만 1-bit 양자화를 진행하고, 파라미터들이 이에 맞추어 어느 정도 학습이 이루어진 후에 파라미터들을 1-bit 양자화하도록 진행하여 기존보다 양자화된 모델의 성능을 향상시켰습니다.

Inverse Layer-wise Fine-tuning(ILF)

SensiMix에는 여러층의 MP encoder들이 있고, 이들은 각각 1-bit로 양자화된 FFN 모듈을 포함하고 있습니다. 만약 모델 내의 모든 MP encoder들의 FFN 모듈들을 한 번에 1-bit로 양자화하게 된다면, 이에 따른 성능 하락이 커질 수 있습니다. 이를 방지하기 위하여 본 논문에서는 상층부에 있는 MP encoder부터 한 에폭(epoch)에 한 층씩 1-bit 양자화를 진행합니다. 예를 들어 3개의 MP encoder를 가지는 SensiMix 모델의 경우 첫 에폭에서 MP encoder들 중 가장 위 층을 1-bit 양자화하고, 이후에 그 아래 층을, 그리고 3번 째 에폭에서 모든 MP encoder들을 1-bit 양자화하게 됩니다. 이렇게 층별로 순차적으로 양자화를 진행함으로써 아래층들이 이에 맞게 학습될 수 있도록 하여 양자화된 모델의 성능을 향상시켰습니다. 그림 3은 ILF의 진행 과정을 묘사하고 있습니다. 초기에 가장 위 층의 MP encoder만이 1-bit 양자화가 되어있지만, 한 에폭이 지난 후에 그 아래 층까지도 1-bit 양자화가 된 것을 알 수 있습니다.
그림 3. ILF의 진행 과정 예시


Fast matrix multiplications for inference

본 논문에서는 양자화 및 학습이 완료된 SensiMix 모델의 추론 속도를 향상시키기 위해 다음과 같은 방식을 적용하였습니다.
  • 8-bit 인덱스 양자화의 경우 파라미터들은 8-bit로 저장되지만 연산을 진행할 때에는 FP32(full precision)를 사용하게 됩니다. SensiMix에서는 이를 FP16(half precision) GEMM으로 연산하도록하여 빠른 추론을 가능하게 하였습니다.
  • 1-bit 양자화가 진행된 값들의 경우 GEMM을 XNOR 및 count 연산을 이용하여 빠르게 수행할 수 있습니다. 이를 바탕으로 XNOR-Count GEMM을 구현하여 추론 속도를 향상시켰습니다.

Experiment

본 실험에서는 GLUE (General Language Understanding Evaluation) benchmark를 이용하여 SensiMix와 기존의 다른 압축 기법들의 성능을 비교하였으며 압축 대상 모델로는 BERT-base를 사용하였습니다. 실험 결과는 아래 표와 같습니다.

표 1. SensiMix와 경쟁 기법들의 BERT 압축 결과 비교

표 1을 보면, BERT-base를 단순한 pruning 기법만을 사용하여 압축하였을 때, 모델의 크기가 기존의 30%만 되더라도 평균 성능이 2.9 가량 크게 하락하며, DistilBERT 기법의 경우 모델을 크게 압축하지 않았음에도 1 이상의 성능 하락을 보인 것을 확인할 수 있습니다. ALBERT-large 모델의 경우에는 16% 크기의 작은 모델로 높은 성능을 보였지만, weight sharing을 기반으로 한 기법이기 때문에 기학습된 모델보다 훨씬 긴 추론 시간을 보이는 것을 알 수 있습니다. SensiMix 모델들의 경우 3개의 8-bit encoder들과 6개의 MP encoder들로 이루어진 SensiMix (3+6) 모델과 3개의 8-bit encoder들과 3개의 MP encoder들로 이루어진 SensiMix (3+3) 모델의 실험 결과가 소개되어있습니다. SensiMix (3+6)의 경우 기학습된 BERT-base 모델의 17%밖에 되지 않는 크기를 가짐에도 기존 모델과 거의 비슷한 성능을 보였고, 추론 시간 또한 효과적으로 줄어든 것을 확인할 수 있었으며, SensiMix (3+3) 또한 가장 작은 모델임에도 불구하고 높은 성능을 보여주고 있습니다. 이를 통해 SensiMix를 이용할 경우 기존의 압축 기술들보다 효과적으로 BERT 모델을 압축할 수 있음을 알 수 있습니다.

Conclusion

본 문서에서는 2022년 PLOS ONE 저널에 발표된 "SensiMix: Sensitivity-Aware 8-bit Index & 1-bit Value Mixed Precision Quantization for BERT Compression" 논문을 소개하였습니다. 본 논문에서는 기학습된 언어 모델을 양자화를 통해 압축하는 기법인 SensiMix를 제안합니다. SensiMix에서는 모듈, 층별 민감도를 반영한 양자화 기법과 1-bit 양자화의 성능을 향상시키기 위한 세 가지 기법들(ABWR, PT, ILF)을 제안하였고, FP16 GEMM 및 XNOR-Count GEMM을 이용하여 양자화된 모델의 추론 속도 또한 향상시켰습니다. 또한, 실험을 통해 SensiMix가 기존 기법들보다 더욱 작고, 빠르면서도 성능이 높은 모델을 만들 수 있다는 것 또한 보여주었습니다.
SensiMix를 이용할 경우 기존의 크고, 느렸던 언어 모델들의 성능을 최대한 유지하면서 작고, 빠르게 만들 수 있기 때문에 기존 언어 모델을 이용한 서비스들의 지연 속도를 개선하거나, 모바일 기기들을 위한 작은 언어 모델을 만드는데 큰 도움이 될 수 있습니다. 논문에서는 본 문서에서 다룬 내용 이외에도 더욱 자세한 설명들과 실험 결과들이 있으니 관심 있으신 분들은 논문을 참고해주시면 감사하겠습니다. (링크)