From b99a215e4a329fa56fcc9a84aa82495b128c5ddc Mon Sep 17 00:00:00 2001
From: Kauna <16511995+klei22@users.noreply.github.com>
Date: Sat, 4 Apr 2026 08:50:47 -0700
Subject: [PATCH 1/2] Add dSiLU and custom Mish activation variations

---
 variations/activation_variations.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/variations/activation_variations.py b/variations/activation_variations.py
index e253c2ce48..fc0b212633 100644
--- a/variations/activation_variations.py
+++ b/variations/activation_variations.py
@@ -51,6 +51,29 @@ def forward(self, x):
         # Apply the shifted GELU activation
         return self.gelu(x - self.shift)
 
+class dSiLU(nn.Module):
+    """
+    Derivative of SiLU:
+    d/dx[x * sigmoid(x)] = sigmoid(x) * (1 + x * (1 - sigmoid(x)))
+    """
+    def __init__(self, config):
+        super().__init__()
+
+    def forward(self, x):
+        sig = torch.sigmoid(x)
+        return sig * (1 + x * (1 - sig))
+
+
+class Mish(nn.Module):
+    """
+    Mish activation: x * tanh(softplus(x)).
+    """
+    def __init__(self, config):
+        super().__init__()
+
+    def forward(self, x):
+        return x * torch.tanh(torch.nn.functional.softplus(x))
+
 
 class PiecewiseLearnableActivation(nn.Module):
     def __init__(self, config):
@@ -349,6 +372,7 @@ def __init__(self, config=None):
     "glu": GLU_Config,
     "leaky_relu": LeakyReLU_Config,
     "mish": Mish_Config,
+    "mish_custom": Mish,
     "piecewise": PiecewiseLearnableActivation,
     "pfla": PiecewiseFullyLearnableActivation,
     "pfla_le": PiecewiseFullyLearnableActivationLearnedEnds,
@@ -361,6 +385,7 @@ def __init__(self, config=None):
     "selu": SELU_Config,
     "sigmoid": Sigmoid_Config,
     "silu": SiLU_Config,
+    "dsilu": dSiLU,
     "softplus": Softplus_Config,
     "softsign": Softsign_Config,
     "softshrink": Softshrink_Config,

From 50f8ba149ca8eaa81481200c6bede5046547028d Mon Sep 17 00:00:00 2001
From: Kauna <16511995+klei22@users.noreply.github.com>
Date: Sat, 4 Apr 2026 09:01:12 -0700
Subject: [PATCH 2/2] Add default_inf activation sweep config and wire new
 activation choices

---
 explorations/all_activations_default_inf.yaml | 100 ++++++++++++++++++
 train_args.py                                 |   2 +
 2 files changed, 102 insertions(+)
 create mode 100644 explorations/all_activations_default_inf.yaml

diff --git a/explorations/all_activations_default_inf.yaml b/explorations/all_activations_default_inf.yaml
new file mode 100644
index 0000000000..97f8c890c4
--- /dev/null
+++ b/explorations/all_activations_default_inf.yaml
@@ -0,0 +1,100 @@
+# all_activations_default_inf.yaml
+# Based on explorations/default_inf.yaml, but sweeps MLP activation variants.
+---
+
+named_static_groups:
+  # QK Norm
+  - named_group: "qk_norm"
+    use_qk_norm: [true]
+    use_qk_norm_scale: [true]
+
+  # Norm Type
+  - 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: "softmax"
+    softmax_variant_attn: ["softmax"]
+
+  # 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]
+
+common_group:
+  dataset: ["minipile"]
+  eval_interval: [2500]
+  max_iters: [10000]
+  never_save_checkpoint: [true]
+  compile: [true]
+  log_rankme: [true]
+  log_areq: [true]
+
+  # Parameters needed by activation implementations that read config fields directly.
+  relu_power: [2.0]
+  shifted_gelu_learnable_shift: [true]
+  shifted_gelu_initial_shift: [0.0]
+  softshrink_lambda: [0.5]
+  pla_num_points: [7]
+  pla_left_bound: [-2.0]
+  pla_right_bound: [2.0]
+  pfla_num_points: [50]
+  pfla_left_bound: [-10.0]
+  pfla_right_bound: [10.0]
+  lsa_num_knots: [30]
+
+parameter_groups:
+  - named_group_static:
+      - "qk_norm"
+      - "peri_ln"
+      - "rotary"
+      - "softmax"
+      - "infinite"
+      - "mqa"
+      - "hd_100"
+    activation_variant:
+      - "celu"
+      - "elu"
+      - "gelu"
+      - "gelu_shifted"
+      - "glu"
+      - "leaky_relu"
+      - "learned_spline"
+      - "mish"
+      - "mish_custom"
+      - "piecewise"
+      - "pfla"
+      - "pfla_le"
+      - "prelu"
+      - "relu"
+      - "relu_power"
+      - "relu6"
+      - "rrelu"
+      - "disco"
+      - "selu"
+      - "sigmoid"
+      - "silu"
+      - "dsilu"
+      - "softplus"
+      - "softsign"
+      - "softshrink"
+      - "squared_relu"
+      - "squared_gelu"
+      - "tanh"
+      - "identity"
diff --git a/train_args.py b/train_args.py
index 9553b7d92b..3db58861a2 100644
--- a/train_args.py
+++ b/train_args.py
@@ -815,6 +815,7 @@ def parse_args():
             "leaky_relu",
             "learned_spline",
             "mish",
+            "mish_custom",
             "piecewise",
             "pfla",
             "pfla_le",
@@ -827,6 +828,7 @@ def parse_args():
             "selu",
             "sigmoid",
             "silu",
+            "dsilu",
             "softplus",
             "softsign",
             "softshrink",