sonar-sae / ops /CODEX_PREREVIEW_COMP.md
nickypro's picture
Upload ops/CODEX_PREREVIEW_COMP.md with huggingface_hub
56973dd verified
|
Raw
History Blame Contribute Delete
19.8 kB
# Pre-Launch Review: SONAR Composition Experiments
Scope: `parascopes/layerwise/src/tae_comp_anatomy.py`, `tae_comp_stack.py`, `tae_comp_deletion.py`, `tae_comp_language.py`, `tae_comp_features.py`, with shared deps `tae_composition.py`, `tae_probe.py`, `tae_erasure.py`, `sonar_interp_lab.py`, `train_sae_topk.py`.
## Findings
1. `parascopes/layerwise/src/tae_comp_stack.py:57` - MAJOR - `W2` is fit on all sentence pairs, and `pair_cos` at `:60` is a self-fit score, not a held-out pair-map score.
One-line fix: split pairs before fitting (`ntr=int(.85*n)`), fit `W2` on train pairs, and report pair/chain metrics only on held-out rows.
2. `parascopes/layerwise/src/tae_comp_stack.py:84-88` - MAJOR - `left_fold_cos`, `right_fold_cos`, and `mean_baseline_cos` are computed on all chains, while `direct_ridge_cos_heldout` is computed only on `ntr:`; these numbers are not comparable and include train-like rows.
One-line fix: compute all depth metrics on the same held-out slice (`left[ntr:]`, `right[ntr:]`, `meanb[ntr:]`, `Zfull[ntr:]`).
3. `parascopes/layerwise/src/tae_comp_stack.py:67` - MAJOR - chain evaluation reuses the same paragraphs used to train the pair map, so consecutive pairs inside the evaluated chains can overlap the pair-map training distribution exactly.
One-line fix: build pair-map training pairs and chain-eval chains from disjoint shards/text indices, or split `texts` before constructing either dataset.
4. `parascopes/layerwise/src/tae_comp_stack.py:90-94` - MINOR - decode examples use `left[:n_decode]`, not held-out chain rows, so qualitative examples can be train-contaminated.
One-line fix: decode `left[ntr:ntr + args.n_decode]` and compare to `full_txt[ntr:ntr + args.n_decode]`.
5. `parascopes/layerwise/src/tae_comp_anatomy.py:82-83` - MAJOR - the PC basis is fit using all `ZA`/`ZB`, including rows later treated as test rows for PC-transfer metrics.
One-line fix: compute `mean` and `V` from train rows only, e.g. `torch.cat([ZA[:ntr], ZB[:ntr]])`.
6. `parascopes/layerwise/src/tae_comp_anatomy.py:104-122` - MAJOR - `pc0_adder` and `norm_model` fit their regressions on the same held-out rows used to report R2, so the reported arithmetic fit is in-sample.
One-line fix: fit coefficients on train-derived PC/norm features and report R2 on `ntr:` only.
7. `parascopes/layerwise/src/tae_comp_anatomy.py:125-137` - MINOR - the axis dictionary uses all rows for both PC basis/projections and attribute correlations, with no held-out confirmation or baseline.
One-line fix: choose axis names on train rows and report the winning attribute correlation/R2 again on held-out rows.
8. `parascopes/layerwise/src/tae_comp_deletion.py:50-51` - MAJOR - PC corruptions use a mean and SVD basis fit with evaluation rows included, so `top64_pcs_only`/`tail_pcs_only` leak test distribution into the corruption operator.
One-line fix: fit `mean` and `V` on `ZA[:ntr]` only, and apply that fixed basis to `ZA[ntr:]`.
9. `parascopes/layerwise/src/tae_comp_deletion.py:77-79` - MINOR - `te` selects only the first decode window, while cosine metrics use the full held-out set; decode-locality metrics are therefore a small-window qualitative measure, not the same population as `cos_to_true_AB`.
One-line fix: either label these fields as decode-window metrics or compute SBERT locality in batches over the full held-out set.
10. `parascopes/layerwise/src/tae_comp_language.py:60-64` - FATAL - `n = args.n_texts` assumes the filtered `keep` list reached the requested length; if fewer texts pass the length filter, labels from `repeat_interleave(n)` no longer match `Zall` block boundaries and can become wrong or empty-class.
One-line fix: set `n = Z0.shape[0]` immediately after `Z0 = X[keep].to(device)` and optionally assert `n > 0`.
11. `parascopes/layerwise/src/tae_comp_language.py:67` - MAJOR - classifier standardization is fit on all language rows before train/test split, leaking held-out distribution into `train_linear`.
One-line fix: compute mean/std on the training split inside each classifier evaluation and apply those stats to held-out rows.
12. `parascopes/layerwise/src/tae_comp_language.py:75-78` - MAJOR - binary `train_linear` calls filter each eng-vs-L subset after global permutation but rely on the helper's first-85-percent split; there is no explicit stratification, so small/default runs can produce unstable or class-skewed test slices.
One-line fix: create explicit stratified train/test indices per binary language task and pass already-split data or add split support to `train_linear`.
13. `parascopes/layerwise/src/tae_comp_language.py:97-105` - MAJOR - mean-shift swap metrics use class means fit from `mus` and then report before/after cosine over all rows, including rows that contributed to those means.
One-line fix: fit language means on train rows and report swap cosine/decode agreement on held-out rows only.
14. `parascopes/layerwise/src/tae_comp_language.py:117-120` - MAJOR - LEACE before/after probes standardize with `Z2.mean/std` from all rows, including probe test rows; this leaks test distribution into both reported accuracies.
One-line fix: compute standardization stats on the probe train split only, reusing the same stats for erased and unerased held-out rows.
15. `parascopes/layerwise/src/tae_comp_language.py:132-149` - MAJOR - cross-language composition fits `W2` on all pairs and evaluates `pure`/`mixed` on `pairs[:50]`, so the English baseline is self-fit and optimistic.
One-line fix: fit `W2` on train pairs and evaluate both pure and mixed composition on held-out pairs, e.g. `pairs[ntr:ntr+50]`.
16. `parascopes/layerwise/src/tae_comp_features.py:45-51` - MAJOR - default SAE path is BatchTopK (`*_btk.pt`), and `TopKSAE.encode()` still uses batch-global top-k at eval; active feature sets depend on batch composition and are not per-sample stable across separate `CA`, `CB`, and `CAB` passes.
One-line fix: for analysis, force per-sample top-k from `sae.preacts()`/`relu.topk(sae.k)` or implement the promised fixed eval threshold for BatchTopK checkpoints.
17. `parascopes/layerwise/src/tae_comp_features.py:70-81` - MAJOR - shared-feature additivity fits and reports R2 on the same flattened activations, so the `r2` is in-sample and can overstate additivity.
One-line fix: split pairs first, fit `sol` on shared train activations, and report R2 on shared held-out activations.
18. `parascopes/layerwise/src/tae_comp_features.py:84-91` - MINOR - `top_droppable_feats` only counts features active in A and dropped in AB, despite the stated goal of checking whether A or B survives less.
One-line fix: compute separate `drop_rate_A` and `drop_rate_B` tables and add a position-asymmetry summary.
19. `parascopes/layerwise/src/tae_comp_features.py:85` - MINOR - `base = mA.float().sum(0).clamp_min(20)` suppresses low-support features in the denominator but still allows them into `topk`, making support unclear for reported droppable features.
One-line fix: mask features with support `<20` out of the ranking and include `support` and `dropped` counts in the output rows.
20. `parascopes/layerwise/src/tae_composition.py:36-40` - MINOR - ridge regularizes the bias row because the identity penalty covers the appended constant feature.
One-line fix: use a penalty matrix with the final diagonal entry set to zero before solving.
## Checked Non-Finding
`parascopes/layerwise/src/tae_comp_anatomy.py:78-80` correctly accounts for the ridge/apply convention. `ridge(X, Y)` returns `W` shaped `[features + 1, 1024]`, and `apply_w(X, W)` computes `[X, 1] @ W`; therefore `Wmat[:D]` is the row-vector block for `z_A @ W_A_row`, and the script's `.t()` converts it to the documented column-vector convention `output = WA @ z_A + WB @ z_B`.
Also, `parascopes/layerwise/src/tae_comp_deletion.py:78-79` uses Python list slicing with a `slice` object (`A_txt[te]`, `B_txt[te]`), which is valid and should not crash.
## CHANGES
- `tae_comp_anatomy.py`: fit the PC basis/mean on train rows only, fit `pc0_adder` and `norm_model` regressions on train rows, report their R2 on held-out rows, and add held-out confirmation for train-selected axis dictionary attributes.
- `tae_comp_stack.py`: split source texts before dataset construction so pair-map training and chain evaluation use disjoint text sets; compute `pair_cos` on held-out pairs; report left/right/direct/mean depth metrics on the same held-out chain slice and decode from that slice.
- `tae_comp_deletion.py`: fit PC corruption mean/SVD basis and noise scale from train rows only before applying corruptions to held-out rows.
- `tae_comp_language.py`: set `n` from filtered texts with a clear minimum-size error, use per-language train/test splits, fit standardization/class means/swaps/LEACE/cross-language pair maps on train rows only, and report metrics on held-out rows/pairs.
- `tae_comp_features.py`: force SAE eval to per-sample top-k by setting `sae.batchtopk = False`, and fit the shared-feature additivity regression on a random half of shared activations while reporting R2 on the other half.
## Severity Table
| Severity | Count | Findings |
| --- | ---: | --- |
| FATAL | 1 | 10 |
| MAJOR | 13 | 1, 2, 3, 5, 6, 8, 11, 12, 13, 14, 15, 16, 17 |
| MINOR | 6 | 4, 7, 9, 18, 19, 20 |
## WAVE 2
Reviewed and patched the three new pre-launch scripts:
1. `parascopes/layerwise/src/tae_canonicity_fixoff.py:131-135` - MAJOR - input normalization statistics were fit on the concatenation of train shards plus the final validation shard. This leaked held-out distribution into FVU and dictionary matching metrics.
Fix applied: load train shards and validation shard separately, fit `x_mean`/`scale` on train rows only, and normalize validation with those train-only stats.
2. `parascopes/layerwise/src/tae_canonicity_fixoff.py:37` - MAJOR - the archetypal decoder included an unconstrained relaxation term `relax * R`, but `R` had no penalty. Because the effective decoder is normalized after adding `R`, `R` could grow until the archetypal convex-combination constraint became mostly cosmetic.
Fix applied: add an archetypal-only `R` penalty during training using the existing `--l2` coefficient, preserving the effective decoder `softmax(P) @ anchors + relax * R` and its differentiable gradient path.
3. `parascopes/layerwise/src/tae_canonicity_fixoff.py:82-110` - MAJOR - `match_metrics` crashed when either model had no features above `fire_min`, because empty decoder/correlation matrices were reduced with `max`/`median`.
Fix applied: return the live counts and `NaN` match metrics for this degenerate case instead of crashing.
4. `parascopes/layerwise/src/train_composer_textloss.py:36` - FATAL - default `--device cuda` crashed immediately on CPU-only hosts, unlike the adjacent experiment scripts.
Fix applied: use `cuda` only when `torch.cuda.is_available()`, otherwise default to CPU.
5. `parascopes/layerwise/src/train_composer_textloss.py:43-54` - MAJOR - the composer split assumes the final shard is held out, but a one-shard run made the train corpus empty and later failed with obscure tensor concatenation errors.
Fix applied: require at least two shards and raise clear errors if train or eval sentence-pair extraction yields zero pairs.
6. `parascopes/layerwise/src/train_composer_textloss.py:72-78` - MAJOR - if all train pairs were removed by token-length filtering, training silently ran zero optimizer steps and still wrote a "finetuned" artifact.
Fix applied: raise if no train pairs survive the token filter.
7. `parascopes/layerwise/src/train_composer_textloss.py:79-81` - MAJOR - if all evaluated held-out pairs were filtered by token length, `decoder_ce` reported `0.0` via `max(1, m)`, which is a false success metric.
Fix applied: precompute held-out token IDs, require at least one eval pair under the token cap, and reuse those IDs in `evaluate()`.
8. `parascopes/layerwise/src/train_composer_textloss.py:126-141` - MINOR - wandb failures after successful init could still abort training/finalization.
Fix applied: wrap periodic and final wandb logging in `try` blocks and disable/continue on failure.
9. `parascopes/layerwise/src/tae_comp_unbind.py:52` - MINOR - unbinding split dimensions were hard-coded as `1024`, despite the dependency already exposing `D_SONAR`.
Fix applied: import and use `D_SONAR` for the recovered A/B split.
10. `parascopes/layerwise/src/tae_comp_unbind.py:39-51` - MINOR - very small or empty pair sets produced downstream decode/index/SVD failures.
Fix applied: add minimum-pair and non-empty train/eval split checks.
11. `parascopes/layerwise/src/tae_comp_unbind.py:80-87` - MINOR - length-quartile reporting could emit NaNs for empty buckets on tiny eval sets.
Fix applied: only include non-empty quartile buckets.
## WAVE 2 Checked Non-Findings
- `parascopes/layerwise/src/tae_comp_unbind.py:93-101` already fits the order-subspace SVD basis and `mean_delta` from train rows only (`D[:ntr]`), then applies them to held-out rows. No train/test leakage found there.
- `parascopes/layerwise/src/train_composer_textloss.py:65-70` fits ridge initialization on all train pairs, and `evaluate()` uses only `ev_pairs` from the separate final shard. This matches the intended train/eval split.
- `parascopes/layerwise/src/tae_canonicity_fixoff.py:83` encodes `Xval` in one batch inside `match_metrics`; this preserves BatchTopK's batch-global behavior for a single fixed evaluation batch, as requested.
## WAVE 2 Changes
- `tae_canonicity_fixoff.py`: removed validation leakage from normalization, kept anchors train-only, added archetypal relaxation regularization, and made matching robust to no-live-feature runs.
- `train_composer_textloss.py`: added CPU-safe default device, explicit split/data/token-filter validation, precomputed eval token IDs, and resilient wandb logging.
- `tae_comp_unbind.py`: added small-data guards, used `D_SONAR` for recovery slicing, skipped empty quartile buckets, and bounded the rank-8 projection for tiny training sets.
- Verification: `python3 -m py_compile parascopes/layerwise/src/tae_comp_unbind.py parascopes/layerwise/src/tae_canonicity_fixoff.py parascopes/layerwise/src/train_composer_textloss.py` passes.
## WAVE 2 Severity Table
| Severity | Count | Findings |
| --- | ---: | --- |
| FATAL | 1 | 4 |
| MAJOR | 6 | 1, 2, 3, 5, 6, 7 |
| MINOR | 4 | 8, 9, 10, 11 |
## WAVE 3
Reviewed and patched the three pre-launch scripts:
1. `parascopes/layerwise/src/tae_weight_svd.py:21-27` - FATAL - `get_cross_attn_kv()` assumed the object passed in had `decoder.layers`, but callers pass `SonarDecoderCE`; per `sonar_grad.py`, `SonarDecoderCE.decoder` is the fairseq2 model whose own `.decoder` holds the transformer `.layers`.
Fix applied: unwrap both `SonarDecoderCE -> fairseq2 model` and `model -> transformer decoder`, then discover cross-attention modules at runtime with fallbacks for common fairseq2 names and a clear error listing available child names.
2. `parascopes/layerwise/src/tae_weight_svd.py:36,48,69` - MAJOR - decode naming used `mean + alpha * sigma * sqrt(d) * v` with default `alpha=3.0`; since `sigma` was per-dim RMS, this made the perturbation about three typical centered vector norms and could decode extreme off-manifold directions.
Fix applied: default `--alpha` is now `0.5`, and the dose is `alpha * mean((X - mean).norm()) * v`. The JSON records the dose convention and numeric RMS norm.
3. `parascopes/layerwise/src/tae_weight_svd.py:74` - MINOR - the SAE-span proxy uses the first 4096 decoder rows, not top-firing features, so the reported overlap can be basis/order dependent.
Fix applied: kept the proxy to avoid adding activation dependencies, but added `sae_span_proxy` metadata in JSON documenting `first_4096_decoder_rows_normalized_not_top_firing`.
4. `parascopes/layerwise/src/tae_ica_axes.py:23-38` - MAJOR - FastICA returned only `W`, so non-convergence silently produced apparently valid axes.
Fix applied: `fastica()` now returns convergence metadata (`converged`, `n_iter`, `max_delta`), prints a warning on non-convergence, and writes the metadata under the existing JSON root without removing existing keys.
5. `parascopes/layerwise/src/tae_ica_axes.py:66-70` - MAJOR - whitening used an untyped CPU identity and unclamped eigenvalues; this was brittle if inputs moved devices/dtypes and could amplify tiny negative numerical eigenvalues.
Fix applied: build the jitter identity on the covariance dtype/device, clamp eigenvalues before inverse square root, record kept whitening dimensions, and fail clearly if no directions survive.
6. `parascopes/layerwise/src/tae_ica_axes.py:112-116` - MINOR - causal deltas for continuous heuristic attributes were raw-unit differences, unlike binary attributes where raw mean difference is directly interpretable.
Fix applied: binary deltas remain raw mean differences; continuous deltas are now standardized by the held-out reference attribute mean/std. The `bestk is None` path still returns `null` without crashing.
7. `parascopes/layerwise/src/tae_lang_inert.py:65-68` - MAJOR - LEACE inputs and language deltas relied on mixed device/dtype tensors from encode/decode outputs; this was unnecessarily fragile around CPU erasure and amplification.
Fix applied: force LEACE inputs, labels, held-out embeddings, and amplified embeddings to CPU `float32` before erasure or shifting.
8. `parascopes/layerwise/src/tae_lang_inert.py:73-98` - MAJOR - amp and English-flag probes recomputed English references inline, making the intended held-out slice alignment easy to break during later edits.
Fix applied: cache `ref_en = txt_en[ntr:]` and its SBERT embeddings once, add a `probe()` length assertion, and reuse the same English held-out references for all English-flag conditions.
9. `parascopes/layerwise/src/tae_lang_inert.py:55-56` and `parascopes/layerwise/src/tae_ica_axes.py:60-61` - MINOR - tiny filtered datasets could leave empty held-out slices, causing divide-by-zero or decode-index errors later.
Fix applied: both scripts now require at least two train rows and one held-out row with a clear `RuntimeError`.
## WAVE 3 Checked Non-Findings
- `tae_weight_svd.py` uses `torch.linalg.qr(Wd[:4096].t())` with the default reduced/economy mode, so the resulting `Q` has compatible shape for `(Vr @ Q)` and does not require a shape fix.
- `tae_ica_axes.py` already names axes and computes attribute correlations on held-out rows (`s[ntr:]`, `y[ntr:]`) while drawing extreme examples from train rows only.
- `tae_lang_inert.py` uses `leace_eraser(X, y)` in the expected order: embeddings `[n, d]` first and scalar labels `[n]` second.
- `tae_lang_inert.py` English-flag amp probes compare decodes to `txt_en[ntr:]`, which is semantically the right reference for decoding held-out French-origin embeddings with the English flag.
## WAVE 3 Changes
- `tae_weight_svd.py`: robust fairseq2 decoder/layer/cross-attention discovery, clearer runtime diagnostics, safer default dose scaling, and JSON metadata for dose and SAE-span proxy.
- `tae_ica_axes.py`: FastICA convergence reporting, safer whitening eigenspectrum handling, tiny-data guards, bounded axis count, and standardized continuous causal deltas.
- `tae_lang_inert.py`: explicit float32 CPU erasure inputs, held-out size guard, cached aligned English references, and probe length assertions.
- Verification: `python3 -m py_compile parascopes/layerwise/src/tae_weight_svd.py parascopes/layerwise/src/tae_ica_axes.py parascopes/layerwise/src/tae_lang_inert.py` passes.
## WAVE 3 Severity Table
| Severity | Count | Findings |
| --- | ---: | --- |
| FATAL | 1 | 1 |
| MAJOR | 5 | 2, 4, 5, 7, 8 |
| MINOR | 3 | 3, 6, 9 |