loose-bits/uq-hiddenstates

uq-hiddenstates — residual-stream states of reasoning traces at a fixed depth

Every token position of a reasoning trace, recorded at relative model depth 0.75, with correctness labels. Built for studying whether uncertainty is legible in the residual stream while the model reasons, rather than only at the answer.

Layout

Per model and dataset:

• —{model}_{ds}_L{idx}.part{k}.npy — fp16 [rows, hidden], traces concatenated, raw states, not normalized. Sharded at ~20 GB on item boundaries, so no trace ever spans two files.
• —{model}_{ds}_L{idx}.idx.json — the map.

import numpy as np, json
idx = json.load(open("gemma_omnimath_L45.idx.json"))
it  = idx["items"][0]
X   = np.load(idx["shards"][it["shard"]]["name"], mmap_mode="r")
gen = np.asarray(X[it["off"] + it["prompt_len"] : it["off"] + it["n"]], np.float32)

Per item: id (frozen item id), shard, off (row offset within that shard), off_global (offset in the notional concatenation), n (total tokens), prompt_len (generation starts at off + prompt_len), pre_tok and eot_tok (pre-answer and end-of-think positions, relative to generation start), and y — 1 = incorrect, 0 = correct, −1 = unscored. Items truncated at the token budget are excluded.

Models and depth

The index is the output of decoder block idx−1. For a Jacobian lens keyed by block, that is key idx−1.

Datasets

mmlupro, logiqa2, supergpqa, omnimath, lcb (LiveCodeBench), triviaqa — N≈1000 each, seeded stratified samples, one reasoning trace per item.

Two things to know before using it

Generation length predicts error on its own — within-dataset AUROC 0.58–0.85, mean 0.716 on this suite. Any score that aggregates over token positions inherits that. Report length-matched numbers (AUROC within deciles of generation length) alongside raw ones.

The token-direction manifold is anisotropic. Random token axes are far from chance here, so any vocabulary-derived direction needs a dimension-matched random control, and contrasts should be bipolar — the shared component is large and carries no signal.