6502プロセッサ上の自己回帰言語モデル

tl;dr - 私は小さなMambaベースの自己回帰言語モデルを訓練し、それを8ビットの6502プロセッサ（1975年製、RAM 32KB）で実行するための推論エンジンを作成しました。私の父親のBBC Microで実行したところ、以下のテキストが生成されました。 once upon a time tom and lily saw things lily were sad her house he heartd them ilily and tom said yes she saw a little girl smiled tom was so excited her mom said yes MOS 6502は1975年にリリースされた8ビットマイクロプロセッサで、BBC MicroやApple IIを搭載していました。私は80年代の父親のBBC Model Bにアクセスできる幸運に恵まれ、最新の機械学習を使用して、このマシンに搭載できる最も強力な言語モデルがどのようなものになるかを見てみたかったのです。驚くことではありませんが、これには大きな課題が伴います。モデルの重みと推論コードは、25KBのユーザー空間メモリ内に収まる必要があります。私の最終構成は、9KBの推論コードと13KBのモデル重みでした。CPUは8ビット整数データ型のみを操作し、命令セットに乗算を含んでいません。推論コードにはCC65が使用されており、C言語から6502命令セットへのコンパイルが可能になります。その後、MacBookで訓練されたモデルのバイナリを、PlayUEFと私がDIYしたカスタム3.5mm-テープケーブルを使用してBBC Microに書き込むことができます。これにより、BBCはテープドライブを接続していると認識し、私のラップトップはヘッドフォンジャックからオーディオを出力します。sim65エミュレータは、C言語の推論バイナリと参照Pythonモデル実装との間のパリティチェックを可能にします。完全な推論エンジンは、BBC Microで実行する前にjsbeebエミュレータでテストできます。あなたも自分で実行できます。以下のリンクは、ブラウザでBBC Microを起動し、GitHubから直接UEFテープイメージをロードし、モデルを実行するためのコマンドを自動入力します。エミュレータのインストールは不要です（生成には数分かかります）。BBC MicroでBitNetを実行する → モデリング このプロジェクトの目標は、自己回帰言語モデルを構築することです。これは、最先端の言語モデルと同様に、トークンを1つずつ生成する言語モデルです。モデルは、既存のコンテキストから次のトークンを生成する関数 $f$ です。 $$ f: \text{'the cat sat on the ma'} \mapsto \text{'t'} $$ 大規模言語モデリングでは、「トークン」は単語またはサブワードの一部になります。この投稿では、使用する語彙（トークンのリスト）は26文字と「 」（スペース）文字になります。この規模の小さなモデルでは、より大きな語彙（例：単語またはサブワードの語彙）は、語彙エンコーダー/デコーダー層がパラメータ予算の大部分を消費することになります。埋め込み層は、トークンをモデルの隠れ次元（この場合はdim=56）にマッピングします。3Blue1Brownのニューラルネットワークに関するビデオは、空間的なトークン埋め込みがどのように機能するかを理解するのに役立ちます。 $$ g: \{\text{a}, \text{b}, ... , \text{z}, '\text{ }'\} \to \mathbb R^{56} $$ トークンが高次元空間にマッピングされた後、混合層がトークン間の再帰的な依存関係をモデル化するために使用されます（再帰層を参照）。 BitNet BitNetは、CPUでの高速推論のための方法として導入されました。行列乗算 $Y = XW$ は、Xの行とWの列の内積のセットです。 $$Y_{ij} = \sum_k X_{ik} W_{kj}$$ BitNetはWを量子化して、その値が三項集合 $\{-1, 0, 1\}$ に収まるようにします。これにより、内積は加算/減算操作のシーケンスに削減されます。 $$ \begin{align*} Y_{ij} &= X_{i1} W_{1j} + X_{i2} W_{2j} + \cdots + X_{in} W_{nj} \ &= X_{i1} - X_{i2} + ... - X_{in}\ \end{align*} $$ 6502プロセッサの命令セットには乗算が含まれていません。代わりに、乗算は繰り返しビットシフトと加算操作から構築されます。単一の8x8乗算と累積には150クロックサイクルかかります。対照的に、三項累積は30クロックサイクルで済むため、高量子化重みでの推論が大幅に高速になります。各BitNetパラメータは、int8の場合は8ビット、float32の場合は32ビットと比較して、$\\log_2(3) = 1.58$ ビットのストレージしか必要としません。バイトあたり4つまたは5つのパラメータをパックできます。バイトあたり5つのパラメータはデータ効率が高いですが、三項値への展開には繰り返し除算/3操作が必要であり、6502のネイティブ命令ではないため展開にコストがかかります。対照的に、バイトあたり4つのパラメータをパックすると、各パラメータに2ビットのチャンクが割り当てられ、展開には右シフトが必要なだけです。推論速度のために、バイトあたり4つのパラメータを選択します。バイトあたり4つのパラメータで13KBの場合、52kのBitNetパラメータを格納できます。実験的検証によると、低量子化での高いパラメータ数は、少ない高精度パラメータよりも優れたコストパフォーマンスを提供します。BitNetパラメータを訓練するには、他の量子化訓練と同様の方法が使用されます。パラメータは完全なfloat32精度で格納され、フォワードパス中に三項値に量子化されますが、バックワードパスでは勾配がフル精度で流れます。 def ternary_quantize(w: torch.Tensor) -> torch.Tensor: """Round to {-1, 0, +1} with straight-through estimator on the backward pass.""" q = torch.clamp(torch.round(w), -1.0, 1.0) # quantize forwards return w + (q - w).detach() # full precision gradient for backwards pass 実際には、最終的なLMヘッドはint4に保たれます。出力射影は、確率を語彙全体にきれいに広げるためにより高い解像度が必要です。モデル内の他のすべての行列パラメータは三項値です。 再帰層アーキテクチャ アテンション GPT-3のような従来の言語モデルは、シーケンシャルな関係をモデル化するためにアテンションを使用しました。アテンションは各トークンにベクトルを割り当て、各ベクトルのペア間の内積を使用してトークンが情報を交換できるようにします。トークンあたり次元 $d$、コンテキスト内の合計トークン数 $s$ の場合、トークンとコンテキストウィンドウ内のすべてのトークンとの間の内積が必要です。合計 $O(ds)$ FLOPs。推論フォワードパス（BBC Microで実際に実行するもの）は、コンテキストサイズによって変化します。これは6502にとって困難です。なぜなら、テンソルサイズは推論中に増加し、トークンを生成するにつれてメモリ使用量が増加するからです。トランスフォーマーがトークンを生成すると、KVキャッシュは生成されるトークンあたり $O( ext{レイヤー数} imes ext{隠れ次元})$ 増加します。32KBのRAM予算では、これはモデル重みの格納に使用したいスペースを食い尽くしてしまいます。 図1：アテンションはすべての前のトークンを振り返るため、計算とKVキャッシュのメモリはコンテキストとともに増加します。 アテンションが提供する利点は何でしょうか？それは各トークンからの正確な想起を可能にします。これにより、トランスフォーマーはコピーやニードルインアヘイスタックの問題でうまく機能します。しかし、このレベルの精度が必要でしょうか？私たちのデータセットでは、単語のスペルを学習し、単純な文法形式を生成するのに十分な短期記憶があれば十分です。 再帰モデル 他のアーキテクチャは、各モデルフォワードパスが同じ計算形状を持つという特性を持っています。モデルによって格納される状態は、固定サイズの状態ベクトル $h$ 内にあります。 $$f: (t_i, h) \to t_{i+1}$$ ここで、$t_i$ はトークン番号 $i$ であり、$h$ は固定サイズのモデル隠れ状態です。SSM（例：S4、Mamba）とRNNスタイルのモデル（GRU、LSTM）の両方がこの特性を満たしており、6502での制約付きメモリ推論により適しています。 図2：再帰モデルは固定サイズの状態を保持し、各ステップでフィードバックされます。計算は各トークンで同じ形状になります。 なぜGRUではないのか？ 再帰モデルは、勾配消失/爆発問題の犠牲になりやすいことが知られており、しばしば不安定な訓練につながります。再帰の各ステップは、フォワード行列の最大固有値の大きさによって誤差を累積します。1よりわずかに大きい固有値は壊滅的になる可能性があります。シーケンス長 $s$ が増加すると、誤差は $\\lambda^s$ として累積します。ここで、$\\lambda = 1 + \varepsilon$ です。BitNetの領域では、フォワード行列のスペクトル半径は通常1をはるかに超えています。約1のスペクトル半径を持つためには、重みの98％が0である必要があります。GRUからのこれらの不安定性により、すべての訓練実行で発散が見られます。これらの問題を回避する唯一の方法は、主要な重み行列をより高い量子化（例：int4）で格納することです。これらはモデルの最大の行列であるため、int4で重みを格納すると、モデルの可能な次元が大幅に減少します。