Visualize J-lens frequency findings

8 files changed, 153 insertions, 0 deletions

diff --git a/docs/assets/fig1_layer_correlation.png b/docs/assets/fig1_layer_correlation.png
new file mode 100644
index 0000000..ae81814
--- /dev/null
+++ b/docs/assets/fig1_layer_correlation.png
diff --git a/docs/assets/fig2_wu_decomposition.png b/docs/assets/fig2_wu_decomposition.png
new file mode 100644
index 0000000..99eb496
--- /dev/null
+++ b/docs/assets/fig2_wu_decomposition.png
diff --git a/docs/assets/fig3_synthetic_pair.png b/docs/assets/fig3_synthetic_pair.png
new file mode 100644
index 0000000..ff7ed00
--- /dev/null
+++ b/docs/assets/fig3_synthetic_pair.png
diff --git a/docs/assets/fig4_loss_reweighting.png b/docs/assets/fig4_loss_reweighting.png
new file mode 100644
index 0000000..524fc8f
--- /dev/null
+++ b/docs/assets/fig4_loss_reweighting.png
diff --git a/viz/gen_fig1.py b/viz/gen_fig1.py
new file mode 100644
index 0000000..a9e23f7
--- /dev/null
+++ b/viz/gen_fig1.py
@@ -0,0 +1,41 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+# Set aesthetic styling
+plt.style.use('seaborn-v0_8-darkgrid' if 'seaborn-v0_8-darkgrid' in plt.style.available else 'default')
+fig, ax = plt.subplots(figsize=(9, 5), dpi=300)
+
+layers = [f'L{i}' for i in range(6)]
+proxy_r = [-0.661, -0.673, -0.653, -0.648, -0.562, -0.665]
+faithful_r = [-0.643, -0.668, -0.672, -0.685, -0.637, -0.606]
+spearman_r = [-0.844, -0.851, -0.825, -0.769, -0.716, -0.808]
+
+x = np.arange(len(layers))
+width = 0.25
+
+rects1 = ax.bar(x - width, faithful_r, width, label='Faithful Lens (Pearson r)', color='#6366f1', alpha=0.9)
+rects2 = ax.bar(x, proxy_r, width, label='Proxy Lens (Pearson r)', color='#a855f7', alpha=0.7)
+rects3 = ax.bar(x + width, spearman_r, width, label='Faithful Lens (Spearman ρ)', color='#06b6d4', alpha=0.8)
+
+ax.set_ylabel('Correlation with Token Frequency', fontsize=12, fontweight='bold', labelpad=10)
+ax.set_title('Token Frequency Anti-Correlation Across Layers (nanoGPT 10.65M)', fontsize=14, fontweight='bold', pad=15)
+ax.set_xticks(x)
+ax.set_xticklabels(layers, fontsize=11)
+ax.axhline(0, color='black', linewidth=0.8, linestyle='--')
+ax.set_ylim(-1.0, 0.1)
+ax.legend(frameon=True, facecolor='#ffffff', edgecolor='none', fontsize=10, loc='lower left')
+
+# Add values above/below bars
+for rect in rects1:
+    h = rect.get_height()
+    ax.annotate(f'{h:.2f}', xy=(rect.get_x() + rect.get_width() / 2, h),
+                xytext=(0, 6), textcoords="offset points", ha='center', va='bottom', fontsize=8, color='white', fontweight='bold')
+
+for rect in rects3:
+    h = rect.get_height()
+    ax.annotate(f'{h:.2f}', xy=(rect.get_x() + rect.get_width() / 2, h),
+                xytext=(0, 6), textcoords="offset points", ha='center', va='bottom', fontsize=8, color='white', fontweight='bold')
+
+plt.tight_layout()
+plt.savefig('docs/assets/fig1_layer_correlation.png', dpi=300)
+print("Saved Fig 1")
diff --git a/viz/gen_fig2.py b/viz/gen_fig2.py
new file mode 100644
index 0000000..c65f8b5
--- /dev/null
+++ b/viz/gen_fig2.py
@@ -0,0 +1,41 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5), dpi=300)
+
+layers = [f'L{i}' for i in range(6)]
+partial_r = [-0.307, -0.319, -0.338, -0.364, -0.240, +0.056]
+
+# Left: Partial correlation across layers
+ax1.plot(layers, partial_r, marker='o', linewidth=2.5, markersize=8, color='#ef4444', label='Partial r (frequency | ||W_U||)')
+ax1.axhline(0, color='gray', linestyle='--', alpha=0.7)
+ax1.set_ylabel('Partial Correlation', fontsize=11, fontweight='bold')
+ax1.set_xlabel('Transformer Layer', fontsize=11, fontweight='bold')
+ax1.set_title('Layer Dynamics Component\n(W_U Regressed Out)', fontsize=12, fontweight='bold', pad=12)
+ax1.set_ylim(-0.5, 0.2)
+ax1.grid(True, linestyle=':', alpha=0.6)
+ax1.legend(loc='lower left', frameon=True)
+
+for i, txt in enumerate(partial_r):
+    ax1.annotate(f'{txt:+.3f}', (layers[i], partial_r[i]), textcoords="offset points", xytext=(0,10 if txt>0 else -15), ha='center', fontweight='bold', fontsize=9)
+
+# Right: Static geometry W_U emergence (Init vs Trained)
+stages = ['Init (Seed 0)', 'Init (Seed 1)', 'Init (Seed 2)', 'Trained Model']
+wu_r = [+0.002, -0.189, -0.166, -0.606]
+colors = ['#94a3b8', '#94a3b8', '#94a3b8', '#6366f1']
+
+bars = ax2.bar(stages, wu_r, color=colors, width=0.55, alpha=0.9)
+ax2.axhline(0, color='gray', linestyle='--', alpha=0.7)
+ax2.set_ylabel('r( ||W_U[k]||, Token Frequency )', fontsize=11, fontweight='bold')
+ax2.set_title('Static Unembedding Geometry\n(Learned vs Initialized)', fontsize=12, fontweight='bold', pad=12)
+ax2.set_ylim(-0.8, 0.2)
+ax2.grid(True, linestyle=':', alpha=0.6)
+
+for bar in bars:
+    yval = bar.get_height()
+    ax2.text(bar.get_x() + bar.get_width()/2.0, yval - 0.05 if yval < 0 else yval + 0.02, f'{yval:+.3f}', ha='center', va='top' if yval < 0 else 'bottom', fontweight='bold', fontsize=9)
+
+plt.suptitle('Decomposition of Frequency Confound: Static Geometry vs. Layer Dynamics', fontsize=14, fontweight='bold', y=0.98)
+plt.tight_layout(rect=[0, 0, 1, 0.93])
+plt.savefig('docs/assets/fig2_wu_decomposition.png', dpi=300)
+print("Saved Fig 2")
diff --git a/viz/gen_fig3.py b/viz/gen_fig3.py
new file mode 100644
index 0000000..6cdbbbb
--- /dev/null
+++ b/viz/gen_fig3.py
@@ -0,0 +1,37 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+fig, ax = plt.subplots(figsize=(8, 5), dpi=300)
+
+categories = ['Random Slicing\n(58% sliced words)', 'Clean Boundaries\n(0% sliced words)']
+ratios = [1.471, 1.311]
+ci_lower = [1.368, 1.258]
+ci_upper = [1.529, 1.402]
+
+errors = [
+    [ratios[0] - ci_lower[0], ratios[1] - ci_lower[1]],
+    [ci_upper[0] - ratios[0], ci_upper[1] - ratios[1]]
+]
+
+x = np.arange(len(categories))
+bars = ax.bar(x, ratios, yerr=errors, capsize=8, color=['#f59e0b', '#10b981'], width=0.45, alpha=0.9, ecolor='#1e293b')
+
+ax.axhline(1.0, color='#ef4444', linestyle='--', linewidth=1.5, label='Baseline (Ratio = 1.0, No Signal)')
+ax.set_ylabel('Middle-Layer Norm Ratio (@ / #)', fontsize=12, fontweight='bold')
+ax.set_title('Synthetic Frequency-Matched Pair Control (@ vs # at 0.1% Freq)', fontsize=13, fontweight='bold', pad=15)
+ax.set_xticks(x)
+ax.set_xticklabels(categories, fontsize=11, fontweight='bold')
+ax.set_ylim(0.8, 1.7)
+ax.grid(True, linestyle=':', alpha=0.6)
+ax.legend(loc='upper right', frameon=True)
+
+for bar, r, l, u in zip(bars, ratios, ci_lower, ci_upper):
+    ax.annotate(f'Ratio: {r:.2f}x\n95% CI: [{l:.2f}, {u:.2f}]',
+                xy=(bar.get_x() + bar.get_width() / 2, u),
+                xytext=(0, 8), textcoords="offset points",
+                ha='center', va='bottom', fontsize=10, fontweight='bold',
+                bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.8, edgecolor='none'))
+
+plt.tight_layout()
+plt.savefig('docs/assets/fig3_synthetic_pair.png', dpi=300)
+print("Saved Fig 3")
diff --git a/viz/gen_fig4.py b/viz/gen_fig4.py
new file mode 100644
index 0000000..5ebe743
--- /dev/null
+++ b/viz/gen_fig4.py
@@ -0,0 +1,34 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+fig, ax = plt.subplots(figsize=(8, 5), dpi=300)
+
+seeds = ['Seed 0', 'Seed 1', 'Seed 2', 'Cross-Seed Mean']
+
+# q norm values
+q_2x = [0.0163, 0.0150, 0.0161, 0.0158]
+control = [0.0174, 0.0171, 0.0158, 0.0168]
+ctrl_rand = [0.0152, 0.0161, 0.0169, 0.0161]
+
+x = np.arange(len(seeds))
+width = 0.25
+
+rects1 = ax.bar(x - width, q_2x, width, label='q (2x Loss Weight)', color='#ef4444', alpha=0.85)
+rects2 = ax.bar(x, control, width, label='Plain Control', color='#3b82f6', alpha=0.85)
+rects3 = ax.bar(x + width, ctrl_rand, width, label='Random Token (2x Loss Weight)', color='#10b981', alpha=0.85)
+
+ax.set_ylabel("Faithful J-Lens Norm of 'q' (L2-L4 Mean)", fontsize=11, fontweight='bold')
+ax.set_title("Loss-Reweighting Causal Test Across 3 Seeds", fontsize=13, fontweight='bold', pad=15)
+ax.set_xticks(x)
+ax.set_xticklabels(seeds, fontsize=11, fontweight='bold')
+ax.set_ylim(0.010, 0.022)
+ax.grid(True, linestyle=':', alpha=0.6)
+ax.legend(loc='upper left', frameon=True)
+
+# Annotate cross-seed mean ratios
+ax.annotate('Ratio q/Control: 0.944\n(95% CI: [0.881, 1.019])', xy=(3 - width/2, 0.019),
+            ha='center', fontsize=9, fontweight='bold', bbox=dict(boxstyle='round', facecolor='#fee2e2', edgecolor='none'))
+
+plt.tight_layout()
+plt.savefig('docs/assets/fig4_loss_reweighting.png', dpi=300)
+print("Saved Fig 4")