{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Mechanistic Interpretability of Grokking in MLPs\n", "\n", "Combined analysis notebook for the paper *Mechanistic Interpretability of Grokking in MLPs* (Paul Dubrulle, ISYE 4803).\n", "\n", "This file merges two notebooks — `MI.ipynb` (main analysis) and `loss_landscape.ipynb` (loss-landscape analysis). Cell **outputs were cleared** to keep the file lightweight; re-run top to bottom to reproduce the figures.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "---\n", "\n", "## Part 1 — Main analysis\n", "\n", "*Source: `MI.ipynb`*\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import polars as pl\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import time\n", "from pathlib import Path\n", "\n", "import torch\n", "import torch.nn as nn\n", "import torch.optim as optim\n", "import torch.nn.functional as F\n", "\n", "from IPython.display import clear_output, HTML, display\n", "\n", "torch.cuda.is_available()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\"\"\"\n", "setup a 2-layer nn with one hot encoded inputs with the mod 113 dataset for modular addition and try to recreate the grokking phenomenon. Potentially set up live graphs with weight and biases if thats\n", "the only way.\n", "\"\"\"\n", "#bias=False to match gromov/doshi closed form, which has no additive biases (only W1, W2 and the activation)\n", "class baseMLP(nn.Module):\n", " def __init__(self, p: int, hidden_dim: int):\n", " super().__init__()\n", " self.fc1 = nn.Linear(2*p, hidden_dim, bias=False) #one hot encoded inputs, p inputs for input a and b\n", " self.relu = nn.ReLU()\n", " self.fc2 = nn.Linear(hidden_dim, p, bias=False)\n", "\n", " def forward(self, x):\n", " return self.fc2(self.relu(self.fc1(x)))\n", " " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "p = 113\n", "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#dataset for any of the 4 modular ops mod p\n", "#add/sub use all p^2 pairs, mul/div drop b=0 since it has no inverse so p*(p-1) pairs\n", "\n", "def make_dataset(operation, p=113, train_frac=0.5, seed=0, device='cuda'):\n", " if operation in ('add', 'sub'):\n", " a = torch.arange(p).repeat_interleave(p)\n", " b = torch.arange(p).repeat(p)\n", " elif operation in ('mul', 'div'):\n", " a = torch.arange(p).repeat_interleave(p - 1)\n", " b = torch.arange(1, p).repeat(p)\n", " else:\n", " raise ValueError(f\"unknown operation: {operation}\")\n", "\n", " if operation == 'add':\n", " y_labels = (a + b) % p\n", " elif operation == 'sub':\n", " y_labels = (a - b) % p\n", " elif operation == 'mul':\n", " y_labels = (a * b) % p\n", " elif operation == 'div':\n", " b_inv = torch.tensor([pow(int(bi), p - 2, p) for bi in b]) #fermat: b^(p-2) mod p, valid since p prime\n", " y_labels = (a * b_inv) % p\n", "\n", " x = torch.cat([F.one_hot(a, p), F.one_hot(b, p)], dim=1).float() #shape (n_pairs, 2p)\n", " y = F.one_hot(y_labels, p).float() #shape (n_pairs, p)\n", "\n", " g = torch.Generator().manual_seed(seed)\n", " perm = torch.randperm(x.size(0), generator=g)\n", " n_train = int(round(x.size(0) * train_frac))\n", "\n", " return {\n", " 'x_train': x[perm[:n_train]].to(device),\n", " 'y_train': y[perm[:n_train]].to(device),\n", " 'x_test': x[perm[n_train:]].to(device),\n", " 'y_test': y[perm[n_train:]].to(device),\n", " 'operation': operation,\n", " }" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#discrete log basis for mul/div: doshi shows the analytical solution is periodic only after reindexing inputs by their discrete log under a primitive root g\n", "#smallest primitive root mod 113 is g=3\n", "#IPR_2 below matches gromov (sum of squared normalized power) and doshi\n", "\n", "def find_primitive_root(p):\n", " if p == 2:\n", " return 1\n", " n = p - 1\n", " factors, d = [], 2\n", " while d * d <= n:\n", " if n % d == 0:\n", " factors.append(d)\n", " while n % d == 0:\n", " n //= d\n", " d += 1\n", " if n > 1:\n", " factors.append(n)\n", " for g in range(2, p):\n", " if all(pow(g, (p - 1) // f, p) != 1 for f in factors):\n", " return g\n", " raise RuntimeError(f\"no primitive root for p={p}\")\n", "\n", "def log_basis_perm(p, g):\n", " #pi[k] = g^k mod p for k in 0..p-2, drops index 0 (log undefined at 0)\n", " pi = torch.empty(p - 1, dtype=torch.long)\n", " val = 1\n", " for k in range(p - 1):\n", " pi[k] = val\n", " val = (val * g) % p\n", " return pi\n", "\n", "def to_log_basis(weight_block, pi_log):\n", " #reorder a (..., p) tensor along last axis to (..., p-1) in log basis\n", " return weight_block.index_select(-1, pi_log)\n", "\n", "def compute_ipr(weight_matrix, dim=-1):\n", " #IPR_2 along given axis, returns scalar mean across leading axes\n", " W_freq = torch.fft.fft(weight_matrix, dim=dim)\n", " power = (W_freq.abs() ** 2)\n", " power_normalized = power / power.sum(dim=dim, keepdim=True).clamp(min=1e-30)\n", " ipr_per_row = (power_normalized ** 2).sum(dim=dim)\n", " return ipr_per_row.mean().item()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#single training function, replaces the 7 duplicated training cells\n", "#full-batch AdamW and MSE on one-hot, IPR tracked on fc1 a-block, fc1 b-block, fc2\n", "#fc2 is transposed to (hidden_dim, p) before IPR so the FFT runs along the output axis q (gromov: W2 is periodic in q, not k)\n", "#for mul/div, IPR is also tracked in the discrete-log basis on all three blocks\n", "\n", "def train_run(operation, p=113, hidden_dim=128, lr=1e-3, weight_decay=1.0,\n", " train_frac=0.3, epochs=50_000, seed=0, snapshot_every=500, log_every=50,\n", " device=None, live_plot=True, g=None):\n", " if device is None:\n", " device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n", " if g is None:\n", " g = find_primitive_root(p)\n", "\n", " torch.manual_seed(seed)\n", " np.random.seed(seed)\n", "\n", " data = make_dataset(operation, p=p, train_frac=train_frac, seed=seed, device=device)\n", " pi_log = log_basis_perm(p, g)\n", "\n", " model = baseMLP(p, hidden_dim).to(device)\n", " optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)\n", " criterion = nn.MSELoss()\n", "\n", " train_losses, test_losses = [], []\n", " train_accs, test_accs = [], []\n", " ipr_history = []\n", " weight_snapshots = []\n", "\n", " y_train_idx = data['y_train'].argmax(dim=1)\n", " y_test_idx = data['y_test'].argmax(dim=1)\n", "\n", " t_start = time.time()\n", " for epoch in range(epochs):\n", " model.train()\n", " optimizer.zero_grad()\n", " loss = criterion(model(data['x_train']), data['y_train'])\n", " loss.backward()\n", " optimizer.step()\n", "\n", " if epoch % log_every == 0 or epoch == epochs - 1:\n", " model.eval()\n", " with torch.no_grad():\n", " train_logits = model(data['x_train'])\n", " train_loss = criterion(train_logits, data['y_train']).item()\n", " train_acc = (train_logits.argmax(dim=1) == y_train_idx).float().mean().item()\n", "\n", " test_logits = model(data['x_test'])\n", " test_loss = criterion(test_logits, data['y_test']).item()\n", " test_acc = (test_logits.argmax(dim=1) == y_test_idx).float().mean().item()\n", "\n", " w1 = model.fc1.weight.data.cpu() #(hidden_dim, 2p)\n", " w2 = model.fc2.weight.data.cpu() #(p, hidden_dim), so transpose to put output axis q last\n", " w2_t = w2.transpose(0, 1).contiguous() #(hidden_dim, p)\n", " w1_a, w1_b = w1[:, :p], w1[:, p:]\n", "\n", " ipr_entry = {\n", " 'epoch': epoch,\n", " 'fc1_a': compute_ipr(w1_a),\n", " 'fc1_b': compute_ipr(w1_b),\n", " 'fc2': compute_ipr(w2_t),\n", " }\n", " if operation in ('mul', 'div'):\n", " ipr_entry['fc1_a_log'] = compute_ipr(to_log_basis(w1_a, pi_log))\n", " ipr_entry['fc1_b_log'] = compute_ipr(to_log_basis(w1_b, pi_log))\n", " ipr_entry['fc2_log'] = compute_ipr(to_log_basis(w2_t, pi_log))\n", "\n", " train_losses.append(train_loss)\n", " test_losses.append(test_loss)\n", " train_accs.append(train_acc)\n", " test_accs.append(test_acc)\n", " ipr_history.append(ipr_entry)\n", "\n", " if live_plot:\n", " xs = [e['epoch'] for e in ipr_history]\n", " clear_output(wait=True)\n", " fig, axes = plt.subplots(1, 3, figsize=(15, 4))\n", "\n", " axes[0].plot(xs, train_accs, label='train')\n", " axes[0].plot(xs, test_accs, label='test')\n", " axes[0].set_xlabel('Epoch')\n", " axes[0].set_ylabel('Accuracy')\n", " axes[0].set_title(f'{operation} accuracy')\n", " axes[0].legend()\n", "\n", " axes[1].plot(xs, train_losses, label='train')\n", " axes[1].plot(xs, test_losses, label='test')\n", " axes[1].set_xlabel('Epoch')\n", " axes[1].set_ylabel('MSE loss')\n", " axes[1].set_yscale('log')\n", " axes[1].set_title(f'{operation} loss')\n", " axes[1].legend()\n", "\n", " axes[2].plot(xs, [e['fc1_a'] for e in ipr_history], label='fc1 a')\n", " axes[2].plot(xs, [e['fc1_b'] for e in ipr_history], label='fc1 b')\n", " axes[2].plot(xs, [e['fc2'] for e in ipr_history], label='fc2')\n", " if operation in ('mul', 'div'):\n", " axes[2].plot(xs, [e['fc1_a_log'] for e in ipr_history], '--', label='fc1 a log')\n", " axes[2].plot(xs, [e['fc1_b_log'] for e in ipr_history], '--', label='fc1 b log')\n", " axes[2].plot(xs, [e['fc2_log'] for e in ipr_history], '--', label='fc2 log')\n", " axes[2].set_xlabel('Epoch')\n", " axes[2].set_ylabel('IPR_2')\n", " axes[2].set_title(f'{operation} IPR (Fourier)')\n", " axes[2].legend()\n", "\n", " plt.tight_layout()\n", " plt.show()\n", "\n", " elapsed = time.time() - t_start\n", " epochs_per_sec = (epoch + 1) / max(elapsed, 1e-6)\n", " eta = (epochs - epoch - 1) / max(epochs_per_sec, 1e-6)\n", " print(f\"Epoch {epoch}: train_acc={train_acc:.4f}, test_acc={test_acc:.4f}, IPR_fc1_a={ipr_entry['fc1_a']:.4f}, IPR_fc2={ipr_entry['fc2']:.4f}\")\n", " print(f\"Time: {elapsed:.1f}s elapsed | ETA: {eta:.0f}s | {epochs_per_sec:.1f} epochs/sec\")\n", "\n", " if epoch % snapshot_every == 0 or epoch == epochs - 1:\n", " weight_snapshots.append({\n", " 'epoch': epoch,\n", " 'fc1': model.fc1.weight.data.cpu().clone(),\n", " 'fc2': model.fc2.weight.data.cpu().clone(),\n", " })\n", "\n", " return {\n", " 'config': {\n", " 'operation': operation, 'p': p, 'hidden_dim': hidden_dim, 'lr': lr,\n", " 'weight_decay': weight_decay, 'train_frac': train_frac, 'epochs': epochs,\n", " 'seed': seed, 'g': g,\n", " },\n", " 'train_losses': train_losses,\n", " 'test_losses': test_losses,\n", " 'train_accs': train_accs,\n", " 'test_accs': test_accs,\n", " 'ipr_history': ipr_history,\n", " 'weight_snapshots': weight_snapshots,\n", " 'pi_log': pi_log,\n", " }" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#three viz functions, each takes a result dict from train_run and displays an animation\n", "#mul/div automatically use discrete-log basis\n", "\n", "from matplotlib.animation import FuncAnimation\n", "import matplotlib as mpl\n", "mpl.rcParams['animation.embed_limit'] = 100 #MB, default is 20\n", "\n", "def viz_fourier_spectrum(result, max_frames=80):\n", " snapshots = result['weight_snapshots']\n", " op = result['config']['operation']\n", " p = result['config']['p']\n", " pi_log = result['pi_log']\n", " use_log = op in ('mul', 'div')\n", "\n", " stride = max(1, len(snapshots) // max_frames)\n", " frames = list(range(0, len(snapshots), stride))\n", " if frames[-1] != len(snapshots) - 1:\n", " frames.append(len(snapshots) - 1)\n", "\n", " spectra_a, spectra_b = [], []\n", " for i in frames:\n", " w1 = snapshots[i]['fc1']\n", " w1_a, w1_b = w1[:, :p], w1[:, p:]\n", " if use_log:\n", " w1_a = to_log_basis(w1_a, pi_log)\n", " w1_b = to_log_basis(w1_b, pi_log)\n", " sa = torch.fft.fft(w1_a, dim=1).abs() ** 2\n", " sb = torch.fft.fft(w1_b, dim=1).abs() ** 2\n", " sa = sa / sa.sum(dim=1, keepdim=True).clamp(min=1e-30)\n", " sb = sb / sb.sum(dim=1, keepdim=True).clamp(min=1e-30)\n", " half = sa.size(1) // 2 #DFT is symmetric for real input, only first half is informative\n", " spectra_a.append(sa[:, :half].numpy())\n", " spectra_b.append(sb[:, :half].numpy())\n", "\n", " sort_idx = np.argsort(spectra_a[-1].argmax(axis=1))\n", " basis_label = 'log basis' if use_log else 'natural basis'\n", " fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n", "\n", " im1 = ax1.imshow(spectra_a[0][sort_idx], aspect='auto', cmap='hot', vmin=0, vmax=0.3)\n", " ax1.set_xlabel('Frequency')\n", " ax1.set_ylabel('Neuron (sorted)')\n", " ax1.set_title(f'fc1 a-block spectrum, {basis_label}')\n", " fig.colorbar(im1, ax=ax1)\n", "\n", " im2 = ax2.imshow(spectra_b[0][sort_idx], aspect='auto', cmap='hot', vmin=0, vmax=0.3)\n", " ax2.set_xlabel('Frequency')\n", " ax2.set_ylabel('Neuron (sorted)')\n", " ax2.set_title(f'fc1 b-block spectrum, {basis_label}')\n", " fig.colorbar(im2, ax=ax2)\n", "\n", " title = fig.suptitle(f'{op}, epoch {snapshots[frames[0]][\"epoch\"]}')\n", "\n", " def update(k):\n", " f = frames[k]\n", " im1.set_data(spectra_a[k][sort_idx])\n", " im2.set_data(spectra_b[k][sort_idx])\n", " title.set_text(f'{op}, epoch {snapshots[f][\"epoch\"]}')\n", " return im1, im2, title\n", "\n", " anim = FuncAnimation(fig, update, frames=len(frames), interval=100, blit=False)\n", " plt.close()\n", " display(HTML(anim.to_jshtml()))\n", "\n", "\n", "def viz_raw_weights(result, max_frames=80):\n", " snapshots = result['weight_snapshots']\n", " op = result['config']['operation']\n", " p = result['config']['p']\n", "\n", " stride = max(1, len(snapshots) // max_frames)\n", " frames = list(range(0, len(snapshots), stride))\n", " if frames[-1] != len(snapshots) - 1:\n", " frames.append(len(snapshots) - 1)\n", "\n", " #sort neurons by final fft peak frequency on the a-block\n", " final_w1_a = snapshots[-1]['fc1'][:, :p]\n", " final_spec = (torch.fft.fft(final_w1_a, dim=1).abs() ** 2).numpy()\n", " sort_idx = np.argsort(final_spec[:, 1:p // 2].argmax(axis=1))\n", "\n", " fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n", " w0 = snapshots[frames[0]]['fc1']\n", " im1 = ax1.imshow(w0[:, :p].numpy()[sort_idx], aspect='auto', cmap='RdBu', vmin=-1, vmax=1)\n", " ax1.set_xlabel('Input index a')\n", " ax1.set_ylabel('Neuron (sorted)')\n", " ax1.set_title(f'{op} fc1 a-block raw weights')\n", "\n", " im2 = ax2.imshow(w0[:, p:].numpy()[sort_idx], aspect='auto', cmap='RdBu', vmin=-1, vmax=1)\n", " ax2.set_xlabel('Input index b')\n", " ax2.set_ylabel('Neuron (sorted)')\n", " ax2.set_title(f'{op} fc1 b-block raw weights')\n", "\n", " title = fig.suptitle(f'{op}, epoch {snapshots[frames[0]][\"epoch\"]}')\n", "\n", " def update(k):\n", " f = frames[k]\n", " w = snapshots[f]['fc1']\n", " im1.set_data(w[:, :p].numpy()[sort_idx])\n", " im2.set_data(w[:, p:].numpy()[sort_idx])\n", " title.set_text(f'{op}, epoch {snapshots[f][\"epoch\"]}')\n", " return im1, im2, title\n", "\n", " anim = FuncAnimation(fig, update, frames=len(frames), interval=100, blit=False)\n", " plt.close()\n", " display(HTML(anim.to_jshtml()))\n", "\n", "\n", "def viz_nanda_2d(result, max_frames=60):\n", " snapshots = result['weight_snapshots']\n", " op = result['config']['operation']\n", " p = result['config']['p']\n", " hidden_dim = result['config']['hidden_dim']\n", " pi_log = result['pi_log']\n", " device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n", " use_log = op in ('mul', 'div')\n", "\n", " if op in ('add', 'sub'):\n", " a = torch.arange(p).repeat_interleave(p).to(device)\n", " b = torch.arange(p).repeat(p).to(device)\n", " n_a, n_b = p, p\n", " else:\n", " a = torch.arange(p).repeat_interleave(p - 1).to(device)\n", " b = torch.arange(1, p).repeat(p).to(device)\n", " n_a, n_b = p, p - 1\n", " x_all = torch.cat([F.one_hot(a, p), F.one_hot(b, p)], dim=1).float()\n", "\n", " if op == 'add':\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p)\n", " bi = torch.arange(p).unsqueeze(0).expand(p, p)\n", " cls = (ai + bi) % p\n", " elif op == 'sub':\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p)\n", " bi = torch.arange(p).unsqueeze(0).expand(p, p)\n", " cls = (ai - bi) % p\n", " elif op == 'mul':\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p - 1)\n", " bi = torch.arange(1, p).unsqueeze(0).expand(p, p - 1)\n", " cls = (ai * bi) % p\n", " elif op == 'div':\n", " b_inv_row = torch.tensor([pow(int(bi), p - 2, p) for bi in range(1, p)])\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p - 1)\n", " bi = b_inv_row.unsqueeze(0).expand(p, p - 1)\n", " cls = (ai * bi) % p\n", "\n", " work_model = baseMLP(p, hidden_dim).to(device)\n", "\n", " stride = max(1, len(snapshots) // max_frames)\n", " frames = list(range(0, len(snapshots), stride))\n", " if frames[-1] != len(snapshots) - 1:\n", " frames.append(len(snapshots) - 1)\n", "\n", " logit_imgs, spec_imgs = [], []\n", " for i in frames:\n", " snap = snapshots[i]\n", " work_model.fc1.weight.data = snap['fc1'].to(device)\n", " work_model.fc2.weight.data = snap['fc2'].to(device)\n", " work_model.eval()\n", " with torch.no_grad():\n", " logits = work_model(x_all).reshape(n_a, n_b, p).cpu()\n", " correct_logit = logits.gather(2, cls.unsqueeze(2)).squeeze(2) #(n_a, n_b)\n", " if use_log:\n", " #a-axis is 0..p-1, reorder so position k holds value g^k. drops a=0 row.\n", " #b-axis stores values 1..p-1 at indices 0..p-2, reorder by pi_log[k]-1.\n", " correct_logit = correct_logit.index_select(0, pi_log).index_select(1, pi_log - 1)\n", " spectrum = torch.fft.fft2(correct_logit).abs() ** 2\n", " spectrum[0, 0] = 0\n", " spectrum = spectrum / spectrum.sum().clamp(min=1e-30)\n", " spectrum = torch.fft.fftshift(spectrum)\n", " logit_imgs.append(correct_logit.numpy())\n", " spec_imgs.append(spectrum.numpy())\n", "\n", " basis_label = 'log basis' if use_log else 'natural basis'\n", " fig, axes = plt.subplots(1, 2, figsize=(13, 5))\n", "\n", " im1 = axes[0].imshow(logit_imgs[0], aspect='equal', cmap='RdBu')\n", " axes[0].set_xlabel('b')\n", " axes[0].set_ylabel('a')\n", " axes[0].set_title(f'{op} logit at correct class, {basis_label}')\n", " fig.colorbar(im1, ax=axes[0])\n", "\n", " H, W = spec_imgs[0].shape\n", " freq_a = np.arange(H) - (H // 2)\n", " freq_b = np.arange(W) - (W // 2)\n", " extent = (freq_b[0] - 0.5, freq_b[-1] + 0.5, freq_a[-1] + 0.5, freq_a[0] - 0.5)\n", " im2 = axes[1].imshow(spec_imgs[0], aspect='equal', cmap='hot', vmin=0, vmax=0.05, extent=extent)\n", " axes[1].set_xlabel('Frequency b (centered at DC)')\n", " axes[1].set_ylabel('Frequency a (centered at DC)')\n", " axes[1].set_title(f'{op} 2D Fourier of correct logit, {basis_label}')\n", " axes[1].axhline(0, color='cyan', linewidth=0.5, alpha=0.4)\n", " axes[1].axvline(0, color='cyan', linewidth=0.5, alpha=0.4)\n", " fig.colorbar(im2, ax=axes[1])\n", "\n", " title = fig.suptitle(f'{op}, epoch {snapshots[frames[0]][\"epoch\"]}')\n", "\n", " def update(k):\n", " f = frames[k]\n", " im1.set_data(logit_imgs[k])\n", " im1.set_clim(vmin=logit_imgs[k].min(), vmax=logit_imgs[k].max())\n", " im2.set_data(spec_imgs[k])\n", " title.set_text(f'{op}, epoch {snapshots[f][\"epoch\"]}')\n", " return im1, im2, title\n", "\n", " anim = FuncAnimation(fig, update, frames=len(frames), interval=100, blit=False)\n", " plt.close()\n", " display(HTML(anim.to_jshtml()))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "RUNS_DIR = Path('runs')\n", "\n", "def save_run(result, runs_dir=RUNS_DIR):\n", " runs_dir.mkdir(parents=True, exist_ok=True)\n", " cfg = result['config']\n", " path = runs_dir / f\"{cfg['operation']}_seed{cfg['seed']}.pt\"\n", " torch.save(result, path)\n", " return path\n", "\n", "def load_run(operation, seed, runs_dir=RUNS_DIR):\n", " return torch.load(runs_dir / f\"{operation}_seed{seed}.pt\", weights_only=False)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "FIGURES_DIR = Path('figures')\n", "\n", "def save_figures(result, figures_dir=FIGURES_DIR):\n", " op = result['config']['operation']\n", " out = figures_dir / op\n", " out.mkdir(parents=True, exist_ok=True)\n", " save_training_summary(result, out / 'training_summary.png')\n", " save_fourier_spectrum_final(result, out / 'fourier_spectrum_final.png')\n", " save_raw_weights_final(result, out / 'raw_weights_final.png')\n", " save_nanda_2d_final(result, out / 'nanda_2d_final.png')\n", " return out\n", "\n", "\n", "def save_training_summary(result, path):\n", " op = result['config']['operation']\n", " ipr = result['ipr_history']\n", " xs = [e['epoch'] for e in ipr]\n", "\n", " fig, axes = plt.subplots(1, 3, figsize=(15, 4))\n", " axes[0].plot(xs, result['train_accs'], label='train')\n", " axes[0].plot(xs, result['test_accs'], label='test')\n", " axes[0].set_xlabel('Epoch')\n", " axes[0].set_ylabel('Accuracy')\n", " axes[0].set_title(f'{op} accuracy')\n", " axes[0].legend()\n", "\n", " axes[1].plot(xs, result['train_losses'], label='train')\n", " axes[1].plot(xs, result['test_losses'], label='test')\n", " axes[1].set_xlabel('Epoch')\n", " axes[1].set_ylabel('MSE loss')\n", " axes[1].set_yscale('log')\n", " axes[1].set_title(f'{op} loss')\n", " axes[1].legend()\n", "\n", " axes[2].plot(xs, [e['fc1_a'] for e in ipr], label='fc1 a')\n", " axes[2].plot(xs, [e['fc1_b'] for e in ipr], label='fc1 b')\n", " axes[2].plot(xs, [e['fc2'] for e in ipr], label='fc2')\n", " if op in ('mul', 'div'):\n", " axes[2].plot(xs, [e['fc1_a_log'] for e in ipr], '--', label='fc1 a log')\n", " axes[2].plot(xs, [e['fc1_b_log'] for e in ipr], '--', label='fc1 b log')\n", " axes[2].plot(xs, [e['fc2_log'] for e in ipr], '--', label='fc2 log')\n", " axes[2].set_xlabel('Epoch')\n", " axes[2].set_ylabel('IPR_2')\n", " axes[2].set_title(f'{op} IPR (Fourier)')\n", " axes[2].legend()\n", "\n", " plt.tight_layout()\n", " fig.savefig(path, dpi=150, bbox_inches='tight')\n", " plt.close(fig)\n", "\n", "\n", "def save_fourier_spectrum_final(result, path):\n", " snapshots = result['weight_snapshots']\n", " op = result['config']['operation']\n", " p = result['config']['p']\n", " pi_log = result['pi_log']\n", " use_log = op in ('mul', 'div')\n", "\n", " w1 = snapshots[-1]['fc1']\n", " w1_a, w1_b = w1[:, :p], w1[:, p:]\n", " if use_log:\n", " w1_a = to_log_basis(w1_a, pi_log)\n", " w1_b = to_log_basis(w1_b, pi_log)\n", " sa = torch.fft.fft(w1_a, dim=1).abs() ** 2\n", " sb = torch.fft.fft(w1_b, dim=1).abs() ** 2\n", " sa = sa / sa.sum(dim=1, keepdim=True).clamp(min=1e-30)\n", " sb = sb / sb.sum(dim=1, keepdim=True).clamp(min=1e-30)\n", " half = sa.size(1) // 2\n", " sa_n, sb_n = sa[:, :half].numpy(), sb[:, :half].numpy()\n", " sort_idx = np.argsort(sa_n.argmax(axis=1))\n", "\n", " basis_label = 'log basis' if use_log else 'natural basis'\n", " fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n", " im1 = ax1.imshow(sa_n[sort_idx], aspect='auto', cmap='hot', vmin=0, vmax=0.3)\n", " ax1.set_xlabel('Frequency')\n", " ax1.set_ylabel('Neuron (sorted)')\n", " ax1.set_title(f'{op} fc1 a-block spectrum, {basis_label}, epoch {snapshots[-1][\"epoch\"]}')\n", " fig.colorbar(im1, ax=ax1)\n", "\n", " im2 = ax2.imshow(sb_n[sort_idx], aspect='auto', cmap='hot', vmin=0, vmax=0.3)\n", " ax2.set_xlabel('Frequency')\n", " ax2.set_ylabel('Neuron (sorted)')\n", " ax2.set_title(f'{op} fc1 b-block spectrum, {basis_label}, epoch {snapshots[-1][\"epoch\"]}')\n", " fig.colorbar(im2, ax=ax2)\n", "\n", " plt.tight_layout()\n", " fig.savefig(path, dpi=150, bbox_inches='tight')\n", " plt.close(fig)\n", "\n", "\n", "def save_raw_weights_final(result, path):\n", " snapshots = result['weight_snapshots']\n", " op = result['config']['operation']\n", " p = result['config']['p']\n", "\n", " final_w1_a = snapshots[-1]['fc1'][:, :p]\n", " final_spec = (torch.fft.fft(final_w1_a, dim=1).abs() ** 2).numpy()\n", " sort_idx = np.argsort(final_spec[:, 1:p // 2].argmax(axis=1))\n", "\n", " fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n", " w = snapshots[-1]['fc1']\n", " im1 = ax1.imshow(w[:, :p].numpy()[sort_idx], aspect='auto', cmap='RdBu', vmin=-1, vmax=1)\n", " ax1.set_xlabel('Input index a')\n", " ax1.set_ylabel('Neuron (sorted)')\n", " ax1.set_title(f'{op} fc1 a-block raw weights, epoch {snapshots[-1][\"epoch\"]}')\n", "\n", " im2 = ax2.imshow(w[:, p:].numpy()[sort_idx], aspect='auto', cmap='RdBu', vmin=-1, vmax=1)\n", " ax2.set_xlabel('Input index b')\n", " ax2.set_ylabel('Neuron (sorted)')\n", " ax2.set_title(f'{op} fc1 b-block raw weights, epoch {snapshots[-1][\"epoch\"]}')\n", "\n", " plt.tight_layout()\n", " fig.savefig(path, dpi=150, bbox_inches='tight')\n", " plt.close(fig)\n", "\n", "\n", "def save_nanda_2d_final(result, path):\n", " snapshots = result['weight_snapshots']\n", " op = result['config']['operation']\n", " p = result['config']['p']\n", " hidden_dim = result['config']['hidden_dim']\n", " pi_log = result['pi_log']\n", " device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n", " use_log = op in ('mul', 'div')\n", "\n", " if op in ('add', 'sub'):\n", " a = torch.arange(p).repeat_interleave(p).to(device)\n", " b = torch.arange(p).repeat(p).to(device)\n", " n_a, n_b = p, p\n", " else:\n", " a = torch.arange(p).repeat_interleave(p - 1).to(device)\n", " b = torch.arange(1, p).repeat(p).to(device)\n", " n_a, n_b = p, p - 1\n", " x_all = torch.cat([F.one_hot(a, p), F.one_hot(b, p)], dim=1).float()\n", "\n", " if op == 'add':\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p)\n", " bi = torch.arange(p).unsqueeze(0).expand(p, p)\n", " cls = (ai + bi) % p\n", " elif op == 'sub':\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p)\n", " bi = torch.arange(p).unsqueeze(0).expand(p, p)\n", " cls = (ai - bi) % p\n", " elif op == 'mul':\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p - 1)\n", " bi = torch.arange(1, p).unsqueeze(0).expand(p, p - 1)\n", " cls = (ai * bi) % p\n", " elif op == 'div':\n", " b_inv_row = torch.tensor([pow(int(bi), p - 2, p) for bi in range(1, p)])\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p - 1)\n", " bi = b_inv_row.unsqueeze(0).expand(p, p - 1)\n", " cls = (ai * bi) % p\n", "\n", " work_model = baseMLP(p, hidden_dim).to(device)\n", " snap = snapshots[-1]\n", " work_model.fc1.weight.data = snap['fc1'].to(device)\n", " work_model.fc2.weight.data = snap['fc2'].to(device)\n", " work_model.eval()\n", " with torch.no_grad():\n", " logits = work_model(x_all).reshape(n_a, n_b, p).cpu()\n", " correct_logit = logits.gather(2, cls.unsqueeze(2)).squeeze(2)\n", " if use_log:\n", " correct_logit = correct_logit.index_select(0, pi_log).index_select(1, pi_log - 1)\n", " spectrum = torch.fft.fft2(correct_logit).abs() ** 2\n", " spectrum[0, 0] = 0\n", " spectrum = spectrum / spectrum.sum().clamp(min=1e-30)\n", " spectrum_shifted = torch.fft.fftshift(spectrum).numpy()\n", " H, W = spectrum_shifted.shape\n", " freq_a = np.arange(H) - (H // 2)\n", " freq_b = np.arange(W) - (W // 2)\n", " extent = (freq_b[0] - 0.5, freq_b[-1] + 0.5, freq_a[-1] + 0.5, freq_a[0] - 0.5)\n", "\n", " basis_label = 'log basis' if use_log else 'natural basis'\n", " fig, axes = plt.subplots(1, 2, figsize=(13, 5))\n", "\n", " im1 = axes[0].imshow(correct_logit.numpy(), aspect='equal', cmap='RdBu')\n", " axes[0].set_xlabel('b')\n", " axes[0].set_ylabel('a')\n", " axes[0].set_title(f'{op} logit at correct class, {basis_label}, epoch {snap[\"epoch\"]}')\n", " fig.colorbar(im1, ax=axes[0])\n", "\n", " im2 = axes[1].imshow(spectrum_shifted, aspect='equal', cmap='hot', vmin=0, vmax=0.05, extent=extent)\n", " axes[1].set_xlabel('Frequency b (centered at DC)')\n", " axes[1].set_ylabel('Frequency a (centered at DC)')\n", " axes[1].set_title(f'{op} 2D Fourier of correct logit, {basis_label}, epoch {snap[\"epoch\"]}')\n", " axes[1].axhline(0, color='cyan', linewidth=0.5, alpha=0.4)\n", " axes[1].axvline(0, color='cyan', linewidth=0.5, alpha=0.4)\n", " fig.colorbar(im2, ax=axes[1])\n", "\n", " plt.tight_layout()\n", " fig.savefig(path, dpi=150, bbox_inches='tight')\n", " plt.close(fig)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#modular addition\n", "#use cached run if available, otherwise train\n", "_cache_path = RUNS_DIR / f'add_seed0.pt'\n", "if _cache_path.exists():\n", " result_add = load_run('add', 0)\n", " print(f'loaded cached run from {_cache_path}')\n", "else:\n", " result_add = train_run(\n", " operation='add',\n", " p=113,\n", " hidden_dim=128,\n", " lr=1e-3,\n", " weight_decay=1.0,\n", " train_frac=0.3,\n", " epochs=10_000,\n", " seed=0,\n", " snapshot_every=200,\n", " log_every=50,\n", " live_plot=True,\n", " )\n", " save_run(result_add)\n", "save_figures(result_add)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#add visualizations: fc1 fourier spectrum, raw weight rows, nanda 2D fourier\n", "\n", "viz_fourier_spectrum(result_add)\n", "viz_raw_weights(result_add)\n", "viz_nanda_2d(result_add)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#modular subtraction\n", "#expected to mirror addition since a-b mod p = a+(-b) mod p\n", "#use cached run if available, otherwise train\n", "_cache_path = RUNS_DIR / f'sub_seed0.pt'\n", "if _cache_path.exists():\n", " result_sub = load_run('sub', 0)\n", " print(f'loaded cached run from {_cache_path}')\n", "else:\n", " result_sub = train_run(\n", " operation='sub',\n", " p=113,\n", " hidden_dim=128,\n", " lr=1e-3,\n", " weight_decay=1.0,\n", " train_frac=0.3,\n", " epochs=10_000,\n", " seed=0,\n", " snapshot_every=200,\n", " log_every=50,\n", " live_plot=True,\n", " )\n", " save_run(result_sub)\n", "save_figures(result_sub)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#sub visualizations\n", "\n", "viz_fourier_spectrum(result_sub)\n", "viz_raw_weights(result_sub)\n", "viz_nanda_2d(result_sub)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#modular multiplication\n", "#analytical solution lives in discrete-log basis (doshi 2024), train_run automatically tracks log-basis IPR alongside raw IPR\n", "#use cached run if available, otherwise train\n", "_cache_path = RUNS_DIR / f'mul_seed0.pt'\n", "if _cache_path.exists():\n", " result_mul = load_run('mul', 0)\n", " print(f'loaded cached run from {_cache_path}')\n", "else:\n", " result_mul = train_run(\n", " operation='mul',\n", " p=113,\n", " hidden_dim=128,\n", " lr=1e-3,\n", " weight_decay=1.0,\n", " train_frac=0.3,\n", " epochs=10_000,\n", " seed=0,\n", " snapshot_every=200,\n", " log_every=50,\n", " live_plot=True,\n", " )\n", " save_run(result_mul)\n", "save_figures(result_mul)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#mul visualizations\n", "#dashed lines in IPR plot during training are log-basis IPR, the gap between solid and dashed is the \"two-stage emergence\" signal\n", "\n", "viz_fourier_spectrum(result_mul)\n", "viz_raw_weights(result_mul)\n", "viz_nanda_2d(result_mul)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#modular division\n", "#a/b mod p = a * b^(-1) mod p, expected to mirror multiplication in log basis with sign flip on b\n", "#use cached run if available, otherwise train\n", "_cache_path = RUNS_DIR / f'div_seed0.pt'\n", "if _cache_path.exists():\n", " result_div = load_run('div', 0)\n", " print(f'loaded cached run from {_cache_path}')\n", "else:\n", " result_div = train_run(\n", " operation='div',\n", " p=113,\n", " hidden_dim=128,\n", " lr=1e-3,\n", " weight_decay=1.0,\n", " train_frac=0.3,\n", " epochs=10_000,\n", " seed=0,\n", " snapshot_every=200,\n", " log_every=50,\n", " live_plot=True,\n", " )\n", " save_run(result_div)\n", "save_figures(result_div)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#div visualizations\n", "\n", "viz_fourier_spectrum(result_div)\n", "viz_raw_weights(result_div)\n", "viz_nanda_2d(result_div)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#multi-seed sweep: train seeds 1-4 across all 4 ops, saved next to seed 0 in runs/\n", "#live_plot is off and per-seed figures are skipped\n", "#log_every=50 to match seed 0 so cross-seed plots are smooth, ~10-15 min total\n", "\n", "extra_seeds = [1, 2, 3, 4]\n", "ops = ['add', 'sub', 'mul', 'div']\n", "\n", "for op in ops:\n", " for seed in extra_seeds:\n", " _cache_path = RUNS_DIR / f'{op}_seed{seed}.pt'\n", " if _cache_path.exists():\n", " print(f\"{op} seed {seed} -> cached at {_cache_path}, skipping\")\n", " continue\n", " print(f\"\\n{op} seed {seed}\")\n", " t0 = time.time()\n", " result = train_run(\n", " operation=op,\n", " p=113,\n", " hidden_dim=128,\n", " lr=1e-3,\n", " weight_decay=1.0,\n", " train_frac=0.3,\n", " epochs=10_000,\n", " seed=seed,\n", " snapshot_every=200,\n", " log_every=50,\n", " live_plot=False,\n", " )\n", " save_run(result)\n", " ipr_final = result['ipr_history'][-1]\n", " print(f\" done in {time.time()-t0:.0f}s, \"\n", " f\"train_acc={result['train_accs'][-1]:.4f} \"\n", " f\"test_acc={result['test_accs'][-1]:.4f} \"\n", " f\"IPR_fc1_a={ipr_final['fc1_a']:.3f} \"\n", " f\"IPR_fc2={ipr_final['fc2']:.3f}\")\n", "\n", "print(\"\\nall multi-seed runs complete, runs/ has 5 seeds per op\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#load all 5 seeds for all 4 ops, align logs to common epoch grid (seed 0 has finer log_every than seeds 1-4)\n", "\n", "ops = ['add', 'sub', 'mul', 'div']\n", "seeds = [0, 1, 2, 3, 4]\n", "\n", "raw = {op: [load_run(op, s) for s in seeds] for op in ops}\n", "\n", "#find common epoch set across all seeds across all ops\n", "common = None\n", "for op in ops:\n", " for r in raw[op]:\n", " eps = set(e['epoch'] for e in r['ipr_history'])\n", " common = eps if common is None else (common & eps)\n", "common = sorted(common)\n", "print(f\"common epochs across all seeds: {len(common)}, range {common[0]} to {common[-1]}\")\n", "\n", "#downsample finer-grained seeds to the common epoch set so trajectories stack\n", "all_results = {}\n", "for op in ops:\n", " all_results[op] = []\n", " for r in raw[op]:\n", " ep_to_idx = {e['epoch']: i for i, e in enumerate(r['ipr_history'])}\n", " indices = [ep_to_idx[e] for e in common]\n", " r2 = dict(r)\n", " r2['ipr_history'] = [r['ipr_history'][i] for i in indices]\n", " r2['train_accs'] = [r['train_accs'][i] for i in indices]\n", " r2['test_accs'] = [r['test_accs'][i] for i in indices]\n", " r2['train_losses'] = [r['train_losses'][i] for i in indices]\n", " r2['test_losses'] = [r['test_losses'][i] for i in indices]\n", " all_results[op].append(r2)\n", "\n", "COMPARISON_DIR = Path('figures/comparison')\n", "COMPARISON_DIR.mkdir(parents=True, exist_ok=True)\n", "\n", "def plot_band(ax, epochs, trajs, label=None, color=None, linestyle='-'):\n", " #plot mean across seeds with +/- 1 std band\n", " mean = trajs.mean(axis=0)\n", " std = trajs.std(axis=0)\n", " line = ax.plot(epochs, mean, label=label, color=color, linestyle=linestyle)[0]\n", " ax.fill_between(epochs, mean - std, mean + std, alpha=0.2, color=line.get_color())\n", " return line" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#figure 1: cross-op IPR trajectories with seed bands\n", "#mul/div shown in log basis (their natural basis), add/sub in raw basis\n", "#expected: add ~ sub trajectories overlap, mul ~ div overlap, two pairs separated\n", "\n", "fig, axes = plt.subplots(1, 3, figsize=(18, 4.5))\n", "op_color_map = {'add': 'tab:blue', 'sub': 'tab:orange', 'mul': 'tab:green', 'div': 'tab:red'}\n", "\n", "for ax_idx, key_base in enumerate(['fc1_a', 'fc1_b', 'fc2']):\n", " ax = axes[ax_idx]\n", " for op in ops:\n", " suffix = '_log' if op in ('mul', 'div') else ''\n", " key = key_base + suffix\n", " epochs = np.array([e['epoch'] for e in all_results[op][0]['ipr_history']])\n", " trajs = np.array([[e[key] for e in r['ipr_history']] for r in all_results[op]])\n", " plot_band(ax, epochs, trajs, label=op, color=op_color_map[op])\n", " ax.set_xlabel('Epoch')\n", " ax.set_ylabel('IPR_2')\n", " ax.set_title(f'{key_base} IPR (mul/div in log basis)')\n", " ax.legend()\n", " ax.set_ylim(0, 0.5)\n", " ax.axhline(y=1/113, color='black', linestyle=':', alpha=0.4)\n", "\n", "plt.tight_layout()\n", "fig.savefig(COMPARISON_DIR / 'fig1_cross_op_ipr.png', dpi=150, bbox_inches='tight')\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#figure 2: hidden progress check, IPR vs train/test acc per op, with memorization plateau shaded\n", "#\"hidden progress\" claim: IPR climbs DURING the plateau, before test acc moves\n", "\n", "fig, axes = plt.subplots(2, 2, figsize=(14, 9))\n", "for i, op in enumerate(ops):\n", " ax = axes[i // 2, i % 2]\n", " epochs = np.array([e['epoch'] for e in all_results[op][0]['ipr_history']])\n", " train_accs = np.array([r['train_accs'] for r in all_results[op]])\n", " test_accs = np.array([r['test_accs'] for r in all_results[op]])\n", "\n", " primary_key = 'fc1_a_log' if op in ('mul', 'div') else 'fc1_a'\n", " ipr_trajs = np.array([[e[primary_key] for e in r['ipr_history']] for r in all_results[op]])\n", " basis_label = 'log' if op in ('mul', 'div') else 'raw'\n", "\n", " plot_band(ax, epochs, train_accs, label='train acc', color='tab:blue')\n", " plot_band(ax, epochs, test_accs, label='test acc', color='tab:orange')\n", "\n", " ax2 = ax.twinx()\n", " plot_band(ax2, epochs, ipr_trajs, label=f'IPR fc1_a ({basis_label})', color='tab:green')\n", "\n", " #shade the memorization plateau region (train_acc > 0.99 AND test_acc < 0.5)\n", " train_mean = train_accs.mean(0)\n", " test_mean = test_accs.mean(0)\n", " in_plateau = (train_mean > 0.99) & (test_mean < 0.5)\n", " if in_plateau.any():\n", " plateau_epochs = epochs[in_plateau]\n", " ax.axvspan(plateau_epochs.min(), plateau_epochs.max(),\n", " color='yellow', alpha=0.2, label='memorization plateau')\n", "\n", " ax.set_xlabel('Epoch')\n", " ax.set_ylabel('Accuracy')\n", " ax2.set_ylabel(f'IPR_2 fc1_a ({basis_label})')\n", " ax.set_title(f'{op}: hidden progress check')\n", " ax.set_ylim(-0.05, 1.05)\n", " ax2.set_ylim(0, 0.5)\n", "\n", " h1, l1 = ax.get_legend_handles_labels()\n", " h2, l2 = ax2.get_legend_handles_labels()\n", " ax.legend(h1+h2, l1+l2, loc='center right', fontsize=8)\n", "\n", "plt.tight_layout()\n", "fig.savefig(COMPARISON_DIR / 'fig2_hidden_progress.png', dpi=150, bbox_inches='tight')\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#figure 3: two-stage emergence test for mul and div\n", "#raw IPR (natural basis) vs log IPR, with test acc as reference\n", "#if log climbs first and raw stays flat, that's the two-stage signal\n", "\n", "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n", "for i, op in enumerate(['mul', 'div']):\n", " ax = axes[i]\n", " epochs = np.array([e['epoch'] for e in all_results[op][0]['ipr_history']])\n", " raw_trajs = np.array([[e['fc1_a'] for e in r['ipr_history']] for r in all_results[op]])\n", " log_trajs = np.array([[e['fc1_a_log'] for e in r['ipr_history']] for r in all_results[op]])\n", " test_accs = np.array([r['test_accs'] for r in all_results[op]])\n", "\n", " plot_band(ax, epochs, raw_trajs, label='fc1_a IPR (raw basis)', color='tab:blue')\n", " plot_band(ax, epochs, log_trajs, label='fc1_a IPR (log basis)', color='tab:red')\n", "\n", " ax2 = ax.twinx()\n", " plot_band(ax2, epochs, test_accs, label='test acc', color='gray', linestyle='--')\n", " ax2.set_ylim(-0.05, 1.05)\n", " ax2.set_ylabel('Test accuracy')\n", "\n", " ax.set_xlabel('Epoch')\n", " ax.set_ylabel('IPR_2')\n", " ax.set_title(f'{op}: raw vs log basis IPR')\n", " ax.set_ylim(0, 0.5)\n", "\n", " h1, l1 = ax.get_legend_handles_labels()\n", " h2, l2 = ax2.get_legend_handles_labels()\n", " ax.legend(h1+h2, l1+l2, loc='center right')\n", "\n", "plt.tight_layout()\n", "fig.savefig(COMPARISON_DIR / 'fig3_two_stage_emergence.png', dpi=150, bbox_inches='tight')\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#figure 4: final IPR values across ops and layers, mean +/- std over 5 seeds\n", "#bar chart with mul/div showing both raw and log basis side by side\n", "#dotted line is the random baseline 1/p, anything above is structured periodicity\n", "\n", "fig, ax = plt.subplots(figsize=(15, 5))\n", "\n", "labels, means, stds, colors = [], [], [], []\n", "op_color_map = {'add': 'tab:blue', 'sub': 'tab:orange', 'mul': 'tab:green', 'div': 'tab:red'}\n", "\n", "for op in ops:\n", " for key, sublabel in [('fc1_a', 'fc1_a'), ('fc1_b', 'fc1_b'), ('fc2', 'fc2')]:\n", " finals = np.array([r['ipr_history'][-1][key] for r in all_results[op]])\n", " labels.append(f'{op}\\n{sublabel}\\nraw')\n", " means.append(finals.mean())\n", " stds.append(finals.std())\n", " colors.append(op_color_map[op])\n", " if op in ('mul', 'div'):\n", " for key, sublabel in [('fc1_a_log', 'fc1_a'), ('fc1_b_log', 'fc1_b'), ('fc2_log', 'fc2')]:\n", " finals = np.array([r['ipr_history'][-1][key] for r in all_results[op]])\n", " labels.append(f'{op}\\n{sublabel}\\nlog')\n", " means.append(finals.mean())\n", " stds.append(finals.std())\n", " colors.append(op_color_map[op])\n", "\n", "x = np.arange(len(labels))\n", "ax.bar(x, means, yerr=stds, color=colors, capsize=3, edgecolor='black', linewidth=0.5)\n", "ax.set_xticks(x)\n", "ax.set_xticklabels(labels, fontsize=8)\n", "ax.set_ylabel('Final IPR_2')\n", "ax.set_title('Final IPR across ops and layers, mean +/- std over 5 seeds')\n", "ax.set_ylim(0, 0.5)\n", "ax.axhline(y=1/113, color='black', linestyle=':', alpha=0.6, label='random baseline 1/p')\n", "ax.legend()\n", "\n", "plt.tight_layout()\n", "fig.savefig(COMPARISON_DIR / 'fig4_endpoint_ipr.png', dpi=150, bbox_inches='tight')\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#per-row Fourier spectrum film strip across training, supports the\n", "#claim that the spectrum sharpens monotonically rather than locking onto\n", "#specific modes first. Loads from runs/ so no retraining needed.\n", "#change op or target_epochs to regenerate for any operation.\n", "\n", "op = 'add'\n", "target_epochs = [0, 500, 1000, 5000, 10000]\n", "\n", "result = torch.load(RUNS_DIR / f'{op}_seed0.pt', weights_only=False)\n", "snapshots = result['weight_snapshots']\n", "p = result['config']['p']\n", "pi_log = result['pi_log']\n", "use_log = op in ('mul', 'div')\n", "\n", "snap_epochs = np.array([s['epoch'] for s in snapshots])\n", "chosen = [int(np.argmin(np.abs(snap_epochs - te))) for te in target_epochs]\n", "\n", "#sort neurons once by their final-epoch peak frequency so the strip is comparable across frames\n", "final_w1_a = snapshots[-1]['fc1'][:, :p]\n", "if use_log:\n", " final_w1_a = to_log_basis(final_w1_a, pi_log)\n", "final_spec = torch.fft.fft(final_w1_a, dim=1).abs() ** 2\n", "final_spec_norm = final_spec / final_spec.sum(dim=1, keepdim=True).clamp(min=1e-30)\n", "half = final_spec_norm.size(1) // 2\n", "sort_idx = np.argsort(final_spec_norm[:, :half].numpy().argmax(axis=1))\n", "\n", "n_frames = len(chosen)\n", "fig, axes = plt.subplots(1, n_frames, figsize=(4.0 * n_frames, 5), sharey=True)\n", "if n_frames == 1:\n", " axes = [axes]\n", "\n", "im = None\n", "for ax, idx in zip(axes, chosen):\n", " snap = snapshots[idx]\n", " w1_a = snap['fc1'][:, :p]\n", " if use_log:\n", " w1_a = to_log_basis(w1_a, pi_log)\n", " spec = torch.fft.fft(w1_a, dim=1).abs() ** 2\n", " spec_norm = spec / spec.sum(dim=1, keepdim=True).clamp(min=1e-30)\n", " sa = spec_norm[:, :half].numpy()\n", " im = ax.imshow(sa[sort_idx], aspect='auto', cmap='hot', vmin=0, vmax=0.3)\n", " ax.set_title(f'epoch {snap[\"epoch\"]}')\n", " ax.set_xlabel('Frequency')\n", "\n", "axes[0].set_ylabel('Neuron (sorted by final-epoch peak)')\n", "fig.colorbar(im, ax=axes[-1], fraction=0.046, pad=0.04)\n", "basis_label = 'log basis' if use_log else 'natural basis'\n", "fig.suptitle(f'{op} fc1 a-block per-row Fourier spectrum across training, {basis_label}', y=1.02)\n", "plt.tight_layout()\n", "out_path = FIGURES_DIR / 'comparison' / f'fig5_spectrum_filmstrip_{op}.png'\n", "out_path.parent.mkdir(parents=True, exist_ok=True)\n", "fig.savefig(out_path, dpi=150, bbox_inches='tight')\n", "plt.show()\n", "print(f'saved {out_path}')\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#2x2 training dynamics, one panel per op, twin axis showing accuracy (left) and loss (right, log).\n", "#loads from runs/ so no retraining.\n", "\n", "ops = ['add', 'sub', 'mul', 'div']\n", "fig, axes = plt.subplots(2, 2, figsize=(13, 9))\n", "for ax, op in zip(axes.flat, ops):\n", " result = torch.load(RUNS_DIR / f'{op}_seed0.pt', weights_only=False)\n", " xs = [e['epoch'] for e in result['ipr_history']]\n", " ax.plot(xs, result['train_accs'], color='tab:blue', label='train acc')\n", " ax.plot(xs, result['test_accs'], color='tab:orange', label='test acc')\n", " ax.set_xlabel('Epoch')\n", " ax.set_ylabel('Accuracy')\n", " ax.set_ylim(-0.05, 1.05)\n", " ax.set_title(op)\n", "\n", " ax2 = ax.twinx()\n", " ax2.plot(xs, result['train_losses'], color='tab:blue', linestyle='--', alpha=0.5, label='train loss')\n", " ax2.plot(xs, result['test_losses'], color='tab:orange', linestyle='--', alpha=0.5, label='test loss')\n", " ax2.set_yscale('log')\n", " ax2.set_ylabel('MSE loss (log)')\n", "\n", " lines1, labels1 = ax.get_legend_handles_labels()\n", " lines2, labels2 = ax2.get_legend_handles_labels()\n", " ax.legend(lines1 + lines2, labels1 + labels2, loc='center right', fontsize=8)\n", "\n", "fig.suptitle('Training dynamics across operations (seed 0)', y=1.0)\n", "plt.tight_layout()\n", "out_path = FIGURES_DIR / 'comparison' / 'fig6_training_dynamics_2x2.png'\n", "out_path.parent.mkdir(parents=True, exist_ok=True)\n", "fig.savefig(out_path, dpi=150, bbox_inches='tight')\n", "plt.show()\n", "print(f'saved {out_path}')\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#2x2 endpoint per-row Fourier spectra of fc1 a-block, one panel per op, in correct basis.\n", "#loads from runs/ so no retraining.\n", "\n", "ops = ['add', 'sub', 'mul', 'div']\n", "fig, axes = plt.subplots(2, 2, figsize=(13, 10))\n", "last_im = None\n", "for ax, op in zip(axes.flat, ops):\n", " result = torch.load(RUNS_DIR / f'{op}_seed0.pt', weights_only=False)\n", " snapshots = result['weight_snapshots']\n", " p = result['config']['p']\n", " pi_log = result['pi_log']\n", " use_log = op in ('mul', 'div')\n", "\n", " w1_a = snapshots[-1]['fc1'][:, :p]\n", " if use_log:\n", " w1_a = to_log_basis(w1_a, pi_log)\n", " spec = torch.fft.fft(w1_a, dim=1).abs() ** 2\n", " spec = spec / spec.sum(dim=1, keepdim=True).clamp(min=1e-30)\n", " half = spec.size(1) // 2\n", " sa = spec[:, :half].numpy()\n", " sort_idx = np.argsort(sa.argmax(axis=1))\n", "\n", " basis_label = 'log basis' if use_log else 'natural basis'\n", " last_im = ax.imshow(sa[sort_idx], aspect='auto', cmap='hot', vmin=0, vmax=0.3)\n", " ax.set_xlabel('Frequency')\n", " ax.set_ylabel('Neuron (sorted)')\n", " ax.set_title(f'{op}, {basis_label}, epoch {snapshots[-1][\"epoch\"]}')\n", "\n", "fig.colorbar(last_im, ax=axes, fraction=0.025, pad=0.04)\n", "fig.suptitle('Endpoint per-row Fourier spectrum, fc1 a-block', y=1.0)\n", "out_path = FIGURES_DIR / 'comparison' / 'fig7_endpoint_spectra_2x2.png'\n", "out_path.parent.mkdir(parents=True, exist_ok=True)\n", "fig.savefig(out_path, dpi=150, bbox_inches='tight')\n", "plt.show()\n", "print(f'saved {out_path}')\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#2x2 Nanda 2D Fourier of the correct logit, one panel per op, with fftshift so anti-diagonal\n", "#energy for sub/div is visible. mul/div in log basis, add/sub in natural basis.\n", "#loads from runs/ so no retraining.\n", "\n", "ops = ['add', 'sub', 'mul', 'div']\n", "fig, axes = plt.subplots(2, 2, figsize=(12, 11))\n", "last_im = None\n", "for ax, op in zip(axes.flat, ops):\n", " result = torch.load(RUNS_DIR / f'{op}_seed0.pt', weights_only=False)\n", " snapshots = result['weight_snapshots']\n", " p = result['config']['p']\n", " hidden_dim = result['config']['hidden_dim']\n", " pi_log = result['pi_log']\n", " use_log = op in ('mul', 'div')\n", "\n", " if op in ('add', 'sub'):\n", " a = torch.arange(p).repeat_interleave(p).to(device)\n", " b = torch.arange(p).repeat(p).to(device)\n", " n_a, n_b = p, p\n", " else:\n", " a = torch.arange(p).repeat_interleave(p - 1).to(device)\n", " b = torch.arange(1, p).repeat(p).to(device)\n", " n_a, n_b = p, p - 1\n", " x_all = torch.cat([F.one_hot(a, p), F.one_hot(b, p)], dim=1).float()\n", "\n", " if op == 'add':\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p)\n", " bi = torch.arange(p).unsqueeze(0).expand(p, p)\n", " cls = (ai + bi) % p\n", " elif op == 'sub':\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p)\n", " bi = torch.arange(p).unsqueeze(0).expand(p, p)\n", " cls = (ai - bi) % p\n", " elif op == 'mul':\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p - 1)\n", " bi = torch.arange(1, p).unsqueeze(0).expand(p, p - 1)\n", " cls = (ai * bi) % p\n", " elif op == 'div':\n", " b_inv_row = torch.tensor([pow(int(bi), p - 2, p) for bi in range(1, p)])\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p - 1)\n", " bi = b_inv_row.unsqueeze(0).expand(p, p - 1)\n", " cls = (ai * bi) % p\n", "\n", " work_model = baseMLP(p, hidden_dim).to(device)\n", " snap = snapshots[-1]\n", " work_model.fc1.weight.data = snap['fc1'].to(device)\n", " work_model.fc2.weight.data = snap['fc2'].to(device)\n", " work_model.eval()\n", " with torch.no_grad():\n", " logits = work_model(x_all).reshape(n_a, n_b, p).cpu()\n", " correct_logit = logits.gather(2, cls.unsqueeze(2)).squeeze(2)\n", " if use_log:\n", " correct_logit = correct_logit.index_select(0, pi_log).index_select(1, pi_log - 1)\n", "\n", " spectrum = torch.fft.fft2(correct_logit).abs() ** 2\n", " spectrum[0, 0] = 0\n", " spectrum = spectrum / spectrum.sum().clamp(min=1e-30)\n", " spectrum_shifted = torch.fft.fftshift(spectrum).numpy()\n", " H, W = spectrum_shifted.shape\n", " freq_a = np.arange(H) - (H // 2)\n", " freq_b = np.arange(W) - (W // 2)\n", " extent = (freq_b[0] - 0.5, freq_b[-1] + 0.5, freq_a[-1] + 0.5, freq_a[0] - 0.5)\n", "\n", " basis_label = 'log basis' if use_log else 'natural basis'\n", " last_im = ax.imshow(spectrum_shifted, aspect='equal', cmap='hot', vmin=0, vmax=0.05, extent=extent)\n", " ax.set_xlabel('Frequency b (centered at DC)')\n", " ax.set_ylabel('Frequency a (centered at DC)')\n", " ax.set_title(f'{op}, {basis_label}, epoch {snap[\"epoch\"]}')\n", " ax.axhline(0, color='cyan', linewidth=0.5, alpha=0.4)\n", " ax.axvline(0, color='cyan', linewidth=0.5, alpha=0.4)\n", "\n", "fig.colorbar(last_im, ax=axes, fraction=0.025, pad=0.04)\n", "fig.suptitle('2D Fourier of the correct-class logit (fftshift, full spectrum)', y=1.0)\n", "out_path = FIGURES_DIR / 'comparison' / 'fig8_nanda_2d_2x2.png'\n", "out_path.parent.mkdir(parents=True, exist_ok=True)\n", "fig.savefig(out_path, dpi=150, bbox_inches='tight')\n", "plt.show()\n", "print(f'saved {out_path}')\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "---\n", "\n", "## Part 2 — Loss-landscape analysis\n", "\n", "*Source: `loss_landscape.ipynb`*\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Loss landscape and 3D Fourier features\n", "\n", "PCA loss-landscape animation for addition, plus a 2x2 of 3D surfaces showing the\n", "learned correct-class logit for all four operations. Loads saved runs from `runs/`,\n", "no retraining.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "from pathlib import Path\n", "from matplotlib.animation import FuncAnimation\n", "from mpl_toolkits.mplot3d import Axes3D\n", "from IPython.display import HTML\n", "\n", "import matplotlib as mpl\n", "mpl.rcParams['animation.embed_limit'] = 200 #MB\n", "\n", "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n", "print(f\"device: {device}\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#redefine the trained architecture and dataset builder so we can load weights and recompute losses\n", "#these must match what was used during training in MI.ipynb (bias=False, train_frac in config, etc.)\n", "\n", "class baseMLP(nn.Module):\n", " def __init__(self, p, hidden_dim):\n", " super().__init__()\n", " self.fc1 = nn.Linear(2*p, hidden_dim, bias=False)\n", " self.relu = nn.ReLU()\n", " self.fc2 = nn.Linear(hidden_dim, p, bias=False)\n", " def forward(self, x):\n", " return self.fc2(self.relu(self.fc1(x)))\n", "\n", "def make_dataset(operation, p=113, train_frac=0.3, seed=0, device='cuda'):\n", " if operation in ('add', 'sub'):\n", " a = torch.arange(p).repeat_interleave(p)\n", " b = torch.arange(p).repeat(p)\n", " else:\n", " a = torch.arange(p).repeat_interleave(p - 1)\n", " b = torch.arange(1, p).repeat(p)\n", " if operation == 'add':\n", " y_labels = (a + b) % p\n", " elif operation == 'sub':\n", " y_labels = (a - b) % p\n", " elif operation == 'mul':\n", " y_labels = (a * b) % p\n", " elif operation == 'div':\n", " b_inv = torch.tensor([pow(int(bi), p - 2, p) for bi in b])\n", " y_labels = (a * b_inv) % p\n", " x = torch.cat([F.one_hot(a, p), F.one_hot(b, p)], dim=1).float()\n", " y = F.one_hot(y_labels, p).float()\n", " g = torch.Generator().manual_seed(seed)\n", " perm = torch.randperm(x.size(0), generator=g)\n", " n_train = int(round(x.size(0) * train_frac))\n", " return {\n", " 'x_train': x[perm[:n_train]].to(device),\n", " 'y_train': y[perm[:n_train]].to(device),\n", " 'x_test': x[perm[n_train:]].to(device),\n", " 'y_test': y[perm[n_train:]].to(device),\n", " }\n", "\n", "def load_run(operation, seed, runs_dir=Path('runs')):\n", " return torch.load(runs_dir / f\"{operation}_seed{seed}.pt\", weights_only=False)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#load addition seed 0, the only run we animate\n", "\n", "result_add = load_run('add', 0)\n", "print(f\"add: {len(result_add['weight_snapshots'])} snapshots, \"\n", " f\"final test_acc={result_add['test_accs'][-1]:.4f}\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#PCA-based loss landscape animation\n", "#1. flatten all snapshot weights into a (n_snap, n_params) matrix\n", "#2. SVD for top 2 principal components, project trajectory to 2D\n", "#3. evaluate train and test loss on a grid in the 2D PC plane\n", "#4. plot 3D surfaces side by side, animate the trajectory growing\n", "#5. save GIF and return anim for inline display\n", "\n", "#margin=4.0 extends the grid 4x past trajectory range on each side, so the trajectory looks small in a big landscape\n", "#z-axis uses log10(loss) so the descent into the basin reads as a real funnel rather than a flat plain\n", "\n", "def loss_landscape_animation(result, n_grid=80, margin=4.0, log_floor=1e-4, save_dir=Path('figures/comparison')):\n", " op = result['config']['operation']\n", " p = result['config']['p']\n", " hidden_dim = result['config']['hidden_dim']\n", " train_frac = result['config']['train_frac']\n", " seed = result['config']['seed']\n", " snapshots = result['weight_snapshots']\n", "\n", " print(f\"{op} landscape\")\n", " data = make_dataset(op, p=p, train_frac=train_frac, seed=seed, device=device)\n", "\n", " fc1_shape = snapshots[0]['fc1'].shape\n", " fc2_shape = snapshots[0]['fc2'].shape\n", " fc1_numel = snapshots[0]['fc1'].numel()\n", "\n", " weight_vecs = torch.stack([\n", " torch.cat([s['fc1'].flatten(), s['fc2'].flatten()]) for s in snapshots\n", " ])\n", " mean_w = weight_vecs.mean(dim=0)\n", " centered = weight_vecs - mean_w\n", " U, S, V = torch.linalg.svd(centered, full_matrices=False)\n", " pc1, pc2 = V[0], V[1]\n", " explained = (S ** 2 / (S ** 2).sum())\n", " print(f\" PC1 {explained[0]*100:.1f}%, PC2 {explained[1]*100:.1f}%, \"\n", " f\"top-2 total {(explained[0]+explained[1])*100:.1f}%\")\n", "\n", " proj = (centered @ V[:2].T).numpy()\n", "\n", " x_min, x_max = proj[:, 0].min(), proj[:, 0].max()\n", " y_min, y_max = proj[:, 1].min(), proj[:, 1].max()\n", " xr = max(x_max - x_min, 1e-8)\n", " yr = max(y_max - y_min, 1e-8)\n", " x_min -= margin * xr; x_max += margin * xr\n", " y_min -= margin * yr; y_max += margin * yr\n", " xs = np.linspace(x_min, x_max, n_grid)\n", " ys = np.linspace(y_min, y_max, n_grid)\n", " X, Y = np.meshgrid(xs, ys)\n", "\n", " work_model = baseMLP(p, hidden_dim).to(device)\n", " work_model.eval()\n", " criterion = nn.MSELoss()\n", " mean_w_d = mean_w.to(device)\n", " pc1_d = pc1.to(device)\n", " pc2_d = pc2.to(device)\n", "\n", " def loss_at(a, b):\n", " with torch.no_grad():\n", " w_flat = mean_w_d + float(a) * pc1_d + float(b) * pc2_d\n", " work_model.fc1.weight.data = w_flat[:fc1_numel].view(fc1_shape)\n", " work_model.fc2.weight.data = w_flat[fc1_numel:].view(fc2_shape)\n", " tr = criterion(work_model(data['x_train']), data['y_train']).item()\n", " te = criterion(work_model(data['x_test']), data['y_test']).item()\n", " return tr, te\n", "\n", " print(f\" computing {n_grid}x{n_grid} grid losses...\")\n", " train_grid = np.zeros((n_grid, n_grid))\n", " test_grid = np.zeros((n_grid, n_grid))\n", " for j, b in enumerate(ys):\n", " for i, a in enumerate(xs):\n", " tr, te = loss_at(a, b)\n", " train_grid[j, i] = tr\n", " test_grid[j, i] = te\n", "\n", " traj_train = np.array([loss_at(*p2)[0] for p2 in proj])\n", " traj_test = np.array([loss_at(*p2)[1] for p2 in proj])\n", "\n", " #log10 transform with floor so log(0) doesn't blow up\n", " log_train_grid = np.log10(np.maximum(train_grid, log_floor))\n", " log_test_grid = np.log10(np.maximum(test_grid, log_floor))\n", " log_traj_train = np.log10(np.maximum(traj_train, log_floor))\n", " log_traj_test = np.log10(np.maximum(traj_test, log_floor))\n", "\n", " fig = plt.figure(figsize=(15, 6))\n", " ax_tr = fig.add_subplot(121, projection='3d')\n", " ax_te = fig.add_subplot(122, projection='3d')\n", "\n", " surf_kwargs = dict(cmap='viridis', alpha=0.7, edgecolor='none', antialiased=True)\n", " ax_tr.plot_surface(X, Y, log_train_grid, **surf_kwargs)\n", " ax_te.plot_surface(X, Y, log_test_grid, **surf_kwargs)\n", " ax_tr.set_xlabel('PC1'); ax_tr.set_ylabel('PC2'); ax_tr.set_zlabel('log10 train loss')\n", " ax_tr.set_title(f'{op} train loss landscape (log scale)')\n", " ax_te.set_xlabel('PC1'); ax_te.set_ylabel('PC2'); ax_te.set_zlabel('log10 test loss')\n", " ax_te.set_title(f'{op} test loss landscape (log scale)')\n", "\n", " line_tr, = ax_tr.plot([], [], [], 'r-', linewidth=2)\n", " point_tr, = ax_tr.plot([], [], [], 'ro', markersize=8)\n", " line_te, = ax_te.plot([], [], [], 'r-', linewidth=2)\n", " point_te, = ax_te.plot([], [], [], 'ro', markersize=8)\n", " title = fig.suptitle(f'{op}, epoch {snapshots[0][\"epoch\"]}')\n", "\n", " def update(frame):\n", " line_tr.set_data_3d(proj[:frame+1, 0], proj[:frame+1, 1], log_traj_train[:frame+1])\n", " point_tr.set_data_3d([proj[frame, 0]], [proj[frame, 1]], [log_traj_train[frame]])\n", " line_te.set_data_3d(proj[:frame+1, 0], proj[:frame+1, 1], log_traj_test[:frame+1])\n", " point_te.set_data_3d([proj[frame, 0]], [proj[frame, 1]], [log_traj_test[frame]])\n", " title.set_text(f'{op}, epoch {snapshots[frame][\"epoch\"]}')\n", " return line_tr, point_tr, line_te, point_te, title\n", "\n", " anim = FuncAnimation(fig, update, frames=len(proj), interval=120, blit=False)\n", "\n", " save_dir.mkdir(parents=True, exist_ok=True)\n", " save_path = save_dir / f'landscape_{op}.gif'\n", " print(f\" saving GIF to {save_path}...\")\n", " anim.save(save_path, writer='pillow', fps=8, dpi=100)\n", " plt.close(fig)\n", " return anim\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#addition landscape\n", "\n", "anim_add = loss_landscape_animation(result_add)\n", "HTML(anim_add.to_jshtml())\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#2x2 of 3D surfaces showing the learned correct-class logit L(a, b), one panel per op\n", "#mul/div use the discrete-log basis (where their analytical solution is periodic)\n", "#we project onto the top 30 Fourier modes by magnitude so the wave structure is visible without high-frequency noise\n", "\n", "ops = ['add', 'sub', 'mul', 'div']\n", "fig = plt.figure(figsize=(16, 13))\n", "for i, op in enumerate(ops):\n", " result = load_run(op, 0)\n", " snapshots = result['weight_snapshots']\n", " p = result['config']['p']\n", " hidden_dim = result['config']['hidden_dim']\n", " pi_log = result['pi_log']\n", " use_log = op in ('mul', 'div')\n", "\n", " if op in ('add', 'sub'):\n", " a = torch.arange(p).repeat_interleave(p).to(device)\n", " b = torch.arange(p).repeat(p).to(device)\n", " n_a, n_b = p, p\n", " else:\n", " a = torch.arange(p).repeat_interleave(p - 1).to(device)\n", " b = torch.arange(1, p).repeat(p).to(device)\n", " n_a, n_b = p, p - 1\n", " x_all = torch.cat([F.one_hot(a, p), F.one_hot(b, p)], dim=1).float()\n", "\n", " if op == 'add':\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p)\n", " bi = torch.arange(p).unsqueeze(0).expand(p, p)\n", " cls = (ai + bi) % p\n", " elif op == 'sub':\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p)\n", " bi = torch.arange(p).unsqueeze(0).expand(p, p)\n", " cls = (ai - bi) % p\n", " elif op == 'mul':\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p - 1)\n", " bi = torch.arange(1, p).unsqueeze(0).expand(p, p - 1)\n", " cls = (ai * bi) % p\n", " elif op == 'div':\n", " b_inv_row = torch.tensor([pow(int(bi), p - 2, p) for bi in range(1, p)])\n", " ai = torch.arange(p).unsqueeze(1).expand(p, p - 1)\n", " bi = b_inv_row.unsqueeze(0).expand(p, p - 1)\n", " cls = (ai * bi) % p\n", "\n", " work_model = baseMLP(p, hidden_dim).to(device)\n", " snap = snapshots[-1]\n", " work_model.fc1.weight.data = snap['fc1'].to(device)\n", " work_model.fc2.weight.data = snap['fc2'].to(device)\n", " work_model.eval()\n", " with torch.no_grad():\n", " logits = work_model(x_all).reshape(n_a, n_b, p).cpu()\n", " Z = logits.gather(2, cls.unsqueeze(2)).squeeze(2)\n", " if use_log:\n", " Z = Z.index_select(0, pi_log).index_select(1, pi_log - 1)\n", "\n", " #project onto top-k Fourier modes by magnitude, drop dc bin so the plot is centered\n", " #conjugate symmetric pairs are kept naturally because their magnitudes match\n", " top_k = 30\n", " spec = torch.fft.fft2(Z)\n", " spec[0, 0] = 0\n", " mag = spec.abs()\n", " flat_thresh = torch.topk(mag.flatten(), top_k).values[-1]\n", " mask = (mag >= flat_thresh).to(spec.dtype)\n", " Z = torch.fft.ifft2(spec * mask).real.numpy()\n", "\n", " H, W = Z.shape\n", " A_grid, B_grid = np.meshgrid(np.arange(W), np.arange(H))\n", "\n", " ax = fig.add_subplot(2, 2, i + 1, projection='3d')\n", " surf = ax.plot_surface(A_grid, B_grid, Z, cmap='viridis', linewidth=0, antialiased=True, rcount=50, ccount=50)\n", " basis_label = 'log basis' if use_log else 'natural basis'\n", " ax.set_title(f'{op}, {basis_label}, epoch {snap[\"epoch\"]}')\n", " ax.set_xlabel('b')\n", " ax.set_ylabel('a')\n", " ax.set_zlabel('logit deviation')\n", " ax.view_init(elev=35, azim=-60)\n", " fig.colorbar(surf, ax=ax, shrink=0.5, pad=0.1)\n", "\n", "fig.suptitle('Learned Fourier features: correct-class logit surface (top 30 modes)', y=1.0)\n", "plt.tight_layout()\n", "out_path = Path('figures') / 'comparison' / 'fig9_fourier_features_3d_clean.png'\n", "out_path.parent.mkdir(parents=True, exist_ok=True)\n", "fig.savefig(out_path, dpi=150, bbox_inches='tight')\n", "plt.show()\n", "print(f'saved {out_path}')\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.13" } }, "nbformat": 4, "nbformat_minor": 4 }