Gemma 2 2b Crosscoder L13 Mu4.1e 02 Lr1e 04

🚀 クロスコーダーモデル

このクロスコーダーは、finewebとlsmsy - chat - 1mデータセットの一部におけるGemma 2 2BとGemma 2 2B ITのレイヤー13の並列活性化を用いて学習されました。

🚀 クイックスタート

📦 インストール

このモデルを使用するには、dictionary_learningライブラリの特定のブランチをインストールする必要があります。以下のコマンドを実行してください。

!pip install git+https://github.com/jkminder/dictionary_learning

💻 使用例

基本的な使用法

from dictionary_learning import CrossCoder
from nnsight import LanguageModel
import torch as th

crosscoder = CrossCoder.from_pretrained("Butanium/gemma-2-2b-crosscoder-l13-mu4.1e-02-lr1e-04", from_hub=True)
gemma_2 = LanguageModel("google/gemma-2-2b", device_map="cuda:0")
gemma_2_it = LanguageModel("google/gemma-2-2b-it", device_map="cuda:1")
prompt = "quick fox brown"

with gemma_2.trace(prompt):
    l13_act_base = gemma_2.model.layers[13].output[0][:, -1].save() # (1, 2304)
    gemma_2.model.layers[13].output.stop()

with gemma_2_it.trace(prompt):
    l13_act_it = gemma_2_it.model.layers[13].output[0][:, -1].save() # (1, 2304)
    gemma_2_it.model.layers[13].output.stop()


crosscoder_input = th.cat([l13_act_base, l13_act_it], dim=0).unsqueeze(0).cpu() # (batch, 2, 2304)
print(crosscoder_input.shape)
reconstruction, features = crosscoder(crosscoder_input, output_features=True)

# print metrics
print(f"MSE loss: {th.nn.functional.mse_loss(reconstruction, crosscoder_input).item():.2f}")
print(f"L1 sparsity: {features.abs().sum():.1f}")
print(f"L0 sparsity: {(features > 1e-4).sum()}")

📄 ライセンス

このモデルはMITライセンスの下で提供されています。

📚 ドキュメント

モデル情報