ロジットレンズは、機械学習の解釈可能性における技法であり、Transformer (architecture)ベースの大規模言語モデルの内部表現を検査するために用いられる。これは、特定の層における隠れ状態を取得し、モデルの最終的なアンエンベディング行列（隠れ状態を語彙上のロジットに写像する）を適用し、その後ソフトマックスを適用して確率分布を得るという仕組みである。この分布は、その層でモデルが停止した場合に何を予測するかを明らかにし、予測が層ごとにどのように進化するかを示す窓を提供する。この手法は2020年にnostalgebraistによって導入され、以来、GPT-2やGPT-3のようなモデルの内部計算を研究するための標準的なツールとなっている。
プロービング分類器が補助モデルの訓練を必要とするのとは異なり、ロジットレンズはモデル自身の出力ヘッドを使用するため、ゼロパラメータで訓練不要のアプローチである。これは、各層の隠れ状態が加算的であり直接射影できるため、残差ストリームを持つモデルに特に効果的である。この技法は、モデルが特定の予測にコミットするタイミングを特定したり、「遅い」デコードと「早い」デコードを検出したり、事実の想起や算術推論などの現象を分析するために使用されてきた。しかし、その適用可能性は、タイドまたは共有エンベディングを持つモデルに限定され、複雑な層正規化や非線形変換を持つモデルではうまく機能しない場合がある。
メカニズムと実装
ロジットレンズは、トランスフォーマーの残差ストリーム上で動作する。典型的なトランスフォーマーでは、各層がその出力を残差ストリームに加算するため、層 \( l \) における隠れ状態は、初期エンベディングと層 \( l \) までのすべての層出力の合計となる。最終ロジットは、層正規化（存在する場合）を適用し、その後アンエンベディング行列 \( W_U \) を適用して計算される。ロジットレンズは、分布シフトを考慮するために、最終層ノルム（または学習されたアフィン変換）を適用した後、中間隠れ状態に同じアンエンベディングを適用する。
タイドエンベディング（入力エンベディングと出力アンエンベディングが重みを共有する）を持つモデルでは、射影は簡単である。別個のアンエンベディング行列を持つモデルでも、アンエンベディングが利用可能であればこの技法は機能する。結果は、各層の語彙上の確率分布の系列であり、ヒートマップとして可視化したり、候補トークン間の「ロジット差」などの指標を計算するために使用できる。
解釈可能性における応用
ロジットレンズは、さまざまな解釈可能性タスクに適用されてきた。一般的な用途の1つは、層をまたぐ予測の発展を追跡することである。例えば、「フランスの首都は__」という入力を完了するように求められたモデルでは、ロジットレンズは、初期の層では「パリ」に高い確率を割り当てず、いくつかの層を経て初めて高い確率を割り当てることを示すかもしれない。これにより、モデルが関連する事実をいつ取得するかが明らかになる。研究者はこれを用いて、特定の行動に寄与する「誘導ヘッド」やその他の注意パターンを特定してきた。
もう1つの応用は、「遅い」デコードの検出であり、モデルが最初にあるトークンを予測しても、後で考えを変える場合がある。各層でロジットレンズを調べることで、最終予測がいつ「固定」され、いつまだ柔軟であるかを見ることができる。これは、モデルの信頼度を理解し、介入を設計するための示唆を与える。
この技法は、算術と推論の研究にも使用されてきた。GPT-3のようなモデルでは、ロジットレンズは中間層が中間の和や繰り上がりを表現していることを示すことができ、多段階計算の証拠を提供する。これは、Anthropicや他の研究グループによる「誘導ヘッド」や「回路」フレームワークなど、機構的解釈可能性に関するさらなる研究につながった。
限界と批判
その有用性にもかかわらず、ロジットレンズにはいくつかの限界がある。第一に、アンエンベディング行列がすべての層にとって意味のある射影であると仮定しているが、モデルの隠れ状態が（層正規化など）考慮されていない重要な変換を受ける場合、これは成立しないかもしれない。特に事前正規化アーキテクチャを持つモデルでは、中間状態に最終層ノルムを適用する必要があるが、これで常に十分とは限らない。
第二に、ロジットレンズは、語彙が大きいモデルや複数トークンの予測を必要とするタスクでは効果が低く、次トークン分布のみを示す。また、注意やフィードフォワードサブ層の寄与をその加算効果以外では無視するため、モデルの完全な状態を捉えることはできない。
第三に、モデルが残差ストリームと整合していない別個の出力ヘッドを使用する場合、この技法は誤解を招く結果を生む可能性がある。そのような場合、射影されたロジットはノイズが多く、情報量が少ないかもしれない。研究者は、「チューニングレンズ」と呼ばれる、層ごとのアフィン変換を学習して整合性を改善する変種を提案しているが、これには追加の訓練データが必要である。
他の解釈可能性手法との関係
ロジットレンズは、プロービング分類器、活性化パッチング、因果トレーシングを含む、より広範な解釈可能性技法のファミリーの一部である。プロービング分類器は、特定の特徴を予測するために隠れ状態に線形または非線形モデルを訓練するが、教師あり学習を必要とし、モデルの実際の計算を反映しない場合がある。対照的に、ロジットレンズはモデル自身の出力ヘッドを使用するため、より直接的である。
活性化パッチングは、隠れ状態に介入して出力への影響を調べるものであり、ロジットレンズからの発見を検証するために使用できる。David Kaplanらによって広められた因果トレーシングは、破損した実行を使用して特定の事実に責任がある層を特定し、ロジットレンズは各層の予測トークンを示すことでこれを補完できる。
もう1つの関連手法は「残差ストリーム」ビューであり、トランスフォーマーを一連の加算的更新として扱う。ロジットレンズは、各点で残差ストリームを射影するため、このビューに自然に適合する。これは、「Transformer Debugger」や他の可視化ライブラリなどのツールに影響を与えた。
拡張と変種
ロジットレンズのいくつかの拡張が提案されている。「チューニングレンズ」は、層ごとのアフィン変換を学習して隠れ状態をアンエンベディングとより良く整合させ、標準的なロジットレンズが失敗するモデルでの性能を向上させる。もう1つの変種は「ソフトマックスレンズ」であり、ロジットに温度を適用して分布を鋭くまたは平坦にし、低確率の予測を見やすくする。
複数の出力ヘッドやデコーダーを持つモデルでは、ロジットレンズを各ヘッドに個別に適用できる。エンコーダー-デコーダーモデルでは、この技法をデコーダーの隠れ状態に適用できるが、エンコーダーに適用する方法はあまり明確ではない。
最近の研究では、画像トークンや他のモダリティを含む語彙を持つマルチモーダルモデルの分析にロジットレンズを使用することも模索されている。しかし、これはまだ活発な研究分野である。
計算上の考慮事項
ロジットレンズは計算効率が高く、各層で順伝播と行列乗算のみを必要とする。層数 \( L \) と語彙サイズ \( V \) のモデルでは、追加コストは \( O(L \cdot V \cdot d) \) であり、ここで \( d \) は隠れ次元である。これは順伝播自体のコストに比べて無視できる。
しかし、すべての層のロジットを保存することは、特に大きな語彙ではメモリ集約的になる可能性がある。実際には、各層で上位kトークンを計算してメモリ使用量を削減できる。この技法は、TransformerLensライブラリやHugging Faceの解釈可能性ツールを含む、いくつかのオープンソースライブラリに実装されている。
今後の方向性
深層学習モデルが規模を拡大し続けるにつれて、ロジットレンズのような解釈可能性ツールはますます重要になる。将来の研究は、Mixture of expertsや状態空間モデルなどの新しいアーキテクチャにこの技法を適応させ、自動化された解釈可能性パイプラインと統合することに焦点を当てるかもしれない。ロジットレンズはまた、モデルが有害な出力について「考えている」ことを検出するための安全性研究でも使用されており、これは整合戦略に情報を提供する可能性がある。
全体として、ロジットレンズは、ニューラルネットワークのブラックボックスを覗き込むためのシンプルでありながら強力なツールであり続けており、その応用は機構的解釈可能性の分野が成熟するにつれて拡大する可能性が高い。