概要
ローカルWindow AttentionとGRU風の回帰構造を組み合わせたHybridAttentionにより、Rust言語生成を大幅に高速化した25.6MパラメータのLLMの実験報告。
詳細内容
本プロジェクトは、Rust言語に特化した2,560万パラメータの小規模言語モデル（SLM）をPyTorchを用いてスクラッチで構築した試みです。最大の特徴は、従来の全結合Attention（O(n²)）を、ローカル・ウィンドウ・アテンションとGRU（Gated Recurrent Unit）のような回帰パスを組み合わせた独自の「HybridAttention」に置き換えた点にあります。この設計により、推論速度が従来の5.6 tok/sから286.6 tok/sへと約51倍に向上し、計算複雑性を準線形へと削減することに成功しました。
モデルの訓練には173.5MBのRustソースコード（公式ドキュメントおよび上位500のcrate）を使用しており、RTX 4060 Tiという一般的なコンシューマ向けGPUで3万ステップの学習を行っています。KVキャッシュの管理においても、直近64トークンをVRAMに保持し、古いトークンを8ビットで圧縮保存する等の工夫が施されています。
生成結果の品質としては、関数のシグネチャやインポート文といったRustの構文構造は高い精度で再現されていますが、長期的な意味論の維持や出力の繰り返しに課題が残ると報告されています。著者によれば、アーキテクチャの変更よりもコーパスの拡充がモデルの品質向上に最も大きく寄与したとのことです。今後は生成コードのコンパイルチェックによる検証や、BPEトークナイザーの導入が検討されています。