README.md

---
library_name: transformers
license: apache-2.0
datasets:
- AI-MO/NuminaMath-CoT
base_model:
- Qwen/Qwen2.5-Math-1.5B
pipeline_tag: text-generation
tags:
- reinforcement-learning
- grpo
- hierarchical
- reasoning
- math
- tree-based
model_name: TreeRPO-Qwen2.5-Math-1.5B
---

# TreeRPO-Qwen2.5-Math-1.5B

**Summary:**  
A 1.5B parameter math reasoning model fine-tuned with **TreeRPO**, a hierarchical extension of GRPO that assigns rewards to “thought” nodes (not just full completions). Achieves higher GSM8K accuracy with just ~10K supervised + RL examples and **no reward model**.

🔎 **Full write-up (method, math, analysis):**  
[TreeRPO: Hierarchical Credit Assignment for Reasoning in Language Models](https://omrisapir.substack.com/publish/post/167273414)

---

## Model Details
- **Base model:** [`Qwen/Qwen2.5-Math-1.5B`](https://huggingface.co/Qwen/Qwen2.5-Math-1.5B)
- **Method:** TreeRPO (tree-structured GRPO;)
- **Reward signal:** Deterministic exact-match checker (binary). Interior node rewards = mean descendant leaf rewards.
- **Domain:** Grade-school and intermediate math word problems (GSM8K style)

## Intended Use
Research on hierarchical RL for reasoning; math tutoring (with human oversight); or as a research baseline for deterministic pass/fail domains (potential to extend to code with unit tests).

**Not intended for:**  
Open-ended or unsafe dialog, general factual QA, or high-stakes applications.

---

## Evaluation (GSM8K Test Set, 1,319 problems)

| Model                          | Greedy (%) | Maj@8 (%) | Notes                                |
|---------------------------------|------------|-----------|--------------------------------------|
| Qwen2.5-Math-1.5B-Instruct      | 84.8       | 89.5      | Reported settings                    |
| **Qwen2.5-Math-1.5B-TreeRPO**   | **86.4**   | **89.6**  | Same decoding (temp 0 / (0.7, 0.8))  |

- **Greedy:** temperature = 0 (deterministic)
- **Maj@8:** 8 completions (temperature 0.7, top-p 0.8); majority vote on final boxed answer

---

## How to Use

If your Transformers version supports chat templates (≥4.38), use:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "omrisap/TreeRPO-Qwen2.5-Math-1.5B"
tok = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")

messages = [
    {"role": "system", "content": "You are a helpful math reasoning assistant. Provide step-by-step reasoning and put the final answer in \\boxed{}."},
    {"role": "user",   "content": "If 3x + 5 = 17, what is x?"}
]

prompt_text = tok.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)

inputs = tok(prompt_text, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=256, temperature=0.0)
print(tok.decode(outputs[0], skip_special_tokens=True))