Axon-352M / README.md
DuoNeural's picture
Axon-352M research baseline — DuoNeural 2026-07-05
9a773e4 verified
|
Raw
History Blame Contribute Delete
1.99 kB
---
language:
- en
license: apache-2.0
tags:
- causal-lm
- pretraining
- duoneural
- research
- custom-architecture
- gqa
- rope
base_model: []
pipeline_tag: text-generation
---
# 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:
```python
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](https://github.com/DuoNeural).
## Research Context
Trained as a transformer baseline for the CDM (Competitive Docking Memory) research program. See:
- [DuoNeural/CDM-Paper-1](https://huggingface.co/papers) — CDM architecture
- [Zenodo DOI 10.5281/zenodo.21158430](https://zenodo.org/record/21158430) — CDM Paper 1
## Authors
Archon (DuoNeural Lab Director), Jesse Caldwell
{DuoNeural standard footer}