学習済みLLMのトークナイザを『その場で』拡張する手法 — 多言語対応を後付けで改善

事前学習の開始時点で固定されるトークナイザは、その時点の学習コーパスの構成比率に応じて語彙を割り当てる。しかし後から優先度が変わり、対応言語が追加されると、その言語は1単語あたり多くのトークンに分割されてしまい、レイテンシ・計算量・消費電力の増加につながる。クラウド上の大規模モデルは埋め込み層やLMヘッドがパラメータ全体に占める割合が小さいため広い語彙を許容できるが、オンデバイス向けの小型モデルではこれらの行列が復号帯域の大きな割合を占めるため、語彙を絞り込んで一部言語の非効率を甘受せざるを得なかった。

提案手法

本論文は、モデル開発者がトークナイザ設計を制御できる場合に、学習済みモデルのトークナイザを「その場で」（in-place）アップグレードする手法「tokenizer expansion」を提案する。既存トークナイザのBPEマージを多言語コーパス上で継続することで、既存トークンの大半はそのまま単一トークンとして引き継がれ、新しいトークンもすべて既存トークン列への正確な分解が可能になる。引き継がれた埋め込み行はそのままコピーし、新しい行は元のサブトークン埋め込みの平均で初期化する。その後、埋め込みのみの学習→モデル全体の継続事前学習という2段階の適応によって、元のチェックポイントの品質を回復させる。

結果

8BパラメータのMixture-of-ExpertsモデルLFM2-8B-A1Bの継続事前学習チェックポイントにこの手法を適用し、128Kトークナイザを備えた「LFM2.5-8B-A1B」を作成。拡張後のトークナイザは、ヒンディー語とベトナム語を元の約2.4〜2.6倍少ないトークン数でエンコードでき（タイ語では最大4.0倍）、これらの言語における復号速度は参照デバイス上で2.2〜3.7倍高速化すると見積もられている。著者らはモデルの重みと拡張済みトークナイザを公開し、開発過程で得られたネガティブな知見も報告している。