Update README with detailed specs

README.md

@@ -0,0 +1,149 @@
+---
+license: mit
+library_name: sae-lens
+tags:
+- sae
+- sparse-autoencoder
+- t5gemma
+- t5gemma2
+- google-t5gemma
+- interpretability
+- mechanistic-interpretability
+- sparse-autoencoders
+---
+
+# T5Gemma 2-270M Sparse Autoencoders (All 36 Layers)
+
+Sparse Autoencoders (SAEs) trained on all layers of `google/t5gemma-2-270m-270m`.
+
+## Model Specifications
+
+| Property | Value |
+|----------|-------|
+| **Base Model** | `google/t5gemma-2-270m-270m` |
+| **Architecture** | T5 Encoder-Decoder |
+| **Encoder Parameters** | ~270M |
+| **Decoder Parameters** | ~270M |
+| **Total Parameters** | ~540M |
+| **Encoder Layers** | 18 |
+| **Decoder Layers** | 18 |
+| **Hidden Size (d_model)** | 640 |
+| **FFN Dimension** | 2,560 |
+| **Attention Heads** | 10 |
+| **Vocabulary Size** | 32,128 |
+
+## SAE Configuration
+
+| Property | Value |
+|----------|-------|
+| **SAE Input Dimension (d_in)** | 640 |
+| **SAE Hidden Dimension (d_sae)** | 4,096 |
+| **Expansion Factor** | 6.4× |
+| **L1 Coefficient** | 0.01 |
+| **Training Epochs** | 5 |
+| **Batch Size** | 2 |
+| **Learning Rate** | 1e-4 |
+| **Optimizer** | Adam |
+| **Activation Hook** | `self_attn.o_proj` |
+| **Precision** | float16 (model), float32 (SAE) |
+
+## Coverage
+
+| Component | Layers | Count |
+|-----------|--------|-------|
+| Encoder | 0-17 | 18 SAEs |
+| Decoder | 0-17 | 18 SAEs |
+| **Total** | **36** | **36 SAEs** |
+
+## Usage
+
+```python
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import torch
+from huggingface_hub import hf_hub_download
+
+# Load base model
+model = AutoModelForSeq2SeqLM.from_pretrained("google/t5gemma-2-270m-270m")
+tokenizer = AutoTokenizer.from_pretrained("google/t5gemma-2-270m-270m")
+
+# Load SAE for a specific layer
+sae_path = hf_hub_download(
+    repo_id="mindchain/t5gemma2-sae-all-layers",
+    filename="encoder/sae_encoder_00.pt"
+)
+sae = torch.load(sae_path, map_location="cpu")
+
+# SAE forward pass
+def forward_sae(sae, activations):
+    features = torch.relu(activations @ sae['W_enc'] + sae['b_enc'])
+    reconstructed = features @ sae['W_dec'] + sae['b_dec']
+    return reconstructed, features
+
+# Example: Extract activations and run through SAE
+inputs = tokenizer("Translate to German: Hello world", return_tensors="pt")
+with torch.no_grad():
+    # Get encoder layer 0 output (you'd need to hook this)
+    activations = ... # hook self_attn.o_proj of encoder layer 0
+    recon, features = forward_sae(sae, activations)
+```
+
+## SAE Checkpoint Structure
+
+Each checkpoint file contains:
+
+```python
+{
+    'model_name': 'google/t5gemma-2-270m-270m',
+    'layer_type': 'encoder' or 'decoder',
+    'layer_idx': int,
+    'd_in': 640,
+    'd_sae': 4096,
+    'W_enc': Tensor([640, 4096]),   # Encoder weights
+    'b_enc': Tensor([4096]),         # Encoder bias
+    'W_dec': Tensor([4096, 640]),   # Decoder weights
+    'b_dec': Tensor([640]),          # Decoder bias
+    'history': {
+        'loss': [...],      # Loss per epoch
+        'l0': [...]         # Sparsity (active features) per epoch
+    }
+}
+```
+
+## Training Details
+
+- **Dataset**: 1500 diverse text samples, repeated across epochs
+- **Hook Point**: Self-attention output projection (`model.{encoder,decoder}.layers.{N}.self_attn.o_proj`)
+- **Training Duration**: ~2 hours on single GPU
+- **Final Metrics** (example, layer 0):
+  - Loss: ~0.0014
+  - L0 (active features): ~1367/4096 (~33% sparsity)
+
+## Files
+
+```
+encoder/
+  ├── sae_encoder_00.pt  # Encoder layer 0
+  ├── sae_encoder_01.pt  # Encoder layer 1
+  ├── ...
+  └── sae_encoder_17.pt  # Encoder layer 17
+
+decoder/
+  ├── sae_decoder_00.pt  # Decoder layer 0
+  ├── sae_decoder_01.pt  # Decoder layer 1
+  ├── ...
+  └── sae_decoder_17.pt  # Decoder layer 17
+```
+
+## License
+
+MIT
+
+## Credits
+
+Trained by [mindchain](https://huggingface.co/mindchain)
+
+## References
+
+- [T5Gemma 2 Model Card](https://huggingface.co/google/t5gemma-2-270m-270m)
+- [SAELens](https://github.com/decoderesearch/SAELens)
+- [TransformerLens](https://github.com/TransformerLensOrg/TransformerLens)