MLP Language Model: Turn a Context Window into One Prediction

One context produces one target

With IDs [0,1,2,0,2] and T=2:

[0,1]->2, [1,2]->0, [2,0]->2.

Batch targets have shape (B,), not (B,T).

Track the dimensions

(T,) IDs -> (T,D) embeddings -> (T*D,) flat vector

Then Linear(T*D,M) -> ReLU -> Linear(M,V).

Order is encoded by distinct concatenation slots.

Predict the hidden representation

E[0]=[1], E[1]=[2]. Context [0,1] flattens to [1,2].

First-layer rows [1,1], [1,-1]; zero bias.

Compute preactivations and ReLU outputs.

Complete the forward trace

Preactivations [3,-1]; hidden [3,0].

Output rows [1,0], [0,1], [-1,1] yield [3,0,-3].

Target 1 gives stable loss about 3.050946.

Reverse the context

[1,0] flattens to [2,1].

Hidden [3,1], logits [3,1,-2].

Order matters without positional embeddings because the first weight matrix has position-specific columns.

Count the capacity

Parameters: V*D + T*D*M + M + M*V + V.

For V=3, T=2, D=1, M=2: 18.

Increasing T changes the first layer, so the constructed model requires exactly T IDs.

Train on a mean loss

Compute one stable loss per context, average over B examples, backward once.

Duplicating the batch should preserve the gradient.

Finite-difference a weight away from a ReLU kink.

Generate with a complete window

Seed length must be at least T; this API has no padding token.

Use the last T IDs each step and preserve the full seed in the result.

Tokens falling outside the window cannot affect predictions through hidden state.