**Bo-Kyeong Kim**
Senior Researcher, Nota AI

소개

대형 언어 모델(LLM)은 세계에 혁신을 가져오고 있습니다. [Kaplan et al, 2020]의 확장 법칙과 추가 증거에 따르면, 모델의 크기가 커질수록 모델의 성능이 더 우수해지는 것으로 알려져 있습니다. 그러나 이러한 모델은 상당한 재정적, 계산적 자원이 필요합니다. 이 연구에서는 구조적 가지치기 방식을 통해 대형 언어 모델의 추론 비용을 줄이고, 이를 통해 개별 하드웨어에서 모델의 처리 속도를 효과적으로 향상시키는 것을 목표로 하고 있습니다.

🍀 추가 정보를 위한 리소스는 여기서 확인하실 수 있습니다. : GitHub, ArXiv.

🍀 ICLR’24 Workshop on ME-FoMo에서 채택되었으며, Daily Papers by AK에 소개되었습니다.

개요

대형 언어 모델은 여러 개의 트랜스포머 블록[Vaswani et al., 2017]이 쌓여서 구성되며, 각 블록은 다중 헤드 어텐션(MHA) 모듈과 순방향 신경망 (FFN)을 포함합니다. 대형 언어 모델에 대한 구조적 가지치기를 위해, 처음에는 너비 가지치기가 시도되었습니다. LLM-Pruner [Ma et al., 2023]와 FLAP [An et al., 2024]는 다중 헤드 어텐션 모듈의 어텐션 헤드와 순방향 신경망의 중간 뉴런을 가지치기하여 네트워크 너비를 줄입니다. 이번 연구에서는 트랜스포머 블록 일부를 제거하는 대형 언어 모델에 대한 깊이 가지치기 방법을 제안합니다. 저희는 대형 언어 모델의 추론 효율성에 미치는 영향을 비교 분석하며 너비와 깊이 두 가지 가지치기 차원을 비교했습니다 (그림 1 참조).

방법

그림 2는 연구의 접근 방식을 보여줍니다. 먼저, 각 블록의 중요도를 계산해 제거해야 할 블록을 식별합니다. 이후 여러 블록을 한 번에 제거하는 일회성 가지치기를 수행합니다. 가지치기된 네트워크를 얻은 후에는 빠르고 메모리 효율적인 방식으로 생성 성능을 회복하기 위해 Low-Rank Adaptation (LoRA) 재훈련 [Hu et al., 2022]을 진행합니다.

llm pruning — **그림 2.** 저희의 깊이 가지치기 접근 방식입니다. 간단한 지표를 사용해 중요하지 않은 블록을 식별한 후, 일회성 가지치기를 수행하고 간단한 재훈련을 진행합니다. 오른쪽에 있는 LoRA 그림은 Hu et al. [2022]에서 인용한 것입니다.

저희는 각 블록의 중요도를 평가하기 위해 다음 가지치기 기준을 사용했습니다. 먼저, 선형 가중치 행렬은 크기가 $(d_{\text{out}}, d_{\text{in}})$ 인 $\mathbf{W}^{k,n} = \left[W_{i,j}^{k,n}\right]$ 로 표시됩니다. 여기서 $k$ 는 $n$ 번째 트랜스포머 블록에서의 연산 유형(예: 다중 헤드 어텐션(MHA)에서의 쿼리 투영 또는 순방향 신경망(FFN)에서의 상향 투영)을 나타냅니다. 가중치 중요도 점수는 출력 뉴런 수준에서 계산되고, 이 점수를 합산하여 블록의 중요도를 평가합니다.

Taylor+ 기준: 이 기준은 주어진 데이터셋 $D$ 에서 가중치가 가지치기되어 0으로 대체될 때의 학습 손실 $\mathcal{L}$ 의 변화를 기반으로 합니다. [LeCun et al.,1989; Molchanov et al., 2019] 식은 다음과 같습니다: $\left| \mathcal{L}(W_{i,j}^{k,n}; D) - \mathcal{L}(W_{i,j}^{k,n} = 0; D) \right| \approx \left| \frac{\partial \mathcal{L}(D)}{\partial W_{i,j}^{k,n}} W_{i,j}^{k,n} \right| $, 연구자들은 이 계산에서 첫 번째 차수 항만 사용해도 충분하다는 것을 입증했습니다. 저희는 블록의 중요도 점수를 다음과 같이 정의합니다: $I_{\text{Taylor}}^n = \sum_k \sum_i \sum_j \left| \frac{\partial \mathcal{L}(D)}{\partial W_{i,j}^{k,n}} W_{i,j}^{k,n} \right|$. 기호 `+' 는 첫 블록과 마지막 블록 몇 개를 가지치기하지 않고 유지하는 방법을 나타냅니다.

Perplexity (PPL) 기준: 각 트랜스포머 블록을 제거한 후, 보정 데이터셋 $D$를 사용해 PPL(혼란도) 지표에 미치는 영향을 확인합니다. $I_{\mathrm{PPL}}^n = \exp \left\{ -\frac{1}{SL} \sum_{s} \sum_{l} \log p_{\theta^{n}}(x_{l}^{(s)} | x_{<l}^{(s)}) \right\}$ 여기서 $\theta^{n}$ 는 $n$ 번째 블록을 제거한 모델을 의미합니다. $s = 1, \ldots, S$ 와 $l = 1, \ldots, L$ 은 보정 데이터셋 $D$ 에서 각각 시퀀스(문장의 순서) 와 토큰 (단어나 기호의 위치) 를 나타냅니다. 그림 3에서 보듯, 여러 블록을 제거해도 성능에 미치는 영향은 미미합니다. 하지만, 처음과 마지막 블록을 제거하면 성능이 크게 저하되므로 이 블록들은 유지해야 합니다.

결과

저희는 하드웨어 제약이 있는 소규모 배치 환경에서 대형 언어 모델의 추론 속도를 향상시키는 데 중점을 두었습니다. 이는 메모리가 제한된 로컬 장치에 대형 언어 모델을 배포할 때 자주 발생하는 상황입니다. 그림 4는 정량적 결과를, 그림 5는 정성적 생성 예시를 보여줍니다.

너비 가지치기를 통해 가중치 크기를 줄이는 것은 대형 언어 모델의 추론 속도 향상에 효과적이지 않습니다. 이는 대형 언어 모델의 추론이 메모리 제한에 크게 영향을 받기 때문입니다. 또한, 너비 가지치기는 가지치기 후의 가중치 크기가 GPU 성능에 맞지 않을 경우 오히려 속도를 저하시킬 수 있습니다.

저희 방법은 너비 가지치기 방법과 유사한 제로샷 성능을 유지하면서 추론 속도를 향상시킵니다. 특히, 일부 모듈을 완전히 제거하는 깊이 가지치기를 통해서만 눈에 띄는 속도 향상을 이룰 수 있음을 입증했습니다.