InfiR-1B-Instruct / README.md
congkai's picture
Update README.md
b8ac2ad verified
metadata
license: llama3.2
language:
  - en
base_model:
  - meta-llama/Llama-3.2-1B
pipeline_tag: text-generation

Model Card for InfiR-1B-Instruct

InfR aims to advance AI systems by improving reasoning, reducing adoption barriers, and addressing privacy concerns through smaller model sizes.

Model Details

Model Description

Model Sources

Uses

Bias, Risks, and Limitations

  • Performance gaps remain vs. 70 B+ models on very hard reasoning (e.g., OlympiadBench).
  • Safety & bias: inherits Llama-3.2 tokenizer & pre-training distribution; may reflect web biases.
  • Knowledge cut-off: mid-2023.
  • Evaluation has focused on English benchmarks; multilingual robustness not verified.

How to Get Started with the Model

Installation

First, install the required dependencies:

pip install torch transformers

For optimal performance, we recommend using PyTorch 2.0+ and CUDA 11.8+.

Basic Usage

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

Advanced Usage Examples

1. Mathematical Reasoning

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

2. Code Generation

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

3. Chain-of-Thought Reasoning

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

Training Details

Training Data

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.

Training Procedure

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

Evaluation

Benchmarks & Results

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

Technical Specifications

Model Architecture and Objective

  • Base: Llama-3.2-1B (32 layers, 32 heads, RoPE, GQA, 2 k ctx → 4 k extended)

Citation

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.


Glossary

  • SLM: Small Language Model (<2 B parameters)
  • CoT: Chain-of-Thought prompting or training
  • REC: Renewable Energy Certificate
  • PUE: Power Usage Effectiveness (ratio of total facility power to IT power)