| --- |
| 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 |
