Exea LabsBlog
日本語Japanese version

Day 2 ·

Pure TONN lost to Llama. The hybrid wins.

Two attention layers for every TONN layer. It beats a same-size Llama at equal compute, up to 0.1B.

Date
2026-11-11
Read
5 min read
Topics
Architecture, Speed, Scaling
By
The Exea team
1.464

bits per byte, ExeaLM at 0.1B, same training time as Llama (1.481)

1.01×

Llama's training cost per token

49

training runs in the sweep, on one Mac

The first fair test went against us. We trained a 13M TONN and a same-size on the same data for the same steps. They ended close on prose and far apart on code. This post follows the 49 runs that turned that into a win. It also covers the speed work that made the win hold at equal compute.

All numbers are on held-out English, Japanese Wikipedia, Japanese web text and code. Lower is better. Every model uses the same 32k tokenizer and the same data mix: 50% English, 30% Japanese, 20% code.

01Pure TONN against Llama

At 13M parameters, 3,000 steps, two seeds each:

ModelEnglishJapanese wikiJapanese webCode
Llama1.4671.5401.4201.789
TONN, with the search1.4761.5641.4201.970
TONN1.4801.5681.4221.979

On prose TONN is close: tied on Japanese web text, 0.009 behind on English. The deficit is almost all code, about 0.18 bits per byte. Code rewards exact copying of names, which attention does well. Llama has full attention in every layer; a TONN layer with a random start holds, on average, less than one small attention circuit.

02The sweep

That evening a queue of runs started on the Mac, one idea at a time, each against the same baseline. It ran until the next evening.

03:0006:0009:0012:0015:0018:00Hybrid layers, random startConvolution-heavy startLearning rates, longer runsBigger models, equal stepsEqual computeLayoutsRetrieval and swarms on the hybridSpeed
Every run in the sweep, on one Mac Studio GPU, from 1:41 in the morning on November 11. Each bar is one training run, in the lane of the question it answered; orange marks the equal-compute runs. Point at a bar for its result. The runs on the right are larger models, which take longer.

What we tried, in order:

  1. Hybrid with a random start. TONN layers alternate with full attention layers (attention at layers 1, 3 and 5). Seed 0 beat Llama. Seed 1 did not: close on prose, 1.950 on code. The seed also picks the TONN layers' starting circuits. Seed 0 drew 7 short convolutions. Seed 1 drew only 3, plus per-position gated units, which duplicate what the MLP already does.
  2. The search helped seed 1 a little (1.594 to 1.585) but did not fix it.
  3. More attention circuits inside pure TONN cut the code gap from 0.17 to 0.08 and stayed behind Llama. Full attention layers do better than many small attention circuits.
  4. A fixed seed 1. Code went from 1.950 to 1.748.

The hybrid with the convolution-heavy start beat Llama on every kind of text, with both seeds:

ModelSeedEnglishJapanese wikiJapanese webCodeMean
Hybrid01.4521.5231.4061.7451.532
Hybrid11.4511.5191.4021.7481.530
Llama01.4611.5341.4131.7751.546
Llama11.4731.5451.4261.8031.562

Its two seeds land within 0.002 of each other. With the retrieval and board circuits kept in, as in the real model, it still wins on every kind of text. The idle memory circuits cost 0.010.

03Bigger models, same steps

At 27M parameters (width 384, 8 layers, two seeds) the hybrid's lead doubled, from 0.023 to 0.045. Llama does better at this size with a lower learning rate, so we checked each model at its best rate. The hybrid still leads by 0.031, on every kind of text. At 48M, with the same 3,000 steps, it leads by 0.013.

04Equal compute: the speed problem

Winning per step is half the test. A model that is better per step but slower can lose when both get the same training time. That is the comparison. It decides whether an architecture is worth using.

The hybrid was 1.3 times slower per step than Llama on the Mac. Attention runs as one fused kernel; TONN's circuits run as plain PyTorch. Profiling one layer showed where the time went:

One layer, forward and backward (ms)
GI127 → 25 ms
SM145 → 45 ms
beforeafter the prefix scanattention block 18 ms
Training cost per token, relative to Llama (0.1B)
Pure TONN1.54× / 2.66×
Hybrid, half attention1.13× / 1.66×
Hybrid 2:1 (default)1.01× / 1.49×
memory circuits idleretrieval onLlama = 1
Where the speed came from. Left: time for one layer's forward and backward pass by circuit type, width 256, before and after the parallel prefix scan; an attention block takes about 18 ms. Right: training throughput of the full 0.1B model on the Mac, relative to a same-size Llama. Lower cost per token is better.
  • The two gated recurrences, GI and SM, took 127 and 145 ms because they ran one step at a time. A computes the same running sums in a logarithmic number of passes: 5 times and 3.3 times faster.
  • One FFT instead of two for the softmax-pooling circuit.
  • Cheaper rotary positions.
  • The TONN layer now permutes its small weight matrices into circuit-type order instead of shuffling activations: the same function to 3e-7, 9% faster per step.

05Equal compute: the results

Each pair gets the same seconds on the same GPU, each model at its best learning rate:

Equal compute (same training time)Equal steps
-0.04-0.020.00+0.02+0.04+0.06hybrid better ↑Llama better ↓13M27M48M90M
48M, first try (half attention)0.1B, hybrid at lr 1.5e-3
The hybrid's lead over a same-size Llama, by model size. Above zero, the hybrid is better. Equal steps: the same 3,000 steps for both. Equal compute: the same training time on the same GPU, so the faster model takes more steps. At 48M the first equal-compute try lost (hollow point); two attention layers per TONN layer won it back. Point at a point for the numbers.
  • 13M, 600 seconds: a tie (1.527 each). Llama is slightly ahead on prose, the hybrid on code.
  • 27M, 1,000 seconds: the hybrid wins by 0.006, on both seeds.
  • 48M, 2,000 seconds: the first try lost by 0.020. Llama got 23% more steps in the same time. At equal steps the hybrid still won, so the loss came from speed, not from the architecture fading.
  • 48M with two attention layers per TONN layer: the hybrid wins again, 1.445 against 1.452, on both seeds and every kind of text. More attention layers make it faster. The remaining TONN layers still carry the gain.

This 2:1 layout is now the default: every third layer is TONN. At the size we plan to train, it costs 1.01 times Llama per token.

06At 0.1B

The real model: width 640, 14 layers, TONN at layers 0, 3, 6, 9 and 12, about 90M parameters with the 32k tokenizer. Both models get 3,000 seconds on the Mac.

ExeaLM hybridLlama
4.05.06.07.08.005001000150020002500training time, secondsExeaLM hybrid 3.835Llama 3.963
Training loss against wall-clock time at 0.1B, seed 0, both models at learning rate 1e-3. Llama takes 11% more steps in the same 3,000 seconds; the hybrid still ends lower. Point at the chart to compare the two at any moment.
Same 3,000 secondsStepsEnglishJapanese wikiJapanese webCodeMean
ExeaLM hybrid, lr 1e-33,1521.3731.4511.3391.6921.464
ExeaLM hybrid, lr 1.5e-33,1511.3861.4621.3471.7021.474
Llama, lr 1e-33,4911.3921.4681.3521.7131.481

At the same learning rate as Llama, the hybrid leads by 0.017 bits per byte. It is ahead on every kind of text.

07What this does not show

  • One seed at 0.1B. Seed 1 is running: Llama now, the hybrid at 1e-3 next. A seed-1 hybrid at 1.5e-3 has finished at 1.482, behind seed 0's 1.474. It took 2,902 steps against 3,151, because the Mac ran slower during that run. There is no seed-1 Llama to compare it with yet. From 13M to 48M, every equal-compute result held on two seeds.
  • Llama ran at one learning rate at 0.1B (1e-3, its best rate at 27M and 48M). A better rate could narrow the gap.
  • Both models are badly under-trained: 50M to 57M tokens for 90M parameters, against the 100B tokens planned. Margins can shrink with longer training. At 13M, doubling the steps cut the lead from 0.014 to 0.006.
  • Speed was measured on one Mac GPU, uncompiled. torch.compile speeds up Llama by 20% and the hybrid by 8% on this GPU, so a compiled comparison would favour Llama today. A compile-friendly TONN layer is on the to-do list.
00 TONN010203 TONN040506 TONN070809 TONN101112 TONN13RESIDUAL STREAM14 LAYERS640 CHANNELS
14 layers. About 90M parameters.

Plates with circuits are TONN layers. The rest are attention. Point at a layer.

  • Memory
  • Short range
  • Running state
  • Gates and products
  • Attention
The 0.1B model, layer by layer. The five plates carrying circuits are TONN layers; the outlined plates are full attention layers. Point at a layer to open it.

The full ranking of every run is on the results page.

08

Who did the work

Avneh Bhatia@AvnehSBhatia
Ran the 49-run sweep and the 0.1B equal-compute runs on the Mac Studio. Wrote the prefix scan and the layer speedups.
09

Sources