This repository contains pre-trained linear probes for predicting task difficulty from LLM activations, as presented in the paper LLMs Encode Their Failures: Predicting Success from Pre-Generation Activations.

These probes are part of the PIKA (Probe-Informed K-Aware Routing) project.

Authors: William Lugoloobi, Thomas Foster, William Bankes, and Chris Russell.

Overview

These probes predict whether a model will correctly solve a problem based on its internal activations before generation.

Use Cases

Routing: Direct easy/hard problems to appropriate models
Calibration: Estimate uncertainty before generation
Analysis: Understand model capabilities across problem types

Installation

bash

pip install pika
# or just
pip install huggingface_hub joblib

Quick Start

python

from huggingface_hub import snapshot_download
from pika.probe import LinearEoiProbe

# Download a probe
probe_path = snapshot_download(
    repo_id="CoffeeGitta/pika-probes",
    allow_patterns=["gpt-oss-20b-high--MATH--linear-eoi-probe--mv-correct--k5-t1.0/*"]
)

# Load and use
probe = LinearEoiProbe.load_from_checkpoint(probe_path)
predictions = probe.predict(texts=["Solve: Find all real x such that x^3 - 3x + 1 = 0"])

Using PIKA's built-in downloader

python

from pika.hub import download_probe

probe = download_probe(
    repo_id="CoffeeGitta/pika-probes",
    model_name="gpt-oss-20b_high",
    dataset="MATH"
)

Available Probes

Probe	Model	Layer	EOI Pos	Test AUC
`gpt-oss-20b-high--MATH--linear-eoi-probe--mv-correct--k5-t1.0`	gpt-oss-20b_high	15	-3	0.860
`gpt-oss-20b-low--MATH--linear-eoi-probe--mv-correct--k5-t1.0`	gpt-oss-20b_low	17	-3	0.845
`gpt-oss-20b-medium--MATH--linear-eoi-probe--mv-correct--k5-t1.0`	gpt-oss-20b_medium	17	-3	0.844
`Qwen2.5-Math-7B-Instruct--MATH--linear-eoi-probe--mv-correct--k5-t0.7`	Qwen2.5-Math-7B-Instruct	17	-1	0.845
`Qwen3-8B--gneubig_aime-1983-2024--linear-eoi-probe--mv-correct--k5-t0.6`	Qwen3-8B	33	-3	0.698

Probe Structure

Each probe folder contains:

best_probe.joblib: Trained sklearn LogisticRegression model
platt_scaler.joblib: Platt calibration scaler (for probability calibration)
probe_metadata.json: Full training metadata
config.json: Clean configuration summary

Activation Extraction

The probes are trained on hidden state activations extracted at a specific layer and token position:

best_layer_idx: Which transformer layer to extract activations from
eoi_token_offset: Token position relative to end of templated input
- -1 = last token of input
- -3 = 3rd token from end
- This is the position before any generation tokens

Example: For eoi_token_offset=-3, extract the hidden state at hidden_states[best_layer_idx][:, -3, :] after tokenizing the input prompt.

Model Details

GPT-OSS-20B Variants

The GPT-OSS-20B models are different training checkpoints (low/medium/high training compute). Probes trained on layer 15 activations at position -3 (3 tokens before end of input).

Qwen2.5-Math-7B-Instruct

Probe trained on layer 17 activations at position -1 (last token of input).

Training Your Own Probes

bash

python -m pika.main \
    --probe linear_eoi_probe \
    --dataset DigitalLearningGmbH_MATH-lighteval \
    --model Qwen/Qwen2.5-Math-7B-Instruct \
    --max_len 3000 --k 5 --temperature 0.7

Citation

If you use these probes in your research, please cite:

bibtex

@misc{lugoloobi_llms_2026,
    title = {{LLMs} {Encode} {Their} {Failures}: {Predicting} {Success} from {Pre}-{Generation} {Activations}},
    shorttitle = {{LLMs} {Encode} {Their} {Failures}},
    url = {http://arxiv.org/abs/2602.09924},
    doi = {10.48550/arXiv.2602.09924},
    publisher = {arXiv},
    author = {Lugoloobi, William and Foster, Thomas and Bankes, William and Russell, Chris},
    month = feb,
    year = {2026},
    note = {arXiv:2602.09924 [cs]},
}

See also our earlier work: LLMs Encode How Difficult Problems Are (2025)

License

MIT License