LLMを幅方向にUp-Scalingする方法を考えて試してみる
1人ローカルLLMアドベントカレンダーの7日目です。
ローカルLLMやマルチモーダルモデルの学習やデータセット周りについて書いていく予定なので、興味がある方は明日以降も読んでいただけると嬉しいです！
要約
- 幅方向へのUp-Scalingを行うことでperplexityが数千倍に増加
- 1B未満の事前学習で日本語を話せるようになるが、ベンチマークスコアは低いまま
目的
数日前にDepth Up-Scalingについて書きました。これは既存のモデルのレイヤーを複製して深さ方向にモデルを大きくする方法です。
この手法は大規模なモデルを効率よく開発することができるため、個人的には非常に有用なものだと考えています。
しかし、深さ方向にスケールアップすることにも課題があり、NVIDIAのpruningに関する報告からもわかるように、推論速度に与える影響は幅より深さのほうが大きく、同じ4Bモデルでも深さ方向でpruningしたモデルのほうがThroughputが高くなっています。
How to Prune and Distill Llama-3.1 8B to an NVIDIA Llama-3.1-Minitron 4B Modelより
そこで、幅方向にスケールアップする方法があれば、バランスよくモデルの大型化をしていけるのではないかと思い、その方法を探っていきたいというのが今回の趣旨になります。
Width Up-Scaling
Up-Scalingの方法
軽く調べた限りでは幅方向へのUp-Scalingに関する文献が見つからなかったので、とりあえず思いついた方法を試していきます。
思いついたのは以下の3パターンです。
- 元の重みをコピーする
- 双線形補間で重みを広げる
- 特異値分解を行い、特異行列を総線形補間で拡張
1番の方法は、サイズを大きくしたレイヤーに元の重みをコピーして、残りの部分も元の重みを分割していく方法です。
2番は、行列を引っ張って広げたように拡張するイメージです。
3番は、特異行列はそのままに、特異値ベクトルを拡張しています。意味があるかは不明です。
これらの方法は理論的に求めるべきものだと思いますが、今回は数学的な厳密性を気にせず、とにかくモデルを拡張できればOKという気持ちで考えています。
3通りのUp-Scalingについて、関数を以下のリポジトリに上げてます。整理する時間がなく、関数部分をpyファイルに貼り付けただけになってますので大変分かりづらいです。アドカレ終わって気が向いたら整理します...
Up-Scaling後のperplexity
上記3つの方法でllm-jp/llm-jp-3-1.8bのhidden sizeとintermediate sizeが1.1倍になるようにスケールアップをしました。
また、それらのモデルに対して、「日本で一番高い山は富士山で、標高は3776mあります。」という文章に対するperplexityを算出します。
| モデル | perplexity |
|---|---|
| llm-jp/llm-jp-3-1.8b | 8.01 |
| 重みコピー | 2078084.25 |
| 双線形補間 | 312985.03125 |
| 特異値分解 | 221372.671875 |
Up-Scale後はperplexityが大幅に上昇しているのが分かります。今回の3パターンの中では特異値分解して特異値ベクトルを拡張する方法が一番良さそうです。
なお、perplexityを見て分かる通り、Up-Scale後のモデルの出力は文章にもならないレベルで支離滅裂になっています。
Up-Scaling後の学習
学習
特異値分解の方法で作ったモデルに対してAbejaさんのコーパスで事前学習を行っていきます。
予算の都合で学習できたのは780Mトークンだけです...
青: ランダムな初期値で始めた学習
グレー: 特異値分解で拡張したモデルの学習
明確な差が出ていたので学習は途中で切り上げました。
ランダム初期値と比べてlossが大幅に低いことから、今回のUp-Scalingには一定の効果がありそうです。
一応、学習後のモデルはこちらです。学習時のhyper parameterもこちらに記載しています。
モデルの出力
次に、学習したモデルの出力を見ていきます。
| 入力 | 出力 |
|---|---|
| 日本の首都は、 | 東京です。 東京は、東京都の中心部を指す言葉で、人口が約1,300万人。 世界で4番目に多い |
| 大規模言語モデルとは、 | 大規模なデータセットを用いて自然言語処理を行うためのアルゴリズムです。 ## AIの活用事例 AIはさまざまな分野で応用 |
| Q: 猫を飼育するために気をつけるべきことを1つ教えて A: |
ノミやダニが猫に寄生すると皮膚炎の原因になります。こまめなお手入れが必要です。 また、猫はトイレと爪とぎの頻度が多く、定期的に掃除する必要があります。 猫用のトイレは、猫が排泄物 |
1Bトークンも学習していないにしては日本語を話せている印象です。
ベンチマーク
最後に、llm-jp-evalで4-shotでの日本語能力を測りましたが、さすがにまだ学習不足という感じでした。
| モデル | JCom | JMMLU | JSQuAD | NIILC |
|---|---|---|---|---|
| llm-jp/llm-jp-3-1.8b | 0.16 | 0.3 | 0.5515 | 0.4359 |
| Up-Scalingモデル(2.3B) | 0.06 | 0.0 | 0.0659 | 0.0142 |
まとめ
今回は、幅方向へのUp-Scalingの検証を行ってみました。
今回行った方法にも一定の効果がありそうで良かったです。
スケールアップの方法としてはDepth Up-Scalingがシンプルで有効だと思いますが、個人的にはモデルアーキテクチャのバランスを崩さないためにも幅方向にスケールアップする方法もあったほうが良いだろうと考えています。ある程度まとまった時間ができたら理論から考えてみようと思います。
Discussion