Upload folder using huggingface_hub

README.md

@@ -0,0 +1,143 @@
+---
+title: FunctionGemma 270M SAE
+language: en
+tags:
+  - sparse-autoencoder
+  - sae
+  - interpretability
+  - functiongemma
+  - gemma
+license: apache-2.0
+---
+
+# FunctionGemma 270M Sparse Autoencoders
+
+Sparse Autoencoders (SAEs) trained on all 18 layers of [google/functiongemma-270m-it](https://huggingface.co/google/functiongemma-270m-it).
+
+## Architecture
+
+- **Base Model**: google/functiongemma-270m-it
+- **Layers**: 18 (decoder-only)
+- **Hidden Size**: 640
+- **SAE Dimension**: 4096 (6.4x expansion)
+- **Hook Point**: `self_attn.o_proj` (output projection of self-attention)
+
+## Training
+
+- **Epochs**: 5 per layer
+- **Batch Size**: 1
+- **Learning Rate**: 1e-4
+- **Optimizer**: AdamW
+- **Loss**: MSE + 0.01 * L1 regularization
+- **Activation Clipping**: [-10, 10]
+- **Gradient Clipping**: max_norm=1.0
+
+## Checkpoints
+
+Each checkpoint contains:
+```python
+{
+    "model_name": "google/functiongemma-270m-it",
+    "layer_idx": int,
+    "d_in": 640,
+    "d_sae": 4096,
+    "W_enc": torch.Tensor,  # (640, 4096)
+    "b_enc": torch.Tensor,  # (4096,)
+    "W_dec": torch.Tensor,  # (4096, 640)
+    "b_dec": torch.Tensor,  # (640,)
+    "history": {
+        "loss": [...],
+        "mse": [...],
+        "l0": [...]
+    }
+}
+```
+
+## Usage
+
+```python
+import torch
+from huggingface_hub import hf_hub_download
+
+# Load SAE for a specific layer
+layer_idx = 0
+ckpt_path = hf_hub_download(
+    "mindchain/functiongemma-270m-sae",
+    f"sae_layer_{layer_idx:02d}.pt"
+)
+sae = torch.load(ckpt_path, map_location="cpu")
+
+# Use SAE
+class JumpReLUSAE(torch.nn.Module):
+    def __init__(self, W_enc, b_enc, W_dec, b_dec):
+        super().__init__()
+        self.W_enc = torch.nn.Parameter(W_enc)
+        self.b_enc = torch.nn.Parameter(b_enc)
+        self.W_dec = torch.nn.Parameter(W_dec)
+        self.b_dec = torch.nn.Parameter(b_dec)
+
+    def forward(self, x):
+        batch, seq, d_in = x.shape
+        x_flat = x.view(-1, d_in)
+        pre_act = x_flat @ self.W_enc + self.b_enc
+        features = torch.relu(pre_act)
+        recon = features @ self.W_dec + self.b_dec
+        return recon.view(batch, seq, d_in), features.view(batch, seq, -1)
+
+sae_model = JumpReLUSAE(
+    sae["W_enc"], sae["b_enc"],
+    sae["W_dec"], sae["b_dec"]
+)
+
+# Get activations from FunctionGemma and encode
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "google/functiongemma-270m-it",
+    torch_dtype=torch.bfloat16,
+    device_map="cuda"
+)
+tokenizer = AutoTokenizer.from_pretrained("google/functiongemma-270m-it")
+
+inputs = tokenizer("What's the weather?", return_tensors="pt").to(model.device)
+
+# Hook to get activations
+acts = []
+def hook(module, inp, out):
+    acts.append(out[0].detach().float())
+handle = model.model.layers[layer_idx].self_attn.o_proj.register_forward_hook(hook)
+with torch.no_grad():
+    _ = model(**inputs)
+handle.remove()
+
+# Run through SAE
+recon, features = sae_model(acts[0])
+print(f"Active features: {(features > 0).sum().item()}")
+```
+
+## Training Results
+
+| Layer | Final Loss | Final MSE | L0 |
+|-------|------------|-----------|-----|
+| 0 | 3.4457 | 3.1244 | 1225 |
+| 1 | 2.0052 | 1.9042 | 1386 |
+| 2 | 0.1182 | 0.0759 | 1546 |
+| 3 | 0.1182 | 0.0758 | 3096 |
+| 4 | 0.0361 | 0.0170 | 1635 |
+| 5 | 0.0414 | 0.0351 | 399 |
+| 6 | 0.0318 | 0.0138 | 1807 |
+| 7 | 0.0877 | 0.0661 | 1120 |
+| 8 | 0.0733 | 0.0445 | 1379 |
+| 9 | 0.0561 | 0.0317 | 1569 |
+| 10 | 0.0997 | 0.0852 | 591 |
+| 11 | 0.0252 | 0.0097 | 3658 |
+| 12 | 0.0565 | 0.0395 | 962 |
+| 13 | 0.0924 | 0.0619 | 1403 |
+| 14 | 0.2711 | 0.2504 | 709 |
+| 15 | 0.1501 | 0.1062 | 1576 |
+| 16 | 0.1670 | 0.1426 | 870 |
+| 17 | 0.0385 | 0.0218 | 1470 |
+
+## License
+
+Apache 2.0

sae_layer_00.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:68841d7a1a9e9a4dc690c0b0bd5296d1e540dda18f273970df465049b1de11a1
+size 20992760

sae_layer_01.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3949f21bce6742ce35e5c258f855a7ae786fa8be7d4d8a4cec20784a589a989
+size 20992760

sae_layer_02.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25608c0fed52bae9065bd35033428a010c0b79a00b16df3ec5c20132a98508be
+size 20992760

sae_layer_03.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:00999bab7630c4492cef40e47d6f2268ca8364fd3a6e3dd335c117239ac571ee
+size 20992760

sae_layer_04.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a5d6402ded655b9e7fb62a8df33d45158af3ffef9e1e31c39c392f218425897
+size 20992760

sae_layer_05.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:13bda23f02320a65ee466b3f88c0d421e652377fa26cc61491a8587964b6f0ba
+size 20992760

sae_layer_06.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ba06b5797ea7ce20522e8d04cbf93085db7333abf9839b8890d13b7971c0364
+size 20992760

sae_layer_07.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d61060bda5f66d5d6c1d9d25900076fd0092b296412fe9b16485c5a7ff4846d
+size 20992760

sae_layer_08.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8cb29be94c5490ad1b3d0ac9746089ab4567e54b27b4beb3a5c98fa57ac9cec9
+size 20992760

sae_layer_09.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:02e4e68c3e040c75b2ce73ef5fb052326926e104722a157319f4905c54c68594
+size 20992760

sae_layer_10.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f6037efd3b492a493858a1389f26c84a412cc37b1331ae6fbbdd6345bd83305f
+size 20992760

sae_layer_11.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d1109e982c2239c84e70e260bbcf48638618256b4b542e38622a1f07016153f
+size 20992760

sae_layer_12.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:62e542f0f548b6027e4a3660ae868f1c44ff96231da7f8b347646ca31dbaa5fc
+size 20992760

sae_layer_13.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67d6bcf3fa7d02c8bff0a484484d6e0dcd460b9ff4184063143410eb475240e0
+size 20992760

sae_layer_14.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8d5f2adabdad0c21d54a4828d4d1aa2848a939db40a97b2e609f5a8b55f1473d
+size 20992760

sae_layer_15.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa91a0170e4ddf1ab82d9d6a582ab2941458a64f4e81367dd95950dfd65ff238
+size 20992760

sae_layer_16.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eb51ce70a5bbf895f469f9c10e1cce34b434cd95852708e2d39b08ce70e07588
+size 20992760

sae_layer_17.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f8f923993944b6b37f2cd8b4e86c18b649d3f18556709feaeea3ab7a38791316
+size 20992760