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