Se00n00/TinyLM2-50M-Reasoning

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|>