「ノートPCで5分でAIをトレーニングする」という条件で最も高性能なAIを作る方法

AIのトレーニングは基本的に高性能なGPUを長時間稼働させて実行します。これに反して、GitHub社員でAI関連の製品開発に携わっているSean Goedeck氏は「MacBook Proで5分間で言語モデルをトレーニングする」というチャレンジに挑み、どんな手法が最も適切なのかを探りました。

What's the strongest AI model you can train on a laptop in five minutes?
https://www.seangoedecke.com/model-on-a-mbp/

◆機械学習ライブラリ
言語モデルのトレーニングに用いるライブラリには「PyTorch」を採用。Appleが開発している「MLX」も試したものの、速度を向上する効果はなかったそうです。

◆学習データ
トレーニング時間を5分間に制限すると、学習データの量も大きく制限され、約50MB(約1000万トークン)のデータしか使えないとのこと。この50MBの中に偏りの少ないデータセットを含める必要があります。

最初は「シンプル英語版ウィキペディア」のデータを使用しましたが、「○○は○○の仲間です」といった百科事典特有のフレーズばかりを繰り返すようになってしまいました。以下のテキストは太字部分がGoedeck氏が入力したもので、その後に言語モデルの出力したテキストが続いています。内容の正誤はともかく、「Wikipediaの各ページの最初の1文」っぽいフレーズを繰り返していることが分かります。

Paris, France is a city in North Carolina. It is the capital of North Carolina, which is officially major people in Bhugh and Pennhy. The American Council Mastlandan, is the city of Retrea. There are different islands, and the city of Hawkeler: Law is the most famous city in The Confederate. The country is Guate.

最終的にGoedeck氏は4歳児レベルの読解力に合わせて書かれた短編物語を収録したデータセット「TinyStories」を採用しました。

◆アーキテクチャ
Goedeck氏は言語モデルのアーキテクチャとして、Googleが開発し現代の大規模言語モデルの多くが採用しているアーキテクチャ「Transformer」を選択し、GPT-2のトレーニング手法を模倣しました。チャレンジ全体のうち、Transformerのパラメーター調整に最も長い時間を費やしたそうです。

Gemini Diffusionなどの登場で言語モデルのトレーニングに拡散モデルを採用するアプローチが注目を集めていますが、今回のトレーニング環境では拡散モデルは役に立ちませんでした。

◆モデルサイズ
言語モデルはパラメーター数が大きいほど高性能なイメージがありますが、「MacBook Proで5分でトレーニングする」という条件ではモデルサイズを過度に大きくすると性能が下がってしまいます。Goedeck氏は複数のパラメーター数の言語モデルを作成して性能を比較しました。

以下のグラフは横軸がパラメーター数、縦軸が困惑度(困惑度が低いほど高性能)を示しています。グラフを見ると、200万パラメーター前後が高性能であることをが分かります。

◆結論
最終的に、Goedeck氏は「MacBook Proで5分でトレーニングする」という条件では「PyTorchで、TinyStoriesを学習データとして用いて、Transformerを採用した180万パラメーターの言語モデルをトレーニングする」という手法が最も高性能な言語モデルを作成できると結論付けています。

作成した言語モデルの出力例が以下。太字部分がプロンプトで、細字部分が言語モデルの出力です。短いストーリーを割と高精度に生成できています。

Once upon a time, there was a little boy named Tim. Tim had a small box that he liked to play with. He would push the box to open. One day, he found a big red ball in his yard. Tim was so happy. He picked it up and showed it to his friend, Jane. “Look at my bag! I need it!” she said. They played with the ball all day and had a great time.