Skip to content

fix: validate reused coordinate decompositions by pursuit algorithm - #1932

Merged
jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
emerardd:fix/coordinate-pursuit-reuse
Oct 22, 2026
Merged

jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
emerardd:fix/coordinate-pursuit-reuse

Conversation

@emerardd

Copy link
Copy Markdown
Contributor

Description

Coordinate patching rejects valid gradient_pursuit decompositions when they are supplied explicitly or reused through the per-position/cache path. Its compatibility validator currently applies NNLS KKT stationarity to both supported algorithms, although gradient pursuit deliberately performs approximate projected-gradient updates.

For example, with dictionary rows [[1, 0, 0], [0.7, 0.7, 0], [0, 0, 1]], activation [1, 0.3, 0], and k=2, a fresh gradient-pursuit patch from atom 0 to atom 2 returns [0, 0.3, 1]. Passing its own decomposition back into the solver instead raises NNLS result violates stationarity on positive coordinates. solve_coordinate_patch_positions() encounters the same error even on its first call because it passes the freshly computed decomposition through the reuse validator.

This change validates gradient-pursuit coordinates by replaying the coefficient updates over the recorded selected atoms, including the solver's final active-support pruning. It avoids a vocabulary-wide selection scan on reuse. The existing NNLS KKT validation and shared activation/dictionary/reconstruction checks remain in place, so self-consistently rescaled coordinates and stale activations are still rejected.

The existing algorithm argument now selects the validation rule for supplied/cached decompositions as well as the fresh solver; its docstrings state that it must match the decomposition's producer. No decomposition fields or public signatures change.

Type of change

  • Bug fix (non-breaking change which fixes an issue)

Validation

The regression coverage includes correlated atoms, explicit reuse, per-position calls, cache misses/hits, both coefficient-update algorithms, activation scales, invalid coordinates, empty supports, unknown algorithms, and hook/offline parity through JacobianLens.

  • Coordinate-patch, sparse-decomposition and coordinate-patch-hook tests: 231 passed.
  • Full mypy .: 389 source files passed.
  • Pinned pycln/isort/Black checks on all four changed files and git diff --check: passed.

The complete unit suite and GPU/MPS tests were not run.

Checklist

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation (API docstrings)
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes (affected files run; full suite not run)
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@jlarson4
jlarson4 merged commit 12cfe63 into TransformerLensOrg:dev Oct 22, 2026
28 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