Centralize soft-cap logic behind a shared helper and a named disabled…

tests/unit/utilities/test_activation_functions.py

@@ -0,0 +1,70 @@
+"""Unit tests for the shared soft-cap helpers.
+
+Covers ``softcap_enabled`` and ``apply_softcap`` in
+``transformer_lens.utilities.activation_functions`` (see issue #1489).
+"""
+
+import torch
+
+from transformer_lens.utilities.activation_functions import (
+    SOFTCAP_DISABLED,
+    apply_softcap,
+    softcap_enabled,
+)
+
+
+class TestSoftcapEnabled:
+    """Test cases for softcap_enabled."""
+
+    def test_none_is_disabled(self):
+        assert softcap_enabled(None) is False
+
+    def test_zero_is_disabled(self):
+        assert softcap_enabled(0.0) is False
+
+    def test_disabled_sentinel_is_disabled(self):
+        assert softcap_enabled(SOFTCAP_DISABLED) is False
+
+    def test_negative_is_disabled(self):
+        assert softcap_enabled(-5.0) is False
+
+    def test_positive_is_enabled(self):
+        assert softcap_enabled(30.0) is True
+
+
+class TestApplySoftcap:
+    """Test cases for apply_softcap."""
+
+    def test_disabled_is_identity(self):
+        x = torch.tensor([-100.0, -1.0, 0.0, 1.0, 100.0])
+        result = apply_softcap(x, SOFTCAP_DISABLED)
+        assert torch.equal(result, x)
+
+    def test_none_is_identity(self):
+        x = torch.tensor([-100.0, 0.0, 100.0])
+        result = apply_softcap(x, None)
+        assert torch.equal(result, x)
+
+    def test_enabled_matches_capped_tanh_formula(self):
+        x = torch.tensor([-50.0, -10.0, 0.0, 10.0, 50.0])
+        cap = 30.0
+        result = apply_softcap(x, cap)
+        expected = cap * torch.tanh(x / cap)
+        assert torch.allclose(result, expected)
+
+    def test_enabled_output_stays_within_cap(self):
+        x = torch.tensor([-1000.0, 1000.0])
+        cap = 30.0
+        result = apply_softcap(x, cap)
+        assert torch.all(result.abs() <= cap)
+
+    def test_disabled_sentinel_does_not_saturate_large_values(self):
+        """Regression guard for the truthiness trap the issue describes:
+        a naive ``if cap:`` check treats -1.0 as enabled, which computes
+        ``-1.0 * tanh(x / -1.0) == tanh(x)`` and saturates every large score
+        to 1.0. The disabled path must instead be a pure identity.
+        """
+        x = torch.tensor([50.0, 100.0])
+        result = apply_softcap(x, SOFTCAP_DISABLED)
+        assert torch.equal(result, x)
+        assert not torch.allclose(result, torch.ones_like(x))

transformer_lens/HookedTransformer.py

@@ -32,7 +32,6 @@
 import numpy as np
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 import tqdm.auto as tqdm
 from jaxtyping import Float, Int
 from transformers import AutoTokenizer, PreTrainedModel, PreTrainedTokenizerBase
@@ -66,12 +65,14 @@
 from transformer_lens.utilities import (
     USE_DEFAULT_VALUE,
     TypedModuleList,
+    apply_softcap,
     get_best_available_device,
     get_device_for_block_index,
     init_kaiming_normal_,
     init_kaiming_uniform_,
     init_xavier_normal_,
     init_xavier_uniform_,
+    softcap_enabled,
 )
 from transformer_lens.utilities.devices import move_to_and_update_config
 from transformer_lens.weight_processing import ProcessWeights
@@ -669,10 +670,7 @@ def forward(
                 return None
             else:
                 logits = self.unembed(residual)  # [batch, pos, d_vocab]
-                if self.cfg.output_logits_soft_cap > 0.0:
-                    logits = self.cfg.output_logits_soft_cap * F.tanh(
-                        logits / self.cfg.output_logits_soft_cap
-                    )
+                logits = apply_softcap(logits, self.cfg.output_logits_soft_cap)
                 if return_type == "logits":
                     return logits
                 else:
@@ -1420,7 +1418,7 @@ def from_pretrained(
                     "architecture. Setting center_writing_weights=False."
                 )
                 center_writing_weights = False
-        if center_unembed and cfg.output_logits_soft_cap > 0.0:
+        if center_unembed and softcap_enabled(cfg.output_logits_soft_cap):
             logging.warning(
                 "You tried to specify center_unembed=True for a model using logit softcap, but this can't be done! Softcapping is not invariant upon adding a constant "
                 "Setting center_unembed=False instead."

transformer_lens/components/abstract_attention.py

@@ -19,6 +19,7 @@
 from transformer_lens.FactoredMatrix import FactoredMatrix
 from transformer_lens.hook_points import HookPoint
 from transformer_lens.utilities import get_offset_position_ids
+from transformer_lens.utilities.activation_functions import apply_softcap
 from transformer_lens.utilities.attention import complex_attn_linear, simple_attn_linear
 
 if is_bitsandbytes_available():
@@ -522,10 +523,7 @@ def calculate_attention_scores(
             k, "batch key_pos head_index d_head -> batch head_index d_head key_pos"
         )
         attn_scores = q_ @ k_ / self.attn_scale
-        if self.cfg.attn_scores_soft_cap > 0:
-            attn_scores = self.cfg.attn_scores_soft_cap * F.tanh(
-                attn_scores / self.cfg.attn_scores_soft_cap
-            )
+        attn_scores = apply_softcap(attn_scores, self.cfg.attn_scores_soft_cap)
         return attn_scores
 
     def calculate_z_scores(

transformer_lens/config/hooked_transformer_config.py

@@ -14,7 +14,10 @@
 import numpy as np
 import torch
 
-from transformer_lens.utilities.activation_functions import SUPPORTED_ACTIVATIONS
+from transformer_lens.utilities.activation_functions import (
+    SOFTCAP_DISABLED,
+    SUPPORTED_ACTIVATIONS,
+)
 from transformer_lens.utilities.devices import get_device
 
 from .transformer_lens_config import TransformerLensConfig
@@ -280,8 +283,8 @@ class HookedTransformerConfig(TransformerLensConfig):
     decoder_start_token_id: Optional[int] = None
     tie_word_embeddings: bool = False
     use_normalization_before_and_after: bool = False
-    attn_scores_soft_cap: float = -1.0
-    output_logits_soft_cap: float = -1.0
+    attn_scores_soft_cap: float = SOFTCAP_DISABLED
+    output_logits_soft_cap: float = SOFTCAP_DISABLED
     use_NTK_by_parts_rope: bool = False
     NTK_by_parts_low_freq_factor: float = 1.0
     NTK_by_parts_high_freq_factor: float = 4.0

transformer_lens/config/transformer_bridge_config.py

@@ -4,6 +4,8 @@
 
 import torch
 
+from transformer_lens.utilities.activation_functions import SOFTCAP_DISABLED
+
 from .transformer_lens_config import TransformerLensConfig
 
 
@@ -77,8 +79,8 @@ def __init__(
         decoder_start_token_id: Optional[int] = None,
         tie_word_embeddings: bool = False,
         use_normalization_before_and_after: bool = False,
-        attn_scores_soft_cap: float = -1.0,
-        output_logits_soft_cap: float = -1.0,
+        attn_scores_soft_cap: float = SOFTCAP_DISABLED,
+        output_logits_soft_cap: float = SOFTCAP_DISABLED,
         use_NTK_by_parts_rope: bool = False,
         NTK_by_parts_low_freq_factor: float = 1.0,
         NTK_by_parts_high_freq_factor: float = 4.0,

transformer_lens/model_bridge/bridge.py

@@ -50,6 +50,7 @@
     VARIANT_SUBMODULE_NAMES,
 )
 from transformer_lens.model_bridge.get_params_util import get_bridge_params
+from transformer_lens.utilities.activation_functions import softcap_enabled
 from transformer_lens.utilities.aliases import resolve_alias
 from transformer_lens.utilities.devices import move_to_and_update_config
 from transformer_lens.utilities.lm_utils import lm_cross_entropy_loss
@@ -920,7 +921,7 @@ def process_weights(
             print(f"Processing weights for {self.cfg.model_name}...")
 
         # Soft capping (tanh) is not translation-invariant; centering would change output.
-        if center_unembed and getattr(self.cfg, "output_logits_soft_cap", -1.0) > 0.0:
+        if center_unembed and softcap_enabled(getattr(self.cfg, "output_logits_soft_cap", None)):
             import logging
 
             logging.warning(

transformer_lens/model_bridge/sources/native/model.py

@@ -18,6 +18,7 @@
 
 from transformer_lens.config import TransformerBridgeConfig
 from transformer_lens.utilities import TypedModuleList
+from transformer_lens.utilities.activation_functions import apply_softcap
 
 # gelu_new = the tanh-approximation HF GPT-2 / HT use; F.gelu(approximate="tanh")
 # is the exact same formula.
@@ -269,9 +270,7 @@ def forward(
 
         scores = torch.matmul(q, k.transpose(-2, -1)) / self.scale
         # Gemma2 soft-cap before the causal mask so masked positions stay -inf.
-        if self.attn_scores_soft_cap > 0:
-            c = self.attn_scores_soft_cap
-            scores = c * torch.tanh(scores / c)
+        scores = apply_softcap(scores, self.attn_scores_soft_cap)
 
         if self.causal:
             block_mask = self.causal_mask[:seq, :seq]
@@ -468,7 +467,5 @@ def forward(
             )
         hidden_states = self.ln_out(hidden_states)
         logits = self.head(hidden_states)
-        if self.output_logits_soft_cap > 0:
-            c = self.output_logits_soft_cap
-            logits = c * torch.tanh(logits / c)
+        logits = apply_softcap(logits, self.output_logits_soft_cap)
         return logits

transformer_lens/utilities/__init__.py

@@ -1,8 +1,11 @@
 from .activation_functions import (
+    SOFTCAP_DISABLED,
     SUPPORTED_ACTIVATIONS,
     ActivationFunction,
+    apply_softcap,
     gelu_fast,
     gelu_new,
+    softcap_enabled,
     solu,
 )
 from .attribute_utils import get_nested_attr, set_nested_attr

transformer_lens/utilities/activation_functions.py

@@ -3,7 +3,7 @@
 Utilities for interacting with all supported activation functions.
 """
 
-from typing import Callable, Dict
+from typing import Callable, Dict, Optional, cast
 
 import numpy as np
 import torch
@@ -118,6 +118,22 @@ def relu2(input: Float[torch.Tensor, "batch pos d_mlp"]) -> Float[torch.Tensor,
     return relu_applied * relu_applied
 
 
+SOFTCAP_DISABLED: float = -1.0
+
+
+def softcap_enabled(cap: Optional[float]) -> bool:
+    """A soft-cap is active only for a strictly positive value; -1.0 / 0 / None mean disabled."""
+    return cap is not None and cap > 0
+
+
+def apply_softcap(x: torch.Tensor, cap: Optional[float]) -> torch.Tensor:
+    """Gemma-style ``cap * tanh(x / cap)`` when enabled, identity otherwise."""
+    if not softcap_enabled(cap):
+        return x
+    cap = cast(float, cap)
+    return cap * torch.tanh(x / cap)
+
+
 # Convenient type for the format of each activation function
 ActivationFunction = Callable[..., torch.Tensor]