AXIOM

An experimental neural architecture by Lee Madover, built and measured end to end on one consumer GPU (RTX 5070, 12 GB).

AXIOM replaces two standard parts of a language model:

• Attention becomes a recurrent working memory: a gated linear recurrence, so each new token costs O(1).
• The feed-forward layer becomes a bank of routed experts called Law Cells. These can grow during training by fission, but only when a paired, held-out trial shows that a split pays for its complexity.

This repository holds the architecture, the training pipeline, every experiment, and honest results, including the negative ones.

Tested with Python 3.13 and PyTorch 2.11 + CUDA 12.8 on an RTX 5070. The RTX 50 series needs a cu128 or newer PyTorch build.

pip install -r requirements.txt

python -c "from huggingface_hub import snapshot_download; snapshot_download('Atlas-AI-research/Axiom-1-Chat', local_dir='hf/axiom1-chat')"

Then chat in a terminal, with tokens streamed as they're generated:

python chat.py --model hf/axiom1-chat --device cuda

Or in Python:

from axiom.hub import load
model, tok, meta = load("Atlas-AI-research/Axiom-1-Chat", device="cuda")

All numbers are measured; the scripts and raw outputs are in results/. The reference is SmolLM2-135M-Instruct, a mature open model about 2× larger that was pretrained on about 650× more text.

What the experiments found:

• The recurrent working memory works. It trains stably and generates faster than a same-size Transformer at batch 1.
• Fission did not earn its keep.
  • On synthetic data, a fixed mixture of experts matched or beat it, and its evidence-derived split direction was no better than a random one.
  • In Axiom-1, all 9 trial splits were rejected.
• The weaknesses are mostly data scale and shape: 3.1B pretraining tokens, long chat answers in the tuning data, and no examples of declining to answer.

Full design, maths, failure analysis and prior art: DESIGN.md. Chat-quality research: rd_coherence.md.

The data is about 9.6 GB: four FineWeb-Edu shards plus the two chat datasets. Download them into data/ as prepare_data.py describes, then:

python -I prepare_data.py

powershell -ExecutionPolicy Bypass -File run_axiom1.ps1

run_axiom1.ps1 runs four steps:

1. Tokenizes the data.
2. Runs a pilot of two memory designs.
3. Pretrains for about 22 hours, with fission rounds.
4. Chat-tunes for about 2 hours.

It's resumable and keeps the PC awake. Supporting tools:

python hwcheck.py

python train.py --config configs/axiom0.yaml --out checkpoints/axiom0

python benchmark.py --config configs/axiom0.yaml --seeds 3

python benchmark.py --config configs/axiom0.yaml --seeds 3 --variants axiom,axiom_no_gate,moe_fixed,moe_k2,moe_k8,moe_k12 --tag wrongk

python -m pytest tests -q

Override any config value with --set key=value.

Code and weights are released under the Apache-2.0 license (see LICENSE). Training data: FineWeb-Edu (ODC-By 1.0), smol-smoltalk (Apache-2.0) and everyday-conversations-llama3.1-2k (Apache-2.0), all by Hugging Face.