How angelazzz-611/reasoning_models_probing works
1. What this repo does
- This is an offline research pipeline for probing reasoning models’ hidden states to predict whether an intermediate reasoning step is correct.
- The system is not a serving stack; it is a data-generation, labeling, representation-extraction, training, and evaluation workflow.
2. Main pipeline stages
- The end-to-end flow is: generate model rollouts → segment rollouts into reasoning chunks → label chunk answers for correctness → extract hidden-state representations → train a probe classifier → evaluate on held-out data.
- Each stage produces a new artifact on disk, and later stages consume those artifacts rather than re-running earlier model inference.
3. Raw reasoning data layer
- initial_cot/ contains pre-generated chain-of-thought rollout files in JSONL format for multiple model families and datasets.
- These files are the starting point for the pipeline and encode the reasoning traces used to derive training examples for the probe.
4. Reasoning segmentation subsystem
- src/get_reasoning_chunks.py and the shell wrapper get_reasoning_chunks.sh break raw CoT traces into smaller reasoning chunks.
- The README indicates this step uses spaCy (en_core_web_sm) for text segmentation, and it supports multiprocessing because this can be CPU-heavy.
- The output is a merged segmented dataset that becomes the unit of labeling and representation extraction.
5. Intermediate answer labeling subsystem
- src/label_answer_correctness.py labels each chunk with whether its intermediate answer is correct.
- This stage uses the Gemini API as an external LLM service for answer extraction and correctness labeling, and can merge chunks that do not contain an answer into later chunks.
- The output is a labeled JSONL dataset that aligns reasoning chunks with correctness targets.
6. Hidden-state extraction subsystem
- src/get_representation.py runs the target reasoning model on the labeled chunks and saves hidden-state representations.
- The README shows the model is loaded from a local path, and extraction is batched and sharded into multiple files via file_id and file_size to manage storage and shuffle quality.
- This is the boundary between the base reasoning model and the downstream probe: the probe never sees text alone, only chunk representations.
7. Probe training subsystem
- src/train_predictor_with_class_weights.py trains the self-verification probe, which is a classifier predicting correctness from hidden states.
- The training code supports class-weight style imbalance handling via alpha_imbalance_penalty, plus standard hyperparameters such as learning rate, weight decay, threshold, and optional hidden size.
- train_probe.sh appears to orchestrate grid search across probe configurations, storing checkpoints and experiment metadata under grid_search/.
8. Probe evaluation subsystem
- src/test_predictor_with_class_weights.py evaluates trained probes on held-out embeddings.
- It can test a single checkpoint or select top-k probes from grid-search results using a validation metric such as best_val_acc.
- Test outputs are stored separately from training artifacts, keeping evaluation isolated from model selection.
9. Shared training/data utilities
- src/dataloader.py, src/probe_model.py, src/compute_metrics.py, and src/utils.py are the supporting runtime layer for loading embeddings/labels, defining the probe network, computing metrics, and sharing helper logic.
- These modules are the internal boundary between data format concerns and model-training concerns.
10. Math-specific grading support
- math_eval_tools/ is a bundled utility package for math answer normalization and grading.
- It likely supports correctness labeling and/or evaluation for math-style tasks where answers need canonicalization before comparison.
11. Orchestration and automation
- The repo relies heavily on shell scripts such as generate_reasoning.sh, get_reasoning_chunks.sh, label_answer.sh, train_probe.sh, and test_probe.sh.
- These scripts define the operational sequencing of the pipeline and make the system reproducible across datasets and model variants.
12. Key storage boundaries
- The pipeline is intentionally file-based: raw rollouts in initial_cot/, segmented/merged chunks in processed directories, labeled examples in labeled_cot/, model embeddings in model_embeds/, trained checkpoints in grid_search/, and evaluation outputs in test_result/.
- That separation is important for architecture diagrams because each directory marks a distinct artifact boundary and a distinct compute stage.
13. External dependencies that materially shape the architecture
- Hugging Face-style local model loading is implied for rollout generation and hidden-state extraction.
- spaCy is required for chunk segmentation.
- Gemini API is required for intermediate answer extraction/labeling.
- PyTorch is the likely training runtime for the probe classifier and checkpointing workflow.
14. Repo-level architectural takeaway
- The repository implements a staged offline ML pipeline, not a monolithic application.
- The central architectural idea is to convert reasoning traces into labeled hidden-state datasets, then learn a lightweight probe that estimates correctness from internal activations rather than from the final answer text alone.
Main entry points: test_predictor_with_class_weights.py Evaluate probe · train_predictor_with_class_weights.py Train probe · label_answer_correctness.py Label correctness · get_reasoning_chunks.py Chunk reasoning
Components 18
Offline pipeline
File-based stages that turn reasoning rollouts into probe-ready artifacts.
- Generate rollouts · Creates raw chain-of-thought rollouts from target reasoning models.
src/generate_reasoning.py - Chunk reasoning · Splits raw reasoning traces into smaller reasoning chunks using spaCy and multiprocessing.
src/get_reasoning_chunks.py - Label correctness · Extracts intermediate answers and assigns correctness labels with Gemini, optionally merging answerless chunks.
src/label_answer_correctness.py - Extract hidden states · Runs the reasoning model on labeled chunks and saves hidden-state embeddings in shards.
src/get_representation.py - Train probe · Trains the self-verification probe to predict correctness from hidden states.
src/train_predictor_with_class_weights.py - Evaluate probe · Scores trained probes on held-out embeddings and selects best checkpoints.
src/test_predictor_with_class_weights.py
Artifact stores
On-disk datasets and checkpoints that separate each stage of the workflow.
- Raw CoT · Pre-generated rollout JSONL files used as the starting point for the pipeline.
initial_cot - Segmented data · Merged chunk-level reasoning dataset produced from raw rollouts.
- Labeled chunks · Chunk-level JSONL dataset with correctness targets for probe learning.
- Model embeds · Stored hidden-state representations used as probe inputs.
- Grid search · Stores probe checkpoints and tuning metadata for training runs.
- Test results · Separated evaluation outputs for held-out probe performance.
- Math grading · Normalizes and grades math answers for correctness-focused tasks.
math_eval_tools
Training runtime
Shared code used by probe training and evaluation.
- Data loader · Loads embeddings, labels, and split artifacts for training and evaluation.
src/dataloader.py - Probe model · Defines the lightweight correctness predictor used by the probe pipeline.
src/probe_model.py - Metrics · Computes evaluation metrics and validation summaries.
src/compute_metrics.py - Shared utils · Common helper logic used across training and evaluation code.
src/utils.py
Orchestration
Shell entrypoints that sequence the offline workflow.
- Pipeline scripts · Shell wrappers that drive generation, chunking, labeling, training, and testing.
Connections 27
- Generate rollouts → Raw CoT: writes · Saves rollout JSONL artifacts for downstream processing.
- Raw CoT → Chunk reasoning: input · Raw reasoning traces are segmented into chunks.
- Chunk reasoning → Segmented data: produces · Creates merged segmented reasoning datasets.
- Segmented data → Label correctness: input · Chunked reasoning becomes the labeling target.
- Label correctness → Labeled chunks: produces · Outputs labeled intermediate-answer examples.
- Labeled chunks → Extract hidden states: input · Labeled chunks are fed to the reasoning model for embedding extraction.
- Extract hidden states → Model embeds: produces · Writes sharded hidden-state representations to disk.
- Model embeds → Train probe: input · Embeddings are the training features for the probe.
- Train probe → Grid search: writes · Stores checkpoints and grid-search artifacts.
- Grid search → Evaluate probe: selects · Evaluation can load top-ranked checkpoints from search results.
- Model embeds → Evaluate probe: input · Held-out embeddings are used for evaluation.
- Evaluate probe → Test results: writes · Saves final evaluation outputs separately from training artifacts.
- Data loader → Train probe: supports · Loads training splits and embedding datasets.
- Data loader → Evaluate probe: supports · Loads evaluation splits and checkpoints.
- Probe model → Train probe: implements · Defines the classifier architecture used during training.
- Probe model → Evaluate probe: implements · Shared model definition is reused for checkpoint loading.
- Metrics → Evaluate probe: supports · Computes accuracy and validation summaries.
- Shared utils → Train probe: supports · Shared training helpers and configuration logic.
- Shared utils → Evaluate probe: supports · Shared evaluation helpers and artifact handling.
- Math grading → Label correctness: supports · Canonicalizes and grades math answers for labeling.
- Math grading → Evaluate probe: supports · May support math-style evaluation normalization.
- Pipeline scripts → Generate rollouts: orchestrates · Shell entrypoints drive the staged workflow.
- Pipeline scripts → Chunk reasoning: orchestrates · Shell entrypoints drive the staged workflow.
- Pipeline scripts → Label correctness: orchestrates · Shell entrypoints drive the staged workflow.
- Pipeline scripts → Extract hidden states: orchestrates · Shell entrypoints drive the staged workflow.
- Pipeline scripts → Train probe: orchestrates · Shell entrypoints drive the staged workflow.
- Pipeline scripts → Evaluate probe: orchestrates · Shell entrypoints drive the staged workflow.