Zenyx-Base-220M: High-Density Nano Foundation Model
Zenyx-Base-220M is a 220 million parameter causal language model built from scratch using JAX/Flax on Kaggle TPU v5e-8.
Unlike typical small models trained on limited data, Zenyx-Base was trained on ~153 Billion tokens—far exceeding the Chinchilla optimal point for this parameter count. This "over-training" strategy was employed to maximize the information density and logic capabilities of the weights, creating a robust foundation for reasoning tasks.
🧠 Model Description
- Architecture: Custom Llama-style Transformer (RoPE, SwiGLU, RMSNorm, Grouped Query Attention).
- Tokenizer: Qwen 2.5 Tokenizer (151,650 Vocab Size) for high compression efficiency.
- Context Window: 2048 Tokens.
- Training Hardware: TPU v5e-8.
- Final Validation Loss: ~2.38 (Exceptional convergence for 220M).
Technical Specifications
| Hyperparameter | Value |
|---|---|
| Layers | 12 |
| Hidden Dim | 768 |
| MLP Dim | 3072 |
| Attention Heads | 12 |
| KV Heads | 4 (GQA) |
| Vocab Size | 151,646 |
📚 Training Curriculum (The "Omni-Mix")
The model was trained using a rigorous 4-stage curriculum designed to layer capabilities sequentially:
- Phase 1: Fundamentals (FineWeb-Edu)
- Focus on high-quality educational English text to establish linguistic baselines.
- Phase 2: Logic & Structure (StarCoder - Python)
- Introduction of code data to enforce logical indentation, syntax, and structured thinking.
- Phase 3: Multilingualism (FineWeb-2)
- Exposure to 6 major languages (Hindi, Chinese, Russian, Japanese, French, Spanish) to expand the semantic embedding space.
- Phase 4: The Infinite Polish (Omni-Mix)
- A weighted interleaving of all previous datasets plus OpenWebMath to converge the model's logic and language capabilities.
💻 Usage
This model is a raw JAX/Flax checkpoint saved in .safetensors format. It uses a custom architecture definition and requires flax and jax to run.
Loading with JAX/Flax
import jax
import jax.numpy as jnp
from flax.training import train_state
from flax import serialization
from safetensors.flax import load_file
from transformers import AutoTokenizer
import flax.linen as nn
# 1. Define Architecture (Must match training config)
class TransformerLM(nn.Module):
vocab_size: int
embed_dim: int = 768
num_layers: int = 12
num_heads: int = 12
num_kv_heads: int = 4
mlp_dim: int = 3072
max_length: int = 2048
dropout_rate: float = 0.0
# ... (Insert full model class definition here from the training script) ...
# 2. Load Resources
repo_id = "Arko007/Zenyx_Base_220M"
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", trust_remote_code=True)
# 3. Initialize & Load Weights
model = TransformerLM(vocab_size=len(tokenizer))
dummy_input = jnp.ones((1, 1), dtype=jnp.int32)
params = model.init(jax.random.PRNGKey(0), dummy_input)['params']
# Load Safetensors
# Ensure model.safetensors is downloaded locally
loaded_params = load_file("model.safetensors")
print("Weights loaded successfully!")
Benchmarks (measured 2026-10-14)
Re-evaluated on a Kaggle T4 in fp32 from model.safetensors in this repository, loaded strictly (every parameter matched by name and shape) into the training
notebook's own TransformerLM definition, with the Qwen2.5 tokenizer named in this card. The siblings are other Zenyx models scored with the identical code and text.
Task suite (the same 19 tasks and extras as Zenyx-V3-Base)
0-shot, full evaluation sets (MMLU 5-shot), log-likelihood scoring of candidate continuations; acc_norm is length-normalised in characters.
Context is 2,048 tokens, so longer prompts keep their last 2,048 tokens. Bold = the conventional metric. Random for BoolQ is its 62.2% majority class.
| Benchmark | Zenyx_Base_220M | Zenyx-v2-base @96,000 | Random | n |
|---|---|---|---|---|
HellaSwag (acc_norm) |
31.13% ± 0.46 | 27.81% | 25.0% | 10,042 |
ARC-Easy (acc_norm) |
46.42% ± 1.02 | 34.13% | 25.0% | 2,376 |
ARC-Challenge (acc_norm) |
25.34% ± 1.27 | 23.98% | 25.0% | 1,172 |
PIQA (acc_norm) |
62.08% ± 1.13 | 54.79% | 50.0% | 1,838 |
WinoGrande (acc) |
51.54% ± 1.40 | 49.96% | 50.0% | 1,267 |
OpenBookQA (acc_norm) |
27.80% ± 2.00 | 23.60% | 25.0% | 500 |
BoolQ (acc) |
58.47% ± 0.86 | 47.43% | 62.2% | 3,270 |
SciQ (acc) |
74.60% ± 1.38 | 51.90% | 25.0% | 1,000 |
LAMBADA (OpenAI) (acc) |
26.22% ± 0.61 | 3.43% | 0.0% | 5,153 |
RACE (acc_norm) |
31.25% ± 0.66 | 26.04% | 25.0% | 4,934 |
CommonsenseQA (acc_norm) |
30.88% ± 1.32 | 26.04% | 20.0% | 1,221 |
COPA (acc) |
59.00% ± 4.92 | 56.00% | 50.0% | 100 |
LogiQA (acc_norm) |
24.27% ± 1.68 | 24.73% | 25.0% | 651 |
WSC273 (acc) |
53.48% ± 3.02 | 50.92% | 50.0% | 273 |
TruthfulQA MC1 (acc) |
19.95% ± 1.40 | 26.44% | 22.8% | 817 |
Arithmetic (acc) |
0.36% ± 0.05 | 0.36% | 0.0% | 14,000 |
WikiText-2 (raw), rolling 1024-token window: 1.0784 bits/byte (token perplexity 25.38; per-token figures depend on the tokenizer).
Zenyx-Fact1000 (internal)
A frozen 1000-question, four-way factual / common-sense completion benchmark (chance 25%): the model must rank the correct completion above three same-type distractors by
log-likelihood (sha256 843416524fb39c5f…; not yet published). The "prior-only" control scores the same candidates after the neutral context "The answer is".
| Model | accuracy | 95% interval | without arithmetic | prior-only control |
|---|---|---|---|---|
| Zenyx_Base_220M | 65.3% | 62.3–68.2% | 73.2% | 29.8% |
| Zenyx-42M | 49.5% | 46.4–52.6% | 54.4% | 29.1% |
| Zenyx-v2-base @96,000 (85M unique) | 45.8% | 42.7–48.9% | 48.1% | 30.7% |
| Zenyx-MCEN-Delta-Mamba @4,100 (565M, still training) | 45.5% | 42.4–48.6% | 45.9% | 29.6% |
| Zenyx-V3-Base @107,200 (1.58B MoE) | 81.1% | 78.6–83.4% | 80.8% | 28.7% |
SciQ (n = 1000, 0-shot, chance 25%, standard error about 1.5 points)
| Model | acc |
acc_norm |
|---|---|---|
| Zenyx_Base_220M | 74.60% | 67.90% |
| Zenyx-v2-base @96,000 | 52.10% | 47.20% |
| Zenyx-MCEN-Delta-Mamba @4,100 | 55.30% | 49.80% |
| Zenyx-V3-Base @107,200 (T4 fp16) | 79.90% | 72.70% |
The Pile, by domain (bits per byte, LOWER is better)
Text: NeelNanda/pile-10k (pinned revision 127bfedcd5), documents grouped by pile_set_name, capped at 1.5M characters per domain, scored with a rolling 1024-token window at
stride 512 (every token scored once with ≥ 512 tokens of left context); bits/byte is normalised by the UTF-8 bytes of the source text, so it is comparable across tokenizers.
| Domain | Zenyx_Base_220M | V2 @96,000 | MCEN @4,100 | V3 @107,200 | text (MB) |
|---|---|---|---|---|---|
| Github | 0.7299 | 1.0832 | 1.0490 | 0.5765 | 1.53 |
| PubMed Central | 0.8031 | 1.3518 | 1.1163 | 0.7426 | 1.51 |
| USPTO Backgrounds | 0.8272 | 1.2530 | 1.0945 | 0.7807 | 1.50 |
| NIH ExPorter | 0.9071 | 1.3249 | 1.1771 | 0.8620 | 0.22 |
| PubMed Abstracts | 0.9245 | 1.4109 | 1.2603 | 0.8749 | 1.50 |
| ArXiv | 0.9524 | 1.5944 | 1.3221 | 0.8565 | 1.50 |
| FreeLaw | 0.9982 | 1.6474 | 1.4213 | 0.9632 | 1.51 |
| StackExchange | 1.0398 | 1.3558 | 1.3336 | 0.9205 | 1.50 |
| Wikipedia (en) | 1.0630 | 1.6001 | 1.3944 | 0.9902 | 1.50 |
| Pile-CC | 1.0947 | 1.5380 | 1.4215 | 1.1001 | 1.51 |
| BookCorpus2 | 1.1226 | 1.5307 | 1.4099 | 1.1426 | 0.61 |
| OpenWebText2 | 1.1322 | 1.6778 | 1.5765 | 1.1295 | 1.53 |
| PhilPapers | 1.2019 | 1.8543 | 1.9045 | 1.3358 | 0.16 |
| Enron Emails | 1.2454 | 1.8693 | 1.7662 | 1.2206 | 0.06 |
| Gutenberg (PG-19) | 1.2861 | 1.8397 | 1.6983 | 1.2715 | 0.53 |
| HackerNews | 1.2953 | 1.6748 | 1.6611 | 1.2806 | 0.24 |
| OpenSubtitles | 1.3275 | 1.7211 | 1.6327 | 1.3392 | 0.86 |
| Books3 | 1.3544 | 2.1966 | 1.9334 | 1.3300 | 1.51 |
| DM Mathematics | 1.4764 | 2.1898 | 1.8535 | 1.3201 | 0.81 |
| Ubuntu IRC | 1.5678 | 2.5754 | 2.5103 | 1.7147 | 0.04 |
| YoutubeSubtitles | 1.7768 | 2.3116 | 2.6923 | 1.8218 | 0.19 |
| EuroParl | 2.0889 | 2.6663 | 3.0108 | 1.9986 | 0.07 |
| All 22, byte-weighted | 1.0547 | 1.5795 | 1.4261 | 1.0015 | 20.40 |
⚠️ Limitations
- Size: At 220M parameters, the model's knowledge retrieval capacity is limited compared to 7B+ models.
- Base Model: This is a pre-trained base. It has not been fine-tuned for chat or instruction following (see Zenyx-DeepSeek-220M for the instruct version).
- Hallucinations: While logically consistent, it may generate factually incorrect statements.
📜 Citation
@misc{ZenyxBase220M,
title = {Zenyx-Base-220M: High-Density Foundation Model},
author = {Arko007},
year = {2025},
publisher = {HuggingFace},
url = {https://huggingface.co/Arko007/Zenyx_Base_220M}
}