axiom1 1.0.6

AXIOM-1: The first EGen Core language model by EGen Labs / ErebusTN

AXIOM-1 | EGen-Core/AXIOM-1

EGen Labs | ErebusTN

AXIOM-1 is a 130M-parameter recurrent-depth language model built on the LoopLLM architecture. A single shared transformer block is applied 12 times to generate text.

Note: This model is built from scratch by editing the existent AXIOM-1 codebase. It serves as a prototype and testing ground for recurrent-depth models.

Quick Reference

• Parameters: ~130M (with weight tying)
• Size: d_model=512, heads=8q/2kv, d_ff=1365, loops=12
• Precision: bfloat16 ONLY (never float16)
• Vocab: 32000 (LLaMA-compatible tokenizer)

Installation (PyPI Ready)

You can install AXIOM-1 as a python package directly from the source repository:

pip install .

Once published to PyPI, you can install it via:

pip install axiom1

Running the Google Colab Prototype

We have provided a ready-to-use Google Colab notebook: axiom1_colab_training.ipynb

Features included in the notebook:

1. Installs the axiom1 package and all dependencies directly.
2. Loads the official dataset: ErebusTN/axiom1 using the datasets library.
3. Instantiates the AXIOM-1 model and runs the mandatory 5-step sanity check.
4. Trains the prototype using PyTorch's native Trainer or loop.
5. Exports the final model weights and pushes the checkpoint directly to the Hugging Face Hub.

How to use:

1. Upload axiom1_colab_training.ipynb to Google Colab.
2. Set your runtime to T4 GPU or higher.
3. Run all cells. (You will be prompted to log in to Hugging Face to push your model).

Manual Usage

Sanity Check (Required Before Training)

from axiom1.model.config import AXIOM1Config
from axiom1.model.model import AXIOM1
from axiom1.evals.sanity_check import run_sanity_check

cfg = AXIOM1Config()
model = AXIOM1(cfg)
run_sanity_check(model, cfg)

Inference

import torch
from axiom1.model.config import AXIOM1Config
from axiom1.model.model import AXIOM1
from axiom1.inference.generate import generate

cfg = AXIOM1Config()
model = AXIOM1(cfg).to('cuda')
prompt_ids = torch.tensor([[1, 234, 555]], device='cuda')
ids = generate(model, prompt_ids, cfg, max_new_tokens=256)

---

"One block. Twelve iterations. The first axiom of EGen Core."