Qwen2.5-3B-UFO
This model is based on Qwen2.5-3B-Instruct and trained with PPO (Proximal Policy Optimization) on the MetaMathQA dataset for mathematical reasoning, as presented in the paper A Simple "Try Again" Can Elicit Multi-Turn LLM Reasoning.
Github: https://github.com/lichengliu03/unary-feedback Website: https://unary-feedback.github.io/
Overview
"Let's Try Again" addresses a critical gap in language model training: while single-turn reinforcement learning (RL) improves reasoning, these models fail in multi-turn interactive scenarios, often repeating the same wrong answers despite feedback.
Key Problem
Single-turn RL models lose the ability to revise reasoning across multiple turns. In 70% of failure cases, they produce identical answers across 5 interaction rounds, unable to incorporate simple feedback like "try again."
Solution: UFO Framework
Unary Feedback as Observation (UFO) transforms static datasets into multi-turn training by:
- Using only minimal feedback signals ("Try Again")
- Treating failure feedback as part of the observation
- Enabling models to learn from historical mistakes
Results
- 14% improvement in multi-turn success rates
- 10% reduction in average interaction turns
- Better performance even in single-turn scenarios
- 90% non-repetitive answers (vs 80% baseline)
Impact
UFO enables effective multi-turn RL training on existing static datasets without expensive annotations, making it practical to train models that can learn from sparse feedback and improve iteratively through trial-and-error, just like humans do.
Model Info
- Base model: Qwen/Qwen2.5-3B-Instruct
- Training method: PPO (full-parameter fine-tuning, not LoRA)
- Training data: MATH_MetaMathQA
- Training steps: 200 steps
- Framework: VERL
- Tensor parallel: 2x GPU distributed training
- Model size: ~6GB
Training Config
- Micro Batch Size: 1 per GPU
- PPO Mini Batch Size: 8
- Actor Learning Rate: auto
- Critic Learning Rate: auto
- KL Penalty: 0.001
- Clip Ratio: 0.2-0.28
- Temperature: 1.0 (train), 0.5 (eval)
UFO Framework Details
The UFO framework transforms static single-turn datasets into multi-turn interactive training through a simple yet effective approach.
The UFO framework flow: Static datasets are transformed into multi-turn episodes where models receive minimal feedback ("Try Again") and learn to revise their reasoning across multiple attempts.
Problem Formulation
We model multi-turn problem solving as a finite-horizon Markov Decision Process (MDP) where:
- State: Encodes the original question and history of past attempts with feedback
- Action: All possible answers the model can generate
- Reward: Binary signal (1 for correct, 0 for incorrect)
- Transition: Agent generates answer, receives feedback, episode continues until success or max turns
Unary Feedback as Observation (UFO)
The core innovation is treating minimal feedback as part of the observation:
Question: What is the value of x + y?
Attempt 1: [wrong answer]
Feedback: Try Again.
Attempt 2: [correct answer]
Key Features:
- Only negative feedback (e.g., "Try Again") is included in context
- No positive confirmation signals are ever shown
- Model must learn to revise based solely on failure history
- Episodes terminate immediately upon correct answer
Training with PPO
We use Proximal Policy Optimization (PPO) to train the policy:
- Agent observes input with full interaction history
- Generates answer and receives binary reward
- Policy updates using clipped surrogate objective
- Value function provides advantage estimates for stable training
Reward Design
Two complementary strategies encourage efficient reasoning:
1. Exponential Reward Decay:
DecayReward(t) = ฮณ^t if correct, 0 otherwise
Favors solving problems in fewer turns.
2. Repetition Penalty:
Penalty(ฯ) = ฮป ยท (1 - E(ฯ)/T)
Penalizes duplicate answers, encouraging diverse reasoning strategies.
This framework enables effective multi-turn RL training on static datasets without requiring expensive annotations or complex environments.
Usage
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("LichengLiu03/qwen2.5-3b-ppo-metamath-full")
model = AutoModelForCausalLM.from_pretrained(
"LichengLiu03/qwen2.5-3b-ppo-metamath-full",
torch_dtype=torch.float16,
device_map="auto"
)
# Example math problem
prompt = "Solve this math problem: If a circle has a radius of 5cm, what is its area?"
inputs = tokenizer(prompt, return_tensors="pt")
# Generate answer
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=512,
temperature=0.7,
do_sample=True,
pad_token_id=tokenizer.eos_token_id
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)
Features
This model is optimized for mathematical reasoning with PPO, and compared to the base model, it improves:
- โ Math problem understanding
- โ Logical reasoning accuracy
- โ Clarity of solution steps
- โ Calculation accuracy
Technical Details
- Tensor parallel training: 2 GPUs, distributed
- Memory optimization: gradient checkpointing and mixed precision
- Reward modeling: based on MetaMathQA correctness and reasoning quality
- Policy optimization: PPO for stable training
Limitations
- Mainly optimized for mathematical reasoning
- May not perform as well on general tasks
- Recommended for math, logic, and reasoning tasks
Key Results
Multi-Turn Reasoning Performance
We compare our multi-turn UFO model against a strong single-turn PPO baseline. For a fair comparison, the baseline is evaluated on 5 independent samples (Pass@5), while our model uses 5 sequential attempts with feedback (Succ@5). Success is recorded if any of the 5 responses is correct. We also analyze the impact of varying the maximum number of interaction turns at training.
Left: Multi-turn (5-turn) RL significantly outperforms single-turn baseline. Right: Performance comparison with different training turns (1, 5, and 10).
Key Findings:
- +14% success rate over single-turn PPO baseline
- Benefits generalize to both multi-turn and single-turn inference
- Best results with 5-turn training; more turns yield diminishing returns
Effectiveness of Unary Feedback
To further investigate the role of unary feedback, we compare model performance under different feedback availability conditions. In scenario (a), unary feedback is provided during both training and validation phases, while in scenario (b), unary feedback is available only during training but not at validation. The results show that access to unary feedback during both phases substantially improves validation success rate. In contrast, providing unary feedback solely during training does not yield improvements, indicating that the benefit of unary feedback is contingent on its availability at inference time.
Success rate comparison under different unary feedback settings: (a) feedback in both training and validation; (b) feedback only in training.
Key Insights:
- Feedback in both training and validation is crucial for improvement
- Feedback only in training phase does not help at inference
Reward Design Impact
Exponential Reward Decay:
- Decreases the average number of actions required to solve problems by ~10%
- Encourages faster and more efficient problem solving
Answer Diversity:
- Non-repetitive answer ratio increases from 79.7% to 92.8%
- Multi-turn RL with UFO encourages answer diversity and strengthens robustness
License
This model is licensed under Apache 2.0.
Acknowledgements
We thank the DeepSeek team for providing the DeepSeek-R1 model and early conceptual inspirations. We are grateful to the veRL team for their infrastructure support and the RAGEN team for their multi-turn RL framework.
- Downloads last month
- 31