From bf0550447eafef1135bc9a204f1ece2c1c00b39b Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Sun, 9 Aug 2026 23:18:05 -0700 Subject: [PATCH] Add linear-tail squared activation variants --- .../linear_squared_variants_pre_norm.yaml | 40 +++++++++++++++ gpt_conf.py | 2 + tests/test_linear_squared_variants.py | 51 +++++++++++++++++++ train_args.py | 6 +++ variations/activation_variations.py | 18 +++++++ variations/softmax_variations.py | 25 +++++++++ 6 files changed, 142 insertions(+) create mode 100644 explorations/linear_squared_variants_pre_norm.yaml create mode 100644 tests/test_linear_squared_variants.py diff --git a/explorations/linear_squared_variants_pre_norm.yaml b/explorations/linear_squared_variants_pre_norm.yaml new file mode 100644 index 0000000000..c5561a34d6 --- /dev/null +++ b/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] diff --git a/gpt_conf.py b/gpt_conf.py index 056b7c7721..d30f9e7982 100644 --- a/gpt_conf.py +++ b/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 diff --git a/tests/test_linear_squared_variants.py b/tests/test_linear_squared_variants.py new file mode 100644 index 0000000000..70f24ffde7 --- /dev/null +++ b/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)) diff --git a/train_args.py b/train_args.py index a6b7491d0f..07f9d03b7d 100644 --- a/train_args.py +++ b/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) diff --git a/variations/activation_variations.py b/variations/activation_variations.py index e253c2ce48..174cc371f4 100644 --- a/variations/activation_variations.py +++ b/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, diff --git a/variations/softmax_variations.py b/variations/softmax_variations.py index 9cbaf2bcf9..f80f3b5b10 100644 --- a/variations/softmax_variations.py +++ b/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,