Centralize soft-cap logic behind a shared helper and a named disabled…#1491
Merged
jlarson4 merged 1 commit intoJul 7, 2026
Merged
Conversation
… sentinel Consolidates the tanh soft-capping (Gemma-2 style) used for attention scores and output logits into apply_softcap()/softcap_enabled() in transformer_lens/utilities/activation_functions.py, replacing five independent cap * tanh(x / cap) implementations across HookedTransformer, abstract_attention, and the native TransformerBridge model, plus the center_unembed compatibility checks in both systems. Also replaces the bare -1.0 disabled sentinel with a named SOFTCAP_DISABLED constant in both HookedTransformerConfig and TransformerBridgeConfig. Pure refactor, no behavior change. A naive 'if cap:' truthiness check would treat -1.0 as enabled (computing -1.0 * tanh(x / -1.0) == tanh(x), saturating every large score to 1.0); softcap_enabled() makes that mistake structurally impossible instead of relying on every call site getting 'cap > 0' right independently. Verified: uv run mypy . clean (307 files), make unit-test clean (3204 passed, 32 skipped, 9 xfailed, 0 failed), docstring tests on changed files clean. Fixes TransformerLensOrg#1489
Collaborator
|
Great work on this @delaidam! I appreciate the thoroughness & additional testing you added. Merging now |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #1489
Consolidates the tanh soft-capping (Gemma-2 style) used for attention scores
and output logits into
apply_softcap()/softcap_enabled()intransformer_lens/utilities/activation_functions.py, replacing fiveindependent
cap * tanh(x / cap)implementations:Also replaces the bare
-1.0disabled sentinel with a namedSOFTCAP_DISABLEDconstant in bothHookedTransformerConfigandTransformerBridgeConfig.Pure refactor, no behavior change. Left the HF-native attention path
(
generalized_components/position_embeddings_attention.py) and theconfig-mapping sites (
sources/_bridge_builder.py,sources/transformers.py)untouched, per the issue's scope notes.
Verification
uv run mypy .- clean, 307 filesmake unit-test- 3204 passed, 32 skipped, 9 xfailed, 0 failedpytest(docstrings) on all changed files - cleantests/unit/utilities/test_activation_functions.py,including a regression guard for the exact truthiness trap the issue
describes (
-1.0must stay identity, not saturate to 1.0)