Skip to content

fix(analysis): reject unfolded final norm in DLA - #1806

Merged
jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
emerardd:fix/dla-unfolded-final-ln
Sep 24, 2026
Merged

jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
emerardd:fix/dla-unfolded-final-ln

Conversation

@emerardd

@emerardd emerardd commented Sep 24, 2026 •

Copy link
Copy Markdown
Contributor

Description

direct_logit_attribution() currently checks compatibility_mode but accepts enable_compatibility_mode(fold_ln=False) and no_processing=True. In those states a learned final-normalization scale is still applied during the forward pass, while ActivationCache.logit_attrs() projects against W_U without that scale. A tiny native Bridge with nonuniform final-LN weights produced an attribution sum of -2.2847 for a token difference whose actual logit difference was -3.0057; the folded control reconstructed the difference.

This change checks the final norm's current affine parameters before DLA: its effective weight must be identity and its bias zero. It fails clearly when either learned parameter remains active, while allowing parameter-free final norms and a folded norm. Checking live weights also catches changes made after compatibility mode was enabled. The API docstring and error messages describe the supported state and how to obtain it. No separate Issue was filed.

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • This change requires a documentation update (the API docstring was updated)

Validation

  • New tiny-native regression tests: 10 passed, covering LN/RMS, fold_ln=False, no_processing=True, an active final-LN bias, post-fold edits, folded numerical reconstruction, and parameter-free LNPre/RMSPre.
  • Existing DLA integration tests with the cached GPT-2 model, run offline: 14 passed.
  • pycln --check --all, pinned isort --check-only, Black --check, git diff --check, and mypy . passed.
  • Broader tests/unit/tools run on this Windows environment: 962 passed, 1 skipped, 2 failed in test_sparse_probing.py (test_default_tolerance_accepts_large_scale_activations and test_stop_reason_distinguishes_converged_from_capped_fits). Both nodes also showed failure in an isolated, unmodified upstream/dev checkout using the same locked environment; their optimizer reached tolerance_grad rather than the test's expected refinement/max_eval route. The DLA patch does not touch that solver or its tests. I am not claiming the complete unit suite passed locally.

Checklist

  • I have made corresponding changes to the documentation (API docstring)
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective
  • New and existing unit tests pass locally with my changes (see the two sparse-probing failures above)
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@jlarson4

Copy link
Copy Markdown
Collaborator

Solid fix to this bug @emerardd, this is good to go and will be merged now!

@jlarson4
jlarson4 merged commit 2e4cdd1 into TransformerLensOrg:dev Sep 24, 2026
26 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants