concept-probe

concept-probe is a small, config-driven Python library for training, evaluating, scoring, and steering concept probes in open-weights language models.

What it includes

• Contrastive probe training from positive vs negative system prompts
• Probe training from labelled example texts (no system prompts required)
• Layer sweeps with effect sizes and optional p-values
• Token-level and sequence-level scoring
• Generation-time steering with learned concept vectors
• Multi-probe scoring on the same model
• Saved artifacts for configs, metrics, tensors, logs, plots, and HTML heatmaps
• Behavioral eval helpers and multi-probe reporting utilities
• Optional generation-logit capture

What it does not include

• Paper datasets
• Paper figures
• Paper analysis outputs
• Repo-level experiment history

Installation

Install from GitHub:

pip install git+https://github.com/mneuronico/concept-probe.git

Local editable install:

pip install -e .

Minimal runtime dependencies:

pip install numpy torch transformers

Optional extras:

pip install -e .[plots,stats]
pip install -e .[coherence]
pip install -e .[full]
pip install -e .[dev,stats,plots]

Notes:

• The library expects a local PyTorch environment and a Hugging Face-compatible causal LM.
• scipy is used for sweep p-values.
• matplotlib is used for plots.
• statsmodels and scikit-learn are used by reporting utilities.
• groq is only needed for coherence rating helpers.
• bitsandbytes is only needed for 4-bit loading.

Quickstart

from concept_probe import ConceptSpec, ProbeWorkspace

SAD_SYSTEM = "You are a helpful assistant with a subdued, melancholic tone."
HAPPY_SYSTEM = "You are a helpful assistant with a warm, optimistic tone."

concept = ConceptSpec(
    name="sad_vs_happy",
    pos_label="sad",
    neg_label="happy",
    pos_system=SAD_SYSTEM,
    neg_system=HAPPY_SYSTEM,
    eval_pos_texts=[
        "I feel heavy and distant today.",
        "It is hard to find hope right now.",
    ],
    eval_neg_texts=[
        "I feel light and excited about today.",
        "Everything feels uplifting and positive.",
    ],
)

workspace = ProbeWorkspace(model_id="meta-llama/Llama-3.2-3B-Instruct")
probe = workspace.train_concept(concept)

probe.score_prompts(
    prompts=["Write a short paragraph about the ocean."],
    alphas=[0.0, 6.0, -6.0],
    alpha_unit="sigma",
)

Outputs are written under outputs/<concept_name>/<timestamp>/.

Training from labelled texts

If you already have labelled example sentences for each pole, set training.train_prompt_mode = "texts" and provide them on the concept itself via train_pos_texts / train_neg_texts. No pos_system / neg_system is needed; pooling falls back to readout.read_mode (default sequence_all_mean), the same path the read-side eval uses.

from concept_probe import ConceptSpec, ProbeWorkspace

concept = ConceptSpec(
    name="raya_animal_vs_geom",
    pos_label="animal",
    neg_label="geometry",
    pos_system="",
    neg_system="",
    train_pos_texts=[
        "La raya nadaba sobre el lecho arenoso.",
        # ... 20+ examples
    ],
    train_neg_texts=[
        "Trazó una raya recta con la regla.",
        # ... 20+ examples
    ],
    eval_pos_texts=[...],   # held-out animal examples
    eval_neg_texts=[...],   # held-out geometry examples
)

workspace = ProbeWorkspace(
    model_id="meta-llama/Llama-3.2-3B-Instruct",
    config_overrides={"training": {"train_prompt_mode": "texts"}},
)
probe = workspace.train_concept(concept)

By default each labelled text is read as a user message after the neutral system prompt, and sequence_* readouts pool over the whole chat-formatted sequence. To read each text as the model's own reply and pool over the text's tokens only, set training.train_text_role = "assistant" with an assistant_* read mode. The optional training.train_text_prompt is a user turn placed before every text. Eval texts are read the same way, so the layer sweep sees the same readout as training. To score other texts with the trained probe under the same readout, use extract_text_reps with the same role, prompt and system: it encodes texts exactly as training does, and can return several readouts from one forward pass.

training.train_text_role = "raw" reads each text alone, with no chat template: the BOS token and the text, with the assistant_* readouts skipping BOS. model.chat_template (a Jinja string) replaces the tokenizer's chat template for everything the bundle does. Use it, for instance, to drop the dated system block that Llama 3.1 inserts when there is no system message.

config_overrides = {
    "training": {
        "train_prompt_mode": "texts",
        "train_text_role": "assistant",
        "train_text_prompt": "What do you think of the painting?",
    },
    "readout": {"read_mode": "assistant_all_mean"},  # or assistant_last / assistant_last_k_mean
}
reps = extract_text_reps(model, tokenizer, new_texts, role="assistant",
                         prompt="What do you think of the painting?", system=neutral_system,
                         pool=["assistant_all_mean", "assistant_last"])

Variance-normalized directions

training.direction_normalization chooses how the per-layer direction is estimated from the pole representations:

• "none" (default): the difference of means, mu_pos - mu_neg.
• "var": the difference of means divided coordinate-wise by the pooled within-pole variance, 0.5 * (var_pos + var_neg). This is the diagonal form of the covariance correction of mass-mean probing. It keeps a few very high-variance coordinates from dominating the direction, as with the massive activations of Gemma 3.
• "sd": the same, divided by the pooled standard deviation.

The direction is always normalized to unit length and stored in tensors.npz. Every later use of the probe (scoring, steering, ConceptProbe.load) reads the stored vector, so a probe trained with "var" stays variance-normalized wherever it is used. The layer sweep, and therefore best_layer, is computed with the normalized directions. concept_direction(pos, neg, normalization) exposes the same estimator for representations computed elsewhere.

config_overrides = {"training": {"direction_normalization": "var", "best_layer_search": [0.2, 0.8]}}

JSON-driven workflow

You can train from a JSON concept spec:

from concept_probe import train_concept_from_json

probe = train_concept_from_json("examples/concepts/bored_vs_interested.json")

The example configs in examples/concepts/ are intentionally general and are not tied to the paper repo.

Repo layout

concept-probe/
  concept_probe/          # library package
  examples/               # general examples only
  pyproject.toml
  README.md
  LICENSE

Public API

Main imports:

from concept_probe import (
    ConceptSpec,
    ProbeWorkspace,
    ConceptProbe,
    concept_direction,
    extract_completion_reps,
    extract_text_reps,
    multi_probe_score_prompts,
    train_concept_from_json,
    generate_multi_probe_report,
    run_scored_eval,
    run_multi_scored_eval,
)

Development

Run the lightweight test suite:

python -m pytest -q

Compile the example entrypoints:

python -m py_compile examples/quickstart.py examples/train_from_json.py

GitHub Actions CI runs the same lightweight checks and also builds source and wheel distributions.

Release notes

• New, opt-in, defaults unchanged: training.train_text_role = "raw" (the text alone, BOS included, no chat template), model.chat_template to replace the tokenizer's template, and extract_text_reps to score new texts exactly as texts-mode training read its own.
• New, opt-in, defaults unchanged: training.direction_normalization ("none", "var", "sd") for variance-normalized concept directions, and training.train_text_role = "assistant" (with training.train_text_prompt) to read labelled texts as the model's reply and pool over their tokens only. With the defaults, training produces exactly the same vectors as before.
• apply_chat now normalizes apply_chat_template output to a (1, seq_len) LongTensor. Recent transformers versions can return a BatchEncoding/dict (or a 1D tensor) instead of a bare 2D tensor; previously this raised AttributeError on .shape during generation and read paths. All generation/read now works regardless of the tokenizer's return shape.
• Steered generations are now scored with the steering hooks removed. Previously the scoring forward pass ran with the hooks still active, so the injected vector "read itself back" and inflated probe scores under steering by roughly alpha * cos(steer_dir, score_dir) at the steered layers — a large effect when steering and scoring with the same probe (and ~zero for random-vector controls, which therefore could not detect it). Probe-score-vs-alpha curves from earlier versions should be regenerated by re-scoring the saved completions. Saved generation logits still reflect the steered model and are labeled generation_logits_source = "steered_model".
• best_d/best_p in metrics.json are post-selection values (the best layer is chosen by optimizing over the layer sweep on the same eval data); confirm effect sizes on held-out data before reporting them.
• About alpha_unit="sigma": alpha is multiplied by the std of the sequence-pooled eval projections at the best layer, and the resulting raw offset is added at every token position of each steered layer (split across layers when steer_distribute=true, compounding through subsequent layers). It is a convenient scale knob, not literally "n standard deviations of behavior".
• score_prompts(..., seed=...) / multi_probe_score_prompts(..., seed=...) seed the RNGs for reproducible sampling (best-effort: GPU kernels are not bit-deterministic across hardware/library versions).
• Coherence rating sends completions to the external Groq API (it runs by default in run_scored_eval when a GROQ_API_KEY is configured; a notice is printed before sending). Pass rate_coherence=False for data that must not leave the machine.
• The default config in concept_probe/defaults.json is usable, but opinionated.
• Behavioral eval helpers depend on optional packages and, for coherence rating, an external API key.
• Example scripts are intended as starting points and will download model weights when you run them.