Team「たぬき」開発振り返りメモ1: Scaling Lawに挑戦しようと準備する話

はじめに

Team「たぬき」でleaderをしているHatakeyamaです｡そろそろチームとしての開発が終わりそうなので､備忘録として､開発の振り返りメモを書いていきます｡
特に公言してこなかったのですが､実は今回の開発は､個人的には「Scaling Lawとの戦い」が裏テーマでした｡本記事では､その準備過程について､記します｡

• 関連記事
  • チーム「たぬき」ゴールデンウィークの現況

背景: Scaling Lawとは?

簡単な説明

Scaling lawとは､大規模言語モデル開発の基礎となる経験則です｡ChatGPTの説明は以下の通り｡

大規模言語モデルにおけるスケーリング法則の一部には、特に「べき乗則（Power Law）」に基づくものがあります。べき乗則は、特定の変数（例えば、モデルのサイズや訓練データの量）が増加する際の性能向上が、対数的な関係で表現されるという法則です。具体的には、以下のような関係が観察されます。

ややこしいことが書かれていますが､要するに､モデルの性能は､投入した計算量のべき乗に比例するということです｡

良さ

開発者にとってのScaling lawの利点は､投入した計算資源を増やせば､(今のところ?)際限無く^[1]､モデルの性能が向上するということです｡そのため､モデルの巨大化がどんどん進んでいます｡

難しさ

性能を上げようとすると､計算コストがバカ高くなるというのが､Scaling lawの難しさです｡例えば簡単のため､性能$x$が投入した計算量$y$の二乗で表されるとすると､必要な計算量は､$y=x^2$となります｡
性能を1上げるために必要なコストが､どんどん上がっていきます｡

• 性能1: 必要コスト=1
• 性能2: 必要コスト=4
• 性能3: 必要コスト=9
• ...

大きなモデルを作るために必要なコストが､べき乗で増えていくという事実は､金銭的には､かなりしんどい戦いです｡
対応策として､Metaは2024年にGPU(H100)を34万枚購入するというニュースもありました｡これに対し､今回のコンペで使えるGPUは､その1万分の1以下の24枚です^[2]｡

マネーゲームでは勝てない｡ではどうするか?

Scaling lawの真っ向勝負では､どう考えても勝てない相手が沢山います｡そこで､どうにかして､この法則の「抜け穴」を探すという取り組みを行ってきました｡

もちろん､Scaling lawそのものを乗り越えることは､極めて困難(ノーベル賞級の研究?)です｡一方で､$y=ax^b$における係数$a,b$を小さくすることは可能です｡
そのような着想に基づき､モデルアーキテクチャ､事前学習のデータセット準備､ファインチューニングのデータセット準備の3過程で､係数を小さくするための試行錯誤をしました｡

モデルアーキテクチャを考える

筆者は真正のAIの専門家ではないので､このプロジェクトにおいて､モデルアーキテクチャについては文献収集をするなどにとどめ､teamとして保守的なアプローチを取りました｡

色々と調べてみたのですが､結論からいえば､2024年5月時点においても､学習効率という観点において､アーキテクチャとしてはTransformerがベストで､それ以外の新興モデル(Mamba, RWKV, ...)は､同等という印象でした｡実績面も鑑みると､Transformerが手堅いモデルということになります｡

もちろん､Transformerの中でも､通常の(denseな)モデルではなく､いわゆるsparse MoE (mixture of experts)の方が学習効率が良いとの情報は把握しています^[3]｡ただ､学習が不安定化したり^[4]､計算効率(FLOPS)が落ちたりするという怖い噂がありました｡また､今回使える計算リソース(10 b modelを作れる程度)の中でMoEを行うことの具体的な効能が見えてこなかった^[5]こともあり､MoEは採用しませんでした｡ただし、後述の通り、BTMというアプローチを使うことにしました。

結論: アーキテクチャとしては普通のTransformerを使うことにした｡

事前学習データベースを工夫する

限られた計算資源下でモデルの性能を上げるためには､モデルだけでなく､データベースを工夫するのが効果的です｡事前学習データベースについては､主に3つの工夫をしました｡

1. きれいなデータセットを作る
2. ドメインを絞る
3. Branch-Train-Mergeという手法を軽く検討する

詳細を以下に述べます｡

きれいなデータセットを作る

事前学習に使われるデータセットからノイズを取り除き､有意義な情報の密度を高めることができれば､投入コストに対する学習効率が向上すると考えました(コーパスの構築コードはこちら)｡

これまでのデータベースに比べ､(自称ながら)かなり真面目にクリーニングを行いました｡
具体的には､Web上のデータには広告があまりにも多すぎたので､ルールベースや機械学習フィルタを活用して､それらを削ったりしました^[6]｡検討の詳細は以下のスライドにまとめています｡
https://www.docswell.com/s/KanHatakeyama/ZYW393-2024-04-08-112244

CommonCrawlからもデータを収集するなどして､それなりにクリーニングをした日本語のデータセット(約500GB, 100 B token)を作ることができました｡

実務面では､とにかくデータサイズがでかいので､処理速度やストレージといった問題が噴出しました｡Google Cloud Platformなどを使い慣れている方々が､並列処理etcを検討してくれました｡

ドメインを絞る

データセットで工夫した2つ目の点は､ドメインを絞ったということです｡今回、作ろうとしているモデルは10 b程度のサイズであり､GPT-4 (1800 bとの噂)よりも遥かに小さいです｡このような小さなモデルに､この世のすべてを学ばせても､中途半端な器用貧乏モデルにしかならない可能性があります｡そこで､学習させるデータは若干､絞ることにしました｡

• 日本語ドメイン (約100 b token)
  • CommonCrawlを中心にしたデータ
• 英語ドメイン (約100 b token)
  • Wikipedia, Wikibook
  • 学術論文
  • コード､数学データ

通常のアプローチとの大きな違いは､英語ドメインにおいて､通常のWebデータ(CommonCrawl)は使わず､きれいな文章(Wiki, 論文, コード)のみを用いた点です｡

この判断は､

• 知識転移(英語で学んだことを日本語で回答する)はあまり期待できないだろうという直感
• このモデルとの上質な英会話を期待するユーザーは殆どいないだろう予測
  に基づいています｡

Branch-Train-Merge (BTM)もどきを画策する

BTMは､データベースをドメイン毎にN分割し､独立に学習させたモデルを最後に統合して推論するテクニックです｡各分野の専門家の集まりからなるモデルを総動員して運用するという意味では､文字通りのmixture of expertsです^[7]｡
投入した学習リソースに対する性能(つまり学習効率)が良い上､学習プロセス自体は通常のdenseモデルと同じなので､この手法には注目していました｡

ただし､今回は学習リソースも限られていることから､最大限に学習効率を高めるために､「ドメインごとに専門家モデルを独立に構築する」のではなく､「一つのモデルに対して､ドメインを切り替えながら継続事前学習を行う」ことにしました｡
「国語」→「算数」→「理科」...のような順番でモデルを学習させ､学習直後のsnapshotを保存して運用するイメージです｡

今回は､教師なしクラスタリングを用いて､日本語を5つのドメインに自動分割した上で､
英語+codeドメイン → 日本語A → 日本語B → ... →日本語Eという順番で学習を行うことにしました^[8]｡

ファインチューニングデータセットを工夫する

データは1000件で十分なのか?

事前学習済みモデルのファインチューニングについては､様々な考え方があります｡
一斉を風靡したのは､いわゆる LIMA論文です｡LIMAはLess Is More for Alignmentの略で､ファインチューニングデータは上質な1000件程度あれば十分という主張をしています｡
これは開発者にとっては､ある意味「福音」です｡頑張って1000件のデータセットを作れば､それでファインチューニングタスクがほぼ完了(!?)するからです｡

1000件では足りないかもしれない

しかし､予備検討においては1000件では全くうまくいきませんでした｡
公開されているモデルを用い､色々なデータセットを試してみたのですが､少なくとも10b程度のモデルでは､1000件程度のデータを学習するのみでは､所望の出力が得られませんでした｡詳細については､以下の記事にまとめています｡
https://note.com/kan_hatakeyama/n/n55e050cb74fe
https://note.com/kan_hatakeyama/n/n2f1a8dfeebfa

ポイントは以下の通りです｡

• 10b程度のモデルは汎化性能が低く､未学習のzero-shot taskは殆ど解けない
• 特に､形式が厳格なタスクに対する指示追従性が低い
• 出力形式にこだわらないタスクであれば､1000件程度でのファインチューニングでも問題ないかもしれない (LIMA論文での検証タスク(?))

1000万件でファインチューニング?

LIMA論文の対極をなすアプローチとして､最近公開されたLlama 3では､1000万件以上のデータでファインチューニングを行った(!?)との報告がありました｡この発想は､どちらかというと筆者寄りです｡「頭の悪い」モデルに対しては､可能な限り徹底的に人類に対する想定問答としての学習を繰り返すという考え方です｡あるいは､猛特訓で受験勉強を行うイメージでしょうか｡

もちろん､alignment taxと呼ばれる考え方があったり､未知の事象をファインチューニングで覚えさせようとすると､モデルが混乱する(こちら)という最近の報告もある通り､ファインチューニングのデメリットは存在するはずです｡
しかし､現実問題として「頭の悪いモデル」(生徒)に対して､開発者ができる(唯一の?)選択肢といえば､とにかくファインチューニングで特訓をするしかないわけです｡

(ちなみに､「賢いモデルに対してはファインチューニングデータは少なくても良い」という旨のことを某社の某氏がおっしゃっていました｡この考え方には､筆者も同意します｡)

ファインチューニングデータをどう集めるか?

上記の経緯もあり､徹底的にモデルを訓練するという構想は､わりと初期からありました｡
MinnadeChatのような､データセットの投稿サイトもチームメンバーが開発してくれました｡
https://zenn.dev/matsuolab/articles/ca67927eaa497a
これによって､数百件程度の指示学習データが集まりました｡

しかし､このペースでは､Llama3の1000万件には到底､及びません｡

ゲームチェンジャー(?)としてのMixtral 8x22b

大量のデータセットを作るための鍵は､AIによる自動生成です｡筆者の考えでは､OpenAI, Google, Metaなどのトップランナーは､指示データセットを完全手作業で作るのではなく､モデル自身に生成させています^[9]｡
これは､AIに対して手取り足取りものごとを教えるのではなく､AI自身が提案してきた回答を人間が修正するというプロセスです｡これはある意味､驚くべきことで､人間がgood/badのようなフィードバックを出すだけで､AIが自ら成長できるようになってきたとことを意味します｡

これまで､このような高度な自動生成のアプローチは､一部のトップ企業しか行うことができませんでした^[10]｡しかし､真のOpen AIカンパニーとも言われるMistralは､自動生成にも展開可能な水準の賢いモデル(Mixtral 8x22b)を寛容なライセンスで4月に公開しました｡プロジェクト開始後の出来事です｡

https://labs.perplexity.ai/

Mixtral 8x22bはゲームチェンジャーとも呼ばれる存在で､日本においても､「人海戦術」によるデータセット作成のゲームが､「人海戦術+計算資源」に変わりつつあることを印象付けました^[11]｡
実際にやってみるとわかりますが､データセットの作成は､かなりの手間とメンタルを消費します｡これをAIによって半自動化できた点は大きいです｡

やや話がそれましたが､Scaling lawの文脈における自動生成データの位置づけを考えます｡
賢くてOpenな言語モデルの登場により､構築したいモデルの最終調整段階で､より教師モデルの出力を参考するというアプローチを取る､すなわち上質な人工データを学習させられるようになってきました｡これにより､より少ない計算・人的コストで､高性能なモデルを作ることができます^[12]｡自動生成データというアプローチの有効性は､有名なTextbooks are all you need論文や､最近ではベンチマークスコアに対するScaling lawの論文などでも示されています｡

結果は?

本記事では､Scaling lawに対抗するための方策を色々と練りました｡
結果はどうだったのかというと､想定外のことばかりで､とても大変でした｡
長くなりましたので､詳細は次の記事で記します｡

この成果は、NEDO（国立研究開発法人新エネルギー・産業技術総合開発機構）の助成事業「ポスト５Ｇ情報通信システム基盤強化研究開発事業」（JPNP20017）の結果得られたものです。

脚注

1. 最先端では､性能が頭打ちしているという噂もあります｡

2. H100は1枚あたり500万円くらいはしますので､それでもかなりの金額です｡

3. 要出典ですが､割愛します｡

4. そうはいっても､Mixtralなどはcodeが公開されてるようで､それに乗っかる分には､意外とOKという雰囲気を感じました｡

5. 例えば8x7b, 7x22bのようなモデルは公開されていますが､8x1bという感じのアーキテクチャはあまり見かけません｡見かけないということは､やるメリットが少ないと判断されているということだと思います｡

6. その結果､どのようなコンテンツバランスになったのかなどについては､本当は詳しく検証する必要があります｡

7. 一般的に使われているMoEは､ドメインレベルではなく､トークンレベルで分業をします｡

8. データセットの分割はかなり大変でした｡数百GBのテキストに対して､FastTextでクラスタリングを行いました｡

9. たとえばLlama3のサイト(https://huggingface.co/blog/llama3)では､指示データはhuman annotated dataと表記されています｡human generatedではない点に注意｡

10. GPT-4, Claude, Llamaなどの高性能なモデルは､出力をAIの学習に使うことが利用規約で禁止されています｡

11. 公開したMistralからすれば､この程度のモデルは公開してしまっても､競争的な優位性は崩れないという判断をした(つまり､彼らはもっと先に進んでいる)ということになります｡この事実については､重く受け止める必要があります｡

12. 広義のモデル蒸留です｡ちなみに､Mixtral 8x22bは日本語ドメインなどでかなりハルシネーションを含む出力を出します｡自動生成されたデータに対しても､人間によるチェックが必要です｡