Project Crucible
Axion Labs is running a 6-week R&D sprint (July 4 – August 17, 2026) to build "Crucible Axon" — an efficient small-model architecture for edge deployment and cost-sensitive inference. The architecture combines five recently proven techniques into a single, tightly integrated design.
TAME — Ternary Adaptive MoE
Nicknamed TAME, the architecture stacks five components designed to work together, not as isolated tricks.
Ternary Weights
1.58-bit BitLinear with straight-through estimation. Weights are ternary (-1, 0, +1), giving extreme memory savings while STE preserves a usable gradient signal through the hard quantization.
MoE (MoTE-style)
A single full-precision shared expert handles common patterns; multiple ternary routed experts specialize via top-1 gating. This keeps the compute profile close to a dense model of equivalent size while expanding capacity.
Multi-Token Prediction
An auxiliary head predicts t+2 alongside the standard t+1 head. The auxiliary loss provides a richer training signal per step without additional autoregressive decoding during inference.
Grouped Query Attention
8 query heads share 2 KV heads, cutting the KV cache footprint by 4× compared to full multi-head attention. This is critical for deployment on memory-constrained hardware.
Aux-Loss-Free Balancing
DeepSeek-V3-style per-expert bias that adjusts routing logits directly, with no auxiliary load-balancing loss term. This eliminates a known source of gradient interference in MoE training.
SubLN
Sub-layer normalization adapted from Microsoft's BitDistill research, added to stabilize training at the larger Axon-Chat scale where ternary quantization noise compounds across more layers.
Day 22 Sprint Status
Real progress, and one expensive mistake. The architecture holds up everywhere we've tested it at small scale — but the 600M run has to be thrown away and done again, because the learning rate was wrong.
14M Proof-of-Concept Trained
We trained a ~14M parameter TAME model on TinyStories for 10,000 steps on a free Colab T4. Despite the tiny scale, the model produces coherent multi-paragraph English with dialogue and narrative structure — promising signal for the architecture.
GQA + Aux-Loss-Free Variant Built
A second variant swapping in Grouped Query Attention and aux-loss-free load balancing was built and validated against the first variant at identical step count for a controlled head-to-head comparison. Result: equivalent quality with ~600K fewer parameters and a 4× smaller KV cache — confirming the efficiency techniques hold up in practice, not just in theory.
102M Trained, 300M Came Up Short
Crucible Axon-Chat (102M) trained on a mixed UltraChat + TinyChat dataset using the full validated stack — topically relevant responses, correct poem formatting, natural stopping points. Scaling the same recipe to 300M was a partial failure: enough of it worked to be worth diagnosing, not enough to keep.
Having fixed what caused it, we skipped retrying at 300M and went straight to 600M, built on Qwen 0.6B — the same base model Microsoft used for BitDistill, which makes the head-to-head below a fair fight instead of an apples-to-oranges one.
The Learning Rate Was 0.009 Too High — Run Scrapped
The training config had the learning rate set 0.009 above where it should have been. Written down on its own, 0.009 looks like a rounding error. As a learning rate it is nothing of the sort: they live down in the ten-thousandths, so an extra 0.009 doesn't nudge the value, it swamps it. The model was taking wildly oversized steps from the first batch onward and never had a chance to settle.
There's no partial credit to salvage. Everything downstream of that number is untrustworthy, and a run you can't trust is worse than no run at all — you end up stacking the next experiment on top of it. So the 600M model goes in the bin and we start it again.
Rerunning the 600M Model
Same architecture, same Qwen 0.6B base, same setup — corrected learning rate. The one upside of losing a run this way is that nothing about the design is in question, so there is nothing to redesign before starting over.
Head-to-Head Against Microsoft's BitDistill
Once the rerun finishes we're evaluating it against BitDistill*, Microsoft's pipeline for taking an off-the-shelf full-precision model down to 1.58-bit. The comparison is deliberately tight: BitDistill's published results start from Qwen 0.6B, so our 600M run starts from the same weights. Same base model, same target precision, two different routes there — which means whatever gap shows up is attributable to the method rather than to a difference in the models being compared.
Either result is worth having. If TAME comes out ahead, we have a reason to keep building it. If BitDistill does, we have a better starting point than the one we were about to build ourselves — which is a cheaper thing to learn now than at 300M of wasted compute.
What Goes Into Crucible Axon
Here's what each component looks like in practice.
Ternary BitLinear
Weights are quantized to ternary values during the forward pass: round(clamp(w / scale, -1, 1)) where scale is the per-channel absolute mean. Gradients pass through the round operation via straight-through estimation, ignoring the quantization threshold in the backward pass. No full-precision copies are kept — the master weights are stored in FP16 and quantized per step.
MoE Routing
Every MoE layer has one full-precision shared expert and 8 ternary routed experts. A learned router produces logits over the 8 routed experts; top-1 gating selects which expert to activate per token. The shared expert processes every token, acting as a stable backbone while routed experts handle specialization.
Aux-Loss-Free Balancing
Instead of adding an auxiliary loss to encourage balanced expert usage (which competes with the main LM loss gradient), we use per-expert bias terms added directly to the routing logits. After each step, if an expert receives fewer tokens than the expected uniform share, its bias is incremented; if it receives more, its bias is decremented. This is a purely feedback-based mechanism — no gradient competition, no tuning of loss coefficients.
Multi-Token Prediction
The model has two LM heads. The primary head predicts t+1 as usual. A secondary head, fed from an intermediate layer, predicts t+2. Both heads contribute to the training loss with a fixed ratio. At inference time, only the primary head is used — the auxiliary head is discarded, so there's no runtime cost.
Grouped Query Attention
With 8 query heads and 2 KV heads, grouped query attention projects each KV head to serve 4 query heads. The KV cache is 4× smaller than full multi-head attention (8 KV heads), and the attention computation is proportionally faster. Positional encoding via RoPE (rotary position embeddings).
Model Configuration
| Property | PoC (14M) | From scratch (300M) |
|---|---|---|
| Parameters | ~14.5M | ~300M |
| Layers | 8 | 24 |
| Hidden dim | 256 | 1024 |
| Attention heads | 8 query / 2 KV | 8 query / 2 KV |
| Expert count | 8 (ternary) + 1 (shared) | 8 (ternary) + 1 (shared) |
| Training data | TinyStories (~10K steps) | TBD (FineWeb-Edu + curated) |
| Training hardware | Colab T4 (free tier) | TBD |
The 300M column is the from-scratch design that came up short. The current 600M run starts from Qwen 0.6B instead, so it inherits that model's configuration rather than this one — its specs will go up here once the rerun finishes.
* BitDistill (BitNet Distillation) is Microsoft's pipeline for converting an off-the-shelf full-precision language model into 1.58-bit ternary weights for a given task, instead of training a ternary model from scratch. It combines three pieces: the SubLN normalisation module from BitNet, multi-head attention distillation from the full-precision teacher, and a continual pre-training warm-up that stops the quality gap between the full-precision and ternary versions widening as models get larger. Microsoft reports quality comparable to the full-precision original with up to 10× memory savings and roughly 2.65× faster CPU inference. Paper: arXiv:2510.13998 · code: github.com/microsoft/BitNet.