🔬

機械論的解釈可能性の紹介

2024/03/29に公開

はじめに

近年AIモデルの性能が急速に向上し、AGI(汎用人工知能)の登場が近いのではないかと議論される中、AIシステムを人間の価値観に合わせる研究分野「AIアライメント」(AI Alignment) がOpenAIやAnthropicといったAIスタートアップを中心に推進されています。「価値観のアライメント」と聞くと、RLHF (Reinforcement Learning from Human Feedback) などを思い浮かべる人が多いかと思われますが、RLHFは基本的にAIモデルをブラックボックスとして扱うため、根本的なアライメントの達成は難しいと考えられます。そこで、深層学習モデルのブラックボックスを明らかにしようとする「機械論的解釈可能性」[1] (Mechanistic Interpretability) の研究が注目されています。
機械論的解釈可能性の目標は、「AIモデルをリバースエンジニアリングすること」と説明されます[2]。言い換えれば、AIモデルを人間が理解できるコンピュータープログラムに変換することが理想です。
この記事では、本分野のキーコンセプトや、主な研究の方向性[3]、代表的な研究を簡単に紹介します。

キーコンセプト

1. 特徴量 (Feature)

  • ニューラルネットワークへの入力の一部。よく「解釈可能な特徴量」の意味で用いられる。
  • 表現空間において、方向によって表現される。
  • 回路, 普遍性と合わせてニューラルネットワークの基本的な性質。

2. 回路 (Circuit)

  • 特徴量を別の特徴量に変換するサブネットワーク。
  • 実用的には、あるタスクを遂行するのに必要十分なサブネットワークという意味で用いられる。
    「窓」、「車のボディ」、「タイヤ」といった特徴量を「車」特徴量に変換する回路("Zoom in: An Introduction to Circuits")

3. 普遍性 (Universality)

  • 同じような回路が異なるタスクやモデルを通じて形成される現象。
  • このように普遍的に見られる回路のパターンをモチーフ (motif) と呼ぶ。

4. 重ね合わせ (Superposition)

  • ニューラルネットワークが表現空間において、その次元数よりも多くの特徴量を表現する現象。
  • 一つのニューロンが複数の特徴量に対して発火することになり、解釈が難しくなる。
  • 一つの特徴量に対応するニューロンを一義的ニューロン (Monosemantic Neuron) 、複数の特徴量に対応するニューロンを多義的ニューロン (Polysemantic Neuron) と呼ぶ。
    特徴量のスパース性が上がると、特徴量が重ね合わされるようになる("Toy Models of Superposition")

主な研究の方向性

1. トイ言語モデルの分析

  • 大規模言語モデルの基本メカニズムを明らかにするために、数層のTransformerブロックからなる言語モデルのリバースエンジニアリングを目指す。
  • 関連研究
    • A Mathematical Framework for Transformer Circuits (Anthropic, 2021)
      • 一層/二層Transformerモデルのリバースエンジニアリング
    • In-context Learning and Induction Heads (Anthropic, 2022)
      • 上の研究で発見された帰納ヘッド (induction heads: [A][B]...[A]という文脈が与えられたときに次に[B]を出力する確率を高めるような二つの注意ヘッドからなる回路)が文脈内学習 (in-context learning) に大きく貢献することを明らかにした
        帰納ヘッドの振る舞い

2. 大規模言語モデル内の回路同定

  • 我々が普段利用しているような大規模言語モデルに含まれる回路の同定を目的とする。現時点では、最新モデルを解釈しようにも、パラメータが公開されていなかったり、規模があまりに大きかったりするため、大規模言語モデルの中でもGPT-2 smallのような小さめのモデルを対象にすることが多い。
  • 回路同定には、アブレーション手法[4]が使われる。AIシステムのある要素が、特定のタスクを遂行するのに不可欠な寄与をしている(= 回路の一部である)かを、その要素を欠落させたり改変したりして判断する。
  • 関連研究

3. 小規模モデルに実装されたアルゴリズムの解釈

  • 小規模モデルをなんらかのタスクで訓練し、モデルが実装したアルゴリズムをリバースエンジニアリングする。大規模モデル解釈の練習にもなる。
  • 関連研究
    • A Mechanistic Interpretability Analysis of Grokking (Nanda et al., 2023)
      • 一層のTransformerをmod加算タスクで訓練し、実装されたアルゴリズムを完全に解明する
      • 解明された回路へのアブレーションを行い、グロッキング (Grokking: モデルの汎化性能が突如急激に向上する現象) の進行を記憶・内挿・汎化・浄化・安定の5つのフェーズで説明
        Transformerに実装されたアルゴリズム。入力を離散フーリエ変換し、三角関数の積和を経てmod加算を行っており、これは著者も当初予想していなかった。

4. 多義性や重ね合わせの研究

  • 重ね合わせ現象はモデルの解釈を難しくする。重ね合わせ現象が生じるメカニズムの理解や、訓練時の重ね合わせを抑制する技術、訓練済みモデル内で生じた重ね合わせを解釈のため事後的に解消する技術が提案されている。
  • 関連研究
    • Toy models of Superposition (Anthropic, 2022)
      • 重ね合わせ現象を、一層の隠れ層を持つReLUネットワークと人工の入力で網羅的に調査
      • 入力に含まれる特徴量を操作して、重ね合わせ現象の相転移や、重ね合わせがさまざまな幾何学構造をとることを明らかにした[5]
        特徴量のスパース性を高めていくと、重ね合わされる特徴量がなす多面体の面の数が増えていく(一見減っているようにみえるところは多面体の直積を取ることで矛盾なく説明できる)
    • Softmax Linear Units (Anthropic, 2022)
      • 重ね合わせを抑制する、Softmaxを用いた新しい活性化関数を提案
        提案された活性化関数の式
    • Towards Monosemanticity: Decomposing Language Models With Dictionary Learning (Anthropic, 2023)
      • 言語モデルに対して、スパースオートエンコーダを用いて重ね合わせを展開し、人間に理解可能な特徴量を提示する手法を提案
        512のニューロンからなる層が捉えている、4096の特徴量を可視化

5. 訓練ダイナミクスの研究

  • AIモデルでの創発現象の理解に役立つ。
  • 既に上で挙げた研究でも訓練ダイナミクスが調べられているものが多い(Grokking, Superpostion, and etc)。
    6つの人工的な特徴量(相関した3つの特徴量*2組)が、初めに組内で分かれて、次にねじれて組間の干渉を減少させている("Toy Models of Superposition")
    mod加算の回路(汎化回路)以外を除いて計算したロスをTrig lossとして学習曲線を可視化すると、汎化回路の形成とgrokkingのタイミングが一致していることがわかる("Progress measures for grokking via mechanistic interpretability")

6. 解釈ツールの開発、自動化

  • 解釈ツールの定式化、自動化を行うことで、AI解釈のスピードアップを目指す。
  • 関連研究
    • Causal Scrubbing: a method for rigorously testing interpretability hypotheses (Redwood Research, 2022)
      • ニューラルネットワークの内部構造についての仮説を厳密にテストする手法を定式化
        解釈対象のニューラルネットワークG、機械論的解釈仮説を表すグラフI、その間の対応cが与えられたとき、Iのノードに対応するGのノードを出力から順に見ていき、Iから矢印が伸びている親ノードの活性値を同じ活性値をもたらすはずの別の入力によるものに変更し、Iから矢印が伸びていない親ノードをランダムな入力によるものに変更していって、ロスが大きく増えるかを見る
    • Towards Automated Circuit Discovery for Mechanistic Interpretability (Conmy et al., 2023)
      • ニューラルネットワークから自動で回路を見つける手法を提案
        出力側から一つずつモデルのエッジを試しに切っていき、タスク性能が変わるかでエッジが求めたい回路に含まれるかを判定
    • The Building Blocks of Interpretability (Google Brain, 2018)
      • 視覚モデルの解釈のための視覚的UIについてまとめている
      • Feature visualization (どのような特徴量を見ているか), Attribution (それがどれだけ出力に影響するか), Neuron grouping (人間が理解しやすいようなニューロンのグルーピング) を組み合わせている

7. 視覚モデルの解釈

  • Transformerの導入によって言語モデルの性能が急激に向上する前から、視覚モデルの解釈研究が行われてきた。視覚モデルには、異なるアーキテクチャやタスクに共通してみられるモチーフがいくつも確認されている。
  • 関連研究
    • Feature Visualization (Olah et al., 2017)
      • 視覚モデル内の特定のニューロンやレイヤーが反応する特徴量の可視化手法についてのサーベイ(リバースエンジニアリングをしているわけではないので正確には機械論的解釈ではない)
      • データセットを用いて画像の集合から特徴量を可視化する手法の限界と、モデルの特定の要素の活動を目的関数として入力空間まで誤差逆伝播する手法[6]のバリエーションと課題をまとめている
        深層であればあるほど複雑な特徴量を捉えている
    • Curve Detectors (OpenAI, 2021)
      • 視覚モデルに現れる曲線検出器についての網羅的な調査
      • データサンプルの可視化、特徴量可視化手法、人工カーブを用いた分析、等によって方向選択性をもつ曲線検出器を明らかにしている
        曲線を検出するニューロンの、特徴量可視化(左)と、人工カーブ(上)に対するアクティベーションを可視化
    • Curve Circuits (OpenAI, 2021)
      • 上の研究の知見から、視覚モデル内の曲線検出器をリバースエンジニアリング
    • High-Low Frequency Detectors (OpenAI, 2021)
      • 視覚モデルがどのように異なる周波数成分を処理しているかを調査
    • Multimodal neurons in artificial neural networks (OpenAI, 2021)
      • CLIPにおいて、文字・記号・概念といった異なる提示方法にかかわらず、同じ概念に反応するニューロンを発見

8. 強化学習モデルの解釈

  • AIアライメントの文脈で懸念されるような汎用人工知能 (AGI) はおそらく強化学習エージェントであるため、強化学習モデルの解釈は不可欠である。
  • 関連研究

9. 学習される特徴量の研究

  • AIモデルによって学習される特徴量には、「車」のような想像が容易なものもあれば、「人数」といった予想が難しいものもある。適切な解釈仮説のためには、AIシステムが学習する特徴量の理解を深める必要がある。
  • 関連研究
    • Neuroscope (Neel Nanda, 2022)
      • いくつかの大規模言語モデル内の各ニューロンが特に反応する入力サンプルを一覧できるツール

さいごに

機械論的解釈可能性は比較的新しい分野ですが、Neel Nandaをはじめとする研究者らの多大なる貢献もあり、多くの専門家や学生が参入し始めているようです。本記事執筆時点では、"Mechanistic interpretability"とGoogle検索をしても日本語の情報はほとんどヒットしませんが、本記事を見てこの分野に興味を持たれたら、ぜひ原論文や以下に挙げる補足資料をチェックしていただければと思います。今後、日本で機械論的解釈可能性やAIアライメントに興味を持つ人が増えれば嬉しいです。

補足資料

脚注
  1. 比較的新しい分野なので、和訳がまだ確定していないようです。ここでは、ほぼ直訳を採用します。また、本記事に登場する用語の多くも和訳が確定してるわけではないことに注意してください。 ↩︎

  2. この記事では、本研究分野の創始者とされているChris Olahが、AIの機械論的解釈をバイナリコードのリバースエンジニアリングとのアナロジーで説明しています。 ↩︎

  3. こちらはNeel Nandaの記事を参考にしています。 ↩︎

  4. zero ablation, mean ablation, interchange ablationといったバリエーションがありますが、執筆時点(2024年3月)ではinterchange ablationが回路にとって訓練時の状況と一番近いため最も適切であるとされる傾向にあるようです。 ↩︎

  5. 元論文では、重ね合わせ現象とトムソン問題との類似性が指摘されています。 ↩︎

  6. 悪夢のような画像を生成すると話題になったDeep Dreamにも適用されている技術です。 ↩︎

Discussion