[24’] Not All Layers of LLMs Are Necessary During Inference
카테고리: Language
🔍 Abstract
굉장히 도발적인 제목이다. 이 논문은 Inference 시 LLM의 모든 Layer가 필요하지 않다는 것을 주장한다. 즉, 어떤 task에서는 LLM이 이미 intermediate layer에서 final output과 거의 동일한 output을 얻을 수 있다는 것이다. 저자들의 실험 결과를 정리하면 다음과 같다.
- Inference 시에 LLM의 모든 Layer를 사용하지 않아도 거의 동일한 성능을 얻을 수 있다. 즉, Early Stopping을 통해 LLM의 Inference 시간을 줄일 수 있다.
- Simple task는 적은 layer가 필요하고, Complex task는 많은 layer가 필요하다. 즉, task의 난이도에 따라 필요한 layer의 수가 달라진다.
그러나 하나의 task 내부에서도 token에 따라 필요한 layer의 수가 다를 수 있다. 따라서 저자들은 token probability를 활용하여 SVM을 훈련시키고, 각 token에 대해 충분한 probability가 된 경우 Early Stopping을 진행하는 AdaInfer를 제안한다. 이를 통해 17.8%의 Pruining Rate를 달성하였다.
1. Not all Layers are Necessary
저자들은 LLaMA2-7B 모델을 활용하여 Early Stopping을 했을 때의 정확도를 측정한다. 저자들은 Sentiment Analysis(실선)라는 비교적 쉬운 문제와 MMLU(점선)라는 비교적 어려운 문제를 사용하였다. 결과적으로 (1) Early Stopping 시의 정확도와 Final Layer의 정확도에는 큰 차이가 없었고, (2) 어려운 문제일수록 더 나중 layer가 필요했다. 굉장히 직관적이고, 놀라운 결과이다.
2. AdaIner: Adaptive Inference
따라서 저자들은 AdaInfer라고 하는 Early Stopping Module을 도입한다. AdaInfer는 크게 2가지 단계로 이루어지는데, 각각 (1) Feature Selection, (2) Binary Classification이다. 각각을 살펴보자.
2.1. Feature Selection
먼저 저자들은 Feature Selection을 위해 어떤 Feature를 사용할 수 있을지에 대해 조사하였다. 즉, token에서 어떤 정보를 사용해야 해당 token이 필요한지를 판단할 수 있을지 조사하였다. 저자들은 가능성 있는 Feature로 다음과 같은 것들을 제안한다.
- Gap: P(top token) - P(second token)으로, 해당 token에 대한 일종의 confidence score로 볼 수 있다.
- Top Prob: 해당 token의 probability, 즉 P(top token)을 의미한다.
- Cosine Similarity: token의 current layer, previous layer 간 cosine similarity를 의미한다. Attention value, MLP output, hidden state 등을 사용할 수 있다.
위 그림을 통하여 알 수 있듯이 Observation 때와 유사한 경향성을 보이는 Feature로 Gap, Top Prob을 꼽을 수 있다. 즉 (1) Deep Layer에서 높아지고 (2) Task의 난이도에 따라 Saturated되는 Layer가 다르다는 점에서 유사한 경향성을 보인다. 따라서 저자들은 Feature vector $x^ d = [\text{Gap}, \text{Top Prob}]$을 사용하였다.
2.2. Binary Classification
Feature Selection을 통해 얻은 Feature를 바탕으로 Binary Classification을 수행한다. 데이터는 $(x^ d, y_ c)$ 쌍을 사용하였고, $y_ c$는 해당 layer에서의 output token이 final decoding token과 동일한지 여부를 나타낸다. 이를 통해 SVM을 훈련시키고, 각 token에 대해 충분한 probability가 된 경우 Early Stopping을 진행한다.
3. Experiments
3.1. Main Results
논문에서는 Accuracy에 추가로 Pruning Ratio와 Average Number of Activated Layers를 보고하여 효율성이 얼마나 되는지를 확인하였다.
그 결과는 위와 같고, Similarity score 기반으로 pruning을 진행하는 ShortGPT에 그 성능을 비교하였다. 그 결과로 (1) AdaInfer는 1% 미만의 성능 감소를 보이면서 (2) 평균 17.8%의 Pruning Rate를 달성하였다. 그리고 SVM의 훈련 결과 어려운 task일수록 더 많은 layer가 필요하다는 것을 확인할 수 있었다.
3.2. Factor Study
저자들은 2.1절에서 결정한 Feature Selection에 대해 다양한 실험을 진행하였다. 2.1절에서는 결국 Gap, Top Prob이 가장 효과적일 것이라 생각하여 이 두 가지만을 선택하여 SVM 학습에 사용하였는데, 추가로 Cosine Similarity를 사용하였을 때의 결과를 확인하였다.
결과적으로 Attention은 성능에 영향을 미치지 않았고, Hidden state와 MLP output은 오히려 성능을 저하시켰다. 따라서 Gap, Top Prob이 가장 효과적인 Feature라는 것을 재확인할 수 있었다.
💡 Summary
Not All Layers of LLMs Are Necessary During Inference는 LLM의 Inference 시에 모든 Layer를 사용하지 않아도 거의 동일한 성능을 얻을 수 있다는 것을 보여주는 논문이다. 해당 논문의 내용을 요약하면 다음과 같다.
- Inference 시에 Layer Early Stopping을 통해서도 거의 동일한 성능을 얻을 수 있으며, 어려운 task일수록 더 많은 layer가 필요함
- Early Stopping을 언제 진행할지 결정하기 위해 가장 중요한 Feature는 Token Probability이며, 이를 활용하여 AdaInfer를 제안함
그런데 그렇다면 late layer의 역할은 무엇인지 궁금하다. 이 논문에서는 이에 대한 해답을 제시하지 않았다. 다른 연구들에서도 late layer가 불필요하며, 이를 오히려 pass할 때의 성능이 좋다는 보고들이 있다. 한 가지 가능성을 생각해보면 훈련 과정에서 유추하기 어려운 Fine-grained information을 학습하다 보니 이와 관련된 Hallucination이 late layer에서 학습되어 오히려 문제를 일으킨다는 것이다. 관련하여서는 아직 Background가 부족하여 더 많은 논문들을 찾아볼 필요가 있다.
댓글 남기기