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