scale-aware compression: rep_selection + merge strategy#34
Merged
Conversation
Proposal, design, specs, and tasks for scale-aware representative selection (norm-proximity + KL-ablation + log n_fires scoring) and a merge compression strategy that freq-weight-averages cluster norms to preserve decoder norm mass rather than hard-zeroing. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Adds two opt-in extensions to Compressor while keeping defaults unchanged (strategy=zero, rep_selection=n_fires). rep_selection=scale_aware scores each cluster candidate by norm-proximity to median (0.4) + normalised KL-ablation (0.4) + log n_fires (0.2). Falls back to n_fires when every kl_ablate value is NaN (geometry-only confirmer) and emits a UserWarning. strategy=merge rescales the surviving representative W_dec row to a freq-weighted (or simple-mean) average of cluster norms while zeroing non-rep encoder columns, biases, and decoder rows. This preserves total norm mass (scale_compression_ratio == 1.0 under simple_mean) instead of discarding it. Per-cluster scale stats land on ClusterPlan (cluster_norm_mean, cluster_norm_std, merged_norm) and aggregate scale_compression_ratio sits on CompressionReport. JSON round-trip preserves all new fields; missing fields default cleanly so older reports still load. 29 new tests cover rep selection, merge arithmetic, scale stats, JSON round-trip, and an end-to-end integration run. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Two short comments per PR review: (1) note the ~36 MB cost of loading W_dec in plan() for scale_aware, (2) explain why merge skips encoder columns (would expand the surviving feature firing footprint). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
apply_merge now emits a UserWarning if a cluster's merged_norm collapses to ~0 (every member had a zero-norm decoder row in the source). Rare in practice but surfaces silent reps that would otherwise pass through unnoticed. 2 new tests cover the warning path and confirm healthy clusters stay quiet. Also adds a comment in _compute_scale_compression_ratio pointing at the design.md open question about singleton exclusion, so a future reader hits the rationale without grepping the change folder. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
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
1 participant
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.
Summary
Compressor(rep_selection="scale_aware")— norm-proximity (0.4) + KL-ablation (0.4) + log n_fires (0.2) scoring, withUserWarningfallback to n_fires when every kl_ablate is NaN (geometry-only confirmer path).Compressor(strategy="merge", merge_mode="freq_weighted"|"simple_mean")— rescales the surviving representative W_dec row to a weighted average of cluster norms instead of discarding non-rep mass.ClusterPlan(cluster_norm_mean,cluster_norm_std,merged_norm) andCompressionReport.scale_compression_ratio(= 1.0 undermerge+simple_mean).strategy="zero",rep_selection="n_fires"); JSON round-trip is additive and old reports load with safe defaults.openspec/changes/scale-aware-compression/for proposal, design, specs, and tasks.Test plan
pytest tests/compression/— 79 passed (including 29 new)pytest tests/— 534 passedruff check polygram/compression/ tests/compression/🤖 Generated with Claude Code