TLDR: pytorch と triton の内部をフォークしました。 attention を変更し、最初の層は線形、中央の層は二次、最後の層は線形にしました
推論は、テストでのパープレキシティの低下が小さいまま、かなり高速化しました。
主要な結果は、データセットサイズを増やすことが、あらゆるアーキテクチャ変更よりも重要だったという点です。
25.6M パラメータの Rust 向け言語モデルを、バイトレベルの GPT スタイルデコーダを使ってスクラッチから学習させました。
コアとなる Rust ソースの約 31MB から、数百個のクレートを追加して約 173MB までコーパスを拡張したところ、他のどんな要素よりも大きな改善が得られました。学習はより速く収束し、検証損失もより低くなりましたが、アーキテクチャ変更の影響は小さめでした。
最終的な検証損失は 0.82、パープレキシティは 2.15 です。最良のチェックポイントはステップ 18.5k 付近に見え、その後は軽い過学習が起きています。
各層では、標準的な attention を、ローカルなウィンドウ付き attention と GRU のような再帰的状態を組み合わせたハイブリッド機構に置き換えます。これは学習されたゲートで混合されます。ローカル経路は短距離の文法情報を捉え、再帰経路は圧縮された長距離の情報を運びます。
このハイブリッド attention は、標準的なセットアップと比べて生成品質を明確に改善しませんでした。しかし、推論効率には大きな影響がありました。
VRAM 上に小さな直近ウィンドウを保持し、より古いトークンを圧縮する KV キャッシュにより、4060 Ti では推論が 5.6 tokens per second から 286 tokens per second に改善しました。これは出力品質が明確に落ちることなく、約 50 倍の高速化です。
モデルはもっともらしい Rust の構文や構造を生成しますが、意味的一貫性はまだ弱く、繰り返しがよく起きます。
次のステップは、ハイブリッド、ローカルのみ、再帰のみの各バリアントを比較するアブレーションを実行すること、生成品質の観点でより早いチェックポイントを評価すること、パースやコンパイルなどコード固有の評価を追加すること、そしてより長いコンテキストと BPE トークン化をテストすることです。
パープレキシティ以外の、小さなコードモデル向け評価方法についてのフィードバックが欲しいです。また、コード生成においてハイブリッドなローカル attention と再帰 attention が実際にうまく機能したかどうか、そしてこのスケールでのさらなる改善が、より多くのデータ、より長いコンテキスト、あるいはアーキテクチャ変更から得られる可能性が高いのはどれかについても知りたいです。
[link] [comments]
