"""Learn four dictionary directions from synthetic 3D activations on CPU. Run without flags to print metrics and check the geometry. --export writes the article's JSON. The original house-price experiment remains in ../experiment/. """ import argparse import itertools import json from pathlib import Path import torch from torch import nn SEEDS = [7, 11, 23, 42, 101] STEPS = 2000 TRAIN_SIZE = 4096 TEST_SIZE = 2048 DICT = torch.cat((torch.eye(3), torch.ones(3, 1) / 3**0.5), dim=1) def dataset(n, active, seed): generator = torch.Generator().manual_seed(seed) # Select a uniformly random subset of exactly `active` features per example. scores = torch.rand(n, 4, generator=generator) support = torch.zeros_like(scores).scatter_(1, scores.topk(active, dim=1).indices, 1.0) features = support * (0.5 + torch.rand(n, 4, generator=generator)) return features @ DICT.T, features class SAE(nn.Module): def __init__(self, top_one): super().__init__() self.top_one = top_one self.encoder = nn.Linear(3, 4, bias=False) self.decoder = nn.Linear(4, 3, bias=False) # A simple positive initialization for this positive-octant dataset. # It uses neither generating directions nor per-example feature labels. with torch.no_grad(): self.decoder.weight.copy_(torch.rand(3, 4)) self.normalize_decoder() self.encoder.weight.copy_(self.decoder.weight.T) def normalize_decoder(self): self.decoder.weight.div_(self.decoder.weight.norm(dim=0).clamp_min(1e-8)) def forward(self, x): code = self.encoder(x).relu() if self.top_one: winner = code.argmax(dim=1, keepdim=True) code = code * torch.zeros_like(code).scatter_(1, winner, 1.0) return self.decoder(code), code def geometry_checks(): raw = torch.cat((torch.eye(3), torch.ones(3, 1)), dim=1) assert torch.linalg.matrix_rank(raw).item() == 3 assert torch.equal(raw @ torch.tensor([0., 0., 0., 1.]), raw @ torch.tensor([1., 1., 1., 0.])) assert torch.equal(raw @ torch.tensor([-1., -1., -1., 1.]), torch.zeros(3)) projection = torch.tensor([[1., 0., 1.], [0., 1., 1.]]) assert torch.equal(projection @ torch.tensor([2., 3., 4.]), torch.tensor([6., 7.])) assert torch.equal(projection @ torch.tensor([1., 1., -1.]), torch.zeros(2)) # Projecting on D also responds to A, although D's coefficient is zero. assert torch.isclose(DICT[:, 3] @ DICT[:, 0], torch.tensor(1 / 3**0.5)) def run(active, top_one, seed): train, _ = dataset(TRAIN_SIZE, active, 100 + active) test, truth = dataset(TEST_SIZE, active, 200 + active) torch.manual_seed(seed) model = SAE(top_one) initial_decoder = model.decoder.weight.detach().clone() optimizer = torch.optim.Adam(model.parameters(), lr=0.01) for _ in range(STEPS): indices = torch.randint(len(train), (256,)) batch = train[indices] reconstruction, _ = model(batch) loss = (reconstruction - batch).square().sum(dim=1).mean() optimizer.zero_grad() loss.backward() optimizer.step() with torch.no_grad(): model.normalize_decoder() with torch.no_grad(): reconstruction, code = model(test) decoder = model.decoder.weight # Feature numbers are arbitrary. Match one-to-one, with no sign flip: # codes are nonnegative. This oracle comparison is evaluation only. cosines = DICT.T @ decoder permutation = max(itertools.permutations(range(4)), key=lambda p: sum(float(cosines[j, p[j]]) for j in range(4))) matched_code = code[:, list(permutation)] baseline_error = (test - train.mean(dim=0)).square().sum() normalized_error = float((reconstruction - test).square().sum() / baseline_error) coefficient_rmse = float((matched_code - truth).square().mean().sqrt()) matched_cosines = [float(cosines[j, permutation[j]]) for j in range(4)] result = { "seed": seed, "normalized_reconstruction_error": normalized_error, "mean_matched_cosine": sum(matched_cosines) / 4, "matched_cosines": matched_cosines, "coefficient_rmse": coefficient_rmse, "mean_active_code_entries": float((code > 0).float().sum(dim=1).mean()), "exact_support_accuracy": float(((matched_code > 0) == (truth > 0)).all(dim=1).float().mean()), "activation_counts": (code > 0).sum(dim=0).tolist(), "matched_indices": list(permutation), "initial_decoder": initial_decoder.tolist(), "encoder": model.encoder.weight.tolist(), "decoder": decoder.tolist(), } assert torch.isfinite(reconstruction).all() assert torch.allclose(decoder.norm(dim=0), torch.ones(4), atol=1e-5) if top_one: assert ((code > 0).sum(dim=1) <= 1).all() return result def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--export", action="store_true") args = parser.parse_args() torch.set_num_threads(1) torch.use_deterministic_algorithms(True) geometry_checks() configurations = [ (1, True, "top1-one", "Top-1", "One"), (1, False, "relu-one", "ReLU, no sparsity constraint", "One"), (2, True, "top1-two", "Top-1", "Two"), (4, True, "top1-four", "Top-1", "Four"), ] groups = [] for active, top_one, key, label, activity in configurations: runs = [run(active, top_one, seed) for seed in SEEDS] metrics = ["normalized_reconstruction_error", "mean_matched_cosine", "coefficient_rmse", "mean_active_code_entries", "exact_support_accuracy"] summary = {metric: { "mean": sum(r[metric] for r in runs) / len(runs), "min": min(r[metric] for r in runs), "max": max(r[metric] for r in runs), } for metric in metrics} groups.append({"key": key, "label": label, "active_features": active, "activity": activity, "summary": summary, "runs": runs}) print(key, json.dumps(summary), flush=True) result = { "protocol": {"torch": torch.__version__, "dtype": "float32", "device": "cpu", "seeds": SEEDS, "steps": STEPS, "batch_size": 256, "train_size": TRAIN_SIZE, "test_size": TEST_SIZE, "optimizer": "Adam", "learning_rate": 0.01, "feature_strength": [0.5, 1.5], "train_seed": "100 + active_features", "test_seed": "200 + active_features"}, "dictionary": DICT.tolist(), "groups": groups, } if args.export: path = Path(__file__).resolve().parent.parent / "data" / "sparse-results.json" path.write_text(json.dumps(result, indent=2) + "\n") print(f"Wrote {path}") if __name__ == "__main__": main()