Tiny Specialized Models

Can tiny specialized models outperform frontier LLMs on very specific tasks? Yes, of course!

While the world pursues general intelligence with large models, we take a step in the other direction -- towards tiny specialized models, that serve a particular outcome, and run at a fraction of the cost.

We begin by exploring a proof-of-concept on the differentiation task, i.e. find the derivative $\frac{d}{dx} f(x)$ of a given expression $f(x)$. For example,

Our goal is to find the smallest model that can solve this simple, yet interesting, reasoning task.

📄 Paper: Coming soon!

We build 4 differentiation models, Diff-TSMs, spanning 135M to 1.5B parameters, trained using a combination of supervised finetuning (SFT) and reinforcement learning (RL). More details to follow in the paper and training code.

# pip install transformers
from transformers import AutoModelForCausalLM, AutoTokenizer
checkpoint = "AvataarAI/Diff-TSM-Qwen-0.5B"
device = "cuda" # for GPU usage or "cpu" for CPU usage
model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
messages = [
    {"role": "system", "content": "Please reason step by step, and put your final answer within \\boxed{{}}."},
    {"role": "user", "content": r"What is the derivative of the following function w.r.t. x: e^(2x) \sin(14x)"},
]
input_text=tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
outputs = model.generate(inputs, temperature=0.1, top_p=0.7, top_k=50, do_sample=True, max_new_tokens=2048)
print(tokenizer.decode(outputs[0]))

# Create a virtual environment
python -m venv venv
source venv/bin/activate

# Download the veRL submodule
git submodule update --init --recursive   

# Install dependencies
pip install hydra-core==1.4.0dev1 \
            antlr4-python3-runtime==4.11 \
            sympy flash-attn typer
pip install -e external/verl

Create pairs of (function, derivative) without repetition, format them into LLM prompts, and split them into train and test sets. .parquet files are saved in the --local-dir directory.

python -m tslm.data main-generate-and-format \
    --local-dir "./data/calc-diff" \
    --simplicity 0.5 \
    --n-train 512 \
    --n-test 512 \
    --workers 40

# Outputs are saved under the key 'responses' in the output parquet file.
MODEL_PATH="AvataarAI/Diff-TSM-Qwen-0.5B" # HuggingFace model or local path
OUTPUT_PATH=./outputs/diff-tsm-qwen-0.5b.out.parquet

python -m verl.trainer.main_generation \
    trainer.nnodes=1 \
    trainer.n_gpus_per_node=1 \
    data.path="./data/calc-diff/test.parquet" \
    data.prompt_key=prompt \
    data.n_samples=1 \
    data.batch_size=1000 \
    data.output_path=$OUTPUT_PATH \
    model.path=$MODEL_PATH \
    rollout.temperature=0.1 \
    rollout.top_k=50 \
    rollout.top_p=0.7 \
    rollout.prompt_length=256 \
    rollout.response_length=2048 \
    rollout.tensor_model_parallel_size=1 \
    rollout.gpu_memory_utilization=0.8 \
    +actor.ulysses_sequence_parallel_size=1 \
    +rollout.n=1

The tslm/eval.py script symbolically evaluates the predicted derivatives using the reward function in tslm/reward.py and reports average accuracy.

python -m tslm.eval ./outputs/diff-tsm-qwen-0.5b.out.parquet --workers 40

We evaluate the performance of base-line models upto 8B parameters, including those trained for the math domain without task specificity, and OpenAI’s o-series of models as representatives of large “intelligent” LLM Results are in decreasing order of model size, with ours at the bottom.

Results:

Coming soon!

• Harder hand-crafted examples to test OOD generalization.
• Update reward to incentivize simpler expressions.
• More tasks!

• We use the veRL library for both training and running inference.
• We use SymPy extensively for all symbolic operations.
• We thank HuggingFace and the Qwen team for open-sourcing the SmolLM and Qwen2.5 series of models, which we finetune for our experiments.
• Last but not least, Typer is awesome for building CLIs!

If you find this work helpful, please cite us as follows:

@misc{tslm,
  author       = {Shubham Goel and Shubham Jain and Gaurav Baid and Sravanth Aluru},
  title        = {Tiny Specialized Models},
  howpublished = {https://github.com/AvataarAILabs/tsm},
  year         = {2025}
}