埋め込みは高すぎる？固定トークンコードで小型LMを46MBに圧縮

概要

小型言語モデルでは、語彙を表すEmbedding表が本体より大きくなることがあります。Gemma 3 270Mでは、全268Mパラメータのうち約168MをEmbedding表が占めます。さらに他の層だけを量子化すると、Embeddingは高い精度のまま残り、最終ファイルの大部分を占める場合があります。

Hugging Face Daily Papersで紹介された研究は、この設計を見直しました。研究者は学習可能なEmbedding行列の代わりに、各トークンへ512ビットの固定コードを割り当てています。コードは事前学習済みEmbeddingを主成分投影し、その値の符号を二値化して作成します。一度生成した後、語彙表は学習中に更新されません。

主な結果

• 学習パラメータを削減：通常の語彙表を使うモデルと固定表のモデルを、同じGemma型の構造、同じ20.5億トークンでゼロから学習しました。学習可能な表は検証損失を0.20 nat改善しましたが、必要な学習パラメータは約2.6倍でした。
• 下流評価で競争力：LAMBADA、HellaSwag、PIQAでは固定表モデルが上回るか同等でした。一方、ARC-Easyでは学習可能な表が優位でした。
• 希少語と追加語に強い：希少語の評価では固定表が1.7〜2.3 nat上回りました。学習後に1万語を追加した場合の性能低下は約3%でしたが、学習可能な表では約43%低下しました。
• 端末向けの小型化：BitNet b1.58の三値化本体と組み合わせると、モデル全体は約46MBになります。ノートPCのCPU上で毎秒615〜740トークンを生成し、メモリ使用量は約62MBでした。llama.cppの浮動小数点モデル用8ビットGGUFとの比較では、速度は約5〜6倍、メモリは約7分の1と報告されています。

意義と限界

この研究は、小型モデルが大規模で学習可能な語彙表を必ず必要とするという前提に疑問を投げかけます。固定コードに置き換えれば、ストレージの余力をTransformer本体へ回せます。また、希少語や学習後の語彙追加でも安定しやすく、スマートフォンやノートPCでのオフライン推論に適した設計になり得ます。

ただし、Embeddingが一般に不要だと結論づけるものではありません。今回の実験は小型モデルと特定の学習トークン数を対象にしています。モデル規模やデータ量が増えたとき、調整可能なEmbeddingが持つ柔軟性を固定コードが維持できるかは未検証です。言語、文字体系、ドメイン、希少文字ごとの分析も必要でしょう。

今後は、新しいトークンのコードをどう作るのか、射影の再計算が必要か、学習量の増加で固定表がボトルネックになるかを確認する必要があります。現時点では、これはEmbedding全般を否定する結果ではなく、低ビット小型モデルの端末向け設計として有力な選択肢と見るのが妥当です。

出典：Hugging Face Daily Papers