Add linear-tail squared activation variants

explorations/linear_squared_variants_pre_norm.yaml

@@ -0,0 +1,40 @@
+# Sweep C1-continuous linear-tail squared variants using the default recipe.
+---
+
+named_static_groups:
+  - named_group: "qk_norm"
+    use_qk_norm: [true]
+    use_qk_norm_scale: [true]
+
+  - named_group: "pre_ln"
+    use_pre_ln: [true]
+    use_peri_ln: [false]
+    use_post_ln: [false]
+
+  - named_group: "rotary"
+    named_group_settings:
+      use_rotary_embeddings: [true]
+      use_abs_pos_embeddings: [false]
+
+named_variation_groups:
+  - named_group: "wte_norm"
+    norm_variant_wte: [default, "hyperspherenorm", null]
+
+common_group:
+  dataset: ["minipile"]
+  eval_interval: [1000]
+  max_iters: [10000]
+  never_save_checkpoint: [true]
+
+parameter_groups:
+  - compile: [true]
+    named_group_static: ["qk_norm", "pre_ln", "rotary"]
+    named_group_variations: ["wte_norm"]
+    softmax_variant_attn: ["relu2max_linear"]
+    relu2max_linear_cutoff: [2, 5, 10, 20, 50]
+
+  - compile: [true]
+    named_group_static: ["qk_norm", "pre_ln", "rotary"]
+    named_group_variations: ["wte_norm"]
+    activation_variant: ["squared_relu_linear"]
+    squared_relu_linear_cutoff: [2, 5, 10, 20, 50]

gpt_conf.py

@@ -204,6 +204,7 @@ class GPTConfig:
     activation_transition_start_iter: int = 0
     activation_transition_end_iter: int = None
     relu_power: float = 2.0
+    squared_relu_linear_cutoff: float = 10.0
 
     # MLP Options
     use_parallel_mlp: bool = False
@@ -301,6 +302,7 @@ class GPTConfig:
 
     ## ReLUMax options
     relu2max_divisor: float = 256.0
+    relu2max_linear_cutoff: float = 10.0
 
     ## SigmoidMax options
     sigmoidmax_divisor: float = 256.0

tests/test_linear_squared_variants.py

@@ -0,0 +1,51 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from variations.activation_variations import SquaredReLULinear
+from variations.softmax_variations import ReLU2MaxLinear
+
+
+def test_squared_relu_linear_matches_square_and_tangent():
+    activation = SquaredReLULinear(
+        SimpleNamespace(squared_relu_linear_cutoff=2.0)
+    )
+    x = torch.tensor([-1.0, 0.0, 1.0, 2.0, 3.0], requires_grad=True)
+
+    output = activation(x)
+
+    torch.testing.assert_close(output, torch.tensor([0.0, 0.0, 1.0, 4.0, 8.0]))
+    output.sum().backward()
+    torch.testing.assert_close(x.grad, torch.tensor([0.0, 0.0, 2.0, 4.0, 4.0]))
+
+
+def test_relu2max_linear_applies_divisors():
+    activation = ReLU2MaxLinear(
+        SimpleNamespace(
+            relu2max_divisor=2.0,
+            relu2max_linear_cutoff=2.0,
+            div_by_seq_len=True,
+        )
+    )
+
+    output = activation(torch.tensor([[1.0, 2.0, 3.0, -1.0]]))
+
+    torch.testing.assert_close(output, torch.tensor([[0.125, 0.5, 1.0, 0.0]]))
+
+
+@pytest.mark.parametrize(
+    ("variant", "attribute"),
+    [
+        (SquaredReLULinear, "squared_relu_linear_cutoff"),
+        (ReLU2MaxLinear, "relu2max_linear_cutoff"),
+    ],
+)
+def test_linear_squared_variants_reject_nonpositive_cutoffs(variant, attribute):
+    config = {
+        attribute: 0.0,
+        "relu2max_divisor": 1.0,
+        "div_by_seq_len": False,
+    }
+    with pytest.raises(ValueError, match="must be greater than zero"):
+        variant(SimpleNamespace(**config))

train_args.py

@@ -900,6 +900,7 @@ def parse_args():
             "softsign",
             "softshrink",
             "squared_relu",
+            "squared_relu_linear",
             "squared_gelu",
             "tanh",
             "identity",
@@ -922,6 +923,9 @@ def parse_args():
     ## ReLUPower
     model_group.add_argument("--relu_power", type=float, default=2.0)
 
+    ## Squared ReLU with a linear tail
+    model_group.add_argument("--squared_relu_linear_cutoff", type=float, default=10.0)
+
     ## Shifted Gelu
     model_group.add_argument("--shifted_gelu_learnable_shift",  type=bool, default=True, action=argparse.BooleanOptionalAction)
     model_group.add_argument("--shifted_gelu_initial_shift", type=float, default=0.0)
@@ -1313,6 +1317,7 @@ def parse_args():
         "polymax",
         "relumax",
         "relu2max",
+        "relu2max_linear",
         "sigmoidmax",
         "vpolymax",
         "exppolymax",
@@ -1362,6 +1367,7 @@ def parse_args():
 
     ### ReLU2Max Options
     model_group.add_argument("--relu2max_divisor", type=float, default=256.0)
+    model_group.add_argument("--relu2max_linear_cutoff", type=float, default=10.0)
 
     ### SimgoidMax Options
     model_group.add_argument("--sigmoidmax_divisor", type=float, default=256.0)

variations/activation_variations.py

@@ -10,6 +10,23 @@ def __init__(self, config):
     def forward(self, x):
         return torch.pow(torch.relu(x), 2)
 
+
+class SquaredReLULinear(nn.Module):
+    """Squared ReLU with its tangent line used above a configurable cutoff."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.cutoff = config.squared_relu_linear_cutoff
+        if self.cutoff <= 0:
+            raise ValueError("squared_relu_linear_cutoff must be greater than zero")
+
+    def forward(self, x):
+        relu_x = torch.relu(x)
+        # The tangent to x^2 at c is 2cx-c^2.  Matching both value and slope
+        # makes this transition C1 continuous while preventing quadratic growth.
+        linear_tail = 2 * self.cutoff * relu_x - self.cutoff**2
+        return torch.where(relu_x <= self.cutoff, relu_x**2, linear_tail)
+
 class ReLUPower(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -366,6 +383,7 @@ def __init__(self, config=None):
     "softshrink": Softshrink_Config,
     "relu_power": ReLUPower,
     "squared_relu": SquaredReLU,
+    "squared_relu_linear": SquaredReLULinear,
     "squared_gelu": SquaredGELU,
     "tanh": Tanh_Config,
     "identity": Identity_Config,

variations/softmax_variations.py

@@ -577,6 +577,30 @@ def forward(self, x):
         return result
 
 
+class ReLU2MaxLinear(nn.Module):
+    """ReLU2Max with a C1-continuous linear tail above a cutoff."""
+
+    def __init__(self, config, dim=-1):
+        super().__init__()
+        self.dim = dim
+        self.relu2max_divisor = config.relu2max_divisor
+        self.cutoff = config.relu2max_linear_cutoff
+        self.div_by_seq_len = config.div_by_seq_len
+        if self.cutoff <= 0:
+            raise ValueError("relu2max_linear_cutoff must be greater than zero")
+
+    def forward(self, x):
+        relu_x = torch.relu(x)
+        linear_tail = 2 * self.cutoff * relu_x - self.cutoff**2
+        result = torch.where(relu_x <= self.cutoff, relu_x**2, linear_tail)
+        result = result / self.relu2max_divisor
+
+        if self.div_by_seq_len:
+            result = result / x.shape[self.dim]
+
+        return result
+
+
 class Softplus2Max(nn.Module):
     """ Softmax variant based on arxiv 1805.10829 with added handles for base """
     def __init__(self, config, dim=-1):
@@ -870,6 +894,7 @@ def forward(self, x):
     "sigsoftmax": SigSoftmax,
     "relumax": ReLUMax,
     "relu2max": ReLU2Max,
+    "relu2max_linear": ReLU2MaxLinear,
     "sigmoidmax": SigmoidMax,
     "softshrink": Softshrink,
     "gelumax": Gelumax,