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.