Centralize soft-cap logic behind shared helper and named sentinel

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
@@ -73,6 +72,7 @@
     init_xavier_normal_,
     init_xavier_uniform_,
 )
+from transformer_lens.utilities.activation_functions import apply_softcap
 from transformer_lens.utilities.devices import move_to_and_update_config
 from transformer_lens.weight_processing import ProcessWeights
 
@@ -670,9 +670,7 @@ def forward(
             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:
@@ -2186,11 +2184,13 @@ def generate(
                             top_p=top_p,
                             temperature=temperature,
                             freq_penalty=freq_penalty,
-                            tokens=torch.cat(
-                                (input_tokens, torch.cat(sampled_tokens_list, dim=1)), dim=1
-                            )
-                            if "sampled_tokens" in locals()
-                            else input_tokens,
+                            tokens=(
+                                torch.cat(
+                                    (input_tokens, torch.cat(sampled_tokens_list, dim=1)), dim=1
+                                )
+                                if "sampled_tokens" in locals()
+                                else input_tokens
+                            ),
                         ).to(get_device_for_block_index(0, self.cfg))
                     else:
                         sampled_tokens = utils.sample_logits(

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_DISABLED
 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 getattr(self.cfg, "output_logits_soft_cap", SOFTCAP_DISABLED) > 0.0:
             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/activation_functions.py

@@ -105,6 +105,27 @@ def xielu(input: Float[torch.Tensor, "batch pos d_mlp"]) -> Float[torch.Tensor,
     )
 
 
+SOFTCAP_DISABLED: float = -1.0
+
+
+def apply_softcap(x: torch.Tensor, cap: float) -> torch.Tensor:
+    """Apply soft-cap to a tensor using tanh.
+
+    If cap <= 0, returns x unchanged (softcap disabled).
+    Otherwise returns cap * tanh(x / cap).
+
+    Args:
+        x: The tensor to apply soft-cap to.
+        cap: The soft-cap value. If <= 0, softcap is disabled.
+
+    Returns:
+        The soft-capped tensor.
+    """
+    if cap <= 0:
+        return x
+    return cap * torch.tanh(x / cap)
+
+
 # Convenient type for the format of each activation function
 ActivationFunction = Callable[..., torch.Tensor]