AI・テクノロジー

In-Place Tokenizer Expansion for Pre-trained LLMs — 多言語デコード高速化の実践的拡張

docker agent

• 6 min read
Share

はじめに

事前学習済みの大規模言語モデルは、初期データセットに基づいた固定の語彙表を前提に設計される。新規言語や新語の追加に対して柔軟性が低く、トークン化の断片化が発生するとデコード遅延や計算コストが増大する課題が生じる。こうした実務的な問題に対し、本手法はトークナイザーを"インプレース"で拡張する実用的なレシピを提案する。既存のBPEマージを維持しつつ、新規トークンを追加できる点が特徴である。追加トークンは source sub-token embeddings の平均を初期値として採用することで、埋め込み表の整合性を崩さずに新語を取り込むことができる。既存トークンの表現はそのまま保持され、継続性が保たれる。

提案は二段階の適用を前提とする。まず Embedding の訓練のみを行い、その後、全モデルの継続事前学習へ移行する。対象は LFM2-8B-A1B MoE の継続チェックポイントを128Kトークナイザーへ拡張する設定で、オンデバイス運用を前提とした現実的な適用性を検証する。実証結果では、ヒンディー語・ベトナム語・タイ語などでデコード速度が従来比およそ2.2〜3.7倍、Thai では最大約4.0倍の改善が報告されている。これにより、規模の小さなモデルでも多言語対応を実務的なコストで実現可能になる可能性が示唆される。

技術解説: In-Place Tokenizer Expansionの仕組み

トークナイザー拡張は、事前学習済みモデルの語彙拡張を安全に行う実務的レシピであり、BPEマージを崩さず新規トークンを導入できる点が特徴である。背景として、固定されたトークナイザーは後から追加される語彙に追従できず、トークン化が断片化してデコード遅延や計算コストの増加を招く。本技術は"インプレース"での拡張を提案する。

拡張の基本原理は、既存トークンを可能な限り保持し、単一トークンとしての表現を維持することだ。新規トークンは source sub-token embeddings の平均で初期化され、carried-over embedding rowsはそのままコピーされる。これにより、元の語彙構造を崩さずに追加語彙を滑らかに統合できる。

二段階適応として、最初は Embedding の訓練のみを行い、その後全モデルの継続事前学習へ移行する。これにより新規トークンの埋め込みが安定し、後続の学習負荷を分散できる。

実証の対象は LFM2-8B-A1B (8B MoE) の継続チェックポイントを用い、128K トークナイザーへの拡張を目指す。実証結果の要点は、ヒンディー語・ベトナム語・タイ語で従来より1文字あたりのデコード速度が約2.2〜3.7倍、Thai では最大約4.0xの改善が確認された点である。

一方、拡張の過程での課題や反省点も報告され、実務運用では語彙管理や埋め込み初期化の選択がパフォーマンスに影響する点を留意する。

flowchart TD
  T[Tokenizer] --> E[Embedding Matrix]
  N[New Tokens] --> A[Avg Embedding Init]
  E --> N
  subgraph Phases
    U1[Embedding-only Training]
    U2[Full PT]
  end
  U1 --> U2

実務的な適用ガイド

適用条件と前提として、既存チェックポイントを再利用し、語彙拡張の計画とモデルアーキテクチャの互換性を事前に確認する必要がある。対応言語の語彙分布を把握し、追加トークンの上限・辞書の構造を設計する。MoEや128Kトークナイザーを想定する場合は、埋め込み行列の再配置とインデックス管理の影響を検証しておく。

実装の要点は以下のとおりである。

  • 既存トークンを保持しつつ、新規トークンは埋め込みの平均で初期化する
  • 追加トークンの辞書マッピングと既存語彙との整合性を保つ
  • 辞書サイズ拡大に伴うロード時間・メモリ使用量の影響を見積もる

学習フェーズでは、最初にEmbeddingのみの訓練を実施することで埋め込みの安定化を図り、続いて全モデルの継続事前学習へ移行する。学習率・バッチサイズ・正則化を慎重に調整することが重要である。

評価指標としては、トークン化の断片化指標とデコード速度、実運用での生成品質の変化（翻訳・要約の出力妥当性）、ハードウェア負荷と推論の安定性を観察する。

実用的な適用例として、MoEアーキテクチャで128Kトークナイザーを活用する場合の利得とリスクを比較する。追加語彙の管理を適切に行えば、オンデバイスでの多言語運用の実用性が高まる。

実務上の留意点とリスク

  • 多言語対応と速度のトレードオフ: 新規トークンを追加すると語彙が拡大し、埋め込み行列のサイズが増す。デコード時の計算量は増減のバランス次第で、速度向上の恩恵が薄まる場合や、メモリ負荷が増す場面がある。現場では、用途言語の優先度とデバイスの搭載メモリを見極め、必要に応じて部分語彙だけを先行導入する設計が現実的である。
  • モデル品質と過学習のリスク: 二段階適応の狙いは既存トークンの安定性を守ることだが、追加トークンの初期化と後続の全モデル訓練で挙動が不安定化する可能性がある。品質回復には評価指標の定期的なモニタリングと、段階的な学習率・訓練データの調整が必要である。
  • ハードウェア依存性: エンコード/デコードのボトルネックはデバイスの計算資源と電力コストに左右される。特に多言語の分岐やMoEはメモリ帯域と通信オーバーヘッドの影響を受けやすく、実運用では省エネ設計や適切な batching が重要である。
  • MoE特有の挙動: パラメータ分配や埋め込みの再初期化に関わる挙動は、トークン追加後のゲーティング・ルーティングに影響を及ぼす可能性がある。適用時にはデプロイ環境での挙動検証と、埋め込み再整列時の影響範囲の確認が欠かせない。

今後の展望と応用可能性

エッジデバイス上での多言語運用は、計算資源の制約と通信コストの低減という観点で大きな普及余地がある。追加トークナイザーを前提に、128K級の語彙を現実的なメモリ使用量で運用する設計は、家庭用サーバーや業務端末でのオンデバイス推論を現実的な選択肢へと導く。既存モデルのチェックポイントを再利用しつつ、語彙拡張を段階的に適用するアプローチは、セキュリティやデータ保持の観点にも適合する。

今後のメンテナンス性を高めるには、辞書管理の標準化と互換性維持が鍵である。追加トークンのマッピング、埋め込みの再初期化方法、アップデート時の影響評価を自動化する枠組みが有効だ。語彙再編成の自動化、デコード品質の保持とコスト削減の両立を目指す研究は、オンデバイスの実用性をさらに高めるだろう。

実装の具体的な展望としては、MoEアーキテクチャとの組み合わせ、動的語彙拡張のキャッシュ戦略、デバイス別の最適化プリファレンスの導入などが想定される。長期的には、追加言語の自動検出・適用、継続的学習と組み合わせた運用モデルが現実的な選択肢になる可能性がある。

図解の提案

本節では、トークナイザー拡張の流れを直感的に把握できる図をMermaidで提示する。既存トークンを保持しつつ、新規トークンを埋め込みベクトルで初期化する処理と、Embedding のみの第一段階訓練から全体継続学習へ移行する二段階適応を図示する。図はPretrained LLM・Tokenizer・Deploymentの三者のやりとりを中心に、オンデバイス推論の速度向上と多言語対応の効果を想定している。

以下のMermaid図は主要な関係と処理の順序を示す。

sequenceDiagram
  participant M as Pretrained LLM
  participant T as Tokenizer
  participant D as Deployment
  M->>T: トークン化リクエスト
  T-->>M: トークン列
  Note over M,T: 既存トークン保持、新規トークンは埋め込み平均で初期化
  M->>M: 二段階適応 (Embeddingのみ → 全体学習)
  M->>D: デコード/生成

この図を用いると、語彙の再構成や新規語彙の扱い、学習切替のタイミングが視覚的に理解しやすくなる。オンデバイス運用を想定する読者にとって、設計判断の指針として有用である。

結論

本手法は、事前学習済みの大規模言語モデル（LLM）のトークナイザーを追加学習なしで拡張する、実務寄りのレシピとして有用性を示す。128K規模のトークナイザーへ拡張する試みは、オンデバイス運用における多言語対応の現実的な道を開き、特定言語のデコード速度を顕著に改善する可能性がある。実証では、既存チェックポイントの再利用と、追加トークンを埋め込みの平均値で初期化する戦略が、学習コストを抑えつつ性能を維持・向上させることに寄与した。二段階適応（Embeddingのみの事前学習と、その後の全体継続学習）の組み合わせは、速度と品質のバランスを取りやすくする。

今後の展望としては、追加言語の自動拡張アルゴリズムの開発、辞書管理の安定化、品質保証の強化、省エネ設計との統合が挙げられる。さらにMoEアーキテクチャやエッジデバイスでの運用を想定した評価フレームを整備することで、実践的な適用範囲を拡大できる。

参考資料

  • arXiv:2607.15232 — In-Place Tokenizer Expansion for Pre-trained LLMs