Zenyx-v2-base (Nano-Titan Architecture)
Zenyx-v2-base is a highly optimized, 85-million parameter language model built on the custom Nano-Titan architecture. Engineered entirely in JAX/Flax for pure BF16 execution on TPU v5e-8 hardware, this model is designed to push the boundaries of parameter efficiency. The pretraining corpus is heavily skewed towards reasoning, utilizing a meticulously balanced stream of mathematics, programming languages, and high-quality educational text targeting a 200 billion token lifecycle.
Architectural Specifications
The Nano-Titan architecture employs parameter sharing and advanced attention mechanisms to achieve an effective depth far exceeding its raw parameter count.
| Component | Specification |
|---|---|
| Total Parameters | ~85M (Unique weights) |
| Effective Depth | 32 Layers (8 unique blocks × 4 recurrences) |
| Embedding Dimension ($D_{model}$) | 576 |
| Attention Mechanism | Multi-Head Latent Attention (MLA) |
| Heads | 9 Query Heads, 3 KV Heads |
| MLP Architecture | ConvSwiGLU (Hidden Dim: 1536) |
| Context Window | 8192 Tokens (YaRN-scaled RoPE) |
| Vocabulary Size | 32,768 |
Core Innovations
Multi-Token Prediction (MTP) The model leverages a 3-head Multi-Token Prediction objective during pretraining. By predicting multiple future tokens simultaneously with decaying loss weights (1.0, 0.3, 0.1), the network develops stronger long-term contextual representations and improves sample efficiency without increasing inference latency.
Multi-Head Latent Attention (MLA) Memory overhead during training and inference is drastically reduced via MLA. The key-value cache is compressed into a 128-dimensional latent space, while queries are projected through a 384-dimensional latent bottleneck. This structural optimization allows for extended context windows within constrained hardware environments.
ConvSwiGLU Feedforward Networks Standard feedforward blocks are replaced with Convolutional SwiGLU networks. A 1D convolution (kernel size 3) is applied to the gating mechanism prior to the SiLU activation, imparting a localized inductive bias that enhances the processing of structured data formats like code and mathematics.
Recurrent Depth via Block Sharing
To maximize parameter utility, the network consists of 8 distinct TitanBlock modules that are sequentially recurred 4 times. This yields an effective representational depth of 32 layers while maintaining an 85M parameter memory footprint.
YaRN-Scaled Rotary Position Embeddings Extrapolation to an 8192-token context is stabilized using YaRN (Yet another RoPE extensioN). The scaling factor is set to 32.0 with dynamic wavelength interpolation, ensuring uniform attention distribution across both local and extended sequence lengths.
Training Data & Optimization
The dataset pipeline is fully streaming and dynamic, reading sharded Parquet files directly from the Hugging Face Hub to prevent local disk bottlenecks.
| Domain | Mixture Weight | Source Datasets | Filtration Criteria |
|---|---|---|---|
| Mathematics | 45% | finemath-4plus, infiwebmath |
Statically routed across 3+ and 4+ streams |
| Code | 35% | starcoderdata |
24 strictly filtered languages, min length 20 chars |
| English / Edu | 20% | fineweb-edu |
Educational score $\ge$ 3.0 |
Training utilizes a Warmup-Stable-Decay (WSD) learning rate schedule, peaking at 3e-4 before executing a cosine decay to 3e-5 over the final 18,000 steps. Global gradient accumulation is managed via chunked cross-entropy to strictly cap per-chip HBM usage at ~12.2 GB.
Evaluation Metrics (Checkpoint: Step 96,000)
The model is currently evaluated against a dedicated high-quality educational validation split encompassing 128 micro-batches.
| Metric | Value |
|---|---|
| Best Validation Loss | 3.7206 |
| Recent Training Loss (Avg) | ~2.10 |
| Throughput | ~200,000 tokens/sec |
| Tokens Processed | ~200 Billion |
Benchmarks (checkpoint step 96,000)
All tasks use the standard base-model protocol: the model scores the log-likelihood of each candidate continuation and the
highest-scoring one is the answer, so nothing depends on instruction following. 0-shot, full evaluation sets (MMLU 5-shot).
acc_norm normalises by continuation length in characters. These are the project's own implementation of the
lm-evaluation-harness conventions (modal_base_evals.py), not a harness run, and they are computed with the module's native bf16 compute.
Bold = the conventional metric for the task. Random is the chance level (BoolQ: its 62.2% majority class).
| Benchmark | acc |
acc_norm |
Random | Δ vs random | n |
|---|---|---|---|---|---|
| HellaSwag | 26.36% ± 0.44 | 27.81% ± 0.45 | 25.0% | +2.8 | 10,042 |
| ARC-Easy | 37.16% ± 0.99 | 34.13% ± 0.97 | 25.0% | +9.1 | 2,376 |
| ARC-Challenge | 18.86% ± 1.14 | 23.98% ± 1.25 | 25.0% | -1.0 | 1,172 |
| PIQA | 55.39% ± 1.16 | 54.79% ± 1.16 | 50.0% | +4.8 | 1,838 |
| WinoGrande | 49.96% ± 1.40 | — | 50.0% | -0.0 | 1,267 |
| OpenBookQA | 13.80% ± 1.54 | 23.60% ± 1.90 | 25.0% | -1.4 | 500 |
| BoolQ | 47.43% ± 0.87 | 59.91% ± 0.86 | 62.2% | -14.8 | 3,270 |
| SciQ | 51.90% ± 1.58 | 47.20% ± 1.58 | 25.0% | +26.9 | 1,000 |
| LAMBADA (OpenAI) | 3.43% ± 0.25 | — | 0.0% | +3.4 | 5,153 |
| MMLU (5-shot) | 24.61% ± 0.36 | — | 25.0% | -0.4 | 14,042 |
| RACE | 22.38% ± 0.59 | 26.04% ± 0.62 | 25.0% | +1.0 | 4,934 |
| CommonsenseQA | 19.08% ± 1.12 | 26.04% ± 1.26 | 20.0% | +6.0 | 1,221 |
| COPA | 56.00% ± 4.96 | 57.00% ± 4.95 | 50.0% | +6.0 | 100 |
| LogiQA | 20.89% ± 1.59 | 24.73% ± 1.69 | 25.0% | -0.3 | 651 |
| WSC273 | 50.92% ± 3.03 | — | 50.0% | +0.9 | 273 |
| TruthfulQA MC1 | 26.44% ± 1.54 | — | 22.8% | +3.6 | 817 |
| Arithmetic | 0.36% ± 0.05 | — | 0.0% | +0.4 | 14,000 |
At 85M unique parameters (32 effective layers) the model is more than two standard errors above chance on: SciQ 51.90%, LAMBADA (OpenAI) 3.43%, ARC-Easy 34.13%, HellaSwag 27.81%, CommonsenseQA 26.04%, PIQA 54.79%. It is within two standard errors of chance (or below it) on: ARC-Challenge, WinoGrande, OpenBookQA, BoolQ, MMLU (5-shot), RACE, COPA, LogiQA, WSC273. (TruthfulQA is left out of both lists: it penalises stronger language models, so it is not a capability signal.) Arithmetic (0.36%, 14,000 problems) is generated in-harness from a fixed seed and is not interchangeable with published EleutherAI/arithmetic numbers.
Language modelling. WikiText-2 (raw), rolling 1024-token window at stride 512: 1.7676 bits/byte (token perplexity 116.47). LAMBADA: accuracy 3.43%, perplexity 174.70. Perplexities are per-token and this model's 32,768-entry tokenizer is not comparable to a larger vocabulary's, so bits/byte is the number to compare across models.
The Pile, by domain (bits/byte, LOWER is better). The text is NeelNanda/pile-10k (pinned revision 127bfedcd5), documents grouped by
pile_set_name, joined with blank lines, capped at 1.5M characters per domain, scored with a rolling 1024-token window at stride 512 so every
token is scored exactly once with at least 512 tokens of left context. Bits/byte is normalised by the UTF-8 bytes of the source text.
The first column is the project harness run above; the second is an independent re-run (T4, fp32 compute, a separate implementation of
the same protocol): identical token counts in all 22 domains and a largest difference of 0.0008 bits/byte. The last two columns score the
sibling models on the identical text (T4, fp16; V3 is a 1.58B-parameter MoE trained on ~81B tokens, MCEN is a 565M model still training and only 0.54B tokens in).
| Domain | V2 (harness) | V2 (T4 fp32 re-run) | V3-Base @107,200 | MCEN @4,100 | text (MB) |
|---|---|---|---|---|---|
| Github | 1.0833 | 1.0832 | 0.5765 | 1.0490 | 1.53 |
| USPTO Backgrounds | 1.2532 | 1.2530 | 0.7807 | 1.0945 | 1.50 |
| NIH ExPorter | 1.3250 | 1.3249 | 0.8620 | 1.1771 | 0.22 |
| PubMed Central | 1.3521 | 1.3518 | 0.7426 | 1.1163 | 1.51 |
| StackExchange | 1.3561 | 1.3558 | 0.9205 | 1.3336 | 1.50 |
| PubMed Abstracts | 1.4112 | 1.4109 | 0.8749 | 1.2603 | 1.50 |
| BookCorpus2 | 1.5308 | 1.5307 | 1.1426 | 1.4099 | 0.61 |
| Pile-CC | 1.5381 | 1.5380 | 1.1001 | 1.4215 | 1.51 |
| ArXiv | 1.5952 | 1.5944 | 0.8565 | 1.3221 | 1.50 |
| Wikipedia (en) | 1.6002 | 1.6001 | 0.9902 | 1.3944 | 1.50 |
| FreeLaw | 1.6477 | 1.6474 | 0.9632 | 1.4213 | 1.51 |
| HackerNews | 1.6751 | 1.6748 | 1.2806 | 1.6611 | 0.24 |
| OpenWebText2 | 1.6779 | 1.6778 | 1.1295 | 1.5765 | 1.53 |
| OpenSubtitles | 1.7208 | 1.7211 | 1.3392 | 1.6327 | 0.86 |
| Gutenberg (PG-19) | 1.8402 | 1.8397 | 1.2715 | 1.6983 | 0.53 |
| PhilPapers | 1.8544 | 1.8543 | 1.3358 | 1.9045 | 0.16 |
| Enron Emails | 1.8697 | 1.8693 | 1.2206 | 1.7662 | 0.06 |
| DM Mathematics | 2.1898 | 2.1898 | 1.3201 | 1.8535 | 0.81 |
| Books3 | 2.1967 | 2.1966 | 1.3300 | 1.9334 | 1.51 |
| YoutubeSubtitles | 2.3117 | 2.3116 | 1.8218 | 2.6923 | 0.19 |
| Ubuntu IRC | 2.5762 | 2.5754 | 1.7147 | 2.5103 | 0.04 |
| EuroParl | 2.6664 | 2.6663 | 1.9986 | 3.0108 | 0.07 |
| All 22, byte-weighted | 1.5797 | 1.5795 | 1.0015 | 1.4261 | 20.40 |
V2 is ahead of MCEN@4,100 in 3 of 22 domains and behind V3-Base@107,200 in 22 of 22. The thin domains (EuroParl, Ubuntu IRC, YoutubeSubtitles, Enron, PhilPapers, NIH ExPorter, HackerNews; each under ~250 KB of text) are noisier.
SciQ cross-check. An independent T4/fp32 re-run of SciQ gives acc 52.10% / acc_norm 47.20%, against 51.90% / 47.20% from the harness.
For reference, the same test gives 55.30% / 49.80% for MCEN@4,100: 3.2 points apart on acc, with a standard error of about 1.6 points each, i.e. 1.4σ -- not a significant difference.
Zenyx-Fact1000. An internal, frozen 1000-question, four-way factual / common-sense completion benchmark (geography, science, history,
language, arithmetic, computing; chance 25%; the model must rank the correct completion above three same-type distractors by log-likelihood;
file sha256 843416524fb39c5f…; not yet published). A "prior-only" control scores the same candidates after the neutral context "The answer is":
a model whose score sits near the control is doing little more than preferring common words.
| Model | accuracy | 95% interval | without arithmetic | prior-only control |
|---|---|---|---|---|
| Zenyx-v2-base (step 96,000) | 45.8% | 42.7–48.9% | 48.1% | 30.7% |
| Zenyx-V3-Base (step 107,200, 1.58B MoE) | 81.1% | 78.6–83.4% | 80.8% | 28.7% |
| Zenyx-MCEN-Delta-Mamba (step 4,100, 565M) | 45.5% | 42.4–48.6% | 45.9% | 29.6% |
| Zenyx_Base_220M (Qwen tokenizer, ~153B tokens per its card) | 65.3% | 62.3–68.2% | 73.2% | 29.8% |
| Zenyx-42M (FineWeb-Edu only, 6.5B tokens) | 49.5% | 46.4–52.6% | 54.4% | 29.1% |
V2 and MCEN@4,100 are statistically indistinguishable overall (paired exact McNemar test on the 1000 items: V2 alone right 129, MCEN alone right 126, p = 0.9), and they answer different questions correctly.
Running on a Kaggle T4. The module hard-codes bfloat16 matmuls, which a T4 cannot execute (UNIMPLEMENTED: ALG_DOT_BF16_BF16_F32). The
T4 runs above use the identical module source with jnp.bfloat16 replaced by jnp.float32 (and, as a check, float16: 999 of 1000 Fact1000 answers were unchanged).
Implementation & Usage
Until the weights are explicitly converted to PyTorch safetensors and mapped to AutoModelForCausalLM, the model must be instantiated using the native JAX/Flax module defined in the training environment.
import jax
import jax.numpy as jnp
from flax import serialization
from transformers import PreTrainedTokenizerFast
from zenyx_model import ZenyxV2 # Import your model class here
# 1. Initialize Tokenizer
tokenizer = PreTrainedTokenizerFast.from_pretrained("Arko007/zenyx-v2-tokenizer")
# 2. Instantiate Architecture
model = ZenyxV2(
vocab_size=32768,
d_model=576,
n_heads=9,
n_kv_heads=3,
head_dim=64,
hidden_dim=1536,
n_unique_blocks=8,
n_recurrences=4,
max_seq_len=8192,
kv_latent=128,
q_latent=384,
mtp_heads=3,
dropout_rate=0.0
)
# 3. Load Checkpoint Weights
# Note: Ensure params_step96000.msgpack is downloaded from the Hub
with open("params_step96000.msgpack", "rb") as f:
params_bytes = f.read()
dummy_input = jnp.ones((1, 8192), dtype=jnp.int32)
rng = jax.random.PRNGKey(0)
variables = model.init(rng, input_ids=dummy_input, train=False)
loaded_params = serialization.from_bytes(variables["params"], params_bytes)
# 4. Forward Pass (Inference)
input_text = "The derivative of f(x) = x^2 is"
input_ids = tokenizer(input_text, return_tensors="jax")["input_ids"]
# Model returns a list of MTP logits; index 0 is the primary next-token prediction
logits = model.apply({"params": loaded_params}, input_ids=input_ids, train=False)[0]
next_token = jnp.argmax(logits[0, -1, :])
print(tokenizer.decode([int(next_token)]))
Limitations and Biases
Zenyx-v2-base functions strictly as a foundational pre-trained model. It has not been subjected to instruction fine-tuning, Reinforcement Learning from Human Feedback (RLHF), or Direct Preference Optimization (DPO). As a direct result of operating exclusively on next-token and multi-token prediction objectives, the model does not possess conversational alignment. Outputs may exhibit hallucinations, lack structural coherence in zero-shot dialogue scenarios, or reflect inherent biases present within the constituent datasets (FineMath, StarCoderData, and FineWeb-Edu). Downstream utilization requires domain-specific fine-tuning and the implementation of robust safety guardrails prior to production deployment.
Hardware and Environmental Setup
The pre-training lifecycle was executed on Google TPU v5e-8 topology utilizing pure bfloat16 (BF16) arithmetic. The distributed execution framework leverages jax.pmap for cross-device data parallelism. The software stack mandates JAX $\ge$ 0.4.16, optimized via libtpu, with gradient transformations managed through Optax. The asynchronous data pipeline is heavily optimized for high-throughput streaming directly from Hugging Face Parquet shards to circumvent local disk I/O bottlenecks.
License
The Zenyx-v2-base model weights, architecture code, and associated tokenizer are distributed under the Apache License 2.0. You are free to use, modify, distribute, and utilize this model for commercial applications, provided you adhere to the attribution and liability limitations specified within the license framework.
Citation
If you utilize the Zenyx-v2-base model, the Nano-Titan architecture, or the associated multi-token prediction training methodologies in your research or production environments, please cite the repository using the following BibTeX entry:
@misc{zenyx2026,
author = {Anamitra},
title = {Zenyx-v2-base: A Parameter-Efficient Nano-Titan Language Model},
year = {2026},
publisher = {Hugging Face},
howpublished = {\url{https://huggingface.co/Arko007/zenyx-v2-base}},
note = {Pre-trained on TPU v5e-8 using JAX/Flax}
}