[논문 리뷰] Extremely Low Bit Transformer Quantization for On-Device Neural Machine Translation

온디바이스 AI

ChatGPT의 등장으로 인해 많은 일반 사용자들이 AI에 관심을 갖게 되었다. 이러한 관심은 자연스레 기업들이 자신들의 제품에 AI 기술들을 적용하도록 만들었다. 예를 들어, 2024년 2월 삼성전자에서 새로운 갤럭시 S24 시리즈를 출시하였는데, 온디바이스 번역(Neural Machine Translation) 모델을 탑재하여 인터넷 연결 없이도 빠른 속도의 번역 기능을 제공하고 있다. 뿐만 아니라, 애플에서는 ‘LLM in a flash’ 라는 논문을 게재하여, 플래시 메모리를 이용한 디바이스 내 LLM 인퍼런스가 가능하다고 주장하고 있는 상태이다.

그렇다면, 왜 이런 기업들은 온디바이스 AI 에 목을 메는 것일까? 흔히들 생각할 수 있는 이유는 바로 보안 문제이다. 사실 SKT에서 출시한 ‘에이닷’이나 네이버 ‘파파고’와 같은 AI 모델들은 서버향 AI 모델이기 때문에 나의 정보가 서버로 전송되는 것을 피할 수 없다. 통화요약을 위해서는 나의 통화내용이 반드시 서버로 전송되어야하고, 파파고 번역을 위해서는 내가 번역하고 싶은 문장이 서버로 전송되어야한다. 사용자들은 이러한 개인 정보가 외부로 노출되는 것을 꺼리기 때문에, 온디바이스 AI를 이용한다면 사용자 보안 문제를 해결할 수 있게 되는 것이다.

그러나 온디바이스 AI를 이용하는 더 큰 이유는 따로 있는데, 바로 비용 문제이다. 삼성이나 애플같은 회사는 기본적으로 ‘상품’을 파는 회사이다. 휴대폰을 팔고 난 뒤에 서버향 AI를 지속적으로 제공한다면 어떻게 될까? 기업은 서버를 계속 운용해야하며 이는 지속적인 비용을 의미한다. 따라서, 갤럭시를 판매한 뒤 사용자가 빅스비를 많이 이용할수록 오히려 비용이 증가해버리는 문제가 생긴다. (이러한 서비스들로 발생하는 브랜드 차원의 이익 등은 배제한다.) 그러므로 온디바이스 AI를 이용하여 서버 비용을 줄이고자 하는 노력이 끈임없이 생기고 있는 것이다.

Binary Quantization

하지만, 온디바이스 AI는 생각만큼 쉬운 기술이 아니다. 우리가 흔히 아는 서버향 AI (앞서 언급한 에이닷, 파파고 등) 는 모델의 크기가 매우 크기 때문에 성능은 더 좋을지언정 휴대폰에서 이를 직접 이용할 수 없다. 휴대폰을 샀는데, 이런 모델들이 램을 1GB씩 먹고 있다고 생각해보라. 기분이 벌써 나빠진다. 또한, 서버와 달리 휴대폰은 배터리도 한정적이며, 발열을 해결하기도 어렵다는 문제가 있다.

따라서 가장 중요한 것은 모델을 성능을 유지한 채 작게 만드는 것인데, 대표적으로는 Quantization 방법이 있다. 오늘 리뷰할 논문 Extremely Low Bit Transformer Quantization for On-Device Neural Machine Translation 은 번역 모델에서 많이 사용되는 Transformer 아키텍처를 효율적으로 Quantization 하는 방법을 다루고 있다. 여러 가지 Quantization 기법 중에서도 Binary Quantization 을 이용하고 있으며, Transformer 내 각 Layer 별로 Quantization 정도를 개별로 적용했다는 점에서 흥미롭다고 할 수 있다.

본 포스트에서는 Quantization 개념 자체를 다루는 것이 아니므로 간단하게만 설명하고 넘어가겠지만, 관심 있는 사람들은 잘 정리된 이 곳 에서 공부를 할 수 있을 것이다.

먼저, Binary Quantization 은 아래 그림과 같이 표현될 수 있는데,

Weight Matrix가 기본적으로 +1 혹은 -1 의 부호로 저장되는 것을 확인할 수 있다. (이 역시도 어떤 weight을 +1, 혹은 -1로 해야하는 지에 대해 많은 방법론이 존재하는데 본 논문에서는 Greedy Approximation 을 이용했다.) 물론, 이렇게 간단하게 해버리면 모델 성능이 엉망이 되어버리기 때문에 Binary Quantization에서는 두 가지 고려사항이 생긴다. 먼저, ‘k bits’를 활용하게 된다. Weight Matrix 를 +1, -1로 이루어진 Quantized Matrix로 표현하면 하나의 Parameter는 1bit 로 구성되는데, 단순히 한 개의 Quantized Matrix를 사용하지 않고 k개를 사용한다는 것이다. 그림에서는 2bits ($B_1$, $B_2$) 를 사용한 모습이다. 그리고 Scaling Factor ($\alpha_1$, $\alpha_2$)를 사용한다. 즉, 모든 Row마다 하나의 FP32 실수를 사용하여 각 부호를 Scaling 시키는데 사용하여 조금 더 높은 정밀도를 가져가는 것이다.

정리하자면, $n \times n$ Weight Matrix 를 표현할 때, 기존에는 $n \times n \times 32$ bits가 필요하다. 그에 반해 k-bits Binary Quantization 을 한다면 $n \times n \times 1 \times k + n \times 32 \times k$ bits만 필요하게 ($k « n$) 변하는 것이다.

Quantization (Embedding)

서론이 너무 길었는데, 본 논문에서는 크게 두 가지 부분에서 Quantization을 조금 ‘다르게’ 적용하고 있다. 먼저, Embedding Layer이다. 알다시피, Embedding Layer는 각 Token 을 표현하는 Vector들로 이루어져있다. 그런데 모든 Token이 과연 모두 ‘동등하게’ 중요할까?

저자들은 WMT14 데이터셋에서 모델의 32768개의 Token들이 등장한 빈도가 power-law 분포를 따르고 있다는 것을 확인하였다. 즉, 특히 중요한 몇 백, 몇 천 개의 단어가 문장의 대부분에서 사용되고 있다는 것이다. 그러므로, 중요한 단어들에게는 조금 더 많은 bits 를 투자하고, 중요하지 않은 단어들에게는 더 적은 bits를 투자하여 Embedding Layer 전체적인 관점에서 모델 크기를 줄이려는 시도를 하게된다.

중요하게 볼 것은 4, 5번째 라인이다. 단어 분포가 power-law 를 따르기 때문에, 단어들을 exponential 하게 클러스터링하여 각 클러스터마다 다른 k-bits 를 부여하는 것이다. 예를 들어, cluster 수 $b$ 를 4로 하고 ratio $r$ 을 2로 하는 경우를 생각해보자. 총 4개의 클러스터가 중요도(빈도)에 따라 1:2:4:8 의 비율로 형성되며, 각 클러스터에 포함된 단어들은 각각 4, 3, 2, 1 bits Binary Quantization 된다고 볼 수 있다.

Quantization (Encoder & Decoder)

두 번째로 저자들은 Encoder, Decoder, 내부의 Attention, Feed Forward Network 등 각각의 레이어에 따라 Quantization이 다르게 적용되어야한다고 생각했다. 다시 말해, 기존의 연구들은 Transformer 레이어 전체에 k-bits Binary Quantization을 일괄적으로 적용하였는데, 이보다는 각 Layer가 k 에 대해 얼마나 민감한가를 관찰하고 이에 따라 Layer 별로 다른 k를 써야한다는 것이다.

위 표에서 볼 수 있듯, Encoder ($Enc$) 와 Decoder ($Dec$) 그리고 내부에 존재하는 Attention, Feed Forward Network 들은 사용 되는 bit 수에 따라 모델의 성능에 미치는 영향이 각각 다르다. 예컨대 Encoder 내의 Feed Forward Network $Enc_{ffn}$ 와 Decoder 내의 Encoder-Decoder Attention Layer $Dec_{ed}$ 가 quantization bit 수에 특히 더 민감하게 반응하는 것을 확인할 수 있다. (bit 수가 줄어듦에 따라 성능 하락이 심함.) 따라서, 저자들은 각 sub-layer들을 이 민감도에 따라 다른 bit 수를 이용하여 Quantize 하였다.

Experiments & Results

먼저, Overall Performance를 살펴본다. Average Bits 는 각각의 Sub-Layer를 Quantize한 Bit 수를 평균낸 것을 의미한다. 즉, 2.5, 1.8, 3.7 모델은 Embedding Layer는 평균 2.5, Decoder Layer는 1.8, Encoder Layer는 3.7 bits로 Quantize한 모델이다. $Dec_{dd}$, $Dec_{ed}$, $Dec_{ffn}$ 은 각각 2, 3, 1 bits 로 Quantize 되었고, $Enc_{ee}$, $Enc_{ffn}$ 는 각각 3, 4 bits로 Quantize 되었다. 여기서 확인할 수 있듯, bit 수에 민감하게 반응했던 $Dec_{ed}$ 와 $Enc_{ffn}$ Sub Layer들은 더 높은 bit수를 할당받았다. 또 흥미로운 점은, Encoder 에 비해 Decoder 에 더 적은 bit 수를 사용했다는 점이다. 이는 Transformer 구조에서 Encoder 는 병렬화 특성을 쉽게 이용할 수 있지만, Decoder 는 Auto-regressive 한 방법으로 token-by-token을 생성하는 아키텍처이기 때문이다. 온디바이스 모델에서는 Latency가 중요하기 때문에 Decoder에 사용된 bit수를 의도적으로 줄인 것을 확인할 수 있다.

표에서 볼 수 있는 것처럼, 저자들의 방법을 이용하여 Transformer를 Quantize한 결과 BLEU 점수는 최대 0.5점밖에 떨어지지 않았지만 모델의 크기가 10배 이상 줄어든 것을 확인할 수 있다.

온디바이스 AI 모델에서 또 한가지 중요한 것은 Latency 이다. 상대적으로 좋은 Computation 환경을 가진 서버보다 디바이스 내의 CPU (혹은 GPU) 가 더 열악하기 때문에 큰 Latency를 보일 수 있고, 이는 사용성에서의 큰 하락을 의미하기 때문이다. 저자들은 Quantize 된 모델을 갤럭시 노트 10+ 에서 Inference 해보았는데, Quantize 이전의 FP baseline 모델보다 약 3.5배 빨라진 속도를 확인할 수 있었다. 저자들이 Decoder 에서 Bit 수를 특히 줄였는데, 이를 통해 디바이스 내에서 빠른 Latency를 달성할 수 있었다.

References

• https://arxiv.org/pdf/2312.11514.pdf
• https://jin-choi.tistory.com/29
• https://openaccess.thecvf.com/content_cvpr_2017/papers/Guo_Network_Sketching_Exploiting_CVPR_2017_paper.pdf
• https://aclanthology.org/2020.findings-emnlp.433.pdf