ksopyla/cogito-probe-arith

CogitoProbe-Arith: nested arithmetic with mixed brackets

Synthetic nested + - * expressions with mixed brackets ()[]{}. Three question types: the final number (eval, an easy shortcut), internal-node values (subexpr, the real test), and which closer matches an opener (match). Use it to test whether a model stored the tree, not just a calculator.

Author: Krzysztof Sopyła · License: Apache-2.0 · Seed: 20260916 · Tokenizer: HuggingFaceTB/SmolLM3-3B

In 60 seconds

Several independent expressions are written up front, then filler, then a question about the first (farthest) expression.

Mixed ()[]{} do not change the numeric value. They colour the brackets so matching is a real question. Do not report eval accuracy as “the model understands arithmetic.”

Load it

from datasets import load_dataset

ds = load_dataset("ksopyla/cogito-probe-arith")
row = ds["validation"][0]

print(row["seq_len"], row["task"], row["variant"])
print("query: ", row["query"])
print("answer:", row["answer"])
print("prize bits:", row["prize_bits"], "gap:", row["gap"])

# Loss only on the answer span (already marked).
# input_ids / labels are lists of int, length == seq_len.
loss_tokens = [t for t in row["labels"] if t != -100]

# Start small on a laptop: 1,024-token rows, fixed fact count.
small = ds.filter(lambda r: r["seq_len"] == 1024 and r["variant"] == "fixed")

You can ignore input_ids and train from context / query / answer as text. If you do use the provided ids, they are already tokenized for HuggingFaceTB/SmolLM3-3B (Llama-3 vocab) and must not be re-tokenized.

The four CogitoProbe datasets

Length ladder (every family): 1024 → 4096 → 8192 → 16384 → 32768. Half the rows are fixed (same fact count as at 1k, longer haystack), half are scaled (more facts as the row grows).

Why these exist

Web text is locally predictable: a language model can look strong by guessing nearby words without remembering a fact from thousands of tokens earlier. These four datasets hide a known set of facts in a long padded haystack so you can measure whether a model (or a small latent memory) actually stored them.

Each row tells you how many bits the answer is worth (prize_bits) and how far the question sits from the last fact (gap). That is the whole point: the information content is labelled, the distractor text is not the prize, and the length is a ladder rather than a single context size.

Real rows use random single-token English-ish pieces from the Llama-3 / SmolLM3 vocabulary (gonzalez, oslo, validators, …), not the toy names alice / bob in the examples above. The grammar of the task is the same.

How to score

Train or evaluate only on the answer span. Teacher-forced token accuracy on labels != -100 is the main number. Recovered bits against the labelled prize:

max(0, prize_bits + Σ log2 p(gold_t))

A decoder that cannot see tokens more than gap away must sit at chance — the evidence is that far from the answer.

Score subexpr and match as the real tasks. Score eval only as a shortcut control (~one integer). If only eval moves, the model is a calculator, not a structure memory. Digit strings are space-separated (3 9 for 39, - 7 for −7).

Schema

This build

Per-rung means:

Example rows (truncated):

• —arith/seq1024/scaled/train/00000 task=eval prize=6.64 bits gap=1004 query=Q eval answer=3 9
• —arith/seq1024/scaled/train/00001 task=subexpr prize=15.27 bits gap=1005 query=Q sub 2 . 1 . answer=1 8 . 1 4
• —arith/seq1024/scaled/train/00002 task=subexpr prize=15.27 bits gap=988 query=Q sub 5 . 2 . answer=- 1 . - 1 2

Split leakage

Within-split duplicate input_ids counts: {'train': 0, 'validation': 0, 'test': 0}.

Train / validation / test use disjoint random streams. A fingerprint of the facts is checked for overlap. Shared answer strings (for example the same 8 colours) are expected and are not a leak.

Rebuild

Deterministic rebuild (does not upload):

uv run python scripts/build_concept_probe_datasets.py \
  --scale full --seed 20260916 \
  --tokenizer HuggingFaceTB/SmolLM3-3B \
  --families arith \
  --out_dir Cache/concept_probes/full

Ids are composed from a verified 1-token atom table of HuggingFaceTB/SmolLM3-3B. Arithmetic rows inject bare digit and bracket ids; they do not BPE-encode a glued string such as (1+2)*[3-4] (that merge path is not a well-defined alphabet).

Limitations

• —Not natural language. Atoms are verified 1-token pieces of the SmolLM3 / Llama-3 vocab, chosen so each symbol is one id. Do not treat this as a human corpus.
• —Answers are packed (several values in one span). Single-token labels are too sparse for a small latent channel to learn from.
• —prize_bits is a counting lower bound on the answer, not a cross-entropy floor of a local language-model window.
• —Arithmetic mixed brackets colour the tree; they do not change + - * meaning. eval-only accuracy is not evidence of rich structure.
• —Rows are padded with a repeating filler cycle, so gzip of the full text looks tiny. Compare prize_bits, not compressibility of the padded row.

Origin

These files were built for a research project on compressing long context into a small set of latent vectors (“concepts”), so the author could ask what those vectors actually store. You do not need that project, its training code, or its internal experiment log to use the datasets.

Project page: ai.ksopyla.com · author: Krzysztof Sopyła. Generator: data/concept_probes/ in the public research repo (MIT).

License

Apache-2.0 for this synthetic dataset. No web scrapes, no personal data. Generator code is MIT.

Citation

@misc{cogitoprobe2026,
  title  = {CogitoProbe: synthetic long-haystack probes for memory and compression},
  author = {Sopyła, Krzysztof},
  year   = {2026},
  url    = {https://huggingface.co/datasets/ksopyla/cogito-probe-arith},
  note   = {Seed 20260916. Four families: bits, bind, arith, props.},
}