CLT-Forge: A Scalable Library for Cross-Layer Transcoders and Attribution Graphs

以下は、論文「CLT-Forge: A Scalable Library for Cross-Layer Transcoders and Attribution Graphs」の技術的な詳細な要約です。

1. 背景と課題 (Problem)

機械的解釈性（Mechanistic Interpretability）の分野では、大規模言語モデル（LLM）がどのように情報を表現・処理するかを理解することが目標です。近年、辞書学習（Dictionary Learning）とトランスコーダー（Transcoders）を用いたアプローチが注目されており、モデルの計算を「疎で解釈可能な特徴量」とその相互作用として表現し、「特徴量アトリビューショングラフ（Feature Attribution Graph）」を生成できます。

しかし、既存のアプローチには以下の重大な課題がありました：

• グラフの肥大化と冗長性: 従来のトランスコーダーでは、類似した特徴量が複数の層に重複して出現しやすく、アトリビューショングラフが数百ノードに膨れ上がり、実用的な分析が困難です。
• クロスレイヤートランスコーダー（CLT）の限界: 層を超えて特徴量を共有し、冗長性を減らす「クロスレイヤートランスコーダー（CLT）」はよりコンパクトで解釈しやすいグラフを提供しますが、パラメータ数が膨大になるため、大規模なスケールでのトレーニングや分析が困難でした。
• ツールの断絶: 業界（Anthropic など）のフレームワークは非公開であり、オープンソースの既存ライブラリは、効率的な分散トレーニング、メモリ効率化のための活性化値（Activation）の圧縮キャッシュ、自動化された解釈性パイプライン、可視化インターフェースを統合的に提供していませんでした。

2. 提案手法：CLT-Forge (Methodology)

著者らは、CLT のエンドツーエンドのトレーニングと解釈性分析を可能にする、オープンソースの統合ライブラリ「CLT-Forge」を提案しました。このフレームワークは以下の主要コンポーネントで構成されています。

2.1 スケーラブルなトレーニング基盤

• 活性化値のキャッシュと圧縮: 大規模データセット（数億トークン）からの MLP 入出力活性化値を事前に計算・保存します。ストレージ要件を削減するため、層ごとの対称量子化（int8, int4, int2）と zstd エントロピー符号化を適用し、保存容量を 20TB から 4TB 程度（LLaMA 1B モデルの場合）に削減しています。これにより、トレーニング品質への影響は 2-3% 程度に抑えられています。
• GPU シャーディング: 単一 GPU では収まらない大規模 CLT に対応するため、特徴量次元（Feature Dimension）にわたるシャーディング（Feature-wise Sharding）を実装しました。各 GPU が特徴量のサブセットを担当し、フォワードパスの最後に出力を集約します。これにより、DDP や FSDP に比べてメモリ効率が大幅に向上します。
• トレーニング目標: MLP の出力を再構成する損失関数に加え、特徴量の疎性（Sparsity）と「死んだ特徴量（Dead features）」の抑制を正則化項として組み込み、JumpReLU 活性化関数を使用します。

2.2 自動化された解釈性パイプライン (AutoInterp)

• トレーニング済みモデルの特徴量を解釈するため、各特徴量を最も強く活性化させる入力シーケンスを特定し、統計情報を集計します。
• 特徴量次元全体で並列計算を行い、GPU 上でトップ K の活性化シーケンスを直接追跡・更新することでメモリ使用量を最小化します。
• 収集した統計情報（最も活性化されるトークンやシーケンス）を用いて、LLM をプロンプトに組み込み、自然言語による特徴量の説明（解説）を自動生成します。

2.3 アトリビューショングラフと可視化

• Circuit-Tracer 統合: 生成された CLT を Circuit-Tracer ライブラリと統合し、特徴量レベルのアトリビューショングラフを効率的に計算・剪定します。
• インタラクティブ可視化: Dash ベースの Web インターフェースを提供し、アトリビューショングラフの探索、特徴量のクラスタリング、層間や特徴量間への介入（Intervention）を視覚的に行えるようにしています。これにより、研究ワークフローへの統合が容易になります。

3. 主要な貢献 (Key Contributions)

1. スケーラブルなトレーニング基盤: 最適化されたモデルシャーディングと、量子化・圧縮を伴う活性化値キャッシュによる、大規模 CLT の効率的なトレーニング環境を提供。
2. 効率的なアトリビューション計算: Circuit-Tracer とのネイティブ統合により、特徴量レベルのグラフ計算と剪定を高速化。
3. 統合された解釈性ワークフロー: 自動化された解釈（AutoInterp）、アトリビューショングラフの計算、インタラクティブな可視化を単一のライブラリで完結させ、研究のハードルを低下させた。

4. 結果と評価 (Results)

• 性能: GPT-2 などのモデルにおいて、既存の研究（Hanna et al., 2025 など）と同等の性能を達成しました（L0=100 で約 0.8 の説明分散、置換スコア約 0.8、グラフ完全性約 0.95）。
• 効率性: 特徴量シャーディングを実装したことで、DDP 比較においてメモリ効率を大幅に改善し、LLaMA 3.2 1B モデルに対して 8 GPU で拡張係数 48（約 150 万特徴量）のトレーニングを可能にしました。
• 実用性: 多言語分析や「greater-than」メカニズムの再現など、既存の研究を拡張・再現する能力を実証しました。また、可視化インターフェースを通じた直感的な分析が可能であることを示しました。

5. 意義と将来展望 (Significance & Future Work)

• 意義: CLT-Forge は、オープンソース界隈で初めて、大規模な CLT のトレーニングから解釈性分析、可視化までを統合的にサポートするフレームワークです。これにより、以前は Anthropic などのクローズド環境に限定されていた大規模スケールの機械的解釈性研究が、学術界や一般の開発者にもアクセス可能になりました。
• 将来展望:
  • アトリビューショングラフ計算におけるアテンションマップの凍結を解除し、アテンションの寄与を統合する。
  • パラメータ数とメモリフットプリントの削減に向けた、より効率的なアーキテクチャの研究。
  • 置換スコア（Replacement Score）の直接的な最適化など、より堅牢な学習目的関数の開発。

このライブラリは、機械的解釈性の研究における再現性を高め、大規模モデルの内部メカニズムを解明するための重要なインフラとして機能することが期待されています。