From c968bdbd3f0686e029be1b2ce54bc179b11d4832 Mon Sep 17 00:00:00 2001
From: Kauna <16511995+klei22@users.noreply.github.com>
Date: Sun, 14 Jun 2026 11:38:00 -0700
Subject: [PATCH 1/3] Add relu2max capped hypersphere exploration

---
 .../relu2max_capped_hypersphere_peri_ln.yaml  | 97 +++++++++++++++++++
 1 file changed, 97 insertions(+)
 create mode 100644 explorations/relu2max_capped_hypersphere_peri_ln.yaml

diff --git a/explorations/relu2max_capped_hypersphere_peri_ln.yaml b/explorations/relu2max_capped_hypersphere_peri_ln.yaml
new file mode 100644
index 0000000000..10e61319a3
--- /dev/null
+++ b/explorations/relu2max_capped_hypersphere_peri_ln.yaml
@@ -0,0 +1,97 @@
+# relu2max_capped_hypersphere_peri_ln.yaml
+# Sweep ReLU2Max attention with CappedHyperSphereNorm in peri-LN mode.
+---
+
+named_static_groups:
+  # QK Norm
+  - named_group: "qk_norm"
+    use_qk_norm: [true]
+    use_qk_norm_scale: [true]
+
+  # Norm Type: peri-LN mode
+  - named_group: "peri_ln"
+    use_pre_ln: [true]
+    use_peri_ln: [true]
+    use_post_ln: [false]
+
+  # Position Embeddings
+  - named_group: "rotary"
+    use_rotary_embeddings: [true]
+    use_abs_pos_embeddings: [false]
+
+  # Attention Softmax
+  - named_group: "relu2max"
+    softmax_variant_attn: ["relu2max"]
+
+  # Infinite Attention
+  - named_group: "infinite"
+    attention_variant: ["infinite"]
+    use_concat_heads: [true]
+
+  # MQA
+  - named_group: "mqa"
+    n_kv_group: [1]
+
+  # Head Dimension
+  - named_group: "hd_100"
+    n_qk_head_dim: [100]
+    n_v_head_dim: [100]
+
+  - named_group: "hd_150"
+    n_qk_head_dim: [150]
+    n_v_head_dim: [150]
+
+  - named_group: "hd_200"
+    n_qk_head_dim: [200]
+    n_v_head_dim: [200]
+
+  # Capped hypersphere norms for the residual-stream-facing projections.
+  - named_group: "capped_hypersphere_pair"
+    norm_variant_attn: ["cappedhyperspherenorm"]
+    norm_variant_output: ["cappedhyperspherenorm"]
+
+  - named_group: "capped_hypersphere_wte"
+    norm_variant_wte: ["cappedhyperspherenorm"]
+
+  - named_group: "no_wte_norm"
+
+named_variation_groups:
+  - named_group: "head_dim_var"
+    named_group_alternates: ["hd_100", "hd_150", "hd_200"]
+
+  - named_group: "wte_norm_var"
+    named_group_alternates: ["no_wte_norm", "capped_hypersphere_wte"]
+
+common_group:
+  dataset: ["minipile"]
+  eval_interval: [2500]
+  max_iters: [10000]
+  never_save_checkpoint: [true]
+  compile: [true]
+  log_rankme: [true]
+  log_areq: [true]
+
+parameter_groups:
+  - named_group_static:
+      - "qk_norm"
+      - "peri_ln"
+      - "rotary"
+      - "relu2max"
+      - "infinite"
+      - "mqa"
+      - "capped_hypersphere_pair"
+    n_head:
+      range:
+        start: 1
+        end: 12
+        step: 1
+    relu2max_divisor: [128.0, 256.0, 512.0]
+    # HyperSphereNorm-family knobs kept in this exploration so capped-hypersphere
+    # runs are tagged with the intended radius/gain/learning settings.
+    hsnorm_scale: [1.0]
+    hsnorm_gain: [false]
+    hsnorm_radius: [null]
+    hsnorm_radius_learning: [false]
+    named_group_variations:
+      - "head_dim_var"
+      - "wte_norm_var"

From f90c90cafbce2cd0c80fb47d207de6816894ca3e Mon Sep 17 00:00:00 2001
From: klei22 <kauna@avalonwest.tech>
Date: Sun, 14 Jun 2026 16:49:19 -0700
Subject: [PATCH 2/3] Add sweep for relu2max and capped norm compat

---
 .../relu2max_capped_hypersphere_peri_ln.yaml  | 69 +++++++++++++------
 variations/norm_variations.py                 | 12 +++-
 2 files changed, 57 insertions(+), 24 deletions(-)

diff --git a/explorations/relu2max_capped_hypersphere_peri_ln.yaml b/explorations/relu2max_capped_hypersphere_peri_ln.yaml
index 10e61319a3..d40970c9bf 100644
--- a/explorations/relu2max_capped_hypersphere_peri_ln.yaml
+++ b/explorations/relu2max_capped_hypersphere_peri_ln.yaml
@@ -14,6 +14,11 @@ named_static_groups:
     use_peri_ln: [true]
     use_post_ln: [false]
 
+  - named_group: "pre_ln"
+    use_pre_ln: [true]
+    use_peri_ln: [false]
+    use_post_ln: [false]
+
   # Position Embeddings
   - named_group: "rotary"
     use_rotary_embeddings: [true]
@@ -27,6 +32,7 @@ named_static_groups:
   - named_group: "infinite"
     attention_variant: ["infinite"]
     use_concat_heads: [true]
+    n_head: [3]
 
   # MQA
   - named_group: "mqa"
@@ -46,21 +52,21 @@ named_static_groups:
     n_v_head_dim: [200]
 
   # Capped hypersphere norms for the residual-stream-facing projections.
-  - named_group: "capped_hypersphere_pair"
+  - named_group: "capped_pair"
     norm_variant_attn: ["cappedhyperspherenorm"]
     norm_variant_output: ["cappedhyperspherenorm"]
 
-  - named_group: "capped_hypersphere_wte"
-    norm_variant_wte: ["cappedhyperspherenorm"]
+  - named_group: "capped_rmsnorm"
+    norm_variant_attn: ["cappedhyperspherenorm"]
+    norm_variant_output: ["rmsnorm"]
 
-  - named_group: "no_wte_norm"
+  - named_group: "rmsnorm"
+    norm_variant_attn: ["rmsnorm"]
+    norm_variant_output: ["rmsnorm"]
 
 named_variation_groups:
-  - named_group: "head_dim_var"
-    named_group_alternates: ["hd_100", "hd_150", "hd_200"]
-
   - named_group: "wte_norm_var"
-    named_group_alternates: ["no_wte_norm", "capped_hypersphere_wte"]
+    named_group_alternates: ["capped_rmsnorm", "capped_pair", "rmsnorm"]
 
 common_group:
   dataset: ["minipile"]
@@ -72,26 +78,45 @@ common_group:
   log_areq: [true]
 
 parameter_groups:
+  # Peri ln
+  - named_group_static:
+      - "qk_norm"
+      - "peri_ln"
+      - "rotary"
+      - "relu2max"
+      - "infinite"
+      - "hd_150"
+    hsnorm_scale: [1.0]
+    hsnorm_gain: [false, true]
+    hsnorm_radius_learning: [true]
+    named_group_variations:
+      - "wte_norm_var"
+  # Peri-LN WTE Norm
   - named_group_static:
       - "qk_norm"
       - "peri_ln"
       - "rotary"
       - "relu2max"
       - "infinite"
-      - "mqa"
-      - "capped_hypersphere_pair"
-    n_head:
-      range:
-        start: 1
-        end: 12
-        step: 1
-    relu2max_divisor: [128.0, 256.0, 512.0]
-    # HyperSphereNorm-family knobs kept in this exploration so capped-hypersphere
-    # runs are tagged with the intended radius/gain/learning settings.
+      - "hd_150"
+    norm_variant_wte: ["hyperspherenorm"]
+    norm_wte_gain: [false]
+    norm_wte_radius_learning: [true]
+    hsnorm_scale: [1.0]
+    hsnorm_gain: [false, true]
+    hsnorm_radius_learning: [true]
+    named_group_variations:
+      - "wte_norm_var"
+  # Pre ln Baseline
+  - named_group_static:
+      - "qk_norm"
+      - "pre_ln"
+      - "rotary"
+      - "relu2max"
+      - "infinite"
+      - "hd_150"
     hsnorm_scale: [1.0]
-    hsnorm_gain: [false]
-    hsnorm_radius: [null]
-    hsnorm_radius_learning: [false]
+    hsnorm_gain: [false, true]
+    hsnorm_radius_learning: [true]
     named_group_variations:
-      - "head_dim_var"
       - "wte_norm_var"
diff --git a/variations/norm_variations.py b/variations/norm_variations.py
index 23eee6f867..12ea5c9abd 100644
--- a/variations/norm_variations.py
+++ b/variations/norm_variations.py
@@ -208,12 +208,20 @@ class CappedHyperSphereNorm(nn.Module):
 
     def __init__(self, config):
         super().__init__()
-        self.radius = math.sqrt(config.n_embd)
+
+        ndim = config.n_embd
+
+        self.radius = math.sqrt(ndim)
+
+        if config.hsnorm_gain:
+            self.gain = nn.Parameter(torch.ones(ndim))
+        else:
+            self.gain = 1.0
 
     def forward(self, x):
         norms = x.norm(2, dim=-1, keepdim=True)
         scale = torch.where(norms > self.radius, self.radius / (norms + 1e-8), torch.ones_like(norms))
-        return x * scale
+        return x * scale * self.gain
 
 class IdentityNorm(nn.Module):
     def __init__(self, config=None):  # Accept config for API consistency

From b65c3dbc208bc376653651c7f5a4a2e93b66ad39 Mon Sep 17 00:00:00 2001
From: klei22 <kauna@avalonwest.tech>
Date: Sun, 14 Jun 2026 17:15:10 -0700
Subject: [PATCH 3/3] Add muon adamw options

---
 .../relu2max_capped_hypersphere_peri_ln.yaml        | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/explorations/relu2max_capped_hypersphere_peri_ln.yaml b/explorations/relu2max_capped_hypersphere_peri_ln.yaml
index d40970c9bf..ddd62a9187 100644
--- a/explorations/relu2max_capped_hypersphere_peri_ln.yaml
+++ b/explorations/relu2max_capped_hypersphere_peri_ln.yaml
@@ -3,6 +3,14 @@
 ---
 
 named_static_groups:
+  # Optimizer
+  - named_group: "muon"
+    optimizer: ["muon"]
+    weight_decay: [0.0]
+
+  - named_group: "adamw"
+    optimizer: ["adamw"]
+
   # QK Norm
   - named_group: "qk_norm"
     use_qk_norm: [true]
@@ -67,6 +75,8 @@ named_static_groups:
 named_variation_groups:
   - named_group: "wte_norm_var"
     named_group_alternates: ["capped_rmsnorm", "capped_pair", "rmsnorm"]
+  - named_group: "optimizers"
+    named_group_alternates: ["muon", "adamw"]
 
 common_group:
   dataset: ["minipile"]
@@ -91,6 +101,7 @@ parameter_groups:
     hsnorm_radius_learning: [true]
     named_group_variations:
       - "wte_norm_var"
+      - "optimizers"
   # Peri-LN WTE Norm
   - named_group_static:
       - "qk_norm"
@@ -107,6 +118,7 @@ parameter_groups:
     hsnorm_radius_learning: [true]
     named_group_variations:
       - "wte_norm_var"
+      - "optimizers"
   # Pre ln Baseline
   - named_group_static:
       - "qk_norm"
@@ -120,3 +132,4 @@ parameter_groups:
     hsnorm_radius_learning: [true]
     named_group_variations:
       - "wte_norm_var"
+      - "optimizers"