feat(e7-13): extended activations (mish/geglu/reglu/xielu) + catalog entries

apstenku123 · apstenku123 · commit 12779b2a32a6 · 2026-05-21T16:42:54.000+02:00
Stage E7-13 of E2E Coverage Matrix v2 epic (cppmega-mlx-bb0.13).
Brings the activation registry from 6 → 10 entries: adds dense mish
and gated geglu/reglu/xielu.

cppmega_mlx/nn/activations.py:
- ActivationName Literal extended.
- IS_GATED map covers all 10 entries (4 gated: swiglu/geglu/reglu/xielu;
  6 dense: gelu/relu/relu2/sqrelu/silu/mish).
- apply_activation gains four new branches:
    * mish: x * tanh(softplus(x)); softplus via log1p(exp) for
      bf16 stability
    * geglu: gelu_approx(gate) * up
    * reglu: max(gate, 0) * up
    * xielu: gelu_approx(gate) * silu(up) — extended xGLU variant
      seen in Megatron ablations.

cppmega_v4/explain/catalog.py: +4 ExplainEntry items with paper
references (Misra 2019 for Mish; Shazeer 2020 for GeGLU/ReGLU).
xIELU is marked research-only with no canonical paper ref.

vbgui/src/components/BrickContextPanel.tsx: activation dropdown
extended to 11 options (10 activation names + the legacy 'glu' default
in _build_mlp).

Tests:
- tests/v4/test_extended_activations.py (+10 pytest):
    * ACTIVATION_NAMES has exactly 10
    * gated set has 4 entries (swiglu/geglu/reglu/xielu)
    * mish classified as dense
    * each new gated activation rejects missing gate
    * mish forward parity vs x * tanh(softplus(x))
    * geglu/reglu/xielu forward parity vs reference impls
    * every activation has a catalog entry
- tests/v4/test_optim_spec_lion.py: reused tests updated (6 → 10
  expected names; "mish" replaced by "banana" in unknown-name probe).
- tests/v4/test_catalog.py: dispatch_catalog_list_options_activation
  expects 10 names instead of 6.

Regression: pytest 2290 (was 2270; +10 extended +
2 churn-from-rename) / 2 skip / 15 xfail / 0 fail (excl. pre-existing
path_d). vbgui vitest 150 / 0 fail.

Closes cppmega-mlx-bb0.13.
diff --git a/cppmega_mlx/nn/activations.py b/cppmega_mlx/nn/activations.py
@@ -23,12 +23,15 @@
 
 
 ActivationName = Literal[
-    "gelu", "relu", "relu2", "sqrelu", "silu", "swiglu",
+    "gelu", "relu", "relu2", "sqrelu", "silu", "mish",
+    "swiglu", "geglu", "reglu", "xielu",
 ]
-"""Recognised activation names — E7-12 set (6 entries).
+"""Recognised activation names — E7-12 + E7-13 set (10 entries).
 
-Gated entries (require a ``gate`` companion projection): swiglu.
-Dense entries (single projection input): gelu, relu, relu2, sqrelu, silu.
+Gated entries (require a ``gate`` companion projection):
+  swiglu, geglu, reglu, xielu.
+Dense entries (single projection input):
+  gelu, relu, relu2, sqrelu, silu, mish.
 """
 
 
@@ -38,7 +41,11 @@
     "relu2":  False,
     "sqrelu": False,
     "silu":   False,
+    "mish":   False,
     "swiglu": True,
+    "geglu":  True,
+    "reglu":  True,
+    "xielu":  True,
 }
 
 
@@ -93,9 +100,24 @@ def apply_activation(
         return mx.square(mx.maximum(x, 0))
     if name == "silu":
         return nn.silu(x)
+    if name == "mish":
+        # x * tanh(softplus(x)); softplus(z) = log(1 + exp(z))
+        # Use log1p for numerical stability.
+        return x * mx.tanh(mx.log1p(mx.exp(x)))
     if name == "swiglu":
-        assert gate is not None  # narrowed for type-checkers
+        assert gate is not None
         return nn.silu(gate) * x
+    if name == "geglu":
+        assert gate is not None
+        return nn.gelu_approx(gate) * x
+    if name == "reglu":
+        assert gate is not None
+        return mx.maximum(gate, 0) * x
+    if name == "xielu":
+        # Extended xGLU variant: gelu(gate) * silu(x); not a single
+        # published paper but appears in some Megatron-style ablations.
+        assert gate is not None
+        return nn.gelu_approx(gate) * nn.silu(x)
 
     raise AssertionError(f"unreachable: {name!r}")
 
diff --git a/cppmega_v4/explain/catalog.py b/cppmega_v4/explain/catalog.py
@@ -265,6 +265,55 @@ def list_options(category: str) -> list[ExplainEntry]:
                  "projection; verify will block the mismatch",
                  "Convention: intermediate_size = 8/3*H rounded to 256"),
     ),
+    ExplainEntry(
+        category="activation", name="mish",
+        summary="Mish — `x * tanh(softplus(x))` (Misra 2019). Smooth, "
+                "self-regularizing.",
+        when_to_use="Vision tasks; some LLM ablations show +0.5-1% "
+                    "accuracy vs ReLU/GELU.",
+        when_to_avoid="LLM pretraining at scale — SwiGLU dominates.",
+        recommended_params={},
+        paper_ref="Misra, 2019",
+        paper_url="https://arxiv.org/abs/1908.08681",
+        gotchas=("More compute than GELU; not significantly better on "
+                 "LLM downstream metrics",),
+    ),
+    ExplainEntry(
+        category="activation", name="geglu",
+        summary="Gated GELU — `gelu(gate) * up` (Shazeer 2020).",
+        when_to_use="GLM / Falcon-family pretraining; fine-tuning from "
+                    "GELU-pretrained checkpoints with gating.",
+        when_to_avoid="Same as SwiGLU memory caveat; pick SwiGLU if "
+                      "you're starting fresh in 2025.",
+        recommended_params={},
+        paper_ref="Shazeer, 2020",
+        paper_url="https://arxiv.org/abs/2002.05202",
+        gotchas=("Requires gated_mlp brick",),
+    ),
+    ExplainEntry(
+        category="activation", name="reglu",
+        summary="Gated ReLU — `relu(gate) * up` (Shazeer 2020).",
+        when_to_use="Hardware-constrained inference where ReLU is "
+                    "cheaper than GELU/SiLU on the target accelerator.",
+        when_to_avoid="Modern LLM pretraining — quality lags SwiGLU.",
+        recommended_params={},
+        paper_ref="Shazeer, 2020",
+        paper_url="https://arxiv.org/abs/2002.05202",
+        gotchas=("Dead-gate problem: some gate units may emit 0 "
+                 "indefinitely",),
+    ),
+    ExplainEntry(
+        category="activation", name="xielu",
+        summary="Extended xGLU — `gelu(gate) * silu(up)`. Niche, "
+                "appears in Megatron ablations.",
+        when_to_use="Research / ablation runs; not a default choice.",
+        when_to_avoid="Production pretraining — under-studied.",
+        recommended_params={},
+        paper_ref=None,
+        paper_url=None,
+        gotchas=("Combined gating + activation cost > SwiGLU; quality "
+                 "rarely better",),
+    ),
 ]
 
 
diff --git a/tests/v4/test_catalog.py b/tests/v4/test_catalog.py
@@ -190,9 +190,11 @@ def test_dispatch_catalog_list_options_activation():
     })
     assert resp.error is None
     options = resp.result["options"]
-    assert len(options) == 6
+    # E7-13 extended the activation registry from 6 → 10 entries.
+    assert len(options) == 10
     names = {o["name"] for o in options}
-    assert names == {"gelu", "relu", "relu2", "sqrelu", "silu", "swiglu"}
+    assert names == {"gelu", "relu", "relu2", "sqrelu", "silu", "mish",
+                     "swiglu", "geglu", "reglu", "xielu"}
 
 
 def test_dispatch_catalog_explain_rejects_extra_params():
diff --git a/tests/v4/test_extended_activations.py b/tests/v4/test_extended_activations.py
@@ -0,0 +1,74 @@
+"""E7-13 tests: extended activations (mish/geglu/reglu/xielu)."""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+import pytest
+
+from cppmega_mlx.nn.activations import (
+    ACTIVATION_NAMES, IS_GATED, apply_activation, is_gated,
+)
+from cppmega_v4.explain import get_entry
+
+
+def test_activation_names_have_ten_entries():
+    assert set(ACTIVATION_NAMES) == {
+        "gelu", "relu", "relu2", "sqrelu", "silu", "mish",
+        "swiglu", "geglu", "reglu", "xielu",
+    }
+
+
+def test_gated_set_has_four_entries():
+    gated = {n for n, g in IS_GATED.items() if g}
+    assert gated == {"swiglu", "geglu", "reglu", "xielu"}
+
+
+def test_mish_is_dense():
+    assert not is_gated("mish")
+
+
+@pytest.mark.parametrize("name", ["geglu", "reglu", "xielu"])
+def test_new_gated_activations_require_gate(name):
+    x = mx.random.normal((2, 4))
+    with pytest.raises(ValueError, match="is gated"):
+        apply_activation(name, x)
+
+
+def test_mish_forward_matches_reference():
+    """Mish = x * tanh(softplus(x))."""
+    x = mx.random.normal((4, 8), key=mx.random.key(0))
+    got = apply_activation("mish", x)
+    expected = x * mx.tanh(mx.log1p(mx.exp(x)))
+    assert mx.allclose(got, expected).item()
+
+
+def test_geglu_uses_gelu_on_gate():
+    x = mx.random.normal((2, 4), key=mx.random.key(0))
+    g = mx.random.normal((2, 4), key=mx.random.key(1))
+    got = apply_activation("geglu", x, gate=g)
+    expected = nn.gelu_approx(g) * x
+    assert mx.allclose(got, expected).item()
+
+
+def test_reglu_uses_relu_on_gate():
+    x = mx.random.normal((2, 4), key=mx.random.key(0))
+    g = mx.random.normal((2, 4), key=mx.random.key(1))
+    got = apply_activation("reglu", x, gate=g)
+    expected = mx.maximum(g, 0) * x
+    assert mx.allclose(got, expected).item()
+
+
+def test_xielu_uses_gelu_gate_silu_value():
+    x = mx.random.normal((2, 4), key=mx.random.key(0))
+    g = mx.random.normal((2, 4), key=mx.random.key(1))
+    got = apply_activation("xielu", x, gate=g)
+    expected = nn.gelu_approx(g) * nn.silu(x)
+    assert mx.allclose(got, expected).item()
+
+
+@pytest.mark.parametrize("name", ACTIVATION_NAMES)
+def test_every_activation_has_catalog_entry(name):
+    entry = get_entry("activation", name)
+    assert entry is not None, f"missing catalog entry for {name}"
+    assert entry.summary
diff --git a/tests/v4/test_optim_spec_lion.py b/tests/v4/test_optim_spec_lion.py
@@ -160,21 +160,24 @@ def test_existing_factories_unchanged(factory):
 # ---------------------------------------------------------------------------
 
 
-def test_activation_names_contains_six_entries():
+def test_activation_names_contains_ten_entries():
+    # E7-13 extended the set from 6 → 10 (added mish + geglu/reglu/xielu).
     assert set(ACTIVATION_NAMES) == {
-        "gelu", "relu", "relu2", "sqrelu", "silu", "swiglu",
+        "gelu", "relu", "relu2", "sqrelu", "silu", "mish",
+        "swiglu", "geglu", "reglu", "xielu",
     }
 
 
-def test_is_gated_only_swiglu_is_gated():
-    assert is_gated("swiglu") is True
-    for name in ("gelu", "relu", "relu2", "sqrelu", "silu"):
+def test_is_gated_four_gated_six_dense():
+    for name in ("swiglu", "geglu", "reglu", "xielu"):
+        assert is_gated(name) is True
+    for name in ("gelu", "relu", "relu2", "sqrelu", "silu", "mish"):
         assert is_gated(name) is False
 
 
 def test_is_gated_rejects_unknown_name():
-    with pytest.raises(ValueError, match="unknown activation 'mish'"):
-        is_gated("mish")
+    with pytest.raises(ValueError, match="unknown activation 'banana'"):
+        is_gated("banana")
 
 
 @pytest.mark.parametrize("name", ["gelu", "relu", "relu2", "sqrelu", "silu"])
diff --git a/vbgui/src/components/BrickContextPanel.tsx b/vbgui/src/components/BrickContextPanel.tsx
@@ -10,7 +10,8 @@ import { Tooltip } from "@/components/Tooltip";
 import { ExplainModal } from "@/components/ExplainModal";
 
 const ACTIVATION_OPTIONS = [
-  "glu", "gelu", "relu", "relu2", "sqrelu", "silu", "swiglu",
+  "glu", "gelu", "relu", "relu2", "sqrelu", "silu", "mish",
+  "swiglu", "geglu", "reglu", "xielu",
 ];
 const NORM_OPTIONS = ["rmsnorm", "layernorm", "none"];
 
