イーロン・マスク氏が2024年7月に立ち上げたAI企業「xAI」が3月17日、「Grok-1」のベースモデルの重み（weights）とアーキテクチャをリリースしました。GitHubにて、Apache 2ライセンスで配布されています。

当初、LLMだと思われていたこの「Grok-1」ですが、OpenAIの ChatGPT のようなAPIサービスや、Mistral「Mixtrial-7B」, Stability AI「Japanese Stable LM-7B」, CyberAgent「OpenCALM-7B」のようなLLMとも異なるようです。

カバー画像は、Grokによって提案

Mixture of expertsモデルについては後ほど解説します。
以下、しらいはかせによる解説でお送りします。

まずモデルダウンロードから重い

ダウンロードには分散ファイルシステムの torrentを使います。
296GBあります。

ダウンロードが終わる気配がありません。
現時点で、だいたい4日～1週間ぐらいかかりそうです。
またモデルのサイズが大きい（314Bパラメータ）ため、サンプルコードでモデルをテストするには、十分なGPUメモリを持つマシンが必要です。もし、このファイルサイズそのままをGPUにロードするのであれば296GB以上のVRAM環境が必要です。
つまりダウンロードに成功したとしてもtorrentネットワークの分散ファイルシステムに協力者として放置するぐらいしかできることがなさそうです。

OpenAIの最初のインターンであり、スタンフォード大学博士、現NVIDIAリサーチマネージャー兼エンボディドAI（GEARグループ）リーダー。エージェント、ロボティクス、ゲームのための基礎モデルを作成する ジム・ファン @DrJimFan のツイートが興味深いです。

ソースコードを読んでみる

torrentでのダウンロードの方法については、こちらを参考にしつつ、

ダウンロードに数日掛かりそうなのと、一般的なPC環境などで利用できるモデルではなさそうなので、以下はリポジトリからの調査になります。

README 和訳

このリポジトリには、Grok-1 open-weightsモデルをロードして実行するためのJAXサンプルコードが含まれています。

checkpointをダウンロードし、ckpt-0 ディレクトリを checkpoint に置き、テストコードを実行してください。

pip install -r requirements.txt
python run.py

スクリプトは checkpoint をロードし、テスト入力にモデルからサンプルをロードします。

モデルのサイズが大きい（314Bパラメータ）ため、サンプルコードでモデルをテストするには、十分なGPUメモリを持つマシンが必要です。このリポジトリにおける MoE レイヤーの実装は効率的ではありません。これは、モデルの正しさを検証するためのカスタムカーネルの必要性を避けるために実装しています。

torrentクライアントと、このマグネットリンクを使ってウェイトをダウンロードできます。

magnet:?xt=urn:btih:5f96d43576e3d386c9ba65b883210a393b68210e&tr=https%3A%2F%2Facademictorrents.com%2Fannounce.php&tr=udp%3A%2F%2Ftracker.coppersurfer.tk%3A6969&tr=udp%3A%2F%2Ftracker.opentrackr.org%3A1337%2Fannounce

ライセンス：このリリースのコードと関連するGrok-1の重みは、Apache 2.0ライセンスの下でライセンスされています。このライセンスは、このリポジトリのソースファイルとGrok-1のモデル重みにのみ適用されます。

Mixture of experts (MoE)とは

日本語での原著論文解説を if001 さんがZennで寄稿されていますので引用させていただきます。

もともとの論文ではスパースなMoEについてですが、if001さんのZenn記事ではスパースではない実装についてもコードとともに解説されています。
https://arxiv.org/abs/2101.03961

TransformerのFFN (フィードフォーワードネットワーク）についてはニツオさんのZennが詳しいです。

▼AI界を席巻する「Transformer」をゆっくり解説(5日目) ～Model Architecture編 3～｜ニツオ 
https://zenn.dev/attentionplease/articles/1a01887b783494

GPTやStable Diffusionの内部で使われている「Transformer」の原著論文「Attention is all you need」で解説されているアーキテクチャー内での青い部分が「Feed Forward」です。

右側のデコーダーの線形変換の部分で、灰色の部分は6回繰り返しになっています。アテンション・サブレイヤーに加え、エンコーダーとデコーダーの各レイヤーには全連結のフィード・フォワード・ネットワークが含まれており、各ポジションに別々に同じように適用されます。これは2つの線形変換とその間のReLU活性化で構成されています。通常は、元の画像や文章をベクトル変換して、1つのデータとしてまとめて行列として扱うのですが、そのサイズや次元数を減らすためにフィルターを通して、サイズや次元数を減らします。これは畳み込み（CNN: convolutional neural network）を使います。線形変換はそのフィルターサイズが1x1カーネルで畳み込みをした場合と同じとみなせます。

MoEについてHuggingFaceの解説

MoEとは

• 高密度モデルと比較して、事前トレーニングがはるかに高速

• 同じ数のパラメータを持つモデルと比較して推論が高速になります

• すべてのエキスパートがメモリにロードされるため、高 VRAMが必要

• 微調整では多くの課題に直面しているが、MoE の命令調整に関する最近の研究は有望である

MoE では、Transformerモデルのすべての FFN 層を、ゲートネットワークと一定数のエキスパート（専門家）で構成される MoE 層に置き換えます。

DrJimFanの「mixture of export (2 out of 8 active)」とは8つのエキスパートのうち2つが出力しているよ、という意味と捉えました。

Grok-1のリポジトリには

たしかに、モデルの正しさに対して、チートする目的があれば、MoEをスパース（まばら）を最適化して、カスタムカーネルを用意することでサイズを減らし、特定目的に特化させることはできると考えます。おそらく今後、このGrok-1を削る形で様々なモデル（例えばチャットボットや文書検索、翻訳タスクなど）が出てくるのではないでしょうか。
ただし現状は、このファイルをダウンロードして、Transformerにかけることができる演算基盤を持っている方々に限られます。
イーロン・マスクからの挑戦状、というところでしょうか。

dm-haiku

内部での必要なライブラリについても目を向けていきます

[requirements.txt]
dm_haiku==0.0.12
jax[cuda12_pip]==0.4.25 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
numpy==1.26.4
sentencepiece==0.2.0

Google Deepmindによる「dm-haiku」というJAXベースのニューラルネットワークライブラリを使っています。

JAXベースのニューラルネットワークライブラリです。フレームワークではありません。そしてこれは Google内のDeepMindとGoogle Brainの"仁義なき戦い"でもあります。

JAX/Flaxでの日本語でLLMを作る実装について、fukugawaさんによるGCPでの解説がありましたので引用させていただきます。

JAX/Flaxで作るTransformer言語モデル ❷ cc100編
https://zenn.dev/fukugawa/articles/6f27513f5dce1c

GCP上でテンソルプロセッサを使って0.1BのLLMを作ろうとしています。だいたいこれの3万倍ぐらいの演算基盤が必要になりそうですね。
（計算が間違ってたらすみません）

xAIは「宇宙を理解する」企業

せっかくなのでxAIについても調査してみました。

「Our Mission: Understand the Universe」（宇宙を理解する）とあります。

xAIの技術チーム

xAIのチームは、 Tesla と SpaceX のCEOであるイーロン・マスクが率いています。過去に、DeepMind, OpenAI, Google Research, Microsoft Research, Tesla, University of Toronto. Adam optimizer, 特に Batch Normalization, Layer Normalization, adversarial examples といった発見に貢献しています。

さらにこの分野で最大のブレークスルーのいくつか Transformer-XL, Autoformalization, Memorizing Transformer, Batch Size Scaling, μTransfer, SimCLR といった革新的な技術や分析を導入してきした。AlphaStar, AlphaCode, Inception, Minerva, GPT-3.5, GPT-4 に取り組み、開発をリードしてきました。

・xAI社の社内AIチューターは、技術チームと協力して、当社のモデルをトレーニングおよび評価するための高品質なデータを作成し、管理しています。当社のチューターは幅広い分野の専門家であり、専門的なライティングスキルを備えています。

・xAI社のチームは、現在 Center for AI Safetyのディレクターを務めるダン・ヘンドリックスがアドバイザーを務めています。

パートナーシップ
私たちはX Corpと密接に提携し、Xアプリの5億人以上のユーザーに私たちの技術を提供しています。

開発者 イゴール・バブーシキン

grok-1のリポジトリでメインコントリビューターである イゴール・バブーシキン（Igor Babuschkin）についても注目していきます。

Google DeepMind、OpenAIを経験、スタークラフトIIの研究者でもあります

DIY Automated Papers：論文を自動で書くPythonスクリプトを中心としたGitHubリポジトリです。9年前に書かれています。

https://babushk.in/research.html

最近の論文ではこのような研究も発表されています。

"Formal Mathematics Statement Curriculum Learning" (2022)
(形式化数学記述カリキュラム学習）
Stanislas Polu, Jesse Michael Han, Kunhao Zheng, Mantas Baksys, Igor Babuschkin, Ilya Sutskever

https://arxiv.org/abs/2202.01344

形式化数学、つまり数学の証明などの問題をコンピュータに記述させる課題です。高校のオリンピックから抽出した複数の難問を自動的に解くという、最先端の成果を発表しています。

ダウンロード、終わりそうにないですね！
以上、速報でした。

Grokについての技術情報はこちら
https://x.ai/blog/grok

Grokのような巨大なモデルを必要としないチャットボット開発などはAICU AIDX Labにご相談ください！