바야흐로 대규모 언어모델의 시대이다. 이번에 리뷰할 논문은 2022년 4월에 구글에서 공개한 따끈따끈한 모델, PaLM(Pathways Language Model)이다.

GPT3가 공개된 이후, 많은 대규모 autoregressive 언어 모델들이 공개되었다. 가장 대표적인 post-GPT3 모델들은, GLaM (2021), Gopher (2021), Chinchilla ( 2022), Megatron–Turing NLG (2022), LaMDA (2022)인데, 이들은 다음 4개의 흐름을 가진다.

1. 모델 사이즈 scaling (depth & width)
2. 데이터 추가
3. cleaner & diverse 데이터
4. model capacity 증가

PaLM도 이 흐름에 편승한다. PaLM은 한 마디로 다음의 특징을 가지는 모델이다.

540B parameter, densely activated, autoregressive Transformer, 780 billion tokens, high-quality text

GPT3가 175B 파라미터였다는 것을 생각해보면, 구글이 얼마나 거대한 모델을 만든 것인지 비교 가능하다.

이 거대한 모델은 Pathways라고 명명한 새로운 ML system에 의해 학습 가능해졌는데, 이는 매우 큰 인공신경망을 두 대의 TPU v4 Pods(하나당 4096개 chip), 구체적으로는 6144 TPU v4 chips를 이용해 학습의 효율성을 올린다.

해당 논문의 주요 의의는 다음과 같다:

1. Efficient scaling

   • 모델들은 일반적으로 셋 중 하나로 학습된다
     • single TPU system
     • pipeline parallelism to scale across GPU clusters
     • multiple TPU v3 pods with a maximum scale of 4096 TPU v3 chips
   • PaLM은 최초로 Pathways를 사용한다.
     • pipeline-free training
     • 6144 chips across two TPU v4 Pods

2. Continued improvements from scaling

   • natural language, code, mathematical reasoning 태스크들에서 sota!

3. Breakthrough capabilities

   • NLU & NLG에서 뛰어난 성능을 보인다.
   • 특히 multi-step mathematical/commonsense reasoning이 요구되는 tasks에서 좋은 성능 보인다.
   • 만약 모델이 chain-of-thought prompting과 합쳐진다면, 단순한 few-shot evaluation이 finetuned sota를 능가하는 성능을 낼 수 있다.
     

4. Discontinuous improvements

   • 일반적으로 power law의 관계를 가진다고 알려져 있는데, 이와 달리 8B -> 62B에 비해 62B -> 540B에서 엄청난 성능 향상이 이루어졌다는 것이 관찰되었다.

5. Multilingual understanding

   • 전체 데이터셋에서 비영어 데이터의 비중은 22%인데, 이것들이 다양한 multilingual benchmark들에서 얼만큼의 성능을 보이는지 확인함.

6. Bias and toxicity

   • gender and occupation bias
     • Winogender coreference task improves with model scale
   • co-occurence analysis
     • stereotypes 확인됨
   • toxicity analysis
     • 8B에 비해 62B and 540B 모델에서 higher toxicity

우리는 rough하게 다음의 세 가지를 이해하면 PaLM을 안다고 이야기할 수 있다.

1. Model architecture
2. Pathways
3. Results

이 모든 것들을 한 번에 이해하기는 조금 어려우므로, 이번 내용은 2번에 걸쳐서 진행하려고 한다. 오늘 이야기할 것은 모델 부분이다.

Model

1. Hyperparameters

먼저, 학습된 모델들의 하이퍼 파라미터는 다음과 같다. 배치 사이즈를 학습을 진행시킴에 따라 점진적으로 증가시켰다.

이때, 이 모델은 standard dense Transformers이기 때문에, 토큰 당 FLOPs는 파라미터의 수와 같다.

2. Training Dataset

학습 데이터로 social media conversations, webpages,books, Wikipedia, news, source code를 사용했다. 총 780 billion으로 구성되어 있으며, 모두 깔끔하게 클렌징되었다는 특징이 있다.

학습 데이터에서 특이한 점은 자연어 이외에 프로그래밍 코드를 학습시켰다는 것이다. 해당 코드들은 깃헙에서 긁어왔다.

3. Architecture

PaLM은 gpt 시리즈와 동일하게 Transformer의 decoder의 변이이다. 변화된 부분은 다음과 같다:

• SwiGLU Activation
  • MLP intermediate activation으로 SwiGLU activations 사용
  • standard ReLU, GeLU, Swish activation 보다 확연히 좋은 성능
• Parallel Layers
  • Transformer block에 “serialized” formulation이 아닌 “parallel” formulation 적용
• Multi-Query Attention
  • key/value projections -> 모든 head가 공유
    • (head_num, head_size) => (1, head_size)
    • neutral effect on model quality & training speed
• RoPE Embeddings
  • long sequence에 더 유리함
• Shared Input-Output Embeddings
• No Biases
  • bias term을 없앰 -> 큰 모델의 학습 안정성을 높임
• Vocabulary
  • SentencePiece

논문에서는 가볍게 언급하고 넘어갔지만, 우리는 하나 하나 이것이 무엇인지 확인하고 넘어가보자. 하단의 세 가지는 안다고 생각해 넘어간다.

3.1 SwiGLU Activation

GLU Variants Improve Transformer

스탠다드한 트랜스포머의 FFN 레이어는 다음과 같이 생겼다:

$FFN(x, W1, W2, b1, b2) = max(0, xW1 + b1)W2 + b2$

즉, ReLU 함수를 activation으로 사용한다. T5에서는 여기서 bias 텀을 없애 사용한다.

$FFN_{ReLU}(x, W1, W2) = max(0, xW1)W2$

그런데, 최근 ReLU가 아닌 더 좋은 activation function을 적용하여 트랜스포머의 성능을 높이기 위한 연구들이 진행되었다. 제안된 activation function들은 다음과 같다:

$GELU(x) = xΦ(x)$
$Swish_β(x) = xσ(βx)$

$Bilinear(x, W, V, b, c) = (xW + b) ⊗ (xV + c)$
$GLU(x, W, V, b, c) = σ(xW + b) ⊗ (xV + c)$
$ReGLU(x, W, V, b, c) = max(0, xW + b) ⊗ (xV + c)$
$GEGLU(x, W, V, b, c) = GELU(xW + b) ⊗ (xV + c)$
$SwiGLU(x, W, V, b, c, β) = Swish_β(xW + b) ⊗ (xV + c)$

우리의 목표는 SwiGLU를 이해하는 것이므로, swish와 GLU에 대해 더 자세히 알아보자.

3.1.1 Swish

$Swish_β(x) = xσ(βx)$

Swish 함수는 수식을 통해 확인할 수 있듯, 입력 x와 x에 시그모이드를 거친 값을 곱해준 형태이다. 이때 곱해지는 값 β는 학습 가능한 파라미터이다. 만약 β가 0이라면 Swish 함수는 identity function이 되고, β가 무한대라면 Swish 함수는 ReLU가 된다. 그래프를 통해 확인해보자.

Swish 함수는 unbounded above, bounded below 속성을 가진다는 점에서 ReLU의 장점을 그대로 가진다. 또한, 모든 점에서 미분 가능하며 ReLU와 다르게 단조 증가 함수가 아니라는 점에서 추가적인 장점을 가진다.

3.1.2 GLU

Gated Linear Units (GLU)는 입력에 대한 두개의 선형 변환의 결과를 component-wise product한 신경망 레이어인데, 이 중 하나는 시그모이드를 거친다.

$GLU(x, W, V, b, c) = σ(xW + b) ⊗ (xV + c)$

여기서, $W$와 $V$는 학습 가능한 가중치들이다.
GLU는 activation function을 거치는 단순 linear model $σ(xW + b)$에 $(xV + c)$를 component별로 곱해준 형태이다. 즉, 후자는 filter로의 기능을 수행한다. 즉, filter내의 값이 어떻게 구성되었는지에 따라 어떤 엔트리들은 더욱 영향력이 강해지고, 어떤 것들은 미미해진다.

3.1.3 SwiGLU

따라서, 우리가 최종적으로 알고자 하는 SwiGLU는 다음의 형태를 가진 비선형적 신경망이다.

$SwiGLU(x, W, V) = Swish_1(xW) ⊗ xV$
$SwiGLU(x, W, V) = xWσ(xW) ⊗ xV$

저자들은 activations이 기존의 ReLU, GeLU, Swish activations보다 훨씬 성능이 좋다는 것을 강조한다.

3.1.4 FFNs

$FFN_{ReLU}(x, W_1, W_2) = max(xW_1, 0)W_2$
$FFN_{GLU}(x, W, V, W_2) = (σ(xW) ⊗ xV )W_2$
$FFN_{Bilinear}(x, W, V, W_2) = (xW ⊗ xV )W_2$
$FFN_{ReGLU}(x, W, V, W_2) = (max(0, xW) ⊗ xV )W_2$
$FFN_{GEGLU}(x, W, V, W_2) = (GELU(xW) ⊗ xV )W_2$
$FFN_{SwiGLU}(x, W, V, W_2) = (Swish_1(xW) ⊗ xV )W_2$

SwiGLU 논문에서 비교한 FFN 레이어의 변형이다. 이때, original FFN에서는 두 개의 가중치 행렬을 사용하지만, 나머지에는 세 개의 행렬을 사용한다는 점이 특징이다. 학습해야하는 행렬의 수를 늘려도 computation은 일정하게 가져가기 위해, hidden units의 개수는 2/3으로 줄인다.

다음은 computation과 performance 측면을 평가해본 결과이다.

• Computation
  

• Performance
  

즉, SwiGLU를 사용한 MLP는 성능도 좋은데 속도까지 빠르다. 이 때문에 PaLM의 저자들은 SwiGLU를 이용해 비선형성을 추가한다.

3.2 Parallel formulation

GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model

PaLM은 기존의 “serialized” formulation을 사용하지 않고, 각 트랜스포머 블락에 “parallel” formulation을 적용한다.

• serialized formulation
  $y = x + MLP(LayerNorm(x + Attention(LayerNorm(x)))$

• parallel formulation
  $y = x + MLP(LayerNorm(x)) + Attention(LayerNorm(x))$

즉, 어텐션 블락을 거친 후 MLP를 거치는 방식이 아니라, MLP와 어텐션을 병렬적으로 처리하겠다는 것이다.

이러한 방식은 오로지 모델 병렬화를 통한 학습의 효율 증가를 목적으로 한다. 모델을 병렬화하기 때문에 forward/backward시 all-reduce가 필요하다. all-reducee란 심플하게 해당 블록의 결과물을 모든 프로세서가 Return하는 것을 의미한다.

Megatron-LM의 model parallelism을 확인해보면 대략적인 감이 잡힌다.

즉, 각자 forward/backward 시키고 결과물을 다른 프로세서들과 공유하여 모든 프로세서로 하여금 결과물을 반환시킨다고 생각하면 된다.

이 방식으로, 대략 15%의 throughput 향상이 이루어진다고 한다!

3.3 Multi-Query Attention

Fast Transformer Decoding: One Write-Head is All You Need

해당 논문의 저자들은 Transformer inference의 속도가 "keys"와 "values" 텐서들을 reload하는데 필요한 memory bandwidth, 즉 거대한 텐서를 전달하기 위한 메모리 전달 속도 때문에 제한된다고 주장한다. 이 때문에, 헤드별로 공유된 key와 value를 제안한다.

Multi-query attention은 각 어텐션 레드들이 단일 key, value 값을 가진다는 것 외에는 기본적인 Multi-head attention과 동일하다.

head 개수를 $k$ 크기를 $h$라고 했을 때, Multi-head attention의 key, value, query는 모두 $(k,h)$의 크기를 가진다. 하지만 Multi-query attention에서는 key, value를 공유하기 때문에, key, value=$(1, h)$, query=$(k,h)$의 크기를 가지게 된다.

이러한 방식은 학습 속도와 모델 퀄리티에는 영향이 없었지만, autoregressive decoding을 할 때 굉장한 속도 개선이 있었다. 표를 확인해보자.

속도가 개선된 이유는, multi-headed attention에서는 key/value 텐서가 예제 간에 공유되지 않고 한 번에 하나의 토큰만 디코딩되어, auto-regressive decoding시 accelerator의 사용 효율성이 떨어지기 때문이다.

3.4 RoPE Embeddings

RoFormer: Enhanced Transformer with Rotary Position Embedding

RoPE embeddings (2021)은 absolute/relative position embeddings에 비해 긴 시퀀스에 대해 강건하다고 알려져있다. RoPE embeddings의 목표는 다음의 수식을 해결하는 것이다.

$<f_q(x_m, m), f_k(x_n, n)> = g(x_m, x_n, m − n)$

즉, 단어 $x_m, x_n$과 둘의 relative position($m − n$)이 input으로 들어가는 함수 $g$와 동일해지는 함수 $<f_q(x_m, m), f_k(x_n, n)>$를 찾아보자는 것이 해당 임베딩의 목표이다.

저자들은 다음의 식들을 제안한다.

$f_q(x_m, m) = (W_qx_m)e^{imθ}$
$f_k(x_n, n) = (W_kx_n)e^{inθ}$
$g(x_m, x_n, m − n) = Re[(W_qx_m)(W_kx_n)^∗e^{i(m−n)θ}]$

• $Re[·]$: real part of a complex number
• $(W_kx_n)^∗$: $(W_kx_n)$의 켤레복소수
• $(e^{iθ}) ^m = e^{imθ} = \cos(mθ) + i\sin(mθ)$

최종적인 공식은 다음과 같다:

$f_{q|k}(x_m, m) = R^d_{Θ,m}W_{q|k}x_m$

즉, 해당 단어와 key 혹은 query와 R matrix의 선형결합으로 이루어진다. 이때, R은 다음과 같이 생겼다.

즉, $R^d_{Θ,m}W_{q|k}x_m$은 아핀 변환(Wx+b)된 워드 임베딩 벡터를 그것의 위치 인덱스의 배수 각도의 양만큼 회전한 것이다.

RoPE를 self-attention에 적용하면 다음과 같이 작성할 수 있다:

이때, $R^d_{Θ,n−m} = (R ^d_{Θ,m})^TR^d_{Θ,n}$이 되게 되는데, $R$이 직교 행렬이기 때문에 위치 정보를 인코딩하는데의 안정성을 보장할 수 있다.

---

(2)편에서는 초대규모 언어모델의 효율적인 학습을 가능하게 한 Pathways와 그를 이해하기 위해 필요한 선행 지식들에 대해 공부할 것이다. 또한, 모델을 실험하기 위해 사용된 새로운 벤치마크 및 방법들과, few-shot의 성능을 크게 올린 chain-of-thought prompting 등에 대해 공부할 것이다.