はじめに
大規模言語モデル（LLM）は、自然言語処理（NLP）の分野で重要な進歩を遂げています。しかし、これらのモデルのサイズと複雑さは、リソース制約のある環境では利用が困難になっています。
この課題に対処するため、軽量なLLMの開発が進んでいます。モデルが軽量であればメモリや計算能力に制限のあるエッジデバイスに配置でき、インターネット接続なしでリアルタイムなLLMの実行が可能です。
本記事では、軽量なLLMの一つであるTinyLlamaについて説明します。
TinyLlamaの概要
TinyLlamaは、Llamaとまったく同じアーキテクチャとトークナイザーを持つ1.1Bのモデルです。TinyLlamaにはTinyLlama/TinyLlama-1.1B-intermediate-step-480k-1Tを筆頭に現在12のモデルが存在しています。
TinyLlamaの詳細
アーキテクチャ
アーキテクチャについては以下です。(GitHub情報)
| 設定 | 説明 |
|---|---|
| Parameters | 1.1B |
| Attention Variant | Grouped Query Attention |
| Model Size | Layers: 22, Heads: 32, Query Groups: 4, Embedding Size: 2048, Intermediate Size (Swiglu): 5632 |
| Sequence Length | 2048 |
冒頭でも述べましたが。TinyLlamaは1.1Bと非常に軽量です。4ビット量子化されたTinyLlama-1.1Bは、550MBのRAMで動くのが特長です。
データセット
訓練に用いたデータセットについては以下です。(GitHub情報)
| 設定 | 説明 |
|---|---|
| Training Data | Slimpajama & Starcoderdata |
| Data Preprocessing | Excluded GitHub subset of Slimpajama; Sampled all code from Starcoderdata |
| Combined Dataset Size | Around 950B tokens |
| Total Tokens During Training | 3 trillion (最終目標) |
| Natural Language to Code Ratio | 7:3 |
2023年11月17日時点では1.5Tのデータでの学習が完了しています。
高速化
TinyLlamaでは上述の通り、学習データが膨大なため高速化の工夫を取り入れています。
- multi-gpu and multi-node distributed training with FSDP
- flash attention 2 [1]
- fused layernorm [1]
- fused swiglu [2]
- fused cross entropy loss [1]
- fused rotary positional embedding [1]
以下の表より訓練速度について、他の同サイズのモデルと比べて高速化できていることがわかります。
表. A100 GPUで300Bトークンにかかる時間
| モデル名 | トレーニング時間 |
|---|---|
| TinyLlama-1.1B | 3456時間 |
| Pythia-1.0B | 4830時間 |
| MPT-1.3B | 7920時間 |
性能
GPT4All Benchmarksと同様のデータセットで評価が行われています。
以下タスクについて説明します。
- HellaSwag
コモンセンス推論のタスクで、AIに4つの選択肢の中から最も論理的な続きを選ばせるものです。各選択肢は、与えられた文脈に基づいて物語や日常のシナリオを完成させることを目的としています。 - Obqa (Open Book Question Answering)
「オープンブック質問応答」とも呼ばれ、一般的な知識を用いて質問に答えるタスクです。このタスクは、外部情報源へのアクセスなしに、質問に対する答えを見つける能力をテストします。 - WinoGrande
Winograd Schema Challengeに似た大規模なコモンセンス推論タスクで、モデルが文脈に基づいて代名詞の参照を解決する能力を評価します。 - ARC (AI2 Reasoning Challenge)
ARCは2つの部分に分かれています。ARC_c (Challenge) と ARC_e (Easy) です。これらは科学的推論のための問題集で、生徒が中学レベルの科学の質問に答える能力を模倣することを目的としています。 - BoolQ (Boolean Questions):
単純な「はい」または「いいえ」で答えられる質問に答えるタスクです。質問は自然言語であり、しばしば文脈的な理解が必要です。 - PiQA (Physical Interaction Question Answering)
物理的相互作用に関する問題に答えるタスクです。モデルは、日常生活の物理的なシナリオに関する質問に答えることが求められます。
ベンチマークの結果は以下です。
| Model | HellaSwag | Obqa | WinoGrande | ARC_c | ARC_e | boolq | piqa | avg |
|---|---|---|---|---|---|---|---|---|
| TinyLlama-1.1B-intermediate-step-715k-1.5T | 53.68 | 35.20 | 58.33 | 29.18 | 51.89 | 59.08 | 71.65 | 51.29 |
| TinyLlama-1.1B-Chat-v0.4 | 58.59 | 35.40 | 58.80 | 30.80 | 54.04 | 57.31 | 71.16 | 52.30 |
| Llama 7B(参考) | 73.1 | 77.4 | 73.0 | 66.9 | 52.5 | 41.4 | 42.4 | 61.0 |
| Llama 13B(参考) | 68.5 | 79.1 | 76.2 | 70.1 | 60.0 | 44.6 | 42.2 | 63.0 |
まだ平均のスコアを見ると、Llama 7B(参考)にも性能は及ばないですが、BoolQやPiQAについてはTinyLlamaの方が精度が高いことがわかります。
終わりに
本記事では、TinyLlamaという軽量なLLMについて解説しました。
少しでも参考になれば幸いです。