Linear probes on LLM hidden states predict whether a prompt requires a tool call, with F1_macro ≈ 0.92 on schema-disjoint BFCL v4 splits and 14–22× lower wall-clock latency than full autoregressive generation.
A small linear classifier on a single layer's last-token hidden state answers the question "should this prompt invoke a tool?" before autoregressive decoding starts. On the same training data the probe matches LoRA fine-tuning; with 5× more training negatives LoRA wins by ~1.4 pp. Probe directions discovered independently in 5 models (Llama, Qwen-3B, Phi, Qwen-7B, Mistral-7B) correlate at Pearson r = 0.89–0.93 on a shared test set, and activation patching demonstrates a direction-specific causal effect on Qwen.
Models studied: Llama-3.2-3B, Qwen2.5-3B (native + sysprompt), Phi-3.5-mini, Qwen2.5-7B, Mistral-7B-v0.3.
| Llama-3.2-3B | Qwen2.5-3B (native) | |
|---|---|---|
| Probe macro F1 (5-fold CV) | 0.918 ± 0.009 | 0.939 ± 0.005 |
| Probe macro F1 (GroupKFold, schema-disjoint) | 0.898 ± 0.026 | 0.926 ± 0.025 |
| Generation macro F1 (zero-shot, BFCL) | 0.690 [0.540 worst-case parser bound] | 0.807 [0.733] |
| LoRA fine-tune (BFCL+Glaive) macro F1 | 0.940 | 0.952 |
| Probe vs FT — same training data, same n=522 test | 0.926 vs 0.885 (probe +4 pp) | — |
| Cross-architecture probe-direction correlation (alignment-free) | r = 0.89–0.93 across 18 cross-arch pairs | |
| Activation patching at α=20 (strict, vs random max) | 6.7% vs 2.7% (CIs overlap) | 12.7% vs 2.7% (CIs disjoint) |
| Latency (synthetic short prompt) | 14× | 22× |
| Latency (real BFCL prompts, medium bucket, n=10) | 34× | 70× |
See docs/results.md for the full canonical tables. The four plots below are produced by notebooks/results.ipynb; rerun it to refresh them after a new run.
Per-model test-F1 by layer (3-seed mean, last-token pooling). The
canonical headline numbers are computed at L26 (Llama) and L24 (Qwen),
hard-coded in scripts/probe_robustness.py
under DEFAULT_MODELS. The figure is a layer-search aid; per-model
val-F1 was not stored at the time these CSVs were produced, so do not
cite specific layer choices as "val-optimal" without re-running the
sweep with val tracking.
Bars are mean ± bootstrap 95% CI across 3 seeds × 3 layers. MLP (786K params on Llama, 256× the linear probe) gives no statistically significant gain on Qwen and only ~1 pp on Llama. The shuffled-label control's AUC sits at chance (≈0.49) and its macro-F1 lands at 0.50 ± 0.10 — selectivity (real − shuffled F1) is 0.486 ± 0.068 for Llama and 0.436 ± 0.078 for Qwen, with a 10000-permutation p ≈ 0.0025 in both cases (LOSO cross-validated layer choice).
20 pairs across 4 architecture families × 2 scales (3B Llama / Qwen / Phi + 7B Qwen / Mistral). Per-model probes trained independently on their own hidden states discover directions that correlate at Pearson r = 0.89–0.93 on a shared n=522 test set — this is the alignment-free finding and does not depend on Ridge regression. With a Ridge mapping between hidden spaces, transferred F1 retains 96.5–101.3% of target in-model F1; 5/20 pairs (notably those with Mistral-7B as target) exceed 100% retention, which is consistent with Ridge regularization on a noisy target rather than super-universal transfer. The retention number lacks a random-orthogonal-map null baseline, so it should be read as a qualitative ordering claim, not as a calibrated separation.
The Qwen-3B ↔ Qwen-7B pair retains 99.5% / 100.6%, but those models share a tokenizer — alignment for that pair is post-embedding only and is the easiest case in the matrix. Excluding the same-tokenizer pair, the 18 cross-architecture pairs retain 98.76% mean / 96.49% min.
F1 vs number of shared CCA components, per pair. CCA at d=128 retains 55–69% of in-model F1 across cross-arch pairs. The orthogonal-subspace constraint discards non-shared variance that Ridge captures, so the tool-use signal is not fully expressible as a low-rank shared subspace of size ≤ 128 — full transfer requires the unconstrained Ridge map.
# 1. Install uv (https://astral.sh/uv/)
curl -LsSf https://astral.sh/uv/install.sh | sh
# 2. Clone, create env, fill .env (HF_TOKEN required for gated models)
uv sync --extra dev
cp .env.example .env # edit to add HF_TOKEN
# 3. Smoke-test GPU + PyTorch
uv run python scripts/smoke_test.py
# 4. Walk through every experiment with y/n prompts
./run.shrun.sh is the recommended entry point for reproduction. It walks
through 27 steps from data download → hidden-state extraction (~3.5–9 min
per model) → probes → cross-model analysis → activation patching →
length-robustness ablation → optional LoRA fine-tune. Each step is
skipped automatically if its output already exists. Use ./run.sh --yes
for non-interactive end-to-end (long: hidden-state extraction × 6 +
cross-model 5-way + LoRA × 2 ≈ 12 hours on RTX 5080), ./run.sh --skip-lora to skip the LoRA pipeline, or ./run.sh --dry-run to
preview commands.
export PYTHONHASHSEED=0 # belt-and-braces for any subprocesssrc/training/determinism.py:set_deterministic_seed(42) is called at the
top of every script and turns on torch.use_deterministic_algorithms(True).
Per-experiment seeds are derived via seed_from(...) (md5-based) so
results are stable regardless of PYTHONHASHSEED. Caveat: PyTorch
deterministic mode does not eliminate cross-driver / cross-CUDA flutter
on small test sets — std on 5-fold CV typically lands at 0.005–0.013 on
F1_macro across compatible drivers.
See docs/reproduction.md for the full determinism contract, hardware requirements, and known gotchas.
- GPU with ≥ 16 GB VRAM (3B models fit; Mistral-7B / Qwen-7B need
--cpu-offloadon 16 GB cards, ~9 min extraction) - ~100 GB disk total (60 GB HF model caches, ~10 GB hidden-state HDF5, ~5 GB processed data)
- The original work was done on an NVIDIA RTX 5080 (Blackwell sm_120, 16 GB VRAM). Standard PyTorch wheels do not yet support sm_120 — a cu128+ build is required.
configs/ — experiment configs (yaml)
docs/ — methodology, results, reproduction guide
notebooks/ — results visualisation (4 plots)
results/ — pre-computed CSV/JSON outputs (16 MB, no hidden states)
exploratory/ — non-headline experiments; see results/exploratory/README.md
src/ — library code (data, models, training, evaluation)
data/splits.py — stratified + group-split helpers
models/probes.py — Linear / MLP / Attention probes
training/trainer.py — legacy npy-format trainer (HDF5 scripts have own loops)
training/determinism.py — seed_from() + set_deterministic_seed()
scripts/ — CLI entry points
tests/ — pytest (Wilson CI, splits, determinism, classifier regressions)
data/ — gitignored: raw datasets + extracted hidden states
models/ — gitignored: HF model caches + LoRA adapters
run.sh — interactive reproduction pipeline
- docs/methodology.md — experimental design, recipes, statistical methodology.
- docs/results.md — all canonical tables.
- docs/reproduction.md — env setup, seeds, hardware deepdive, known gotchas.
- notebooks/results.ipynb — 4 interactive visualisations: layer sweep, architecture comparison, cross-model transfer heatmap, CCA dimensionality ablation.
- results/exploratory/README.md — non-headline experiments and why they are not in the canonical results.
uv run pytest # all
uv run pytest tests/test_wilson_ci.py # Wilson CI smoke
uv run pytest tests/test_splits.py # stratified / group split
uv run pytest tests/test_determinism.py # seed_from contract
uv run pytest tests/test_classify_response.py # output classifier regressions
uv run pytest tests/test_probe_reproducibility.py # cached probe F1 floorThe probe-reproducibility test requires cached H5 + probe weights. It
skips on a fresh checkout and runs once run.sh has produced the
artefacts.
MIT — see LICENSE.