MLP Language Model

hard · language-models, mlp, context-window

MLP Language Model

Build a fixed-context MLP language model. The model predicts the next token using a window of previous tokens.

Model

Token embedding: (V, D)
Context length: T (block size)
Flattened input size: T * D
MLP: T*D -> hidden -> V

Tasks

1) `MLPLanguageModel`

class MLPLanguageModel(Module):
    def __init__(self, vocab_size: int, block_size: int, embed_dim: int, hidden_dim: int):
        pass

    def forward(self, token_ids: List[int]) -> List[Value]:
        # token_ids length == block_size
        # return logits of length vocab_size
        pass

2) `cross_entropy_loss(logits, target)`

Softmax over logits
Return -log(p_target)

3) `train_step(model, xs, ys, lr)`

xs: list of contexts (each length block_size)
ys: list of target token IDs
Return mean loss value

4) `generate(model, start_ids, max_new_tokens, temperature=1.0)`

Use a rolling context window of size block_size
Sample next tokens autoregressively

Notes

Use Value operations to keep gradients intact.
block_size must be consistent across training and generation.

from __future__ import annotations
import math
import random
from typing import List, Iterator, Tuple, Callable, Set

class Value:
    """Scalar Value with autograd support."""

def __init__(self, data: float, _children: Tuple["Value", ...] = (), _op: str = "") -> None:
        self.data = float(data)
        self.grad = 0.0
        self._backward: Callable[[], None] = lambda: None
        self._prev: Set["Value"] = set(_children)
        self._op = _op

def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), "+")
        def _backward():
            self.grad += out.grad
            other.grad += out.grad
        out._backward = _backward
        return out

def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), "*")
        def _backward():
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out._backward = _backward
        return out

def __pow__(self, n):
        out = Value(self.data ** n, (self,), f"**{n}")
        def _backward():
            self.grad += n * (self.data ** (n - 1)) * out.grad
        out._backward = _backward
        return out

def __neg__(self):
        return self * -1

def __sub__(self, other):
        return self + (-other)

def __truediv__(self, other):
        return self * (other ** -1)

def __radd__(self, other):
        return self + other

def __rmul__(self, other):
        return self * other

def exp(self):
        out = Value(math.exp(self.data), (self,), "exp")
        def _backward():
            self.grad += out.data * out.grad
        out._backward = _backward
        return out

def log(self):
        out = Value(math.log(max(self.data, 1e-10)), (self,), "log")
        def _backward():
            self.grad += (1.0 / max(self.data, 1e-10)) * out.grad
        out._backward = _backward
        return out

def relu(self):
        out = Value(max(0.0, self.data), (self,), "relu")
        def _backward():
            self.grad += (1.0 if self.data > 0 else 0.0) * out.grad
        out._backward = _backward
        return out

def backward(self):
        topo = []
        visited = set()
        def build_topo(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev:
                    build_topo(child)
                topo.append(v)
        build_topo(self)
        self.grad = 1.0
        for node in reversed(topo):
            node._backward()

class Module:
    """Module base class."""

def __init__(self) -> None:
        object.__setattr__(self, "_parameters", {})
        object.__setattr__(self, "_modules", {})
        object.__setattr__(self, "training", True)

def __setattr__(self, name: str, value) -> None:
        if isinstance(value, Value):
            self._parameters[name] = value
        elif isinstance(value, Module):
            self._modules[name] = value
        elif isinstance(value, (list, tuple)):
            self._register_list(name, value)
        object.__setattr__(self, name, value)

def _register_list(self, name: str, values) -> None:
        for i, v in enumerate(values):
            if isinstance(v, Value):
                self._parameters[f"{name}.{i}"] = v
            elif isinstance(v, Module):
                self._modules[f"{name}.{i}"] = v
            elif isinstance(v, (list, tuple)):
                self._register_list(f"{name}.{i}", v)

def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

def forward(self, *args, **kwargs):
        raise NotImplementedError()

def parameters(self) -> List[Value]:
        params = list(self._parameters.values())
        for module in self._modules.values():
            params.extend(module.parameters())
        return params

class Linear(Module):
    """Fully connected layer."""

def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
        super().__init__()
        k = 1.0 / (in_features ** 0.5)
        self.weight = [
            [Value(random.uniform(-k, k)) for _ in range(in_features)]
            for _ in range(out_features)
        ]
        self.bias = [Value(0.0) for _ in range(out_features)] if bias else None
        self.out_features = out_features

def forward(self, x: List[Value]) -> List[Value]:
        out = []
        for i in range(self.out_features):
            val = sum(w * xi for w, xi in zip(self.weight[i], x))
            if self.bias is not None:
                val = val + self.bias[i]
            out.append(val)
        return out

class Embedding(Module):
    """Embedding lookup table."""

def __init__(self, num_embeddings: int, embedding_dim: int) -> None:
        super().__init__()
        k = 1.0 / (embedding_dim ** 0.5)
        self.weight = [
            [Value(random.uniform(-k, k)) for _ in range(embedding_dim)]
            for _ in range(num_embeddings)
        ]
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim

def forward(self, indices: List[int]) -> List[List[Value]]:
        return [self.weight[i] for i in indices]

def softmax(scores: List[Value]) -> List[Value]:
    """TODO: Numerically stable softmax."""
    raise NotImplementedError()

def cross_entropy_loss(logits: List[Value], target: int) -> Value:
    """TODO: Cross-entropy for a single example."""
    raise NotImplementedError()

class MLPLanguageModel(Module):
    """TODO: Implement fixed-context MLP language model."""

def __init__(self, vocab_size: int, block_size: int, embed_dim: int, hidden_dim: int) -> None:
        raise NotImplementedError()

def forward(self, token_ids: List[int]) -> List[Value]:
        raise NotImplementedError()

def train_step(model: MLPLanguageModel, xs: List[List[int]], ys: List[int], lr: float) -> float:
    """TODO: One SGD step over a batch; return mean loss."""
    raise NotImplementedError()

def generate(
    model: MLPLanguageModel,
    start_ids: List[int],
    max_new_tokens: int,
    temperature: float = 1.0,
) -> List[int]:
    """TODO: Autoregressively sample tokens with a rolling window."""
    raise NotImplementedError()

from __future__ import annotations
import math
import random
from typing import List, Iterator, Tuple, Callable, Set

class Value:
    """Scalar Value with autograd support."""

def __neg__(self):
        return self * -1

def __sub__(self, other):
        return self + (-other)

def __truediv__(self, other):
        return self * (other ** -1)

def __radd__(self, other):
        return self + other

def __rmul__(self, other):
        return self * other

def exp(self):
        out = Value(math.exp(self.data), (self,), "exp")
        def _backward():
            self.grad += out.data * out.grad
        out._backward = _backward
        return out

class Module:
    """Module base class."""

def __init__(self) -> None:
        object.__setattr__(self, "_parameters", {})
        object.__setattr__(self, "_modules", {})
        object.__setattr__(self, "training", True)

def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

def forward(self, *args, **kwargs):
        raise NotImplementedError()

class Linear(Module):
    """Fully connected layer."""

class Embedding(Module):
    """Embedding lookup table."""

def forward(self, indices: List[int]) -> List[List[Value]]:
        return [self.weight[i] for i in indices]

def softmax(scores: List[Value]) -> List[Value]:
    """TODO: Numerically stable softmax."""
    raise NotImplementedError()

def cross_entropy_loss(logits: List[Value], target: int) -> Value:
    """TODO: Cross-entropy for a single example."""
    raise NotImplementedError()

class MLPLanguageModel(Module):
    """TODO: Implement fixed-context MLP language model."""

def __init__(self, vocab_size: int, block_size: int, embed_dim: int, hidden_dim: int) -> None:
        raise NotImplementedError()

def forward(self, token_ids: List[int]) -> List[Value]:
        raise NotImplementedError()

def train_step(model: MLPLanguageModel, xs: List[List[int]], ys: List[int], lr: float) -> float:
    """TODO: One SGD step over a batch; return mean loss."""
    raise NotImplementedError()

Auto-advance

Run tests to see results

No issues detected

MLP Language Model

Model

Tasks

1) MLPLanguageModel

2) cross_entropy_loss(logits, target)

3) train_step(model, xs, ys, lr)

4) generate(model, start_ids, max_new_tokens, temperature=1.0)

Notes

1) `MLPLanguageModel`

2) `cross_entropy_loss(logits, target)`

3) `train_step(model, xs, ys, lr)`

4) `generate(model, start_ids, max_new_tokens, temperature=1.0)`