2026-06-28

Which tokens does a hybrid model predict better?

概要 (Overview):

本記事は、同一条件で学習されたTransformer(Olmo 3)とハイブリッドモデル(Olmo Hybrid)の予測性能をトークンレベルで比較・分析したものです。 ハイブリッドは状態追跡を必要とする「内容語」の予測に優れる一方、Transformerは過去の入力をそのまま複製する「コピー」タスクで圧倒的な強みを持つことが明らかになりました。 この結果から、単一の総合損失ではなく、特定のトークンタイプごとに損失を評価する「フィルタリングされた損失」の重要性を提唱しています。

原文ソースを表示 ↗

日本語訳

モデルはどのようなトークンを精度よく予測し、どのようなトークンを苦手とするのでしょうか。この問いは、標準的なTransformerに挑戦し始めている言語モデルアーキテクチャである「ハイブリッドモデル」において、特に興味深いテーマです。
ハイブリッドモデルは標準的なベンチマークにおいてTransformerと同等以上の性能を達成できますが、主要な評価指標の数値だけでは、ハイブリッドモデルがTransformerに対して具体的にどのような優位性を持っているのかについては多くを明らかにしません。
これらのトークンレベルでの挙動を解明するため、私たちは最近、自社で最も強力な7B TransformerモデルであるOlmo 3と、ハイブリッドモデルであるOlmo Hybridを直接比較する実験を行いました。具体的には、LLMへの入力として現れる情報の単位であるトークンの異なるタイプにわたって、モデルの予測の差異を詳細に比較しました。
Olmo 3とOlmo Hybridは、アーキテクチャ以外の部分(データ、トークナイザー、学習レシピなど)ができる限り共通化されて構築されているため、予測におけるあらゆる差異は主にアーキテクチャ自体の違いを反映しています。これらの差異をトークンレベルで観察することで、Transformerに対するハイブリッドモデルの具体的な強みについて洞察を得ることができます。
実験結果は、多くのトークンにおいてハイブリッドの優位性が実在することを示していますが、すべてのトークンでそうであるわけではありません。Olmo Hybridは、名詞、動詞、形容詞などの意味を持つトークンや、代名詞がどの人物を指しているかといった文脈の推移を追うことでしか予測できないトークンにおいて最も強みを発揮します。しかし、入力内に既に存在する内容を単に繰り返すトークン(以前登場した単語やフレーズの逐語的な再現)では、ハイブリッドの優位性はほぼ消失します。そこでは、答えが参照元としてそこに存在しており、これこそがTransformerの真骨頂です。
言語モデルは反復するレイヤーのスタックから構築されており、各レイヤーは周囲のトークンを用いて各トークンの表現を洗練させます。Transformerはすべてのレイヤーでアテンションを使用します。モデルは過去のすべてのトークンに一度に直接アクセスし、現在の予測に対する各トークンの関連性を重み付けできます。これにより、アテンションは入力のはるか手前に現れた特定のトークンであっても、正確に想起することを得意とします。欠点は、すべてのトークンが過去のすべてのトークンと比較されるため、入力が長くなるにつれて計算コストが急激に上昇することです。さらに、アテンションは情報の想起や集約には強いものの、時間とともに逐次的に変化する情報を表現することは苦手としています。
ハイブリッドモデルは、少数のアテンションレイヤーを維持しつつ、残りをリカレントレイヤー(再帰レイヤー)に置き換えます。アテンションレイヤーとは異なり、リカレントレイヤーはトークンを左から右へと読み込み、固定サイズのメモリを保持します。新しいトークンを処理するたびにそれをメモリに畳み込んでいくため、入力がどれほど長くなってもトークンあたりの処理コストは一定のままです。そのメモリは圧縮されており不可逆(ロスあり)であるため、リカレントレイヤーはアテンションのように過去の正確なトークンを直接手繰り寄せることはできません。しかし、モデルがトークンを読み進めるにつれて変化する情報を継続的に追跡することに適しており、アテンションを補完する強みを提供します。
アテンションレイヤーとリカレントレイヤーのそれぞれの強みと弱みを特定するため、Olmo 3とOlmo Hybridに、記事、Wikipediaのエントリ、書籍、学術論文、さらにはPython、HTML、LaTeXなどの構造化テキストを含む文章を与えました。与えられたサンプル内の先行するトークンから、各モデルが各トークンをどれだけ正確に予測できたかをスコア化しました。両モデルは同じ先行トークンを参照し、次に続く可能性のあるすべてのトークンに対して確率を割り当てました。そして、実際に後に続いたトークンに対して各モデルが割り当てた確率を記録しました。その後、2つのモデル間の損失の差である「損失ギャップ(loss gap)」を計算することにより、トークンごとにモデル間の差異をまとめました。ギャップが正の場合はハイブリッドの方が実際の次のトークンをより良く予測できたことを意味し、負の場合はTransformerの方が優れていたことを意味します。
損失ギャップがどこに集中しているかを突き止めるため、いくつかの分析を行いました。まず、各トークンをカテゴリに分類し、これらのカテゴリ内での平均損失ギャップを算出しました。単純な平均値は、カテゴリの希少性やテキストサンプル内でのトークンの重複頻度といった他の要因によって偏る可能性があるため、他の要因を一定に保ちながらカテゴリ自体の影響を推定する回帰分析を用いて、各パターンを再検証しました。その結果、Olmo Hybridはほとんどの種類のトークンにおいてOlmo 3よりも損失が低いことが分かりましたが、その差はトークンの種類によって異なります。
散文において最も明確な境界線は、意味を担う「内容語」(名詞、動詞、形容詞など)と、「the」「of」「is」のような「機能語」の間にあります。ハイブリッドモデルはTransformerよりも内容語の予測に優れている一方、機能語における差はより小さくなります。要するに、ハイブリッドの優位性は文が何について述べているかを示す単語で最も大きく、どのモデルでも構文からほぼ推測できる文法的な単語で最も小さくなります。
対照的に、Transformerに対するハイブリッドモデルの優位性が消失する特定の文脈も存在します。1つ目は、開始カッコではなく「閉じカッコ」の予測であり、このパターンは言語、コード、マークアップのすべてのカッコにおいて堅牢に見られました。その理由は、アテンションがカッコの対応関係を表現するのに十分であることが知られており、閉じカッコの予測にはアテンションだけで十分であることを示唆しているからです。ハイブリッドの優位性がほぼ完全に消失する2つ目のケースは、次のトークンが単に文章中に既に存在する内容を繰り返している場合です。私たちは、これらのケースを重複するn-gramを探すことで特定しました。重複する区間が長くなるほどハイブリッドのリードは小さくなり、最終的にはほぼゼロに近づきます。
最後に、これらの知見に着想を得て、事前学習実験において異なるアーキテクチャをより適切に比較するための評価手法として、特定のタイプのトークンに対する「フィルタリングされた損失(filtered losses)」の利用を検討しました。私たちは、先行する実験から得られた3つの1Bパラメータモデル(Transformer、ハイブリッド、アテンションを一切持たない純粋なリカレントモデル)を使用しました。重複のない意味を担うトークンにおいては、ハイブリッドと純粋なリカレントモデルがTransformerを凌駕し、ハイブリッドが最も優れた性能を示しました。一方で、重複するトークンにおいては、コピー元を手繰り寄せるアテンションを持たない純粋なリカレントモデルが、ハイブリッドとTransformerの両方に後れを取る結果となりました。
本研究からは2つの教訓が得られます。第1に、すべてのトークンにわたるモデルの平均誤差である「単一の総合的な損失」は、Transformerとハイブリッドアーキテクチャを比較するにはあまりにも大雑把すぎます。特定のモデル能力をテストするトークンのみに絞って損失をスコア化することで、重要な差異が浮き彫りになります。第2に、特にハイブリッドモデルに関して、オープンクラス(開かれた語類)のトークンにおいて特定の優位性がある証拠が見つかりました。これはおそらく、RNNレイヤーの状態追跡能力に関連していると考えられます。
素晴らしい研究ですね。AllenAIによるハイブリッドアーキテクチャの解釈や改善には非常に期待しています。1点質問させてください。論文中における唯一の純粋なデータ一致比較は、6つのフィルタリング損失評価に使用されている1Bラダー(Transformer / ハイブリッド / 純粋RNN、同一のデータミックス)ですが、これは総合的な損失としてのみ扱われており、品詞(POS)やカッコ、コピーといった分解は行われていません。公開されたチェックポイント上でのフォワードパスのみであることを考えると、これらのモデルに対して同様のタグ層別化分析を実行した、あるいは実行することは可能でしょうか?データが実際に一定に保たれている場合に、内容語/開始・閉じカッコ/コピーといった構造が維持されるかどうかを示すのに役立つと思います(~7Bのケースと比較して)。内部的にこの点を確認されているかどうかも気になります。