Metis-1.5-base

[!IMPORTANT] This is a base model. It is trained only for next-token prediction and is not instruction-tuned — it continues text, it does not follow instructions or hold a conversation. Prompt it with a stem ("The three laws of motion are") rather than a question. An instruction-tuned Metis-1.5-think variant is trained separately.

At a glance

Why Metis exists — the philosophy

1. Sparsity should be cheap, not just powerful. Standard MoE layers duplicate a full FFN per expert, so memory balloons with expert count. Metis routes every expert through one shared low-rank latent space, so adding experts costs far less — you buy specialization without paying full-FFN memory for each one.
2. Data quality beats data quantity at small scale. The 50B-token corpus is aggressively filtered, deduplicated, and decontaminated, and bucketed by capability (web, encyclopedic, scientific, math, reference, books, Q&A, synthetic textbooks) so the mixture is deliberate rather than incidental. The blend leans education- and STEM-heavy on purpose.
3. Own the whole stack. Tokenizer, data pipeline, model, optimizer, sharding, and checkpointing are all custom JAX. The point is understanding, not abstraction.

Architecture

Single-latent MoE

• Router: linear projection of the hidden state → 32 expert logits, softmax scores, top-4 selection. Load is balanced with an auxiliary-loss-free learned bias (DeepSeek-style: the bias steers selection only; combine weights come from the unbiased scores).
• Routed experts: 32 experts, each a squared-ReLU MLP operating in the 512-dim latent (expert intermediate size 1024). Each token is processed by its top-4 experts, combined by normalized router weights.
• Shared expert: 1 always-on expert applied to every token, providing a stable backbone alongside the sparse routed capacity.

Attention & backbone

• Grouped-query attention — 24 query heads, 8 key/value heads (4:1 grouping), head_dim 64.
• Rotary position embeddings (RoPE, NeoX/Llama half-split convention, θ = 10,000).
• RMSNorm (pre-norm), tied input/output embeddings, no biases.
• 19 layers, d_model 1536, context length 1024, vocab 32,768.

Training data

• English-only language ID (confidence ≥ 0.90, paragraph-level English ratio ≥ 0.85, Latin script).
• Exact + near-duplicate deduplication, globally across the corpus.
• Benchmark decontamination against MMLU, MMLU-Pro, ARC, HellaSwag, OpenBookQA, SciQ, TruthfulQA, GPQA, GSM8K, MATH, AIME, AMC, OlympiadBench, Minerva-Math, SVAMP, ASDiv, IFEval, MT-Bench, Arena-Hard, and RewardBench.

Training procedure

• From scratch, single TPU v6e-8, pure JAX/XLA with a custom data-parallel training loop (pmap, replicate-once state, donated buffers).
• Optimizer: AdaMuon — Newton–Schulz-orthogonalized momentum for 2-D weight matrices, Adam-style second-moment scaling elsewhere.
• Mixed precision: weights are kept as fp32 master copies and cast to bf16 for compute each step, so tiny optimizer updates survive while matmuls stay on the fast bf16 path.
• Schedule: linear warmup → cosine decay.
• Throughput knobs: bf16 attention scores (fp32 softmax reductions), streaming cross-entropy, softmax router with expert capacity factor 1.5, sort-free cumsum dispatch — sustaining ≈245,000 tokens/sec on the pod.
• Sequence length 1024, global batch 64 × 8 grad-accum across 8 devices (≈524k tokens/step).

Tokenizer

Intended uses & limitations

• Research on small-scale / efficient MoE language models.
• A base for continued pretraining, domain adaptation, and instruction/preference fine-tuning.
• Studying the single-latent MoE architecture and JAX/TPU training.

• Drop-in chat or instruction following — use a fine-tuned variant (e.g. Metis-1.5-think) for that.
• Any safety-critical, factual, or high-stakes use.

• Small (≈340M active params): limited world knowledge and reasoning depth versus large models.
• Base only: no instruction tuning, no RLHF, no safety alignment — it can produce inaccurate, biased, or otherwise undesirable text.
• English-only training; ~1024-token context.
• Knowledge is bounded by the training corpus and cutoff.

How to use

[!NOTE] The weights use JAX-native tensor names (embed, final_norm.scale, layers.N.q, layers.N.router, layers.N.expert_w1, …) and the architecture is custom, so this checkpoint does not load via transformers.AutoModel out of the box. It ships as a self-describing safetensors release; config.json carries every dimension, and the Architecture section above specifies the forward pass exactly.

Evaluation

Held-out pretrain-distribution LM loss ≈ 2.34 nats/token (~10.4 ppl) at end of pretraining.

Compute & environmental footprint

License

Citation

Acknowledgements