Uncertainty Signals in a Reasoning LLM.

A reproducible SLURM pipeline for studying how uncertainty signals change throughout SmolLM3-3B reasoning trajectories on MMLU.

AINLPSCHOOL
uncertainty signals in a reasoning llm

Experience

Motivation

Technical Pipeline

The project is organized as a reproducible research pipeline that runs on Mila’s SLURM cluster. It streams seeded MMLU subsets instead of downloading the full dataset, generates a complete reasoning chain for each question, and then probes the model at decile checkpoints to measure how its uncertainty changes over time.

The GPU experiment runs as a data-parallel SLURM job array. Each shard writes checkpoint and token-entropy results independently, after which a validation and merge step combines them into a single run. Separate CPU steps export the questions and generate the figures, tables, and findings used for analysis.

The repository also includes tests for the pure-logic helpers, saved run artifacts, design notes, and conventions for keeping model and dataset downloads on scratch storage rather than login nodes.

Learned

Built with
PythonPyTorchTransformersHugging Face DatasetsSLURMpytest