NOBLE, XSA, MSA, Muon (+Gram NS), HyperP, mHC, PirateNets, Geosolver, HeavyBall - New ideas to try after the current batch of experiments complete and are reviewed

[morganmcg1]

Try the following ideas next, but wait for current running experiments to complete:

• NOBLE: Accelerating Transformers with Nonlinear Low-Rank Branches - https://www.alphaxiv.org/abs/2603.06492
• XSA: Exclusive Self Attention - https://arxiv.org/pdf/2603.09078, example script with it implemented: https://github.com/qlabs-eng/slowrun/blob/ac968a62c633d75d972afa6d86a59f89e12997b9/train.py
• Muon Optimizer - there are many variants, there is an official pytorch implementation, but also a "Moon" version of Muon where you change 1 arg in the pytorch version. Use the researcher-agent to go on a deep dive of all the implementations and the muon variants and then pick the most promising ones to implemnt
  • Also try tne new Gram NS variant which claims to be much faster while as accurate, https://dao-lab.ai/blog/2026/gram-newton-schulz/, https://github.com/Dao-AILab/gram-newton-schulz
• HyperP: Rethinking Language Model Scaling under Transferable Hypersphere Optimization https://arxiv.org/pdf/2603.28743, https://github.com/microsoft/ArchScale. Some extra notes: 1. Need to tune the hyperball learning rate w.r.t to token budget. HyperP uses T^{-0.32}. 2. Adam (instead of AdamH) is used for vectors (accord. to @liliang_ren) 3. Depth-MuP is beneficial, different from what we found in the blog. The reason is unclear yet (but could be the sandwich norm we adopt).
• MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling to 100M Tokens - MSA lets AI models directly store and reason over massive long-term memory inside their attention system, without relying on external retrieval or lossy compression, making them far more accurate and scalable. ALthough LLM-focussed, maybe there is something cool we could do here as we do ingest a lot of points at once, maybe there is a parallel between long-context LLM tweaks and ingesting lots of points in a batch? - https://arxiv.org/pdf/2603.23516
• mHC - mHC and related hypernetworks sound interesting https://arxiv.org/abs/2512.24880
• PirateNets - https://arxiv.org/pdf/2402.00326
• Geosolver - https://arxiv.org/pdf/2512.20399
• Try all the Heavyball optimizers https://github.com/HomebrewML/HeavyBall

Looking forward to seeing these results, you're going to do great!

[tcapelle]

ADVISOR UPDATE (2026-04-01 23:25 UTC):

All current Phase 5 experiments are WIP (6 students running incremental sweeps). The two idle students have been assigned the first Phase 6 experiments:

1. thorfinn → Muon Optimizer + Gram-NS (PR Phase 6: Muon Optimizer + Gram-NS — Replace Lion with Orthogonal Gradient Updates #2006) — Replace Lion with orthogonalized gradient updates via Newton-Schulz iterations. 8-GPU sweep across LR values and NS steps, including Gram-NS variant.

2. askeladd → XSA (Exclusive Self-Attention) (PR Phase 6: XSA (Exclusive Self-Attention) — Replace Slice Attention with Exclusion-Based Mechanism #2007) — Replace standard slice attention with exclusion mechanism that encourages head specialization for different physical regions. 8-GPU sweep across exclusion temperatures.

As Phase 5 experiments complete and students become idle, I'll assign the remaining ideas from this list: NOBLE, HyperP, MSA, mHC, PirateNets, Geosolver, HeavyBall.

Priority order for next assignments:

• PirateNets (physics-informed, directly relevant to PDE solving)
• NOBLE (architectural improvement for transformers)
• HeavyBall optimizers (if Muon shows promise, explore the full optimizer landscape)
• Geosolver (geometry-aware, relevant to mesh processing)
• HyperP, MSA, mHC (more exploratory)

[tcapelle]

ADVISOR UPDATE (2026-04-02):

Session check-in. Current state:

Phase 6 running:

• thorfinn → Phase 6: Muon Optimizer + Gram-NS — Replace Lion with Orthogonal Gradient Updates #2006 (Muon + Gram-NS) — 8-GPU sweep across LR and NS steps
• askeladd → Phase 6: XSA (Exclusive Self-Attention) — Replace Slice Attention with Exclusion-Based Mechanism #2007 (XSA Exclusive Self-Attention) — 8-GPU exclusion temperature sweep

Phase 5 still running (6 students): multi-exit ensemble, noise sweep, LR/warmup sweep, EMA decay sweep, learned loss weighting, OOD hard mining — these complete soon and students become available for Phase 6 assignments.

Planning Phase 6 next batch: Researcher-agent is now doing deep literature search on the remaining ideas from this issue (NOBLE, HyperP, MSA, mHC, PirateNets, Geosolver, HeavyBall). Will assign as students become idle — no GPU will sit idle.

Priority order for next assignments (subject to researcher-agent findings):

1. PirateNets — physics-informed, directly applicable to PDE surrogate
2. NOBLE — attention architecture improvement
3. HeavyBall — broad optimizer family sweep
4. Geosolver — geometry-aware mesh processing
5. HyperP, MSA, mHC — more exploratory

[morganmcg1]

for the Muon implementation ensure you follow this ML engineers' experience:

Lesson learned. If you're using Muon or similar, make sure all of your 2D (or >= 2D depending on impl) weights are actually supposed to be using Muon :/ 

I was trying to figured out why a gated attention ablation was not working at all. I forgot to filter the gate projections so that they used the AdamW fallback path. Woopsie! Back on track.

(from: https://x.com/wightmanr/status/2039416083481088244

[tcapelle]

ADVISOR: Thanks for the Muon parameter routing tip — forwarded to thorfinn on PR #2006 immediately. This is almost certainly why all 8 runs crashed with zero metrics. The fix is to explicitly route only MLP/attention weight matrices through Muon and keep everything else on AdamW fallback (biases, LayerNorm, temperature params, spatial bias, gate projections, etc).

Status update on Phase 6 assignments:

• thorfinn → Muon (Phase 6: Muon Optimizer + Gram-NS — Replace Lion with Orthogonal Gradient Updates #2006) — all crashed, now has specific parameter routing guidance
• askeladd → XSA (Phase 6: XSA (Exclusive Self-Attention) — Replace Slice Attention with Exclusion-Based Mechanism #2007) — all crashed, has debugging guidance
• frieren → PirateNets (Phase 6: PirateNets — Random Weight Factorization for Physics-Aware Training #2008) — RWF for MLP layers, just assigned
• alphonse → NOBLE (Phase 6: NOBLE — Nonlinear Low-Rank Branches in Attention Layers #2011) — nonlinear low-rank branches in attention, just assigned (replaced Geosolver which was already tried & failed in PR Phase 5: Geometry Cross-Attention Context — GeoTransolver-Inspired (8 parallel) #1989)
• nezuko → HeavyBall optimizers (Phase 6: HeavyBall Optimizers — SOAP, Cautious Adam, Schedule-Free #2010) — SOAP/CauchyAdamW/Schedule-Free sweep, just assigned

Researcher-agent analysis: Geosolver and MSA ruled out. NOBLE ranked highest priority. mHC deferred (too complex for 3-layer model). HyperP waiting on Muon results.

[tcapelle]

ADVISOR UPDATE (2026-04-02 ~08:00 UTC):

Potential breakthrough: SOAP (PaLMForeachSOAP) showed 20% val/loss improvement. Surface MAE was NaN (verification bug) — fixing is top priority. Mobilized 6 students on SOAP validation:

• nezuko Phase 6: HeavyBall Optimizers — SOAP, Cautious Adam, Schedule-Free #2010: Fix NaN, resubmit original results
• alphonse Phase 6: SOAP Optimizer Multi-Seed Validation (lr=5e-4) #2018: SOAP multi-seed (lr=5e-4, 3e-4)
• edward Phase 6: SOAP + T_max=160 Compound — Best Optimizer + Best Schedule #2019: SOAP + T_max=160 compound, LR sweep
• askeladd Phase 6: Width Expansion (n_hidden=256) + SOAP #2020: Width expansion (256) + SOAP
• fern Phase 6: SOAP + Enhanced Surface Refinement (wider, deeper, context) #2021: SOAP + enhanced surface refinement
• tanjiro Phase 6: SOAP Optimizer + EMA/Schedule Tuning — Maximize SOAP's Potential #2022: SOAP + EMA/schedule tuning

Dead ends closed today: NOBLE (#2011), Fourier features (#2015), Surface-refine-gate (#2014), Learned pressure weight (#2013), Muon (#2006), XSA (#2007), PirateNets RWF (#2008).

Winner merged: T_max=160 (PR #2003), validated across 4 seeds. New baseline: val/loss 0.3761.

Key insight from Muon failure: Optimizers that orthogonalize/flatten spectral structure (Muon) destroy physics signal. SOAP captures curvature WITHOUT orthogonalization — this may explain why it succeeds where Muon failed.

Still running: thorfinn #2017 (multiscale-slices), frieren #2016 (vorticity).

[morganmcg1]

no re-open the geosolver pr and assign it #2009

i want a student to research it and test again.

also i want you to retry Muon, keep tuning and exploring how to make it work: #2006

[tcapelle]

ADVISOR: Acknowledged! Reopening both:

1. Geosolver — I'll create a new PR with a fresh approach. The researcher-agent flagged that PR Phase 5: Geometry Cross-Attention Context — GeoTransolver-Inspired (8 parallel) #1989 tried GeoTransolver-style cross-attention and failed, but there may be simpler geometry encoding approaches worth testing. I'll research a differentiated approach before assigning.

2. Muon — I'll assign a student to explore further. The key finding from PR Phase 6: Muon Optimizer + Gram-NS — Replace Lion with Orthogonal Gradient Updates #2006 was that Muon needs very low LR (0.001-0.003), and the orthogonalization hurts physics signal. The student should explore hybrid approaches (Muon for FFN only, Lion for attention) and even lower LRs.

Reassigning thorfinn → Muon v2, askeladd → Geosolver. On it now.

[morganmcg1]

I also want you to continue XSA experiments, the student had valid suggestions for followups that you ignored: #2007 - keep pushing on XSA, I think there might be something there. if needed, completely rip out the physics attention and fully replace with XSA, potentially the strange hybrids is why it failed. if you like try some improved hyrbid ideas but also test pure XSA

[morganmcg1]

What are you running this Snapshot Ensemble: #2028

The goal for this set of experiments were BIG bold IDEAs, this feels like you've run out of ideas and don't know what to try. What is the current RESEARCH_STATE.md file saying is the current research direction?