This repo investigates whether structured nonlinear probes — quadratic and low-rank quadratic probes whose gradients have closed-form expressions — produce better activation-steering directions than linear probes for concepts in language models.
| Probe | Score | Steering gradient |
|---|---|---|
| Linear | s(h) = wᵀh + b |
w (state-independent) |
| Quadratic | s(h) = hᵀAh + wᵀh + b |
(A + Aᵀ)h + w |
| Low-rank quadratic | s(h) = Σᵣ λᵣ (uᵣᵀh)² + wᵀh + b |
2 Σᵣ λᵣ (uᵣᵀh) uᵣ + w |
| Degree-m tensor (library only) | s(h) = Σᵣ λᵣ Πq (uᵣ,qᵀh) |
Σᵣ λᵣ Σq (Πp≠q uᵣ,pᵀh) uᵣ,q |
All nonlinear entries have a closed-form gradient that is interpretable in terms of the learned directions. The primary experiment configuration evaluates the linear, quadratic, and low-rank quadratic probes; the degree-m tensor probe is available as a tested library component.
nlsteer/ # importable library: probes, hooks, eval, transfer, utils
stages/ # runnable experiment entry points
configs/ # YAML configs for models, datasets, probes
registry/ # hypothesis registry + methodology
analysis/ # figure generation + hypothesis validation + power calcs
scripts/ # rerun and figure-generation helpers
tests/ # pytest unit tests for the domain library
paper/ # LaTeX source, compiled PDF, figures, and tables
# Install (CPU/dev only)
pip install -e ".[dev]"
# Install with LM and analysis extras
pip install -e ".[lm,analysis,dev]"
# Run CPU tests
pytest -q -m "not integration and not gpu"
# Lint + type-check
ruff check . && mypy nlsteerThe authoritative writeup is paper/main.pdf, with source
in paper/main.tex. The repository includes the compiled
paper, paper figures, and generated paper tables. Raw JSONL run outputs are
not tracked; they are generated locally under the ignored results/ tree.
The evaluation-grid manifest is registry/experiment_grid.yaml. The
evaluation uses prompt-last-token steering, deterministic stratified
train/match/eval slices, shared prompt-suffix tokenization for activation
caching and completion scoring, explicit control upper bounds, and two
probe-selection modes. Generated JSONL files include realized_inputs
metadata recording the effective run configuration after stage-specific
wrappers are applied.
Main verdict: structured nonlinear probes show one replicated
control-calibrated win among the five primary layer-12 settings, but the
broader pattern does not support robust nonlinear superiority. The secondary
best_trained condition has the same SST-2 pythia-410m win.
After generating a complete local results/ tree, validate the hypothesis
verdicts and render the paper figures and tables with:
python -m analysis.hypotheses --results results/matched_accuracy
RESULTS_ROOT=results/matched_accuracy \
BEST_RESULTS_ROOT=results/best_trained \
bash scripts/run_all_figures.sh
cd paper && latexmk -pdf -interaction=nonstopmode -halt-on-error main.texComparison protocol locked in registry/methodology.yaml. Highlights:
- Probe-selection modes:
matched_accuracyis the controlled comparison;best_trainedlets nonlinear probes use their full classification capacity. - Match/eval separation: the match slice is used for probe matching and alpha calibration; the held-out eval slice is used for steering effects.
- Prompt policy: caches and completion scoring both use
add_special_tokens=Falseand the same left-truncated prompt suffix, with room reserved for the longer class completion. - α calibration: the primary operating point is matched perturbation norm; alpha sweeps are reported separately as sensitivity analyses.
- Causal metrics: per-dataset primary metric is independent of the probe family: the implemented datasets use log-prob shift on contrastive completions, never the probe score itself.
- Controls: norm-matched random-direction steering and shuffled-label probes trained under the same checkpoint-selection policy without exposing real match labels.
- Safety outputs: generated ToxicChat JSONL files do not store raw prompt
text; numeric per-prompt effect vectors are retained for paired inference.
ToxicChat is listed as CC-BY-NC-4.0 in
configs/dataset/toxicchat.yaml; downstream users should follow the dataset license and terms.
See registry/methodology.yaml for full details.
| Stage | Module | Output |
|---|---|---|
| 0 | repository setup | repository skeleton |
| 1 | stages.stage1_cache |
cached residual-stream activations |
| 5 | stages.stage5_mvp_sst2 |
SST-2 steering run (Pythia-410M) |
| 6 | stages.stage6_multi_dataset |
CoLA and related binary-classification runs |
| 8 | stages.stage8_safety |
ToxicChat safety run |
| 9 | stages.stage9_transfer |
exploratory cross-model probe transfer |
Probe training, alpha calibration, hook construction, and evaluation metrics
live in the importable nlsteer/ package. Stage runners write JSONL files.
Generated run outputs are ignored by Git.
scripts/rerun_experiments.sh and scripts/run_full_program.sh replay an
existing generated result tree by reading JSONL metadata from
SOURCE_RESULTS (default: results/matched_accuracy) and writing fresh local
outputs under results/{matched_accuracy,best_trained}/. These helpers are
metadata-driven replay schedulers; a fresh clone does not include the required
source tree. For exact replays, point SOURCE_RESULTS at an existing generated
tree and set OUT_ROOT to a separate destination so replay outputs do not
overwrite the source inputs.
- All randomness is seeded; seeds are recorded in JSONL
metarows. - Cache keys are content-hashed so identical configs are guaranteed to reuse the same activations. Dataset semantics and optional model revisions are included in the activation-cache identity.
- JSONL metadata records package versions; newly generated runs also include the current git commit when available.
- Raw JSONL outputs, logs, activation caches, and checkpoints are ignored by Git.
- GPU non-determinism is unavoidable; expect numerical match to ~3 digits.
- The hypothesis registry (
registry/) records the evaluation criteria used by the analysis scripts.
No open-source license is currently granted.