|
|
--- |
|
|
license: apache-2.0 |
|
|
datasets: |
|
|
- kolerk/TON-Math-SFT |
|
|
language: |
|
|
- en |
|
|
metrics: |
|
|
- accuracy |
|
|
base_model: |
|
|
- Qwen/Qwen2.5-VL-7B-Instruct |
|
|
pipeline_tag: image-text-to-text |
|
|
--- |
|
|
|
|
|
|
|
|
# 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} |
|
|
} |
|
|
``` |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
