InfiR: Reasoning-Enhanced Low-Resource Training Pipeline
Collection
InfiR : Crafting Effective Small Language Models and Multimodal Small
Language Models in Reasoning
โข
9 items
โข
Updated
InfR aims to advance AI systems by improving reasoning, reducing adoption barriers, and addressing privacy concerns through smaller model sizes.
First, install the required dependencies:
pip install torch transformers
For optimal performance, we recommend using PyTorch 2.0+ and CUDA 11.8+.
Here's a simple example to get started with InfiR-1B-Instruct:
from transformers import AutoTokenizer, AutoModelForCausalLM
# Define messages in chat format
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "A new program had 60 downloads in the first month. The number of downloads in the second month was three times as many as the downloads in the first month, but then reduced by 30% in the third month. How many downloads did the program have total over the three months? Think step by step."},
]
# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("InfiX-ai/InfiR-1B-Instruct")
model = AutoModelForCausalLM.from_pretrained("InfiX-ai/InfiR-1B-Instruct")
# Apply chat template and generate
raw_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(raw_prompt, return_tensors="pt")
outputs = model.generate(inputs["input_ids"], max_new_tokens=2048)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
# Mathematical problem solving with chat format
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "If a rectangle has a length of 8 units and a width of 6 units, what is its area and perimeter? Solve this step by step."},
]
raw_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(raw_prompt, return_tensors="pt")
outputs = model.generate(
inputs["input_ids"],
max_new_tokens=512,
temperature=0.1,
do_sample=True
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
# Code generation example with chat format
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Write a Python function to calculate the factorial of a number."},
]
raw_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(raw_prompt, return_tensors="pt")
outputs = model.generate(
inputs["input_ids"],
max_new_tokens=256,
temperature=0.2,
do_sample=True
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
# Chain-of-thought reasoning with chat format
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "A train travels 120 km in 2 hours. What is its speed in km/h? Let's approach this step by step."},
]
raw_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(raw_prompt, return_tensors="pt")
outputs = model.generate(
inputs["input_ids"],
max_new_tokens=300,
temperature=0.3,
do_sample=True
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
Stage | Tokens | Composition |
---|---|---|
Pre-training | 900 B | 52 % code, 48 % high-quality web (math, science, encyclopedic) |
Annealing | 40 B | extra math & code + synthetic samples |
SFT | ~4 M | Infinity-Instruct, Orca-AgentInstruct-1M, NuminaMath, ScaleQuest (filtered) |
Data cleaning: heuristic filters, MinHash de-duplication, 10-gram benchmark decontamination, reward-model rejection sampling.
Hyper-parameter | Value |
---|---|
Precision | bf16 mixed |
Optimizer | AdamW |
LR (pre-train) | 1.4 e-3, cosine โ 0 |
LR (SFT) | 2 e-5, cosine w/ 10 % warm-up |
Batch size | 2048 (pre-train), 128 (SFT) |
Sequence len | 4096 |
Epochs | 1 (pre-train), 1 (anneal), 4 (SFT) |
GPUs | 64 ร H800, 5760 GPU-hours total |
Benchmark | InfiR-1B-Instruct | Llama-3.2-1B-Instruct | Qwen-2.5-1.5B-Instruct |
---|---|---|---|
MMLU | 50.22 | 46.27 | 61.78 |
GSM8K | 70.9 | 47.9 | 74.3 |
MATH | 46.4 | 30.0 | 53.4 |
HumanEval | 58.54 | 39.63 | 51.83 |
MBPP | 56.03 | 49.03 | 56.81 |
BibTeX:
@misc{xie2025infir,
title={InfiR: Crafting Effective Small Language Models and Multimodal Small Language Models in Reasoning},
author={Xie, Congkai and Cai, Shuo and Wang, Wenjun and others},
year={2025},
eprint={2502.11573},
archivePrefix={arXiv},
primaryClass={cs.CL}
}
APA:
Xie, C., Cai, S., Wang, W., et al. (2025). InfiR: Crafting Effective Small Language Models and Multimodal Small Language Models in Reasoning. arXiv:2502.11573.