sae: add scorer correction adapter, ported to the canonical SAE stack - #26
sae: add scorer correction adapter, ported to the canonical SAE stack#26mgagvani wants to merge 2 commits into
Conversation
Trains a lightweight residual adapter over the frozen planner's proposal
scorer, predicting the scorer's own error from SAE latents and correcting
scores additively at inference. The base planner and proposal set are
untouched - only the ranking changes.
Ported off the legacy sparseAE L1 stack onto models/sae.py +
sae_utils.py. sparseAE.py, its colliding 88-line sae_utils.py, and
test_sparseAE.py are deliberately not merged: main now has exactly one SAE
implementation and one checkpoint format.
Port surface was narrow - SparseAutoencoder already exposes .encoder, so
infer_residual_dims needed only a dim source, and Ben had already ported
compute_hidden_activations and load_model_and_sae in new_sae_utils. The
substantive change is activation capture: sae.internal_acts (state hung off
the SAE) becomes an explicit ActivationCapture threaded through
extract_batch_targets.
Adds --latent_source {sae,raw}. The paper's control - the same adapter fed
the dense activation the SAE reconstructs instead of its sparse latents -
was previously run by hand and left no reproducible path. It is now a flag,
persisted into the residual checkpoint and read back by the evaluator, so
the control arm can be re-run exactly.
sae_utils gains three things this needs:
* ActivationCapture.register/remove - the hook handle was previously
discarded, so hooks could not be detached.
* freeze_module / set_eval_mode, carried over from the legacy utils.
* normalize_compiled_state_dict, salvaged from sparseAE.py. A checkpoint
saved from a compiled SAE stores encoder._orig_mod.weight; loading it
into an eager module under strict=False silently dropped every tensor
and yielded a randomly-initialised SAE that still appeared to load.
Verified on CPU against the real block-3 checkpoint: hook fires, captured
activation is (B,1,384), latents are (B,384), and the handle detaches.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5a69e65d1e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| backbone_model, sae, capture, _ = load_model_and_sae( | ||
| model_checkpoint_path=model_checkpoint_path, | ||
| sae_checkpoint_path=sae_checkpoint_path, | ||
| block_idx=block_idx, | ||
| ) |
There was a problem hiding this comment.
When --sae_checkpoint_path points directly to a .pt file whose saved block_index differs from --block_idx, this call hooks the requested backbone block but loads the incompatible SAE without checking its metadata. Because every current block has the same 384-dimensional activation, the mismatch passes all shape checks and silently trains and evaluates on semantically invalid latents, corrupting experiment results; compare the checkpoint's recorded block_index with block_idx before proceeding.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Pull request overview
This PR adds an SAE-informed residual “scorer correction” adapter that learns to predict the proposal scorer’s error from SAE latents (or a reproducible “raw activation” control arm) and applies an additive correction at inference, without changing the base planner or proposal set.
Changes:
- Introduces
sae_scorer_correction.pywith training/eval utilities, residual model, ranking loss, and checkpoint I/O (including persisted--latent_source {sae,raw}). - Updates
sae_utils.pyto make activation capture explicit and removable, adds freeze/eval helpers, and normalizes compiled vs eager SAE checkpoint key formats. - Adds a CLI evaluation script for scorer residual checkpoints.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| src/camera-based-e2e/test_sae_scorer_residual.py | Adds a CLI evaluator for scorer residual checkpoints and metrics reporting. |
| src/camera-based-e2e/sae_utils.py | Adds hook handle lifecycle management and compiled-checkpoint key normalization for SAE loading. |
| src/camera-based-e2e/sae_scorer_correction.py | Implements the residual correction model, training loop, metrics, and checkpoint export/load. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| from __future__ import annotations | ||
|
|
||
| import argparse | ||
| import json | ||
| from collections import defaultdict |
Trains a lightweight residual adapter over the frozen planner's proposal
scorer, predicting the scorer's own error from SAE latents and correcting
scores additively at inference. The base planner and proposal set are
untouched - only the ranking changes.
Ported off the legacy sparseAE L1 stack onto models/sae.py +
sae_utils.py. sparseAE.py, its colliding 88-line sae_utils.py, and
test_sparseAE.py are deliberately not merged: main now has exactly one SAE
implementation and one checkpoint format.
Port surface was narrow - SparseAutoencoder already exposes .encoder, so
infer_residual_dims needed only a dim source, and Ben had already ported
compute_hidden_activations and load_model_and_sae in new_sae_utils. The
substantive change is activation capture: sae.internal_acts (state hung off
the SAE) becomes an explicit ActivationCapture threaded through
extract_batch_targets.
Adds --latent_source {sae,raw}. The paper's control - the same adapter fed
the dense activation the SAE reconstructs instead of its sparse latents -
was previously run by hand and left no reproducible path. It is now a flag,
persisted into the residual checkpoint and read back by the evaluator, so
the control arm can be re-run exactly.
sae_utils gains three things this needs:
discarded, so hooks could not be detached.
saved from a compiled SAE stores encoder._orig_mod.weight; loading it
into an eager module under strict=False silently dropped every tensor
and yielded a randomly-initialised SAE that still appeared to load.
Verified on CPU against the real block-3 checkpoint: hook fires, captured
activation is (B,1,384), latents are (B,384), and the handle detaches.