Merge pull request #73 from cybertronai/feat/muon-review-prompt

docs: Add agent prompt system + Task 9 (Muon optimizer review)

Author: zh4ngx

docs/agent-prompts/muon-review.md

@@ -0,0 +1,73 @@
+# Muon Optimizer Literature Review
+
+You are contributing to the **Sutro Group**, a research lab studying energy-efficient AI training. Your task is a literature review of the **Muon optimizer**.
+
+## Project context
+
+Read these files first:
+- `DISCOVERIES.md` — what's already known
+- `LAB.md` — lab rules (important: do NOT modify measurement code, tracker.py, cache_tracker.py, data.py, config.py, harness.py)
+- `findings/_template.md` — reference for expected report format
+
+- We optimize for **energy efficiency**, not just accuracy or speed
+- Our benchmark: **sparse parity** (n=20 bits, k=3 secret, 17 noise)
+- Our metric: **ByteDMD** (https://github.com/cybertronai/ByteDMD) — successor to DMC. Tracks memory at **byte level** (not per-element), using integer arithmetic. Reference implementation has been hardened against escape hatches (agents were bypassing TrackedArray via np.asarray() → Python ints). See also: https://github.com/cybertronai/ByteDMD-examples for test cases.
+- **Baseline SGD**: ~0.12s wall-clock, high DMD (exact ByteDMD score unmeasured — you'd need to run the ByteDMD tracer)
+- **Best known**: GF(2) Gaussian Elimination (~509us, low byte-level access count), KM-min (~1ms)
+
+## Task: Research the Muon optimizer
+
+### 1. Read the Muon paper
+
+Primary source: https://kellerjordan.github.io/posts/muon/
+
+Key facts to extract:
+- The algorithm: Newton-Schulz iteration on SGD momentum updates to orthogonalize gradients
+- How it replaces Adam's element-wise adaptation with matrix orthogonalization
+- Performance claims (CIFAR-10, NanoGPT, 1.5B parameter models)
+- Computational overhead (<1% FLOP overhead claimed)
+- Scope: only for 2D hidden-layer weights, paired with AdamW for embeddings/biases
+
+### 2. Analyze relevance to energy efficiency
+
+Study the ByteDMD metric: https://github.com/cybertronai/ByteDMD
+
+Investigate:
+- Does Newton-Schulz iteration **reduce byte-level data movement** compared to Adam's moment tracking (first + second moment buffers)?
+- What's the **FLOP/memory tradeoff**? Does it reduce byte accesses at the cost of more compute?
+- Does orthogonalization change **reuse distance patterns** compared to Adam or SGD?
+- Would Muon help on **small networks** (hidden=200, our sparse parity setup) or only large LLMs?
+
+### 3. Write findings
+
+Create `findings/exp_muon_review.md` following this structure:
+
+```markdown
+# Muon Optimizer — Literature Review
+
+## Hypothesis
+[What we expect Muon to help with (or not) for energy efficiency]
+
+## Key Facts from Paper
+- [5-10 bullets on algorithm, results, limitations]
+
+## Relevance to Sutro Group
+[Would Muon improve energy efficiency on sparse parity? Why/why not?]
+
+## Comparison to Our Methods
+[How Muon relates to our best methods: SGD, GF(2), KM-min]
+
+## Open Questions
+[What we'd need to test experimentally to know for sure]
+
+## References
+- [Links to paper, code, follow-ups]
+```
+
+## Rules
+
+- **DO NOT** modify any experiment/run code (tracker.py, cache_tracker.py, data.py, config.py, harness.py, fast.py, src/)
+- **DO NOT** run experiments — this is a literature review only
+- Write findings to `findings/exp_muon_review.md`
+- Check `DISCOVERIES.md` first to avoid repeating known results
+- Do NOT change LAB.md, DISCOVERIES.md, CLAUDE.md, CODEX.md, or any other project configuration

docs/findings/exp_muon_review.md

@@ -0,0 +1,33 @@
+# Muon Optimizer — Literature Review
+
+## Hypothesis
+If we substitute standard optimizers with Muon's Newton-Schulz matrix orthogonalization for the sparse parity task, we expect it to *worsen* energy efficiency (increase ByteDMD) compared to SGD. Newton-Schulz iterations require multiple heavy matrix multiplications per step, generating high byte-level data movement through intermediate matrices. Furthermore, its overhead scales inversely with batch size, making it extremely inefficient for our single-sample or small-batch regimes.
+
+## Key Facts from Paper
+- **Algorithm:** Muon orthogonalizes the update matrix from SGD-momentum using a Newton-Schulz (NS) iteration, replacing the update with the nearest semi-orthogonal matrix ($UV^\top$ from SVD).
+- **Scope:** It applies only to 2D hidden layer parameters (e.g., $W_1$). Embedding, bias, and output parameters must still be optimized by AdamW/standard methods.
+- **Memory footprint:** Muon uses SGD-momentum (one buffer), substituting Adam's second moment accumulator with NS iteration matrix temporaries ($A, B, X$) computed in bfloat16.
+- **Computational overhead:** The overhead is bounded by $Tm/B$ where $T \approx 5$ is NS steps, $m$ is model dimension, and $B$ is batch size in tokens. For large transformers ($B > 500,000$), FLOP overhead is < 1%.
+- **Performance:** Set speedrun records for training NanoGPT and CIFAR-10. It reduced the training time of a 1.5B parameter transformer on HellaSwag from 13.3 to 10 hours.
+- **Limitations:** Muon has a slower per-step wallclock time than standard optimizers; its success hinges on reducing the overall number of training steps (epochs) significantly.
+
+## Relevance to Sutro Group
+- **Energy efficiency (ByteDMD):** Muon is highly detrimental to ByteDMD. The 5-step NS iteration performs operations like $A = XX^\top$ and $B = bA + cA^2$, aggressively reading and writing matrices. Each matmul natively incurs a massive ByteDMD cost (as intermediate values flush the LRU stack), ruining the cache locality that makes single-sample SGD efficient.
+- **Small batch scaling:** Our sparse parity setups use small networks ($m=200$ to $1000$) with tiny batch sizes ($B=1$ to $32$). The FLOP overhead $Tm/B$ becomes astronomical, e.g., $5 \times 1000 / 32 \approx 156\times$ the baseline FLOPs.
+- **Problem structure:** Sparse parity's difficulty is primarily searching for exactly $k$ secret bits (a combinatorial feature extraction problem), not resolving a poorly conditioned loss landscape in hidden layer activations, which is what orthogonalization fixes. 
+
+## Comparison to Our Methods
+- **SGD:** Single-sample SGD has excellent L1-cache locality yielding low ByteDMD/ARD. Muon would require accumulating a batch update and running an $O(m^3)$ matrix orthogonalization, destroying our $100\%$ L1 hit rates and resulting in a much higher DMC.
+- **EGD (Egalitarian Gradient Descent):** Existing experiments (`exp_egd`) show that gradient SVD-normalization cuts epochs by ~2x but its per-step SVD overhead makes wall time 12% worse. Muon's NS iteration is an approximation to SVD, meaning it would run into the same exact bottleneck for this task: optimizing the step count at the expense of per-step computational & data movement overhead.
+- **GF(2) & KM-min:** These perform extremely well on parity (DMC ~3500-8600, ~1ms wall-clock). Muon, as a gradient-based iterative optimizer, fundamentally has no path to beating $O(n)$ or deterministic matrix algebra solvers on purely logical tasks.
+
+## Open Questions
+- Does Newton-Schulz iteration yield *any* net epoch reduction for 20-bit sparse parity, or does orthogonalizing the gradient matrices corrupt the needle-in-a-haystack subsets of weights that are actually learning?
+- Can the Newton-Schulz step be computed iteratively in a block-tiled manner to fit completely within the L1 cache, thereby bounding its otherwise catastrophic ByteDMD cost?
+- What is the empirical ByteDMD cost of running 1 step of Muon vs. 1 step of SGD on a $1000 \times 20$ matrix using `bytedmd`'s tracker?
+
+## References
+- Primary Source: [Muon Blog Post](https://kellerjordan.github.io/posts/muon/)
+- Code: [KellerJordan/Muon GitHub](https://github.com/KellerJordan/Muon)
+- Metric: [ByteDMD Metric Definition](https://github.com/cybertronai/ByteDMD)
+- Internal Context: `DISCOVERIES.md` (`exp_egd` and scaling properties)

docs/tasks/009-muon-review.md

@@ -0,0 +1,29 @@
+# Task 9: Muon Optimizer Literature Review
+
+**Priority**: MEDIUM
+**Status**: DONE
+**Agent**: Antigravity
+**Source**: Yaroslav's Information Bottleneck podcast interview (April 1), search_space.yaml (listed but untested)
+
+## Context
+
+Yaroslav mentioned Muon in the Information Bottleneck podcast as an example. The lab's learning-guide.md notes: "Muon (first optimizer to beat Adam in 10 years) was discovered on 2-second CIFAR runs."
+
+Muon appears in `research/search_space.yaml` but was never tested. The question: does an optimizer that orthogonalizes gradients (Newton-Schulz iteration) reduce memory access patterns compared to Adam's moment tracking? Relevant to our ByteDMD metric.
+
+## Tasks
+
+- [x] Read the Muon paper (https://kellerjordan.github.io/posts/muon/)
+- [x] Study ByteDMD metric (https://github.com/cybertronai/ByteDMD)
+- [x] Analyze whether Newton-Schulz iteration reduces byte-level data movement vs Adam
+- [x] Assess whether Muon helps on small networks (hidden=200) or only large LLMs
+- [x] Write findings to `docs/findings/exp_muon_review.md` using the agent prompt scaffold
+
+## References
+
+- Agent prompt: [docs/agent-prompts/muon-review.md](../agent-prompts/muon-review.md)
+- Muon paper: https://kellerjordan.github.io/posts/muon/
+- ByteDMD: https://github.com/cybertronai/ByteDMD
+- ByteDMD test cases: https://github.com/cybertronai/ByteDMD-examples
+- Learning guide mention: [docs/learning-guide.md](../learning-guide.md)
+- Search space: [research/search_space.yaml](../../research/search_space.yaml)

docs/tasks/INDEX.md

@@ -12,5 +12,5 @@ Feedback from Meetings #8-11, Telegram, and Google Docs. Updated 2026-04-19.
 | 6 | Prep homework for next Monday (present results/process) | HIGH | SUPERSEDED | [006-homework-monday.md](006-homework-monday.md) |
 | 7 | Homework for Meeting #10: DMC baseline + optimization | HIGH | DONE | [007-homework-meeting10.md](007-homework-meeting10.md) |
 | 8 | Pre-experiment plan: DMC comparison, RL env, license | MEDIUM | IN PROGRESS | [008-pre-experiment-plan.md](008-pre-experiment-plan.md) |
-| 9 | Muon optimizer review (ByteDMD-aware agent prompt) | MEDIUM | OPEN (PR #73) | docs/agent-prompts/muon-review.md |
+| 9 | Muon optimizer literature review (ByteDMD-aware) | MEDIUM | DONE | [009-muon-review.md](009-muon-review.md) |
 | 10 | ASI-Evolve paper review — autonomous loop lessons | HIGH | IN PROGRESS | [010-asi-evolve-integration.md](010-asi-evolve-integration.md) |