# Reproduction Plan: ICML 2026 Paper #22 — "Deep sequence models tend to memorize geometrically; it is unclear why"

**Paper:** arXiv:2510.26745 | **OpenReview:** 2NuCrYf8Ap (note: challenge references CT2tSmahVQ, but the correct forum ID is 2NuCrYf8Ap)  
**Authors:** Shahriar Noroozizadeh, Vaishnavh Nagarajan, Elan Rosenfeld, Sanjiv Kumar  
**Official Code:** https://github.com/shahriarnz14/geometric_memory  
**Hardware:** GB10 GPU (~85GB usable), 121GB RAM, Linux

---

## 0. Setup & Infrastructure

| Step | Action | Status |
|------|--------|--------|
| 0.1 | Clone repo from GitHub | ✅ Done |
| 0.2 | Set up conda environment (`geometric_memory.yml`) | Pending |
| 0.3 | Create Trackio logbook with metadata (arxiv_id, tags) | ✅ Logbook scaffolded, metadata pending |
| 0.4 | Verify CUDA/GPU availability | ✅ GB10 detected, CUDA 13.0 |
| 0.5 | Smoke test: run tiny graph notebook locally (CPU or small GPU) | Pending |

---

## 1. Claim 1 — Path-star memorization at scale

**Claim:** Deep sequence models (Transformers and Mamba) solve an in-weights path-star (l-fold composition) task by memorizing a fixed graph's edges during training and then predicting complete root-to-leaf paths, achieving near-perfect accuracy on graphs with up to 50,000 nodes and path lengths up to l=10 hops even when trained on only 75% of paths.

### Reproduction Strategy

| Step | Action | Scale | Method |
|------|--------|-------|--------|
| 1.1 | Reproduce Figure 3: Transformer accuracy on path-star graphs | **Full scale** (50k nodes, l=6, l=10) | Run `train_in_weights.py` with `mixed_full_path` recipe on star graphs |
| 1.2 | Reproduce Mamba results (Fig. 14) | **Full scale** (50k nodes, l=6) | Same config with `--model_family mamba` |
| 1.3 | Verify 75% path coverage claim | **Full scale** | Use existing dataset files (`train_7500`, `test_2500` naming convention) |
| 1.4 | Reproduce first-token-only training (hardest token) | **Full scale** | Use `mixed_hardest_token` recipe |
| 1.5 | Reproduce staged training baseline | **Full scale** | Use `staged_full_path` recipe |

### Commands (full scale, GPU)

```bash
# Claim 1.1: Transformer, 50k nodes, path length 6
python train_in_weights.py \
  --training_recipe mixed_full_path \
  --model_family gpt \
  --graph_type star \
  --star_degree 10000 \
  --star_subtree_degree 1 \
  --path_length 6 \
  --add_forward_edges \
  --add_backward_edges \
  --edge_memorization_epochs 2500 \
  --path_finetuning_epochs 10000 \
  --experiment_log_root ./experiment_logs \
  --no-use_wandb

# Claim 1.2: Transformer, 50k nodes, path length 10
python train_in_weights.py \
  --training_recipe mixed_full_path \
  --model_family gpt \
  --graph_type star \
  --star_degree 10000 \
  --star_subtree_degree 1 \
  --path_length 10 \
  --add_forward_edges \
  --add_backward_edges \
  --edge_memorization_epochs 2500 \
  --path_finetuning_epochs 10000 \
  --experiment_log_root ./experiment_logs \
  --no-use_wandb

# Claim 1.3: Mamba variant
python train_in_weights.py \
  --training_recipe mixed_full_path \
  --model_family mamba \
  --graph_type star \
  --star_degree 10000 \
  --star_subtree_degree 1 \
  --path_length 6 \
  --add_forward_edges \
  --add_backward_edges \
  --edge_memorization_epochs 2500 \
  --path_finetuning_epochs 10000 \
  --experiment_log_root ./experiment_logs \
  --no-use_wandb
```

### Expected results
- First-token accuracy near 100% on unseen paths (Fig. 3 shows ~100% at l=6, ~95%+ at l=10)
- Mamba should show similar behavior (Fig. 14)
- Staged training should also succeed but potentially with different dynamics

### HF Jobs plan
If local GB10 proves insufficient for full 50k-node runs:
- Submit GPU job with A100-80GB or H100
- Script: `train_in_weights.py` with same args
- Upload results to HF Bucket

---

## 2. Claim 2 — Contradiction with associative-memory prediction

**Claim:** Near-perfect performance contradicts the associative-memory prediction that l-fold composition should require exponentially Ω(exp(l)) more capacity or data.

### Reproduction Strategy

| Step | Action | Scale | Method |
|------|--------|-------|--------|
| 2.1 | Build associative-memory baseline model | **Small scale** (N=22 nodes) | Use `--freeze_token_embeddings` to force associative behavior |
| 2.2 | Measure accuracy vs. path length l for both models | **Small → medium** (l=2 to l=6) | Compare geometric vs. associative at each l |
| 2.3 | Quantify the exponential gap | **Small scale** | Plot accuracy vs. l for both models |
| 2.4 | Reproduce Section 2.3 analysis (hardest token contradiction) | **Small scale** | Train only on first-token loss, verify failure of associative model |

### Key analysis

The associative model stores edges as a lookup table: `f(u)[v] = Φ(u)^T W_assoc Φ(v)`. With orthogonal Φ, this only captures direct edges. To perform ℓ-hop reasoning, it would need to compose the lookup ℓ times — which requires exponential parameters.

The geometric model stores edges as: `f(u)[v] = Φ_geom(u)^T Φ_geom(v)` where Φ_geom encodes global structure. This makes ℓ-hop navigation a single dot-product.

**Verification:** Show that the associative model's accuracy drops exponentially with l, while the geometric model stays near-perfect.

### Commands

```bash
# Associative baseline: freeze embeddings (no geometry can form)
python train_in_weights.py \
  --training_recipe mixed_full_path \
  --model_family gpt \
  --graph_type star \
  --star_degree 3 \
  --star_subtree_degree 2 \
  --path_length 4 \
  --add_forward_edges \
  --add_backward_edges \
  --freeze_token_embeddings \
  --edge_memorization_epochs 2500 \
  --path_finetuning_epochs 10000 \
  --experiment_log_root ./experiment_logs \
  --no-use_wandb
```

### HF Jobs plan
Small-scale runs should fit on CPU or small GPU. Use HF Jobs only if local GPU memory is constrained.

---

## 3. Claim 3 — Embeddings organize into global geometric structure

**Claim:** Learned embeddings reflect global path identity rather than only local co-occurrence, visible as diagonal clustering in heatmaps and separated clusters in UMAP projections (Figure 6).

### Reproduction Strategy

| Step | Action | Scale | Method |
|------|--------|-------|--------|
| 3.1 | Train tiny Transformer on small graph | **Tiny** (N≤30 nodes) | Use `tiny_graphs_notebooks/experiment_notebooks/tiny_transformer.ipynb` |
| 3.2 | Generate embedding similarity heatmap (Fig. 1 equivalent) | **Tiny** | Extract final embeddings, compute dot-product matrix, plot |
| 3.3 | UMAP projection of embeddings (Fig. 6 equivalent) | **Tiny** | Run UMAP on embedding matrix, color by path membership |
| 3.4 | Compare Transformer vs. Node2Vec geometry strength | **Tiny** | Run Node2Vec notebook, compare embedding quality metrics |
| 3.5 | Associative vs. geometric memorization curves | **Tiny** | Plot accuracy of edge prediction vs. path prediction over training |

### Key notebooks to run

```bash
# Primary geometry notebook
jupyter nbconvert --to script --execute \
  tiny_graphs_notebooks/experiment_notebooks/tiny_transformer.ipynb \
  --output /tmp/repro_tiny_transformer.py

# Node2Vec comparison
jupyter nbconvert --to script --execute \
  tiny_graphs_notebooks/experiment_notebooks/tiny_node2vec.ipynb \
  --output /tmp/repro_tiny_node2vec.py
```

### Output artifacts
- Embedding similarity heatmaps (PNG)
- UMAP projections (PNG)
- Memorization curves (CSV + PNG)
- Comparison table: Transformer vs. Node2Vec geometry metrics

### HF Jobs plan
Tiny graph experiments should run locally (CPU or small GPU). No HF Jobs needed unless UMAP on larger graphs.

---

## 4. Claim 4 — Spectral bias in learned representations

**Claim:** Fiedler vectors of the graph align with learned Node2Vec-style embeddings (Section 4, Figure 7). Geometry stems from spectral bias that arises naturally from cross-entropy loss minimization.

### Reproduction Strategy

| Step | Action | Scale | Method |
|------|--------|-------|--------|
| 4.1 | Compute graph Laplacian eigenvectors (Fiedler vectors) | **Tiny** | Use numpy/scipy on the same graph used in training |
| 4.2 | Extract Node2Vec-style embeddings from 2-layer model | **Tiny** | Run `tiny_graphs_notebooks/experiment_notebooks/tiny_nn.ipynb` |
| 4.3 | Measure alignment between Laplacian eigenvectors and embeddings | **Tiny** | Compute projection/correlation |
| 4.4 | Reproduce spectral bias analysis (Appendix F) | **Tiny** | Run dynamics notebooks to verify eigenvalue convergence |
| 4.5 | Verify that Transformer embeddings also show spectral alignment | **Small** | Compare Transformer embeddings to Laplacian eigenvectors |

### Key analysis

The paper's Appendix F derives that in a 2-layer Node2Vec model, gradient descent on cross-entropy loss causes embeddings to converge to the top eigenvectors of the negative graph Laplacian. This is the "spectral bias" mechanism.

**Verification:** Show that:
1. Learned embeddings correlate with top Laplacian eigenvectors
2. This alignment increases over training
3. The effect is present in Node2Vec models and partially in Transformers

### Commands

```bash
# Spectral analysis notebook
python -c "
import sys
import nbformat
from nbconvert.preprocessors import ExecutePreprocessor

nb = nbformat.read('tiny_graphs_notebooks/experiment_notebooks/tiny_node2vec.ipynb', as_version=4)
ep = ExecutePreprocessor(timeout=600, kernel_name='python3')
ep.preprocess(nb, {'metadata': {'path': '.'}})
with open('tiny_graphs_notebooks/output_node2vec.ipynb', 'w') as f:
    nbformat.write(nb, f)
"
```

### HF Jobs plan
Tiny graph experiments should run locally. No HF Jobs needed.

---

## 5. Additional Experiments & Robustness Checks

| Step | Action | Rationale |
|------|--------|-----------|
| 5.1 | Reproduce results on non-star graphs (grid, cycle, irregular) | Test generality of claims |
| 5.2 | Reproduce self-edge ablation (experiment_notebooks_self_edges/) | Verify importance of self-edge supervision |
| 5.3 | Reproduce regularizer ablation (experiment_notebooks_regularizers/) | Test claim that geometry arises without explicit pressure |
| 5.4 | Scale path length to l=8 on medium graphs | Extend beyond paper's reported range |

---

## 6. Trackio Logbook Publishing

| Step | Action |
|------|--------|
| 6.1 | Set metadata.json with arxiv_id and tags |
| 6.2 | Create pages for each claim |
| 6.3 | Log each experiment run with `trackio logbook run` |
| 6.4 | Add markdown cells with findings |
| 6.5 | Add figure cells for plots |
| 6.6 | Write Executive Summary on Conclusion page |
| 6.7 | Build poster with Posterly |
| 6.8 | Pin executive summary and poster |
| 6.9 | Publish to HF Space |
| 6.10 | Upload artifacts to HF Bucket |

---

## 7. Risk Assessment & Mitigations

| Risk | Likelihood | Mitigation |
|------|-----------|------------|
| GB10 memory insufficient for 50k-node training | Medium | Use staged training, reduce batch size, or use HF Jobs with A100-80GB |
| Training takes too long locally | Medium | Use smaller graphs for initial runs, scale up on HF Jobs |
| Conda environment conflicts | Low | Use `geometric_memory.yml` as-is, verify dependencies |
| Tiny graph results don't generalize | Low | Paper's analysis is mathematically rigorous (Appendix F); should hold |
| OpenReview page behind Cloudflare | High | Use arXiv HTML version (already extracted) |

---

## 8. Timeline (estimated)

| Phase | Duration | Dependencies |
|-------|----------|-------------|
| Setup (env, logbook, smoke test) | ~30 min | None |
| Claim 1 (full-scale path-star) | ~2-4 hours | Setup |
| Claim 2 (associative baseline) | ~1 hour | Setup |
| Claim 3 (embedding geometry) | ~1 hour | Setup |
| Claim 4 (spectral bias) | ~1 hour | Setup |
| Additional experiments | ~2 hours | Setup |
| Logbook compilation & publishing | ~1 hour | All claims |
| **Total** | **~6-8 hours** | |

---

## 9. Deliverables

1. **Trackio logbook** published to HF Space
2. **Reproduction scripts** uploaded to HF (or kept in logbook artifacts)
3. **Plots**: accuracy curves, embedding heatmaps, UMAP projections, spectral alignment
4. **Executive summary** with outcome-first paragraph and scope/cost table
5. **Reproduction poster** (Posterly)
6. **HF Bucket** with all experiment artifacts