Vashtra-0.6B

Vashtra-0.6B

A 0.6B model tuned to answer machine learning questions.

The base model, Qwen3-0.6B, can discuss ML but its knowledge comes from general web text. Vashtra is trained further on the ML literature and on ML questions that practitioners actually asked and answered. It is small enough to run in a browser tab, which is what the demo Space does.

Weights are not up yet. This repo currently holds the full recipe: corpus builder, training script, ONNX export, and a Colab notebook that runs all of it. I will add weights and measured numbers when the training run finishes. There are no benchmark claims below because I have not measured anything yet.

Training

Two stages, both full-parameter. At 0.6B a full fine-tune beats LoRA for domain adaptation and still fits in a free Colab T4, using about 10GB of 16GB.

Stage A, continued pretraining. Full-text arXiv ML papers from common-pile (openly licensed only) plus 117k arXiv cs.LG titles and abstracts, packed into dense 1024-token blocks so no compute goes to padding. LR 5e-5, cosine, one epoch. The LR is low on purpose. The goal is to absorb the domain without wrecking the base model's general ability.

Full text matters here. Abstracts teach the model what papers claim. Full text teaches it how the argument is built, which is what you want when someone asks it to explain a method.

The corpus is about 90M tokens, more than a free Colab session will finish, so stage A takes a --cpt_tokens budget and pushes checkpoints to the Hub. Passing --resume picks up where a dropped session stopped.

Stage B, supervised fine-tuning. ML, stats and CS StackExchange answers that were accepted with score 3 or higher, arXiv explanation and titling tasks, and three subsets of smoltalk2: science reasoning, instruction following, and general chat. LR 1e-5, two epochs, loss masked so only assistant turns count.

The general chat anchor

Fine-tuning a 0.6B model only on domain data reliably destroys its ability to hold a normal conversation. About 22k general examples from smoltalk2 stay in the SFT mixture to stop that. It costs a little domain sharpness and keeps the model usable.

A tokenisation detail

Qwen3's chat template adds an empty <think>\n\n</think> block to the last assistant turn. If you tokenise message by message, that block disappears, and you end up training on a format the model never sees at inference. The encoder here is checked against apply_chat_template(...) and round-trips to it exactly. If you fork this, keep that check. It is easy to get wrong and it fails silently.

Reproducing it

python build_corpus.py    # builds and pushes vigneshwar234/vashtra-ml-corpus
python train_vashtra.py   # stage A, stage B, push
python export_onnx.py     # ONNX + quantisation for the browser demo

Or open Vashtra_training.ipynb in Colab, set the runtime to T4, and run all. About two to three hours on the free tier.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained(
    "vigneshwar234/Vashtra-0.6B", dtype=torch.bfloat16)
tok = AutoTokenizer.from_pretrained("vigneshwar234/Vashtra-0.6B")

msgs = [{"role": "user", "content": "When should I use focal loss instead of cross-entropy?"}]
ids = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt")
out = model.generate(ids, max_new_tokens=256)
print(tok.decode(out[0][ids.shape[1]:], skip_special_tokens=True))

Evaluation

Not run yet. The plan is held-out perplexity on arXiv cs.LG text against the base model, which is cell 8 of the notebook. Numbers go here when I have them.

What it is for, and what it is bad at

Useful for explaining ML concepts, summarising papers, and on-device or edge cases where 0.6B is the size you can afford.

It is still a 0.6B model. It gets specifics wrong, especially exact numbers, named papers, and anything recent. Do not cite it and do not use it for anything that matters without checking the answer.

Licence

Apache-2.0, same as the Qwen3-0.6B base. The corpus mixes licences: full-text papers are CC-BY, CC-BY-SA or CC0 through common-pile/arxiv_papers_filtered, abstracts are AFL-3.0 through CShorten/ML-ArXiv-Papers, StackExchange content is CC-BY-SA-4.0, and smoltalk2 is Apache-2.0. Anything derived from StackExchange stays CC-BY-SA-4.0.

Built by vigneshwar234. Code on GitHub.