Accurate Sublayer Pruning for Large Language Models by Exploiting Latency and Tunability Information

Posted by 날짜:
본 문서에서는 IJCAI'25 학회에서 발표될 "Accurate Sublayer Pruning for Large Language Models by Exploiting Latency and Tunability Information" 논문을 소개합니다. 논문에 대한 상세 정보는 다음과 같습니다.
  • Title: Accurate Sublayer Pruning for Large Language Models by Exploiting Latency and Tunability Information
  • Authors: Seungcheol Park, Sojin Lee, Jongjin Kim, Jinsik Lee, Hyunjik Jo, and U Kang
  • Conference: The 34th International Joint Conference on Artificial Intelligence (IJCAI 2025)

Sublayer Pruning for Large Language Models

최신 거대 언어 모델(large language model)들은 챗봇, 번역, 코드 생성 등 실생활의 다양한 작업을 수행하는 데에 있어서 뛰어난 성능을 보이고 있습니다. 하지만, 거대 언어 모델을 구성하는 방대한 양의 파라미터 수는 모델의 실질적인 운용을 어렵게 만드는 요소입니다. 서브레이어 가지치기(sublayer pruning)는 거대 언어 모델에서 불필요한 서브레이어를 삭제함으로써 모델의 크기로 인한 문제점을 효과적으로 해결합니다. 거대 언어 모델은 multi-head attention (MHA) 서브레이어와 multi-layer perceptron (MLP) 서브레이어가 번갈아서 층을 구성하고 있으며, 서브레이어들의 연산은 병렬적으로 수행할 수 없기 때문에 서브레이어 가지치기는 단순하지만 효과적으로 거대 언어 모델의 추론 속도를 극대화할 수 있습니다.

그림 1. 기존 가지치기 기법과 본 논문의 제안 기법 (SPRINT) 비교

서브레이어 가지치기의 핵심 목표는 핵심 목표는 가지치기할 대상 서브레이어를 정확하게 선별하는 것입니다. 하지만 기존의 서브레이어 가지치기 기법들은 각 서브레이어의 특징을 적절하게 반영하여 선별하지 못하기 때문에 가지치기 후 모델의 정확도를 효과적으로 보존하지 못합니다. 거대 언어 모델을 구성하는 두 가지 종류의 서브레이어는 추론 속도 측면에서 서로 다른 양상을 보이지만, 기존 기법은 이를 고려하지 못합니다. 또한, 기존 기법들은 가지치기 후 미세 조정에 따라 서브레이어의 중요도가 달라진다는 점을 고려하지 않기 때문에, 정확한 가지치기 대상 서브레이어를 선별하지 못합니다. 그림 1의 예시는 기존 기법이 가지는 한계점을 보여주고 있습니다.
본 논문에서 해결하고자 하는 거대 언어 모델 가속 문제의 구체적인 정의는 다음과 같습니다.
  • 주어진 정보
    • 학습이 되어 있는 고성능의 거대 언어 모델
    • 적은 수의 샘플 데이터셋
    • 추론 시간 제약
  • 목표
    • 주어진 추론 시간 제약을 만족하는 정확한 언어 모델을 생성
거대 언어 모델의 정확도를 유지하면서 서브레이어를 가지치기하는 알고리즘을 설계하기 위해서는 아래와 같은 문제점을 해결해야 합니다.
  • (서브레이어들의 추론 시간 차이) 기존 기법들은 거대 언어 모델을 구성하는 두 가지 서브레이어의 추론 시간 차이를 고려하지 않습니다. 거대 언어 모델의 추론을 효과적으로 가속하기 위해, 서로 다른 추론 시간을 가진 서브레이어들의 중요도를 어떻게 계산해야 할까요?
  • (미세 조정에 따른 서브레이어 민감도 변화) 기존 기법들은 미세 조정에 따라 서브레이어의 민감도가 변할 수 있다는 사실을 간과합니다. 어떻게 하면 서브레이어의 민감도를 더욱 정확하게 측정할 수 있을까요?
  • (반복 가지치기로 인한 가지치기 비용) 반복 가지치기는 높은 정확도를 보장하지만, 모든 서브레이어의 중요도를 반복적으로 예측해야 하므로 비용이 매우 비쌉니다. 어떻게 하면 반복 가지치기로 인한 비용을 줄일 수 있을까요?

Proposed Method

그림 1. SPRINT의 작동 방식
그림 2. SPRINT의 작동 방식.

해당 논문에서는 앞서 설명한 문제점을 해결하기 위한 기법인 SPRINT를 제안합니다. SPRINT는 아래와 같은 핵심 아이디어로 구성되어 있습니다.
  • (추론 시간을 고려한 중요도 계산) 서브레이어를 삭제했을 때 모델의 추론 속도가 향상되는 정도를 반영하여 서브레이어의 중요도를 계산합니다.
  • (회복 가능성을 고려한 민감도 측정) 서브레이어를 삭제한 후 미세 조정을 했을 때 모델의 성능이 복구되는 정도를 반영하여 서브레이어의 민감도를 측정합니다.
  • (불필요한 연산 삭제) 중단값 저장과 빠른 가지치기 후보 선택을 통해 반복 가지치기에서의 불필요한 연산을 삭제합니다.
그림 2는 6개의 서브레이어(sublayer)로 이루어진 거대 언어 모델에 대한 SPRINT의 동작 방식을 도식화한 것입니다. SPRINT는 우선 두 가지 종류의 서브레이어에 대해, 서브레이어가 삭제되었을 때 모델의 추론 속도가 얼마나 빨라지는지를 각각 측정합니다. 이렇게 측정한 추론 속도 차이에 기반하여, (1) 각 서브레이어의 민감도를 측정하여 중요도를 계산하고, (2) 가장 덜 중요한 서브레이어를 지우는 과정을 반복합니다. 각 과정의 핵심적인 내용에 대해서는 아래에서 자세하게 설명하겠습니다.

Latency-aware Importance Scoring

기존 서브레이어 가지치기 기법은 서브레이어를 가지치기했을 때 모델의 추론 속도가 가속되는 정도를 고려하지 않기 때문에, 동일 가속량에서 모델의 정확도를 하락시키지 않는 서브레이어를 정확하게 고를 수 없다는 한계점이 있습니다. SPRINT에서는 이러한 한계점을 극복하기 위해, 서브레이어의 중요도를 계산할 때 해당 서브레이어가 입력값을 처리하는 데에 걸리는 시간 정보를 고려함으로써 높은 가속량과 적은 정확도 손실을 일으키는 서브레이어를 선별하여 가지치기할 수 있게 합니다.

Tunability-aware Sensitivity Evaluation

일반적으로 가지치기 이후 발생하는 모델의 정확도 손실을 복구하기 위해 추가적인 미세 조정(fine-tuning)이 수반됩니다. 이때, 각 서브레이어의 민감도는 미세 조정에 따라 변화할 수 있습니다. 결과적으로 미세 조정 이후 민감도가 가장 낮은 서브레이어를 선택하는 것이 최종 모델의 성능을 가장 높게 유지할 수 있는 방법이지만, 기존 기법은 미세 조정에 따른 민감도 변화를 고려하지 않기 때문에 정확한 가지치기 대상 서브레이어를 선별할 수 없습니다. SPRINT는 이러한 문제점을 해결하기 위해, 가지치기 과정에서 미세 조정에 따라 서브레이어의 출력값이 회복되는 정도를 고려하여 민감도를 측정함으로써 더욱 정확한 가지치기 대상 서브레이어를 선별할 수 있게 합니다.

Avoiding Unnecessary Computations

반복 가지치기 기법은 압축된 모델의 성능을 높일 수 있지만, 모든 서브레이어의 민감도를 반복적으로 측정해야 하기 때문에 가지치기 비용이 많이 든다는 단점이 있습니다. SPRINT에서는 반복 가지치기의 고비용 문제를 해결하기 위해 (1) 중단값 저장과 (2) 빠른 가지치기 후보 선택을 통해 반복 가지치기에서의 불필요한 연산을 줄입니다. 각각에 대한 자세한 설명은 아래에서 하겠습니다.

Activation Checkpointing

SPRINT에서는 중단값 저장을 통해 반복 가지치기 과정에서 불필요한 계산을 삭제합니다. 예컨대 10개의 서브레이어로 구성된 모델을 가지치기하는 상황을 가정했을 때, 첫 번째 스텝에서 5번째 서브레이어를 가지치기했다면 1~4번째 서브레이어의 민감도는 변화하지 않을 것이므로, 두 번째 스텝에서 민감도를 다시 계산하지 않아도 됩니다. 이러한 특징을 이용해, SPRINT는 중단값을 구간마다 저장하여 재사용함으로써 불필요한 연산을 수행하지 않도록 합니다.

Fast Candidate Selection

SPRINT는 가지치기 대상 서브레이어를 선별할 때 미세 조정에 따른 민감도 변화를 고려하기 때문에, 모든 서브레이어에 대해 미세 조정 후의 민감도를 측정할 경우 시간이 오래 걸릴 수 있습니다. SPRINT는 빠른 가지치기 후보 선택을 통해 민감도 측정에 걸리는 시간을 단축합니다. 빠른 가지치기 후보 선택은 가지치기 전후의 서브레이어 민감도가 큰 폭으로 변동하지는 않는다는 가정을 기반으로, 모든 서브레이어에 대해 '미세 조정 전'의 민감도를 통해 '미세 조정 후'의 민감도를 측정할 일부 서브레이어를 선별하고, 선별된 서브레이어에 대해서만 '미세 조정 후'의 민감도를 측정합니다.

Experiments

해당 논문은 제안한 기법의 우수성을 입증하기 위해 기존 기법들과의 비교 실험을 수행하였습니다. 논문에서는 거대 언어 모델의 일반 상식 추론 능력을 평가하기 위한 5개의 zero-shot commonsense reasoning (CSR) 벤치마크에 대해 압축된 모델의 정확도를 보고하였습니다. 아래 그림 3은 오픈소스 거대 언어 모델인 Llama-2 및 Llama-3 모델에 기존 가지치기 기법 및 SPRINT를 적용한 결과입니다. SPRINT는 기존 기법 대비 동일 가속 성능에서 최대 23.88%p 높은 정확도를 달성하였습니다.

그림 3. 기존 가지치기 기법들과 SPRINT의 가속 대비 정확도.

그림 4는 1.4배 가속 상황에서 기존 서브레이어 가지치기 기법들과 SPRINT의 가지치기 비용 대비 CSR 벤치마크의 평균 정확도를 나타냅니다. SPRINT는 기존 서브레이어 가지치기 기법들에 비해 높은 정확도 및 효율성을 달성하였습니다.

그림 4. 기존 서브레이어 가지치기 기법들과 SPRINT의 가지치기 비용 대비 정확도.

그림 5는 Llama 모델들에 대한 SPRINT의 가지치기 패턴을 보여줍니다. 노란색은 MHA 서브레이어, 파란색은 MLP 서브레이어, 그리고 회색은 가지치기된 서브레이어를 나타냅니다. SPRINT는 MHA 서브레이어를 MLP 서브레이어보다 많이 가지치기합니다. 또한, 중상부에 위치한 서브레이어를 주로 가지치기하는 양상을 보이는데, 이를 통해 중단값 저장이 가지치기 비용을 줄이는 데에 효과적으로 기여하였음을 확인할 수 있습니다.

그림 5. Llama 모델들에 대한 SPRINT의 가지치기 패턴

Conclusion

본 문서에서는 2025년 IJCAI 학회에서 발표될 "Accurate Sublayer Pruning for Large Language Models by Exploiting Latency and Tunability Information" 논문을 소개하였습니다. 해당 논문에서는 거대 언어 모델의 추론 속도를 효과적으로 향상할 수 있는 서브레이어 가지치기의 정확도를 높인 SPRINT 기법을 제안하였습니다. 구체적으로는, 서브레이어별 가지치기에 따른 가속량 차이를 기반으로 서브레이어의 중요도를 계산하고, 가지치기 후 미세 조정에 따른 서브레이어 민감도 변화를 반영하여 더욱 정확한 가지치기 대상 서브레이어를 선별함으로써 더욱 정확한 압축 모델을 만들어낼 수 있습니다. SPRINT는 별도의 가속 커널을 이용하지 않고도 거대 언어 모델의 추론 속도를 효과적으로 향상하면서도 높은 정확도를 유지할 수 있게 해주어, 제한된 자원에서도 거대 언어 모델을 운용할 수 있도록 합니다. 본 논문에 대한 자세한 정보는 [링크]에서 확인할 수 있습니다.