Metis-1.5-base

Metis-1.5-base

Metis-1.5-base is an 898M-parameter (≈340M active/token) single-latent Mixture-of-Experts language model, pretrained from scratch on 50B tokens of curated, decontaminated English text. It was trained end to end in pure JAX on a single TPU v6e-8 with a fully custom training stack — no PyTorch, no Megatron, no Hugging Face Trainer.

[!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

Metis is an independent, from-scratch exploration of how much capability you can pack into a small, efficient Mixture-of-Experts model trained on a single TPU pod — with a custom JAX stack you fully understand, top to bottom, rather than a black-box framework.

Three convictions shape it:

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.

Metis-1.5 is a proof of concept for that architecture and pipeline — small enough to iterate on quickly, complete enough to be genuinely useful as a base for fine-tuning and continued pretraining.

---

Architecture

Metis-1.5 is a decoder-only transformer whose feed-forward sublayer is replaced by a single-latent Mixture-of-Experts (LatentMoE) block.

Single-latent MoE

Instead of giving each expert its own full-width FFN, every token is first projected down into a shared 512-dim latent space (latent_down: 1536 → 512). Routing and all expert computation happen in that compact latent space, then a single latent_up projection (512 → 1536) maps the combined result back. Concretely, per MoE layer:

• 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.

This is why total parameters (898M) and active parameters (≈340M/token) diverge so much: only 4 of 32 routed experts fire per token.

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

50B tokens of English text, assembled from open corpora and bucketed by capability target so that under-filled sources can only fall back within the same bucket — keeping the intended capability mix intact.

Every source passes the same gate before tokenization:

• 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.

All credit for the underlying corpora belongs to their original authors and curators (see Acknowledgements).

---

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

The released checkpoint is the bf16 export of the fp32 master weights at end of pretraining (step 95,367 ≈ 50B tokens).

Tokenizer

A custom byte-level BPE tokenizer, vocabulary 32,768, trained on ~12M sampled documents (min pair frequency 2). Special tokens: <pad>=0, <bos>=1, <eos>=2, <unk>=3. The full tokenizer is included in this repository (tokenizer.json).

---

Intended uses & limitations

Intended for

• 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.

Not intended for

• 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.

Limitations

• 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.

Load the raw tensors with safetensors (framework-agnostic):

from safetensors import safe_open

weights = {}
with safe_open("model.safetensors", framework="numpy") as f:
    for k in f.keys():
        weights[k] = f.get_tensor(k)   # e.g. weights["layers.0.q"], weights["embed"]

# config.json holds d_model, n_layer, moe_num_experts, head_dim, rope_theta, ...
# Wire these tensors into the forward pass described in the Architecture section
# (single-latent MoE: latent_down -> softmax top-4 routing -> squared-ReLU experts
#  in 512-dim latent + shared expert -> latent_up; GQA + NeoX RoPE; RMSNorm; tied embeddings).

A reference JAX forward/generation implementation is part of the Metis training stack; a standalone single-file loader and an optional transformers wrapper are planned follow-ups.

Evaluation

0-shot accuracy on the full test split of each benchmark, scored with a custom JAX harness — multiple-choice by length-normalized loglikelihood (acc_norm) or plain loglikelihood (acc); GSM8K by greedy generation (chat template + flexible numeric extraction). Training data was decontaminated against these benchmarks, so these are clean held-out numbers.

¹ BoolQ majority-class baseline ≈ 62%. GSM8K is not run for the base model (non-instruct).

How to read these. Metis-1.5 is a ~340M-active model (898M total, MoE) trained on only 50B tokens — far fewer than the 0.3–18T behind modern sub-2B models — so it lands around GPT-2-medium tier: clearly above chance on ARC-Easy, modestly so on HellaSwag / PIQA / WinoGrande, and at chance on MMLU (which needs more knowledge capacity than this scale holds). Supervised fine-tuning leaves raw knowledge roughly unchanged (base ≈ think on multiple choice) while adding instruction-following — visible on BoolQ (+9.5) and GSM8K (7.6%), the latter notably strong for the scale (TinyLlama-1.1B / Pythia-1B sit ~2–3%), reflecting the math/reasoning-heavy data mix. The clearest lever for higher scores is more training tokens, not a different architecture.

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

Compute & environmental footprint

Trained on a single Google Cloud TPU v6e-8 in JAX. At ≈245k tokens/sec, the 50B-token run corresponds to on the order of ~2.4 days of single-pod compute. Training on one small pod (rather than a large cluster) keeps the footprint modest and the experiment reproducible.

---

License

Released under CC0-1.0 — a public-domain dedication. You may use, modify, redistribute, fine-tune, and build upon Metis-1.5-base for any purpose, commercial or otherwise, with no restrictions and no attribution required. The model is provided as-is, without warranty of any kind.

(Note: the underlying training corpora are governed by their own respective licenses; CC0 applies to these released model weights.)

Citation

@misc{metis15base2026,
  title  = {Metis-1.5-base: A single-latent Mixture-of-Experts language model trained from scratch on one TPU pod},
  author = {Lernex},
  year   = {2026},
  howpublished = {Hugging Face},
  note   = {898M parameters (340M active), pretrained on 50B tokens in JAX on TPU v6e-8}
}

Acknowledgements

Deep thanks to the open-data community whose corpora made this possible — including DCLM, FineWeb / FineWeb-Edu / FineMath, Proof-Pile-2, peS2o, OpenWebMath, Nemotron-CC, MegaMath, Cosmopedia, Project Gutenberg / PG-19, Common-Pile, Common-Corpus, OpenStax, and the maintainers of the StackExchange and Wikimedia data dumps — and to the JAX and Cloud TPU teams for the training stack.