How to train a Small Language Model
We will train a 19.4-million-parameter Transformer with TensorFlow. It learns to predict the next byte from short English stories. After training, the model can generate text in a browser using TensorFlow.js.
The implementation uses Keras building blocks for attention without using a ready-made Transformer implementation.
View the complete Python and browser code in a new tab.
In this article
01Training setup
The main training run used a single A100 on Verda.
We use the following hyperparameters:
| Setting | Value |
|---|---|
| Vocabulary | 256 byte values |
| Context length | 512 bytes |
| Transformer blocks | 6 |
| Model width | 512 |
| Attention heads | 8 |
| Dimension per head | 64 |
| Feed-forward width | 2,048 |
| Parameters | 19,439,872 |
Training uses batches of 32 sequences and float32 weights. The optimizer is AdamW with a learning rate of 0.0003, weight decay of 0.01 and global gradient clipping at 1.0.
02Prepare the data
We use SimpleStories, a corpus of over two million synthetic English stories generated using larger language models. The project's paper describes its construction. The available story text in our corpus snapshot amounted to about 2.62 GB.
Hugging Face Datasets provides streaming access to the stories. We split them by a hash of the normalized text, holding out approximately 1% for validation. The rest supplies the training examples. This keeps a story in the same partition across runs. We evaluate a fixed 1,000-story subset of the held-out partition.
We then use tf.data to encode the stories, append separators, concatenate their bytes and form fixed-length chunks. The text cache uses RAM:
texts = tf.data.Dataset.from_generator(
lambda: (row["story"] for row in training_stories),
output_signature=tf.TensorSpec((), tf.string),
).cache()
data = texts.map(
lambda text: tf.io.decode_raw(text + "\n\n", tf.uint8)
)
data = data.unbatch().batch(513, drop_remainder=True)
data = data.shuffle(4096, seed=42).batch(32, drop_remainder=True)
training = data.map(
lambda chunk: (chunk[:, :-1], chunk[:, 1:])
).prefetch(tf.data.AUTOTUNE)Each 513-byte chunk supplies 512 input bytes and 512 targets, offset by one position. The model predicts the byte following each input position. Chunks can contain part of a story or cross a separator between stories.
The model's vocabulary is simply the 256 possible byte values. The byte values serve directly as input IDs, so this experiment needs no trained tokenizer. Context length and generated length are consequently measured in bytes rather than words or subword tokens.
03Define the Transformer
Each input byte is mapped to a learned 512-component embedding. We add a learned position embedding, then apply six Transformer blocks. Each block contains causal multi-head self-attention and a feed-forward network, with normalization before each operation and residual additions around them.
The block's forward pass is:
def call(self, x, training=False):
normalized = self.norm1(x)
x = x + self.attention(
normalized, normalized,
use_causal_mask=True, training=training)
return x + self.contract(self.expand(self.norm2(x)))expand is a dense layer with 2,048 outputs and a GELU activation. contract projects back to 512 components. The attention layer has eight heads, each with a dimension of 64. Both normalization layers use an epsilon of 0.00001.
After the last block, a final normalization and dense layer produce 256 logits for each position. The input embedding and output projection have separate weights.
The causal mask prevents an input position from attending to later positions. This matters during training: the entire chunk is present, including bytes that are targets for earlier positions. Masking lets us calculate predictions across the sequence in parallel while preserving the next-byte task.
04Train and validate
Sparse categorical cross-entropy compares the logits with the target byte IDs. TensorFlow differentiates the loss through the model, and AdamW updates the weights.
The core training call is:
optimizer = tf.keras.optimizers.AdamW(
learning_rate=0.0003,
weight_decay=0.01,
global_clipnorm=1.0,
)
optimizer.exclude_from_weight_decay(var_names=["bias", "beta", "gamma"])
model.compile(
optimizer=optimizer,
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
)
model.fit(
training.repeat(),
validation_data=validation,
steps_per_epoch=500,
epochs=40,
shuffle=False,
callbacks=[Progress(checkpoints)],
)Here, each reported epoch is a block of 500 updates; it is not a complete pass through the corpus. Forty such blocks give a target of 20,000 updates. A callback saves model and optimizer checkpoints and prints a generated sample at the end of each block. Keeping the optimizer state allows training to resume with its accumulated moments.
Our completed small-model checkpoint reached a validation loss of 0.622 nats per byte, approximately 0.898 bits per byte. We use held-out text for this measurement and inspect generated continuations separately.
For example, the model completed this opening, shown in bold:
Far far away, he found a letter in the corner of the park.
05Generate text
To generate a continuation, encode the prompt as bytes and run the model. The logits at the last position give scores for the next byte. Divide them by temperature, sample from the resulting distribution, append the selected byte and repeat.
Lower temperature concentrates sampling on the higher-probability choices. A fixed seed makes the random-number sequence repeatable, which is useful when comparing prompts or checkpoints.
The browser implementation processes the prompt together in one prefill pass. It then retains each layer's keys and values in a KV cache. Subsequent generation steps process one new byte against that cached state instead of recalculating the whole prefix.
For this architecture, a cache covering 512 positions occupies 12 MiB in float32. The cache belongs to the current passage; the model weights are shared across generations. The demo stops when prompt and output reach 512 bytes in total.
Try the model. The controls expose the prompt, temperature, seed and output length, together with prefill time, generation rate and KV-cache size.
06Run it in the browser
The original float32 weights occupy 77.8 MB. For download, we quantize each tensor to unsigned eight-bit values and store its minimum and scale. TensorFlow.js reconstructs the float32 values when loading the model.
The export code rounds each tensor’s weights to one of 256 values and records the information needed to reconstruct them:
array = value.numpy().astype("<f4")
minimum = float(array.min())
scale = (float(array.max()) - minimum) / 255 or 1.0
codes = np.clip(
np.rint((array - minimum) / scale), 0, 255
).astype(np.uint8)
stream.write(codes.tobytes())
specs.append({
"name": name,
"shape": list(array.shape),
"dtype": "float32",
"quantization": {
"dtype": "uint8",
"min": minimum,
"scale": scale,
},
})This reduces the weight payload to 19.4 MB, or 16.15 MB with gzip. On our fixed held-out comparison, the eight-bit version increased loss from 0.92279 to 0.92308 bits per byte. This is a separate evaluation sample from the training run's validation set.
The download reduction does not reduce the resident float32 weight data: it still occupies about 78 MB. Runtime buffers, intermediate tensors and the KV cache use additional memory.
Generation runs in a worker through TensorFlow.js, using WebGPU where an appropriate adapter is available or WebAssembly otherwise. Prompts and generated text stay on the device.
In the browser, TensorFlow.js decodes the binary using the tensor shapes, minima and scales in the manifest:
const config = await fetch("model/model-int8.json")
.then(response => response.json());
const buffer = await fetch(`model/${config.weight_file}`)
.then(response => response.arrayBuffer());
const weights = tf.io.decodeWeights(buffer, config.weights);
const model = new TinyLM(config, weights);The inference code below is handwritten JavaScript using TensorFlow.js operations. It reproduces the trained Keras model; TensorFlow does not generate this code. Calling step() with the whole prompt performs prefill. Subsequent calls supply one new byte and the returned cache.
class TinyLM {
constructor(config, weights) {
this.config = config;
this.weights = weights;
}
norm(x, name) {
const {mean, variance} = tf.moments(x, -1, true);
return x.sub(mean).mul(variance.add(this.config.epsilon).rsqrt())
.mul(this.weights[`${name}.gamma`]).add(this.weights[`${name}.beta`]);
}
dense(x, name) {
return tf.matMul(x, this.weights[`${name}.kernel`]).add(this.weights[`${name}.bias`]);
}
step(bytes, cache = []) {
const {width, heads, layers, context} = this.config;
const offset = cache.length ? cache[0].key.shape[1] : 0;
const length = bytes.length;
if (!length || offset + length > context) throw new Error("Context limit reached");
return tf.tidy(() => {
const ids = tf.tensor1d(bytes, "int32");
let x = tf.gather(this.weights.tokens, ids)
.add(tf.gather(this.weights.positions, tf.range(offset, offset + length, 1, "int32")));
const next = [];
for (let i = 0; i < layers; i++) {
const name = `b${i}`;
const normalized = this.norm(x, `${name}.norm1`);
const project = part => this.dense(normalized, `${name}.${part}`)
.reshape([length, heads, width / heads]).transpose([1, 0, 2]);
const query = project("q");
let key = project("k"), value = project("v");
if (cache.length) {
key = tf.concat([cache[i].key, key], 1);
value = tf.concat([cache[i].value, value], 1);
}
next.push({key, value});
let scores = tf.matMul(query, key, false, true).mul(1 / Math.sqrt(width / heads));
if (length > 1) {
const future = tf.range(0, offset + length, 1, "int32").expandDims(0)
.greater(tf.range(offset, offset + length, 1, "int32").expandDims(1));
scores = scores.add(future.cast("float32").mul(-1e9));
}
const attention = tf.matMul(tf.softmax(scores, -1), value)
.transpose([1, 0, 2]).reshape([length, width]);
x = x.add(this.dense(attention, `${name}.o`));
const expanded = this.dense(this.norm(x, `${name}.norm2`), `${name}.expand`);
const gelu = expanded.mul(0.5).mul(tf.erf(expanded.div(Math.SQRT2)).add(1));
x = x.add(this.dense(gelu, `${name}.contract`));
}
const last = this.norm(x.slice([length - 1, 0], [1, width]), "norm");
return {logits: this.dense(last, "head").reshape([256]), cache: next};
});
}
}
self.TinyLM = TinyLM;The returned logits supply the next-byte distribution. The caller samples a byte, feeds it back with the cache, and disposes the previous logits and cache after each step. tf.tidy() releases intermediate tensors while retaining the returned tensors.
07A larger experiment
A separate run doubles the width to 1,024 and the block count to twelve. It has 152.2 million parameters, approximately 7.8 times the small model's count. Context remains 512 bytes, with 16 attention heads and a feed-forward width of 4,096.
At 20,000 updates, the larger model reached a validation loss of 0.5477 nats per byte, compared with 0.6222 for the completed smaller model: a 12% reduction. This comparison measures next-byte prediction on the same validation setup.
You can select either model in the browser demo. The larger version uses the checkpoint at 20,000 updates. Its eight-bit weights download as 123.5 MB with gzip and expand to 608.8 MB of float32 weights, plus runtime buffers and up to 48 MiB of KV cache. The 19.4M model remains the default.