CoffeeGitta/pika-probes

PIKA Probes — Pre-trained Difficulty Prediction Probes

This repository contains pre-trained linear probes for predicting task difficulty from LLM activations. These probes are part of the PIKA (Probe-Informed K-Aware Routing) project.

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

pip install pika
# or just
pip install huggingface_hub joblib

Quick Start

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

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

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:

@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