TinyLM2-50M-Reasoning
TinyLM2-50M-Reasoning is a compact decoder-only Transformer language model designed for efficient instruction following and conversational AI. The model has approximately 50M parameters and has been fine-tuned using Reasoning Fine-Tuning (RFT) on reasoning-trace data to improve chain-of-thought and multi-step problem-solving capabilities while maintaining a lightweight footprint suitable for local inference and resource-constrained environments.
Evaluation
All evaluations are zero-shot unless stated otherwise, and i used lm_eval to run them.
Versus the SupraLabs/Supra-50M-Reasoning reference at the same ~50M scale, TinyLM2-50M-Reasoning lands roughly par — beating Supra-50M-Reasoning on BLiMP, PIQA, COPA, BoolQ, TruthfulQA MC2, SWAG, RACE and WikiText-2 perplexity, while trailing on WinoGrande, HellaSwag, CommonsenseQA, SciQ, ARC-Easy, OpenBookQA, ARC-Challenge, MMLU and LAMBADA. It is a competitive, comparable checkpoint rather than a clear winner.
Model Architecture & Hyperparameters
TinyLM2-50M-Reasoning is built on a custom ALiBi Decoder-Only Transformer architecture with pre-normalization and gated feedforward networks:
| Hyperparameter | Value | Description |
|---|---|---|
| Architecture | ALiBi Decoder-Only Transformer | Autoregressive Decoder-Only Transformer |
| Total Parameters | ~50.96M (53,430,272) | Compact and ultra-fast for edge & local CPU/GPU |
| inference | ||
vocab_size | 50,271 | Includes special chat tags (`< |
hidden_size (d_model) | 512 | Model hidden dimension |
intermediate_size (ff_hidden_d) | 819 | SwiGLU Gated Feedforward hidden dimension |
num_hidden_layers | 12 | Number of Transformer block layers |
num_attention_heads | 8 | Attention heads (Head dim = 64) |
max_position_embeddings | 2,048 | Maximum context sequence length |
| Normalization | RMSNorm (eps=1e-8) | Scale normalization for accelerated throughput |
| Activation Function | SwiGLU (SiLU) | Gated Feedforward activation |
| Positional Encoding | ALiBi | Attention with Linear Biases |
| Tie Word Embeddings | True | Tied input embedding and LM head projection |
Tokenizer & Chat Template
The model uses a custom Byte-Level BPE Tokenizer equipped with special tokens and a pre-configured Jinja2 chat_template for multi-turn conversations, including <|THINK|> reasoning-trace markers.
| Property | Value |
|---|---|
| Tokenizer Type | GPT2Tokenizer (Byte-Level BPE) |
| Vocabulary Size | 50,271 |
| Special Tokens | `< |
| Chat Control Tokens | `< |
| Extra Special Tokens | `< |
| Chat Template | Native Jinja2 support via tokenizer.apply_chat_template() |
Training Configuration
| Parameter | Value |
|---|---|
| Pipeline Process | Reasoning Fine-Tuning (RFT) |
| Dataset | Se00n00/reason_test (reasoning-trace Q&A pairs) |
| Total Examples | reasoning-trace Q&A pairs, 6 epochs |
| Learning Rate | 6e-5 |
| Learning Rate Schedule | Cosine (warmup_ratio=0.01) |
| Micro-Batch Size | 2 per device |
| Gradient Accumulation | 16 steps |
| Effective Batch Size | 32 × 2,048 tokens |
| Optimizer | AdamW (weight_decay=0.1) |
| Max Sequence Length | 2,048 tokens |
| Precision | float16 |
| Hardware | NVIDIA Tesla T4 x 2 GPU |
Inference
# pip install torch transformers
import torch
from transformers import pipeline
pipe = pipeline(
"text-generation",
model="Se00n00/TinyLM2-50M-Reasoning",
trust_remote_code = True
)
messages = [
{"role": "system", "content": "You are a helpful AI assistant."},
{"role": "user", "content": "Explain artificial intelligence in simple terms."}
]
prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
result = pipe(
prompt,
max_new_tokens=120,
do_sample=True,
temperature=0.4,
top_k=40,
top_p=0.9
)
print(result[0]['generated_text'])
──────
Sample Outputs
Chat-templated generation showing the model's thinking trace (temperature 0.4, top-p 0.9, repetitive middle trimmed with …):
System: You are a helpful AI assistant
User: A train travels 120 km in 2 hours. What is its average speed?
<|THINK|>Okay, the user is asking about its average speed. Let me start by recalling what I know. The average speed is when a train is traveling at a speed that's higher than the speed of the train. So, the main factors are speed, speed, and the train's speed.First, the train's speed is the speed of the train itself. … I should also note that the train's speed is the key factor, and to keep the explanation concise.
<|/THINK|>The average speed of the train is measured by its speed …
System: You are a helpful AI assistant
User: Why is the sky blue?
<|THINK|>Okay, the user is asking about the sky blue. Let me start by recalling what I know. The sky is blue because it's a white, almost white. It's a combination of colors, like red, yellow, and blue. The sky is a lot of blue, so I need to make sure I'm not missing anything. Maybe the key points: the blue sky is a lot of blue, so the key points are the blue sky, the blue sky, and the blue sky. … I should also mention the key points: the blue sky is a lot of blue, so the key points are the blue sky, the blue sky, and the blue sky. Keep it concise, under 150 words. …<|/THINK|>