File size: 7,959 Bytes
8fcb6eb 6974e24 8fcb6eb 6974e24 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 | ---
license: mit
pipeline_tag: text-generation
tags:
- esp32
- esp32-s3
- microcontroller
- embedded
- edge
- per-layer-embeddings
- tinystories
- int4
datasets:
- roneneldan/TinyStories
language:
- en
---
# esp32-ai-tinystories
A 28.9M-parameter language model that runs entirely offline on an ESP32-S3
microcontroller, generating text at 9.88 tokens/second.
This is the TinyStories model from the
[esp32-ai](https://github.com/slvDev/esp32-ai) project, a demonstration of
**Per-Layer Embeddings (PLE)** on a microcontroller. It is not a
general-purpose chat model.
It is also not a `transformers` model. It is a raw binary for a small C
inference runtime that runs on the device.
## What it does
Continues a story in the style of
[TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories): simple
English at roughly a 3 to 4 year old's vocabulary. Given `Once upon a time`, it
writes a short children's story one token at a time, on the device.
It cannot answer questions, follow instructions, or hold a conversation. It was
trained from scratch on story text and nothing else.
## Why it fits on a microcontroller
The model is 28.9M parameters but only 556K of them are dense transformer core.
The rest is a Per-Layer Embedding table, read one row per token straight from
memory-mapped flash.
| component | params | share | lives in |
|---|---|---|---|
| PLE table | 25,165,824 | 87% | flash, memory-mapped |
| token embedding, also the tied head | 3,145,728 | 11% | flash, staged to PSRAM |
| dense transformer core | 556,416 | 2% | flash, staged to PSRAM |
| **total** | **28,867,968** | | **14.9 MB at int4** |
The constraint this addresses is fast memory, not total memory. The board has
8 MB of PSRAM, so size alone is not the problem, but the ESP32-S3 has only
512 KB of internal SRAM. A conventional model of this width would spend that
budget on embedding tables it reads once per token. PLE moves those to flash and
leaves the fast pool for what is read constantly.
## Architecture
```
architecture PLE
format version 1, TIED_HEAD
input vocab 32,768 stored embedding and PLE table rows
output vocab 25,353 logits the model produces
d_model 96
layers 6
heads 4
ffn_hidden 66
ple_dim 128
seq_len 256
rope_theta 10000.0
weights int4, group size 128
```
The two vocabulary sizes differ on purpose. The embedding and PLE table store
32,768 padded rows, but the tokenizer has 25,353 entries, so only those can ever
be produced or decoded. `TIED_HEAD` means the output head is the **first 25,353
rows** of the token embedding; tying does not require the row counts to match.
The header states both, so the runtime does not have to be told separately how
many logits to score.
Parameter counts above include all 32,768 stored rows, because the binary
stores them.
## Runtime placement
| tier | holds |
|---|---|
| flash, memory-mapped | PLE table and token embedding |
| PSRAM | per-position core and head, staged to int8 at boot; KV cache; logits (99 KiB) |
| SRAM | float scratch buffers and RMSNorm vectors, 29,320 B |
Activations are quantized to int8 for each staged matvec. The head is split
across both LX7 cores and is PSRAM-bandwidth-bound. int8 activations cost
+0.0003 nats of validation cross-entropy over 32,768 predictions
(2.4793 to 2.4796, perplexity 11.93 / 11.94).
## Result it demonstrates
Against a same-core, SRAM-fitting baseline at equal core parameters:
- PLE wins by 0.098 nats, 2 seeds, +/-0.006, roughly 16x the seed noise
- perplexity 12.58 to 11.41
- the gain survives 4-bit post-training quantization, 2 seeds
Full ablations, including the vocab-4096 control where the edge shrinks to
+0.025 nats, are in
[RESULTS.md](https://github.com/slvDev/esp32-ai/blob/main/RESULTS.md).
## Measured speed
| | value |
|---|---:|
| compute | **94.9 ms/token** |
| attached serial | **9.88 tok/s** |
Measured on the board with the runtime described above: 44 staged tensors,
29,320 B managed SRAM, 4.19 MB PSRAM, compiled at `-O3`.
## Files
| file | what it is |
|---|---|
| `model.bin` | int4 weights and header, flashed to the `model` partition |
| `tokenizer.json` | canonical 25,353-entry BPE, trained on the same TinyStories slice |
| `metadata.json` | architecture, parameters, runtime placement, SHA-256 of the model and tokenizer |
| `LICENSE` | MIT |
Verify a download before trusting it:
```bash
shasum -a 256 model.bin
# 1d8326c05c383ccfa615f5455575802817cb453dbc7ab28875d41a9dbb45477e
```
The firmware also prints an FNV-1a fingerprint of the mapped image at boot,
`a9bdd778`, which must match `device_fingerprint_fnv1a` in `metadata.json`.
The firmware's `vocab.h` is generated from `tokenizer.json` by the source
repository, so it is not distributed here.
## Verification
Reference logits are not shipped in this bundle. Verification lives with the
runtime, in the source repository, and covers two distinct things:
- `runtime/host_verify/verify.c` against `golden.txt` checks the **exact int4,
float-activation** path against PyTorch, to 1e-5.
- `runtime/host_verify/staging_verify.c` checks int8 weight staging, scale alignment,
ranged matvec equivalence, platform hook dispatch, header validation and the
untied-head format branch.
The device path enables int8 activations and is therefore **not bit-identical to
the host golden**. It is validated separately for output quality and throughput,
by the perplexity figure above and by on-device measurement.
## Usage
These weights are not usable on their own. The firmware also needs a decode
header generated from `tokenizer.json`, and it has to be compiled and flashed
alongside the model. The
[esp32-ai repository](https://github.com/slvDev/esp32-ai) does both steps:
```bash
scripts/fetch_model.sh tinystories # downloads and verifies these files
scripts/deploy.sh tinystories # generates the header, runs host gates, compiles, flashes
```
`fetch_model.sh` checks the assets above against a SHA-256 and byte size pinned
in the script, and cross-checks `metadata.json` against those same pins. It
installs nothing unless every check passes. `deploy.sh` never reaches the
network. Use `deploy.sh` rather than writing `model.bin` by hand: it regenerates
the decode table the firmware compiles against, and writes both the model and the
firmware.
## Training data
The first 300 MB of
[roneneldan/TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories).
The tokenizer was trained on that same slice. **The dataset is not redistributed
in this repository.**
Reproducibility is approximate. The preparation script downloads from the
dataset's `main` branch without pinning a revision and records no hash of the
raw slice, so a re-run reproduces the method rather than the same bytes.
Training sets no determinism flags, so retraining yields an equivalent model
rather than this file.
The training checkpoint, 110 MB, is not distributed. The deployable binary plus
the recipe is the public contract.
## Limitations
- Simple children's-story English only. No instruction following, no question
answering, no chat.
- Will produce fluent nonsense outside its distribution. 556K dense parameters
do not store facts.
- 256-token context, greedy decoding.
- Not bit-reproducible from the recipe, see above.
## License
| what | license |
|---|---|
| model weights (`model.bin`) | **MIT** |
| tokenizer (`tokenizer.json`) | **MIT** |
| training dataset (TinyStories) | CDLA-Sharing-1.0, **not redistributed here** |
TinyStories is licensed under CDLA-Sharing-1.0 and is not redistributed here.
The model weights and tokenizer are released under MIT. They were trained from
scratch and contain no third-party weights.
## Credits
The PLE design is reproduced from Google's published Gemma 3n Per-Layer
Embeddings work. No model, checkpoint or method here derives from it beyond the published
description.
|