みなさん、こんにちは。
BitNet 1.58b の論文に沿って、極端なLLM量子化をいろいろ試していました。3値量子化 {-1, 0, 1} なら、高価な行列乗算を単純な加算に置き換えられて素晴らしいのですが、重みを過度に制限してしまうことで、モデルの能力を取りこぼしているのではないかと気になりました。
そこで、ゼロから PentaNet を構築して学習しました。重みの状態を5値に広げるための、カスタムアーキテクチャです:{-2, -1, 0, +1, +2}。
なぜ ±2 なのでしょう? 2 を掛けるのにハードウェア乗算は不要だからです! 単なる左ビットシフト(x << 1)です。これにより、PentaNet は BitNet が持つ「ゼロ乗算」推論の恩恵を完全に保持しつつ、重み1つあたりの情報量を 47% 増やします(対数で、log₂(5) ≈ 2.32 ビット vs 3値の log₂(3) ≈ 1.58 ビット)。つまり知識を符号化するための情報が増えるのです。
ベンチマーク
同じ計算予算とセットアップで、2つの124Mパラメータモデル(GPT-2アーキテクチャ)を WikiText-103 で学習し、正面から比較しました。統計的な有意性を確保するため、各モデルにつき独立した3つのシードで実行しています。
結果(WikiText-103):
計算オーバーヘッドの観点では実質「タダ」みたいなもので、約6.4%のパープレキシティ改善が得られました。そして Straight-Through Estimator(STE)も完全に安定したままでした。
重み分布 & 崩壊しないこと
私が最も心配していたのは、モデルが ±2 のバケットを無視して、何も言わずに ternary の BitNet に戻って(崩壊して)しまうことでした。学習中にバケットを追跡したところ、実際には見事に安定しています:
️ テキスト生成の例
論文上では PPL の差が小さく見えるかもしれませんが、124Mパラメータの規模だと、その差は英語の「どもり」と「首尾一貫した英語」を分けます。以下はシード42からの、手作業で選んでいないサンプルです(プロンプト:「The history of the internet began with」):
BitNet:
The history of the internet began with the <unk> to be a way , <unk> , which was the first recent of the <unk> , and the city and the <unk> . The French army was the first to be the first @- * scale*
PentaNet:
The history of the internet began with the original level of the other . The term of the original world was to the public court of the United States in July 2013 in February 15 , 2015 , as well as the team of $ 2 @,@ 000 . In the same year , the
(もちろん、このモデルは20分学習した小さなモデルなので、事実としてはハルシネーションしています。しかし、PentaNet が実際に流暢な文法を学習し、<unk> の崩壊を回避していることに注目してください!)
リンク & コード
学習コード、PyTorch の PentaLinear レイヤの実装、そして NeurIPS 風の技術ドラフトをオープンソースにしました。
- HuggingFace(Weights):Kyworn/pentanet-
現時点では、PyTorch のレイヤが学習時のために量子化をシミュレーションしています。次の自然なステップは、bit-shift 操作を実際に活用して、現実の速度向上につなげるためのカスタム Triton/CUDA カーネルを書くことです。
ぜひ皆さんのご意見を聞きたいです。特に、この種の量子化推論のために低レベルのカーネルを書いた経験がある方がいれば、なおうれしいです!
[link] [comments]




