NL-258, LLM Pruning and Distillation in Practice: The Minitron Approach, Preprint 2024

◼ Comment

  • Minitron 논문에서 간과했던 한 부분이, teacher 모델을 distillation할때 사용되는 학습 데이터다
    • 보통 distillation할때, teacher 모델 학습에 사용됐던 pretraining 데이터를 그대로 사용한다고 한다
    • 하지만 이는 항상 접근이 가능한게 아니기 때문에 데이터를 다른거를 써야할때 문제점이 생기고, 
    • 이 논문은 이 문제점을 해결한 논문이다
  • 방법
    • teacher correction으로 teacher모델을 사용할 수 있는 데이터로 살짝 학습한다
    • 그리고 여기서부터 distillation하고, 데이터도 이것을 사용하는 것이다
      • 이 논문에서는 Nemotron-4 curated continued training (CT) 데이터을 사용함
    • 방법론은 되게 직관적인데 성능은 좋다고 한다
    • 나머지는 minitron과 크게 다를건 없다
    • 근데 사실 여기서 사용되는 데이터에 따라 영향을 분석하지 않았는데.. 영향이 있지 않을까? 싶음
    • 마지막에 Instruct-tuning하여 Instruction model도 만드는데, 마찬가지로 성능이 좋다고 리포트한다
  • minitron에서 NAS로 찾은 아키텍쳐를 그대로 차용했다고 하고, 여기서 리포트하니 실제 사용하려면 비슷하게 구성하면 될것 같다?
    • 아키텍쳐에 맞게 프루닝할때, Layer 중요도 계산이 조금 달라짐 (본문 참조)
  • distillation의 최근 트렌드인거 같고, 엄청 중요하다고 개인적으로 느낀건
    • 세계 각각에서 pretraining을 다 따로하는데, findings에 따르면 사실 거대한 모델한번 만들때만 필요한게 아닐까? 싶음
    • 나머지는 그 모델로부터 시작하면 되지 않을까..? pruning, distillation, further pretraining이든 뭐든
  • 이런건 가능할까? 영어-한국어가 되는 모델에서 한국어에 포커싱해서 distillation하는것임
    • 그러면 영어 성능은 떨어져도 한국어 성능은 유지되는 작은모델이 되는것
    • 즉 한국어 테스크에서 오픈소스 모델을 쓰고 싶을때 heavy한 다국어모델을 가져와서 한국어 위주의 samll model을 만든것을 쓰면 인퍼런스 비용이 매우 절감되지 않을까?

Abstract

구조적 프루닝(structured pruning)과 지식 증류(knowledge distillation)는 소형 언어 모델(SLMs)을 효율적으로 얻을 수 있는 강력한 방법으로, 처음부터 모델을 학습하는 것보다 훨씬 적은 훈련 토큰과 계산 자원을 요구한다. 본 연구에서는 원래의 사전 학습 데이터셋에 대한 접근이 제한된 상황에서 이러한 전략을 효과적으로 적용하는 방법을 탐구한다. 이를 위해, 증류 전에 새로운 교사 수정(teacher correction) 단계를 도입하여, 교사 모델이 가벼운 파인튜닝 과정을 통해 특정 데이터 분포에 적응할 수 있도록 한다.

이 접근법을 사용하여 Mistral NeMo 12B 및 Llama 3.1 8B 모델을 각각 8B 및 4B로 압축하며, 이를 위해 두 가지 프루닝 기법을 적용한다:

  1. 깊이(depth) 프루닝
  2. 은닉/어텐션/MLP(너비) 프루닝

이후, LM Evaluation Harness의 대표적인 벤치마크를 활용해 성능을 평가한다. 추가적으로, NeMo Aligner를 활용한 정렬(alignment) 과정을 거친 후, 모델의 명령어 수행 능력, 롤플레이, 수학, 코딩 및 함수 호출(function calling) 역량을 테스트한다.

이 방법을 통해 Mistral NeMo 12B 모델을 기반으로 한 최첨단(state-of-the-art) Mistral-NeMo-Minitron-8B (MN-Minitron-8B) 모델을 제작하였으며, Llama 3.1 8B 모델을 기반으로 한 우수한 4B 모델 또한 생성하였다. 본 연구에서 생성한 기본 모델 가중치는 Hugging Face에서 **자유로운 라이선스(permissive license)**로 공개한다.

Introduction

LLM(대형 언어 모델) 제공업체들은 종종 **서로 다른 크기(예: Llama 3.1의 8B, 70B, 405B 매개변수 등)**를 가진 모델 계열을 처음부터 학습한다. 이는 다양한 배포 규모, 모델 크기, 그리고 계산 예산을 고려하는 사용자들에게 맞춤형 선택지를 제공하기 위함이다. 

하지만 수십억 개 이상의 매개변수를 가진 여러 모델을 처음부터 학습하는 것은 엄청난 시간, 데이터, 그리고 자원이 소모되는 작업이다.

최근 연구들은 **가중치 프루닝(weight pruning)**과 **지식 증류(knowledge distillation)**를 결합하여 LLM 모델 계열을 학습하는 비용을 획기적으로 절감할 수 있음을 보여주었다. 이 접근법에서는 가장 큰 모델만 처음부터 학습하고, 이후 더 작은 모델들은 큰 모델을 점진적으로 프루닝한 후 지식 증류를 수행하여 정확도를 회복하는 방식을 따른다. 

그러나 이러한 방법은 일반적으로 원본 사전 학습 데이터셋에 접근할 수 있다는 가정을 기반으로 한다. 하지만 최근 최첨단 LLM 모델들(오픈소스 포함)이 사유(private) 또는 독점 데이터셋을 활용해 학습되고 있기 때문에, 이 가정이 성립하지 않는 경우가 많다.

이 논문이 기존 연구와 가장 큰 다른점인듯 (기존의 pretraining 데이터에 접근 못해도 괜찮다. 가장 현실적인 시나리오긴 함)

우리의 접근 방식

본 연구에서는 기존 Minitron 압축 기법을 두 가지 방향으로 확장하여 원본 사전 학습 데이터 접근 없이 LLM을 효과적으로 압축하는 방법을 제안한다:

  1. 새로운 교사 수정(teacher correction) 단계 도입

    • 프루닝되지 않은 교사(teacher) 모델을 특정 데이터 분포에 적응시키는 과정
    • 원본 사전 학습 데이터 없이도 지식 증류를 수행할 수 있도록 함

  2. 더 효과적인 다운스트림 작업 기반 중요도(saliency) 기준을 적용한 깊이(depth) 프루닝 기법 도입

이러한 새로운 압축 전략을 최첨단 LLM 모델인 Llama 3.1 8B 및 Mistral NeMo 12B에 적용하여, 각각 4B 및 8B 크기로 압축하는 데 성공했다. 특히 Llama 3.1 8B 모델의 경우, 4B 크기의 두 가지 변형 모델을 생성하였다:

  1. Llama 3.1-Minitron-4B-Width (너비 축만 프루닝)
  2. Llama 3.1-Minitron-4B-Depth (깊이 축만 프루닝)

실험 결과

  • MN-Minitron-8B 모델 (Mistral NeMo 12B에서 압축된 모델)은 동일한 크기의 다른 모델들과 비교했을 때 최고 수준(state-of-the-art)의 성능을 달성함
  • Llama 3.1-Minitron-4B 모델(Depth 및 Width 프루닝 버전)은
    • 원본 Llama 3.1 8B 모델과
    • 기존 Minitron-4B 모델(이전 세대)보다 우수한 성능을 보임
  • 너비(Width) 프루닝 모델이 깊이(Depth) 프루닝 모델보다 전반적으로 높은 정확도를 기록
  • TensorRT-LLM을 이용한 실행 성능 평가 결과:
    • Depth 프루닝 모델: 2.7배 속도 향상
    • Width 프루닝 모델: 1.8배 속도 향상
    • 원본 Llama 3.1 8B 모델과 비교하여 뛰어난 효율성을 보여줌

이 연구는 대형 LLM을 보다 효율적으로 압축하는 새로운 전략을 제시하며, 특히 원본 데이터셋 없이도 강력한 성능을 유지할 수 있는 방법론을 확립했다.

Methodology

우리의 접근 방식을 전체적으로 나타낸 개요는 그림 1에 제시되어 있다. 
먼저, 교사(teacher) 모델은 대상 데이터셋에서 경량 조정 과정을 거쳐 증류(distillation)에 사용될 준비를 한다. 이 단계를 **교사 보정(teacher correction)**이라 한다. 이후, 모델을 압축하기 위해 가지치기(pruning)를 수행한 다음, 증류 과정을 통해 모델의 정확도를 복구한다.

교사 보정 (Teacher Correction)

증류는 더 정확한 교사 모델로부터 덜 정확한 학생(student) 모델로 지식을 압축하여 향상시키는 효과적인 기법이다 [3][2]. 

일반적으로 교사 모델을 훈련한 것과 동일한 데이터셋을 사용하여 지식을 증류한다. 

  • 그러나 원본 학습 데이터에 접근이 제한되는 경우, 다른 데이터셋을 이용해 증류를 수행하면 교사 모델이 최적의 지도를 제공하지 못한다는 실험적 관찰이 있다.
  • teacher correction을 해야하는 이유!

이는 교사 모델이 원래 훈련된 데이터셋과 증류에 사용되는 데이터셋 간의 서브워드 토큰(sub-word token) 분포 차이로 인해 발생한다고 가설을 세웠다. 

  • 이를 해결하기 위해, 교사 보정(teacher correction) 단계를 새롭게 제안한다(그림 2 참고). 
  • 이 과정에서는 교사 모델을 새로운 증류 데이터셋에 적응시키기 위해 약 **1000억 개의 토큰(∼100B tokens)**을 사용하여 경량 미세 조정(fine-tuning)을 수행한다.

  • 실험 결과(특히 그림 4 참고)에서 우리는 이 절차가 증류 성능을 크게 향상시켜, 보다 정확한 학생 모델을 만드는 데 기여함을 확인하였다. 
  • 또한, 증류와 병렬로 교사를 보정하는 방법도 탐색하였으며, 이는 완전히 보정된 교사를 이용하는 방식과 유사한 성능을 보였다.

가지치기 (Pruning)

가중치 가지치기(weight pruning)는 모델 크기를 줄이는 강력하고 널리 알려진 기법이다. 본 연구에서는 **구조적 가지치기(structured pruning)**에 집중하는데, 이는 모델 가중치에서 특정 블록(또는 채널) 단위로 요소들을 제거하는 방식이다. 구조적 가지치기의 대표적인 기법으로는 뉴런(neuron), 어텐션 헤드(attention head), 합성곱 필터(convolutional filter), 깊이(depth) 가지치기 등이 있다 [6,7,8,9].

우리는 Minitron [2]에서 제안된 **가지치기 절차(pruning recipe)**를 따른다(그림 3 참고).

먼저, 모델의 각 **레이어(layer), 뉴런(neuron), 헤드(head), 임베딩 차원(embedding dimension)**의 중요도를 계산한 후, 이를 정렬하여 중요도 순위를 결정한다.

중요도 추정 (Importance Estimation)

우리는 순전히 활성화(activation) 기반의 중요도 추정 기법을 사용하며, 이는 깊이(depth), 뉴런(neuron), 헤드(head), 임베딩 채널(embedding channel) 등 여러 축(axis)에 대한 민감도(sensitivity) 정보를 동시에 계산한다. 

이 과정에서 소규모 **보정 데이터셋(1024개 샘플)**을 이용하며, 순방향 전파(forward propagation)만 수행하여 중요도를 측정한다.

  • 어텐션 헤드(head) 중요도: 다중 헤드 어텐션(Multi-Head Attention, MHA) 층의 활성화를 분석하여 평가
  • 뉴런 중요도: 다층 퍼셉트론(Multi-Layer Perceptron, MLP) 층의 활성화 분석
  • 임베딩 채널 중요도: MLP, MHA, LayerNorm 층의 활성화 분석

레이어 중요도 (Layer Importance)

깊이(depth) 가지치기의 경우, 레이어의 중요도를 평가하는 두 가지 기준을 사용한다.

  1. 언어 모델(ML) 검증 손실/퍼플렉서티(PPL)
  2. 다운스트림 작업에서의 정확도

이전 연구에서 제안된 블록 중요도(Block Importance, BI) 지표 [8]는 검증 손실/퍼플렉서티(PPL) 기준보다 성능이 낮다고 보고되었기 때문에 사용하지 않는다 [2].

다운스트림 정확도는 이번에 추가된거 같음

중요도 순위를 매길 때는 연속된(layer block) 레이어 단위로 제거하는 방식을 따른다. 이는 개별 레이어를 제거하는 것보다 더 나은 성능을 제공한다는 Gromov et al. [11]의 연구 결과를 반영한 것이다. 실험적 분석(그림 8 및 9 참고) 결과, 우리는 **Winogrande 평가 지표 [10]**를 활용하여 연속된 레이어를 가지치기하는 전략을 도출하였다.

모델 트리밍 (Model Trimming)

Minitron [2]의 접근 방식을 따라, 특정 모델 아키텍처 설정에서 각 축(axis)의 요소를 중요도 순으로 정렬한 뒤, 해당 가중치 행렬을 직접 조정하여 불필요한 부분을 제거한다.

  • 뉴런 및 헤드 가지치기: MLP 및 MHA 층의 가중치를 조정하여 수행
  • 임베딩 채널 가지치기: MLP, MHA, LayerNorm 층의 가중치 행렬에서 임베딩 차원을 줄이는 방식

기존 연구에서는 **신경망 아키텍처 탐색(Neural Architecture Search, NAS)**을 통해 최적의 아키텍처를 찾았으나, 본 연구에서는 해당 단계를 생략하고 기존 연구에서 도출된 네트워크 구조적 학습 결과를 그대로 활용하였다.

증류를 활용한 재훈련 (Retraining with Distillation)

우리는 **재훈련(retraining)**이라는 용어를 가지치기 후 모델의 정확도를 회복하는 과정을 의미하는 데 사용한다. 

본 연구에서는 두 가지 재훈련 전략을 탐색하였다.

  1. 일반적인 학습(conventional training) – 정답 라벨(ground truth labels)을 활용한 훈련
  2. 지식 증류(Knowledge Distillation, KD) – 가지치기되지 않은 모델(교사)로부터 감독을 받아 학습

**지식 증류(KD) [3]**는 더 크거나 복잡한 교사(teacher) 모델에서 더 작거나 단순한 학생(student) 모델로 지식을 전달하는 과정이다. 이 과정에서 학생 모델은 교사 모델의 출력 및/또는 중간 상태를 모방하도록 학습한다.

본 연구에서는 비압축 모델을 교사 모델, 가지치기된 모델을 학생 모델로 설정하여 증류를 수행하였다. Minitron 연구 [2]에서 제안한 최적의 방식에 따라, 우리는 **교사와 학생 모델의 로짓(logits) 사이에서 전방 KL 발산(forward KL Divergence) 손실 [12]**을 사용하여 학습을 진행하였다. 이 과정은 그림 2에 나타나 있다.

학습 세부 사항 (Training Details)

사전 학습 (Pre-training)

Llama 3.1 8B [1] 및 Mistral NeMo 12B [5] 모델은 서로 다른 **독점 데이터셋(proprietary datasets)**에서 사전 학습되었으며, 우리는 해당 데이터에 접근할 수 없다. **Llama 3.1 기술 보고서 [1]**에 따르면, 8B 모델은 15조(15T) 토큰으로 사전 학습되었다. 본 연구에서는 Hugging Face에서 공개적으로 이용 가능한 기본(Base) 모델을 출발점으로 사용하였다.

Llama 3.1 8B [1] 및 Mistral NeMo 12B [5] 모델을 teacher 모델로 사용한듯

데이터셋 (Dataset)

우리의 가지치기 및 증류 실험에서는 Nemotron-4 curated continued training (CT) 데이터셋 [13][14]을 사용하였다.

교사 보정 (Teacher Correction)

원본 Mistral NeMo 12B 또는 Llama 3.1 8B 모델을 교사 모델로 직접 사용하면, 우리가 사용하는 데이터셋에서 최적의 성능을 내지 못했다. 이를 해결하기 위해, 이전 섹션에서 설명한 **교사 보정(teacher correction)**을 두 모델 모두에 적용하였다.

  • 교사 보정 과정: 약 1000억 개(∼100B) 토큰을 사용하여 수행
  • 목적: 교사 모델을 증류 데이터셋에 적응시키는 것
  • 설정:
    • 120 스텝의 워밍업(warm-up)
    • 낮은 학습률 (기존 모델에서 사용된 최고 학습률의 1/5 수준)
    • 동일한 배치 크기(batch size)
    • 기존 모델과 동일한 최소 학습률(minimum learning rate) 및 감소 스케줄(decay schedule)

교사 보정 과정이 다운스트림 작업에서 교사 모델의 정확도에 미치는 영향은 크지 않으며, 일부 작업에서는 성능이 향상되었고 일부에서는 감소하는 결과를 보였다(표 1 참고).

teacher correction 해도 성능 자체가 떨어지거나 그렇지는 않는다고 말함 (표의 맨오른쪽 FT가 그걸 말하는듯)

LoRA을 사용하거나 등의 방법은 탐구하지 않았음

이는 미세 조정(fine-tuning) 과정에서 사용한 데이터셋의 특성에 기인한 것으로 가설을 세울 수 있다. 향후 연구에서는 1000억 개 미만의 토큰을 사용하거나, LoRA (Low-Rank Adaptation) [15]와 같은 가벼운 미세 조정 기법, 혹은 LayerNorm [16] 파라미터만 조정하는 방법 등을 활용하여 이 과정을 더욱 최적화하는 것이 흥미로운 연구 주제가 될 것이다.

가지치기 (Pruning)

우리의 **가지치기 기법(pruning recipe)**은 **Minitron 논문 [2]**에서 제시된 최적의 방법론을 기반으로 하며, 이전 섹션에서 설명한 내용을 따른다.

너비 가지치기 (Width Pruning)

  • L2-노름(L2-norm)과 평균(mean) 함수를 각각 배치(batch) 및 시퀀스(sequence) 차원에서 집계(aggregation) 함수로 사용
  • 단일 실행 방식(single-shot pruning) 사용, 반복적인 가지치기(iterative approaches)는 피함

깊이 가지치기 (Depth Pruning)

  • Gromov et al. [11]의 연구를 참고하여 연속된 서브그룹(layer block) 단위로 제거
  • Winogrande 평가 기준 [10]을 사용하여 정확도 손실이 최소화되는 구조 선택

이번 연구에서는 경량 신경망 아키텍처 탐색(NAS, Neural Architecture Search) 단계를 생략하고, Llama 3.1-Minitron-4B 및 MN-Minitron-8B 모델에 대해 수동으로 아키텍처를 설정하였다.
이러한 모델 구조는 Minitron-4B 및 Minitron-8B 모델 [2]에서 영감을 얻었으며, 세부적인 가지치기 설정은 표 3에 정리되어 있다.

Minitron 방식을 따라서 했다는 것 (MN=Mistral Nemo을 말하는듯)

LLaMA 8B -> 4B, Mistral 12B -> 8B로 압축

가지치기 적용 사례 (Pruning Recipes)

Llama-3.1-Minitron-4B (너비 가지치기)

  • 기본 모델: Llama 3.1 8B
  • 히든 차원(hidden dimension): 4096 → 3072
  • MLP 히든 차원(MLP hidden dimension): 14336 → 9216
  • 어텐션 헤드(attention heads): 변경 없음
  • 깊이(depth): 변경 없음

Llama-3.1-Minitron-4B (깊이 가지치기)

  • 기본 모델: Llama 3.1 8B
  • 히든 차원(hidden dimension): 변경 없음
  • MLP 히든 차원(MLP hidden dimension): 변경 없음
  • 어텐션 헤드(attention heads): 변경 없음
  • 깊이(depth): 32 → 16

MN-Minitron-8B

  • 기본 모델: Mistral NeMo 12B
  • 히든 차원(hidden dimension): 5120 → 4096
  • MLP 히든 차원(MLP hidden dimension): 14336 → 11520
  • 어텐션 헤드(attention heads): 변경 없음
  • 깊이(depth): 변경 없음


증류 (Distillation)

우리는 로짓(logit) 기반 증류 방식을 선택하여, 전방 KL 발산(Forward KL Divergence) [12] 손실을 최소화하는 방식으로 학습을 진행하였다. 이 과정에서 LM(언어 모델) 교차 엔트로피 손실(cross-entropy loss)은 무시하였다.

  • 교사 모델(teacher): 가지치기되지 않은 모델
  • 학생 모델(student): 가지치기된 모델

증류 과정에서 사용된 하이퍼파라미터는 표 4에 정리되어 있으며, 학습을 위해 32개의 NVIDIA DGX H100 노드를 활용하였다.

인스트럭션 튜닝 (Instruction Tuning)

증류된 모델의 명령어 수행(instruction-following) 능력을 평가하기 위해, NeMo-Aligner [17]을 이용한 정렬(alignment) 과정을 수행하였다.

모든 모델에 대해 동일한 방법을 적용하였으며, 다음 단계를 따랐다.

  1. 수학 및 코드(supervised fine-tuning, SFT) 감독 학습
  2. 명령어 SFT (instruction SFT) 수행
  3. 보상 기반 선호 최적화(Reward-aware Preference Optimization, RPO) [18] 2회 반복

분석 (Analysis)

우리는 증류, 교사 보정, 새로운 깊이 가지치기 중요도 평가 지표의 영향을 보다 정확히 이해하기 위해 **소거 연구(ablation studies)**를 진행하였다.

교사 보정(Teacher Correction)

  • 그림 4: MN-Minitron-8B 모델에서 교사 보정의 효과를 비교
    • 교사 보정을 수행한 경우, 보정하지 않은 교사 모델을 직접 증류한 경우보다 분명한 성능 향상이 나타남
  • 그림 5: 교사 보정 방법 두 가지 비교
    1. 보정된 교사 모델을 가지치기 후 증류
    2. 원래(보정되지 않은) 교사 모델을 가지치기한 후, 지속적으로 보정하며 증류
    • 결과: 교사 보정과 증류를 병렬로 수행해도 가지치기된 학생 모델의 정확도를 효과적으로 회복할 수 있음
    • 그림상으로는 그냥 맨처음에 correction하고 쭉 진행하는게 나아보이긴함

가지치기 및 증류 (Pruning and Distillation)

  • 그림 6: 가지치기와 증류가 기존 방법 대비 어떤 이점을 제공하는지 비교

    1. 랜덤 가중치 초기화 후 증류

    2. 중요도 점수를 무시하고 무작위 가지치기 후 증류

    3. 제안된 가지치기 방식 + 기존 LM 교차 엔트로피 손실 학습

    4. 제안된 가지치기 방식 + 증류 기반 재훈련

    • 결과:
      • 가지치기를 적용하면 랜덤 초기화보다 훨씬 더 나은 시작점을 제공
      • 증류 기반 학습이 기존 학습 방식보다 성능이 우수
      • **훈련에 필요한 토큰 수를 크게 절감(최대 40배, 15T → 380B 토큰 사용)**하면서도 최첨단 MN-Minitron-8B 모델을 생성 가능

너비 vs. 깊이 가지치기 (Width vs. Depth Pruning)

  • 그림 7: Llama 3.1-Minitron-4B 모델에서 너비(Width)와 깊이(Depth) 가지치기 비교
    • 너비 가지치기가 더 낮은 초기 손실(loss)과 더 안정적인 학습 곡선을 보임
    • 동일한 파라미터 개수를 유지하면서도 너비 가지치기가 깊이 가지치기보다 일관된 성능 향상을 제공

깊이 가지치기 중요도 평가 지표 (Depth Pruning Metrics)

  • 그림 8: 연속된 블록(layer block)을 제거할 때 LM 검증 손실 증가율 분석

    • 모델의 앞부분과 뒷부분(layer 시작과 끝)이 가장 중요한 역할을 함
    • 비연속적인 레이어를 제거하면 LM 검증 손실을 더 잘 최적화할 수도 있음 (점선 참조)

  • 그림 9: Winogrande 정확도를 기준으로 깊이 가지치기 방식 비교

    1. 레이어별 중요도 기반 가지치기(16개 개별 레이어 선택) → 정확도 0.5

    2. 연속적인 16~31번째 레이어 제거 [11] → 정확도 0.595

    • 결과: 증류 기반 재훈련 시에도 연속된 레이어를 제거하는 방법이 더 나은 성능을 유지
    • 따라서 본 연구에서는 연속적인 블록(layer block) 제거 방식을 선택함.

즉 LLM PPL 기준에서는, 비 연속된 layer을 제거하는게 좋으나

  • 실제 다운스트림 테스크에선, 연속된 layer을 제거하는게 더 성능이 높다고 한다
  • 따라서 연속적인 layer (16-31)을 제거했다고 함

평가 (Evaluation)

벤치마크(Benchmarks)

Llama [20] 연구를 따라, 우리의 압축된 베이스(Base) 모델과 정렬된(Aligned) 모델을 여러 다운스트림 작업에서 평가하였다.

  • 일반적 평가(Benchmarks for Base Models):

    • MMLU [21] (일반 지식 평가)
    • HumanEval [22] (파이썬 코드 생성)
    • MBPP [23] (코드 생성)
    • GSM8K [24] (수학 문제 해결)
    • 질문-응답(QA) 및 상식 추론:
      • ARC-Challenge [25]
      • HellaSwag [26]
      • TruthfulQA [27]
      • WinoGrande [10]
    • XL-Sum English [28] (문서 요약)

  • 명령어 튜닝(Instruct-Tuned Models) 평가:

    • 질문 응답(QA): GPQA [29]
    • 함수 호출(Function Calling): BFCL [30]
    • 명령어 수행(Instruction Following): IFEval [31]
    • 멀티턴 대화(Multi-turn Conversations): MT-Bench (GPT-4 Turbo) [32]
      → 이 MT-Bench는 기존 MT-Bench [33]의 수정된 버전임.

평가 설정 (Evaluation Settings)

  • 베이스 모델:

    • MMLU: 5-shot
    • Winogrande: 5-shot
    • ARC-Challenge: 25-shot
    • HellaSwag: 10-shot
    • XL-Sum: 0-shot (20% 샘플링)
    • HumanEval 및 MBPP: pass@1 점수 (Temperature = 0.2, Nucleus Sampling (top-p = 0.95))

  • 정렬(Aligned) 모델:

    • 0-shot 평가 + 탐욕적 샘플링(Greedy Sampling) 사용 (적용 가능 시)

베이스 모델 평가 (Base Model Evaluation)

  • 결과 (표 1 참고)
    • MN-Minitron-8B 모델은 유사한 크기의 모델보다 전반적으로 더 높은 정확도를 기록
    • Llama 3.1 8B 대비, 40배 적은 학습 토큰(380B vs. 15T)으로 더 우수한 성능
    • Llama 3.1-Minitron-4B 모델도 Llama 3.1 8B 대비 150배 적은 학습 토큰(94B vs. 15T)으로 더 나은 성능
    • 너비 가지치기(Width Pruning) 모델이 깊이 가지치기(Depth Pruning) 모델보다 성능이 우수함

➡️ 우리의 접근 방식이 최첨단 정확도를 제공하면서도 학습 효율성을 획기적으로 개선했음을 입증함.

명령어 튜닝 모델 평가 (Instruct Model Evaluation)

  • 결과 (표 2 참고)
    • 대부분의 벤치마크에서 유사한 크기의 모델을 능가
    • 예외적으로 HumanEval [35] 및 MBPP [23]에서는 다소 낮은 성능
    • Llama 3.1-Minitron-4B 모델이 MT-Bench에서 Gemma2보다 낮은 성능
    • 하지만 MMLU, GSM8K, GPQA, IFEval, BFCLv2 벤치마크에서 일관된 우수한 성능을 보임

➡️ 이 결과는 우리의 정렬 모델이 강력한 명령어 수행 및 질문-응답(QA) 능력을 갖추었음을 보여줌.

실행 성능 분석 (Runtime Performance Analysis)

Llama 3.1 8B 및 Llama 3.1-Minitron-4B 모델을
NVIDIA TensorRT-LLM을 사용하여 최적화하고 실행 속도를 평가함.

FP8(부동소수점 8비트) 기반 요청 처리 속도(Throughput)

  • 실험 환경:

    • H100 80GB GPU 단일 사용
    • 입력 시퀀스 길이(ISL) / 출력 시퀀스 길이(OSL) 조합 증가
    • 배치 크기:
      • 8B-12B 모델 → 32
      • 4B 모델 → 64 (메모리 사용량이 적어 더 큰 배치 가능)

  • 결과 (그림 10 참고):

    • Llama 3.1-Minitron-4B (Depth) 모델 → Llama 3.1 8B 대비 2.7배 더 높은 처리량
    • Llama 3.1-Minitron-4B (Width) 모델 → Llama 3.1 8B 대비 1.8배 더 높은 처리량
    • BF16(16비트) 대비 FP8(8비트) 사용 시 성능 1.4배 향상

➡️ 작은 모델이 더 빠르고 효율적으로 실행됨을 확인. 특히 깊이 가지치기 모델이 실행 성능 면에서 가장 빠름.

결국 이렇게 큰 모델에서 만든 작은모델이, scratch부터 학습한 모델보다 낫다는 것이다

  • 12B->8B로 압축한 모델 > scratch 8B모델

주요 인사이트 (Insights)

이번 평가를 통해 얻은 흥미롭고 예상치 못한 발견들을 정리하였다.

일반적인 관찰 (General Observations)

  1. 교사 보정(Teacher Correction)은 증류의 최적 성능을 위해 필수적이다.

    • 새로운 데이터셋에서 증류를 효과적으로 수행하려면 교사 모델을 해당 데이터셋으로 미세 조정(fine-tuning)하는 과정이 필요함.
    • 이를 수행하면 LM 검증 손실(LM validation loss)이 6% 이상 감소함.
    • 또한, 교사 보정은 가지치기의 최적성(optimality)에 영향을 미치지 않으며, 증류와 병렬로 수행할 수도 있음.

  2. Minitron 논문의 결과와 일치하는 점

    • 가지치기 및 증류 후에도 최첨단(state-of-the-art) 정확도를 유지하면서, 10배 적은 학습 토큰(380B vs. 15T)만으로도 성능을 확보할 수 있음.

  3. 너비 가지치기(Width Pruning)에서는 특정 차원만 줄이는 것이 더 효과적이다.

    • 어텐션 헤드(attention heads)는 유지하고,
    • MLP 중간 차원(intermediate dimension)과 임베딩 채널(embedding channels)만 가지치기하면 더 높은 정확도를 달성함.

Mistral NeMo 12B → MN-Minitron-8B 모델 압축 결과

  1. 가지치기 및 증류 후, 교사 모델보다 더 나은 성능을 달성한 벤치마크 존재
    • GSM8K: 55.7% → 58.5% (수학 문제 해결 성능 향상)
    • HumanEval: 23.8% → 36.2% (코드 생성 성능 대폭 향상)
    • 이러한 개선은 주어진 데이터셋의 영향이 클 것으로 보이며, 재훈련은 증류 손실(distillation loss)만을 사용하여 수행됨.

Llama 3.1 8B → Llama 3.1-Minitron-4B 모델 압축 결과

  1. 너비 가지치기가 깊이 가지치기보다 MMLU 정확도가 더 높음

    • 너비 가지치기 모델(Width Pruned): MMLU 60.5%
    • 깊이 가지치기 모델(Depth Pruned): MMLU 58.7%

  2. 깊이 가지치기 모델(Depth Pruned)의 추론 능력(Reasoning Ability)이 크게 저하됨

    • GSM8K 정확도
      • 너비 가지치기 모델: 41.24%
      • 깊이 가지치기 모델: 16.8%
    • 하지만 명령어 튜닝(Instruct Tuning)을 적용하면 이 성능 차이가 줄어듦.

  3. 실행 속도(Throughput) 비교

    • 깊이 가지치기(Depth Pruning) 모델 → Llama 3.1 8B 대비 2.7배 속도 향상
    • 너비 가지치기(Width Pruning) 모델 → Llama 3.1 8B 대비 1.7배 속도 향상

  4. 깊이 가지치기에서 연속적인 레이어 제거가 더 효과적임

    • 비연속적인(importance-based) 레이어 제거보다, 연속된 블록(block) 단위로 제거하는 것이 더 나은 성능을 제공.

Reference

댓글

댓글 쓰기