Update README.md

README.md

@@ -10,4 +10,96 @@ base_model:
 - Qwen/Qwen2.5-VL-7B-Instruct
 pipeline_tag: image-text-to-text
 ---
-This is the model cited in the paper: [Think or Not? Selective Reasoning via Reinforcement Learning for Vision-Language Models](https://arxiv.org/abs/2505.16854).
+
+
+# TON-Math
+TON is a series of large language models trained using our efficient algorithm, which automatically decides whether to think or not, based on Qwen2.5-VL. 
+We apply Group Relative Policy Optimization (GRPO) for reinforcement learning with "thought dropout" supervised finetuning as a preliminary step.
+## Introduction
+
+Reinforcement Learning (RL) has proven to be an effective post-training strategy for enhancing reasoning in vision–language models (VLMs). Group Relative Policy Optimization (GRPO) is a recent prominent method that encourages models to generate complete reasoning traces before answering, leading to increased token usage and computational cost. Inspired by the human-like thinking process—where people skip reasoning for easy questions but think carefully when needed—we explore how to enable VLMs to first decide *when reasoning is necessary*. To realize this, we propose *TON*, a two-stage training strategy:
+
+1. **(i)** A supervised fine-tuning (SFT) stage with a simple yet effective “**thought dropout**” operation, where reasoning traces are randomly replaced with empty thoughts. This introduces a think-or-not format that serves as a cold start for selective reasoning.
+2. **(ii)** A GRPO stage that enables the model to freely explore when to think or not, while maximizing task-aware outcome rewards.
+
+Experimental results show that *TON* can *reduce the completion length by up to **90%** compared to vanilla GRPO, without sacrificing performance or even improving it*. Further evaluations across diverse vision-language tasks—covering a range of reasoning difficulties under both 3B and 7B models—consistently reveal that the *model progressively learns to bypass unnecessary reasoning steps as training advances*. These findings shed light on the path toward human-like reasoning patterns in reinforcement learning approaches.
+
+## Quickstart
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+example={
+  "image": "./Geo170K/images/test/0.png", ### your image path
+  "problem": "As shown in the figure, in triangle ABC, it is known that angle A = 80.0, angle B = 60.0, DE parallel  BC, then the size of angle CED is ()",
+        
+}
+
+def make_conversation_image(example):
+        return  {
+                'image': example['image'],  # Store path instead of loaded image
+                'prompt': [{
+                    'role': 'user',
+                    'content': [
+                        {'type': 'image', 'text': None},
+                        {'type': 'text', 'text': example['problem']}
+                    ]
+                }]
+        }
+
+model_name = "kolerk/TON-3B-AITZ"
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="auto",
+    device_map="auto"
+)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+
+text = tokenizer.apply_chat_template(
+    make_conversation_image(example),
+    tokenize=False,
+    add_generation_prompt=True
+)
+model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+generated_ids = model.generate(
+    **model_inputs,
+    max_new_tokens=4096,
+    top_p=0.95,
+    top_k=1,
+    temperature=0.6
+)
+generated_ids = [
+    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+]
+
+response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+print(response)
+```
+
+## Evaluation
+
+Run our test Python file in the [code repository](https://github.com/kokolerk/TON/blob/main/src/eval/test_qwen25vl_geoqa.py) for more details.
+
+
+## Citation
+
+If you find our work helpful, feel free to give us a cite.
+
+```
+@misc{wang2025think,
+    title={Think or Not? Selective Reasoning via Reinforcement Learning for Vision-Language Models},
+    author={Jiaqi Wang and Kevin Qinghong Lin and James Cheng and Mike Zheng Shou},
+    year={2025},
+    eprint={2505.16854},
+    archivePrefix={arXiv},
+    primaryClass={cs.AI}
+}
+```
+
+
+
+