quark28/mini-llm-thinking-sft

mini-llm · Thinking SFT (research variant)

The same masked-diffusion encoder–decoder as quark28/mini-llm-classic-sft, with an experimental Thinking module between the encoder and the decoder. Kept for research; for chat use the classic model.

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

Model

Parameters	19.8M (embeddings tied with the vocabulary projections)
Tokenizer	byte-level BPE, 12,288 tokens
Layout	encoder 4 blocks → Thinking → decoder 2 blocks, d = 384, 6 heads
Input	context 256 tokens (left-padded) + answer slot 64 tokens
Objective	masked (absorbing-state) diffusion on the answer slot

Thinking projects encoder states onto the (tied) vocabulary embeddings and replaces the encoder output with a per-position channel plus one context-level vector. Under the same training pipeline it does not beat the plain encoder: SciQ −15.0 points, BLiMP +2.9, other differences are within noise.

Training data and procedure are the same as for the classic model (FineWeb-Edu pretraining → dialogue SFT).

Usage

PyTorch only (no ONNX export for this variant):

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

d = snapshot_download('quark28/mini-llm-thinking-sft')
sys.path.insert(0, d)
from modeling_thinking import NN_Thinking
from sampler import generate, 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 = NN_Thinking(**cfg['MODEL_SETTINGS'])
safetensors.torch.load_model(model, f'{d}/model.safetensors')
model.eval()

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

Sampling is the same as in the classic model card, with margin = 0.3.

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

Same as the classic model; in addition, replies drift off-topic and repeat words more often.

License

Weights: CC BY-NC-SA 4.0 (see the classic model card). Code: see the GitHub repository.