TONN だけでは Llama に負けた。ハイブリッドなら勝つ

最初の公平な比較は、私たちの負けでした。13M の TONN と、同じサイズの Llama 型トランスフォーマーを、同じデータで同じステップ数だけ学習させました。文章では僅差でした。コードでは大差がつきました。この記事では、それを勝ちに変えた 49 回の実行を追います。計算量をそろえた比較でも勝ちが保たれるようにした、速度改善も扱います。

数値はすべて、英語、日本語 Wikipedia、日本語ウェブテキスト、コードの評価用テキストでのバイトあたりビット数（bpb）です。低いほど良い値です。どのモデルも、同じ 32k トークナイザーと同じデータ配合（英語 50%、日本語 30%、コード 20%）を使います。

01TONN だけのモデル対 Llama

13M パラメータ、3,000 ステップ、それぞれ 2 シードでの結果です。

文章では、TONN は接戦です。日本語ウェブテキストでは同点、英語では 0.009 の差です。差のほとんどはコードで、バイトあたりビット数で約 0.18 です。コードでは、名前を正確に写すことが点につながります。アテンションはこれが得意です。Llama は、すべての層にフルアテンションを持ちます。ランダムな初期配置の TONN 層が持つ小さなアテンション回路は、平均で 1 つ未満です。

02スイープ

その晩、Mac で実行のキューを動かし始めました。1 回に 1 つのアイデアを試し、毎回同じベースラインと比べます。キューは翌日の夜まで動きました。

試したことを順に挙げます。

畳み込み寄りの初期配置のハイブリッドは、両方のシードで、すべての種類のテキストで Llama に勝ちました。

2 つのシードの差は 0.002 以内です。実際のモデルと同じく、検索回路と掲示板回路を残しても、すべての種類のテキストで勝ちます。この設定では働かない記憶用の回路が、0.010 のコストになります。

03大きなモデル、同じステップ数

27M パラメータ（幅 384、8 層、2 シード）では、ハイブリッドのリードが 0.023 から 0.045 へと 2 倍になりました。このサイズでは、Llama は学習率を下げたほうが良くなります。そこで、各モデルをそれぞれの最適な学習率で確かめました。それでもハイブリッドは、すべての種類のテキストで 0.031 リードしています。48M では、同じ 3,000 ステップで 0.013 リードしています。

04計算量をそろえた比較：速度の問題

ステップあたりで勝つのは、テストの半分にすぎません。ステップあたりでは良くても遅いモデルは、両方に同じ学習時間を与えると負けることがあります。これが計算量をそろえた比較です。アーキテクチャを使う価値があるかどうかは、この比較で決まります。

Mac では、ハイブリッドは 1 ステップに Llama の 1.3 倍の時間がかかっていました。アテンションは、1 つの融合カーネルで動きます。TONN の回路は、素の PyTorch で動きます。1 つの層をプロファイルすると、時間の使い道がわかりました。

• 2 つのゲートつき再帰回路 GI と SM は、1 ステップずつ動いていたため、127 ms と 145 ms かかっていました。並列プレフィックススキャンは、同じ累積和を対数回のパスで計算します。それぞれ 5 倍と 3.3 倍速くなりました。
• ソフトマックスプーリング回路では、FFT を 2 回から 1 回にしました。
• 回転位置埋め込みの計算を軽くしました。
• TONN 層は、活性化を並べ替える代わりに、小さな重み行列を回路の種類の順に並べ替えるようにしました。計算する関数は 3e-7 の精度で同じです。1 ステップあたり 9% 速くなりました。

05計算量をそろえた比較：結果

各ペアに、同じ GPU で同じ秒数を与えます。各モデルは、それぞれの最適な学習率で学習します。

• 13M、600 秒： 引き分け（どちらも 1.527）。文章では Llama がわずかに上で、コードではハイブリッドが上です。
• 27M、1,000 秒： 両方のシードで、ハイブリッドが 0.006 差で勝ちます。
• 48M、2,000 秒： 最初の試行は 0.020 差で負けました。同じ時間で、Llama は 23% 多くステップを進めました。ステップ数をそろえた比較では、ハイブリッドがまだ勝っていました。したがって、負けた原因は速度です。アーキテクチャの強みが薄れたからではありません。
• 48M、TONN 層 1 つにつきアテンション層 2 つ： ハイブリッドが再び勝ちます。1.445 対 1.452 で、両方のシード、すべての種類のテキストで上回ります。アテンション層を増やすと、速くなります。残った TONN 層が、引き続き改善を担っています。

この 2:1 の構成が、現在の標準です。3 層に 1 層が TONN です。学習予定のサイズでは、トークンあたりのコストは Llama の 1.01 倍です。

060.1B での結果

実際のモデルは、幅 640、14 層です。TONN は第 0・3・6・9・12 層に置きます。32k トークナイザーでは約 90M パラメータになります。両方のモデルに、Mac で 3,000 秒を与えます。

Llama と同じ学習率で、ハイブリッドはバイトあたりビット数で 0.017 リードしています。すべての種類のテキストで上回っています。

07この記事が示していないこと

• 0.1B では 1 シードだけです。 シード 1 を実行中です。いまは Llama を回しています。次は学習率 1e-3 のハイブリッドです。学習率 1.5e-3 のシード 1 のハイブリッドは、すでに 1.482 で終わりました。シード 0 の 1.474 より悪い値です。ステップ数は、3,151 に対して 2,902 でした。その実行中、Mac の動作が遅かったためです。比べるためのシード 1 の Llama は、まだありません。13M から 48M では、計算量をそろえた比較の結果は、すべて 2 シードで成り立ちました。
• 0.1B では、Llama は 1 つの学習率でしか実行していません（1e-3。27M と 48M での最適な学習率です）。より良い学習率なら、差が縮まる可能性があります。
• どちらのモデルも、学習が大きく足りていません。90M パラメータに対して 50M から 57M トークンで、予定の 100B トークンには遠く及びません。学習を長くすると、差は縮む可能性があります。13M では、ステップ数を 2 倍にすると、リードが 0.014 から 0.006 に縮みました。
• 速度は Mac の GPU 1 つで、コンパイルなしで測りました。 この GPU では、torch.compile で Llama は 20%、ハイブリッドは 8% 速くなります。そのため、いまコンパイルありで比べると Llama が有利です。コンパイルしやすい TONN 層は、to-do リストに入っています。

すべての実行の順位は、結果のページにあります。