Spaces:
Sleeping
Sleeping
Commit History
-
Fix Pythia QKV extraction layout; add W_QK correctness tests
extract_pythia_qkv split the combined query_key_value.weight with qkv.chunk(3, dim=0), but HF GPT-NeoX stores QKV output features per-head-interleaved as [h0_Q, h0_K, h0_V, h1_Q, h1_K, h1_V, ...] (confirmed by reading GPTNeoXAttention.forward). The naive chunk scrambled Q, K, V across heads; downstream W_QK was the wrong bilinear form. Re-extract via per-head view + slice. Same fix applied to extract_pythia_biases. New tests under tests/test_wqk_* build tiny HF models with known weights, then verify (a) extracted W_Q/K/V match an architecture- correct reference, (b) x_i^T W_QK x_j equals the score HF's own projection computes on the same inputs, and (c) the factored d_head x d_head SVD path matches a full d_model x d_model SVD. Covers gpt2 and pythia-70m-deduped; same fixture pattern extends to Llama/Mistral/Gemma2 in a follow-up. tests/conftest.py prepends the worktree's src/ to sys.path so tests exercise the local source rather than the env-installed package (which resolves to the …5ecd8f7
-
Wire QK decomposition weight types into primary pipeline
W_Q_gram, W_K_gram, QK_alignment are already computed in head_analyzer.py. This commit exposes weight_types through config.yaml → run_primary.py → process_model() → primary.smk so they can be overridden at runtime. Tests updated to cover all 6 weight types; histogram test restricted to raw weight types since gram entries exceed standard bin range. compute_qk_decomposition.py removed (data now lives in primary dataset). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>a07e79a -
Rename src modules for head/pair symmetry; add pair_figures.smk; split eval_metrics
src/transformer_analysis/: weight_analysis → head_pipeline attn_head_analysis → head_analyzer metrics → head_metrics correlation_analysis → pair_pipeline head_correlations → pair_analyzer eval_metrics.py (new — implementation extracted from compute_perplexity.py) scripts/: plot_correlations.py → plot_pair_correlations.py plot_corr_figures.py → plot_pair_correlations_batch.py compute_perplexity.py → deleted (impl in eval_metrics; thin wrapper run_eval.py) run_eval.py (new thin wrapper) rules/: correlations.smk — switched to sentinel output (consistent with other rules) pair_figures.smk (new — regenerate pair figures from .npz without re-running correlations) eval.smk — updated to call run_eval.py Snakefile: added correlations_all, pair_figures_all convenience targets Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>a81fd74 -
feat: add Snakemake pipeline, MLflow tracking, metrics refactor, tests
Pipeline / workflow: - Snakefile + config.yaml orchestrate primary → transform → merge stages - rules/primary.smk, transform.smk, correlations.smk, merge.smk - Sentinel-file pattern (done/*.done) avoids directory-output conflicts - Incremental cross-model merge: appends new models without clobbering; --refresh flag for explicit replacement after re-transform - Per-model checkpoint merge rule for Pythia step-collapse CLI wrappers (scripts/): - run_primary.py — process_model() + MLflow logging - run_transform.py — reprocess_metrics() + MLflow logging (git sha tagged) - run_correlations.py — thin wrapper for run_multi_circuit_analysis() - run_merge.py — cross-model incremental merge + checkpoint merge Metrics refactor: - src/transformer_analysis/metrics.py: canonical home for all stat functions, bin configs, normality_metrics, singular_value_metrics, and adaptive binning helpers (make_weight_bins, make_sv_bins) - histogram_utils.py: backward-compat re-export shim; no logic c…24e1932