LLMの中身を知りたくて、小さな日本語LLMを一つずつ作っています。
まだTransformerを勉強している途中なのですが、AWSのブログでPathway’s brain-inspired architecture development on Amazon SageMaker HyperPodという記事を見かけました。
紹介されていたのは、BDH（Dragon Hatchling）という言語モデルです。Transformerとは異なる仕組みで作られているそうです。
次を理解する前に、もうTransformerの次が出てきました。
公式の実装がGitHubで公開されています。今回は詳しい数式の理解はいったん置いて、Macで学習と文章生成まで動かしてみます。
BDHについて調べてみる
現在のLLMでは、Transformerという仕組みが広く使われています。
僕も作りながら勉強しているところなので、BDHとの違いを説明できるほど理解していません。
そこで、まずAWSの記事とBDHのREADMEを読んでみました。
BDHは、脳の神経回路を参考にしたPost-Transformerのアーキテクチャとして紹介されています。
論文とAWSの記事を自分なりに読むと、一部の要素が近くの要素とやり取りし、そのつながりに情報を残す仕組みのようです。
AWSの記事では、BDHを拡張したBDH-CQも紹介されています。こちらは、考える途中の文章を一つずつ生成するのではなく、モデル内部の状態を繰り返し更新して答えを探す、と書かれています。
今日はこのBDHを試したいと思います。
今回試すこと
公式リポジトリには、BDHへ小さな英語の文章を学習させるコードが入っています。
今回は、次の順番で試します。
学習前の予測のずれを測る
↓
小さな文章を繰り返し読み込ませる
↓
学習後にもう一度ずれを測る
↓
文章の続きを生成する公式実装は、文章をバイト単位の数値へ変え、次のバイトを予測します。
予測のずれはlossという値で表示します。自分なりに調べた範囲では、モデルが予測した次のバイトと、実際に続くバイトの違いをまとめた値のようです。
今回は、この値が学習後に下がるかを見ます。
なお、BDHのREADMEには、難しい数独で97.4%正解した結果も掲載されています。ただし、その結果はPathwayの内部実装によるもので、公開リポジトリだけでは再現できないと明記されています。
今回は数独へ挑戦する記事ではありません。GitHubで公開されている基礎的な実装を、小さく動かすところまでです。
Macへ実行環境を作る
今回使ったのはApple SiliconのMacです。
まず、公式リポジトリを取得します。
git clone https://github.com/pathwaycom/bdh.git bdh-mac
cd bdh-macこの実験だけで使うPythonの仮想環境を作りました。
python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -r requirements.txt必要なライブラリは、torch、numpy、requestsの3つです。
公式のtrain.pyだけでも学習と生成はできます。ただ、CPUとMPSを切り替える引数や、実行時間をJSONへ保存する処理はありません。
そこで今回は、計測用のrun_experiment.pyと、結果を並べるcompare_results.pyを追加しました。BDH本体のbdh.pyは変更していません。
bdh-mac/
├── bdh.py # 公式のBDH本体
├── train.py # 公式の学習例
├── run_experiment.py # 今回追加した計測用スクリプト
└── compare_results.py # CPUとMPSの結果を並べるスクリプト公式設定を小さくする
公式のtrain.pyは、3000回学習する設定になっています。
モデルの設定も次の内容でした。
BDHConfig(
n_layer=6,
n_embd=256,
dropout=0.1,
n_head=4,
mlp_internal_dim_multiplier=128,
vocab_size=256,
)まず動くところを見たいので、いきなり同じ規模では実行しませんでした。
今回は、BDH本体の処理は変えず、計算を繰り返す回数や内部で扱う数値の量を減らします。
BDHConfig(
n_layer=2,
n_embd=64,
n_head=4,
mlp_internal_dim_multiplier=8,
)学習回数は300回です。一度に読み込む文章も、公式設定より小さくしました。
この設定では、学習によって調整される数値は合計131,072個になりました。
数値の多さだけを見ても、まだあまり実感はありません。少なくとも、LLMという言葉から想像する巨大なモデルではなさそうです。
CPUで学習させる
CPUを使い、次のコマンドで実行しました。
python run_experiment.py \
--device cpu \
--profile quick \
--output results/cpu.json初回だけ、学習用のtiny Shakespeareがダウンロードされます。シェイクスピア作品をまとめた小さな英語データです。
実行結果はこちらです。
device=cpu profile=quick parameters=131,072 iterations=300
initial validation loss=5.5881
step=1/300 training loss=5.6055
step=25/300 training loss=4.5044
step=50/300 training loss=3.3885
step=75/300 training loss=3.0051
step=100/300 training loss=2.8044
step=125/300 training loss=2.6181
step=150/300 training loss=2.5182
step=175/300 training loss=2.4488
step=200/300 training loss=2.3834
step=225/300 training loss=2.3241
step=250/300 training loss=2.2792
step=275/300 training loss=2.2555
step=300/300 training loss=2.2305
final validation loss=2.1463
elapsed=5.83s tokens/sec=52651
generated text:
To be or the this a are the arden the there an with he to thing to me sond arte he arries and and and herend,
result=results/cpu.json学習前は5.5881だったlossが、学習後には2.1463まで下がりました。
少なくとも、渡した文章から何かを学習している様子は確認できました。
軽量設定では、300回の学習にかかった時間は5.83秒です。思っていたよりあっさり終わりました。
文章を生成する
学習が終わると、To be or に続く文章が自動で生成されます。
今回の結果はこちらです。
To be or the this a are the arden the there an with he to thing to me sond arte he arries and and and herend,訳版↓
あるか、それともこれか、あそこか、そこか、彼と、そのこと、私と、音と、芸術と、彼はリスクを冒し、そして、そして、そして、彼女と、正しい文章ではありません。
学習前のランダムな状態で出てきそうな文字列とは違い、英単語らしきまとまりや空白が並んでいます。
わずか300回、約13万個の調整対象しかない設定なので、今回は文章の品質を見る実験ではありません。BDHのコードで、学習から生成まで一通り動いたことを確認できました。
Apple SiliconのGPUでも動かす
公式のtrain.pyには、MacでMPSを使うための記述もありました。
MPSは、PyTorchからApple SiliconのGPUを使うための仕組みです。
CPUとの速度を比べるため、同じモデル設定、学習回数、乱数の初期値で実行します。
python run_experiment.py \
--device mps \
--profile quick \
--output results/mps.jsonこちらも最後まで動きました。
device=mps profile=quick parameters=131,072 iterations=300
initial validation loss=5.5881
step=1/300 training loss=5.6028
step=25/300 training loss=4.5115
step=50/300 training loss=3.3918
step=75/300 training loss=2.9973
step=100/300 training loss=2.7893
step=125/300 training loss=2.6046
step=150/300 training loss=2.5048
step=175/300 training loss=2.4397
step=200/300 training loss=2.3709
step=225/300 training loss=2.3111
step=250/300 training loss=2.2762
step=275/300 training loss=2.2516
step=300/300 training loss=2.2234
final validation loss=2.1657
elapsed=3.37s tokens/sec=91154
generated text:
To be or ward the he and the has the ans of thant to the strat and the therer tof thell wer hound werer thit
result=results/mps.jsonCPUとMPSの結果を並べます。
python compare_results.pycpu: loss 5.5881 -> 2.1463, time 10.69s, speed 28745 tokens/s
mps: loss 5.5881 -> 2.1657, time 3.92s, speed 78370 tokens/stokens/sは、1秒間に処理したバイト単位のトークン数です。今回の学習データは英語なので、ほとんどの場合は1文字が1トークンになります。
軽量設定では、MPSの実行時間はCPUの約58%でした。今回は約1.7倍速くなりました。
CPUとMPSへ同じ乱数の初期値を指定しましたが、学習後のlossと生成された文章は完全には一致しませんでした。
PyTorchの再現性に関するドキュメントを確認すると、同じ乱数の初期値を使っても、CPUとGPUの実行結果は一致しない場合があると書かれていました。どちらもlossは同じくらいまで下がっているため、今回は文章の違いではなく、実行時間を比べます。
動かしてみた感想
公式コードが公開されており、設定を小さくすると、CPUでは10.69秒、MPSでは3.92秒で学習と文章生成まで動きました。
新しいAIアーキテクチャをニュースとして読むだけでなく、Pythonのコードとして動かせたのは面白かったです。
まだbdh.pyの計算を読んでも、何をしているのか理解できない箇所がかなりあります。時間を見つけて入力した文章がBDHの中をどのように通っているのか、少しずつ追ってみたいと思います。
参考
- Pathway’s brain-inspired architecture development on Amazon SageMaker HyperPod
- pathwaycom/bdh
- The Dragon Hatchling: The Missing Link between the Transformer and Models of the Brain
ここまで読んでいただき、ありがとうございます。もしこの記事の技術や考え方に少しでも興味を持っていただけたら、ネクストのエンジニアと気軽に話してみませんか。
- 選考ではありません
- 履歴書不要
- 技術の話が中心
- 所要時間30分程度
- オンラインOK
