2023년 PLOS ONE 저널에 발표된 "PET: Parameter-efficient Knowledge Distillation on Transformer" 논문을 소개합니다. 논문에 대한 상세한 정보는 다음과 같습니다.
- Title: PET: Parameter-efficient Knowledge Distillation on Transformer
- Authors: Hyojin Jeon, Seungcheol Park, Jin-Gee Kim, and U Kang
- Journal: PLOS ONE (2023)
Transformer Compression
어떻게 하면 기학습된 Transformer 모델을 효과적으로 압축할 수 있을까요? Transformer는 시퀀스 투 시퀀스(Sequence to Sequence) 모델로써 번역, 챗봇, 요약 등 다양한 자연어 처리(Natural Langage Processing, NLP) 분야에서 높은 성능을 이루었을 뿐만 아니라, BERT와 GPT 등 혁신적인 성능을 보이는 언어 모델들의 모태가 되는 매우 중요한 모델입니다. 그림 1은 영어를 불어로 번역하는 Transformer의 구조를 나타내고 있습니다. 그림에서 알 수 있듯이 Transformer는 입력 문장을 받아들이고 해석하는 인코더(Encoder)와 인코딩된 정보를 바탕으로 출력을 생성해내는 디코더(Decoder)로 이루어져 있습니다.
그림 1. Transformer의 구조 예시 (출처)
Transformer는 성능이 우수하지만, 그 크기가 매우 크고 연산량(FLOPs)이 매우 많다는 단점이 있습니다. 그래서, 이들의 높은 성능을 유지하면서 크기를 줄이고, 연산량을 줄이는 Transformer 압축(Compression) 연구가 활발하게 이루어지고 있습니다. 하지만, 기존의 연구들은 상대적으로 압축하기 쉬운 인코더(Encoder) 모듈만을 압축하는 경우가 많았고, 많은 연산량을 필요로 하는 디코더(Decoder) 모듈을 압축하는 연구는 많이 없었습니다. 본 논문에서는 인코더 뿐만아니라, 디코더 모듈까지도 효과적으로 압축하여 트랜스포머 모델의 크기를 크게 줄이는 것을 목표로 합니다.
PET: Parameter-efficient Knowledge Distillation on Transformer (Overview)
본 논문에서는 효과적인 Transformer 압축 기법인 PET(Parameter-Efficient knoweldge distillaion on Transformer)를 제안합니다. PET는 가중치 공유(Parameter Sharing) 기법과 지식 증류(Knowledge Distillation) 기법을 효과적으로 사용하여 기학습 Transformer의 정확도 손실을 최소화하며 효율적으로 압축합니다. PET의 주요 아이디어들은 아래와 같습니다.
- Finding replaceable pairs in Encoder and Decoder
- 가중치 공유를 하기 위해서는 어떤 가중치들을 공유해도 되는지를 알아내야 합니다. PET에서는 실험을 통해 가장 공유하기 좋은 가중치 조합을 찾아냅니다.
- Warming up with a simplified task
- 트랜스포머 모델의 추론 결과는 모든 종류의 토큰들에 대한 분류(Classification) 확률 로써 트랜스포머 모델은 매우 복잡한 다중 레이블 분류(Multi-label Classification) 문제를 풀게 되며, 작은 모델이 학습하기에 어렵습니다. 그러므로 더욱 간단한 문제(simplified task)를 설계하여 효과적으로 가중치를 초기화 한 후에 학습을 시작하여 성능을 극대화하고자 합니다.
- Mapping prediction probabilities to labels of simplified task
- 앞서 설계한 간단한 문제의 경우 기존 트랜스포머 모델로써는 바로 풀수가 없기 때문에, 이를 풀기 위한 새로운 방법을 제시합니다.
지금부터는 위 아이디어들에 대해 자세하게 설명해드리도록 하겠습니다.
Finding replaceable pairs in Encoder and Decoder
첫 번째 아이디어는 가중치 공유를 효과적으로 하기 위한 방법입니다. 트랜스포머에서는 어텐션(attention) 연산을 위한 쿼리(query), 키(key), 벨류(value) 전사(projection)를 위한 가중치들이 대부분을 차지하며, 서로 다른 가중치들이 층층이 쌓여있는 형태를 띄고 있습니다. 이러한 구조에서 만약 가중치를 공유해서 사용한다면, 저장해야할 가중치의 숫자를 크게 줄일 수 있어 모델의 크기가 줄어들게 됩니다. 이때, 각 모듈별로 어떤 가중치들을 공유하는 지에 따라 모델의 성능이 크게 달라 질 수 있는데, PET에서는 각 조합별로 가중치를 공유해서 학습을 해봄으로써 가장 좋은 가중치 공유 조합을 찾습니다.
그림 2. PET에서의 가중치 공유 패턴
그림 2는 PET에서 찾은 가중치 공유 패턴을 나타내고 있으며, Q는 쿼리, K는 키, V는 벨류를 의미하고 self는 셀프-어텐션(self-attention)을, ed는 인코더-디코더 어텐션(Encoder-Decoder Attention)을 의미합니다. PET에서는 모듈의 특성을 고려하여 인코더(왼쪽)와 디코더(오른쪽)에서 각각 효율적인 공유 패턴을 찾았습니다. 인코더에서는 첫 번째 레이어와 세 번째 레이어의 쿼리와 키 전사 가중치들을 서로 교대로 공유하고 있으며, 두 번째 레이어와 네 번째 레이어 또한 가중치를 같은 방식으로 공유합니다. 디코더의 경우 첫 번째 레이어와 세 번째 레이어의 파라미터들을 그대로 공유하되, 셀프 어텐션과 인코더-디코더 어텐션의 쿼리 가중치 또한 공유함으로써 필요한 가중치의 크기를 최소화 합니다. PET에서는 이러한 효율적인 가중치 공유 조합을 찾고, 이들의 가중치를 공유함으로써 모델의 크기를 효과적으로 줄였습니다.
Warming up with a simplified task
트랜스포머는 시퀀스-투-시퀀스(sequence-to-sequence) 모델로써, 새로운 문장을 만들어내는 모델입니다. 생성된 문장의 각 토큰을 생성해내는 문제는 토큰 집합(vocabulary) 속의 모든 종류의 토큰들에 대해 등장할 확률을 계산하고, 이들 중 가장 확률이 높은 토큰을 선택하는 일로써 매우 복잡한 분류(classification) 문제를 매번 풀어내는 일과 같습니다. 그러므로, 압축된 작은 모델이 처음부터 이러한 학습하기에는 매우 어려우며, 이러한 문제를 해결하기 위해 PET에서는 간단한 문제(simplified task)를 설계하여 모델을 선행 학습(warming up)하고자 합니다.
기존 분류 문제가 풀기 어려운 이유는 분류 문제의 레이블(label) 개수가 너무 많기 때문이었기에, 기학습 모델의 추론 결과를 반영하여 아래와 같이 각 추론의 결과를 네 개의 레이블로 단순화합니다.
- Confidently correct: 기학습 모델의 추론 결과가 정답과 같으며, 기학습 모델이 추론에 대해 확신하는 경우를 의미합니다.
- Correct, but not confident: 기학습 모델의 추론 결과가 정답과 같으나, 기학습 모델이 추론에 대해 확신하지 못하는 경우를 의미합니다.
- Confidently wrong: 기학습 모델의 추론 결과가 정답과 다르나, 기학습 모델이 추론에 대해 확신하는 경우를 의미합니다.
- Wrong, but not confident: 기학습 모델의 추론 결과가 정답과 다르며, 기학습 모델이 추론에 대해 확신하지 못하는 경우를 의미합니다.
각 토큰별로 추론 결과를 위와 같은 4개의 레이블에 대하여 분류하는 문제로 수정하면, 레이블 수가 매우 적어져서 문제가 단순화됩니다. 하지만, 압축 모델의 추론 결과는 수정되지 않아 아직 모든 토큰들에 대한 분류 문제를 풀고 있기 때문에 학습이 불가능합니다. PET는 이 문제를 이어지는 아이디어를 활용하여 해결합니다.
Mapping prediction probabilities to labels of simplified task
그림 3은 앞서 소개한 문제 상황(왼쪽)과 이를 해결하기 위해 PET에서 선택한 방법(가운데), 그리고 이후의 본 학습 과정(오른쪽)을 순차적으로 나타내고 있습니다. 이때, 그림에서 6632는 예시 상황에서의 레이블의 수, 즉, 토큰 집합 속의 모든 토큰의 종류를 의미합니다. 왼쪽 그림에서는 앞서 말했듯이 압축 모델의 출력 결과는 모든 토큰 종류의 개수와 같으므로 6632개가 되고, 이는 단순화된 문제의 레이블의 수인 4개와 달라 학습이 불가능한 형태입니다. PET에서는 이를 해결하기 위하여 6632개의 추론 결과를 활용하여 확률 모델링을 진행하여 각 레이블별 확률을 만들어 냅니다. 그리고, 선행 학습이 끝나고 나면, 기존의 본 학습 과정으로 넘어가게 됩니다.
그림 3. 단순화된 문제의 레이블을 활용했을 때의 문제점과 PET의 해결 방법
- Confidently correct: (모든 토큰 중 가장 높은 확률) $\times$ (정답 토큰의 확률)
- Correct, but not confident: (1-모든 토큰 중 가장 높은 확률) $\times$ (정답 토큰의 확률)
- Confidently wrong: (모든 토큰 중 가장 높은 확률) $\times$ (1-정답 토큰의 확률)
- Wrong, but not confident: (1-모든 토큰 중 가장 높은 확률) $\times$ (1-정답 토큰의 확률)
즉, 확신의 정도를 모든 토큰 중 가장 높은 확률로, 정답일 확률을 정답 토큰의 확률로 해석하여 계산한 것입니다. 해당 방법의 경우 새로운 모듈이 추가되지 않고, 기존 모델의 추론 결과를 그대로 사용하기 때문에, 선행 학습 후에 학습된 모듈을 삭제하는 등의 과정으로 인한 정확도 손실이 없다는 장점이 있습니다.
Experiment
본 논문에서는 독일어를 영어로 번역하는 테스크(task)인 IWSLT'14 DE-EN benchmark를 이용하여 PET를 기존의 다른 압축 기법들의 성능을 비교하였으며 압축 대상 모델로는 Transformer-base를 사용하였습니다.
그림 4. PET와 경쟁 기법들의 Transformer 압축 결과 비교
그림 4는 실험 결과를 나타내고 있습니다. 이때, Uncompressed는 기학습 모델로써 Transformer-base 모델의 성능을 의미하며, PET 및 경쟁상대로 압축했을 때, 크기 및 추론 속도에 따른 성능 변화를 각각 나타내고 있습니다. 실험 결과 PET는 기학습 모델의 성능을 최대한 유지하면서 크기를 약 81.20%까지 줄였으며, 추론 속도를 약 45.15% 가량 가속화하는데 성공한 것을 알 수 있습니다. 해당 결과는 경쟁상대인 PKD와 CAN에 비해 매우 우수한 결과이며, 압축률이 높아질수록 PET와 경쟁상대들 간의 성능 차이가 급격히 커지는 것 또한 확인할 수 있습니다. 이를 통해 PET를 이용할 경우 기존의 압축 기술들보다 효과적으로 Transformer 모델을 압축할 수 있음을 알 수 있습니다.
Conclusion
본 문서에서는 2023년 PLOS ONE 저널에 발표된 "PET: Parameter-efficient Knowledge Distillation on Transformer" 논문을 소개하였습니다. PET는 가중치를 공유하기 위한 가장 좋은 조합을 찾고, 압축 모델의 성능을 극대화하기 위한 효율적인 선행 학습 기법을 제안합니다. 그리고, 해당 기법들을 통해 경쟁 기법들보다 더욱 우수한 압축 성능을 보이는 것을 실험적으로 입증하였습니다. PET를 사용할 경우 기존의 크고, 느린 언어 모델들의 성능을 최대한 유지하면서 작고 빠르게 만들 수 있기 때문에, 현재 매우 많은 수요가 있는 ChatGPT 등의 인공지능 기반의 자연어 처리 서비스들의 서비스 유지 비용을 효과적으로 줄이고, 결과물의 출력 속도를 향상 시키는 등의 긍정적인 효과를 볼 수 있습니다. 논문에서는 본 문서에서 다룬 내용 이외에도 더욱 자세한 설명들과 실험 결과들이 있으니 관심 있으신 분들은 논문을 참고해주시면 감사하겠습니다. (링크)