This repository implements the methods behind the COLM 2026 paper "Reasoning Fine-Tuning Induces Persistent Latent Policy States".

TL;DR: given a reasoning model and its sentence-level hidden-state trajectories, this code discovers discrete latent policy states, measures how persistent they are, and tests whether those states can be transplanted onto a base model to steer reasoning behavior.

• src/contrastive_gen/ — reusable CEBRA-based encoder, dynamics model, training, and API helpers
• analysis/ — scripts for dataset extraction, regime discovery, causal steering, and transplantation analysis
• experiments/ — structured paper experiments, ablations, and evaluation utilities
• generate_data/ — dataset creation and feature extraction helpers
• scripts/ — Hugging Face utilities and repository management helpers
• prefixguard/ — scores partial reasoning traces with the learned SDS/CEBRA states

All commands in this README are run from the repository root.

git clone https://github.com/withmartian/mi-cot.git
cd mi-cot
python3 -m venv .venv
source .venv/bin/activate
python3 -m pip install --upgrade pip
python3 -m pip install -r requirements.txt
python3 -m pip install -e .

Minimum recommended Python version: >= 3.10.

1. contrastive_gen.load_features() loads hidden-state trajectories and builds training triplets.
2. contrastive_gen.standardize_hidden_states() scales activations into a model-ready tensor.
3. contrastive_gen.train_k_regime_model() trains a discrete latent regime encoder plus dynamics model.
4. The learned regime assignments are evaluated for persistence, transition structure, and causal transplant effects.

This is a paper repository, not a model hub: it provides the code and analysis pipeline used to reproduce the reasoning policy controller experiments.

This repo releases code only.

The main runnable components are:

• contrastive_gen/ — API for training and saving regime models
• analysis/ — dataset assembly, feature extraction, steering, and evaluation scripts
• experiments/ — paper experiments and ablations
• scripts/ — utility scripts for repository management and Hugging Face interaction

from contrastive_gen import (
    load_features,
    standardize_hidden_states,
    train_k_regime_model,
    save_model_artifacts,
)

from contrastive_gen import load_features, standardize_hidden_states, train_k_regime_model, save_model_artifacts

features, triplets = load_features("rpc_dataset/all_sentences_features.pkl")
X_torch, scaler = standardize_hidden_states(features)
result = train_k_regime_model(K=4, features=features, triplets=triplets, epochs=50)

print("Persistence:", result["persistence"])
print("Sequence count:", len(result["sequences"]))

save_model_artifacts("out/cebra_model", result)

Use analysis/gen_rpc_base.py to extract hidden-state features from a base model.

Use contrastive_gen.train_k_regime_model() to learn discrete latent states from sentence activations.

Recover state sequences and persistence scores with:

from contrastive_gen import build_state_sequences, compute_persistence

sequences = build_state_sequences(features, result["states"])
persistence = compute_persistence(sequences)
print(persistence)

Use scripts in analysis/ to run transplant and steering studies, including:

• analysis/cebra_EM.py
• analysis/rpc.py
• analysis/cebra_em_steering_inputDep.py

• experiments/ contains structured paper experiments and ablations
• generate_data/ contains dataset creation utilities
• analysis/ contains example pipelines for feature extraction and steering

The code in prefixguard/ asks a simple question: can the policy states tell us that a reasoning trace is going wrong before the model finishes? It fits state scores from labeled traces, uses them to rank several candidate answers, and repeats the evaluation after seeing 25%, 50%, 75%, and 100% of each trace.

The folder contains the reranking and analysis scripts along with small result summaries. Model weights, activation pickles, and raw generations are left out because they are too large for the repository. Setup and usage are in the prefixguard README.

The core API is exposed through contrastive_gen:

• load_features(path, ...)
• standardize_hidden_states(features)
• train_k_regime_model(K, features, triplets, ...)
• build_state_sequences(features, states)
• compute_persistence(sequences)
• save_model_artifacts(path, artifacts)

If you find this code useful, please cite:

Harrasse, A., et al. (2026). Reasoning Fine-Tuning Induces Persistent Latent Policy States.