技術的概要：TIDE（Token Identity Delivered Everywhere）

1. 問題定義

本論文は、現代の大規模言語モデル（LLM）における根本的かつ未検証の設計選択、すなわち単一注入仮説を特定している。標準的なトランスフォーマー・アーキテクチャでは、トークンインデックスが入力埋め込み層で 1 回だけ参照され、初期の隠れ状態を生成した後に、明示的なトークン識別子が恒久的に破棄される。その後の$L$層にわたるすべての計算は、文脈化された隠れ状態のみに依存して行われる。著者らは、この仮説が 2 つの明確な構造的欠陥を引き起こすと主張している。

1. 希少トークン問題：ジップの法則により、自然言語の語彙はべき乗則分布に従い、希少トークン（例：専門用語、まれな固有名詞）は頻繁に出現しない。確率勾配降下法（SGD）の下では、トークン埋め込みに対する累積勾配信号は、その出現頻度に比例する。その結果、希少トークンの埋め込みは、一般的なトークンと比較して勾配信号の断片しか受け取らず、慢性的な未学習状態に陥る。実証的証拠は、トレーニングが進むにつれて、希少トークンの埋め込みは一般的なトークンと比較して低い L2 ノルムを持ち、成長率が低下することを示している。
2. 文脈的崩壊問題：意味的に異なる 2 つのトークンが、ほぼ同一の構文環境（例：同音異義語である"their/there"や数値トークン）に現れる場合、文脈残差ストリームは限られた識別信号しか提供しない。フィードフォワードネットワーク（FFN）は連続的な隠れ状態空間におけるリプシッツ連続性に依存しているため、隠れ状態が区別不能なベクトルに崩壊したトークンを区別することができない。破棄されたトークンインデックスは、この識別を回復するメカニズムを提供しない。

2. 手法：TIDE アーキテクチャ

これらの欠陥に対処するため、著者らはTIDE（Token Identity Delivered Everywhere）を提案する。これは、標準的なトランスフォーマーを並列かつグローバルに共有されるEmbeddingMemoryモジュールで拡張するアーキテクチャ的変更である。このモジュールは、トークン識別子が文脈残差ストリームとは独立して、ネットワークのすべての深さでアクセス可能であることを保証する。

コアコンポーネント

• EmbeddingMemory アンサンブル：TIDE は、$K$個の独立したMemoryBlocksを導入する。各ブロック$k$は、生トークンインデックス$v$を直接、文脈に依存しない意味ベクトル$M_k(v)$にマッピングする専用埋め込みテーブル$E_k \in \mathbb{R}^{|V| \times d_b}$を含む。これらのベクトルは、単一の埋め込み参照と RMSNorm を通じて、1 フォワードパスあたり 1 回計算される。
• 深さ条件付きルーター：各トランスフォーマー層$\ell$において、軽量な線形ルーターが、アテンション後の正規化された隠れ状態$\tilde{n}_\ell$を受け取り、softmax 関数を介して合成比率$\alpha^\ell_k$を生成する。このルーターは、各 MemoryBlock の信号をどの程度注入するかを決定する。
• Null Bank：ルーターには$(K+1)$番目のスロット、すなわち"Null Bank"が含まれており、ここでは$M_{K+1}(v) = 0$である。これにより、モデルは「オフ」スイッチを学習することができ、不要な場合はメモリ注入を抑制することで、TIDE が標準的なトランスフォーマーに漸近的に一般化することを可能にする。
• 加法的融合：最終的なメモリベクトル$m^\ell(v)$は、$K$個のメモリブロックと Null Bank の加重和である。このベクトルは、すべての層で残差ストリームに加法的に追加される：
  $$h_\ell = \tilde{h}_\ell + \text{FFN}(\tilde{n}_\ell) + m^\ell(v)$$
  重要なのは、メモリ経路が FFN 経路とは独立しており、残差ストリームを共有通信チャネルとして維持しつつ、永続的でトークンを識別する信号を提供することである。

計算効率

メモリへの寄与は、埋め込み参照と加重和のみを含むため、ベースラインの FFN に比べて無視できるオーバーヘッドしか追加しない。さらに、EmbeddingMemory テーブルは離散的なトークン識別子でインデックス付けされ、トレーニング後に静的であるため、量子化（例：4 ビット）が可能であり、SSD ストレージにオフロードできる。これにより、アクティブな VRAM フットプリントを標準モデルと同等に保つことができる。

3. 主要な貢献

理論的貢献

• 漸近的一般化：著者らは、すべての信号を Null Bank にルーティングすることで、TIDE が任意の精度で標準的なトランスフォーマーを近似できることを証明している。これにより、TIDE が標準アーキテクチャの超集合であることが保証される。
• K 経路勾配増幅：本論文は、TIDE がトークン埋め込みに対する累積勾配信号を$K$倍増幅することを形式化している。$K$個の独立した MemoryBlocks のそれぞれが、トークンがバッチ内に出現するたびに勾配を受け取るため、希少トークンはこの並列勾配蓄積から不均衡に恩恵を受ける。
• 文脈的崩壊の解決：著者らは、TIDE が FFN のリプシッツ制約を回避することを証明している。MemoryBlocks は連続的な隠れ状態ではなく、離散的なトークン識別子でインデックス付けされるため、隠れ状態の近接度$\delta$に関係なく、崩壊したトークンペアに対して明確な表現を維持できる。

実証的貢献

• 希少トークンの性能：WikiText-103、PubMed、DCLM における実験は、TIDE が一般的なトークンと比較して、希少トークン（ビン 0-2）の交差エントロピー損失を大幅に低減することを示している。性能向上は MemoryBlocks の数（$K$）と単調に増加し、最も絶対的な改善は最も希少なトークンデシールで観察される。
• 文脈的崩壊の緩和：層ごとの分析は、TIDE がすべてのトランスフォーマー層にわたって、意味的に異なるが文脈的に類似したトークンペア（例：文法上の同音異義語、数値トークン）の隠れ状態間の L2 距離を大幅に増加させ、崩壊を効果的に防止することを示している。
• 下流タスクの改善：TIDE のバリアント（7 億から 30 億パラメータ）は、ARC、BoolQ、HellaSwag、PIQA などのゼロショットベンチマークにおいて、標準的な LLaMA-Base モデルを一貫して上回っている。

4. 結果

• パープレキシティ：TIDE モデルは、ベースラインと比較して、保持されたデータセット（Wikitext-2、PubMed、DCLM）において低いパープレキシティを達成する。例えば、TIDE-24E-1B は Wikitext-2 において 4.60 の PPL を達成し、LLaMA-Base-1B（5.19）を上回っている。
• 勾配ダイナミクス：分析は、TIDE における希少トークンの埋め込みが、標準モデルと比較して高い成長率とノルムを示すことを確認しており、$K$倍の勾配増幅理論を検証している。
• ルーターの動作：学習されたルーターの重みは明確な専門化を示す：Null Bank の重みはトークン頻度とともに増加し（一般的なトークンではメモリを少なく許容する）、特定の MemoryBlocks は希少または中程度の頻度領域に専門化している。
• 圧縮：本論文は、MemoryBlocks が 4 ビット量子化または低ランク分解（SVD）によって圧縮可能であり、パフォーマンスの低下は最小限であることを実証している。これにより、リソース制約のある環境でのアプローチの実用性が確保される。

5. 意義と主張

本論文は、TIDE が LLM 設計において「普遍的に受け入れられているが未検証の」欠陥に対処すると主張している。トークン識別子を文脈残差ストリームから切り離し、すべての層で再注入することにより、TIDE は外部検索メカニズムや事後分析を必要とすることなく、勾配枯渇と文脈的崩壊に対する構造的解決策を提供する。

著者らは、TIDE が単なる検索拡張システムではなく、トークン識別子を永続的で学習可能な信号として扱う根本的なアーキテクチャの転換であると強調している。その意義は、以下の点を示すことにある：

1. 希少トークンは、勾配枯渇により標準的なトランスフォーマーにおいて構造的に不利であり、TIDE は並列メモリ経路を通じてこれを解決する。
2. 文脈的崩壊は、連続的な隠れ状態表現の固有の限界であり、TIDE はすべての層に離散的なトークンインデックス入力を導入することでこれを回避する。

この研究は、トークン識別子でインデックス付けされた専用で文脈に依存しない意味メモリを維持することが、語彙の長尾分布の処理や、同一の文脈における意味的に類似したトークンの識別に特に必要不可欠なコンポーネントであることを示唆している。