README.md

---
title: ThinkWhileThinking
emoji: 🧠
colorFrom: purple
colorTo: indigo
sdk: gradio
sdk_version: 5.33.0
app_file: app.py
pinned: true
license: mit
tags:
  - build-small-hackathon
  - thousand-token-wood
  - tiny-titan
  - track:backyard
  - sponsor:nvidia
  - sponsor:modal
  - achievement:offgrid
  - achievement:welltuned
  - achievement:offbrand
  - achievement:sharing
---

# 🧠 ThinkWhileThinking

> **Real-time reasoning failure detection in small language models via step-level process supervision**


* 🎥📝 **Demo Video & Technical Write-up:** [Watch the Demo & Read the Architecture on X](https://x.com/krishna__pahuja/status/2066275233028591925?s=46)


## What is this?

Most interpretability research is **post-hoc** — you find out *why* a model was wrong *after* it's already wrong. **ThinkWhileThinking** flips this:

A fine-tuned **Step Probe** model watches a reasoner's chain-of-thought in real time, and predicts — at each reasoning step — whether the logic is about to go wrong, **before the final answer is revealed.**

## How it works

```
User inputs a math / logic problem
        ↓
Reasoner (Qwen2.5-1.5B-Instruct) generates chain-of-thought step by step
        ↓
Step Probe (fine-tuned Qwen2.5-0.5B) scores each step in real time
        ↓
UI highlights exactly where reasoning starts to break down
        ↓
Reasoning Health score computed across all steps
```

## The Research Contribution

**Research Question:** *Can a lightweight probe model, trained only on step-level correctness annotations, predict mid-reasoning whether the final answer will be correct — before seeing the answer?*

Unlike prior work on CoT monitoring (which focuses on safety, deception, or adversarial attacks), TWT addresses **correctness failure prediction at inference time** using lightweight process supervision — making real-time interpretability accessible without frontier-scale models.

## Models Used

| Role | Model | Parameters |
|------|-------|-----------|
| Reasoner | Qwen/Qwen2.5-1.5B-Instruct | 1.5B |
| Step Probe | realArceus/twt-probe | 0.5B |
| **Total** | | **2B ≤ 32B ✅** |

## Training

The Step Probe (`realArceus/twt-probe`) is a fine-tuned `Qwen2.5-0.5B-Instruct` trained as a binary sequence classifier on **PRM800K** (trl-lib/prm800k) — OpenAI's process reward model dataset with step-level human correctness annotations.

| Detail | Value |
|--------|-------|
| Dataset | PRM800K (trl-lib/prm800k) |
| Samples | 40,000 step-level pairs |
| Task | Binary classification: correct vs. faulty step |
| Hardware | Modal A10G GPU |
| Epochs | 3 |
| **eval_accuracy** | **98.15%** |
| **eval_f1** | **99.05%** |

## How is this different from PRMs?

Process Reward Models (PRMs) assign quality scores to steps for **training signal**. TWT uses step-level supervision for **real-time inference-time interpretability** — a fundamentally different use case. We also demonstrate this works at 0.5B scale, far smaller than typical PRM 
## Built By

**Krishna Pahuja** — Final year AI student · HF Build Small Hackathon 2026

## License

MIT