PaLM (2022) 논문 리뷰

PaLM 모델의 논문 이름은 PaLM: Scaling Language Modeling with Pathways다. (링크)

 

저자는 Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, Noah Fiedel다. 

 

구글에서 공개한 LLM이다. 

 

 

Abstract

 LLM은 퓨샷 학습(Few-shot learning)을 사용하여 다양한 자연어 태스크에서 놀라운 성능을 달성했니다. 이는 특정한 애플리케이션에 모델을 적용하는 데 필요한 작업별 학습의 사례 수를 대폭 줄여준다. 퓨샷 학습에 대한 규모의 영향을 더 잘 이해하기 위해, Pathways Language Model (PaLM)이라고 명명한 540 billion 개의 매개변수를 가진 Densely Activated 트랜스포머 언어 모델을 학습했다.
 여러 TPU Pod에서 고효율 학습을 지원하는 새로운 머신러닝 시스템인 Pathways를 사용하여 6,144개의 TPU v4 칩에서 PaLM을 학습했다. 수백 개의 언어 이해 및 생성 벤치마크에서 SOTA 퓨샷 학습 결과를 달성함으로써 확장의 지속적인 이점을 입증했다. 이러한 여러 작업에서 PaLM 540B는 획기적인 성능을 달성하여 multi-step reasoning task 다단계 추론 작업에서 파인 튜닝된 최첨단 기술을 능가하고, 최근 출시된 BIG-bench 벤치마크에서 평균적인 인간의 퍼포먼스를 능가했다. 상당수의 BIG-bench 작업에서 모델 규모에 따른 비연속적인 개선이 나타났는데, 이는 가장 큰 모델로 확장할수록 성능이 급격히 향상 됨을의미한다. PaLM은 다국어 작업 및 소스 코드 생성에도 강력한 역량을 갖추고 있으며, 이를 다양한 벤치마크에서 입증 되었다. 또한, bias and toxicity 편향 및 독성에 대한 포괄적인 분석을 제공하고, 모델 규모에 따른 훈련 데이터 기억의 정도를 연구한다. 마지막으로, LLM과 관련된 윤리적 고려 사항과 잠재적인 mitigation strategies 완화 전략에 대해 논의한다.

 

Introduction

PaLM은 540 B의 파라미터 수를 가지는 autoregressive Transformer 구조로 780 B 크기의 고품질의 텍스트 토큰으로 학습했다. 이때 Pathways의 방법으로 학습하여 TPU v4 Pods를 이용해서 효율적인 학습을 가능하게 했다. 

 

 

 

주요 takeaways 시사점을 다음와 같다. 

 

• Efficient scaling – 처음으로 라지 스케일의 Pathways을 선보인다. 6144개의 TPU v4 칩으로 540 B 크기의 언어 모델을 학습시킨다. 

• Continued improvements from scaling – 저자들은 PaLM을 몇백개가 넘는 자연어와 코드, 그리고 수학적 추론에서 평가했는데 많은 부분에서 SOTA를 달성했다. 이는 GPT-3, GLaM, Megatron-Tuning NLG, Gopher, Chinchilla, LaMDA와 겨룰만하다.  

• Breakthrough capabilities – 수 많은 어려운 문제들에 대해서 돌파구가 될만한 능력을 보였다. 이는 섹션 6.3에서 제시되는데 multi-step mathematical or commonsense reasoning과 같은 어려운 문제를 포함한다.  

• Discontinuous improvements – Scaling behavior를 더 잘 이해하기 위해서, 서로 다른 세 가지 스케일인 8B, 62B, and 540B의 모델을 실험했다. 8B에서 62B로 키울 때는 Power law 룰 모양에 근거한 일관적인 성능의 향상을 보였다. 하지만 62B를 540B로 늘렸을 때는 급격한 정확성의 향상이 있었다. 

 

• Multilingual understanding – Machine translation, summarization, question answering을 포함한 다양한 분야에서, 다국어에 대해서 성능을 측정했다. 비 영어 데이터가 22%에 불과했지만 퓨샷 성능 평가를 통해서 540B 모델이 비영어 요약 태스크에서

기존의 파인 튜닝된 SOTA와의 성능 격차를 메우고 번역 작업에서는 기존 SOTA를 능가한다. 영어 및 다국어 작업에 대한 다국어 데이터 비중 증가의 영향을 이해하기 위해서는 추가 연구가 필요하다.

 

 

•  Bias and toxicity – 또한 분포 편향과 독성에 대한 모델 성능을 평가하여 몇 가지 통찰력을 얻게 되었다 (섹션 10). 첫째, 성별 및 직업 bias 편향의 경우, Winogender coreference task 상호 참조 과제의 정확도가 모델 규모에 따라 향상되고, PaLM 540B는 원샷 및 퓨샷에서 새로운 SOTA를 제시한다. 둘째, 인종 / 종교 / 성별 prompt continuation 프롬프트 연속 과제에 대해 수행된 co-occurence analysis 동시 발생 분석은 모델이 고정관념을 잘못 긍정할 가능성을 보여준다. 예를 들어, 무슬림을 테러리즘, 극단주의, 폭력과 연관시키게 된다. 이러한 현상은 여러가지 규모의 모델 전반에 걸쳐 일관되게 나타났다. 마지막으로, 프롬프트 연속 과제에 대한 toxicity 독성 분석은 8B 모델에 비해 62B 및 540B 모델의 전반적인 독성 수준이 약간 더 높음을 보여준다. 그러나 모델에서 생성된 연속 문장의 독성은 프롬프트 텍스트의 독성과 높은 상관관계를 보이는 반면, 인간 생성 연속 문장은 독성과 강한 상관관계를 보이지 않았다. 이는 모델이 인간 생성 텍스트보다 프롬프트 스타일의 영향을 더 크게 받는다는 것을 시사한다. 향후 연구에서는 이러한 기준을 영어가 아닌 언어로 확장하고 잠재적 위험을 더욱 철저하게 고려할 계획이다.

 

 

2. Model Architecture

 

PaLM 표준적인 디코더 온리 트랜스포머 구조를 따르지만 다음의 수정을 가했다. 

• SwiGLU Activation –  MLP intermediate activations로 SwiGLU를 사용하였으며 이는 standard ReLU, GeLU, or Swish activations에 비해 좋은 성능을 보여준다. 

• Parallel Layers – Standard “serialized” formulation  대신에 개별 Transformer block에서 “parallel” formulation을 사용한다.

 

Standard “serialized” formulation: 

y = x + MLP(LayerNorm(x + Attention(LayerNorm(x)))

"Parallel” formulation
y = x + MLP(LayerNorm(x)) + Attention(LayerNorm(x))

 

The parallel formulation은 대략 15% 더 빠르게 학습할 수 있다. 8B에서는 약간의 성능 저하가 있지만 62B와 540B의 규모에서는 성능의 저하가 없었다. 

 

• Multi-Query Attention – Query, Key, Value 중에서 K와 V를 공유하여 학습 속도를 높이면서도 성능에 영향을 주지 않는 Multi-head Attention 방법이다. 전에 공부한 포스트 (링크). 

• RoPE Embeddings – RoPE embeddings는 absolute or relative position embeddings에 비해서 long sequence에 대해서 좋은 성능임을 보였기에 채택했다.  

• Shared Input-Output Embeddings – Shared iput and output embedding matrices을 사용한다. 

 

•  No Biases – No biases were used in any of the dense kernels or layer norms. 이는 대규모 모델의 학습에서 안정성을 향상시킴을 발견했다. 

• Vocabulary – SentencePiece를 사용하며 vocabulary 사이즈는 256k tokens이다. 

 

PaLM 모델에 대한 자세한 내용은 아래 Table 1에 소개되어있다. 

 

 

 

3. Training Dataset

고품질의 코퍼스로 구성된 780 billion tokens로 PaLM을 학습했다. 데이터응 Filtered webpages, books, Wikipedia, news articles, source code와 social media converstaions로 구성된다. LaMDA와 GLaM을 학습할 때 썼던 데이터셋에 기반한 데이터셋이다.

 

소스 코드의 경우 공개된 Github에서 크롤링했으며 라이센스에 따라서 쓸 수 없는 내용들은 필터링했다. 24개의 흔히 사용되는 프로그래밍 언어를 포함하며 Java, HTML, Javascript, Python, PHP, C#, XML, C++, and C도 여기에 속한다. 총 196 GB의 소스 코드를 확보했다. Levenshtein distance에 기반하여 중복된 소스 코드 파일을 제거했다.

 

아래 Table 2에 데이터셋에 대한 비중이 나와있다.

 

4. Training Infrastructure

JAX와 T5X, TPU v4 Pods를 이용해서 PaLM을 학습하고 평가했다.

 

 

Pathways는 2개의 Pods이 각각 별도로 데이터를 받은 다음 그라디언트를 계산한다. 그리고 이를 서로 공유하여 파라미터를 업데이트 하게 된다.

 

4. Training Setup

• Weight initialization – The kernel weights는 “fan-in variance scaling”, i.e., W ∼ N (0, 1 / $n_{in}$ )로 초기화되며 Input embeddings E는 N (0, 1)으로 초기화한다. Pre-softmax output logits는 1 / $\sqrt{n}$로 스케일링하며, 이때 $n$은 embedding size다. 

• Optimizer – Adafactor 옵티마이저를 사용한다. 

• Optimization hyperparameters –  Adafactor의 learning rate는 첫 10,000 스텝에서 10**−2 이고 1 / $\sqrt{k}$의 비율로 decay하는데 이때 k는 step number다.  Momentum of $\beta_1$ = 0.9이고 세컨드 오더 모멘트 $\beta_2 = 1.0 - $ k**-0.8 이다. The second-order moment interpolation value is computed as β 2 = 1.0 − k −0.8 . 일반적으로 많이 사용하는 $\beta_2 = 0.99$ 보다 더 안정적임을 확인했다. Global norm gradient clipping을 사용했으며 이때 기준 value는 1.0이다. Dynamic weight decay를 사용했으며 이는 lr**2.0이다. 

 

• Loss function – Standard language modeling loss function인 average log probability of all tokens without label smoothing을 loss funciton으로 사용했다.  Auxiliary loss인 z_loss = 10 −4 · $\text{log}^2$ Z 를 추가해서 softmax normalizer log(Z) 를 0에 근사하게 만들어서 학습을 보다 안정적으로 만든다.

• Sequence length – A sequence length는 2048다. 

• Batch size – 모든 모델에 대해서, 학습 중에 batch size를 증가시킨다. 가장 큰 모델에 대해서 step 50k까지는 batch size 512 (1M tokens)을, 그 다음 115k step까지는 이를 두배로 늘려서 1024 (2M tokens) 으로, 그리고 마지막 255k step까지는 다시 2배로 늘려서 2048 (4M tokens)을 만든다. 더 작은 모델들도 비슷하게 배치사이즈를 키워나간다.두 가지 이유 때문에 이런 배치 사이즈 전략을 구사한다. (1) 더 작은 배치 사이즈는 초기 학습 단계에서 더 샘플 효율적이다 (보여진 토큰의 함수로서 더 나은 loss를 가진다). 반면에 더 큰 배치 사이즈는 후반부의 학습에서 더 gradient estimate를 지닌다. (2) 더 큰 배치 사이즈는 더 큰 행렬 곱 차원을 가지므로 TPU의 효율성을 증가시킨다. 

• Bitwise determinism – 이 모델은 모든 체크 포인트에 대해서 fully bitwise reproducibl다. 다시 말해서 싱글런으로 17,000 step을 학습할 때 체크포인트 15,000에서 다시 시작해도 17,000까지 동일한 결과물을 얻을 수 있다. 이는 다음의 두 가지 방법으로 가능하다. (1) JAX+XLA+T5X로 a bitwise-deterministic modeling framework가 가능하다. (2) a deterministic dataset pipeline 을 이용하는데, 데이터를 셔플할 때 step number에만 의존하는 함수를 이용해서 random access format을 만들면 된다.

• Dropout – Dropout을 사용하지 않는다. 

 

 

 

6. Evaluation

6.1. English NLP tasks

아래의 영어 NLP 태스크들에 대한 벤치마크를 표기한다.

 

 

 

 

 

 

 

 

 

6.2. Big-bench

 

 

 

사람과 Prior SOTA인 GPT-3 175B 1-shot, Gopher 5-shot, Chinchilla 5-shot과 PaLM 5-shot을 비교했다. 

 

 

6.3. Reasoning

Arithmetic reasoning과 Commonsense reasoning을 실험했다.

 

그리고 Chain-of-Thought도 적용했다. 

 

 

 

위 Figure 10은 GSM8K, MAWPS, ASDiv, AQuA, SVAMP, CommonsenseQA, StrategyQA에 대한 실험 결과다.

 

 

 

6.4. Code Tasks

Text-to-code와 Code-to-code의 성능을 평가한다.

 

 

 

코드에 대해서 추가적으로 파인튜닝한 PaLM-Coder도 소개한다.

 

 

DeepFix라는 틀린 코드를 제대로 고치는 데이터셋에 대해서도 테스트했다.

 

 

 

 

6.5. Translation

 

 

 

 

 

6.6. Multilingual Natural Language Generation

Czech (cz), English (en), German (de), Russian (ru), Spanish (es), Turkish (tr), and Vietnamese (vi)의 언어들에 대한 성능을 평가했다. 데이터셋과 그에 해당되는 언어들은 다음과 같다.

 

 

 

 

 

 

6.6. Multilingual Question Answering

TyDiQA-GoldP 벤치마크를 사용해서 다언어 QA 성능을 평가한다.

 

 

 

 

10. Representational Bias Analysis

 

 

 

Toxicity probability of the continuation (TPC)은 Toxicity probability of the prompt (TPP)의 함수인데 TPP가 증가할 수록 TPC도 증가하는 현상을 확인할 수 있다. 따라서 유해한 프롬프트를 작성하지 않도록 하는게 중요하다. 

 

 

 

여담

 

Pathways는 gradient accumulation이 연상되었다.

 

개인적으로는 Bitwise determinism이 흥미로웠는데 오직 스텝 숫자에만 의존하는 함수로 만든 랜덤 함수를 통해서 가능하게 한 재현성 때문이다.  

 

 

 

 

 

References:

https://devhwi.tistory.com/37