[OUTTA Alpha팀 논문 리뷰] Part 8-7. GPT-NeoX-20B: An Open-Source Autoregressive Language Model

논문 링크: https://arxiv.org/abs/2204.06745

 

GPT-NeoX-20B: An Open-Source Autoregressive Language Model

---

1. Introduction

• 최근 대규모 언어 모델(LLM)의 급격한 발전
• Transformer 기반 언어 모델의 파라미터 수를 늘릴수록 성능이 예측 가능하게 향상되는 현상 확인
• 대형 모델들이 존재하지만, 대부분 비공개 API 형태 혹은 제한된 접근만 가능
• GPT-NeoX-20B는 공개적으로 가중치를 배포해 연구·해석·안전성 등 다양한 연구에 기여 목표
• 학습 과정에서 얻은 체크포인트를 모두 공개해 훈련 동학(dynamics) 연구에도 도움
• 학습 결과 일부가 기존 문헌과 다른 양상 확인
  • 데이터 중복에 따른 손실 증가가 발견되지 않음
  • 특정 과제에서 few-shot이 성능에 큰 기여

---

2. Model Design and Implementation

2.1 Model Architecture

• GPT-3 스타일의 오토레그레시브 Transformer 디코더 구조를 기반
• 총 44개 레이어, 히든 차원 6144, 헤드 수 64, 파라미터 약 200억 개

• 2.1.1. Rotary Positional Embeddings

  • 

  • 기존 GPT 계열 모델에서 사용하던 학습형 positional embedding 대신, rotary embeddings 적용
    • embedding 공간을 회전(rotate)하는 방식으로, 토큰 위치 $m$과 $n$의 attention이 선형적으로 $m-n$에 의존하도록 설계
    • 일반적인 multiheaded attention 공식 $$\mathrm{softmax}\Bigl(\frac{1}{\sqrt{d_{\mathrm{nm}}}}\; x_m^\top W_q^\top \;W_k \; x_n \Bigr)$$에서, rotary를 적용한 경우 $$\mathrm{softmax}\Bigl(\frac{1}{\sqrt{d_{\mathrm{nm}}}}\; x_m^\top W_q^\top \; R^{d_x}_{(n-m)} \;W_k \; x_n \Bigr)$$ 로 변경
      • $x_m, x_n$​: 각각 위치 $m$, $n$에 해당하는 임베딩(배치 형태)
      • $W_q, W_k$​: 쿼리/키 가중치
      • $R^{d_x}_{(n-m)}$​​: $d \times d$ 블록 대각 행렬(block diagonal) 형태. 내부 각 블록은 2D 회전 행렬로 구성
      • 하이퍼파라미터 $\theta_i$ 설정 등으로, $\theta = \frac{i}{10000^{2i/d}}$ 형태를 가질 수 있음
  • 적용 차원
    • 전체 임베딩 벡터에 rotary를 적용하기보다, 첫 25% 차원에만 적용
    • 실험 결과, 성능과 계산 효율 사이에서 균형이 가장 좋았다는 결론

• 2.1.2. Parallel Attention + FF Layers
  • 일반적으로 Attention 레이어와 Feed-Forward(FF) 레이어를 직렬(series)로 쌓는 대신, 병렬(parallel)로 계산 후 합산
  • 병렬화 이유
    • residual addition(잔차 추가) 시 all-reduce 통신이 여러 번 발생하지 않도록 줄이는 방안
    • 한 번에 Attention과 FF 출력을 합쳐서 한 번의 all-reduce로 처리 가능
    • 내부 실험에서 최대 15% 처리량(throughput) 증가 확인
  • Layer Norm 이슈
    • 코드상의 실수로 두 개의 독립된 Layer Norm을 사용하게 됨
    • 원래 의도: $$x + \mathrm{Attn}(\mathrm{LN1}(x)) + \mathrm{FF}(\mathrm{LN1}(x))$$
    • 실제 코드: $$x + \mathrm{Attn}(\mathrm{LN1}(x)) + \mathrm{FF}(\mathrm{LN2}(x))$$
    • 소규모 추가 실험 결과, 성능 차이는 미미했다고 보고

• 2.1.3. Initialization
  • Feed-Forward(FF) 출력 레이어(Residual 이전)에 대해 특정 초기화 스킴 사용
    • $\sqrt{\frac{2}{L \cdot d}}$ 형태 (병렬 Attention+FF 구조 고려해 $\times 2$ 보정)
    • 깊이와 너비가 큰 모델에서 활성값이 급증하지 않도록 방지
  • 나머지 레이어는 $\sqrt{\frac{2}{d + 4d}}$​​ 같은 ‘small init’ 스킴 활용

• 2.1.4. All Dense Layers
  • GPT-3는 sparse layer와 dense layer를 번갈아 쓰는 방식 존재
  • 본 모델에서는 구현 복잡도 감소를 위해 모든 레이어를 dense로만 구성

2.2 Software Libraries

• Megatron과 DeepSpeed 기반으로 대규모 모델 학습을 효율화
• PyTorch v1.10.0 사용
• 분산 통신은 NCCL 활용

2.3 Hardware

• 8개의 NVIDIA A100-SXM4-40GB GPU가 장착된 서버 12대(총 96 GPU)로 구성
• 각 서버에는 2개의 AMD EPYC 7532 CPU, InfiniBand 스위치 등 고속 통신망
• 모델 병렬화(텐서 병렬 + 파이프라인 병렬) 및 데이터 병렬을 조합
• GPU 당 최대 약 117 TFLOPS 효율 유지

---

3. Training

• 초대형 모델에 대한 하이퍼파라미터 탐색이 어렵기에 기존 대규모 모델 설정을 참고
• 최종 학습 스펙
  • 학습률: 두 모델 스케일(예: 13B, 175B) 사이를 보간해 결정
  • 배치 크기: 약 3.15M 토큰(컨텍스트 길이 2048, 배치 1538 등)
  • 총 15만 스텝 진행, 코사인 스케줄로 학습률 90% 감쇠
  • 옵티마이저: AdamW(베타, 감마 등 특정 값 설정) + ZeRO
  • 텐서 병렬 크기 2, 파이프라인 병렬 크기 4로 설정 → 노드 내부에서 텐서·파이프라인 병렬, 노드 간에는 데이터 병렬
• GPU당 117TFLOPS 달성

3.1 Training Data

• The Pile이라는 약 825GiB 규모 텍스트 데이터 사용
• 구성: 학술논문(논문 초록, 법률, 특허 등), 웹 스크레이핑, 인터넷 리소스, 소설·프로즈, 대화 로그, 프로그래밍 코드 등 다양
• 일부 데이터는 비중 조정을 위해 업샘플링
• StackExchange(Q/A 형태) 포함 → 코드, 수학 문제 등에 좋은 영향 기대
• 과거 대규모 모델들과는 달리, 웹만이 아니라 다채로운 도메인 통합

3.2 Tokenization

• BPE 기반 어휘 크기 50257
• The Pile에 맞게 새로 학습한 토크나이저
• 공백 처리 일관화 → 문자열 시작부 공백을 안정적으로 구분
• 반복 공백 토큰(1~24칸) 포함 → 소스코드처럼 공백 많은 텍스트 처리 시 토큰 수 절약 가능

3.3 Data Duplication

• 일괄적으로 1 epoch만 학습하거나, 중복 데이터 제거(dedup)를 권장하는 추세가 있었음
• 그러나 이 모델은 중복 제거 없이 원본 Pile 사용 → 여러 데이터가 1 epoch 넘도록 노출
• 실제로 1 epoch 지점 이후에도 검증 손실이 떨어지는 현상을 확인 → 중복으로 인한 성능 저하 징후 없음
• 다만, 실제 프로덕션 환경에서 데이터 중복 제거는 정보 유출 방지 등으로 유용할 수 있음

---

4. Performance Evaluations

• EleutherAI Language Model Evaluation Harness로 다양한 벤치마크 수행
• 비교 대상
  • 특정 API 기반 모델(GPT-3 시리즈)
  • 공개된 FairSeq 밀집 모델
  • 동일 데이터셋으로 학습된 GPT-J-6B 등
• 일부 모델(API 등)는 0-샷만 평가, 그 외 모델들은 0-샷 및 5-샷 함께 평가

4.1 Tasks Evaluated

1. 자연어 과제(Natural Language Tasks)
   • 예: ANLI, ARC, HellaSwag, LAMBADA, LogiQA, TriviaQA, Winogrande 등
2. 수리 과제(Mathematical Tasks)
   • MATH 테스트 세트, 여러 계산 문제(덧셈·뺄셈 등)
3. 고급 지식 기반 과제(Advanced Knowledge-Based Tasks)
   • 특정 분야(인문·사회·STEM 등)에 대한 심화형 문제(객관식 등)

---

5. Discussion

5.1 Performance Results

• 자연어 과제:
  • GPT-NeoX-20B가 FairSeq 13B보다 우수한 경우가 많으나(HellaSwag 제외) 일부 과제에서는 미세하게 뒤지기도 함
  • HellaSwag에서 매우 큰 격차로 뒤처짐 → 원인은 불명확
  • LAMBADA에서는 오히려 FairSeq 모델보다 크게 앞섬
• 수리 과제:
  • GPT-NeoX-20B가 GPT-3, FairSeq 대비 산술 능력이 강한 편
  • 다만, 데이터셋 중 수학식이 많이 등장해 생긴 현상일 수 있음
• 고급 지식 과제:
  • 기존 연구에서는 GPT-3가 few-shot으로 이득을 못 본다고 했지만, GPT-NeoX-20B 및 FairSeq는 5샷에서 성능이 상당히 상승
  • 모델마다 few-shot 반응이 달라, 특정 모델 기준으로 일반화하면 안 됨을 시사

5.2 Powerful Few-Shot Learning

• 0샷 → 5샷 전환 시 GPT-NeoX-20B는 성능이 크게 향상
• 동일 크기의 다른 모델보다 few-shot 성능 증가 폭이 큼
• GPT-J-6B도 마찬가지 경향 → 같은 데이터(Pile) 영향을 추정

5.3 Limitations

• 대규모 모델에서 하이퍼파라미터 최적화가 어려워, 완벽한 설정인지 확신 없음
• 코딩 문제 전용 평가를 제대로 수행하지 못함(비용·복잡도)
• 데이터 중복 제거를 적용하지 않아, 실제로 성능에 어떤 영향이 있었는지 추가 연구 필요

5.4 Releasing a 20B Parameter LLM

• 기존에는 대형 모델 대부분을 비공개로 유지
• 해당 모델을 공개함으로써, 안전성·해석 가능성·LLM 연구 활성화 기대
• 큰 조직·정부만 모델을 독점하는 것보다, 공개해 다양한 연구자가 접근 가능하게 하는 것이 더 바람직하다는 입장
• 충분한 GPU 리소스가 필요하지만, 추후 인퍼런스 비용 절감 기법 등을 추가 연구 예정

---

6. Summary

• GPT-NeoX-20B는 20B 파라미터를 가진 오토레그레시브 Transformer 모델
• The Pile 데이터셋으로 학습했으며, 기존 모델과 다른 Rotary Embeddings, 병렬 Attention+FF 구조 등 주요 차이 반영
• 다양한 자연어·지식·수리 과제에서 공개된 모델 중 높은 수준의 성능
• few-shot에서 성능 편차가 크게 늘어, 특히 5-샷에서 발전폭이 큼
• 학습 및 평가 코드를 모두 공개하며, 모델 가중치도 전(全) 학습 단계별로 배포
• 대형 모델의 오픈 액세스가 가져올 여러 이점을 중요하게 고려해 공개함