quark28/mini-llm-classic-sft

mini-llm · classic SFT

A 20M-parameter chat model trained from scratch with a masked-diffusion objective: an encoder reads the conversation, and a small decoder fills a 64-token answer slot by iterative unmasking instead of writing left to right. English small talk only.

Code: https://github.com/quark28/mini-llm

Model

Parameters	20.1M (input and output embeddings tied)
Tokenizer	byte-level BPE, 12,288 tokens
Layout	encoder 6 blocks → decoder 2 blocks, d = 384, 6 heads
Blocks	pre-norm RMSNorm, SwiGLU FFN (1024), RoPE in self-attention; decoder cross-attends to the encoder
Input	context 256 tokens (left-padded) + answer slot 64 tokens
Objective	masked (absorbing-state) diffusion on the answer slot

Training

• Pretraining — FineWeb-Edu (sample-10BT).
• SFT — DailyDialog, EmpatheticDialogues, PersonaChat, MBPP and ~1.3k synthetic small-talk / polite-refusal dialogues. Prompt template: <user> text <bot> reply <eos>. The loss covers the whole answer slot including padding, so the model learns where to stop.

Trained on Kaggle T4 GPUs.

Sampling

Generation is not left-to-right. For every reply:

1. Length. One decoder pass over a fully masked slot; the expected reply length is Σ (1 − p(<eos>) − p(<pad>)) over positions. Four candidates are generated with lengths ×0.8 … ×1.2.
2. Block-wise unmasking. The slot is filled in blocks of 8 tokens; inside a block, one masked position with the highest confidence is revealed per step (temperature 0.5, top-k 50, repetition penalty, n-gram bans against the reply and the user's last message).
3. Remasking. After each block and at the end, every token is re-scored with itself masked (leave-one-out); a token whose probability is below margin × p(best alternative) is replaced. This is a training-free variant of G-Star remasking (arXiv:2510.08369).
4. Selection. The candidate with the highest mean leave-one-out log-probability is returned.

The full sampler is in sampler.py; its settings are in config.json → SAMPLER_SETTINGS.

Usage

ONNX (no PyTorch model code needed, only the sampler, which uses torch for tensors):

import json, sys, torch
from huggingface_hub import snapshot_download
from tokenizers import Tokenizer

d = snapshot_download('quark28/mini-llm-classic-sft')
sys.path.insert(0, d)
from sampler import generate, OrtBackend, PAD_ID, USER_ID, BOT_ID, EOS_ID

tok = Tokenizer.from_file(f'{d}/tokenizer.json')
cfg = json.load(open(f'{d}/config.json'))
model = OrtBackend(f'{d}/onnx/encoder.onnx', f'{d}/onnx/decoder.onnx')

ids = [USER_ID] + tok.encode('Hello! How are you?').ids + [BOT_ID]
ids = [PAD_ID] * (256 - len(ids)) + ids
out = generate(model, torch.tensor(ids), device='cpu', **cfg['SAMPLER_SETTINGS']).tolist()
print(tok.decode(out[:out.index(EOS_ID)]))

PyTorch: modeling_classic.py + model.safetensors (safetensors.torch.load_model(NN_Classic(**cfg['MODEL_SETTINGS']), 'model.safetensors')); the same generate works on it.

The ONNX graphs: encoder.onnx (context_ids [1, 256] → memory, mask) runs once per reply; decoder.onnx (memory, mask, answer_ids [batch, 64] → logits [batch, 64, 12288]) runs on every step, batch is dynamic for the remasking pass.

Benchmarks

Model	Params	PIQA	HellaSwag*	ARC-Easy*	ARC-Challenge*	WinoGrande	SciQ	LAMBADA (OpenAI)	BLiMP
mini-llm classic SFT	20.1M	54.5	26.7	31.9	21.0	49.2	48.7	0.0	56.0
mini-llm Thinking SFT	19.8M	52.6	26.0	30.3	22.4	51.9	33.7	0.0	58.9
Pythia-14m	14M	56.9	26.5	32.7	21.0	51.3	50.4	10.3	67.7
GPT-2	124M	62.9	31.1	39.5	22.7	51.6	75.2	32.6	82.1
random	—	50.0	25.0	25.0	25.0	50.0	25.0	0.0	50.0

Zero-shot, lm-evaluation-harness 0.4.13, accuracy in %. * — length-normalized accuracy (acc_norm). Our models are scored with an exact diffusion likelihood: the answer slot is unmasked left to right and the log-probability of each revealed token is summed. Answer options longer than the 64-token slot are truncated (PIQA 8.5% of options, HellaSwag 1.7%, others 0%).

These are base-level language benchmarks: they measure likelihood on raw text, not dialogue. At 20M parameters both models are close to chance on reasoning tasks (ARC, WinoGrande, HellaSwag).

Limitations

• English only; replies are short (usually one sentence).
• No real world knowledge; answers questions in the right form but not with the right content.
• Often misreads the emotion of a message, e.g. "I lost my job yesterday." → "Oh no! I am glad you lost your job."
• Not instruction-following, not safe for any real use. It is a learning / research project.

License

Weights: CC BY-NC-SA 4.0 (the strictest license among the SFT datasets: DailyDialog is CC BY-NC-SA 4.0, EmpatheticDialogues is CC BY-NC 4.0). Code: see the GitHub repository.