---
library_name: pika
license: mit
pipeline_tag: text-classification
tags:
- probe
- difficulty-prediction
- routing
- llm-routing
- interpretability
---

# PIKA Probes — Pre-trained Difficulty Prediction Probes

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](https://huggingface.co/papers/2602.09924).

These probes are part of the [PIKA (Probe-Informed K-Aware Routing)](https://github.com/KabakaWilliam/llms_know_difficulty) 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](https://arxiv.org/abs/2510.18147) (2025)

## License

MIT License