cogito-3b

<div align="center">

Cogito-3B

GRPO reasoning training on Qwen2.5-3B — Countdown and math, verifiable rewards only.

</div>

---

Cogito-3B is Qwen/Qwen2.5-3B fine-tuned with GRPO (Group Relative Policy Optimization) using only verifiable-correctness reward no reasoning demonstrations, no learned reward model, and no preference data. Training follows a two-stage curriculum (Countdown, then math) and reproduces the reinforcement-learning-from-verifiable-rewards (RLVR) setup associated with DeepSeek-R1-Zero at the 3B scale. Over training, mean reward increases and mean completion length grows; the model produces an explicit <think> … </think> trace before a final <answer>.

This card documents the training setup and held-out evaluation. All numbers are reproducible from the recipe in the linked repository.

Results (held-out)

pass@1, greedy decoding. Base is the unmodified Qwen2.5-3B evaluated under the same prompts and decoding.

The largest change is on Countdown, a task under-represented in standard pretraining, where the base model rarely emits a valid solution. GSM8K sits near this base model's ceiling under the given prompt, so the change there is small. Sampled reasoning traces are provided in results/aha_transcripts.md.

Model description

• Base model: Qwen/Qwen2.5-3B
• Method: GRPO via TRL, with vLLM serving rollouts
• Objective: verifiable-correctness reward + a small format reward; no supervised reasoning data, no reward model
• Format: base completion model; the prompt ends with Assistant: <think>\n and the model completes … </think> <answer> … </answer>
• Precision / compute: bf16 with gradient checkpointing; A100 80GB (1 GPU for rollouts, 3 for training)

Intended use and format

Cogito-3B is a base completion model, not an instruction/chat model. It expects the training prompt format shown below.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Bluebox85033/cogito-3b"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="bfloat16", device_map="auto")

def countdown_prompt(nums, target):
    return (
        "A conversation between User and Assistant. The User poses a Countdown number puzzle, and the Assistant solves it.\n"
        "The Assistant first reasons step by step inside <think> </think> tags, and then gives ONLY the final arithmetic expression inside <answer> </answer> tags.\n"
        "Rules: use each given number exactly once; the only allowed operations are + - * / and parentheses; the expression must evaluate exactly to the target.\n"
        "Example answer format: <answer>(3 + 5) * 2</answer>\n\n"
        f"User: Numbers: {nums}. Target: {target}. Find an expression that uses each number exactly once and equals {target}.\n"
        "Assistant: <think>\n"
    )

prompt = countdown_prompt([3, 5, 2], 16)
inputs = tok(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=1024, do_sample=False)  # greedy
print("<think>\n" + tok.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True))

For math, use the same structure with the math template (final answer as \boxed{...}):

def math_prompt(problem):
    return (
        "A conversation between User and Assistant. The User asks a math question, and the Assistant solves it.\n"
        "The Assistant first reasons step by step inside <think> </think> tags, and then gives the final answer inside <answer> </answer> tags, written as \\boxed{...}.\n"
        "Example answer format: <answer>\\boxed{42}</answer>\n\n"
        f"User: {problem}\n"
        "Assistant: <think>\n"
    )

Training procedure

Two-stage curriculum:

1. Countdown. A self-generated, deduplicated arithmetic-puzzle dataset; train and test are disjoint by (sorted numbers, target). Reward: the model's expression must use each number exactly once with only + - * / ( ), and evaluate exactly to the target.
2. Math. Continued from the Countdown checkpoint on GSM8K (and a MATH train mirror). Reward: \boxed{} answer equivalence to ground truth.

Across steps, mean reward increases and mean completion length grows. Full data generation, reward functions, and training/evaluation code are in the GitHub repository and under recipe/ in this repo.

Evaluation

pass@1 with greedy decoding on held-out splits. The Countdown test set is self-generated and disjoint from training; GSM8K uses the official test split; MATH-500 is HuggingFaceH4/MATH-500. Sample sizes are given in the results table; the Countdown set (n=64) is small and its figure should be read as indicative.

Relation to prior work

The method (RLVR / GRPO) and the Countdown demonstration follow DeepSeek-R1-Zero and the TinyZero reproduction. This model is a reproduction and extension rather than a new method: it provides an end-to-end reproducible pipeline and a Countdown→math curriculum with a controlled before/after measurement on held-out splits.

Limitations

• Base completion model: requires the specific prompt format; not a general-purpose assistant.
• 3B scale: outputs can be verbose and may be confidently incorrect.
• The Countdown evaluation set is small (n=64).

Citation

@misc{cogito3b,
  title  = {Cogito-3B: GRPO reasoning training on Qwen2.5-3B (Countdown and math)},
  author = {Bluebox85033},
  year   = {2026},
  howpublished = {\url{https://huggingface.co/Bluebox85033/cogito-3b}}
}