Tokenization → embeddings → causal Transformer → LM head → softmax → loss → backpropagation. No TensorFlow, no PyTorch, and no hidden autograd.
Repository: ** tiny-language-model-neuro-js**.
Most explanations of language models present correct formulas but hide the path between them inside a framework. I wanted the opposite: one small scenario where every scalar is visible and where the terminal clearly shows incorrect answers before learning and correct answers after it.
The project now has one command:
node src/train.js --generalize --adaptive-teach
It requires Node.js 18.19+ and has no dependencies.
The model is queried immediately after random initialization:
BEFORE TRAINING — random, usually wrong answers
> can human read ?
model: ? <unk> ...
expected: human can read. [WRONG]
> can fish swim ?
model: ? <unk> ...
expected: fish can swim. [WRONG]
> can cat read ?
model: ? <unk> ...
expected: cat cannot read. [WRONG]
After pre-training, SFT, and adaptive SFT, the same model produces:
FINAL ANSWERS AFTER ADAPTIVE SFT
> can human read ?
model: human can read. [CORRECT]
> can fish swim ?
model: fish can swim. [CORRECT]
> can bird fly ?
model: bird can fly. [CORRECT]
> can cat read ?
model: cat cannot read. [CORRECT]
Rehearsal controls preserved: 14/14.
Stable criterion reached 11 times in a row.
The initial text varies because initialization is random. The final acceptance criterion does not: all answers must be correct, every target token must have at least 95% probability, and the complete check must pass more than ten times consecutively.
The code previously contained several debug and training modes. They were useful while experimenting but obscured the main idea. The final version keeps one educational pipeline:
text → word tokenization → token IDs
→ token + position embeddings
→ two causal Transformer blocks
→ multi-head self-attention
→ two-hidden-layer FFN
→ LM head → softmax → next-token probabilities
→ cross-entropy → backpropagation → Adam
train.js
now reads as one story rather than a command-line framework.
Every number participating in learning is a Value
:
class Value {
constructor(data, children = [], backward = () => {}) {
this.data = data;
this.grad = 0;
this.children = children;
this._backward = backward;
}
}
For multiplication:
y = a × b
dy/da = b
dy/db = a
The operation stores these local derivatives. backward()
sorts the graph topologically and applies the chain rule from the final loss back to embeddings and weights.
The neuron formula is not hidden behind a tensor API:
output = activation(sum(input[i] × weight[i]) + bias)
Its implementation follows the formula:
forward(input) {
let output = sum(
input.map((value, i) => value.mul(this.weights[i]))
);
if (this.useBias) output = output.add(this.bias);
if (this.activation === 'relu') return output.relu();
return output;
}
A Linear
layer is just an array of neurons receiving the same input. This is slower than matrix multiplication but far easier to inspect.
Each token ID selects one trainable vector:
token representation = tokenEmbedding[id] + positionEmbedding[position]
Embeddings contain random values initially. They acquire useful relations only because gradients repeatedly change them in training contexts. No meaning
property is assigned to cat
, read
, or cannot
.
For every token:
Q = X × Wq
K = X × Wk
V = X × Wv
score = dot(Q, K) / sqrt(headSize)
attention = softmax(score)
output = attention × V
The implementation loops only while past <= position
. That is the causal mask: the model can attend to the current token and its history but never to a future target.
After attention, every token passes through a two-hidden-layer feed-forward network:
dModel → hidden ReLU → hidden ReLU → dModel
LayerNorm and residual paths preserve stable information flow around attention and FFN.
The most important code in the project is only a few lines:
function learnOneToken({ model, optimizer, input, targetId }) {
const loss = model.loss(input, targetId);
optimizer.zeroGrad();
loss.backward();
optimizer.step();
return loss.data;
}
The loss is ordinary next-token cross-entropy:
loss = -log(P(target | previous tokens))
If the correct token has low probability, loss is large. Backpropagation computes dLoss/dWeight
; Adam changes each parameter; the next forward pass gives a different distribution.
The tiny world contains 14 ability relations:
human can read .
fish can swim .
bird can fly .
dog cannot read .
The cat + read
relation is missing deliberately. Pre-training samples positions from this text and learns ordinary next-token prediction.
The same relations are converted into 42 prompt-answer examples:
can fish swim ?
is fish able to swim ?
does fish know how to swim ?
Only answer tokens contribute to SFT loss. The implementation visits every pair and every answer position on each epoch, making the training loop deterministic and readable.
The missing answer is represented only by target tokens:
['cat', 'cannot', 'read', '.']
Six question variants receive those targets. This is direct supervision: the model did not discover a zoological fact on its own. The teacher introduced the fact through loss, and backpropagation distributed that information across embeddings, attention, FFN, LayerNorm, and the LM head.
Why not stop after one correct answer? Because one generation can be fragile. The loop continues until every target token exceeds 95% probability and the whole evaluation succeeds 11 times in a row.
An early implementation trained only the six new cat prompts. It successfully learned the new answer and destroyed old behavior:
can human read ? → cat cannot read.
can fish swim ? → cat cannot read.
That is catastrophic forgetting in miniature. The fix is rehearsal: adaptive epochs also repeat the 14 older can ... ?
examples. The final criterion evaluates both new and old examples, so training cannot finish by overwriting everything with one response.
The command automatically creates:
logs/training-log.txt
It is a sequential ASCII diagram rather than a raw JSON dump. It includes every
forward/loss/backward/update event, followed by the complete matrices at three
checkpoints:
initial random matrices
|
v
matrices after pre-training + SFT
|
v
final matrices after adaptive SFT
For every transition, the log prints the AFTER matrix and its exact DELTA matrix.
Linear rows are named neuron[n]
, columns are named weight[n]
, and biases
are shown beside their neuron. It also points out the largest concrete change as
layer / neuron / weight: before -> after -> delta
.
The architecture and learning rule are real; the scale is intentionally tiny.
| This model | Production model |
|---|---|
| 24 word tokens | Large subword/byte vocabulary |
| 2,160 parameters | Millions or billions |
| Two Transformer blocks | Tens or hundreds |
| Scalar JavaScript graph | Batched tensor graph on accelerators |
| Small structured corpus | Massive curated datasets |
| Narrow trained behavior | Broad language and reasoning |
The project is not a GPT competitor. It is a causal language model reduced until the complete path fits in one repository and one mental model:
token → embedding → attention → FFN → probability
→ loss → gradient → updated weight → changed answer
That path is the point. Once it is visible, frameworks stop looking magical: they execute the same classes of operations at a scale and speed this scalar implementation deliberately avoids.
Repository: ** tiny-language-model-neuro-js**.
Author: ** Maksim Sekretov**.