d0rj/q-prefixlm-51M-base

prefixlm-51M-base

A 50,866,688-parameter English base model trained from random initialization in the Tiny llm ablation experiment.

Architecture and references

10 decoder layers, width 512, SwiGLU 1792, GQA 8 query / 2 KV heads, RoPE, RMSNorm, tied embeddings. Bidirectional prefix fractions are sampled from {0.25, 0.5, 0.75}; only causal suffix targets contribute to training loss. Sequences shorter than 100 tokens fall back to causal training. Evaluation conditions on the entire bidirectional text context, with a causal answer. Context length: 2048. The architecture/tokenizer baseline is Q-50M-Base; these weights are not a fine-tune of that checkpoint. The unchanged tokenizer has 32,768 entries. The prefix objective follows the prefix language-modeling family described in T5.

Training

Exactly 3,932,160,000 processed source tokens, 15,000 optimizer steps, FineWeb-Edu sample-10BT. This counts input blocks, not unique text or supervised target tokens. Effective batch: 8 × 16 accumulation × 2048 = 262,144 source tokens per step. Shuffle buffer 100,000; seed 2026.

BF16 compute / FP32 weights, one RTX 5070 Ti 16 GB. Fused AdamW, peak LR 0.001, 150-step linear warmup then cosine decay to 0.0001; betas (0.9, 0.95), weight decay 0.1 excluding bias/norm/1D parameters, gradient clipping 1.0. Exact configuration. Equal source-token budgets do not imply equal supervision or FLOPs.

Evaluation

Full splits, zero-shot, lm-eval 0.4.12, BF16, RTX 5070 Ti, max context 2048, no chat template. Protocol: Bidirectional text context, causal answer; answer-only conditional likelihood. Accuracy is percent; ± means one standard error; brackets are 95% Wilson intervals. LAMBADA requires every final-word token to match. acc_norm normalizes option likelihood by harness length. All metrics and provenance.

WikiText-2 raw test continuation: 291 nonoverlapping 1024-token blocks, 512 prefix + 512 scored suffix tokens; 148,992 scored tokens, 335 tail tokens excluded. GPU FP32 PPL 39.596 [38.055, 41.170], NLL 3.678741 [3.639031, 3.717709]. Percentile block bootstrap, 10,000 resamples, seed 2026. BF16 results are stored separately. This is conditional continuation PPL, not standard rolling or word PPL. Intervals do not capture training-seed variability or all within-document dependence.

model-index contains author-reported scores; no official leaderboard submission or verified badge is claimed.

Usage

Install requirements.txt. Custom code requires trust_remote_code=True.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
repo = "d0rj/prefixlm-51M-base"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True).eval()
inputs = tokenizer("The purpose of science is", return_tensors="pt", add_special_tokens=False)
inputs["prefix_lengths"] = torch.tensor([inputs["input_ids"].shape[1]])
with torch.no_grad():
    output = model.generate(**inputs, max_new_tokens=32, do_sample=False)
print(tokenizer.decode(output[0], skip_special_tokens=True))

Reproduce core scores from a downloaded repository after installing evaluation/requirements.txt:

python evaluation/run_core.py --device cuda:0 --dtype bfloat16 --batch-size 16 --output evaluation-rerun

TensorBoard and limitations

Training and evaluation event files contain 750 training-loss points (steps 20–15,000) and eval/ scalars at step 15,000, including CI bounds. Looped training resumed from a saved checkpoint; TensorBoard retains the valid history across the restart.

Small English base model, not instruction-tuned. Single training seed; benchmark contamination was not audited. Reference models with other tokenizers, prompts or data are not directly interchangeable. These checkpoints are research ablations.