NL-257, Compact language models via pruning and knowledge distillation, NeurIPS 2024

◼ Comment

• 모델 경량화에서 유명한 논문이라고 한다
• 예로) 8B 모델을 만들때, 이전에는 scratch부터 학습하는 방식이 많았는데, 여기서 말하길 더 큰 모델 15B에서 distill해서 8B만드는게 훨씬 좋다는 것이다
• 성능도 좋고, 추가 학습에 들어가는 비용등등
• 과정은 크게 다음과 같다
• 15B->8B로 갈때 어디부분을 경량화 할지 정한다. 정확히는 탐색하는 과정이 있음. 
• layer 수를 줄여서 depth를 줄인다는지, hidden dim을 줄인다는지, attention 개수를 줄이다는지 등등을 정하는 과정이다
• 각각의 모듈에서 중요도를 계산한다. 
• 너비 가지치기: (hidden, embedding) 중요도라는 것은 별거 아니고 2.2에 보면나오는데, 해당 값의 합으로 보는게 basic한 방법이다. 
• 예를 들면 배치=5, 시퀀스길이=2, embedding dim이 768이라 하면, 5*2*768만큼 임베딩 채널이라고 보는데
• 5*2개의 값을 다 더하면 768개의 스칼라 값이 나올텐데, 여기서 점수가 높을수록 중요도가 높다는 것이다. 즉 이렇게 임베딩 차원을 줄일 수 있다는 것
• 어텐션헤드나, 뉴런중요도도 비슷하게 하는 것이고, 
• 중요도를 계산하는 여러 방법을 시도했는데 배치에서는 L2 norm, 시퀀스는 mean 방식으로 점수를 sum해서 중요도를 계산하는게 가장 좋다고 한다.
• 블록 중요도는 cosine 유사도를 활용한다고 함
• 깊이 가지치기: 레이어 중요도
• 첫 번째 메트릭은 layer을 제거한후, 모델의 PPL을 계산함으로써 중요도를 측정한다
• 두 번째 메트릭은, 블록의 입출력의 코사인유사도를 계산한다. 이 점수가 낮으면 값이 많이 바뀐것이므로 중요도가 높다고 간주하는듯
• 이렇게 모델 사이즈를 조금씩 줄여나간다. 
• 즉 한번에 중요도에 따라 모델을 한번에 경량화하지 않고, 일부 프루닝하고 다시 중요도 계산하고 ... (반복)
• 이렇게 여러 8B 모델 후보가 생성되는데
• 이를 lightweight retraining을 통해 가장 좋은 모델을 pick 한다.
• pick된 8B 모델은 제대로된 retraining 한다
• 여러 loss을 쓰는데, 일반적인 CLM loss
• teacher 모델를 통해 KD loss
• 중간 logits을 비교한 loss
• 여기서 KD가 가장 중요한 loss라고 보면됨
• 여기서 사용되는 학습 데이터는 Nemotron-4 모델을 학습할 때 사용했던 대규모 사전학습(pretraining) 데이터셋과 추가 학습 데이터(CT) 중 일부를 사용합니다. 
• 이 데이터는 1.8B 토큰으로 원래 pretraining 하는 데이터량 기준에서 소량의 수준이라고 보면 됨
• 어텐션 헤드 재분배 공식
• 이는 어텐션 헤드를 프루닝했을 때, 중요한 어텐션 헤드를 그대로 가져가는건 아니고, 사라진 어텐션 헤드와의 차이만큼 더해줘서 가져감
• layer collapse? 개념이라고함
• 15B->4B 모델을 만들땐
• 15B->8B->4B 이런식으로 가는게 좋다고 함 (8B에서도 당연히 retraining이 이뤄져야함)

Abstract 

대형 언어 모델(LLM)은 현재 각 배포 규모 및 크기에 맞춰 개별적으로 학습되어 생산되며, 이는 매우 높은 연산 비용을 요구합니다. 

본 논문에서는 기존 LLM을 가지치기(Pruning)한 후 원래 학습 데이터의 일부(<3%)만을 사용하여 재학습하는 방법이 반복적인 전체 학습을 대체할 수 있는지 조사합니다. 

• 이를 위해, 깊이(depth), 너비(width), 어텐션(attention) 및 MLP 가지치기를 지식 증류(knowledge distillation) 기반의 재학습과 결합하는 실용적이고 효과적인 압축 최적화 기법을 개발합니다. 
• 우리는 각 축에서의 가지치기 전략, 축을 결합하는 방법, 증류 전략, 최적의 압축 아키텍처를 찾기 위한 탐색 기법을 상세히 실험적으로 분석하여 이러한 최적화 기법을 도출합니다.

이 가이드를 활용하여 Nemotron-4 LLM 계열을 2~4배 압축하고, 다양한 언어 모델링 과제에서 유사한 크기의 모델과 비교 평가하였습니다. 사전 학습된 15B 모델을 활용하여 8B 및 4B 모델을 도출하는 과정에서, 본 접근법은 처음부터 학습하는 것과 비교해 모델당 최대 40배 적은 학습 토큰을 요구하며, 전체 모델 계열(15B, 8B, 4B)의 학습 연산 비용을 1.8배 절감할 수 있습니다.

MINITRON 모델은 MMLU 점수에서 최대 16% 향상된 성능을 보였으며, Mistral 7B, Gemma 7B, Llama-3 8B와 같은 커뮤니티 모델들과 유사한 성능을 유지하면서도, 기존 문헌에서 제시된 최신 압축 기법들을 능가하는 성능을 달성하였습니다. 우리는 MINITRON 모델 가중치를 Huggingface에서 오픈소스로 공개했으며, 예제 코드와 함께 관련 자료를 GitHub에서 제공하고 있습니다.

1 Introduction

대형 언어 모델(LLM)은 이제 실제 자연어 처리(NLP) 분야에서 지배적인 역할을 하며, 복잡한 맥락을 이해하는 데 뛰어난 능력을 보이고 있다 [7, 40, 52, 49, 48]. 다양한 배포 크기와 규모를 목표로 하는 사용자들을 지원하기 위해, 모델 제공자들은 종종 여러 크기의 모델을 처음부터 개별적으로 학습한다. 예를 들어, LLaMA-2 모델 패밀리 [49]는 70억, 130억, 700억 개의 파라미터를 가진 세 가지 변형을 제공하며, Pythia 패밀리 [6]는 8천만에서 120억 개의 파라미터를 가진 총 여덟 가지 모델을 포함한다. 그러나 여러 개의 수십억 개 파라미터를 가진 모델을 처음부터 학습하는 것은 엄청난 시간, 데이터, 자원이 소모되는 과정이다.

이 논문에서는 다음과 같은 질문을 제기한다. 

• 우리는 하나의 대형 모델을 학습한 후, 가중치 가지치기(weight pruning)와 재학습(retraining)을 결합하여, 원래 학습 데이터의 일부만을 사용하면서도 더 정확한(처음부터 학습한 것과 비교하여) 소형 모델을 얻을 수 있을까? 

이러한 목표를 달성할 수 있다면, 다양한 배포 규모에 맞춘 LLM 제작 비용을 대폭 절감할 수 있을 것이다.

가중치 가지치기는 모델 크기를 줄이는 강력하고 널리 알려진 기법이다 [51, 21]. 본 연구에서는 구조적 가지치기(structured pruning) 에 초점을 맞춘다. 

• 구조적 가지치기는 모델 가중치에서 특정 블록의 요소를 한 번에 제거하는 방식으로, 대표적인 방법으로는 뉴런 가지치기, 어텐션 헤드 가지치기, 합성곱 필터 가지치기, 깊이(depth) 가지치기 등이 있다 [32, 18, 53, 4, 34, 55, 26]. 
• 현재까지 구조적 가지치기에 대한 연구가 활발하게 진행되었지만, 사용자 입장에서는 어떤 방법을 언제, 어떻게 조합해야 좋은 성능의 가지치기 모델을 얻을 수 있는지 명확하지 않다.
• 가지치기에도 단계가 있는데, 구조 가지치기가 가장 강도높은 가지치기라 생각하면됨

또한, 가지치기는 종종 모델 정확도를 회복하기 위한 일부 재학습을 동반한다 [51]. 하지만 현대 LLM에서는 이러한 재학습 과정이 매우 비용이 많이 들며, 대량의 정제된 데이터를 요구하는 경우가 많다. 현재까지, 구조적 가지치기 연구 중에서 증류(distillation)와 같은 데이터 효율적인 재학습 기법을 활용하여 재학습 비용을 최소화하는 방법을 탐구한 연구는 없다.

이 논문에서는 뉴런 가지치기, 다중 어텐션 헤드 가지치기, 임베딩 채널 가지치기, 모델 깊이 가지치기 등 여러 측면에서 구조적 가지치기와 재학습을 종합적으로 탐구 한다. 실험을 통해 각 가지치기 방식에 적절한 평가 지표(metric) 및 하이퍼파라미터를 파악하고, 가지치기 기법을 효과적으로 조합하여 더 높은 압축률을 달성하는 방법을 분석한다.

예를 들어, 우리는 초기에는 뉴런 및 어텐션 헤드 가지치기가 뉴런, 헤드, 임베딩 채널을 함께 가지치기하는 것보다 더 성능이 좋지만, 몇 번의 재학습이 진행되면 그 순서가 뒤바뀐다는 사실을 발견했다. 

• 또한, 너비(width) 가지치기가 깊이(depth) 가지치기보다 더 효과적이지만, 일부 재학습이 진행된 후에야 그렇다는 점도 확인했다(자세한 수치는 Table 1 참조).
• 더 나아가, 최소한의 추가 데이터를 활용하여 가지치기된 모델을 효율적으로 재학습하는 방법도 자세히 조사했다. 연구 결과를 바탕으로 LLM 압축 및 재학습을 위한 실용적인 가이드라인을 도출했다.

마지막으로, 본 연구의 결과를 실제 모델에 적용하여, Nemotron-4 15B 모델 [43]을 가지치기하고, 새로운 소형 모델 패밀리인 MINITRON을 제작했다. 실험 결과, MINITRON 8B는 Nemotron-3 8B [39]보다 40배 적은 학습 데이터를 사용하면서도 더 높은 정확도를 기록했으며, LLaMA-2 7B [49]보다 높은 성능을 보였다. 또한, MINITRON 8B는 Mistral-7B [25], Gemma 7B [48], LLaMA-3 8B와 비슷한 성능을 달성했다. 한편, MINITRON 4B는 Gemma2 모델을 능가하며, Phi-2 모델과 유사한 성능을 보였다.

이를 통해, 본 연구는 다양한 규모의 LLM을 보다 비용 효율적으로 생성하는 방법을 제시하며, LLM 가지치기 및 재학습의 실용적 가능성을 입증한다.

이 논문은 다음과 같은 주요 기여를 한다:

1. 대형 언어 모델(LLM)의 구조적 가지치기(structured pruning)와 재학습(retraining)에 대한 최초의 종합적인 실증적 연구를 제공한다. 이를 통해 다양한 가지치기 축(뉴런, 어텐션 헤드, 임베딩 채널, 깊이)에 대한 평가 지표(metric) 및 하이퍼파라미터 설정, 가지치기 순서, 여러 가지치기 방법의 조합 효과, 데이터 효율적인 재학습 기법에 대한 유용한 인사이트를 제시한다.

2. 실제 활용 가능한 LLM 압축 및 재학습의 효과적인 베스트 프랙티스(best practices)를 제시한다. 이 가이드는 광범위한 실험적 증거에 기반하여 구성되었다.

3. Nemotron-4 15B 모델을 직접 가지치기하여 생성한 MINITRON 모델 패밀리를 소개한다.

   • MINITRON 모델들은 처음부터 학습하는 것보다 최대 40배 적은 학습 데이터(토큰)를 사용하면서도 높은 성능을 유지한다.
   • 다양한 크기의 주요 오픈소스 LLM들과 비교했을 때 경쟁력 있는 성능을 보이며,
   • 기존 연구에서 제안된 최첨단 깊이 및 너비 가지치기 모델보다 더 우수한 성능을 발휘한다.

2 Pruning Methodology 

Figure 2에서 볼 수 있듯이, 가지치기 과정은 먼저 각 **레이어(layer), 뉴런(neuron), 어텐션 헤드(head), 그리고 임베딩 차원(embedding dimension)**의 중요도를 계산하는 것으로 시작된다. 

이후, 이 중요도 점수를 정렬하여 **중요도 순위(importance ranking)**를 계산한 후, 해당 순위를 기반으로 가지치기된 모델을 생성한다.

---

2.1 배경 및 표기법(Background and Notation)

먼저, 몇 가지 공식적인 정의를 설명한다.

(1) 다층 퍼셉트론(MLP) 레이어

다층 퍼셉트론(MLP, Multi-Layer Perceptron) 레이어는 두 개의 선형(fully connected) 층과 그 사이의 **비선형 활성화 함수(non-linear activation function)**로 구성된다. 수식으로 표현하면 다음과 같다:

$$MLP(X) = δ (X \cdot W_1^T) \cdot W_2$$

여기서,

• $X$ : 입력 데이터
• $W_1$ , $W_2$ : MLP 층의 가중치 행렬
• $W_1, W_2 \in \mathbb{R}^{d_{\text{hidden}} \times d_{\text{model}}}$
  • $d_{\text{model}}$ : 임베딩 차원
  • $d_{\text{hidden}}$ : MLP의 숨겨진(hidden) 차원
• $δ(·)$ : 비선형 활성화 함수

---

(2) 다중 헤드 어텐션(Multi-Head Attention, MHA)

MHA 연산은 입력 $X$에 대해 다음과 같이 정의된다:

$$MHA(X) = \text{Concat}(\text{head}_1, ... , \text{head}_L) \cdot W_O$$

각 개별 어텐션 헤드 $\text{head}_i$는 아래의 어텐션 함수(Attn)로 계산된다:

$$\text{head}_i = \text{Attn}(X W_{Q,i}, X W_{K,i}, X W_{V,i})$$

여기서,

• $W_{Q,i}, W_{K,i}, W_{V,i} \in \mathbb{R}^{d_{\text{head}} \times d_{\text{model}}}$ : Query, Key, Value 가중치 행렬
• $W_O \in \mathbb{R}^{L d_{\text{head}} \times d_{\text{model}}}$ : 출력 변환 가중치 행렬
• $d_{\text{head}}$ : 하나의 어텐션 헤드 크기
• $L$ : 전체 어텐션 헤드 수

---

(3) 층 정규화(Layer Normalization, LayerNorm)

LayerNorm은 입력 $X$에 대해 다음과 같이 정의된다:

$$LN(X) = \frac{X - \mu}{\sqrt{\sigma^2 + \epsilon}} \odot \gamma + \beta$$

여기서,

• $\mu$, $\sigma^2$ : 임베딩 차원에 대한 평균(mean)과 분산(variance)
• $\epsilon$ : 수치 안정성을 위한 작은 값
• $\gamma$, $\beta$ : 학습 가능한 파라미터

---

이러한 수학적 정의를 기반으로, 가지치기 대상이 되는 MLP 뉴런, 어텐션 헤드, 그리고 임베딩 차원을 평가하고 중요도를 계산한 후, 효과적인 가지치기 전략을 설계한다.

2.2 중요도 분석(Importance Analysis)

뉴런, 어텐션 헤드, 레이어와 같은 개별 신경망 구성 요소의 중요도(또는 민감도)를 추정하는 것은 잘 연구된 분야이다 [9, 13, 41]. 특히 LLM(대형 언어 모델)에서는 전통적인 가중치 크기(weight magnitude) 기반 중요도 측정 방법이 비효율적이라는 연구 결과가 있다 [33]. 대신, 최근 연구들은 Gradient/Taylor 기법 [33], 코사인 유사도(cosine similarity) [34], 그리고 검증 데이터셋(perplexity on a calibration dataset)을 활용한 방법 [26]과 같은 새로운 지표를 활용하여 LLM 가지치기를 수행하고 있다.

그러나 최신 LLM은 크기가 너무 커서 경사(gradient) 정보를 직접 계산하는 것은 메모리와 연산 비용 측면에서 매우 부담스럽다. 따라서 본 연구의 주요 목표 중 하나는 이러한 비용이 큰 연산 없이 중요도를 추정하는 방법을 개발하는 것이다. 이를 위해, 우리는 오직 활성화 값(activation)만을 이용한 중요도 추정 전략을 제안한다. 이 방법은 다음과 같은 특징을 가진다:

• 뉴런, 어텐션 헤드, 임베딩 채널, 깊이(레이어) 등 모든 가지치기 축에 대한 중요도를 동시에 계산 가능
• 작은 크기의 검증 데이터셋(1024개 샘플)을 활용하여 중요도를 추정
• 순전파(forward propagation)만을 수행하여 중요도를 계산하므로 계산 비용이 낮음

이제 각 가지치기 축(axis)에 대한 중요도 추정 방식을 설명한다.

---

(1) 너비(Width) 기반 가지치기: 뉴런, 어텐션 헤드, 임베딩 채널 중요도 측정

우리는 MHA(다중 헤드 어텐션), MLP(다층 퍼셉트론), LayerNorm(층 정규화) 층에서 생성된 활성화 값을 분석하여 각 어텐션 헤드, 뉴런, 임베딩 채널의 중요도를 평가한다.

이를 위해 작은 검증 데이터셋 $D$를 사용하며, 각 축에 대한 활성화 기반 중요도 점수는 다음과 같이 정의된다:

• 어텐션 헤드(Head) 중요도

  $$F_{\text{head}}^{(i)} = \sum_{B,S} \| \text{Attn}(X W_{Q,i}, X W_{K,i}, X W_{V,i}) \|_2$$

• 뉴런(Neuron) 중요도

  $$F_{\text{neuron}}^{(i)} = \sum_{B,S} X W_{1,i}^T$$
  • 여기서 $W_{1,i}$는 가중치 행렬 $W_1$의 $i$번째 행을 의미

• 임베딩 채널(Embedding Channel) 중요도

  $$F_{\text{emb}}^{(i)} = \sum_{B,S} \text{LN}(X)_{i}$$

$\sum_{B,S}$ 표기는 배치(batch)와 시퀀스(sequence) 차원을 따라 합산(aggregation)하는 것을 의미한다.

그러나 실험 결과, 단순 합산(sum)을 수행하는 것이 항상 최적의 방법은 아니었다. 따라서 우리는 다양한 **집계 함수(aggregation functions)**를 평가하여 가장 효과적인 방법을 찾았다(결과는 Table 11에 정리됨).

검토한 집계 함수는 다음과 같다:

1. 평균(mean(abs)): $$\frac{1}{n} \sum_{i=1}^{n} |S_i|$$
2. L2 norm: $$\sqrt{\sum_{i=1}^{n} S_i^2}$$
3. 분산(variance): $$\frac{1}{n} \sum_{i=1}^{n} (S_i - \bar{S})^2$$

각 층에서 계산된 중요도 점수들은 합산되어 네트워크 전체의 중요도 점수를 생성한다.

테이블 11에서는, batch는 L2, seq는 mean으로 중요도를 판단하는게 제일 좋다는 뜻인거 같은데?

---

(2) 깊이(Depth) 기반 가지치기: 레이어(Layer) 중요도 측정

깊이 가지치기(depth pruning)의 경우, 각 레이어(layer)의 중요도를 평가하는 두 가지 지표를 사용한다:

1. Perplexity(PPL) 변화 분석 [26]

   • 개별 레이어를 제거한 후, 모델의 혼란도(perplexity) 변화를 측정
   • 혼란도가 크게 증가할수록 해당 레이어의 중요도가 높다고 판단

2. 블록 중요도(Block Importance, BI) [34]

   • 특정 레이어의 입력과 출력을 비교하여 **코사인 거리(cosine distance)**를 기반으로 중요도를 계산
   • 수식: $$BI_i = 1 - \mathbb{E}_{X,t} \left[ \frac{X_{i,t}^T X_{i+1,t}}{\|X_{i,t}\|_2 \|X_{i+1,t}\|_2} \right]$$
     • $X_i$ : 레이어 $i$의 입력
     • $X_{i,t}$ : 입력 행렬의 $t$번째 행

BI 점수는 한 번의 순전파(forward pass)만으로 전체 레이어에 대한 계산이 가능하므로, PPL 기반 방법보다 속도가 빠르다는 장점이 있다. 또한, Gromov et al. [14]의 연구를 참고하여 여러 개의 연속된 레이어를 동시에 평가하는 방법도 적용할 수 있다.

BI는 입출력의 코사인 유사도가 높으면 (즉 유사하면) 중요도가 낮다는 것 같은데? 아마 벡터값(정보값)을 많이 바꿔서 중요하다고 판단한듯

---

(3) 반복적 중요도 분석(Iterative Importance Analysis)

이 방법에서는 가지치기와 중요도 분석을 여러 번 반복(iterative)하여 최적의 압축을 수행한다.

1. 초기 모델에서 중요도를 계산
2. 가장 중요도가 낮은 요소(레이어, 헤드, 뉴런 등)를 일부 제거
3. 가지치기된 모델에서 다시 중요도를 계산
4. 위 과정을 반복하며 점진적으로 가지치기 수행

이를 공식화하면, T번의 반복(iteration)을 수행한다고 할 때:

• 시작 차원 $d_s$, 목표 차원 $d_t$가 주어졌을 때,
• $i$번째 반복에서 가지치기할 차원 수: $$d_s - i \cdot \frac{d_s - d_t}{T}$$
• $i+1$번째 반복 후 남은 차원 수: $$d_s - (i+1) \cdot \frac{d_s - d_t}{T}$$
• $i \in [0, T-1]$

이를 통해, 한 번에 많은 양을 가지치기하는 것이 아니라 점진적으로 줄여나가는 방식을 적용할 수 있다.

조금씩 가지치기 한다는 의미

---

결론

• 우리는 비용이 큰 경사 기반 방법(gradient-based method) 없이, 활성화 값만을 이용해 가지치기 대상의 중요도를 평가하는 방법을 제안하였다.
• 너비(뉴런, 어텐션 헤드, 임베딩 채널) 가지치기는 활성화 값과 다양한 집계 함수(mean, L2 norm, variance)를 사용하여 평가할 수 있다.
• 깊이(레이어) 가지치기는 Perplexity 변화 분석 및 Block Importance(BI) 점수를 활용하여 평가할 수 있다.
• **반복적 중요도 분석(iterative importance analysis)**을 적용하면 가지치기를 더욱 효과적으로 수행할 수 있다.

2.3 가지치기된 모델 생성 (Obtaining a Pruned Model)

Figure 2는 가지치기된(Pruned) 모델을 생성하는 전체 과정을 보여준다. 특정 모델 구조(architecture) 설정이 주어지면, 각 축(axis)에 대해 중요도를 평가하여 요소들을 정렬한 후, 해당하는 가중치 행렬(weight matrix)을 직접 조정(reshaping)하여 모델을 가지치기한다.

• 뉴런(Neuron) 및 어텐션 헤드(Head) 가지치기

  • 뉴런 가지치기 시 MLP 층의 가중치(Weights)를 조정
  • 어텐션 헤드 가지치기 시 MHA(다중 헤드 어텐션) 층의 가중치를 조정

• 임베딩 채널(Embedding Channel) 가지치기

  • MLP, MHA, LayerNorm 층에서 임베딩 차원을 줄이는 방식으로 적용

---

(1) 어텐션 헤드 가지치기: 정보 보존 기법

어텐션 헤드를 가지치기할 때, 잘려나간(pruned) 헤드의 정보를 남은 헤드에 다시 추가하는 방식을 적용하여 가지치기된 모델의 성능을 유지한다.
이 방법은 Layer Collapse [55] 개념을 어텐션 헤드 구조에 적용한 것으로, 실험 결과 모델의 정확도를 향상시키는 효과가 있었다.

어텐션 헤드 재분배 공식

$L$개의 원래 어텐션 헤드(head1, head2, ..., headL)를 $K$개로 가지치기하는 경우,
각 새로운 헤드는 다음과 같이 변형된다:

$$\text{head}_i = \text{head}_i + (\text{head}_i - \text{head}_{2K-i+1}) \quad \text{for} \quad i \in [K - (L - K), K]$$

즉, 일부 제거된 헤드의 정보를 남아 있는 헤드에 더해주는 방식을 사용한다.

여기서 보면, 기존 헤드를 그대로 가져가는게 아니라, 사라지는 헤드와의 값차이를 더해준다는 것 (이 값차이가 중요한 정보라고 가정하는 듯)

• Grouped Query Attention (GQA) [3] 구조를 사용할 경우,
  • 이 전략을 쿼리(query) 헤드에만 적용하여 가지치기한다.

---

(2) 경량 신경망 아키텍처 탐색 (Lightweight Neural Architecture Search)

Figure 3은 최적의 모델 구조를 찾는 탐색(search) 전략을 보여준다.

1. 탐색 공간(search space)과 파라미터 예산(parameter budget)이 주어지면, 가능한 모든 모델 구조를 나열(enumerate)한다.

   • 예를 들어, 주어진 예산 내에서 뉴런, 어텐션 헤드, 임베딩 차원의 조합을 탐색

2. 탐색 공간을 줄이는 기법(유전 알고리즘, 베이지안 최적화 등)을 사용할 수도 있으나, 일반적으로 사용되는 뉴런/헤드/임베딩 차원을 기반으로 하면 탐색 공간을 효과적으로 줄일 수 있다.

   • 1억 개 이하의 파라미터 범위를 설정하면 후보 모델 개수를 20개 이하로 제한 가능

3. 후보 모델들은 가벼운 재학습(lightweight retraining)을 수행한다.

   • **약 18억 개의 토큰(∼1.8B tokens)**을 사용하여 빠르게 재학습
   • Figure 9에서 확인할 수 있듯이, 이 재학습 과정이 상대적인 모델 성능을 안정화시키고, 최적의 모델을 추가 학습할 수 있도록 도와준다.

4. **LoRA (Low-Rank Adaptation) [23]**와 같은 파라미터 효율적 미세 조정(parameter-efficient fine-tuning) 기법을 추가적으로 적용할 수도 있다.

• 본 연구에서는 이를 다루지 않았으며, 향후 연구에서 탐색할 예정

즉 여러 방식으로 모델 사이즈를 줄일 수 있는데, 그 방법을 위처럼 탐색하여 후보들을 추린다는 것

추린 방법으로 중요도 계산하여 프루닝하고, 이를 가볍게 재학습하여 최적의 모델을 찾는다

---

결론

• 가지치기된 모델을 생성하는 과정에서 뉴런, 어텐션 헤드, 임베딩 차원을 조정한다.
• 어텐션 헤드 가지치기 시, 제거된 헤드의 정보를 남은 헤드에 추가하는 방식으로 성능을 유지한다.
• 최적의 모델 구조를 찾기 위해 경량 신경망 탐색(Neural Architecture Search)을 수행하며, 빠른 재학습을 통해 모델 성능을 안정화한다.
• LoRA와 같은 파라미터 효율적 미세 조정 기법을 추가적으로 적용할 가능성이 있다.

3. 재학습 (Retraining)

**재학습(Retraining)**이란, 가지치기 후 성능 저하를 보정하여 모델 정확도를 회복하는 과정을 의미한다.
본 논문에서는 두 가지 재학습 전략을 탐색한다:

1. 일반적인 재학습 (Conventional Training)
   • 원본 데이터셋과 실제 정답(ground truth label)을 사용하여 모델을 다시 학습하는 방식.
2. 지식 증류(Knowledge Distillation, KD)
   • 가지치기 전 모델(teacher)에서 **출력 확률 분포 및 중간 표현(intermediate representations)**을 학습하여, 가지치기된 모델(student)이 이를 모방하도록 유도하는 방식.

---

1. 지식 증류(Knowledge Distillation, KD) 기반 재학습

**지식 증류(KD)**는 더 크거나 복잡한 모델(teacher)에서 **더 작은 모델(student)로 지식을 전이(transfer)**하는 과정이다 [20].
이 연구에서는 원래의 LLM(가지치기 전 모델)을 teacher, 가지치기된 모델을 student로 두고, student 모델이 teacher의 출력을 모방하도록 학습한다.

---

2. 출력 확률 분포 기반 지식 증류

LLM이 특정 입력 토큰 $x_i$에 대해 생성하는 출력 확률 분포는 아래와 같이 계산된다:

$$p(x_i, \tau) = \frac{\exp\left(\frac{x_i}{\tau}\right)}{\sum_{j=1}^{|V|} \exp\left(\frac{x_j}{\tau}\right)}$$

여기서,

• $\tau$ = 소프트맥스 온도(softmax temperature)
  • 온도가 높을수록 확률 분포가 부드러워지고, 낮을수록 확률이 극단적으로 변함.
• $|V|$ = 어휘 집합 크기(vocabulary size)

즉, 가지치기된 student 모델이 teacher 모델의 확률 분포를 학습하도록 유도함.

이를 위한 손실 함수(로짓 기반 KD 손실, $L_{\text{logits}}$)는 다음과 같다:

$$L_{\text{logits}} = \frac{1}{l} \sum_{k=1}^{l} \text{Loss} \left( p_k^t(x, \tau), p_k^s(x, \tau) \right)$$

• $p_k^t(x, \tau)$ = teacher 모델의 $k$번째 토큰 확률 분포
• $p_k^s(x, \tau)$ = student 모델의 $k$번째 토큰 확률 분포
• $l$ = 시퀀스 길이

즉, teacher와 student의 출력 확률 분포 차이를 최소화하는 방식으로 학습이 진행된다.

---

3. 중간 상태(hidden states) 기반 지식 증류

지식 증류는 단순히 출력 확률 분포만 모방하는 것뿐만 아니라, 중간 숨겨진 표현(hidden state)도 모방할 수 있다.

이때 사용하는 중간 상태(hidden state) 기반 KD 손실은 다음과 같다:

$$L_{\text{is}} = \frac{1}{l} \sum_{k \in H} \sum_{i=1}^{l} \text{Loss}_k(h_{ki}^t, h_{ki}^s)$$

• $h_{ki}^t$ = teacher 모델의 $k$번째 hidden state, $i$번째 토큰에 대한 값
• $h_{ki}^s$ = student 모델의 $k$번째 hidden state, $i$번째 토큰에 대한 값
• $H$ = 선택된 hidden state 집합

즉, 특정 Transformer 계층에서 teacher 모델과 student 모델의 hidden state가 유사하도록 학습하는 방식이다.

---

4. Hidden State 차원을 맞추는 보정 기법

• 가지치기된 모델(student)의 hidden state 크기와 teacher 모델의 hidden state 크기가 다를 경우,
  • 학습 가능한 선형 변환(linear transformation)을 도입하여 student 모델의 hidden state를 teacher 모델의 hidden state 차원으로 변환함.
  • 모든 hidden state는 LayerNorm 이후의 값을 사용.

---

5. 최종 손실 함수(Loss Function)

전체 학습 과정에서 최종 손실 함수는 아래와 같이 정의된다:

$$L = L_{\text{CLM}} + L_{\text{logits}} + \alpha \times L_{\text{is}}$$

• $L_{\text{CLM}}$ = student 모델이 ground truth 정답과 비교하여 학습하는 교차 엔트로피 손실(cross-entropy loss)
• $L_{\text{logits}}$ = teacher 모델의 출력 확률 분포를 모방하는 KD 손실
• $L_{\text{is}}$ = teacher 모델의 중간 hidden state를 모방하는 KD 손실
• $\alpha$ = KD 손실의 가중치 계수(weighting coefficient)

Loss을 KD만 쓴것인줄 알았는데 그건 아니고 모든 가능한 loss을 계산하여 활용했군

이에 대한 내용은 4.1에 나옴

가중치 계수( $\alpha$ ) 동적 조정

• 일반적으로 $L_{\text{logits}}$와 $L_{\text{is}}$는 손실 크기가 다르므로,
  • $\alpha$를 고정값이 아니라 동적으로 조정하는 방법을 사용하면 성능이 더 좋아짐.
  • 논문에서는 $\alpha = \frac{L_{\text{logits}}}{L_{\text{is}}}$로 설정하여,
    • 두 손실이 비슷한 크기를 유지하도록 자동 조정.

---

결론

• 가지치기 후 모델(student)의 정확도를 회복하기 위해 지식 증류(KD) 기법을 활용함.
• teacher(가지치기 전 모델)와 student(가지치기 후 모델)의 출력 확률 분포와 hidden state 차이를 줄이는 방식으로 학습.
• 출력 확률 기반 KD 손실( $L_{\text{logits}}$ )과 hidden state 기반 KD 손실( $L_{\text{is}}$ )을 함께 사용하여 성능을 개선.
• 손실 함수의 가중치( $\alpha$ )는 동적으로 조정하여 더 안정적인 학습을 유도.

이를 통해, 모델을 압축(가지치기)하더라도 최대한 성능을 보존하는 방식으로 재학습을 진행할 수 있음! 🚀

4. 실험 및 결과 (Experiments and Results)

본 연구에서는 Nemotron-4 모델 패밀리 [43]에 대한 가지치기 전략을 평가하며,
Nemotron-4 15B 모델(총 156억 개 파라미터)를 두 가지 크기로 압축한다:

1. 8B (80억 개 파라미터)
2. 4B (40억 개 파라미터)

모델 압축 및 재학습을 위해 NVIDIA Megatron-LM 프레임워크 [47]를 사용하여 가지치기 및 지식 증류(KD) 알고리즘을 구현하였다.

---

1. 데이터 및 학습 하이퍼파라미터

• 학습 데이터:

  • Nemotron-4 사전학습 데이터셋(총 8조 토큰, 8T) + 추가 학습 데이터(CT) [42, 44, 43]
  • **가지치기 후 경량 재학습(1.8B 토큰, 400 스텝)**을 수행하여 모델을 안정화
  • 최종 모델 재학습 시에는 8T 학습 데이터와 추가 CT 데이터를 혼합하여 사용

• 중요도 추정(calibration dataset) 데이터:

  • 전체 데이터셋에서 랜덤 샘플링된 1,024개 샘플을 사용하여 가지치기 대상 중요도를 평가

• 최적화 및 학습 설정:

  • 기존 Nemotron-4 모델과 동일한 최적화 설정 및 데이터 분할 전략을 사용 [43]
  • Cosine learning rate decay 사용,
    • 초기 학습률: $2^{-4}$
    • 최종 학습률: $4.5^{-7}$

---

2. 다운스트림 평가 (Downstream Tasks)

Nemotron-4 8B 및 4B 모델을 기존 크기가 비슷한 모델들과 비교하여 성능을 평가하였다.
평가 기준은 **Touvron et al. [49]**의 방식에 따라 다음과 같은 다운스트림 작업을 포함한다:

(1) 일반 벤치마크

• MMLU (5-shot, 대규모 다중태스크 학습 평가) [19]
• HumanEval (Python 코드 생성) [8]
• 상식 추론 및 질의응답(QA) 데이터셋
  • Arc-C (25-shot, 추론 문제 해결) [10]
  • HellaSwag (10-shot, 상식적 텍스트 완성) [56]
  • TruthfulQA (진실성 기반 QA) [29]
  • WinoGrande (5-shot, 대명사 해석) [45]
  • XL-Sum (0-shot, 영어 텍스트 요약) [17]

(2) 코드 생성 성능 평가 (pass@1 Score)

• HumanEval 및 MBPP 데이터셋 사용
  • pass@1 스코어 계산
  • 샘플링 설정: 온도 0.2, top-p = 0.95 (nucleus sampling) [22]

(3) 명령어 기반 튜닝 모델 평가

• MT-Bench (Instruction-Tuned LLM 성능 평가) [57]
• Instruction-Following Eval (IFEval) [59]
• ChatRAG-Bench (대화형 검색 및 응답 평가) [30]
• Berkeley Function Calling Leaderboard (BFCL) [54]

---

정리

• Nemotron-4 15B 모델을 8B 및 4B로 압축한 후, 다양한 다운스트림 작업에서 성능을 평가함.
• **8조 토큰(8T) 데이터셋과 추가 학습 데이터(CT)**를 활용하여 모델을 사전학습 및 재학습 진행.
• 1.8B 토큰(400 스텝) 규모의 경량 재학습을 수행하여 가지치기 후 성능 회복.
• 다양한 MMLU, 코드 생성, QA, 요약 및 명령어 기반 평가 벤치마크에서 실험을 진행하여 압축된 모델의 성능을 검증.

결과적으로, 가지치기된 모델이 기존 모델 대비 성능을 유지하면서도 더 효율적인 크기로 동작하는지 평가하는 실험을 수행하였다. 🚀

4.1 주요 가지치기 결과 (Main Pruning Results)

1. 구조적 압축을 위한 베스트 프랙티스 (Best Practices for Structured Compression)

본 연구를 통해 다음과 같은 효과적인 LLM 가지치기 및 압축 전략을 도출하였다:

1. LLM 패밀리를 학습할 때, 가장 큰 모델을 먼저 학습한 후, 가지치기+지식 증류를 반복적으로 수행하여 작은 모델을 생성하는 것이 최적.
2. 너비(width) 축에서 중요도 추정 시, 배치 방향에서는 L2 norm을 사용하고, 시퀀스 방향에서는 평균(mean)을 사용. 깊이(depth) 축에서는 PPL(perplexity) 또는 BI(Block Importance) 지표를 활용.
3. 중요도 추정을 반복적으로(iterative) 수행하는 것은 별다른 이점을 제공하지 않으므로, 단일 평가(single-shot importance estimation)를 사용하는 것이 효율적.
4. 15B 이하의 모델 크기에서는 깊이(depth) 가지치기보다 너비(width) 가지치기가 더 효과적.
5. 일반적인 학습 대신, KLD(Kullback-Leibler Divergence)를 이용한 지식 증류 손실을 사용하여 재학습하는 것이 가장 효과적.
6. 깊이(depth)를 크게 줄이는 경우, 로짓(logit) + 중간 상태(hidden state) + 임베딩(embedding) 정보를 포함한 지식 증류를 사용.
7. 깊이를 크게 줄이지 않는 경우, 로짓(logit) 정보만을 활용한 지식 증류를 사용하는 것이 더 효과적.
8. 가지치기할 때, 목표 크기에 가장 가까운 모델을 선택하여 가지치기하는 것이 최적의 성능을 제공.
9. 검색된 가지치기 모델 후보들의 중요도 순위를 안정화하기 위해, 경량 재학습(lightweight retraining)을 수행.
10. 가장 큰 모델이 다단계 학습(multi-phase training) 전략으로 학습되었다면, 최종 학습 단계의 모델을 가지치기 및 재학습하는 것이 가장 효과적.

위의 베스트 프랙티스는 다양한 실험과 검증을 통해 도출되었으며, 이후 실험 결과에서 이를 뒷받침하는 데이터를 제시한다.
이러한 전략을 활용하여 MINITRON 모델을 생성하고 재학습하였으며, 그 성능은 Tables 2~3에서 확인할 수 있다.

---

2. MINITRON 모델과 기존 모델 비교

MINITRON 모델 성능을 다음과 같은 기준 모델과 비교하였다:

1. 원본 Nemotron-4 15B 모델
2. 이전 세대 Nemotron-3 8B 모델
3. 유사한 크기의 커뮤니티 모델들 (LLaMa-2 7B, Mistral 7B, Gemma 7B, LLaMa-3 8B 등)
   • 이 모델들은 수조 개(trillions) 이상의 토큰으로 처음부터 학습된 모델들임.
4. 깊이 및 너비 가지치기 기반 최신 연구 모델들과 비교 (Table 4)
• LLM-Pruner [33], SliceGPT [4], LaCo [55], ShortGPT [34], Sheared LLaMa [53]

MINITRON 8B 결과 분석 (Table 2)

• Nemotron-3 8B 및 LLaMa-2 7B보다 뛰어난 성능을 보임.
• Mistral 7B, Gemma 7B, LLaMa-3 8B와 비슷한 성능을 달성하면서도, 훨씬 적은 학습 토큰을 사용하여 모델을 생성.
• 10B 파라미터 규모의 깊이 가지치기 모델들보다 높은 성능을 보임 (Table 4 참고).

MINITRON 4B 결과 분석 (Table 3)

• 특정 작업에서만 강한 소형 특화 모델들보다 전체적으로 더 균형 잡힌 성능을 유지.
• Gemma 2B 모델을 능가하며, 깊이 및 너비 가지치기된 기존 연구 모델들(Table 4)보다 뛰어난 성능을 기록.

---

3. 명령어 튜닝 (Instruction Tuning) 및 추가 평가

• MINITRON 4B를 추가로 명령어 튜닝(supervised fine-tuning, SFT)하여 MINITRON 4B-instruct 모델을 생성.

• Nemotron-4 340B 모델이 사용한 명령어 튜닝 데이터로 학습하여 성능을 평가.

• 평가 기준:

  1. 명령어 이해 및 대화 능력 → IFEval, MT-Bench
  2. RAG 기반 질의응답 → ChatRAG-Bench
  3. 함수 호출(function calling) 성능 → BFCL

• MINITRON 4B-instruct 결과(Table 5~8)

• 유사한 크기의 모든 모델보다 뛰어난 명령어 처리 및 대화 성능을 보임.
• 함수 호출 평가(BFCL)에서 Gemma-2B-IT 및 LLaMa-3-8B-instruct를 능가.

---

4. 비용 절감 효과 (Cost Savings for Training a Model Family)

FLOP 계산을 통한 비용 절감 분석

• Nemotron-4 모델 패밀리에서 15B, 8B, 4B 모델의 한 스텝당 FLOPs:

  • 15B 모델: $4.4 \times 10^{17}$ FLOPs
  • 8B 모델: $2.5 \times 10^{17}$ FLOPs
  • 4B 모델: $1.2 \times 10^{17}$ FLOPs

• 각 모델을 처음부터 학습할 경우 총 FLOPs:

  $$(4.4 \times 10^{17} + 2.5 \times 10^{17} + 1.2 \times 10^{17}) \times \text{steps}$$

• 하지만, 본 연구에서 제안한 방식(가지치기 + 재학습)으로 모델을 생성하면, 추가 모델을 만들 때마다 필요한 학습 토큰 수가 40배 줄어듦.

  • 즉, 새로운 모델을 만들 때 필요한 FLOPs가 **$2.5 \times 10^{17} / 40$ 및 $1.2 \times 10^{17} / 40$**로 감소.
  • 최종적으로, Nemotron-4 모델 패밀리를 학습하는 데 필요한 전체 FLOPs가 1.8배 절감됨.

비용 절감의 의미

• 기존 방식 (각 모델을 처음부터 학습하는 방식)
  • 각 모델을 독립적으로 학습해야 하므로, 데이터 비용과 계산량이 매우 큼.
• 본 연구의 방식 (가장 큰 모델을 학습한 후 가지치기 + 재학습으로 작은 모델을 생성)
  • 추가 모델당 40배 적은 학습 토큰을 사용하므로 비용과 계산량이 크게 절감됨.
  • 전체적으로 모델 학습 비용을 1.8배 줄일 수 있음.

---

5. 결론

• 15B 모델을 처음 학습한 후 가지치기 + 재학습을 통해 8B 및 4B 모델을 생성하는 것이, 처음부터 각각의 모델을 학습하는 것보다 훨씬 효율적.
• 가지치기된 MINITRON 모델들이 기존 연구의 깊이/너비 가지치기된 모델들보다 더 나은 성능을 보이며, 커뮤니티 모델들과도 경쟁력 있는 성능을 기록.
• 추가적으로 수행한 명령어 튜닝(MINITRON 4B-instruct) 실험에서도 뛰어난 성능을 발휘.
• 학습 비용 측면에서도 최대 1.8배 절감 가능.

따라서, 대형 언어 모델을 구축하는 데 있어 "하나의 큰 모델을 학습한 후 가지치기+재학습하는 전략"이 매우 효과적임을 입증하였다. 🚀

4.2 최적의 가지치기 모델 선정 (Obtaining the Best Pruned Model)

최적의 가지치기 모델을 얻기 위해 여러 실험을 수행하였으며, 그 결과 가장 효과적인 가지치기 전략을 도출하였다.

---

1. 최적의 중요도 집계 방법 (Best Aggregation Metric, Best Practice #2)

활성화 기반 가지치기(activation-based pruning)에서 어떤 집계 방식(aggregation function)을 사용하는 것이 가장 효과적인지 분석하였다.
(자세한 내용은 Section 2.2 중요도 분석(Importance Analysis) 참고)

실험 설정

• Nemotron-4 15B → Nemotron-3 8B 아키텍처로 가지치기 (재학습 없음).
• 다양한 집계 함수(aggregation function)를 적용한 후, Zero-shot LM Loss 및 Wikitext2 perplexity(WPPL) 평가.
• 사용한 집계 함수 조합:
  1. (batch = L2 norm, sequence = mean)
  2. (batch = mean, sequence = mean)
  3. (batch = L2 norm, sequence = L2 norm)
  4. 기타 조합

실험 결과 (Table 11 참고)

• 집계 방법에 따라 Zero-shot 성능이 상당히 다르게 나타남 → 즉, 가지치기할 때 적절한 집계 함수를 선택하는 것이 매우 중요함.
• (batch = L2, sequence = mean) 및 (mean, mean) 조합이 가장 좋은 성능을 보임.
• 이후 논문에서는 (L2, mean) 방법을 기본 집계 함수로 사용 (8T 데이터셋에서 더 나은 성능을 보였기 때문).

추가 검증 실험

• (L2, mean) 조합 vs. (L2, L2) 조합 비교
• 두 가지 방법으로 가지치기한 후, 1.8B 토큰(400 스텝) 재학습 수행
• Figure 5 결과:
  • 재학습 후에도 (L2, mean) 조합이 계속해서 더 나은 성능을 유지.
  • 즉, 가지치기 후에도 최적의 집계 함수 선택이 중요함을 입증.

---

2. 반복적 중요도 분석(Iterative Importance, Best Practice #3)

반복적으로(Iterative) 중요도를 분석하는 것이 실제로 효과적인지 평가하였다.
(자세한 내용은 Section 2.2 중요도 분석(Importance Analysis) 참고)

실험 설정

• Nemotron-4 15B 모델에서 임베딩 차원(embedding dimension)만 가지치기
• 목표 차원: 4096
• 반복 횟수(T) 설정:
  1. T = 1회(단일 가지치기)
  2. T = 2회 (두 번에 걸쳐 점진적으로 가지치기)
  3. T = 4회 (네 번에 걸쳐 점진적으로 가지치기)
• 이후 모든 가지치기 모델을 1.8B 토큰(400 스텝) 동안 경량 재학습 수행

실험 결과 (Table 12 참고)

• 가지치기 직후(iteration 전)는 반복 가지치기가 조금 더 좋은 성능을 보였음.
• 그러나 1.8B 토큰 재학습 후 모든 모델이 동일한 Loss 값으로 수렴.
• 즉, 반복적(Iterative) 가지치기는 재학습 이후 별다른 이점을 제공하지 않음.
• 따라서 단일-shot 중요도 분석이 가장 효율적임 → Best Practice #3로 선정.

---

3. 너비 vs. 깊이 가지치기 비교 (Combining Depth and Width, Best Practice #4)

너비(width)와 깊이(depth) 가지치기 중 어떤 방법이 더 효과적인지 비교 실험 수행.
(자세한 내용은 Section 2.2 중요도 분석(Importance Analysis) 참고)

실험 설정

• PPL(Perplexity) 및 BI(Block Importance) 지표를 사용하여, Nemotron-4 15B 모델에서 가장 중요도가 낮은 16개 레이어 제거
• 가지치기 모델 3개 생성:
  1. 깊이(Depth) 가지치기:
     • 중요도가 낮은 16개 레이어 제거 (BI/PPL 기준)
  2. 너비(Width) 가지치기:
     • 뉴런, 어텐션 헤드, 임베딩 차원을 가지치기하여 Nemotron-3 8B 크기로 압축
  3. 혼합(Depth + Width) 가지치기:
     • 중요도가 낮은 4개 레이어 제거 + 너비 가지치기

실험 결과 (Table 13 참고)

• 가지치기 직후에는 깊이+너비 혼합 가지치기가 가장 낮은 Loss 값을 기록.
• 그러나 0.8B 토큰(200 스텝) 재학습 후, 성능이 역전됨 (Table 1 및 Figure 6 참고).
• 최종적으로 너비(Width) 가지치기가 가장 안정적인 성능을 유지.
• 즉, 너비 가지치기가 깊이 가지치기보다 더 효과적이며, 혼합 방법도 깊이 가지치기가 많이 포함될 경우 성능이 저하될 가능성이 있음.
• 따라서 15B 이하 모델에서는 너비 가지치기를 우선적으로 고려하는 것이 최적의 방법 → Best Practice #4로 선정.

---

4. 결론

• 최적의 가지치기 모델을 얻기 위해서는 다음과 같은 원칙을 따르는 것이 중요함.
  1. (L2 norm, mean) 집계 방식을 사용하면 가지치기 후에도 성능을 유지할 가능성이 높음.
  2. 반복적인 중요도 분석(iterative importance)은 추가적인 이점이 없으므로, 단일-shot 방식이 가장 효율적임.
  3. 너비(width) 가지치기가 깊이(depth) 가지치기보다 효과적이며, 혼합 가지치기의 경우 깊이 가지치기를 최소화하는 것이 바람직함.

이를 바탕으로 Nemotron-4 15B 모델을 최적의 방법으로 가지치기하여 MINITRON 모델을 생성하였다. 🚀

4.3 재학습 및 탐색 (Retraining and Search)

1. 지식 증류 vs. 일반 학습 비교 (Best Practice #5)

지식 증류(Knowledge Distillation, KD)가 가지치기 후 모델 성능 회복에 얼마나 효과적인지 검증하기 위해,
다음 세 가지 방법으로 Nemotron-4 15B → 4B 모델을 생성하여 비교 실험을 진행하였다:

1. 랜덤 초기화된 4B 모델 학습 (4B-Random-Init)
   • 랜덤 초기화된 4B 모델을 처음부터 새롭게 학습
2. 15B → 4B로 가지치기 후, 일반적인 재학습 수행 (4B-Pruned)
3. 15B → 4B로 가지치기 후, 지식 증류(KD) 기반 재학습 수행 (4B-Pruned-Distill)
   • 가지치기 전 15B 모델을 teacher로 사용하여 student(4B) 모델을 학습

실험 결과 (Table 14 참고)

• MMLU 벤치마크:
  • (3) 지식 증류를 적용한 4B 모델이 압도적으로 높은 성능을 보임.
  • 반면, (1) 랜덤 초기화 및 (2) 일반적인 재학습을 적용한 모델은 성능이 거의 랜덤 수준(random guessing)에 가까움.
• HellaSwag 벤치마크:
  • (3) > (2) > (1) 순으로 성능이 향상됨.
• 즉, 가지치기 후 일반적인 재학습보다는 지식 증류를 활용하는 것이 훨씬 효과적임을 확인.
• 가지치기도 중요하고 distill도 중요하다는 표인듯

---

2. 최적의 손실 함수 선택 (Best Practice #5)

로짓 기반 손실( $L_{\text{logits}}$ )을 계산할 때, 다양한 손실 함수 비교

실험 설정

• Kullback-Leibler divergence (KLD)
• Mean Squared Error (MSE)
• Cosine Similarity
• Reverse KLD (R-KLD)

실험 결과 (Tables 15~16 참고)

• 가지치기된 모델을 학습할 때, KLD 손실이 가장 효과적임.
• R-KLD는 Instruction-Tuning(SFT)에서는 유용하지만, 기본 모델을 학습할 때는 KLD가 더 우수함.
• 따라서 가지치기 후 모델 재학습에는 KLD를 기본적으로 사용.

---

3. 추가적인 손실 함수 조합 연구 (Best Practices #6 & #7)

지식 증류 손실 $L_{\text{is}}$ (Intermediate States) 적용 여부 실험

• 손실 함수 조합: $$L_{\text{is}} = L_{\text{emb}} + L_{\text{att}} + L_{i} + L_{o}$$
• 이전 연구(Lu et al. [31])를 참고하여 실험 수행

실험 결과 (Tables 17~18 참고)

• 깊이(depth)를 크게 가지치기한 경우(예: 레이어 수 감소), $L_{\text{is}}$ 를 추가하는 것이 효과적.
• 그러나 깊이를 크게 줄이지 않은 모델(예: MINITRON 8B, MINITRON 4B)에서는 $L_{\text{logits}}$ 만을 사용한 경우와 성능 차이가 없음.
• 따라서, 깊이를 줄이지 않는 모델은 $L_{\text{logits}}$ 만을 사용하여 학습하는 것이 최적.

어떻게 모델을 줄이냐에 따라 조금씩 다른듯

---

4. 한 번에 가지치기 vs. 단계적 가지치기 (Best Practice #8)

Nemotron-4 15B → MINITRON 4B 변환 시, 한 번에 가지치기(one-shot) vs. 단계적 가지치기(iterative pruning) 비교

실험 설정

• One-shot pruning:
  • 한 번에 73.3%의 가중치를 제거하여 15B → 4B로 변환
• Iterative pruning:
  • 먼저 15B → 8B로 가지치기 후, 다시 8B → 4B로 가지치기
  • 1단계: 15B → 8B (약 46% 가지치기)
  • 2단계: 8B → 4B (약 50% 가지치기)

실험 결과 (Table 14, 마지막 두 줄 참고)

• Iterative Pruning이 One-shot Pruning보다 MMLU 성능이 12% 더 높음.
• 한 번에 모델을 크게 가지치기하면, 원래 모델의 중요한 능력을 상실할 가능성이 높음.
• 단계적으로 가지치기를 수행하면 성능 저하 없이 효과적으로 모델을 압축할 수 있음.
• 따라서, 강한 가지치기가 필요한 경우(예: 15B → 4B)에는 단계적 가지치기를 적용하는 것이 최적.

---

5. 재학습을 포함한 탐색 전략 (Best Practice #9)

경량 신경망 아키텍처 탐색(NAS)과 재학습을 결합하여 최적의 모델 구조를 선정

실험 설정

• MINITRON 8B 및 4B 모델에 대해 가능한 후보 아키텍처 탐색
• 8B 및 4B 모델의 파라미터 범위를 각각 ±5% 허용
• 가능한 후보 모델 수:
  • 8B 모델 → 15개 후보
  • 4B 모델 → 18개 후보
• 모든 후보 모델을 경량 재학습 수행 (Section 2.3에서 설명한 방식 적용)

실험 결과 (Figure 9 참고)

• 재학습 진행 중 약 300 스텝까지 성능 순위가 변화하고 이후 안정화됨.
• 따라서, 최적의 후보 모델을 선택하려면 300 스텝 정도의 경량 재학습이 필요.

---

6. 단일 단계 vs. 다단계 재학습 (Best Practice #10)

웹 데이터와 정제된 데이터로 이루어진 다단계 학습 과정에서 어느 시점에서 가지치기+재학습을 수행하는 것이 가장 좋은지 비교

실험 설정

• 방법 1:
  • 1단계 학습 체크포인트(웹 데이터 기반)를 가지치기 후, 1단계+2단계 데이터를 혼합하여 재학습.
• 방법 2:
  • 2단계 체크포인트(정제된 데이터 기반)를 가지치기 후, 2단계 데이터만 사용하여 재학습.

실험 결과 (Table 20 참고)

• 2단계 체크포인트를 가지치기하고 2단계 데이터로만 재학습하는 것이 성능이 더 좋음.
• 즉, 다단계 학습을 거친 모델의 최종 단계에서 가지치기를 수행하는 것이 가장 효과적.
• 이 방법은 향후 추가적인 모델 정렬(alignment)에도 적용 가능.

---

결론

1. 가지치기 후 재학습은 일반적인 학습보다 지식 증류(KD)를 활용하는 것이 훨씬 효과적임.
2. 손실 함수로는 KLD가 가장 적합하며, 깊이를 크게 줄인 경우에는 추가적인 hidden state 손실을 고려할 수 있음.
3. 한 번에 크게 가지치기(one-shot)하는 것보다, 단계적 가지치기(iterative pruning)가 성능 유지에 유리함.
4. 모델 탐색 과정에서 경량 재학습을 수행하면 최적의 구조를 선정할 수 있음.
5. 다단계 학습을 수행하는 경우, 최종 단계의 체크포인트에서 가지치기를 수행하는 것이 가장 좋은 성능을 제공함.

이러한 전략을 통해 Nemotron-4 15B → MINITRON 8B 및 4B 모델을 최적의 방식으로 압축하고 성능을 유지할 수 있었음. 🚀

5. 관련 연구 (Related Work)

구조적 LLM 가지치기 (Structured LLM Pruning)

최근 LLM(대형 언어 모델)을 대상으로 한 구조적 가지치기(structured pruning) 연구가 활발하게 이루어지고 있으며, 이를 크게 두 가지 범주로 분류할 수 있다:

1. 깊이(depth)만 가지치기하는 연구(레이어 수 감소)
2. 너비(width, 예: 어텐션 헤드, MLP 중간 차원 등)를 가지치기하거나 너비+깊이를 함께 가지치기하는 연구

(1) 깊이 가지치기 (Depth Pruning)

• 최근 연구 사례: ShortGPT [34], LaCo [55], Shortened LLaMa [26]
• 본 연구에서는 MINITRON 모델의 레이어 가지치기를 위해 위 연구에서 제안한 블록 중요도(block importance) 지표 등을 확장하여 사용함.

(2) 너비 가지치기 (Width Pruning)

• 최근 연구들은 **임베딩 채널(embedding channels), 어텐션 헤드(attention heads), MLP 중간 채널(intermediate channels)**을 대상으로 한 중요도 지표 및 가지치기 기법을 제안함 [11, 4, 53, 33].
• 이 범주의 연구 대부분은 학습 가능한 마스크(learnable masks)와 Augmented Lagrangian 손실을 결합하여 최적의 너비 마스크를 찾는 방식을 사용함 [4, 53, 33].

그러나 이러한 방법은 LLM 규모에서 다음과 같은 단점이 있음:

1. 메모리 및 연산량이 매우 높은 경사도(gradient) 계산이 필요함.
2. 합리적인 마스크를 얻기 위해 많은 양의 데이터와 미세 조정(fine-tuning)이 필요함.

이와 같은 한계를 인식한 예외적인 연구가 **Dery et al. [11]**로,

• 경사도를 사용하지 않고 단순 순전파(forward pass)만으로 계산할 수 있는 중요도 지표(saliency metric)를 제안함.

본 연구의 기여

• 너비(width)와 깊이(depth) 가지치기를 동시에 수행하는 최초의 연구
• LLM 규모에서도 적용 가능 (즉, 순전파만으로 중요도를 계산하고, 소량의 사전학습 데이터만 사용함)
• 최신 연구들과 비교했을 때, 가장 높은 압축률과 정확도를 달성

---

가지치기 후 정확도 회복 (Post-pruning Accuracy Recovery)

가지치기 후 모델의 성능을 회복하기 위해, 최근 연구에서는 더 크고 강력한 teacher 모델을 활용하거나(knowledge distillation) teacher 모델이 생성한 합성 데이터(synthetic data)를 사용하는 방법을 제안함.

• Teacher 모델을 활용한 연구: [2, 27]
• Teacher가 생성한 합성 데이터 사용: [1, 16, 36, 37]

본 연구의 차별점

• 기존 너비 및 깊이 가지치기 연구들과 비교했을 때 [26, 34, 53],
  • 가지치기 후 재학습(retraining) 과정에서, 비압축된(원본) teacher 모델을 활용한 지식 증류(distillation)를 최초로 적용한 연구임.
  • 즉, 구조적으로 가지치기된 student 모델의 성능을 향상시키는 데 teacher 모델을 적극적으로 활용한 새로운 접근법을 제안함.

6. 결론 (Conclusions)

본 연구에서는 대형 언어 모델(LLM)에서의 구조적 가지치기(structured pruning) 및 재학습(retraining)을 종합적으로 탐구하였으며,
이를 통해 가지치기 순서, 가지치기 축(axis) 간의 조합 효과, 최소한의 데이터로 수행하는 재학습 기법에 대한 새로운 통찰(insight)을 제공하였다.

우리는 광범위한 실험을 기반으로 가지치기 및 재학습을 위한 최적의 방법론(best practices)을 개발하였으며,
이를 활용하여 Nemotron-4 15B 모델을 2~4배 압축(pruning)하여 MINITRON 모델을 생성하였다.

핵심 성과

• 모델을 처음부터 학습하는 것과 비교하여, 가지치기된 MINITRON 모델은 훨씬 적은 비용으로 생성 가능
  • 최대 40배 적은 학습 토큰(training tokens)으로도 경쟁력 있는 성능 유지
• 유사한 크기의 커뮤니티 모델들과 비교하여 우수한 성능을 발휘
• 기존 연구에서 제안된 최신(depth 및 width) 가지치기 모델들을 능가하는 성과를 달성

---

한계점 및 향후 연구 방향

• 현재까지 Nemotron 모델 패밀리에만 제안된 기법을 적용하였으며,
  • 다른 모델 패밀리에도 적용하여 범용성을 검증하는 것이 향후 연구 과제
• 제안된 기법이 재학습을 필요로 한다는 점
  • 방법론 자체는 효율적이지만, 가지치기 후 모델을 완전히 재학습해야 한다는 점이 단점으로 작용할 수 있음
  • 향후 재학습 비용을 줄일 수 있는 추가적인 최적화 방법을 연구할 계획

최종 결론

본 연구는 LLM의 구조적 가지치기 및 재학습 전략을 최적화하는 새로운 접근법을 제시하였으며,
효율적인 모델 압축을 통해 비용 절감과 성능 유지라는 두 가지 목표를 동시에 달성할 수 있음을 입증하였다. 🚀

Reference

• https://openreview.net/forum?id=9U0nLnNMJ7&referrer=%5Bthe%20profile%20of%20Pavlo%20Molchanov%5D(%2Fprofile%3Fid%3D~Pavlo_Molchanov1)