Upload trained model

README.md

@@ -0,0 +1,163 @@
+---
+tags:
+  - sparse-autoencoder
+  - crosscoder
+  - interpretability
+  - qwen2
+  - mechanistic-interpretability
+  - dictionary-learning
+license: mit
+---
+
+# dfc-D8k-excl3-freeexcl-k90
+
+A **DFC** sparse crosscoder trained to compare layer-13 activations between:
+- **Model A (ToolRL)**: `chengq9/ToolRL-Qwen2.5-3B` — fine-tuned with tool-use reinforcement learning
+- **Model B (Base)**: `Qwen/Qwen2.5-3B` — vanilla base model
+
+## What is this?
+
+This model learns a sparse dictionary of features from the internal representations of two language models. By comparing which features activate for which model, we can identify:
+- **What the ToolRL fine-tuning changed** (A-exclusive features)
+- **What remained the same** (shared features)
+- **What the base model does that ToolRL suppressed** (B-exclusive features)
+
+## Model Architecture
+
+Dedicated Feature CrossCoder with partitioned dictionary (0.03/0.03 A/B exclusive)
+
+| Parameter | Value |
+|-----------|-------|
+| Dictionary size | 8192 |
+| Top-k active features | 90 |
+| Layer | 13 (middle layer of Qwen2.5-3B) |
+| Activation dimension | 2048 |
+| A-exclusive features | 245 (0.03) |
+| B-exclusive features | 245 (0.03) |
+| Shared features | 7702 |
+
+### How it works
+
+1. **Encode**: Takes stacked activations `(batch, 2, 2048)` from both models, applies per-model encoder weights, sums across models, and selects the top-90 features via ReLU + top-k.
+2. **Decode**: Reconstructs per-model activations from the sparse feature vector using per-model decoder weights.
+3. **Partition masks** (DFC only): Hard binary masks zero out encoder/decoder weights to enforce that exclusive features cannot be used by the wrong model.
+
+## Training
+
+| Parameter | Value |
+|-----------|-------|
+| Loss function | MSE + L1 on shared features only (coef: 1e-3); exclusive features unpenalized |
+| Training steps | 9000 |
+| Learning rate | 1e-4 |
+| Batch size | 1024 |
+| Sparsity coefficient (shared) | 1e-3 |
+| Exclusive sparsity coefficient | 0 |
+| Optimizer | Adam (grad clip 1.0) |
+| W&B project | `dfc-crosscoder-sweep` |
+
+### Training Data
+
+- **FineWeb**: ~40,000 general web text samples (from `HuggingFaceFW/fineweb` sample-10BT)
+- **ToolRL**: ~40,000 tool-use conversation samples (from `emrecanacikgoz/ToolRL`, cycled)
+- Activations extracted from layer 13, last token per sample
+- Both datasets concatenated and z-score normalized
+
+## Usage
+
+### Quick Start
+
+```python
+import torch
+from huggingface_hub import hf_hub_download
+
+# Download model files
+repo_id = "antebe1/dfc-D8k-excl3-freeexcl-k90"
+for fname in ["model.pt", "config.json", "dfc.py"]:
+    hf_hub_download(repo_id=repo_id, filename=fname, local_dir="./model")
+
+# Load the crosscoder
+import sys; sys.path.insert(0, "./model")
+from dfc import DFCCrossCoder
+
+dfc = DFCCrossCoder.load("./model", device="cuda")
+print(f"Loaded: dict_size={dfc.dict_size}, k={dfc.k}")
+```
+
+### Extract Features from Real Models
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Load both models
+model_a = AutoModelForCausalLM.from_pretrained("chengq9/ToolRL-Qwen2.5-3B", device_map="cuda:0")
+model_b = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-3B", device_map="cuda:1")
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B")
+
+# Get activations from layer 13
+# NOTE: hidden_states[0] = embeddings, hidden_states[i] = output of layer i-1
+#       so layer 13 activations are at index 13+1
+text = "Use the search tool to find recent papers on RLHF"
+inputs = tokenizer(text, return_tensors="pt")
+
+with torch.no_grad():
+    out_a = model_a(**inputs.to("cuda:0"), output_hidden_states=True)
+    out_b = model_b(**inputs.to("cuda:1"), output_hidden_states=True)
+    act_a = out_a.hidden_states[13 + 1][:, -1, :]  # last token, layer 13
+    act_b = out_b.hidden_states[13 + 1][:, -1, :]
+
+# Stack and encode
+activations = torch.stack([act_a.cpu(), act_b.cpu()], dim=1)  # (1, 2, 2048)
+features = dfc.encode(activations.to(dfc.W_enc.device))
+
+print(f"Active features: {(features > 0).sum().item()} / {dfc.dict_size}")
+```
+
+### Analyze Partitions (DFC only)
+
+```python
+stats = dfc.feature_stats(features)
+print(f"L0 total:    {stats['l0_total']:.1f}")
+print(f"L0 A-excl:   {stats['l0_a_excl']:.1f}")
+print(f"L0 B-excl:   {stats['l0_b_excl']:.1f}")
+print(f"L0 shared:   {stats['l0_shared']:.1f}")
+
+# Check reconstruction quality
+recon, feats = dfc(activations.to(dfc.W_enc.device))
+mse = torch.nn.functional.mse_loss(recon.cpu(), activations)
+print(f"Reconstruction MSE: {mse.item():.6f}")
+```
+
+## Files
+
+| File | Description |
+|------|-------------|
+| `model.pt` | PyTorch state dict (encoder/decoder weights + partition masks) |
+| `config.json` | Architecture config: dict_size, k, partition sizes (n_a, n_b) |
+| `hparams.json` | Full training hyperparameters including loss, lr, steps, etc. |
+| `dfc.py` | `DFCCrossCoder` class definition — required to load model.pt |
+| `demo.py` | Feature extraction demo (works with downloaded model) |
+| `requirements.txt` | Python dependencies |
+
+## Part of a Sweep
+
+This model is one of 48 models in a hyperparameter sweep. See the full collection:
+
+| Axis | Values |
+|------|--------|
+| k (top-k) | 45, 90, 160 |
+| dict_size | 8,192 / 16,384 |
+| Architecture | DFC (partitioned) / CrossCoder (all shared) |
+| Exclusive % (DFC) | 3%, 5%, 10% |
+| Exclusive sparsity | 1e-3 (penalized) / 0 (free) |
+| CrossCoder L1 | with / without |
+
+## Citation
+
+```bibtex
+@misc{dfc-D8k-excl3-freeexcl-k90,
+  title={DFC CrossCoder: ToolRL vs Base Qwen2.5-3B},
+  author={Andre Shportko},
+  year={2026},
+  url={https://huggingface.co/antebe1/dfc-D8k-excl3-freeexcl-k90}
+}
+```

config.json

@@ -0,0 +1,7 @@
+{
+  "activation_dim": 2048,
+  "dict_size": 8192,
+  "k": 90,
+  "n_a": 245,
+  "n_b": 245
+}

demo.py

@@ -0,0 +1,131 @@
+"""
+DFC CrossCoder Demo - Quick Feature Vector Extraction
+=====================================================
+
+This demo shows how to use the DFC CrossCoder to extract feature vectors
+from text using the ToolRL vs Base Qwen2.5-3B comparison.
+
+Works both locally (from checkpoints/) and from a HuggingFace download.
+"""
+
+import sys
+from pathlib import Path
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from dfc import DFCCrossCoder
+
+# Default model path — override with command line arg or set MODEL_PATH env var
+import os
+DEFAULT_MODEL_PATH = os.environ.get("DFC_MODEL_PATH", "./checkpoints/dfc2")
+LAYER = 13
+
+
+def quick_demo(model_path: str = DEFAULT_MODEL_PATH):
+    """10-line core demo: Load models and extract features from text"""
+
+    # Load the trained CrossCoder
+    dfc = DFCCrossCoder.load(model_path, device="cuda" if torch.cuda.is_available() else "cpu")
+
+    # Load tokenizer and models (for demo - in practice you might use cached activations)
+    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B")
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    # Example texts: general vs tool-use
+    general_text = "The French Revolution began in 1789 with widespread social discontent."
+    tooluse_text = "Search for the cheapest flight from London to Tokyo departing next Friday."
+    
+    print("🚀 DFC CrossCoder Demo")
+    print("="*50)
+    
+    for text, label in [(general_text, "General"), (tooluse_text, "Tool-Use")]:
+        print(f"\n{label} Text: '{text}'")
+        
+        # Tokenize and get activations (simplified - normally you'd extract from layer 13)
+        tokens = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+        
+        # For demo: create mock activations (normally extracted from models at layer 13)
+        # Shape: (batch=1, models=2, hidden_dim=2048)
+        mock_activations = torch.randn(1, 2, dfc.activation_dim, device=dfc.W_enc.device)
+
+        # Core CrossCoder usage (the main 3 lines!)
+        features = dfc.encode(mock_activations)  # Get sparse feature vector
+        reconstruction, _ = dfc(mock_activations)  # Full forward pass
+        
+        # Analyze partition breakdowns
+        active_features = (features > 0).sum().item()
+        a_exclusive = (features[0, :dfc.a_end] > 0).sum().item()
+        b_exclusive = (features[0, dfc.a_end:dfc.b_end] > 0).sum().item()
+        shared = (features[0, dfc.b_end:] > 0).sum().item()
+        
+        print(f"  ✅ Active features: {active_features}/{dfc.dict_size}")
+        print(f"  🔧 A-exclusive (ToolRL): {a_exclusive}")
+        print(f"  📝 B-exclusive (Base): {b_exclusive}")
+        print(f"  🤝 Shared: {shared}")
+        
+        # Show top features
+        top_vals, top_idx = torch.topk(features[0], k=5)
+        print(f"  🔝 Top features: {top_idx.tolist()} (values: {top_vals.tolist()})")
+
+def extract_real_activations(model_path: str = DEFAULT_MODEL_PATH):
+    """Extended demo with real model activations (requires more memory)"""
+    print("\n Extended Demo with Real Model Activations")
+    print("="*50)
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Load models (this requires significant memory!)
+    print("Loading ToolRL and Base models...")
+    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B")
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model_a = AutoModelForCausalLM.from_pretrained("chengq9/ToolRL-Qwen2.5-3B").to(device).eval()
+    model_b = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-3B").to(device).eval()
+
+    # Load CrossCoder
+    dfc = DFCCrossCoder.load(model_path, device=device)
+
+    # Extract real activations from layer 13
+    text = "Use the calculator to compute 15% tip on $45.50"
+    tokens = tokenizer(text, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        # hidden_states[0] = embeddings, hidden_states[i] = output of layer i-1
+        # so layer 13 activations are at index LAYER + 1
+        outputs_a = model_a(**tokens, output_hidden_states=True)
+        outputs_b = model_b(**tokens, output_hidden_states=True)
+
+        hidden_a = outputs_a.hidden_states[LAYER + 1][0, -1]  # (hidden_dim,)
+        hidden_b = outputs_b.hidden_states[LAYER + 1][0, -1]  # (hidden_dim,)
+
+        # Stack for CrossCoder: (1, 2, hidden_dim)
+        activations = torch.stack([hidden_a, hidden_b], dim=0).unsqueeze(0)
+
+        # Run CrossCoder
+        features = dfc.encode(activations)
+
+        # Analysis
+        active_count = (features > 0).sum().item()
+        print(f"\n Real Activation Analysis:")
+        print(f"   Text: '{text}'")
+        print(f"   Active features: {active_count}/{dfc.dict_size}")
+
+        # Partition breakdown
+        a_active = (features[0, :dfc.a_end] > 0).sum().item()
+        b_active = (features[0, dfc.a_end:dfc.b_end] > 0).sum().item()
+        s_active = (features[0, dfc.b_end:] > 0).sum().item()
+
+        print(f"   ToolRL-specific: {a_active}")
+        print(f"   Base-specific: {b_active}")
+        print(f"   Shared: {s_active}")
+
+if __name__ == "__main__":
+    model_path = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_MODEL_PATH
+
+    # Run quick demo with mock data
+    quick_demo(model_path)
+
+    # Uncomment to run with real models (requires significant GPU memory)
+    # extract_real_activations(model_path)

dfc.py

@@ -0,0 +1,222 @@
+"""
+dfc.py — Dedicated Feature CrossCoder (DFC) model.
+
+Feature layout in dict_size
+────────────────────────────
+  ┌─────────────────────┬─────────────────────┬──────────────────────────┐
+  │  A-exclusive (n_a)  │  B-exclusive (n_b)  │     Shared (n_shared)    │
+  └─────────────────────┴─────────────────────┴──────────────────────────┘
+  idx:  0 ─────── a_end ──────── b_end ───────────────────── dict_size
+
+Constraints (enforced by gradient masking + _apply_masks every step)
+──────────────────────────────────────────────────────────────────────
+  • Model A cannot encode/decode B-exclusive features
+  • Model B cannot encode/decode A-exclusive features
+  • Shared features are accessible to both
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DFCCrossCoder(nn.Module):
+
+    def __init__(
+        self,
+        activation_dim: int,
+        dict_size: int,
+        k: int,
+        model_a_exclusive_pct: float = 0.05,
+        model_b_exclusive_pct: float = 0.05,
+    ):
+        super().__init__()
+        self.activation_dim = activation_dim
+        self.dict_size = dict_size
+        self.k = k
+
+        self.n_a     = int(dict_size * model_a_exclusive_pct)
+        self.n_b     = int(dict_size * model_b_exclusive_pct)
+        self.n_shared = dict_size - self.n_a - self.n_b
+        self.a_end   = self.n_a
+        self.b_end   = self.n_a + self.n_b
+
+        print(
+            f"[DFC] dict={dict_size} k={k} | "
+            f"A-excl={self.n_a} B-excl={self.n_b} shared={self.n_shared}"
+        )
+
+        # Encoder: W_enc[model, d_in, dict_size]
+        self.W_enc = nn.Parameter(
+            torch.randn(2, activation_dim, dict_size) / (activation_dim ** 0.5)
+        )
+        self.b_enc = nn.Parameter(torch.zeros(dict_size))
+
+        # Decoder: W_dec[dict_size, model, d_in]
+        self.W_dec = nn.Parameter(
+            torch.randn(dict_size, 2, activation_dim) / (dict_size ** 0.5)
+        )
+        self.b_dec = nn.Parameter(torch.zeros(2, activation_dim))
+
+        # ── Partition masks (move with .to(device)) ───────────────────
+        # enc_mask[model, dict_size]
+        enc_mask = torch.ones(2, dict_size)
+        enc_mask[1, : self.a_end] = 0                   # B cannot encode A-excl
+        enc_mask[0, self.a_end : self.b_end] = 0        # A cannot encode B-excl
+        self.register_buffer("enc_mask", enc_mask)
+
+        # dec_mask[dict_size, model]
+        dec_mask = torch.ones(dict_size, 2)
+        dec_mask[: self.a_end, 1] = 0                   # A-excl: B decoder = 0
+        dec_mask[self.a_end : self.b_end, 0] = 0        # B-excl: A decoder = 0
+        self.register_buffer("dec_mask", dec_mask)
+
+        self._apply_masks()
+
+    # ── Weight enforcement ────────────────────────────────────────────
+
+    @torch.no_grad()
+    def _apply_masks(self):
+        """Zero forbidden weights. Call after every optimiser step."""
+        for m in range(2):
+            self.W_enc.data[m] *= self.enc_mask[m].unsqueeze(0)
+        for m in range(2):
+            self.W_dec.data[:, m, :] *= self.dec_mask[:, m].unsqueeze(1)
+
+    # ── Forward ───────────────────────────────────────────────────────
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """x: (B, 2, d) → features: (B, dict_size) sparse top-k."""
+        W = self.W_enc * self.enc_mask.unsqueeze(1)         # (2, d, dict)
+        pre = torch.einsum("bmd,mdf->bf", x, W) + self.b_enc
+        pre = F.relu(pre)
+        topk_vals, topk_idx = torch.topk(pre, self.k, dim=-1)
+        features = torch.zeros_like(pre)
+        features.scatter_(-1, topk_idx, topk_vals)
+        return features
+
+    def decode(self, features: torch.Tensor) -> torch.Tensor:
+        """features: (B, dict_size) → (B, 2, d)."""
+        W = self.W_dec * self.dec_mask.unsqueeze(-1)        # (dict, 2, d)
+        return torch.einsum("bf,fmd->bmd", features, W) + self.b_dec
+
+    def forward(self, x: torch.Tensor):
+        """x: (B, 2, d) → (reconstruction, features)."""
+        features = self.encode(x)
+        recon    = self.decode(features)
+        return recon, features
+
+    def loss(
+        self, 
+        x: torch.Tensor, 
+        sparsity_coef: float = 1e-3,
+        exclusive_sparsity_coef: float = 1e-3  # Lower penalty for exclusive features
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """MSE + weighted L1 sparsity. Returns (total, mse, l1_shared, l1_exclusive)."""
+        recon, features = self.forward(x)
+        mse = F.mse_loss(recon, x)
+        
+        # Split features by partition
+        # fa = features[:, :self.a_end]          # A-exclusive
+        # fb = features[:, self.a_end:self.b_end] # B-exclusive 
+        fs = features[:, self.b_end:]          # Shared
+
+        # A sees: A-exclusive + shared
+        fa = torch.cat([features[:, :self.a_end], features[:, self.b_end:]], dim=-1)   # A-exclusive + shared
+        fb = torch.cat([features[:, self.a_end:self.b_end], features[:, self.b_end:]], dim=-1)  # B-exclusive + shared
+        
+        # Separate sparsity penalties
+        l1_shared = fs.abs().mean()
+        l1_exclusive = (fa.abs().mean() + fb.abs().mean()) / 2
+        total = mse + exclusive_sparsity_coef * l1_exclusive + sparsity_coef * l1_shared
+        
+        return total, mse, l1_shared, l1_exclusive
+
+    # ── Diagnostics ───────────────────────────────────────────────────
+
+    @torch.no_grad()
+    def verify_partition_integrity(self) -> dict[str, float]:
+        """Max absolute value in weights that should be zero."""
+        if self.n_a == 0 and self.n_b == 0:
+            return {"enc_max_violation": 0.0, "dec_max_violation": 0.0}
+        enc_viol  = (self.W_enc.abs() * (1 - self.enc_mask).unsqueeze(1)).max().item()
+        dec_viol_a = self.W_dec[: self.a_end, 1, :].abs().max().item() if self.n_a > 0 else 0.0
+        dec_viol_b = self.W_dec[self.a_end : self.b_end, 0, :].abs().max().item() if self.n_b > 0 else 0.0
+        return {
+            "enc_max_violation": enc_viol,
+            "dec_max_violation": max(dec_viol_a, dec_viol_b),
+        }
+
+    @torch.no_grad()
+    def feature_stats(self, features: torch.Tensor) -> dict[str, float]:
+        """Partition-level activation stats for a batch of features."""
+        fa = features[:, : self.a_end]
+        fb = features[:, self.a_end : self.b_end]
+        fs = features[:, self.b_end :]
+        return {
+            "l0_total":    (features > 0).float().sum(dim=-1).mean().item(),
+            "l0_a_excl":   (fa > 0).float().sum(dim=-1).mean().item(),
+            "l0_b_excl":   (fb > 0).float().sum(dim=-1).mean().item(),
+            "l0_shared":   (fs > 0).float().sum(dim=-1).mean().item(),
+            "mean_a_excl": fa.mean().item(),
+            "mean_b_excl": fb.mean().item(),
+            "mean_shared": fs.mean().item(),
+        }
+
+    # ── Save / Load ───────────────────────────────────────────────────
+
+    def save(self, path: str) -> None:
+        """Save model with disk space checking and error handling."""
+        import shutil
+        import tempfile
+        
+        # Check available disk space
+        free_space = shutil.disk_usage(Path(path).parent or ".").free
+        if free_space < 100_000_000:  # Less than 100MB
+            raise RuntimeError(f"Insufficient disk space: {free_space / 1e9:.2f}GB available. Need at least 0.1GB.")
+        
+        Path(path).mkdir(parents=True, exist_ok=True)
+        
+        # Save to temporary file first, then move to avoid corruption
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.pt') as tmp_file:
+            try:
+                torch.save(self.state_dict(), tmp_file.name)
+                shutil.move(tmp_file.name, f"{path}/model.pt")
+            except Exception as e:
+                # Clean up temp file on error
+                if Path(tmp_file.name).exists():
+                    Path(tmp_file.name).unlink()
+                raise RuntimeError(f"Failed to save model: {e}")
+        
+        # Save config
+        config_data = dict(
+            activation_dim=self.activation_dim,
+            dict_size=self.dict_size,
+            k=self.k,
+            n_a=self.n_a,
+            n_b=self.n_b,
+        )
+        with open(f"{path}/config.json", "w") as f:
+            json.dump(config_data, f, indent=2)
+        
+        print(f"[DFC] Saved → {path}")
+
+    @classmethod
+    def load(cls, path: str, device: str = "cpu") -> "DFCCrossCoder":
+        cfg = json.load(open(f"{path}/config.json"))
+        model = cls(
+            activation_dim=cfg["activation_dim"],
+            dict_size=cfg["dict_size"],
+            k=cfg["k"],
+            model_a_exclusive_pct=cfg["n_a"] / cfg["dict_size"],
+            model_b_exclusive_pct=cfg["n_b"] / cfg["dict_size"],
+        )
+        model.load_state_dict(
+            torch.load(f"{path}/model.pt", map_location=device, weights_only=True)
+        )
+        return model.to(device)

hparams.json

@@ -0,0 +1,19 @@
+{
+  "architecture": "DFC",
+  "model_a": "chengq9/ToolRL-Qwen2.5-3B",
+  "model_b": "Qwen/Qwen2.5-3B",
+  "layer": 13,
+  "activation_dim": 2048,
+  "dict_size": 8192,
+  "k": 90,
+  "model_a_exclusive_pct": 0.03,
+  "model_b_exclusive_pct": 0.03,
+  "steps": 9000,
+  "lr": 1e-4,
+  "train_batch": 1024,
+  "sparsity_coef": 1e-3,
+  "exclusive_sparsity_coef": 0,
+  "loss": "MSE + L1 on shared features only (coef: 1e-3); exclusive features unpenalized",
+  "wandb_project": "dfc-crosscoder-sweep",
+  "trained_at": "2026-03-30T03:53:46+00:00"
+}

model.pt

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

requirements.txt

@@ -0,0 +1,8 @@
+torch>=2.0.0
+transformers>=4.30.0
+datasets>=2.0.0
+wandb>=0.15.0
+tqdm>=4.64.0
+matplotlib>=3.5.0
+huggingface_hub>=0.16.0
+numpy>=1.21.0