事前学習済みLLMのトークナイザーをその場で拡張、多言語処理を最大4倍高速化

事前学習時に固定されるトークナイザーは、その時点の学習データ構成やデプロイ優先度を反映して語彙が決まる。後から重要度が増した言語は1単語あたりのトークン数が多くなりがちで、レイテンシや計算コスト、消費電力の増加につながる。クラウド向けの大規模モデルでは埋め込み層とLMヘッドの割合が小さいため広い語彙を許容できるが、オンデバイス向けの小型モデルではこれらの行列がデコード帯域の大きな部分を占めるため、語彙を絞り込まざるを得ず、対応外言語ではトークンの断片化が発生していた。

本研究は、モデル開発者がトークナイザー設計を管理できる場合に、事前学習済みモデルのトークナイザーをその場で（in-place）アップグレードする手法「tokenizer expansion」を提案する。

手法の概要

既存トークナイザーのBPEマージ処理を多言語コーパス上でそのまま継続することで、既存の多くのトークンは単一トークンとして変わらず引き継がれ、新しいトークンもすべて既存トークンへの厳密な分解が可能になる。引き継がれた埋め込み行はそのままコピーし、新規行はそのトークンを構成する既存サブトークン埋め込みの平均として初期化する。その上で、まず埋め込み層のみを学習し、続いてモデル全体を継続事前学習するという2段階の適応プロセスにより、拡張前のチェックポイントと同等の品質を回復させる。

適用結果と性能向上

この手法を、80億パラメータのMixture-of-ExpertsモデルであるLFm2-8B-A1Bの継続事前学習済みチェックポイントに適用し、128Kトークナイザーを備えたLFM2.5-8B-A1Bを構築した。拡張後のトークナイザーは、ヒンディー語で約2.4倍、ベトナム語で約2.6倍、タイ語では最大4.0倍もトークン数を削減した。これらの削減効果と、拡張後の語彙による1トークンあたりの実測コストを組み合わせると、対象言語において参照デバイス上で1文字あたりのデコード速度が2.2〜3.7倍向上すると見積もられている。

研究チームはモデルの重みと拡張済みトークナイザーを公開し、開発過程で得られたうまくいかなかった手法についても報告している。