extract-features-clamp-inference

name	extract-features-clamp-inference
description	Extract residual-stream features and run clamp steering for open-source language models using an Anthropic-style SAE pipeline. Use when Codex needs to collect intermediate activations, train a sparse autoencoder, inspect candidate features, and run baseline-vs-clamped inference experiments on chat or base checkpoints.

Extract Features And Clamp Inference

Use this skill to run the end-to-end SAE feature workflow with the bundled scripts.

Script Map

• scripts/build_conversations_hf.py: build mixed HF conversation corpus JSONL.
• scripts/collect_activations.py: collect residual activations from a target layer.
• scripts/train_sae_anthropic.py: train an Anthropic-style SAE.
• scripts/feature_inspect_conversations_topk.py: find top candidate features quickly.
• scripts/feature_topic_summary.py: summarize likely topics from top prompts.
• scripts/chat_clamp_feature.py: one-shot baseline vs clamped generation.
• scripts/chat_clamp_eval_batch.py: batch A/B clamp evaluation with hit-rate metrics.
• scripts/chat_steer_feature.py: residual direction steering (non-clamp).

Workflow

1. Build corpus (optional if you already have conversations):

python scripts/build_conversations_hf.py \
  --out-file data/conversations_hf_scale_v3.jsonl \
  --max-conversations 18000 \
  --ultrachat-target 10000 \
  --glaive-target 5000 \
  --belle-target 5000 \
  --max-turns 8

2. Collect activations from an intermediate layer:

python -u scripts/collect_activations.py \
  --model-id zai-org/glm-4-9b-chat-hf \
  --conversations-file data/conversations_hf_scale_v3.jsonl \
  --layer-index -20 \
  --max-prompts 8000 \
  --max-length 768 \
  --tokens-per-prompt 128 \
  --sample-strategy first \
  --collect-from assistant_gt \
  --assistant-span-scope all \
  --truncation-side left \
  --apply-chat-template \
  --chunk-rows 80000 \
  --out-prefix out/chat_assistant_gt_scale_v3/acts \
  --dtype float16 \
  --trust-remote-code

3. Train SAE:

python -u scripts/train_sae_anthropic.py \
  --chunks-glob 'out/chat_assistant_gt_scale_v3/acts_chunk_*.pt' \
  --d-sae 65536 \
  --batch-size 1024 \
  --epochs 10 \
  --lr 1e-4 \
  --optimizer adamw \
  --weight-decay 0.0 \
  --l1-coeff 5e-5 \
  --l1-warmup-steps 1000 \
  --l1-reduction sum \
  --scalar-normalize \
  --save-dir out/chat_assistant_gt_scale_v3/sae_anthropic_v1_d65k_l1_5e5

4. Inspect features:

python -u scripts/feature_inspect_conversations_topk.py \
  --manifest out/chat_assistant_gt_scale_v3/acts_manifest.json \
  --checkpoint out/chat_assistant_gt_scale_v3/sae_anthropic_v1_d65k_l1_5e5/sae_final.pt \
  --conversations-file data/conversations_hf_scale_v3.jsonl \
  --out-json out/chat_assistant_gt_scale_v3/feature_inspect_conversations_topk.json \
  --top-features 300 \
  --candidate-features 1200 \
  --top-conversations-per-feature 12 \
  --dedupe-mode template \
  --label-field last_user \
  --batch-size 1024 \
  --device cuda

Optional quick topic summary:

python scripts/feature_topic_summary.py \
  --inspect-json out/chat_assistant_gt_scale_v3/feature_inspect_conversations_topk.json \
  --out-json out/chat_assistant_gt_scale_v3/feature_topics_top10.json \
  --top-features 20

5. Run clamp steering:

python -u scripts/chat_clamp_feature.py \
  --model-id zai-org/glm-4-9b-chat-hf \
  --checkpoint out/chat_assistant_gt_scale_v3/sae_anthropic_v1_d65k_l1_5e5/sae_final.pt \
  --feature-id 30376 \
  --steer-layer -20 \
  --steer-all-tokens \
  --clamp-quantile 0.999 \
  --clamp-mult-quantile 4.0 \
  --quantile-manifest out/chat_assistant_gt_scale_v3/acts_manifest.json \
  --prompt "Explain DNS like I am 12." \
  --temperature 0.0 \
  --trust-remote-code

6. Evaluate at scale (baseline vs clamped):

python -u scripts/chat_clamp_eval_batch.py \
  --model-id zai-org/glm-4-9b-chat-hf \
  --checkpoint out/chat_assistant_gt_scale_v3/sae_anthropic_v1_d65k_l1_5e5/sae_final.pt \
  --feature-id 30376 \
  --steer-layer -20 \
  --steer-all-tokens \
  --clamp-quantile 0.999 \
  --clamp-mult-quantile 4.0 \
  --quantile-manifest out/chat_assistant_gt_scale_v3/acts_manifest.json \
  --prompts-file data/steer_eval_neutral_v1.txt \
  --max-prompts 25 \
  --topic-regex '(movie|film|cinema|ticket|actor|director|screenplay)' \
  --ignore-case \
  --temperature 0.0 \
  --out-json out/chat_assistant_gt_scale_v3/clamp_eval_feature_30376_q999x4p0.json \
  --trust-remote-code

Guardrails

• Match layer index across collection, training, and steering.
• Use chat template for chat models (--apply-chat-template at collection time; default chat template path at steering time).
• Treat reconstruction quality as a gate before steering. If explained variance is low, do not trust steering results.
• Prefer quantile clamp targets (--clamp-quantile) before aggressive raw max multipliers.
• If steering is weak, increase sparsity pressure (l1) and rerun training; high reconstruction with very high L0 often needs stronger clamp to show takeover.

Quality Check Snippet

Run after training:

python - <<'PY'
import torch
import torch.nn.functional as F

acts=torch.load('out/chat_assistant_gt_scale_v3/acts_chunk_0000.pt',map_location='cpu').float()
ckpt=torch.load('out/chat_assistant_gt_scale_v3/sae_anthropic_v1_d65k_l1_5e5/sae_final.pt',map_location='cpu')
sd=ckpt['state_dict']
scale=float(ckpt.get('activation_scale',1.0))

x=acts*scale
W=sd['encoder.weight'].float()
b=sd['encoder.bias'].float()
D=sd['decoder.weight'].float()

z=F.relu(F.linear(x,W,b))
xhat=F.linear(z,D)
mse=((x-xhat)**2).mean().item()
baseline=(x**2).mean().item()
l0=(z>0).float().sum(dim=1).mean().item()
print('explained_frac',1-mse/(baseline+1e-12))
print('l0',l0)
PY

Target a useful reconstruction/sparsity tradeoff before interpreting steering effects.