5分でノートPC上で学習できる最強のAIモデルは何か？

• MacBook Proで5分のうちに、約180万パラメータのGPTスタイル・トランスフォーマーモデルを約2000万TinyStoriesトークンで学習し、約9.6のパープレキシティを達成
• 5分以内の学習における主な制約はモデルサイズと処理可能なトークン数であり、モデルが大きいほど収束が遅くなり、少ないデータでは効果が下がる
• 性能最適化では、MPSの使用が、コンパイル・量子化・勾配累積やPyTorch代替よりも、小さいモデルを選ぶことのほうが効果的
• TinyStoriesのような単純で一貫性のあるデータセットは、百科事典的なデータよりも小型モデルの性能により良い影響を与える
• トランスフォーマーアーキテクチャは、小さなサイズと短い学習時間という条件で、LSTMやdiffusion方式より優れた結果を示す

---

概要

この記事は、ノートPC（MacBook Pro）で5分以内に学習可能な最大性能のAI言語モデルを実験した結果と、最適なトレーニング戦略、データセット選定、モデルアーキテクチャに関する知見を提供する

実験結果の要約

• 約180万パラメータのGPTスタイル・トランスフォーマーモデルを約2000万のTinyStoriesデータで学習し、9.6のパープレキシティを記録
• 生成例は短いものの一貫した物語の形を保ち、英語の文法も概ね正しく維持されるレベル
• 5分以内で実用的な水準のモデル結果が得られたことは期待以上だと強調

実験の背景と限界

• ノートPC環境で強力なモデルを短時間で学習するのは現実的ではない、という好奇心から始まった実験
• 実際には**クラウド上の高性能GPU（H100など）**でより強力なモデルを学習できるが、この実験の制約条件は時間（5分）である
• モデルサイズが大きくなるほどトークン処理速度が遅くなり、5分以内に良い結果を出すのが難しい
  • 小さすぎるモデル（例: 1万パラメータ）は十分な複雑さを学習できない
  • 実用的な範囲は約100万〜200万パラメータのモデル

スループット最適化

• **MPS（AppleのMetal Performance Shaders）**の使用が最も効果的
• torch.compile、float16、MLXなどのさまざまな数値最適化は、期待したほどの性能改善がなかったり、むしろ悪化したりした
• 勾配累積はメモリ管理のための手法ではあるが、実際には速度低下が深刻
• モデルが内部メモリ内で素早くweight updateできることが効率上重要

データセットの選択

• 限られたトークン数（約1000万〜2000万）でSimple English Wikipediaなどの単純な英語Wikipediaデータを先に使った結果、文法的一貫性は得られたが意味的一貫性が不足
  • 固有名詞中心で、作為的に並べたような事実の羅列になり、有意義な内容生成には限界があった
• TinyStoriesデータセットを使うと、物語構造が明確で言語も単純なため、結果ははるかに一貫性があり意味のあるものになった
  • 4歳児レベルのストーリーであり、小型モデルでも学習が進みやすい

トークナイザーとトークン化

• トークナイザーのトレーニングは5分以内には含まれず、データ規模も小さいため最適化の必要性は低い
• マルチバイトトークンの学習はモデル学習にとってより容易

モデルアーキテクチャの実験

• **トランスフォーマー（GPT-2スタイル）**アーキテクチャを使用

  • 2〜3レイヤー、SwiGLUなどの活性化関数、positional embeddingなどのハイパーパラメータを調整
  • LSTMは性能が近いものの、パープレキシティの面ではトランスフォーマーのほうが優れる
  • Dropout、mixture-of-expertsなどはサイズが小さいため非効率
  • Curriculum learningは学習時間が短すぎて効果が小さい

• **Diffusionモデル（D3PM）**も試行

  • 自然言語は離散トークンであるため、拡散過程では無意味なランダムトークンしか生成されず失敗
  • トランスフォーマーやLSTMと比べて、文の構造を素早く形成するのが難しい

モデルサイズとトークン/秒スループットの関係

• 100万〜200万パラメータのモデルが最も理想的なsweet spot
  • 大きすぎると5分以内に収束できず、小さすぎると学習直後に性能限界へ達する
• Chinchilla scaling lawと実験結果は概ね一致
  • 総学習トークン数/20が理想的なモデルサイズであり、この実験でも確認された

結論と示唆

• 非常に短い時間と小さなハードウェアでも、一貫したストーリーテリングモデルの学習は可能
• 5分学習は強力なモデル開発には適さないが、小規模・超軽量モデル設計やハードウェア・アーキテクチャ最適化の実験には意義がある
• 今後ノートPC GPUやモデル構造が進化すれば、わずか数分で訓練できるモデル性能がさらに向上する可能性がある