Skip to content

feat(projection_kernel): add attention-head subspace affinity - #1721

Open
janmenjayap wants to merge 2 commits into
TransformerLensOrg:devfrom
janmenjayap:feat/projection-kernel
Open

feat(projection_kernel): add attention-head subspace affinity#1721
janmenjayap wants to merge 2 commits into
TransformerLensOrg:devfrom
janmenjayap:feat/projection-kernel

Conversation

@janmenjayap

@janmenjayap janmenjayap commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

Description

Implements #1720.

Adds a basis-invariant Projection Kernel analysis surface in two reviewable commits:

  1. model-independent reduced-SVD basis extraction, principal angles, raw and normalized PK,
    numerical-rank metadata, roundoff-bound handling, and random-subspace moments;
  2. a TransformerBridge OQ/OK/OV head-affinity wrapper with native GQA KV-head identity,
    hybrid-layer indices, forward/all layer masks, bounded-memory tiled scoring, and ranking.

The implementation promotes fp16/bf16 inputs to fp32 for stable SVDs, preserves float64,
detaches model weights from autograd, and reports role/layer/head context for malformed or
rank-deficient weights. The accompanying guide documents orientation, rank semantics,
limitations, GQA behavior, and the distinction from Composition Score, with runnable numerical
and GPT-2 TransformerBridge examples.

Validation

  • Focused PK suite: 67 passed.
  • make check-format: passed.
  • uv run mypy .: passed (392 source files).
  • uv run build-docs: passed with no PK-specific warnings.
  • uv build: source distribution and wheel built successfully.
  • Local PR test surfaces (make test-pr did not exit successfully because of the three
    integration failures described below):
    • unit: 5,363 passed;
    • docstring: 18 passed;
    • acceptance: 209 passed;
    • integration: 1,394 passed, including all PK tests; two gated Gemma loads failed because the
      local HF token lacks access, and one existing GraniteMoeHybrid MPS parity test reproduced the
      investigated PyTorch 2.10 MPS numerical divergence.
  • GitHub CI: all applicable compatibility, coverage, formatting, typing, docstring, benchmark,
    and notebook checks passed.

Checklist

  • I have read and understood the contribution guidelines.
  • I have added tests that prove my fix is effective or that my feature works.
  • I have added necessary documentation (if appropriate).
  • My changes generate no new warnings.
  • No changelog file is required by this repository; the user-facing change is documented in
    docs/source/content/projection_kernel.md.

Type of change

  • Bug fix
  • New feature
  • Breaking change
  • Documentation update

@janmenjayap
janmenjayap marked this pull request as ready for review August 24, 2026 16:06
@janmenjayap janmenjayap changed the title Add projection-kernel attention-head subspace affinity feat(projection_kernel): add attention-head subspace affinity Aug 24, 2026
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.

1 participant