#!/usr/bin/env python3 """Memory Slot Surgery: check alignment between SAE decoder atoms and GDN k*v^T writes. GDN update: S_t = alpha_t * S_{t-1} + beta_t * k_t * v_t^T Bilinear SAE decoder: feature i contributes c_i * outer(V_dec[i,0], W_dec[i,0]) Question: do the SAE decoder atoms align with actual k_t * v_t^T writes at positions where the feature fires? This script: 1. Loads Qwen3.5-0.8B and a bilinear SAE checkpoint for L9H4 2. Runs the model on 100 sequences, hooking the GDN layer to capture k, v, beta, g 3. Simulates the recurrence to get per-position states 4. Encodes each position's state through the SAE to get feature activations 5. For each feature, finds top-50 activating positions, computes cosine alignment between V_dec/W_dec and the k_t/v_t at those positions 6. Reports histograms, top features, and null comparison """ import json import os import sys import time from pathlib import Path import numpy as np import torch import torch.nn.functional as F from tqdm import tqdm sys.path.insert(0, str(Path(__file__).resolve().parent)) from sae import BilinearMatrixSAE, build_sae_from_config # Model + tokenizer def load_model(model_name: str, device: str): from transformers import AutoModelForCausalLM, AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( model_name, dtype=torch.bfloat16, device_map=device, trust_remote_code=True, ) model.eval() return model, tokenizer # SAE loader def load_sae(ckpt_path: str, device: str) -> BilinearMatrixSAE: ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) cfg = ckpt.get("config", {}) sd = ckpt["model_state_dict"] sae = build_sae_from_config(cfg, state_dict=sd) sae.load_state_dict(sd) sae = sae.to(device).eval() assert isinstance(sae, BilinearMatrixSAE), f"Expected BilinearMatrixSAE, got {type(sae)}" return sae # Corpus loading (streaming openwebtext) def load_corpus_batches(tokenizer, n_seqs: int, seq_len: int, batch_size: int): from datasets import load_dataset ds = load_dataset("Skylion007/openwebtext", split="train", streaming=True) token_ids = [] target = n_seqs * seq_len * 2 for example in ds: token_ids.extend(tokenizer.encode(example["text"], add_special_tokens=False)) if len(token_ids) >= target: break n_actual = min(n_seqs, len(token_ids) // seq_len) batches = [] for start in range(0, n_actual, batch_size): end = min(start + batch_size, n_actual) seqs = [token_ids[i * seq_len : (i + 1) * seq_len] for i in range(start, end)] batches.append(torch.tensor(seqs, dtype=torch.long)) return batches # GDN hook: capture k, v, beta, g from layer forward class GDNWriteCapture: """Register a forward hook on a GDN layer to capture k, v, beta, g.""" def __init__(self, model, layer_idx: int): self.layer_idx = layer_idx self.k = None # (batch, seq_len, n_heads, head_k_dim) self.v = None self.beta = None # (batch, seq_len, n_heads) self.g = None # (batch, seq_len, n_heads) self._original_forward = None self._model_layer = None layers = model.model.layers if hasattr(model, "model") else model.layers self._model_layer = layers[layer_idx] gdn = self._model_layer.linear_attn self._gdn = gdn self._original_forward = gdn.forward gdn.forward = self._hooked_forward def _hooked_forward(self, hidden_states, cache_params=None, attention_mask=None): # Reproduce the GDN forward, capturing k, v, beta, g from transformers.models.qwen3_5.modeling_qwen3_5 import apply_mask_to_padding_states gdn = self._gdn hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) batch_size, seq_len, _ = hidden_states.shape mixed_qkv = gdn.in_proj_qkv(hidden_states) mixed_qkv = mixed_qkv.transpose(1, 2) z = gdn.in_proj_z(hidden_states) z = z.reshape(batch_size, seq_len, -1, gdn.head_v_dim) b = gdn.in_proj_b(hidden_states) a = gdn.in_proj_a(hidden_states) # Conv1d path (non-cached) if cache_params is not None: conv_state = F.pad(mixed_qkv, (gdn.conv_kernel_size - mixed_qkv.shape[-1], 0)) conv_state = cache_params.update_conv_state(conv_state, gdn.layer_idx) if gdn.causal_conv1d_fn is not None: mixed_qkv = gdn.causal_conv1d_fn( x=mixed_qkv, weight=gdn.conv1d.weight.squeeze(1), bias=gdn.conv1d.bias, activation=gdn.activation, seq_idx=None, ) else: mixed_qkv = F.silu(gdn.conv1d(mixed_qkv)[:, :, :seq_len]) mixed_qkv = mixed_qkv.transpose(1, 2) query, key, value = torch.split( mixed_qkv, [gdn.key_dim, gdn.key_dim, gdn.value_dim], dim=-1, ) query = query.reshape(batch_size, seq_len, -1, gdn.head_k_dim) key = key.reshape(batch_size, seq_len, -1, gdn.head_k_dim) value = value.reshape(batch_size, seq_len, -1, gdn.head_v_dim) beta = b.sigmoid() g = -gdn.A_log.float().exp() * F.softplus(a.float() + gdn.dt_bias) if gdn.num_v_heads // gdn.num_k_heads > 1: query = query.repeat_interleave(gdn.num_v_heads // gdn.num_k_heads, dim=2) key = key.repeat_interleave(gdn.num_v_heads // gdn.num_k_heads, dim=2) # Store captures (detach to avoid graph retention) self.k = key.detach() # (bs, seq_len, n_heads, head_k_dim) self.v = value.detach() # (bs, seq_len, n_heads, head_v_dim) self.beta = beta.detach() # (bs, seq_len, n_heads) self.g = g.detach() # (bs, seq_len, n_heads) core_attn_out, last_recurrent_state = gdn.chunk_gated_delta_rule( query, key, value, g=g, beta=beta, initial_state=None, output_final_state=cache_params is not None, use_qk_l2norm_in_kernel=True, ) if cache_params is not None: cache_params.update_recurrent_state(last_recurrent_state, gdn.layer_idx) core_attn_out = core_attn_out.reshape(-1, gdn.head_v_dim) z = z.reshape(-1, gdn.head_v_dim) core_attn_out = gdn.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(batch_size, seq_len, -1) return gdn.out_proj(core_attn_out) def remove(self): self._gdn.forward = self._original_forward # Simulate recurrence to get per-position states for one head @torch.no_grad() def simulate_recurrence( k: torch.Tensor, # (seq_len, head_k_dim) v: torch.Tensor, # (seq_len, head_v_dim) beta: torch.Tensor, # (seq_len,) or (seq_len, 1) g: torch.Tensor, # (seq_len,) dtype: torch.dtype = torch.float32, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Simulate the GDN delta rule recurrence. The actual update (from fla naive reference): S_t = exp(g_t) * S_{t-1} + outer(k_t, beta_t * (v_t - S_{t-1} @ k_t)) where k is L2-normalized (use_qk_l2norm_in_kernel=True). Returns: states: (seq_len, d_k, d_v) per-position recurrent states k_normed: (seq_len, d_k) L2-normalized key vectors (as used in the recurrence) write_v: (seq_len, d_v) the corrected value vectors actually written """ seq_len, d_k = k.shape d_v = v.shape[1] k = k.to(dtype) v = v.to(dtype) beta_flat = beta.reshape(seq_len).to(dtype) g_flat = g.reshape(seq_len).to(dtype) # L2-normalize k (matching use_qk_l2norm_in_kernel=True) k_normed = F.normalize(k, dim=-1) states = torch.zeros(seq_len, d_k, d_v, dtype=dtype, device=k.device) write_v = torch.zeros(seq_len, d_v, dtype=dtype, device=k.device) S = torch.zeros(d_k, d_v, dtype=dtype, device=k.device) for t in range(seq_len): # Decay S = g_flat[t].exp() * S # Delta rule correction: v_corrected = v_t - S @ k_t v_corrected = v[t] - S @ k_normed[t] # Scale by beta v_corrected = beta_flat[t] * v_corrected # Write S = S + torch.outer(k_normed[t], v_corrected) states[t] = S write_v[t] = v_corrected return states, k_normed, write_v # Main alignment computation @torch.no_grad() def compute_alignment( model, tokenizer, sae: BilinearMatrixSAE, layer_idx: int, head_idx: int, n_seqs: int = 100, seq_len: int = 512, batch_size: int = 4, top_n: int = 50, device: str = "cuda", ) -> dict: """Compute alignment between SAE decoder atoms and GDN k*v^T writes.""" print(f"Loading {n_seqs} sequences of length {seq_len}...") batches = load_corpus_batches(tokenizer, n_seqs, seq_len, batch_size) total_seqs = sum(b.shape[0] for b in batches) print(f"Loaded {total_seqs} sequences in {len(batches)} batches") # SAE decoder atoms: V_dec (n_features, rank, d_k), W_dec (n_features, rank, d_v) V_dec = sae.V_dec.detach().float().cpu() # (n_features, rank, d_k) W_dec = sae.W_dec.detach().float().cpu() # (n_features, rank, d_v) n_features = V_dec.shape[0] rank = V_dec.shape[1] d_k = V_dec.shape[2] d_v = W_dec.shape[2] print(f"SAE: {n_features} features, rank={rank}, d_k={d_k}, d_v={d_v}") # For rank-1, V_dec[:,0,:] and W_dec[:,0,:] are the decoder directions # Normalize them for cosine computation (on CPU) v_dec_dir = F.normalize(V_dec[:, 0, :], dim=-1) # (n_features, d_k) w_dec_dir = F.normalize(W_dec[:, 0, :], dim=-1) # (n_features, d_v) # Install hook hook = GDNWriteCapture(model, layer_idx) # Accumulators: for each feature, store (activation, k_cos, v_cos) across all positions # We'll collect all data first, then pick top-N per feature # To avoid OOM, we process batch by batch and keep only top-N per feature feature_data = {i: [] for i in range(n_features)} # feature_data[i] = list of (activation, k_cosine, v_cosine) t_start = time.time() for batch_idx, batch in enumerate(tqdm(batches, desc="Processing batches")): input_ids = batch.to(device) bs = input_ids.shape[0] # Forward pass (captures k, v, beta, g via hook) # use_cache=True to also get final recurrent state for verification outputs = model(input_ids=input_ids, use_cache=True) # hook.k: (bs, seq_len, n_heads, head_k_dim) # hook.v: (bs, seq_len, n_heads, head_v_dim) # hook.beta: (bs, seq_len, n_heads) or (bs, seq_len, n_heads, 1) # hook.g: (bs, seq_len, n_heads) k_all = hook.k # (bs, seq_len, n_heads, d_k) v_all = hook.v # (bs, seq_len, n_heads, d_v) beta_all = hook.beta # (bs, seq_len, n_heads[, 1]) g_all = hook.g # (bs, seq_len, n_heads) for seq_idx in range(bs): # Extract head-specific k, v, beta, g and move to CPU for simulation k_head = k_all[seq_idx, :, head_idx, :].cpu() # (seq_len, d_k) v_head = v_all[seq_idx, :, head_idx, :].cpu() # (seq_len, d_v) # beta shape handling if beta_all.ndim == 4: beta_head = beta_all[seq_idx, :, head_idx, 0].cpu() # (seq_len,) else: beta_head = beta_all[seq_idx, :, head_idx].cpu() # (seq_len,) g_head = g_all[seq_idx, :, head_idx].cpu() # (seq_len,) # Simulate recurrence to get per-position states and actual write vectors states, k_normed, write_v = simulate_recurrence( k_head, v_head, beta_head, g_head, dtype=torch.float32, ) # states: (seq_len, d_k, d_v) # k_normed: (seq_len, d_k) -- L2-normalized keys used in recurrence # write_v: (seq_len, d_v) -- corrected value vectors actually written # Verify simulated final state matches actual cache (first seq only) if batch_idx == 0 and seq_idx == 0: cache = outputs.past_key_values actual_state = cache.layers[layer_idx].recurrent_states[0, head_idx].float().cpu() sim_final = states[-1].cpu() cos_sim = F.cosine_similarity( actual_state.reshape(1, -1), sim_final.reshape(1, -1) ).item() rel_err = (actual_state - sim_final).norm() / actual_state.norm() print(f"Recurrence verification: cos_sim={cos_sim:.6f}, rel_error={rel_err:.6f}") if cos_sim < 0.9: print("WARNING: simulated recurrence does not match actual state well") # Encode states through SAE in sub-batches to get activations sub_bs = 128 all_coeffs = [] for s in range(0, states.shape[0], sub_bs): e = min(s + sub_bs, states.shape[0]) chunk = states[s:e].to(device) coeffs = sae.encode(chunk) # (chunk_size, n_features) all_coeffs.append(coeffs.cpu()) all_coeffs = torch.cat(all_coeffs, dim=0) # (seq_len, n_features) # Normalize write vectors for cosine computation # k is already L2-normed; normalize write_v (use eps to avoid NaN for zero writes) k_dir = k_normed.cpu() # (seq_len, d_k) already unit norm wv_raw = write_v.cpu() wv_norms = wv_raw.norm(dim=-1, keepdim=True).clamp(min=1e-8) wv_dir = wv_raw / wv_norms # (seq_len, d_v) # Vectorized: compute cosines for ALL features at ALL positions at once # v_dec_dir: (n_features, d_k), k_dir: (seq_len, d_k) # k_cos_all: (seq_len, n_features) k_cos_all = k_dir @ v_dec_dir.t() # (seq_len, n_features) v_cos_all = wv_dir @ w_dec_dir.t() # (seq_len, n_features) # Collect active (position, feature) pairs efficiently # all_coeffs: (seq_len, n_features), sparse (k=32 active per position) active_pos, active_feat = (all_coeffs > 0).nonzero(as_tuple=True) if active_pos.numel() > 0: active_vals = all_coeffs[active_pos, active_feat].numpy() active_k_cos = k_cos_all[active_pos, active_feat].numpy() active_v_cos = v_cos_all[active_pos, active_feat].numpy() feat_np = active_feat.numpy() # Group by feature for efficient insertion for fi in np.unique(feat_np): mask = feat_np == fi entries = list(zip( active_vals[mask].tolist(), active_k_cos[mask].tolist(), active_v_cos[mask].tolist(), )) feature_data[int(fi)].extend(entries) # Clear GPU memory del k_all, v_all, beta_all, g_all, outputs if device == "cuda": torch.cuda.empty_cache() hook.remove() elapsed = time.time() - t_start print(f"Data collection took {elapsed:.1f}s") # For each feature, pick top-N activations and compute mean alignment print("Computing per-feature alignment scores...") results = [] alive_count = 0 for feat_idx in range(n_features): data = feature_data[feat_idx] if len(data) == 0: results.append({ "feature": feat_idx, "alive": False, "mean_k_cos": 0.0, "mean_v_cos": 0.0, "n_activations": 0, }) continue alive_count += 1 data.sort(key=lambda x: x[0], reverse=True) top_data = data[:top_n] k_cosines = [d[1] for d in top_data] v_cosines = [d[2] for d in top_data] mean_k = float(np.mean(k_cosines)) mean_v = float(np.mean(v_cosines)) results.append({ "feature": feat_idx, "alive": True, "mean_k_cos": mean_k, "mean_v_cos": mean_v, "mean_abs_k_cos": float(np.mean(np.abs(k_cosines))), "mean_abs_v_cos": float(np.mean(np.abs(v_cosines))), "n_activations": len(data), "top_n_used": len(top_data), }) return { "results": results, "alive_count": alive_count, "n_features": n_features, "elapsed_s": elapsed, "n_seqs": total_seqs, "seq_len": seq_len, } # Reporting def report(output: dict, output_dir: str | None = None): results = output["results"] alive_count = output["alive_count"] n_features = output["n_features"] alive = [r for r in results if r["alive"]] dead_count = n_features - alive_count print(f"\n{'='*70}") print("Memory Slot Surgery: Alignment Report") print(f"{'='*70}") print(f"Features: {n_features} total, {alive_count} alive, {dead_count} dead") print(f"Sequences: {output['n_seqs']} x {output['seq_len']} tokens") print(f"Time: {output['elapsed_s']:.1f}s") if not alive: print("No alive features found.") return # Use absolute cosines (direction shouldn't matter for alignment) k_abs = np.array([r["mean_abs_k_cos"] for r in alive]) v_abs = np.array([r["mean_abs_v_cos"] for r in alive]) # Combined alignment = geometric mean of |k_cos| and |v_cos| combined = np.sqrt(k_abs * v_abs) print(f"\n--- Alignment Statistics (absolute cosines, across {len(alive)} alive features) ---") for name, arr in [("key alignment |cos|", k_abs), ("value alignment |cos|", v_abs), ("combined (geometric mean)", combined)]: print(f" {name}:") print(f" mean={arr.mean():.4f} median={np.median(arr):.4f} " f"std={arr.std():.4f} max={arr.max():.4f}") # Null baseline: expected |cosine| between random unit vectors in d dimensions # E[|cos|] = sqrt(2/pi) / sqrt(d) for d-dim random vectors d = 128 null_expected = np.sqrt(2 / np.pi) / np.sqrt(d) print("\n--- Null Baseline ---") print(f" Random 128-dim vectors: E[|cos|] = {null_expected:.4f}") print(f" Random combined: E[geometric mean] = {null_expected:.4f}") # Fraction above thresholds print("\n--- Fraction of alive features above alignment thresholds ---") for thresh in [0.1, 0.2, 0.3, 0.5]: k_frac = (k_abs > thresh).mean() v_frac = (v_abs > thresh).mean() c_frac = (combined > thresh).mean() print(f" >{thresh:.1f}: key={k_frac:.3f} value={v_frac:.3f} combined={c_frac:.3f}") # Top-10 by combined alignment sorted_alive = sorted(alive, key=lambda r: np.sqrt(r["mean_abs_k_cos"] * r["mean_abs_v_cos"]), reverse=True) print("\n--- Top-10 Features by Combined Alignment ---") print(f" {'Feature':>8} {'|k_cos|':>8} {'|v_cos|':>8} {'combined':>8} {'n_act':>6}") for r in sorted_alive[:10]: c = np.sqrt(r["mean_abs_k_cos"] * r["mean_abs_v_cos"]) print(f" {r['feature']:>8} {r['mean_abs_k_cos']:>8.4f} {r['mean_abs_v_cos']:>8.4f} " f"{c:>8.4f} {r['n_activations']:>6}") # Histogram (text-based) print("\n--- Key Alignment Histogram ---") _text_histogram(k_abs, "key |cos|") print("\n--- Value Alignment Histogram ---") _text_histogram(v_abs, "value |cos|") if output_dir: os.makedirs(output_dir, exist_ok=True) with open(os.path.join(output_dir, "alignment_results.json"), "w") as f: json.dump(output, f, indent=2) print(f"\nResults saved to {output_dir}/alignment_results.json") np.savez( os.path.join(output_dir, "alignment_histograms.npz"), k_abs=k_abs, v_abs=v_abs, combined=combined, ) def _text_histogram(values: np.ndarray, label: str, bins: int = 20): counts, edges = np.histogram(values, bins=bins, range=(0, 1)) max_count = max(counts) if max(counts) > 0 else 1 for i, c in enumerate(counts): bar = "#" * int(40 * c / max_count) print(f" [{edges[i]:.2f}-{edges[i+1]:.2f}) {c:>4} {bar}") # Entry point def main(): import argparse p = argparse.ArgumentParser(description="Memory Slot Surgery alignment check") p.add_argument("--model", default="Qwen/Qwen3.5-0.8B") p.add_argument("--sae-checkpoint", required=True, help="Path to bilinear SAE checkpoint") p.add_argument("--layer", type=int, default=9) p.add_argument("--head", type=int, default=4) p.add_argument("--n-seqs", type=int, default=100) p.add_argument("--seq-len", type=int, default=512) p.add_argument("--batch-size", type=int, default=4) p.add_argument("--top-n", type=int, default=50) p.add_argument("--output-dir", default=None) p.add_argument("--device", default="cuda") args = p.parse_args() model, tokenizer = load_model(args.model, args.device) sae = load_sae(args.sae_checkpoint, args.device) output = compute_alignment( model, tokenizer, sae, layer_idx=args.layer, head_idx=args.head, n_seqs=args.n_seqs, seq_len=args.seq_len, batch_size=args.batch_size, top_n=args.top_n, device=args.device, ) report(output, output_dir=args.output_dir) if __name__ == "__main__": main()