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

📚 Training Curriculum (The "Omni-Mix")

The model was trained using a rigorous 4-stage curriculum designed to layer capabilities sequentially:

1. Phase 1: Fundamentals (FineWeb-Edu)
   • Focus on high-quality educational English text to establish linguistic baselines.
2. Phase 2: Logic & Structure (StarCoder - Python)
   • Introduction of code data to enforce logical indentation, syntax, and structured thinking.
3. Phase 3: Multilingualism (FineWeb-2)
   • Exposure to 6 major languages (Hindi, Chinese, Russian, Japanese, French, Spanish) to expand the semantic embedding space.
4. 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

The weights are model.safetensors (JAX/Flax parameter tree, flat names such as embedding and blocks_0.attn...); modeling_zenyx_220m.py in this repository holds the model code (the training-notebook definition, unchanged) and two helpers: load_model downloads the weights and loads them strictly (every parameter must exist with the right shape), generate decodes text. Runs on CPU, GPU (including a T4: it computes in float32) and TPU.

pip install jax flax safetensors transformers huggingface_hub

from huggingface_hub import hf_hub_download
hf_hub_download("Arko007/Zenyx_Base_220M", "modeling_zenyx_220m.py", local_dir=".")      # the model code

from modeling_zenyx_220m import load_model, generate
from transformers import AutoTokenizer

model, params = load_model()                                              # downloads model.safetensors (881 MB)
tok = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")         # the tokenizer of this model

print(generate(model, params, tok, "The capital of France is", max_new=30))                 # greedy
print(generate(model, params, tok, "def fibonacci(n):", max_new=40, temperature=0.8))      # sampling

It is a base model: give it text to continue, not instructions. Greedy decoding tends to repeat itself (the test run printed "The capital of France is located in the center of the country, and the capital of the country is located in the center of the country ..."); use temperature=0.7-0.9 for more varied text. generate re-runs the whole sequence at every step, which is fine for short outputs but slow for long ones.

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.

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

SciQ (n = 1000, 0-shot, chance 25%, standard error about 1.5 points)

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.

⚠️ 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}
}