AIの文章区切りが5倍速く

Perplexity AIの研究チームは、Unigram tokenizerをRust言語でゼロから再実装し、彼らの推論技術リポジトリであるpplx-gardenでそのコードをオープンソース化しました。本番環境の入力長において、新しいエンコーダーはHugging Face tokenizers crateと比較してp50レイテンシを約5倍削減し、SentencePiece（C++）と比較して約2倍、IREE’s tokenizer（C）と比較して約1.5倍削減しました。また、定常状態でのヒープ割り当てはゼロです。本番環境では、Perplexityの推論スタックにおけるCPU使用率を5〜6倍削減し、rerankerのレイテンシを2桁ミリ秒短縮しました。なぜトークン化がボトルネックになったのか。LLMの推論コストは通常、GPUの作業（KV caches、attention kernels、expert routingなど）を中心に語られます。しかし、埋め込みモデル、分類器、rerankerなどの小さなモデルでは、話が異なります。これらのモデルは、最先端のTransformerモデルよりも2〜3桁小さいです。1つのリクエストで数百の候補ドキュメントをスコアリングするrerankerはその明確な例です。小さなモデルでは、GPUの計算はしばしば1桁ミリ秒で完了します。しかし、すべての入力はまずCPU側のトークン化を通過します。バッチサイズが大きい場合、トークン化は合計リクエストレイテンシの重要な部分を占めるようになります。Perplexityの作業は、SentencePieceで学習された25万トークンのUnigram語彙を持つXLM-RoBERTaを対象としています。ファインチューニングされたRoBERTaファミリーのエンコーダーは、ランキング、検索、類似性タスクにおいて一般的な本番環境での選択肢です。Unigramトークン化とは何か。Unigramトークン化は2018年にKudoによって導入され、SentencePieceで実装されています。これは、セグメンテーションを最も可能性の高いパス問題として捉えます。各語彙トークンには学習された対数確率があります。トークナイザーは、トークンスコアの合計が最も高いセグメンテーションを選択します。その最適なパスを見つけるために使用されるアルゴリズムは、1967年の動的計画法技術であるViterbi algorithmです。バイト位置がグラフの層を形成し、語彙トークンは連続するバイト範囲にまたがるエッジとなります。DPの反復はバイト位置を繰り返します。