feat(projection_kernel): add attention-head subspace affinity - #1721
feat(projection_kernel): add attention-head subspace affinity#1721janmenjayap wants to merge 2 commits into
Conversation
janmenjayap
changed the title
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Excellent work here @janmenjayap! This is a solid implementation of your plan from #1720. Just a couple comments to tighten things up before merging
|
|
||
| overlap = first.T @ second | ||
| score = _clamp_projection_scores(overlap.square().sum(), min(subspace_a.rank, subspace_b.rank)) | ||
| cosines = torch.linalg.svdvals(overlap) |
There was a problem hiding this comment.
torch.linalg.svdvals has no MPS kernel, so projection_kernel raises NotImplementedError on MPS tensors while orthonormal_subspace survives via the CPU fallback. Let's route the small overlap matrix through .cpu() or use sqrt(eigvalsh(overlapᵀoverlap)), whichever you prefer.
| input_shape = (matrix.shape[0], matrix.shape[1]) | ||
| requested_rank = _validate_rank(rank, min(input_shape)) | ||
| compute_dtype = _compute_dtype(matrix.dtype) | ||
| effective_rtol = _validate_rtol(rtol, input_shape, compute_dtype) |
There was a problem hiding this comment.
The default rank tolerance derives from the compute dtype, so a genuinely rank-deficient fp16/bf16 head passes the full-column-rank gate with a quantization-noise basis direction. Can we derive rtol from the storage dtype's eps (or warn at _extract_bases for low-precision inputs)?
It would also be nice to add a low-precision case to the rank guard test.
| if validated_layer_order == "forward": | ||
| layer_tensor = torch.tensor(layer_indices, device=result_device) | ||
| layer_mask = layer_tensor[:, None] < layer_tensor[None, :] | ||
| valid_mask = layer_mask[:, None, :, None].expand_as(scores) |
There was a problem hiding this comment.
In the forward branch valid_mask is a stride-0 broadcast view, so a user write to one entry flips its aliased copies; the all branch returns a materialized tensor. Add .contiguous().
| ) from error | ||
| head_bases.append(subspace.basis.to(device=device)) | ||
| head_ranks.append(subspace.measured_rank) | ||
| layer_bases.append(torch.stack(head_bases)) |
There was a problem hiding this comment.
The basis stack is fully materialized per role (~2.15 GB/role at 32L/32H/4096) while only the overlap temp is tiled, and top_pairs builds every valid pair before taking k. A docs note on the practical model-size ceiling now, with stack tiling + torch.topk as a follow-up, would cover it.
| overlap = first.T @ second | ||
| score = _clamp_projection_scores(overlap.square().sum(), min(subspace_a.rank, subspace_b.rank)) | ||
| cosines = torch.linalg.svdvals(overlap) | ||
| angles = torch.acos(cosines.clamp(min=0.0, max=1.0)) |
There was a problem hiding this comment.
With check_orthonormal=False and a hand-built basis, a cosine of 1.3 stays visible in cosines but clamps to angle 0.0, while score raises on the analogous violation. Apply the same tolerance-then-raise rule to the acos clamp, or document that angles assumes orthonormal inputs.
| attribute = f"W_{role}" | ||
| try: | ||
| matrix = getattr(attn, attribute) | ||
| except (AttributeError, RuntimeError, ValueError) as error: |
There was a problem hiding this comment.
require_readable_weight raises NotImplementedError to mean "unsupported", but the except (AttributeError, RuntimeError, ValueError) here catches it via its RuntimeError base and converts it to ValueError. Let NotImplementedError re-raise with the added role/layer context.
| target_layer_indices: Tuple[int, ...] | ||
| source_head_kind: HeadKind | ||
| target_head_kind: HeadKind | ||
| source_ranks: torch.Tensor |
There was a problem hiding this comment.
source_ranks/target_ranks hold measured rank while the bases carry selected rank, and the docstring documents neither. Please add two lines naming which rank each field carries.
| input_shape: Shape of the matrix from which the basis was extracted. | ||
| """ | ||
|
|
||
| basis: torch.Tensor |
There was a problem hiding this comment.
Sibling modules annotate tensors with jaxtyping.Float shape strings; this module is bare torch.Tensor throughout, so runtime shape-checking enforces nothing here. Should we align with its siblings or add a one-line comment recording that the deviation is deliberate?
| if ( | ||
| isinstance(value.rank, bool) | ||
| or not isinstance(value.rank, int) | ||
| or value.rank < 1 |
There was a problem hiding this comment.
#1720's plan lists "deficient" among the explicit-value cases and validates "moments" plural. Please add an exact-value rank-deficient PK case and a variance assertion.
|
|
||
| __all__ = [ | ||
| "DirectLogitAttribution", | ||
| "AttentionHeadRef", |
There was a problem hiding this comment.
The new entries broke __all__'s alphabetical ordering, please restore alphabetical ordering
2 participants
Description
Implements #1720.
Adds a basis-invariant Projection Kernel analysis surface in two reviewable commits:
numerical-rank metadata, roundoff-bound handling, and random-subspace moments;
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
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.make test-prdid not exit successfully because of the threeintegration failures described below):
local HF token lacks access, and one existing GraniteMoeHybrid MPS parity test reproduced the
investigated PyTorch 2.10 MPS numerical divergence.
and notebook checks passed.
Checklist
docs/source/content/projection_kernel.md.Type of change