diff --git a/tests/unit/utilities/test_activation_functions.py b/tests/unit/utilities/test_activation_functions.py new file mode 100644 index 000000000..3a0cf12b0 --- /dev/null +++ b/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)) diff --git a/transformer_lens/HookedTransformer.py b/transformer_lens/HookedTransformer.py index ee2134bd1..2b9810fe9 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 @@ -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." 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..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 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 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..9d638b7b5 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/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( 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/__init__.py b/transformer_lens/utilities/__init__.py index e242e951c..319280256 100644 --- a/transformer_lens/utilities/__init__.py +++ b/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 diff --git a/transformer_lens/utilities/activation_functions.py b/transformer_lens/utilities/activation_functions.py index 3594acae8..0b650b360 100644 --- a/transformer_lens/utilities/activation_functions.py +++ b/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]