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.

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.

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.

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

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

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.

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

Files

The intermediate checkpoints (a params/ and a checkpoints/ folder, one pair every 500 steps) were removed to keep the repository small; only the final step remains.

Usage

pip install jax flax safetensors transformers huggingface_hub

from huggingface_hub import hf_hub_download
from transformers import PreTrainedTokenizerFast

repo = "Arko007/zenyx-v2-base"
for name in ("modeling_zenyx_v2.py", "model.safetensors"):
    hf_hub_download(repo, name, local_dir=".")

from modeling_zenyx_v2 import build_model, load_params, generate

tok = PreTrainedTokenizerFast.from_pretrained("Arko007/zenyx-v2-tokenizer")
model, params = build_model(), load_params("model.safetensors")
print(generate(model, params, tok, "The derivative of f(x) = x^2 is", max_new=30))   # greedy; a base model, it continues the text

The module computes in bfloat16 (the stored float32 weights are cast on the fly), which runs on CPU, TPU and recent GPUs but not on a Tesla T4 (ALG_DOT_BF16_BF16_F32 is unsupported there; replace jnp.bfloat16 by jnp.float32 in the module to run on a T4). A forward pass returns a list of logits, one entry per prediction head (mtp_heads = 3); entry 0 is the usual next-token prediction.

Greedy continuations produced by this exact code (a base model, expect repetition and made-up facts):

The derivative of f(x) = x^2 is 1/x^2. The derivative of f(x) = x^2 is 1/x^2. The derivative of f

def fibonacci(n): number { let fib = 1; let fib1 = 1; let fib2 = 2; let fib3 =

Photosynthesis is a method of processing a text file. It is a simple way to generate a text file. It is a simple way to generate a text file.

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