Daedalus-150M、CPU推論を起点に設計した畳み込み・注意融合モデル

概要

小型言語モデルは、大規模モデルと同じ Transformer 構造を保ったままサイズを縮小し、完成後に量子化やランタイム最適化で CPU に適応させることが多くあります。Daedalus-150M はこの順番を逆にしました。単一ユーザーが一度に1トークンを生成し、4bit重みを普通の CPU で実行するという条件を先に固定し、その制約に合わせてネットワーク構造を選んでいます。焦点はパラメータ数だけでなく、長い文脈で増え続ける KV キャッシュを何度も読むコストです。

アテンションを限定して使う

約1.6億パラメータ、18ブロックのうち、フルアテンションを使うのは6ブロックだけです。残る12ブロックには短い深さ方向畳み込みを配置します。この畳み込みが保持する再帰状態は2タイムステップ分に限られるため、会話が長くなっても拡大する KV キャッシュを繰り返し参照しません。グローバルな依存関係はアテンションに任せ、局所的・連続的な処理の一部を畳み込みが担当する分担です。

この設計の効果は、文脈長によって速度差が変化する点にあります。空の文脈では混合モデルと全アテンションモデルの差はほぼありません。しかし文脈が長くなるほど、混合モデルの優位性が拡大します。これはキャッシュアクセスがボトルネックになるという説明と整合し、単にモデルが小さいだけの場合に予想される挙動とは異なります。

学習結果と比較実験

モデルは599億トークンでゼロから学習されました。HellaSwag、ARC-Easy、PIQA、OpenBookQA、WinoGrande の5タスクで47.31点を獲得し、学習前に設定された42.20点の基準を上回っています。GPT-2 124M、Pythia-160M、OPT-125M、GPT-Neo-125Mも上回りました。これらの比較対象は、おおむね3〜6倍のデータで学習されています。検証時の bits-per-byte は0.8685でした。

アーキテクチャの効果を確かめるため、同じ規模で、同じデータと条件を使った全アテンション版も訓練されました。勝敗条件は評価前に定められています。混合モデルは選定された品質指標で0.81%上回り、下流タスクでは同等、4bitファイルでは6.3%小さくなりました。文脈長2048トークンでは、比較対象より1.76倍高速で、外部の同規模モデルに対しては2.08倍でした。帯域幅だけから見積もると1.17倍にとどまるため、データ移動量の減少だけでは差を説明しきれません。

残る課題

論文はうまくいかなかった試みも報告しています。量子化対応学習では初期段階に4bit化による品質低下が現れ、畳み込みチャネルの約半数が不活性になりましたが、標準ランタイムの形状チェックのため簡単には枝刈りできませんでした。また、モデル規模に対して語彙を大きくしすぎることも適切ではないと判断されています。

実用面では、公開された要約に具体的な CPU 型番、バッチサイズ、絶対的な遅延やトークン毎秒がない点に注意が必要です。相対的な高速化は機構の傾向を示しますが、導入判断には実際のノート PC やサーバーのコア、ランタイム設定での測定が欠かせません。

意義

Daedalus-150M が示すのは、推論制約を圧縮の最後に処理するのではなく、アーキテクチャ設計の出発点にできるということです。畳み込みが一般にアテンションを置き換えられると証明したわけではありません。小型モデル、ストリーミング生成、長い文脈という条件では、均一に縮小した Transformer より CPU のメモリアクセス特性に合う可能性がある、という示唆です。今後は異なるプロセッサ、ランタイム、実アプリケーションでの再現が重要になります。

出典：Hugging Face Daily Papers