TR-MOE-306

TR-MOE-306 is the 306.5M-parameter token-identity residual model from Token Identity as a Routing Signal for Residual MLP Experts. Every token passes through a shared dense SwiGLU branch, while its token ID selects two of four narrow residual experts through a fixed, layer-specific lookup. The experts still transform the contextual hidden state.

It was trained from scratch on approximately 8B FineWeb-Edu tokens. This is a pretrained base model, not an instruction-tuned or safety-aligned assistant.

Architecture

  • 18 transformer layers
  • hidden size 1,024
  • 16 query heads and 4 key/value heads (GQA)
  • QK-Norm and RoPE
  • 2,048-token training context
  • 32,000-token tokenizer
  • shared SwiGLU width 3,840
  • 4 residual experts of width 64
  • fixed top-2 lookup with 0.5/0.5 route weights
  • learned scalar gates on shared and routed outputs
  • 306,486,564 trainable parameters

Token identity controls parameter selection only; it does not replace contextual computation.

Evaluation

Both models were evaluated zero-shot with EleutherAI LM Evaluation Harness 0.4.12, identical seeds, BF16 inference, no sample limit, and a maximum context of 2,048 tokens.

Model ARC-Easy acc_norm PIQA acc_norm HellaSwag acc_norm WikiText-2 word PPL
Dense-306 49.07 63.49 33.67 35.79
TR-MOE-306 48.61 63.49 34.06 35.20

The multiple-choice differences are smaller than one reported standard error. These results support near-parity, not task-level superiority.

At the last common point on a fixed FineWeb-Edu training-split evaluation stream, the token-routed model reached NLL 2.9329 and the dense model reached 2.9482. That stream is not an independent held-out set, and each architecture was trained with one seed. The routed implementation was slower in the reported training setup.

Files and integrity

model.safetensors is the exact FP32 model export used for the reported downstream evaluation. It preserves each layer's saved routing table and learned shared/routed gates.

  • final training step: 7629
  • training checkpoint SHA-256: 361b73873ac0dda66552325cc31328666b9d4c7684d7041e5b37cd43c42712a2
  • exported model.safetensors SHA-256: 877e9bcc1622b313f0598025a44991b58fa8b646fc81ca5127b73594da92d643

The repository also contains the tokenizer, generation configuration, and export_manifest.json.

Runtime

The checkpoint uses the custom DeepForCausalLM format implemented by the Complexity-ML vLLM runtime. Stock Transformers and upstream vLLM do not currently register the deep architecture. Use that runtime for inference, or adapt its loader before loading these weights.

The tokenizer can be loaded independently:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Pacific-i64/TR-MOE-306")

Limitations

This is a small research base model. It may generate inaccurate, biased, offensive, repetitive, or unsafe text. It should not be used for high-stakes decisions. The training corpus is web-derived and may contain errors, copyrighted material, personal information, and social biases.

Paper and reproducibility

The paper reports a matched single-seed observation. Multi-seed replication remains open.

License

The TR-MOE-306 model weights are released under Creative Commons Attribution–NonCommercial 2.0 (CC BY-NC 2.0). See LICENSE.

Downloads last month
53
Safetensors
Model size
0.3B params
Tensor type
I64
·
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Dataset used to train Pacific-i64/TR-MOE-306

Space using Pacific-i64/TR-MOE-306 1