Upload folder using huggingface_hub

app.py

@@ -157,6 +157,23 @@ def _load_image_models():
     return dit, lat, lat_std, sf, ae, tok, clip
 
 
+SAE_REPO = "rootxhacker/HobbyLM-SAE"
+
+
+def _load_sae():
+    key = ("sae",)
+    if key in _cache:
+        return _cache[key]
+    import json
+    from hobbylm.sae import TopKSAE, SAEConfig
+    meta = json.load(open(hf_hub_download(SAE_REPO, "meta.json")))
+    labels = json.load(open(hf_hub_download(SAE_REPO, "labels.json")))
+    sae = TopKSAE(SAEConfig(**meta["cfg"])).eval()
+    sae.load_state_dict(load_file(hf_hub_download(SAE_REPO, "sae.safetensors")))
+    _cache[key] = (sae, meta, labels)
+    return _cache[key]
+
+
 # --------------------------------------------------------------------------- chat
 
 def _build_prompt(repo, message, history):
@@ -380,6 +397,46 @@ def how_it_works(prompt, layer):
         model.to("cpu")
 
 
+# --------------------------------------------------------------------------- how it works (SAE features)
+
+@spaces.GPU(duration=90)
+def sae_features(prompt, topn):
+    dev = _device()
+    enc = _enc()
+    try:
+        sae, meta, labels = _load_sae()
+    except Exception as e:
+        return f"⚠️ SAE not available yet: {e}"
+    model, _ = _load_llm("HobbyLM-Base")
+    model.to(dev); sae.to(dev)
+    layer, scale = meta["layer"], float(meta["scale"])
+    topn = int(topn)
+    try:
+        ids = enc.encode_ordinary(prompt or "I love listening to music while coding software.")[:48]
+        if not ids:
+            ids = enc.encode_ordinary("Hello world")
+        toks = torch.tensor([ids], device=dev)
+        with torch.no_grad():
+            h = model(toks, capture_layer=layer).float() * scale
+            z = sae.encode(h.reshape(-1, sae.cfg.d_in))           # (S, m)
+        md = ("Each token's residual is decomposed into a few **interpretable features** from the SAE "
+              "dictionary. Below: per token, the strongest features (auto-labelled by the tokens they "
+              "fire on most).\n\n| token | top active features  ·  *(label · strength)* |\n|---|---|\n")
+        for s, tid in enumerate(ids):
+            v, f = z[s].topk(min(topn, z.shape[1]))
+            tok_str = enc.decode([tid]).replace("|", "¦").replace("\n", "⏎").strip() or "·"
+            parts = []
+            for val, fi in zip(v.tolist(), f.tolist()):
+                if val <= 1e-4:
+                    continue
+                lab = labels.get(str(int(fi)), {}).get("label") or f"feat#{int(fi)}"
+                parts.append(f"**{lab}** ({val:.1f})")
+            md += f"| `{tok_str}` | {' · '.join(parts) or '—'} |\n"
+        return md
+    finally:
+        model.to("cpu"); sae.to("cpu")
+
+
 # --------------------------------------------------------------------------- UI
 
 INTRO = """# 🪶 HobbyLM Playground
@@ -456,5 +513,17 @@ with gr.Blocks(title="HobbyLM Playground", theme=gr.themes.Soft()) as demo:
         hiw_load = gr.Plot(label="Expert load (balancing)")
         hiw_btn.click(how_it_works, [hiw_prompt, hiw_layer], [hiw_heat, hiw_load, hiw_summary])
 
+    with gr.Tab("🧠 What it represents"):
+        gr.Markdown(
+            "A **sparse autoencoder** (SAE) trained on HobbyLM-Base's layer-8 residual stream pulls apart each "
+            "activation into a handful of **interpretable features** from a 12,288-entry dictionary. Type text and "
+            "see which concepts light up on each token — words, synonym clusters, syntax, formatting. This is "
+            "*mechanistic interpretability*: looking at what the model actually represents inside.")
+        sae_prompt = gr.Textbox(label="Text", value="I love listening to music while coding software.")
+        sae_top = gr.Slider(2, 8, value=4, step=1, label="Features shown per token")
+        sae_btn = gr.Button("Show features", variant="primary")
+        sae_out = gr.Markdown()
+        sae_btn.click(sae_features, [sae_prompt, sae_top], sae_out)
+
 if __name__ == "__main__":
     demo.queue(max_size=20).launch()

hobbylm/sae.py

@@ -0,0 +1,106 @@
+"""Top-k Sparse Autoencoder for mechanistic interpretability of HobbyLM activations.
+
+A top-k SAE (Gao et al. 2024 / EleutherAI `sae`) over the residual stream: it reconstructs an
+activation x as a sparse, non-negative combination of a learned overcomplete dictionary, keeping only
+the k largest latents active. No L1 coefficient to tune (sparsity is exactly k), and an auxiliary loss
+resurrects dead features so the dictionary stays fully used.
+
+  x_centered = x - b_dec
+  z          = TopK( relu(x_centered @ W_enc + b_enc) )        # exactly k non-zeros
+  x_hat      = z @ W_dec + b_dec                                 # W_dec rows are unit-norm
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+@dataclass
+class SAEConfig:
+    d_in: int = 768          # model activation dim (residual stream)
+    d_sae: int = 12288       # dictionary size (16x expansion)
+    k: int = 32              # active latents per token
+    k_aux: int = 256         # dead-latent auxiliary reconstruction width
+    dead_after: int = 2_000_000   # a feature unused for this many tokens is "dead"
+    aux_coef: float = 1.0 / 32.0  # weight on the auxk resurrection loss
+
+
+def _topk(pre: Tensor, k: int) -> Tensor:
+    """Keep the k largest values per row (after the relu in the caller), zero the rest."""
+    vals, idx = pre.topk(k, dim=-1)
+    out = torch.zeros_like(pre)
+    return out.scatter_(-1, idx, vals)
+
+
+class TopKSAE(nn.Module):
+    def __init__(self, cfg: SAEConfig):
+        super().__init__()
+        self.cfg = cfg
+        d, m = cfg.d_in, cfg.d_sae
+        self.b_dec = nn.Parameter(torch.zeros(d))
+        self.b_enc = nn.Parameter(torch.zeros(m))
+        # decoder rows unit-norm; encoder initialised as the decoder transpose (standard tied init)
+        W_dec = F.normalize(torch.randn(m, d), dim=1)
+        self.W_dec = nn.Parameter(W_dec)
+        self.W_enc = nn.Parameter(W_dec.t().clone())
+        # tokens since each latent last fired (for dead-feature tracking); not a learned param
+        self.register_buffer("last_fired", torch.zeros(m, dtype=torch.long))
+        self.register_buffer("n_seen", torch.zeros((), dtype=torch.long))
+
+    # ---- encode / decode ----
+    def encode(self, x: Tensor) -> Tensor:
+        pre = (x - self.b_dec) @ self.W_enc + self.b_enc
+        return _topk(F.relu(pre), self.cfg.k)            # (..., d_sae), exactly k non-zeros
+
+    def decode(self, z: Tensor) -> Tensor:
+        return z @ self.W_dec + self.b_dec
+
+    def forward(self, x: Tensor):
+        """Returns (x_hat, z, loss_dict). x: (N, d_in)."""
+        pre = F.relu((x - self.b_dec) @ self.W_enc + self.b_enc)
+        z = _topk(pre, self.cfg.k)
+        x_hat = self.decode(z)
+        recon = F.mse_loss(x_hat, x)
+
+        # ---- dead-feature bookkeeping + auxk resurrection ----
+        with torch.no_grad():
+            fired = (z > 0).any(dim=0)                    # (d_sae,)
+            self.last_fired += x.shape[0]
+            self.last_fired[fired] = 0
+            self.n_seen += x.shape[0]
+            dead = self.last_fired > self.cfg.dead_after
+        aux = x.new_zeros(())
+        n_dead = int(dead.sum())
+        if n_dead > 0 and self.cfg.k_aux > 0:
+            # reconstruct the residual error using only the top-k_aux DEAD latents (revives them)
+            resid = x - x_hat.detach()
+            pre_dead = pre.masked_fill(~dead[None, :], 0.0)
+            z_aux = _topk(pre_dead, min(self.cfg.k_aux, n_dead))
+            resid_hat = z_aux @ self.W_dec                # no b_dec: model the centred residual
+            aux = F.mse_loss(resid_hat, resid)
+            loss = recon + self.cfg.aux_coef * aux
+        else:
+            loss = recon
+        return x_hat, z, {"loss": loss, "recon": recon.detach(), "aux": aux.detach(),
+                          "n_dead": n_dead}
+
+    # ---- keep decoder rows unit-norm (call after each optimizer step) ----
+    @torch.no_grad()
+    def normalize_decoder(self):
+        self.W_dec.data = F.normalize(self.W_dec.data, dim=1)
+
+    @torch.no_grad()
+    def set_decoder_to_geometric_mean(self, x: Tensor):
+        """Initialise b_dec to the data mean (standard) — call once on the first activation batch."""
+        self.b_dec.data = x.mean(dim=0)
+
+
+def fraction_variance_explained(x: Tensor, x_hat: Tensor) -> float:
+    """1 - ||x - x_hat||^2 / ||x - mean(x)||^2  (per-batch FVU complement)."""
+    num = (x - x_hat).pow(2).sum()
+    den = (x - x.mean(0, keepdim=True)).pow(2).sum().clamp_min(1e-8)
+    return float((1.0 - num / den).detach())