Anthropicにおける解釈性のスケーリングに伴うエンジニアリングの課題

Anthropicは、解釈性研究をより大規模なモデルへとスケールさせており、主要なボトルネックは理論的研究から分散システムエンジニアリングへと移行しています。この移行は、小規模なtransformerの実験から、研究者が数千万もの「特徴量」（意味的概念を表すニューロンの組み合わせ）を特定したClaude 3 Sonnetのようなモデルの分析へと進むために不可欠です。

疎な自己符号化器（Sparse Autoencoders）のための分散データシャッフル

Sparse Autoencoders (SAEs) のトレーニングをスケールさせるには、順序に依存する偽のパターンを学習することを防ぐために、transformerの活性化データの膨大なデータセットをシャッフルする能力が必要です。単一のGPUに収まるデータであればシャッフルは容易ですが、データセットがペタバイト規模に達すると、重大なエンジニアリングの課題となります。

シャッフル・パイプラインの進化

Anthropicは当初、データを $K$ 個のジョブに分割し、各ジョブがトレーニングデータセット全体をストリーミング読み込みして、自身の出力シェアを書き込むという単純な分散シャッフルを使用していました。このアプローチは、データセットが100TB（1000億のデータポイント）に成長するにつれて非効率的になり、シャッフルに数日かかるようになりました。

これを解決するために、Anthropicはマルチパス分散シャッフルを実装しました：

First Pass: $N$ 個のジョブがそれぞれデータセットの $1/N$ を読み込み、ローカルでシャッフルし、データを $K$ 個の別々のファイル（それぞれにデータの $1/NK$ が含まれる）に書き込みます。
Subsequent Passes: すべてのジョブの最初のファイルがグループ化され、最終的なシャッフル済みデータの最初の $1/K$ を表します。これらのシャードがまだメモリ制限を超える場合は、プロセスを繰り返します。

この手法は、1パスあたりのシャッフルサイズを100分の1に削減します（メモリ100GB、出力ファイル1GBと仮定）。これにより、チームは4パスで100GBから100PBまでスケールさせることが可能になりました。

特徴量可視化パイプラインのエンジニアリング

特徴量可視化のためのデータ生成（研究者がどのトークンが特定の機能に対して最も強く活性化するかを確認できるようにするもの）には、1億トークンのデータセットにわたる数百万の特徴量の処理が必要です。

分散データ処理ワークフロー

数百万の特徴量を扱うために、Anthropicはシャード化された処理パイプラインを利用しています：

Initial Sharding: システムはデータセットと特徴量の両方に対してシャード化を行います。各ジョブは、自身の特定の機能およびデータのスライスに対して、上位 $K$ 個の活性化トークンと $10K$ 個のランダムな活性化トークンを追跡します。
Aggregation: システムは特徴量に対してシャード化を行い、前のパスの計算結果を集約し、各特徴量におけるグローバルな最高活性化トークンを特定します。
Contextual Activation Analysis: 特徴量が周囲のトークンに対してどのように発火するかを計算するために、チームはデータセット全体へのランダムな読み込みを避けるための2段階のプロセスを実装しました：
- Dataset Sharding Pass: 各ジョブはtransformerの活性化データのスライスを扱い、特定の機能グループに必要な活性化データを別々のファイルに保存します。
- Feature Sharding Pass: ジョブはその後、これらの統合されたファイルにアクセスして、最終的な活性化を計算し、フロントエンドの可視化ツール用のデータをフォーマットします。

AI Safetyにおける研究とエンジニアリングの統合

Anthropicは、研究とエンジニアリングを切り離せない構成要素として捉えています。チームは、より小さな実験で研究仮説が初期の成功を示した後、インフラストラクチャに集中的に投資するという反復的なサイクルに従っています。

チーム構成と戦略

解釈性チームは、神経科学、数学、生物学、物理学、データ可視化、およびソフトウェアエンジニアリングなど、多様な背景を持つ18名のメンバーで構成されています。彼らの研究ロードマップは、AnthropicのResponsible Scaling Policy (RSP) に基づいており、これは、より高い能力を持つモデルをデプロイする前に、特定の安全性に関するマイルストーンを達成することを義務付けています。この探索的な研究は、これらの安全性マイルストーンに直接的な影響を与えるように設計されており、モデルの内部構造を理解する能力がモデルの能力向上とともにスケールさせることを確実にしています。