Token-by-Token Reasoning Phases in LLMs (HMM over token features)

This repo implements the full pipeline described in the proposal:

1. sample prompts from benchmark datasets
2. run an open-source LLM token-by-token while collecting per-token signals
3. fit a Gaussian-emission HMM to the feature time series
4. build a "reasoning map" (latent states + transitions)
5. evaluate early forecasting + detour states

Install

python -m venv .venv
source .venv/bin/activate  # (Windows: .venv\Scripts\activate)
pip install -r requirements.txt

Quickstart (end-to-end)

1) Trace an LLM on a dataset

python scripts/run_tracing.py \
  --model mistralai/Mistral-7B-Instruct-v0.2 \
  --task gsm8k \
  --split test \
  --n 200 \
  --max-new-tokens 128 \
  --temperature 0.7 \
  --top-p 0.95 \
  --seed 0 \
  --out runs/mistral_gsm8k_t0p7

This creates runs/.../traces/*.npz plus a manifest.jsonl.

2) Train an HMM and pick K by BIC

python scripts/train_hmm.py \
  --run-dir runs/mistral_gsm8k_t0p7 \
  --k-min 3 --k-max 12 \
  --cov full \
  --n-init 5 \
  --max-iter 200

Outputs:

• runs/.../hmm/model.pkl
• runs/.../hmm/bic.json
• runs/.../hmm/states/*.npz (Viterbi + posteriors)

3) Evaluate early forecasting

python scripts/evaluate_forecasting.py \
  --run-dir runs/mistral_gsm8k_t0p7 \
  --ks 5 10 20 40

4) Identify detour states

python scripts/evaluate_detours.py --run-dir runs/mistral_gsm8k_t0p7

Notes

• The tracing loop is custom (not generate) so we can reliably collect: entropy, chosen-token logprob, hidden-state drift/cosine/norm/variance, and attention entropy/max (last layer, averaged over heads).

• "Hallucination" is implemented as a configurable placeholder: by default, it’s not is_correct for QA-style tasks. You can plug in retrieval-verified claim checking later via OutcomeEvaluator.

License

MIT