Gemma: Open Models Based on Gemini Research and Technology

■ 논문에서는 Gemini 모델을 기반으로 한 오픈 모델 제품군인 Gemma를 소개한다. 

[2403.08295] Gemma: Open Models Based on Gemini Research and Technology

Gemma: Open Models Based on Gemini Research and Technology

 

1. Introduction

■ Gemini에서 영감을 받은 아키텍처, data, training recipes을 사용하여 6T tokens의 텍스트로 Gemma models을 학습시켰다. 

■ Gemma는 automated benchmarks과 human evaluation을 포함한 광범위한 도메인 전반에 걸쳐 평가된다. 

■ Gemini와 마찬가지로, 이 모델들은 text domains에서 강력한 generalist capabilities을 가지며, 해당 규모에서 SOTA의 understanding 및 reasoning skills을 보여준다.  

■ Gemma는 두 가지 크기로 제공된다: GPU 및 TPU에서의 효율적인 배포와 개발을 위한 7B 모델과, CPU 및 on-device applications을 위한 2B 모델  

■ 이 technical report는 Gemma의 model architecture, training infrastructure, 그리고 pretraining and fine-tuning recipes에 대한 내용, 다양한 quantitative 및 qualitative benchmarks뿐만 아니라 standard academic benchmarks과 human-preference evaluations을 통한 평가 결과를 제시한다.  

■ 그 후, safety에 대한 저자들의 접근 방식과 Gemma가 가지는 implications, 모델의 한계 및 장점에 대해 설명한다.  

---

---

2. Model Architecture

■ Gemma model architecture는 transformer decoder를 기반으로 한다. 아키텍처의 파라미터는 Table 1과 같다.

[출처] https://arxiv.org/abs/2403.08295

■ Gemma models은 8192 tokens의 context length로 학습되었다. 

■ original transformer paper 이후에 제안된 몇 가지 개선 사항들을 활용하였다. 

Multi-Query Attention

■ 7B 모델은 MHA를 사용한 반면, 2B는 MQA (num_kv_heads=1)를 사용한다. 이는 small scales에서는 MQA가 잘 작동한다는 ablations 결과에 기반한 결정이다.  

RoPE Embeddings

■ absolute positional embeddings을 사용하는 대신, 각 lyaer에 rotary positional embeddings을 사용한다. 

■ 또한, 모델 크기를 줄이기 위해 inputs과 outputs 간에 embeddings을 공유한다. 

GeGLU Activations

■ GeGLU activation function의 approximated version을 사용한다.  

RMSNorm

■ 학습 안정화를 위해 각 transformer sub-layer( attention layer와 feed-forward layer)의 input을 RMSNorm으로 정규화한다. 

---

---

3. Training Infrastructure

■ TPUv5e를 사용하여 Gemma models을 학습시킨다. TPUv5e는 256개의 칩이 16 x 16 2D 토러스(Torus) 구조로 연결된 'pod' 단위로 운용된다.  

■ 7B 모델의 경우 16개의 pods (=총 4,096개의 TPUv5e), 2B는 2개의 pods (=총 512개의 TPUv5e)으로 pretrain한다. 

■ 7B에 대해서는 16-way model sharding과 16-way data replication을 사용하고, 2B의 경우 단순히 256-way data replication을 사용한다. optimizer state는 ZeRO-3와 유사한 기술을 사용하여 sharded된다.  

■ pod과 pod 사이에는 Pathways 접근법을 사용하여 data-center network를 통해 data replica reduce를 수행한다.  

■ Gemini를 따라, Jax와 Pathways의 'single controller' programming paradigm을 활용한다. 이는 하나의 파이썬 프로세스가 전체 training run을 조율할 수 있게 함으로써 개발 과정을 단순하게 만들어 준다.  

■ 또한 training step computation을 위해 GSPMD partitioner와 MegaScale XLA compiler를 사용한다. 

---

---

4. Pretraining

Training Data

■ Gemma 2B와 7B는 각각 3T와 6T tokens로 학습되었으며, 이는 주로 web documents, mathematics, code로 구성된 English data이다.  

■ Gemini와 달리, 이 모델들은 멀티모달이 아니며, multilingual tasks에서 SOTA 성능을 내도록 학습되지 않았다.  

■ Gemma는 Gemini와의 호환성을 위해 SentencePiece 기반의 토크나이저를 사용한다. 

■ 이 토크나이저는 숫자를 분할(숫자를 통째로 하나의 토큰으로 묶지 않고 개별 숫자 단위로 분할)하고, extra 공백(whitespace)을 보존하며, vocabulary에 없는 unknown tokens에 대해서는 byte-level encodings을 수행한다(즉, unknown token을 [UNK]와 같은 special token으로 처리하지 않고, 바이트 단위로 분해하여 인코딩한다). 
- 파이썬 같은 프로그맹 언어는 코드 사이 공백이나 들여쓰기(즉, extra whitespace로 볼 수 있음)가 문법적으로 매우 중요하므로, 코딩 성능을 고려해 이러한 설정을 적용한 것으로 보인다. 

■ Gemma의 vocabulary는 Gemini의 vocabulary를 사용한다. Gemini의 vocabulary는 많은 언어들을 처리하도록 설계되었으며, vocabulary의 size는 256K tokens이다. 

■ 그래서 Table 2와 같이 하나 또는 몇 개의 언어에만 국한된 모델들에 비해 embedding parameters의 수가 더 많다. 

[출처] https://arxiv.org/abs/2403.08295

Filtering

■ 원치 않거나 안전하지 않은 발언의 위험을 줄이고, 특정 개인 정보나 기타 민감한 데이터를 걸러내기 위해 pre-training dataset을 필터링한다.  

■ 이 필터링 과정에는, 유해하거나 품질이 낮은 콘텐츠를 제거하기 위한 heuristics과 model-based classifiers 사용이 모두 포함된다.  

■ 더 나아가, pre-training data mixture에서 모든 evaluation sets을 필터링하고, evaluation set의 leakage를 확인하기 위해 contamination analysis를 수행한다. 

■ 그리고 모델이 민감한 개인정보나 데이터를 그대로 출력하는 현상을 막기 위해, 그런 정보를 결과물로 나오지 않도록 민감한 출력을 억제하여 recitation(즉, 암기)의 위험을 줄인다.  

■ final data mixture는 2B 및 7B 모델 모두 ablations을 통해 결정하였다.  

■ Gemini Team에서 제안한 것과 유사하게, 학습 후반부로 갈수록 관련성이 높고 고품질인 데이터의 가중치(즉, 비중)를 높이기 위해 training 과정 전체에 걸쳐 corpus의 mixture를 변경하는 stage training을 수행한다.  

---

---

5. Instruction Tuning

■ Gemma 2B와 7B에 대해 '텍스트 전용(text-only), English-only의 synthetic' 및 '사람이 생성한 prompt response pairs'의 혼합을 사용하여 supervised fine-tuning (SFT)을 수행한다.  

■ 또한, 라벨링된 English-only preference data로 학습된 reward model과 고품질 prompts의 set에 기반한 policy를 사용하여 RLHF를 수행한다.  

■ downstream 자동 평가와 model outputs에 대한 human preference evaluations에서 향상된 성능을 얻기 위해 두 단계(SFT, RLHF)가 모두 중요하다는 것을 발견했다고 한다.  

Supervised Fine-Tuning

■ LM 기반의 side-by-side evaluations (LLM-as-a-Judge)를 기반으로 supervised fine tuning을 위한 data mixtures을 선택했다.  

■ held-out prompts의 set이 주어지면, test model에서 responses을 생성하고, baseline model에서 동일한 prompts에 대한 responses을 생성한 다음, 이를 무작위로 섞은 후, 더 크고 고성능인 모델에게 두 responses 사이의 preference를 표시하도록 하였다.  

■ 이때 instruction following, factuality, creativity, safety와 같은 specific capabilities을 위해 서로 다른 prompt sets을 사용하였다.  

■ 그리고 LM-based judge를 위해, human preferences와 aligned되기 위해 CoT prompting, rubrics and constitutions과 같은 프롬프팅 전략들을 사용하였다.

Filtering

■ synthetic data를 사용할 때, 여러 단계의 필터링을 실행하여 특정 개인 정보, unsafe하거나 toxic한 model outputs, 잘못된 자아 식별(self-identification) data (예: Google이 만든 모델인데 "나는 OpenAI 모델입니다"라고 하는 오류), 중복된 examples을 제거하였다.  

■ 저자들은 Gemini를 따라, hallucinations을 최소화하기 위해 더 나은 in-context attribution(답변의 근거를 문맥 내에서 명시하도록), hedging(확신하지 않고 조심스럽게 말하는 화법), 거절(refusals)을 encourage할 수 있는 data를 포함시켰을 때, 다른 평가 지표에서의 성능 저하 없이 factuality 지표의 성능을 향상시킨다는 것을 발견하였다.  

■ final data mixtures과 하이퍼파라미터를 포함한 supervised finetuning recipe는 safety와 hallucinations과 관련된 model의 harms를 최소화하면서 helpfulness를 향상시키는 것을 기준으로 선택하였다.  

Formatting

■ Instruction tuned된 models은 training 및 inference time에서 모든 instruction tuning examples에 추가 정보(메타데이터)를 annotate하는 specific formatter와 함께 학습된다.  

■ 이 formatter는 토크나이저의 special control tokens(Table 3의 <start_of_turn>, <end_of_turn>, user, model)을 사용하며, 다음 두 가지 기능을 수행한다. 
- (1) 사용자 역할과 같은 대화 내의 역할(roles)을 표시하는 것
- (2) multi-turn 대화에서 대화의 turns을 구분하는 것 

■ formatter 없이도 일관된 생성을 얻는 것이 가능할 수는 있지만, 이는(즉, formatter 없이 그냥 text만 넣으면) 모델이 학습한 패턴(즉, 학습한 distribution)과 다르기 때문에 모델에게는 out-of-distribution인 입력이 되며, 더 나쁜 생성 결과를 낳을 가능성이 매우 높다.  
- 모델은 specific formatter를 사용한 input format으로 학습되었기 떄문에, 생성 시 해당 formatter를 사용하지 않은 input을 주면, 이는 모델이 학습한 data distribution과 다른 형태(즉, out-of-distribution)가 된다.  

■ control tokens과 대화 예시는 Table 3, 4에서 확인할 수 있다. 

[출처] https://arxiv.org/abs/2403.08295

Reinforcement Learning from Human Feed back

■ SFT model에 RLHF를 사용하여 추가로 finetuning한다. Gemini와 유사하게, human evaluators로부터 선호도 쌍(pairs of preferences)을 수집하고 Bradley-Terry model 하에서 reward function을 학습시킨다. 
- A와 B라는 두 개의 responses가 있다고 하자. human evaluators은 두 responses에 대해 preference를 부여할 것이다.   
- human이 부여한 preference를 학습시키기 위해, 저자들은 Bradley-Terry model을 사용했는데, 간단히 설명하면 A / (A+B)로 A가 B보다 선택될 확률을 계산한다. 이 값을 reward로 사용한 것이다.  

■ policy는 새로운 강화학습 알고리즘을 사용하여 이 reward function을 최적화하도록 학습시킨다. 

■ SF그리고 하이퍼파라미터를 튜닝하고 reward hacking을 완화하기 위해, 고성능 모델을 automatic evaluator로 사용하여 baseline models과 비교한다.  

---

---

6. Evaluation

■ automated benchmarks와 human evaluation를 모두 사용하여 다양한 도메인에 걸쳐 Gemma를 평가한다. 

Human Preference Evaluations

■ standard academic benchmarks 외에도, human evaluation을 통해 Gemma models과 Mistral v0.2 7B Instruct model을 비교하였다.  

■ creative writing tasks, coding, following instructions 전반에 걸쳐 모델이 instructions을 따르도록 요청하는 데 중점을 둔 약 1,000개의 held-out prompt collection에서, Gemma 7B IT(Instruction Tuned)는 Mistral v0.2 7B Instruct 대비 61.2%의 승률을, Gemma 2B IT는 45%의 승률을 기록했다. 

■ 그리고 basic safety protocols을 테스트하는 데 중점을 둔 약 400개의 held-out prompt collection에서는 Gemma 7B IT가 63.5%의 승률을, Gemma 2B IT가 60.1%의 승률을 기록했다.  

[출처] https://arxiv.org/abs/2403.08295

Automated Benchmarks

■ physical reasoning, social reasoning, question answering, coding, mathematics, commonsense reasoning, language modeling, reading comprehension 등을 포함한 다양한 도메인에서 Gemma models을 평가한다. 

[출처] https://arxiv.org/abs/2403.08295

■ Table 6, 7은 academic benchmarks에서 Gemma 2B 및 7B와 open-source LLM들을 비교한 결과이다. 

■ Table 6의 MMLU를 보면, Gemma 7B는 동일하거나 더 작은 크기의 모델뿐만 아니라 더 큰 모델도 능가하는 것을 볼 수 있다. 그러나, 해당 벤치마크 저자들에 따르면 human expert의 성능은 89.8%라고 한다. 

■ 당시 Gemini Ultra는 해당 임계값을 처음으로 넘은 모델이었으며, Gemma는 Gemini 수준이나 인간 수준의 성능에 도달하기까지는 여전히 개선의 여지가 남아 있음을 보여주는 결과이다. 

■ Gemma models은 mathematics 및 coding benchmarks에서 뛰어난 성능을 보여주는데, GSM8K와 MATH에서 Mistral 7B와 최소 10점 이상 차이가 나는 것을 볼 수 있다. HumanEval에서도 최소 6점 이상 차이가 난다. 

■ 심지어 MBPP에서 코드에 특화된  CodeLLaMA-7B 모델의 성능까지 넘어섰다고 한다. (CodeLLaMA는 41.4%, Gemma 7B는 44.4%)

Memorization Evaluations

■ aligned model이라도, alignment을 우회할 수 있는 adversarial attacks에 취약할 수 있음을 보여준 연구가 있다. 이러한 공격들은 모델을 alignment에서 이탈하게 만들거나, 그 과정에서 memorized된 training data를 그대로 출력하게 만들 수 있다.  

■ 저자들은 discoverable memorization에 초점을 맞췄는데, 이는 model의 memorization에 대한 합리적인 상한선 역할을 하며, 여러 연구들(예: Pythia)에서 사용된 방법이다. 

■ 이전 연구에서 수행한 것과 동일한 방법론으로 Gemma pretrained models의 memorization 여부를 테스트하였다. 

■ 측정 방법은, 각 corpus에서 10,000개의 documents을 샘플링하고 처음 50개의 tokens을 모델의 프롬프트로 사용한다. 앞부분의 50개 tokens을 모델에게 주었을 때, 모델이 생성한 뒷부분의 50개의 tokens이 텍스트의 실제 내용(ground truth)과 정확히 일치할 경우, 모델이 해당 텍스트를 memorized한 것으로 분류한다.  

■ 저자들은 두 가지 기준을 설정했는데, 하나는 exact memorization으로 모델이 생성한 텍스트가 원본 텍스트와 정확히 일치하는 경우(즉, 토씨 하나 안 틀리고 똑같이 생성한 경우)이며,  

■ 다른 하나는 10%의 edit distance 임곗값을 사용하는 approximate memorization이다. 즉, 10% 정도 틀리더라도 거의 비슷하게 생성하는 경우를 확인하기 위한 것이다. 이를 통해, 패러프레이징(단어 몇 개만 바꾼 경우 및 유사한 다른 단어를 사용한 경우)도 잡아낼 수 있다.  

Verbatim Memorization

[출처] https://arxiv.org/abs/2403.08295

■ Fig 2는 Gemma models과 유사한 크기의 PaLM 및 PaLM 2 model과 비교한 결과이다.

■ PaLM 2는 PaLM과의 비교를 위해 두 모델이 공통으로 학습한 데이터만을 사용하여 memorization을 측정했다.

■ 그러나, Gemma의 pretraining data와 PaLM 모델들은 다른 데이터셋으로 학습했기 때문에 중복되는 data가 훨씬 적다. 그러므로 공통 부분만 가지고 테스트하면, Gemma가 memorization을 거의 하지 않는 것처럼 수치가 왜곡될 수 있다. (Fig 2 left) 

■ 그래서 저자들은 더 신뢰할 수 있는 추정치를 얻기 위해, 전체 pretraining dataset에 걸쳐 "total memorization"을 측정하였다. Gemma도 PaLM과 비슷한 비율로 training data를 memorize한다는 것을 볼 수 있다. (Fig 2 right) 

■ Gemma와 유사한 크기의 PaLM과 Gemma가 비슷한 " total memorization"을 보인다는 결과는, Gemma가 특별히 training data를 더 많이 외우는 것도 아니고, 너무 못 외우는 것도 아닌, 모델 크기에 비례하는 정상적인 수준의 memorization을 가지고 있는 것으로 해석할 수 있다.  

Personal Data

[출처] https://arxiv.org/abs/2403.08295

■ 중요한 것은 personal data가 memorized될 가능성이다. 

■ 저자들은 Gemma pre-trained models을 safe하고 reliable할 수 있게 만들기 위해, training sets에 특정 개인 정보 및 기타 민감한 데이터를 필터링하는 데 Google Cloud Sensitive Data Protection을 사용했다.  

■ 이 tool은 많은 범주의 개인 데이터(예: 이름, 이메일 등)를 기반으로 세 가지 심각도 레벨을 출력한다. 저자들은 가장 높은 심각도를 "sensitive"로, 나머지 두 개는 "personal"으로 분류하였다. 

■ 그런 다음, memorized outputs 중 얼마나 많은 것이 sensitive하거나 personal한 data를 포함하고 있는지 측정하였다. 

■ Fig 3에서 볼 수 있듯이, memorized된 sensitive data가 단 한 건도 발견되지 않았다. 

■ 단, 위 기준에 따라 "personal"로 분류된 일부 데이터를 모델이 memorize한다는 것은 발견했지만, 그 비율은 매우 낮다.

■ 그리고 이러한 tool들은 contex를 고려하지 않고 패턴 매칭에만 의존하기 때문에 fals positives이 많다. 그러므로 이 결과는 과대평가되었을 가능성이 있다. 즉, tool을 사용해 발견한 개인정보의 양은 결과의 수치보다 훨씬 적을 것으로 추정된다.  
- tool이 문맥을 보지 않고, 단순히 "이메일 형식(@)"이라는 패턴만 보고 ture로 잘못 판단하는 경우가 많다는 것이다. 

Approximate Memorization

[출처] https://arxiv.org/abs/2403.08295

■ Fig 4의 결과는, 이전 섹션에서 저자들이 edit distance 10%로 설정한 approximate memorization 측정 결과이다. 

■ exact하게 memorized하는 것보다, 대략 50% 더 많은 데이터가 approximately memorized된다는 것을 볼 수 있으며, 이 비율은 데이터의 종류(코드, 위키피디아, 과학 논문 등)에 상관없이 일정하게 나타나는 것을 볼 수 있다.  

■ 즉, 모델이 data를 토씨 하나 안 틀리고 외우기도 하지만, 그보다 더 넓은 범위에서 패턴과 의미를 중심으로 느슨하게 기억하는 경향이 있다는 것을 보여주는 결과이다. 

---