From 41d46a9e0e29e3e3a780a5596f4d822d700e1008 Mon Sep 17 00:00:00 2001 From: hiro-nikaitou Date: Tue, 7 Jul 2026 15:42:01 +0800 Subject: [PATCH 1/2] refactor(utilities): centralize soft-cap logic with shared helper Co-Authored-By: Claude Opus 4.8 (1M context) --- transformer_lens/HookedTransformer.py | 6 ++---- .../components/abstract_attention.py | 6 ++---- .../config/hooked_transformer_config.py | 6 +++--- .../config/transformer_bridge_config.py | 6 ++++-- transformer_lens/model_bridge/bridge.py | 4 +++- .../model_bridge/sources/native/model.py | 9 +++----- .../utilities/activation_functions.py | 21 +++++++++++++++++++ 7 files changed, 38 insertions(+), 20 deletions(-) diff --git a/transformer_lens/HookedTransformer.py b/transformer_lens/HookedTransformer.py index ee2134bd1..9e39ac834 100644 --- a/transformer_lens/HookedTransformer.py +++ b/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 @@ -74,6 +73,7 @@ init_xavier_uniform_, ) from transformer_lens.utilities.devices import move_to_and_update_config +from transformer_lens.utilities.activation_functions import apply_softcap from transformer_lens.weight_processing import ProcessWeights SingleLoss = Float[torch.Tensor, ""] # Type alias for a single element tensor @@ -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: diff --git a/transformer_lens/components/abstract_attention.py b/transformer_lens/components/abstract_attention.py index 01a0646eb..9f0e819ba 100644 --- a/transformer_lens/components/abstract_attention.py +++ b/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( diff --git a/transformer_lens/config/hooked_transformer_config.py b/transformer_lens/config/hooked_transformer_config.py index 57ed34745..a0c056dd8 100644 --- a/transformer_lens/config/hooked_transformer_config.py +++ b/transformer_lens/config/hooked_transformer_config.py @@ -14,7 +14,7 @@ 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 +280,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 diff --git a/transformer_lens/config/transformer_bridge_config.py b/transformer_lens/config/transformer_bridge_config.py index 593407451..9938c6313 100644 --- a/transformer_lens/config/transformer_bridge_config.py +++ b/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, diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py index 4800b0150..a3e6aae8a 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -28,6 +28,8 @@ import einops import numpy as np import torch + +from transformer_lens.utilities.activation_functions import SOFTCAP_DISABLED import tqdm from torch import nn @@ -920,7 +922,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( diff --git a/transformer_lens/model_bridge/sources/native/model.py b/transformer_lens/model_bridge/sources/native/model.py index 4cae95a2a..5617459e7 100644 --- a/transformer_lens/model_bridge/sources/native/model.py +++ b/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 diff --git a/transformer_lens/utilities/activation_functions.py b/transformer_lens/utilities/activation_functions.py index b0a280967..b74f9c20d 100644 --- a/transformer_lens/utilities/activation_functions.py +++ b/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] From ef5345e75928e99a79ff817bed239ab3404e0053 Mon Sep 17 00:00:00 2001 From: hiro-nikaitou Date: Tue, 7 Jul 2026 18:07:42 +0800 Subject: [PATCH 2/2] Fix Format Check: apply isort and black formatting - isort: fix import ordering in bridge.py, HookedTransformer.py - black: reformat conditional expression in HookedTransformer.py - black: wrap long import line in hooked_transformer_config.py Format check commands all pass locally (pycln, isort, black). --- transformer_lens/HookedTransformer.py | 14 ++++++++------ .../config/hooked_transformer_config.py | 5 ++++- transformer_lens/model_bridge/bridge.py | 3 +-- 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/transformer_lens/HookedTransformer.py b/transformer_lens/HookedTransformer.py index 9e39ac834..29b5d3bcd 100644 --- a/transformer_lens/HookedTransformer.py +++ b/transformer_lens/HookedTransformer.py @@ -72,8 +72,8 @@ init_xavier_normal_, init_xavier_uniform_, ) -from transformer_lens.utilities.devices import move_to_and_update_config 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 SingleLoss = Float[torch.Tensor, ""] # Type alias for a single element tensor @@ -2184,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( diff --git a/transformer_lens/config/hooked_transformer_config.py b/transformer_lens/config/hooked_transformer_config.py index a0c056dd8..0e4a757ef 100644 --- a/transformer_lens/config/hooked_transformer_config.py +++ b/transformer_lens/config/hooked_transformer_config.py @@ -14,7 +14,10 @@ import numpy as np import torch -from transformer_lens.utilities.activation_functions import SOFTCAP_DISABLED, 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 diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py index a3e6aae8a..f891f92d0 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -28,8 +28,6 @@ import einops import numpy as np import torch - -from transformer_lens.utilities.activation_functions import SOFTCAP_DISABLED import tqdm from torch import nn @@ -52,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