解釈可能性(Interpretability)の技術的課題
Anthropicの解釈可能性研究は、アルゴリズムだけでなくデータ規模のエンジニアリングが律速の一つになっています。分散シャッフルと特徴可視化パイプラインという2つの実例を読み解きます。
Anthropicの解釈可能性チームは2024年6月13日、「The engineering challenges of scaling interpretability」を公開しました。主張は明快です。解釈可能性研究を前に進める律速の一つは、エンジニアリングにある、というものです。関連する検証はデータ重複とスケーリング則で性能はどう壊れるかでも扱っています。
根拠として挙げられたのが、分散シャッフルと特徴可視化パイプラインの2つの実例です。どちらも新しいアルゴリズムの話ではなく、データが大きくなった途端に既存の素朴なやり方が破綻した、という種類の問題でした。
小さな実験からClaude 3 Sonnetまで、規模はどれだけ違うのか
2023年10月の「Towards Monosemanticity」が扱ったのは、1層だけの小さなトランスフォーマーです。辞書学習(dictionary learning)で得た特徴の数は、最大でも131,072個でした。2024年5月の「Scaling Monosemanticity」は同じ手法をClaude 3 Sonnetに適用し、3種類の辞書を学習しています。
2本の論文で扱った辞書の大きさ
Towards Monosemanticity
最大131,072
1層・512ニューロンのMLP。活性化ベクトル100億のうち80億で学習
Scaling Monosemanticity 小
1,048,576
Claude 3 Sonnetの中間層。死んだ特徴は約2%
Scaling Monosemanticity 中
4,194,304
死んだ特徴は約35%
Scaling Monosemanticity 大
33,554,432
死んだ特徴は約65%
Scaling Monosemanticityでは、モデルの途中(中間層)の残差ストリームに辞書を当てています。残差ストリームはMLP層より小さく、学習と推論の計算コストが抑えられるためです。ここで言う「死んだ特徴」は、1,000万トークンのサンプルで一度も活性化しなかった特徴のことです。辞書が大きいほど死んだ特徴の割合が増えるので、規模の拡大は単純な掛け算で済みません。
34M辞書の学習ステップ数は、スケーリング則の分析で決められています。計算予算を固定したとき、学習損失が最小になる値を選んだということです。実験の設計そのものが、計算資源の見積もりに依存しています。また論文は、十億トークンに1回しか現れない概念を専用の特徴として拾うには、十億個規模の有効な特徴が要ると見積もっています。辞書をいくら大きくしても、まれな概念の取りこぼしは残ります。
ブログの数字とは、見え方が少し違う点があります。ブログは学習データを「100TB、1,000億データポイント(1件1KB)」と書いています。Towards Monosemanticityの論文には、Pileの4,000万コンテキストから集めた活性化ベクトル100億のうち、80億余りで学習したとあります。ブログの1,000億データポイントとは数字が違い、ブログは当時の社内の規模を指している可能性があります。ブログの1件1KBは、512ニューロンを16ビットで持つ場合の計算値(512×2バイト)と一致します。
分散シャッフルで、全員が全データを読む方式をどう避けたか
Sparse autoencoder(疎なオートエンコーダー)の学習には、トランスフォーマーの活性化データを事前にシャッフルしておく必要があります。順序に依存した見せかけのパターンを学習させないためです。学習データが1台のGPUに載っているうちは、シャッフルは些細な作業でした。メモリに収まらない規模になると、トランプ1組のシャッフルが、倉庫いっぱいのカードのシャッフルに変わります。
ブログによれば、ペタバイト級まで伸びる分散シャッフルを最初から作る道もありました。チームが選んだのは、素早く実装できて拡張性は劣る方式です。続いて見直した多段方式との違いを並べます。
シャッフルの2方式
最初の方式
K個のジョブが、出力の1/Kずつを担当します。順列を1つ作り、各ジョブが学習データ全体をストリーミングで読み、自分の担当分だけを書き出します。
読み込み量はジョブ数に比例して増えます。所要時間は数時間から数日に伸び、100TBに達したTowards Monosemanticityの頃には大きな悩みの種でした。
多段方式
あるブログ記事を参考に、シャッフルを複数のパスに分けます。1パスでN個のジョブが動き、各ジョブはデータの1/Nを読み、シャッフルして、K個のファイルに1/NKずつ書き出します。
各ジョブの1番目のファイルを集めると最終出力の最初の1/Kになりますが、その中身はまだ混ざっていません。1回のシャッフルを、K分の1の大きさのN回に変えたことになります。メモリに収まる大きさになるまでパスを重ねます。
ブログは、1ジョブが100GBをメモリに持ち、100個の1GBファイルを書く場合を例に挙げています。1パスで必要なシャッフルの大きさは100分の1になるので、扱える規模はパスごとに100倍になります。
| パス数 | シャッフル可能なデータ規模 |
|---|---|
| 1パス | シャッフル可能なデータ規模100GB |
| 2パス | シャッフル可能なデータ規模10TB |
| 3パス | シャッフル可能なデータ規模1PB |
| 4パス | シャッフル可能なデータ規模100PB |
この表の条件を100TBのデータに当てはめると、必要なのは3パスです。2パスで10TB、3パスで1PBなので、100TBは3パス目で収まります。読み込み量は、パスごとにデータセットを1回読む計算で約300TBです。最初の方式でKを100とした場合は、各ジョブが100TBを読むので合計で約10PBになります。K=100は試算のための仮定で、ブログが示した値ではありません。
この方式の導入後、シャッフルは検討事項から外れました。ブログは「より良い方式も速い実装もあるはずだが、律速だった問題は解けた」という趣旨で書いています。分散シャッフル自体は目新しい問題ではないものの、科学的な進捗のために素早く片づけるべき問題の1つだった、とブログは書いています。Scaling Monosemanticityの謝辞にも、活性化に対するスケーラブルなシャッフルの実装が貢献として挙がっています。
特徴可視化は、なぜデータ側と特徴側の両方でシャードし直すのか
特徴可視化(feature visualization)は、個々の特徴がどのトークンで強く活性化するかを、ユーザーが画面で確認できるようにする仕組みです。Towards MonosemanticityのFeature Browserがその実例です。特徴ごとに、活性化が強いものから弱いものまで多様なデータ例を集める必要があり、数百万個の特徴では分散処理の問題になります。当初は1つのジョブで動いていましたが、すぐに収まらなくなりました。
可視化用のデータセットは1億トークンです。トランスフォーマーの活性化はS3にキャッシュ済みなので、再計算はしません。現在の処理は次の4段で組まれています。
特徴可視化データの生成
- 1
データと特徴の両方を分割して走査する
各ジョブが自分の担当スライスを走査し、担当する特徴ごとに、最も強く活性化するK個のトークンと、ランダムに選んだ10×K個のトークンを記録します。
- 2
特徴側で分け直して集約する
前段の結果を特徴ごとに集め、データセット全体での最強活性化トークンとランダムサンプルを確定します。特徴1つにつき保持する例は、合計11K件です。
- 3
周辺トークンでの発火を調べる(最初の方式は失敗)
各例について、周辺のトークンで特徴がどう発火するかを計算します。最初は特徴ごとにシャードして必要な活性化を読み込みましたが、対象の例はデータセット全体に散らばっています。必要な部分だけを読み出す方法がありませんでした。
- 4
データ側のパスを足してから、特徴側で仕上げる
データセット側でシャードするパスを追加し、各ジョブが活性化の一部を読んで、特徴グループごとに別ファイルへ保存します。その後に特徴側のパスを回せば、必要なデータだけを手軽に読めます。最後に、フロントエンドが読める形式で書き出します。
シャッフルと可視化は題材こそ違いますが、解き方は共通しています。全員が全データを読む形をやめ、軸を変えながら分割とまとめ直しを重ねる形に置き換えました。どちらも、次の軸のパスが読む範囲を狭めるための準備です。
研究とエンジニアリングは、どう行き来しているのか
Sparse autoencoderの研究を始めた時点では、手法がうまくいくかどうかは分かっていませんでした。実験が成功するたびにインフラへの投資を増やし、それがScaling Monosemanticityまで続いたという経緯です。多くの研究アイデアは実を結ばないため、成功の兆しが見えるまでは、インフラに過剰投資しない方針でした。
メンバーは研究とエンジニアリングを頻繁に行き来します。現行システムから規模を絞り出して新しい実験を走らせ、終わればまた研究に戻る、という繰り返しです。ここで言う律速は、アルゴリズムの優劣だけでなく、実験を回せる規模の上限にもあります。
求められるのは特定分野の専門家ではなくジェネラリストなのか
採用面では、パイプライン構築・機械学習実験の実行・GPU利用の最適化といった領域を柔軟に横断できる「ジェネラリストのエンジニア」が重視されています。分散シャッフルも特徴可視化も、特定のアルゴリズムの知識より、分散システムを自分で組み立てる力が効く問題でした。
研究の方向性は、AnthropicのResponsible Scaling Policy(責任あるスケーリングポリシー)が定める安全マイルストーンとの関係で決まります。日々の探索は自由度が高い一方で、解釈可能性研究の全体像のどこにギャップがあり、チームがどの問題に取り組める立場にあるかを見て、次の題材が選ばれます。Alignmentチームをはじめ、他の研究チームとも密に連携しています。
拠点はサンフランシスコ・ボストン・ニューヨーク・シアトル・ロンドンで、人数が最も多いのはサンフランシスコです。そこでは週に数日の出社があります。リモート勤務も受け入れており、その場合はAnthropicのオフィスを年間約25%訪れる条件が付きます。
読み終えて何を持ち帰るか
この記事は採用広報として書かれています。結びでは、2つの実例を、チームのエンジニアがいま取り組んでいる問題の例として示し、応募を直接呼びかけています。それでも、100TBの活性化をどう混ぜ、数百万の特徴の例をどう集めるかという具体的な設計は、解釈可能性の成果がどんな基盤の上に載っているかを知る手がかりになります。研究の発表を追うときは、結果の数字だけでなく、その実験を回せる規模に押し上げたインフラにも目を向けると、論文の見え方が変わります。
安全性研究が製品にどう関わるかは、Anthropicの安全性研究がClaude製品のどこに入っているかにあります。同じ分野の別の探索は、「J-space」を探った解釈可能性研究です。