microgpt in idiomatic MLPL
The same capability as microgpt.py, written with sw-MLPL's Model DSL

The faithful port follows microgpt.py line by line: nine hand-declared weight matrices, a hand-written forward pass, attention heads split with selector matrices, and a hand-written sampler. It is exact: with --rs-parity its output is byte-identical to microgpt-rs. But it reads like a transcription.

This variant keeps the capability and the data regime:

• the same 32,033 names;
• one name per training step, 1000 Adam steps with the same hyperparameters and linear learning-rate decay;
• a 1-layer, 4-head, 16-wide causal transformer;
• 20 samples at temperature 0.5.

It lets MLPL's built-ins do the work: the Model DSL for the network, adam over whole models, and sample for generation. The result is 48 lines of code instead of 230 (comments, blank lines and docstrings not counted), and it runs about 1.5x faster (see literate.md).

What changes on purpose, and why the numbers differ from the faithful run:

As in the faithful doc, every program block below runs in one :session and tangles into a script. scripts/check-literate.sh requires that script's output to equal microgpt-idiomatic.mlpl's reg-rs baseline.

MLPL's at / take accept a single index. Selecting many elements is a row gather on a one-column view; u:pick names it once, for any index shape. (A built-in gather(x, idx) is request #2 in sw-mlpl-requests.md.)

def u:pick(v, idx) {
  "Elements of rank-1 v at integer positions idx (any shape).";
  reshape(gather_rows(reshape(v, [len(v), 1]), idx), shape(idx))
}

The names file is lowercase ASCII, one name per line. Subtracting 97 maps a..z to 0..25, and every newline becomes BOS (26), so the whole corpus is one token stream. A mask and an arithmetic blend do the mapping, with no lookup table:

\[ s_i = \begin{cases} b_i - 97 & b_i \ne \text{newline} \\ 26 & b_i = \text{newline} \end{cases} \qquad\Longrightarrow\qquad s = (\mathrm{BOS},\ e,m,m,a,\ \mathrm{BOS},\ o,l,i,v,i,a,\ \mathrm{BOS}, \dots) \]

(The faithful variant discovers the vocabulary from the data instead; this shortcut assumes the ASCII names file, like microgpt-rs's a..z tests do.)

bytes = tokenize_bytes(load("input.txt"));
stream = concat(concat([26], (bytes - 97) * (bytes != 10) + 26 * (bytes == 10)), [26]);
V = 27; d = 16; T = 16;

bos_at = compress(stream == 26, range(len(stream)));
starts = u:pick(bos_at, range(len(bos_at) - 1));
lens = u:pick(bos_at, range(len(bos_at) - 1) + 1) - starts - 1;
print("num docs:", to_string(len(starts)));
print("vocab size:", to_string(V));

num docs: 32033
vocab size: 27

Names are the spans between consecutive BOS positions. The 1000 names training visits become the rows of one matrix in a single gather: row \(i\) holds \(\tau^{(i)}_{0..T}\) (BOS, the name, BOS, then whatever follows, which is never read), and \(n_i = \min(T, L_i + 1)\) positions are trained.

names = u:pick(shuffle(range(len(starts)), 42), range(1000));
rows = u:pick(stream, reshape(u:pick(starts, names), [1000, 1]) + reshape(range(T + 1), [1, T + 1]));
ns = u:pick(lens, names) + 1;
ns = ns - relu(ns - T);
expunge(["bytes", "stream", "bos_at", "starts", "lens", "names"]);

The expunge drops the 228k-token stream before training: every u: call copies the globals in this interpreter (upstream issue e).

The faithful port spends ~60 lines on 9 weight matrices, RMSNorm, linear, 4 selector-based heads and the residual forward pass. The Model DSL says it in one expression:

\[ z = \mathrm{Linear}_{V}\Big(B_2\big(B_1(\mathrm{RMS}(E_{tok}[\tau_t] + E_{pos}[t]))\big)\Big), \qquad t = 0 \dots n-1 \] \[ B_1(x) = x + \mathrm{CausalAttn}_{4\ \mathrm{heads}}(\mathrm{RMS}(x)), \qquad B_2(x) = x + \mathrm{Linear}_{d}\big(\mathrm{ReLU}(\mathrm{Linear}_{4d}(\mathrm{RMS}(x)))\big) \]

residual(b) is \(x \mapsto x + b(x)\). causal_attention(d, 4, seed) is multi-head attention with the causal mask built in (the faithful doc derives the mask and heads by hand). Positions are a second embed applied to range(n): a DSL chain has no position layer, so the two embeddings are added outside the chain. That is the moe-microscope idiom, and it is why this variant cannot use the KV cache (see the compact variant).

tok = embed(V, d, 1);
pos = embed(T, d, 2);
body = chain(rms_norm(d),
             residual(chain(rms_norm(d), causal_attention(d, 4, 3))),
             residual(chain(rms_norm(d), linear(d, 4 * d, 4), relu_layer(), linear(4 * d, d, 5))),
             linear(d, V, 6));
print("num params:", to_string(param_count(tok) + param_count(pos) + param_count(body)));

num params: 4299

The whole forward pass is one function, annotated with its equation:

@formula "z = body( E_tok[τ_t] + E_pos[t] ),   t = 0..n−1"
@ascii "z = body(E_tok[tau_t] + E_pos[t]),   t = 0..n-1"
def u:logits(toks) {
  "Next-token logits [n, V] for a token sequence.";
  apply(body, apply(tok, toks) + apply(pos, range(len(toks))))
}

print(annotations("u:logits").formula);
print("untrained loss on the first training name:", to_string(round(10000 * cross_entropy(u:logits(u:pick(take(rows, 0, 0), range(at(ns, 0)))), u:pick(take(rows, 0, 0), range(at(ns, 0)) + 1))) / 10000), " ln 27 =", to_string(round(10000 * log(27)) / 10000));
0

z = body( E_tok[τ_t] + E_pos[t] ),   t = 0..n−1
untrained loss on the first training name: 20.5445  ln 27 = 3.2958

The untrained loss is far above the uniform-guess value ln 27 = 3.30. The DSL's linear initializes much larger than microgpt's gauss(0, 0.08): it gives output std ~2.4 on unit input, which means weights around 0.6, while embed uses ~0.1. The first logits are therefore large and confidently wrong, and the early steps spend their updates shrinking them. Training recovers (below), but an init-scale option is on the request list.

\[ \mathcal{L} = -\frac{1}{n} \sum_{t=0}^{n-1} \log \mathrm{softmax}(z_t)[\tau_{t+1}], \qquad \theta \leftarrow \mathrm{Adam}_{\eta_k}(\theta, \nabla_\theta \mathcal{L}), \quad \eta_k = 0.01\,(1 - k/1000) \]

adam takes the three models as a list and updates every parameter inside them: there are no names to spell out, unlike the faithful variant's nine. train records each step's loss in last_losses.

train 1000 {
  n = at(ns, step);
  row = take(rows, 0, step);
  adam(cross_entropy(u:logits(u:pick(row, range(n))), u:pick(row, range(n) + 1)), [tok, pos, body], 0.01 * (1 - step / 1000), 0.85, 0.99, 0.00000001)
};
print("mean loss, steps 901-1000:", to_string(round(10000 * reduce_add(last_losses * (range(1000) >= 900)) / 100) / 10000));

mean loss, steps 901-1000: 2.4684

The faithful variant's last-100-step mean is 2.37. The difference comes from the architecture choices in the table above (biases, eps, the extra leading RMSNorm) and, most likely, the DSL's larger linear initialization.

sample(logits, temperature, seed) draws from \(\mathrm{softmax}(z/\tau_s)\) by inverse CDF on one seeded uniform:

\[ p = \mathrm{softmax}(z_{n-1} / 0.5), \qquad i^{*} = \min\Big\{ i : \sum_{j=0}^{i} p_j \ge u \Big\} \]

Each step recomputes the logits for the whole prefix, at most 16 tokens. The outer loop is a repeat over the 20 samples, and the inner while stops at BOS.

print("--- inference (new, hallucinated names) ---");
i = 0;
repeat 20 {
  seq = [26];
  done = 0;
  while (len(seq) <= T) * (1 - done) {
    next = sample(last_row(u:logits(seq)), 0.5, 1000 * i + len(seq));
    if next == 26 { done = 1 } else { seq = concat(seq, [next]) }
  };
  name = if len(seq) > 1 { decode_bytes(97 + u:pick(seq, range(len(seq) - 1) + 1)) } else { "" };
  print(str_concat("sample ", str_concat(to_string(i + 1), str_concat(": ", name))));
  i = i + 1
};
i

--- inference (new, hallucinated names) ---
sample 1: jiafini
sample 2: tamesin
sample 3: jali
sample 4: kair
sample 5: tasen
sample 6: maleri
sample 7: onia
sample 8: airka
sample 9: jamio
sample 10: zarilon
sample 11: dlanyn
sample 12: jaral
sample 13: meliw
sample 14: jarel
sample 15: korli
sample 16: aryle
sample 17: tapal
sample 18: kalia
sample 19: iiiel
sample 20: deidei

• Gained:
  • The network is one expression. The nine weight names, the selector-matrix heads, the mask and the forward function are gone.
  • adam takes models, not name lists.
  • The data pipeline is two gathers.
  • Runs faster: the DSL layers are single native ops.
• Lost:
  • Exact correspondence with microgpt.py's equations: biases, \(\epsilon = 10^{-8}\) and a leading RMSNorm are the DSL's choices.
  • Parameter count parity (4,299 vs 4,192).
  • Byte parity with microgpt-rs.
  • The KV cache, because positions live outside the chain.
• Wished for (see sw-mlpl-requests.md #1, #2, #10):
  • linear(..., {bias: 0, std: 0.08}) and rms_norm(d, {eps: 1e-5}) to match the paper and its initialization exactly;
  • a position layer inside chain that gen_state can cache;
  • gather / slice to replace u:pick;
  • format to replace the str_concat nest.