技術概要：分割木を用いたトークナイズ（ToaST）

問題定義

トークナイズは自然言語処理における基礎的なステップであるが、既存の手法はしばしば圧縮や下流タスクのパフォーマンスを最適化しない貪欲なヒューリスティクスに依存している。バイトペアエンコーディング（BPE）は貪欲なマージ戦略を通じてトークン数を最小化するが、これはヒューリスティックな近似に過ぎない。先行研究は、圧縮を最大化するトークナイザーの発見が、有限アルファベット上であっても NP 完全であることを示している。さらに、圧縮と下流タスクのパフォーマンスとの関係については議論が続いており、いくつかの研究は相関を指摘する一方、他の研究は圧縮だけではトークナイズの品質を説明するには不十分であると主張している。UnigramLM や PathPiece といった既存の代替手法は、異なる推論手順（例えば、除去型のトップダウン推論や最短経路推論）を利用するが、依然として大域的最適化と計算のスケーラビリティにおいて課題を抱えている。

手法

著者は、候補トークンの構築と最終語彙の選択を分離するサブワードトークナイズ手法である**分割木を用いたトークナイズ（ToaST）**を導入する。このアプローチは以下の 3 つの主要な段階で構成される。

1. 分割木の構築

BPE がボトムアップでトークンを構築するのに対し、UnigramLM がトップダウンで確率を学習するのとは異なり、ToaST は訓練コーパス内の各一意のプリトークンに対して、いかなる語彙にも依存せず、完全な二分木を構築する。

• 前処理: 訓練データを正規表現を用いてプリトークンに分割する。
• 分割ロジック: 与えられたプリトークンに対して、アルゴリズムはそれを 2 つの非空部分に再帰的に分割する。分割点は、コーパスから事前計算された n-グラム統計に基づき、生成される左側と右側のバイト n-gram の最小出現回数を最大化するように ($\max \min(c_p, c_{p'})$) 決定論的に選択される。
• 構造: このプロセスは、葉が単一のバイトになるまで継続される。これらの木のノードは、すべての候補トークンの集合を形成する。この構築は最終的な語彙に依存せず、反復的な訓練を必要としない。

2. 再帰的推論

語彙 $V$ が与えられた場合、トークナイズは分割木を再帰的に下降することで実行される。

• 根（完全なプリトークン）から開始し、アルゴリズムは現在のノードが $V$ に含まれるか確認する。
• ノードが $V$ に含まれる場合、それをトークンとして出力する。
• 含まれない場合、アルゴリズムは 2 つの子ノードを再帰的に処理する。
• 特性: この手順により、必要なすべての単一バイトを含む任意の語彙が有効なトークナイズを生成することが保証される。語彙からトークンを削除するだけで、マージベースの手法とは異なり、他のトークンへの連鎖的な影響なく、推論が木をさらに深く下降させることになる。

3. 整数計画法による語彙選択

中核的な革新は、訓練コーパス全体での総トークン数を最小化するために、語彙選択を**整数計画法（IP）**として定式化することである。

• 変数: 二値変数 $x_i$ は、候補トークン $i$ が語彙に含まれるかを示す。二値変数 $z_{jk}$ は、分割木 $j$ 内のノード $k$ が推論中に出力されるかを示す。
• 目的関数: 出力されるトークンの重み付き総数を最小化する：$\min \sum_{j} c_j \sum_{k} z_{jk}$。
• 制約条件:
  • 語彙サイズ: 選択されたトークンの合計は、目標サイズ $m$ に等しくなければならない。
  • カバレッジ: 葉から根へのすべての経路に対して、正確に 1 つのノードが選択されなければならない（プリトークン全体がカバーされることを保証）。
  • 一貫性: ノードが出力される（$z_{jk}=1$）ためには、それが語彙に含まれている（$x_i=1$）かつ、その祖先ノードのいずれも含まれていない場合にのみ可能である。
• 求解: IP は**線形計画法（LP）**に緩和される。実際には、LP 緩和は「ほぼ整数解」であり、解には非常に少数の分数値しか含まれない。単純な丸めヒューリスティクスを適用して整数解を得ると、それは証明可能な近最適解（大域的最適解から無視しうる程度の差以内）となる。

主要な貢献

1. 新しいトークナイズフレームワーク: ToaST は、BPE の貪欲な近似や UnigramLM の特定の推論制約を回避する、分割木ベースのアプローチを導入する。
2. 大域的最適化: 語彙選択を tight な LP 緩和を伴う IP として定式化することで、ToaST は局所的な貪欲な決定に依存するのではなく、与えられた分割木の集合に対して総トークン数を大域的に最小化する語彙を見つける。
3. スケーラビリティ: この手法は分割木の数に対して二次的にスケーリングするため、実用的な訓練時間内で大規模コーパス（プリトークン出現回数の 99% をカバー）の処理が可能である。
4. 分離された設計: 木構築（コーパス統計に基づく）と語彙選択（最適化に基づく）の分離により、最適化モデルを再訓練することなく、形態素や文字レベルの好みを組み込むなどの柔軟な拡張が可能である。

実験結果

著者は CulturaX データセットを用いて英語テキスト上で ToaST を評価し、BPE、WordPiece、UnigramLM と比較した。

• 圧縮: 語彙サイズが 40,960 以上の場合、ToaST はすべてのベースラインと比較してトークン数を**11%**以上削減する。これはバイト/トークンで測定される。
• トークン分布: 語彙サイズ 65,536 において、ToaST はベースラインと比較して「Leaf」トークン（単一バイトのフォールバック）の頻度を 14〜19 倍大幅に削減する。逆に、「Root」トークン（完全なプリトークン）を大幅に多く生成し、意味単位の良い保持を示している。
• Rényi 効率: ToaST は高い Rényi 効率（$\alpha=2.5$）を達成し、より均一なトークン分布と高頻度トークンの減少を示しており、これはしばしば下流タスクのパフォーマンス向上と相関する。
• 下流タスクのパフォーマンス: 15 億パラメータの言語モデル（nanochat）を訓練する実験において、ToaST は最高CORE スコア（22 のベンチマークの平均）を達成し、ベースラインを**2.6% から 7.6%**上回った。BPE や UnigramLM を大幅に上回り、22 の個別タスクのうち 13 で最良のパフォーマンスを示した。
• 最適化の品質: LP 緩和は極めて tight であることが判明した。128 の問題インスタンス全体で、丸められた整数解と LP 境界との最大相対ギャップは $5.76 \times 10^{-7}$ であり、多くのインスタンスは完全に整数解であった。

意義と主張

本論文は、ToaST が貪欲なヒューリスティクスに依存するのではなく、大域的最適化フレームワークを用いて圧縮のために直接トークナイズを最適化することの価値を実証していると主張している。著者らは、優れた圧縮がより大きな実効コンテキストウィンドウをもたらし、その結果として下流タスクのパフォーマンスが向上すると論じている。

この研究の意義は以下の点にある。

• 実用性: 最適なトークナイズの理論的な NP 完全性にもかかわらず、LP 緩和によりこのアプローチは計算的に実行可能となり、大規模訓練に対して実用的である。
• パフォーマンス: 経験的結果は、トークン数の 11% の削減が言語モデルのパフォーマンス（CORE スコア）における測定可能な向上につながり、より良い圧縮がより良いモデル有用性と相関するという仮説を支持している。
• 柔軟性: このアーキテクチャは、語彙最適化ロジックを変更することなく木構築階層を修正することで、マルチバイト文字や形態素分割の処理など、将来の拡張を可能にする。

著者は、実験が英語テキストに限定されており、非スペース区切り言語における上位集約プリトークンでの訓練の有効性は異なる可能性があるという限界を指摘している。また、コードの実装や図の生成に LLM を使用したことを開示している。