Aria

1.58-bit ternary LLM in Rust/Burn. Byte-level/BPE tokenizer, spectral compact factors, native sparse attention.

Training proven at 29M params: stable convergence, 346 tok/s on single GPU, no OOM. Full training stack with HymOpt (Muon+LOTUS), GDN2 delta-rule, NSA, and ATH.

Architecture

Input tokens (vocab=16384 BPE or byte-level 256)
  ↓
Embed (one_hot @ matmul)
  ↓
AnchorBlock — 2× NSA + WaveletKAN
  ↓
HelixCore — N-loop recurrent reasoning
  ├── GDN2Attention (chunked associative scan, delta-rule)
  ├── WaveletKANFFN
  ├── GRU gate
  ├── Value head (confidence)
  └── ATH (adaptive token halting)
  ↓
NSA cross-attention (compression + selection + window)
  ↓
SynthesisBlock — MTP-4 heads
  ↓
FormatHead — format token prediction
  ↓
Logits

Attention

• NSA — compression (block avg + MLP, ~32×), fine-grained selection (top-n), sliding window, gated mixing
• GDN2 (Helix core) — chunked associative scan with delta-rule erase; fused CubeCL kernel available
• No KV cache — NSA recomputes from compressed blocks

Parameters

• SCTLinear: u[s, r] * s[r] * v[d_out, r] — never materializes dense matrix
• Ternary forward pass: STE {-1, 0, +1} with learnable scale
• Hadamard transform on o_proj (BitNet v2)

Optimizer: HymOpt

• Muon (Newton-Schulz 5-step quintic + LOTUS low-rank preconditioner) for 2D weights
• AdamW for 1D / embed
• Spectral gradient clipping + column normalization
• Gradient clipping (clip_norm=1.0) and spectral clipping (max_sigma=1.0) for stability
• EMA shadow weights (BF16)

Requirements

• CUDA-capable GPU (tested on A100 with 11GB)
• Rust edition 2021

Build

# Default CUDA backend
cargo build --release

# With CubeCL kernels (fused GDN2 scan)
cargo build --release --features cubecl

Usage

Train

# Generate training data first
cargo run --release --bin data-download -- --dir data --count 1024

# Pretrain (mode 0)
cargo run --release --bin aria-burn configs/29m.yaml

# SFT (mode 1)
cargo run --release --bin aria-burn configs/test_sft.yaml

# GRPO (mode 2) — needs data_grpo/
cargo run --release --bin aria-burn configs/test_grpo.yaml

Generate

cargo run --release --bin generate -- configs/29m.yaml "Hello" 256 0.7 0.9 20 0.05 false

Profile

cargo run --release --bin profile -- configs/29m.yaml

Tests

cargo test --test smoke_test -- --nocapture

Configuration (YAML)

See configs/29m.yaml (29M params) and configs/test_quick.yaml (smol) for reference.

Field	Default	Description
d_model	1536	Hidden dimension
n_heads	16	Attention heads
n_loops	6	Helix loops (training), up to 48 at inference via ATH
sct_rank	32	SCT low-rank dimension
clip_norm	1.0	Global gradient clipping
max_sigma	1.0	Spectral gradient clipping
schedule_free	false	Schedule-Free wrapper (off by default — interferes with clip_norm)
lr_muon	0.0005	Learning rate for Muon path
lr_adamw	0.0001	Learning rate for AdamW path
wd_muon	0.2	Weight decay for Muon path
gdn_chunk_size	8	GDN2 scan chunk (truncated BPTT window)
ath	true	Adaptive token halting
lcsb_ratio	0.25	Loop-wise curriculum (mask ratio)
forward_ff	true	Forward-Forward local loss
qac_bits	0	Activation quantization: 0=off, 8=FP8

Training Status

29.1M model — verified stable:

• 1000 steps: loss 11.6 (random) → 6.66 (learning)
• 346 tok/s on A100
• ~500 MB GPU memory
• No NaN, no divergence, no OOM
• WSD schedule with 50-step warmup, decay from step 800

CubeCL fused kernels for GDN2 delta-rule scan available (--features cubecl), limited by CubeCL 0.10 cross-thread sync constraints — elementwise erase instead of full matmul.

License

MIT — go wild. Companies: I'd appreciate a heads-up, that's all.