Axon-352M

DuoNeural Research | 2026-07-05 | Archon

Research baseline model trained on smollm-corpus as a counterpart to CDM-based architectures. Used to establish a transformer baseline for comparative CDM studies.

Architecture

Custom GPT-style transformer (not a standard HuggingFace architecture):

Parameter Value
Layers 30
Hidden dim 1024
FFN dim 2560
Attention heads 8Q / 4KV (GQA)
Head dim 128
Vocab size 49,152 (SmolLM tokenizer)
Max seq length 2,048
Activation ReLU²
Normalization RMSNorm + QK-norm
Position encoding RoPE (θ=10000)
Logit cap 30.0
Total params ~352M

Training

  • Data: smollm-corpus (FineWeb-edu-dedup 50%, Cosmopedia-v2 30%, OpenWebMath 10%, Python-edu 10%)
  • Tokens: ~8.5B
  • Optimizer: MuonH (matrix params) + AdamW (embeddings)
  • Peak LR: 3e-4 (trapezoidal: 5% warmup, 85% stable, 10% decay)
  • Hardware: RTX 3090 (1×) on vast.ai

Loading

This model uses a custom architecture not directly loadable via AutoModel. To load:

import torch
from safetensors.torch import load_file

# Load state dict
state_dict = load_file("model.safetensors")

# Architecture must be defined from training script
# See train_axon_300m.py for the full model class

Full training script and architecture code available at DuoNeural GitHub.

Research Context

Trained as a transformer baseline for the CDM (Competitive Docking Memory) research program. See:

Authors

Archon (DuoNeural Lab Director), Jesse Caldwell

{DuoNeural standard footer}

Downloads last month
231
Safetensors
Model size
0.4B params
Tensor type
BF16
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support