CritICL: Inference-Time Weak-to-Strong Generalization from Small Language Model Failure Modes

This is the official code for CritICL (COLM 2026). CritICL improves the reasoning of a strong LLM at inference time by reusing the failure modes of weaker models from the same family. It does not train anything, and CritICL-static needs only one generation from the target model.

Key observation. Models in the same family (e.g. Qwen2.5-1.5B/3B/7B vs. Qwen2.5-72B) show very similar failure-mode distributions. If you aggregate the failure statistics of several weak models, the result is even closer to the strong model's distribution. So critiques of weak-model mistakes are useful in-context guidance for the strong model.

CritICL has two stages:

Stage 1: CritBank construction (offline, once per model family)
  train questions ─► weak SLMs (5 CoT samples each) ─► keep incorrect responses
                  ─► frontier LLM (gpt-4o-mini): failure-mode labels L(q,r) + critique C(q,r)
                  ─► CritBank = {(q, r, L(q,r), C(q,r))}  and family profile P_M(l)

Stage 2: inference with the target LLM
  CritICL-dynamic: target model predicts <=5 likely failure modes for q'  ─┐
  CritICL-static : top-T failure modes of the family profile P_M(l)      ─┴─► Failure Mode-Based Sample
                   Selection (Alg. 1) ─► K critique examples + q' ─► target model ─► answer

This repository contains:

the method: CritBank construction, Failure Mode-Based Sample Selection (Algorithm 1), CritICL-dynamic (Algorithm 2), and CritICL-static (Algorithm 3);
dataset loaders and answer checking for GSM8K, MATH, AMC23, AIME24, AIME25 and GPQA;
every baseline and ablation from the paper, plus scripts for evaluation, inference-cost accounting, significance tests and failure-mode distribution analysis.

No model weights or checkpoints are needed. All models are used off the shelf.

Installation

git clone https://github.com/umwyf/CRITICL.git
cd CRITICL
conda create -n criticl python=3.10 -y && conda activate criticl
pip install -r requirements.txt     # or: pip install -e ".[vllm,retrieval,verify,dev]"

vllm is used to sample from the weak models offline and to serve the target model through an OpenAI-compatible API.
OPENAI_API_KEY must be set for the annotator (gpt-4o-mini) and for the LLM-as-a-judge baseline.
sentence-transformers is only needed for the semantic-retrieval ablations.
math-verify is optional. If it is installed, it adds symbolic answer equivalence on top of the built-in string and numeric matcher.

Repository structure

criticl/
├── taxonomy.py        # failure-mode taxonomies: coarse (8), fine (20, default), full (50; Appendix H)
├── prompts.py         # all prompt templates (Appendix G.1) + baseline prompts
├── llm.py             # LLM backends (OpenAI-compatible / vLLM offline / mock) + token accounting
├── data/
│   ├── datasets.py    # GSM8K, MATH, AMC23, AIME24/25, GPQA loaders; correctness function φ(q, r)
│   └── answer_utils.py# \boxed{} extraction, normalization, equivalence
├── construction.py    # Stage 1: response sampling, annotation, label clustering, CritBank assembly
├── critbank.py        # CritBank: L(q,r), C(q,r), inverse map L⁻¹(l), family profile P_M(l)
├── selection.py       # Algorithm 1 (Failure Mode-Based Sample Selection) + random / semantic selectors
├── methods/
│   ├── criticl.py     # CritICL-dynamic (Alg. 2), CritICL-static (Alg. 3), CritBank selection ablations
│   └── baselines.py   # zero/few-shot ICL, Consistency@k, Self-Reflection, LLM-as-Judge, dense retrieval
├── evaluation.py      # accuracy, cost, bootstrap CIs, paired bootstrap tests
└── analysis.py        # Spearman / Kendall τ / Top-10 overlap / JS distance, Figure-2 style plots
scripts/
├── prepare_data.py          # download benchmarks to local JSONL (for offline compute nodes)
├── generate_responses.py    # Stage 1a
├── annotate_failures.py     # Stage 1b
├── cluster_failure_modes.py # Stage 1c (optional, open-ended labels -> taxonomy)
├── build_critbank.py        # Stage 1d
├── run_inference.py         # Stage 2 (CritICL + all baselines)
├── evaluate.py              # result tables + significance tests
├── analyze_failure_modes.py # Section 4.1 / Table 4 / Figure 2
├── build_critbank.sh        # end-to-end Stage 1 for one model family
└── run_eval.sh              # end-to-end Stage 2 for one target model
tests/                       # unit + end-to-end tests with a mock LLM

Quick start (dry run, no GPU)

Every script accepts --backend mock, a deterministic fake LLM. You can use it to check the whole pipeline end to end without GPUs or API keys:

pytest -q                                   # unit tests + mock end-to-end pipeline

# Mock run of the full CLI pipeline on 20 GSM8K questions
python scripts/generate_responses.py --backend mock --model weak --datasets gsm8k --limit 20 --output /tmp/criticl/resp.jsonl
python scripts/annotate_failures.py --backend mock --inputs /tmp/criticl/resp.jsonl --output /tmp/criticl/ann.jsonl
python scripts/build_critbank.py --annotations /tmp/criticl/ann.jsonl --output /tmp/criticl/critbank.jsonl --profile_out /tmp/criticl/profile.json
python scripts/run_inference.py --backend mock --model target --method criticl_static --dataset gsm8k --limit 10 \
    --critbank /tmp/criticl/critbank.jsonl --profile /tmp/criticl/profile.json --output /tmp/criticl/preds/gsm8k/criticl_static.jsonl

Data

Name	Source (HF Hub)	Role
`gsm8k`	`openai/gsm8k` (`main`)	CritBank (train, 7.4k) + ID eval (test, 1.3k)
`math`	`EleutherAI/hendrycks_math` (7 subjects)	CritBank (train, 7.5k) + ID eval (test, 5k)
`amc23`	`math-ai/amc23`	OOD eval
`aime24`	`math-ai/aime24`	OOD eval
`aime25`	`math-ai/aime25`	OOD eval
`gpqa`, `gpqa_{chemistry,biology,physics,quantum}`	`Idavidrein/gpqa` (gated)	other-domain experiments (Appendix D.2 / E.3)

The scripts download datasets on the fly. On nodes without internet access, export them first:

python scripts/prepare_data.py --out_dir data                    # gsm8k/math/amc23/aime24/aime25
python scripts/prepare_data.py --datasets gpqa --out_dir data    # requires `huggingface-cli login`

Then pass --data_file data/<name>_<split>.jsonl (or --data_files for generate_responses.py). GPQA has only one split, so prepare_data.py holds out half of it (gpqa_test.jsonl, plus per-domain files) for evaluation and leaves the rest (gpqa_train.jsonl) for building a GPQA CritBank. You can use your own data as JSONL with the fields id, question, answer (and optionally solution).

Stage 1: Build CritBank

One command runs all of Stage 1 for a model family:

export OPENAI_API_KEY=...
FAMILY=qwen  bash scripts/build_critbank.sh     # weak: Qwen2.5-{1.5B,3B,7B}-Instruct
FAMILY=llama bash scripts/build_critbank.sh     # weak: Llama-3.2-{1B,3B}-Instruct, Llama-3.1-8B-Instruct
# LIMIT=50 FAMILY=qwen bash scripts/build_critbank.sh   # quick smoke test

The same steps, run one at a time:

# 1a. Sample 5 CoT responses per training question from each weak model; mark correctness φ(q, r).
python scripts/generate_responses.py --backend vllm --model Qwen/Qwen2.5-1.5B-Instruct \
    --datasets gsm8k math --split train --num_samples 5 \
    --output outputs/qwen/responses/Qwen2.5-1.5B-Instruct.jsonl
# (repeat for Qwen2.5-3B-Instruct and Qwen2.5-7B-Instruct)

# 1b. Up to 5 failure-mode labels + a critique for every incorrect response (gpt-4o-mini).
python scripts/annotate_failures.py --model gpt-4o-mini --taxonomy fine \
    --inputs outputs/qwen/responses/*.jsonl --output outputs/qwen/annotations.jsonl

# 1c. Assemble CritBank and the family-level failure-mode profile P_M(l).
python scripts/build_critbank.py --annotations outputs/qwen/annotations.jsonl \
    --output outputs/qwen/critbank.jsonl --profile_out outputs/qwen/profile.json

Deriving your own taxonomy (optional). The paper first collects noisy labels and then clusters them into representative failure modes (following Didolkar et al., 2024). To do the same:

python scripts/annotate_failures.py --open_labels --no_critique \
    --inputs outputs/qwen/responses/*.jsonl --output outputs/qwen/annotations_open.jsonl
python scripts/cluster_failure_modes.py --annotations outputs/qwen/annotations_open.jsonl \
    --num_clusters 20 --taxonomy_out outputs/qwen/taxonomy.json
# then annotate with the new taxonomy:  --taxonomy outputs/qwen/taxonomy.json

There are three built-in taxonomies (Table 6): coarse (8 groups, the list in the prompt templates), fine (20 groups, the default, the failure modes in Figure 2) and full (all 50 failure modes of Appendix H).

CritBank format (critbank.jsonl, one entry per incorrect response):

{"id": "math-train-algebra-12::Qwen/Qwen2.5-1.5B-Instruct::3",
 "question_id": "math-train-algebra-12", "question": "...", "response": "<incorrect CoT>",
 "failure_modes": ["incorrect_formula_application", "logical_step_skipping"],
 "critique": "<concise critique>", "model": "Qwen/Qwen2.5-1.5B-Instruct", "dataset": "math",
 "gold_answer": "...", "predicted_answer": "..."}

Stage 2: Inference with CritICL

Serve the target model with an OpenAI-compatible server and run either variant:

vllm serve Qwen/Qwen2.5-32B-Instruct --tensor-parallel-size 4 --port 8000 &

# CritICL-static: family profile -> top-T failure modes -> K critique examples (1 generation)
python scripts/run_inference.py --method criticl_static \
    --model Qwen/Qwen2.5-32B-Instruct --base_url http://localhost:8000/v1 \
    --dataset math --critbank outputs/qwen/critbank.jsonl --profile outputs/qwen/profile.json \
    --k 5 --profile_size 5 \
    --output outputs/qwen/preds/Qwen2.5-32B-Instruct/math/criticl_static.jsonl

# CritICL-dynamic: predict query-specific failure modes, then retrieve (2 generations)
python scripts/run_inference.py --method criticl_dynamic \
    --model Qwen/Qwen2.5-32B-Instruct --base_url http://localhost:8000/v1 \
    --dataset math --critbank outputs/qwen/critbank.jsonl --taxonomy fine --k 5 \
    --output outputs/qwen/preds/Qwen2.5-32B-Instruct/math/criticl_dynamic.jsonl

To run every method on every benchmark and print the result table:

FAMILY=qwen TARGET=Qwen/Qwen2.5-32B-Instruct bash scripts/run_eval.sh
FAMILY=llama TARGET=meta-llama/Llama-3.1-70B-Instruct bash scripts/run_eval.sh

Useful options:

Option	Meaning
`--k`	retrieval budget K (default 5)
`--profile_size`	T, the number of top failure modes used by CritICL-static (default 5)
`--weighted`	use w(l) = P_M(l) in Algorithm 1 instead of w(l) = 1
`--profile_models`	restrict the profile to some weak models (when `--profile` is not given)
`--critbank_models / --critbank_datasets`	restrict the retrieval pool
`--max_response_chars`	truncate long incorrect solutions in the prompt
`--allow_duplicate_questions`	allow several incorrect responses to the same question among the K examples
`--backend vllm`	run the target model offline with vLLM instead of a server

Each prediction record stores the response, the extracted answer, correctness, token usage, and the selected CritBank entries and failure modes. Runs resume automatically from an existing output file.

Baselines and ablations

Paper name	`--method`	Notes
Zero-shot	`zero_shot`
n-shot (Rand.)	`few_shot_random --shots n`	new random correct training exemplars for each query
n-shot (Fixed)	`few_shot_fixed --shots n`	one fixed set of exemplars
Consistency@k	`self_consistency --num_samples k`	T = 1.0, majority vote, 5-shot
Self-Reflection	`self_reflection --max_rounds 2`	solve → self-critique → refine, stops early when judged correct
LLM-as-Judge	`llm_judge --judge_model gpt-4o-mini`	5 candidates + 1 judge call
Random / Fixed / Semantic selection (Table 3)	`critbank_random` / `critbank_fixed` / `critbank_semantic`	same CritBank and prompt, different selection
Dense correct-exemplar retrieval (Table 7)	`dense_correct_retrieval`	K most similar correct training examples
Weak incorrect response only (Table 7)	`criticl_static --no_critique`
Shuffled failure labels (Table 7)	`criticl_static --shuffle_labels`
Taxonomy granularity (Table 6)	build CritBank with `--taxonomy coarse/fine/full`
Cross-family transfer (Table 10)	e.g. Qwen target + `--critbank outputs/llama/critbank.jsonl`
Cross-domain transfer (Table 11)	build a GPQA CritBank from `data/gpqa_train.jsonl`, then mix or swap CritBanks

If you run several shot counts, give each run its own output file (for example .../math/few_shot_fixed_3.jsonl), because evaluate.py uses the file name as the method name.

Evaluation, cost and significance

# Accuracy table (ID / OOD / overall averages) + average generations / input / output / total tokens
python scripts/evaluate.py --pred_dir outputs/qwen/preds/Qwen2.5-72B-Instruct

# 95% bootstrap CIs and paired bootstrap test against a baseline (Table 9)
python scripts/evaluate.py --pred_dir outputs/qwen/preds/Qwen2.5-72B-Instruct \
    --compare criticl_static self_consistency

Inference cost follows Table 2. Generations counts every model call needed for the final answer, including failure-mode prediction and judge calls. When one call samples k completions, its input tokens are counted once.

Failure-mode analysis

To reproduce Section 4.1 (Figure 2 and Table 4), sample and annotate the strong model on the same training questions, then compare the distributions:

python scripts/generate_responses.py --backend vllm --model Qwen/Qwen2.5-72B-Instruct --tensor_parallel_size 4 \
    --datasets gsm8k math --split train --output outputs/qwen/responses_strong/Qwen2.5-72B-Instruct.jsonl
python scripts/annotate_failures.py --no_critique --inputs outputs/qwen/responses_strong/*.jsonl \
    --output outputs/qwen/annotations_strong.jsonl
python scripts/analyze_failure_modes.py --weak outputs/qwen/annotations.jsonl \
    --strong outputs/qwen/annotations_strong.jsonl --plot outputs/qwen/failure_modes.png

This prints the Spearman correlation, Kendall's τ, Top-10 overlap and Jensen–Shannon distance for each weak model and for the aggregate weak profile.

Python API

from criticl import CritBank, failure_mode_sample_selection
from criticl.llm import OpenAIChatLLM
from criticl.methods import METHODS
from criticl.data import load_examples

bank = CritBank.load("outputs/qwen/critbank.jsonl")
print(bank.failure_mode_profile())                    # P_M(l)
print(len(bank.inverse("incorrect_formula_application")))   # |L^{-1}(l)|

examples = failure_mode_sample_selection(bank, ["logical_step_skipping", "incorrect_equation_setup"], k=5)

llm = OpenAIChatLLM("Qwen/Qwen2.5-72B-Instruct", base_url="http://localhost:8000/v1")
method = METHODS["criticl_static"](llm, critbank=bank, k=5, profile_size=5)
records = method.run(load_examples("aime25", "test"))

Implementation notes

The paper leaves some details open. This is how the code handles them:

Weak-model sampling. 5 CoT samples per (question, model). The paper does not give a sampling temperature; the default here is T = 0.7, top-p = 0.95 (--temperature). Target-model inference uses greedy decoding (T = 0), as in the paper.
Algorithm 1. Ties are broken randomly with a fixed seed. By default the K selected examples come from different questions (--allow_duplicate_questions turns this off). The test question is always excluded from retrieval.
CritICL-static is input-agnostic, so every query of a run gets the same K examples (selected once from the top-T profile modes).
CritICL-dynamic maps predicted labels onto the taxonomy with fuzzy matching. If no valid label is predicted, it falls back to the family profile.
Answer checking. The final answer is the last \boxed{}. It is compared after MATH-style normalization and numeric comparison, and also with math-verify if it is installed. GPQA uses option letters.

Citation

@inproceedings{wu2026criticl,
  title     = {CritICL: Inference-Time Weak-to-Strong Generalization from Small Language Model Failure Modes},
  author    = {Wu, Yufan and He, Yinghui and Hu, Zhengyi and Wei, Lang and Li, Ruichen and Yang, Qifan and Zhu, Ting},
  booktitle = {Conference on Language Modeling (COLM)},
  year      = {2026},
  url       = {https://arxiv.org/abs/2608.27455}
}

License

This project is released under the MIT License. Datasets and models keep their original licenses.

Navigation Menu

Repository navigation

About

Repository files navigation

CritICL: Inference-Time Weak-to-Strong Generalization from Small Language Model Failure Modes

Table of contents

Installation

Repository structure

Quick start (dry run, no GPU)

Data

Stage 1: Build CritBank

Stage 2: Inference with CritICL

Baselines and ablations

Evaluation, cost and significance

Failure-mode analysis

Python API

Implementation notes

Citation

License

About

Releases

Packages

Contributors

Languages

Name		Name	Last commit message	Last commit date
criticl		criticl
scripts		scripts
tests		tests
.gitignore		.gitignore
LICENSE		LICENSE
README.md		README.md
pyproject.toml		pyproject.toml
requirements.txt		requirements.txt

Latest commit

History

Folders and files

Repository files navigation

CritICL: Inference-Time Weak-to-Strong Generalization from Small Language Model Failure Modes

Table of contents

Installation

Repository structure

Quick start (dry run, no GPU)

Data

Stage 1: Build CritBank

Stage 2: Inference with CritICL

Baselines and ablations

Evaluation, cost and significance

Failure-mode analysis

Python API

Implementation notes

Citation

License

About

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages